diff --git a/docs/c-api.rst b/docs/c-api.rst --- a/docs/c-api.rst +++ b/docs/c-api.rst @@ -301,6 +301,11 @@ be careful to check the return code of *both* the entry point itself, and :c:func:`futhark_context_sync`. +For the rules on entry points that consume their input, see +:ref:`api-consumption`. Note that even if a value has been consumed, +you must still manually free it. This is the only operation that is +permitted on a consumed value. + GPU --- @@ -432,7 +437,6 @@ value less than ``1``, then the runtime system will use one thread per detected core. - General guarantees ------------------ @@ -459,3 +463,8 @@ OpenCL) may perform file system operations during startup, and perhaps for caching GPU kernels in some cases. This is beyond Futhark's control. + +Violation the restrictions of consumption (see :ref:`api-consumption`) +can result in undefined behaviour. This does not matter for programs +whose entry points do not have unique parameter types +(:ref:`in-place-updates`). diff --git a/docs/man/futhark-literate.rst b/docs/man/futhark-literate.rst --- a/docs/man/futhark-literate.rst +++ b/docs/man/futhark-literate.rst @@ -83,7 +83,7 @@ DIRECTIVES ========== -A directive is a way to show the result of running a funtion. +A directive is a way to show the result of running a function. Depending on the directive, this can be as simple as printing the textual representation of the result, or as complex as running an external plotting program and referencing a generated image. @@ -205,7 +205,7 @@ Note that empty arrays must be written using the ``empty(t)`` notation, e.g. ``empty([0]i32)``. -Function applications are either of Futhark funtions or *builtin +Function applications are either of Futhark functions or *builtin functions*. The latter are prefixed with ``$`` and are magical (usually impure) functions that could not possibly be implemented in Futhark. The following builtins are supported: diff --git a/docs/server-protocol.rst b/docs/server-protocol.rst --- a/docs/server-protocol.rst +++ b/docs/server-protocol.rst @@ -52,9 +52,20 @@ produce outputs of defined types. The notion of transparent and opaque types are the same as in the C API: primitives and array of primitives are directly supported, and everything else is treated as -opaque. See also :ref:`valuemapping`. When printed, types -follow basic Futhark type syntax *without* sizes (e.g. ``[][]i32``). +opaque. See also :ref:`valuemapping`. When printed, types follow +basic Futhark type syntax *without* sizes (e.g. ``[][]i32``). +Uniqueness is not part of the types, but is indicated with an asterisk +in the ``inputs`` and ``outputs`` commands (see below). +Consumption and aliasing +------------------------ + +Since the server protocol closely models the C API, the same rules +apply to entry points that consume their arguments (see +:ref:`api-consumption`). In particular, consumed variables must still +be freed with the ``free`` command - but this is the only operation +that may be used on consumed variables. + Commands -------- @@ -92,13 +103,14 @@ .................. Print the types of inputs accepted by the given entry point, one per -line. +line. If the given input is consumed, the type is prefixed by `*`. ``outputs`` *entry* ................... Print the types of outputs produced by the given entry point, one per -line. +line. If the given output is guaranteed to be unique (does not alias +any inputs), the type is prefixed by `*`. ``clear`` ......... diff --git a/docs/usage.rst b/docs/usage.rst --- a/docs/usage.rst +++ b/docs/usage.rst @@ -340,6 +340,38 @@ rule does not apply when the entry point has been given a return type ascription that is not syntactically a tuple type. +.. _api-consumption: + +Consumption and Aliasing +~~~~~~~~~~~~~~~~~~~~~~~~ + +Futhark's support for :ref:`in-place-updates` has implications for the +generated API. Unfortunately, The type system of most languages +(e.g. C) is not rich enough to express the rules, so they are not +statically (or currently even dynamically checked). Since Futhark +will never infer a unique/consuming type for an entry point parameter, +this section can be ignored unless uniqueness annotations have been +manually added to the entry points parameter types. The rules are +essentially the same as in the language itself: + +1. Each entry point input parameter is either *consuming* or + *nonconsuming* (the default). This corresponds to unique and + nonunique types in the original Futhark program. A value passed + for a consuming parameter is considered *consumed*, now has an + unspecified value, and may never be used again. It must still be + manually freed, if applicable. + Further, any *aliases* of that value are also considered consumed + and may not be used. + +2. Each entry point output is either *unique* or *nonunique*. A + unique output has no aliases. A nonunique output aliases *every* + nonconsuming input parameter. + +Note that these distinctions are currently usually not visible in the +generated API, and so correct usage requires knowledge of the original +types in the Futhark function. The safest strategy is to not expose +unique types in entry points. + Generating C ^^^^^^^^^^^^ diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,7 +1,6 @@ cabal-version: 2.4 --- Run 'cabal2nix . >futhark.nix' after adding deps. name: futhark -version: 0.19.6 +version: 0.19.7 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to @@ -26,7 +25,7 @@ . <> -category: Language +category: Futhark homepage: https://futhark-lang.org bug-reports: https://github.com/diku-dk/futhark/issues maintainer: Troels Henriksen athas@sigkill.dk @@ -58,12 +57,14 @@ Futhark.Analysis.DataDependencies Futhark.Analysis.HORep.MapNest Futhark.Analysis.HORep.SOAC + Futhark.Analysis.Interference + Futhark.Analysis.LastUse Futhark.Analysis.Metrics Futhark.Analysis.Metrics.Type Futhark.Analysis.PrimExp - Futhark.Analysis.PrimExp.Parse Futhark.Analysis.PrimExp.Convert Futhark.Analysis.PrimExp.Generalize + Futhark.Analysis.PrimExp.Parse Futhark.Analysis.PrimExp.Simplify Futhark.Analysis.Rephrase Futhark.Analysis.SymbolTable @@ -97,8 +98,8 @@ Futhark.CodeGen.Backends.COpenCL.Boilerplate Futhark.CodeGen.Backends.GenericC Futhark.CodeGen.Backends.GenericC.CLI - Futhark.CodeGen.Backends.GenericC.Server Futhark.CodeGen.Backends.GenericC.Options + Futhark.CodeGen.Backends.GenericC.Server Futhark.CodeGen.Backends.GenericPython Futhark.CodeGen.Backends.GenericPython.AST Futhark.CodeGen.Backends.GenericPython.Definitions @@ -111,22 +112,22 @@ Futhark.CodeGen.Backends.SequentialPython Futhark.CodeGen.Backends.SimpleRep Futhark.CodeGen.ImpCode - Futhark.CodeGen.ImpCode.Kernels + Futhark.CodeGen.ImpCode.GPU Futhark.CodeGen.ImpCode.Multicore Futhark.CodeGen.ImpCode.OpenCL Futhark.CodeGen.ImpCode.Sequential Futhark.CodeGen.ImpGen Futhark.CodeGen.ImpGen.CUDA - Futhark.CodeGen.ImpGen.Kernels - Futhark.CodeGen.ImpGen.Kernels.Base - Futhark.CodeGen.ImpGen.Kernels.SegHist - Futhark.CodeGen.ImpGen.Kernels.SegMap - Futhark.CodeGen.ImpGen.Kernels.SegRed - Futhark.CodeGen.ImpGen.Kernels.SegScan - Futhark.CodeGen.ImpGen.Kernels.SegScan.SinglePass - Futhark.CodeGen.ImpGen.Kernels.SegScan.TwoPass - Futhark.CodeGen.ImpGen.Kernels.ToOpenCL - Futhark.CodeGen.ImpGen.Kernels.Transpose + Futhark.CodeGen.ImpGen.GPU + Futhark.CodeGen.ImpGen.GPU.Base + Futhark.CodeGen.ImpGen.GPU.SegHist + Futhark.CodeGen.ImpGen.GPU.SegMap + Futhark.CodeGen.ImpGen.GPU.SegRed + Futhark.CodeGen.ImpGen.GPU.SegScan + Futhark.CodeGen.ImpGen.GPU.SegScan.SinglePass + Futhark.CodeGen.ImpGen.GPU.SegScan.TwoPass + Futhark.CodeGen.ImpGen.GPU.ToOpenCL + Futhark.CodeGen.ImpGen.GPU.Transpose Futhark.CodeGen.ImpGen.Multicore Futhark.CodeGen.ImpGen.Multicore.Base Futhark.CodeGen.ImpGen.Multicore.SegHist @@ -148,19 +149,18 @@ Futhark.FreshNames Futhark.IR Futhark.IR.Aliases - Futhark.IR.Decorations - Futhark.IR.Parse - Futhark.IR.Kernels - Futhark.IR.Kernels.Kernel - Futhark.IR.Kernels.Simplify - Futhark.IR.Kernels.Sizes - Futhark.IR.KernelsMem + Futhark.IR.GPU + Futhark.IR.GPU.Kernel + Futhark.IR.GPU.Simplify + Futhark.IR.GPU.Sizes + Futhark.IR.GPUMem Futhark.IR.MC Futhark.IR.MC.Op Futhark.IR.MCMem Futhark.IR.Mem Futhark.IR.Mem.IxFun Futhark.IR.Mem.Simplify + Futhark.IR.Parse Futhark.IR.Pretty Futhark.IR.Primitive Futhark.IR.Primitive.Parse @@ -174,6 +174,7 @@ Futhark.IR.Prop.Scope Futhark.IR.Prop.TypeOf Futhark.IR.Prop.Types + Futhark.IR.Rep Futhark.IR.RetType Futhark.IR.SOACS Futhark.IR.SOACS.SOAC @@ -207,9 +208,11 @@ Futhark.Optimise.InPlaceLowering.LowerIntoStm Futhark.Optimise.InPlaceLowering.SubstituteIndices Futhark.Optimise.InliningDeadFun + Futhark.Optimise.ReuseAllocations + Futhark.Optimise.ReuseAllocations.GreedyColoring Futhark.Optimise.Simplify Futhark.Optimise.Simplify.Engine - Futhark.Optimise.Simplify.Lore + Futhark.Optimise.Simplify.Rep Futhark.Optimise.Simplify.Rule Futhark.Optimise.Simplify.Rules Futhark.Optimise.Simplify.Rules.BasicOp @@ -224,10 +227,10 @@ Futhark.Pass Futhark.Pass.ExpandAllocations Futhark.Pass.ExplicitAllocations - Futhark.Pass.ExplicitAllocations.Kernels + Futhark.Pass.ExplicitAllocations.GPU + Futhark.Pass.ExplicitAllocations.MC Futhark.Pass.ExplicitAllocations.SegOp Futhark.Pass.ExplicitAllocations.Seq - Futhark.Pass.ExplicitAllocations.MC Futhark.Pass.ExtractKernels Futhark.Pass.ExtractKernels.BlockedKernel Futhark.Pass.ExtractKernels.DistributeNests @@ -236,7 +239,7 @@ Futhark.Pass.ExtractKernels.Interchange Futhark.Pass.ExtractKernels.Intragroup Futhark.Pass.ExtractKernels.StreamKernel - Futhark.Pass.ExtractKernels.ToKernels + Futhark.Pass.ExtractKernels.ToGPU Futhark.Pass.ExtractMulticore Futhark.Pass.FirstOrderTransform Futhark.Pass.KernelBabysitting @@ -247,10 +250,8 @@ Futhark.Pkg.Solve Futhark.Pkg.Types Futhark.Script - Futhark.Server Futhark.Test Futhark.Test.Values - Futhark.Test.Values.Parser Futhark.Tools Futhark.Transform.CopyPropagate Futhark.Transform.FirstOrderTransform @@ -314,6 +315,8 @@ , file-embed >=0.0.9 , filepath >=1.4.1.1 , free >=4.12.4 + , futhark-data >= 1.0.0.1 + , futhark-server >= 1.1.0.0 , gitrev >=1.2.0 , hashable , haskeline @@ -372,6 +375,7 @@ Futhark.IR.PrimitiveTests Language.Futhark.CoreTests Language.Futhark.SyntaxTests + Futhark.Optimise.ReuseAllocations.GreedyColoringTests Paths_futhark hs-source-dirs: unittests diff --git a/rts/c/server.h b/rts/c/server.h --- a/rts/c/server.h +++ b/rts/c/server.h @@ -131,7 +131,9 @@ const char *name; entry_point_fn f; struct type **out_types; + bool *out_unique; struct type **in_types; + bool *in_unique; }; int entry_num_ins(struct entry_point *e) { @@ -459,6 +461,9 @@ int num_ins = entry_num_ins(e); for (int i = 0; i < num_ins; i++) { + if (e->in_unique[i]) { + putchar('*'); + } puts(e->in_types[i]->name); } } @@ -475,6 +480,9 @@ int num_outs = entry_num_outs(e); for (int i = 0; i < num_outs; i++) { + if (e->out_unique[i]) { + putchar('*'); + } puts(e->out_types[i]->name); } } diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs --- a/src/Futhark/Actions.hs +++ b/src/Futhark/Actions.hs @@ -29,12 +29,12 @@ import qualified Futhark.CodeGen.Backends.PyOpenCL as PyOpenCL import qualified Futhark.CodeGen.Backends.SequentialC as SequentialC import qualified Futhark.CodeGen.Backends.SequentialPython as SequentialPy -import qualified Futhark.CodeGen.ImpGen.Kernels as ImpGenKernels +import qualified Futhark.CodeGen.ImpGen.GPU as ImpGenGPU import qualified Futhark.CodeGen.ImpGen.Multicore as ImpGenMulticore import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGenSequential import Futhark.Compiler.CLI import Futhark.IR -import Futhark.IR.KernelsMem (KernelsMem) +import Futhark.IR.GPUMem (GPUMem) import Futhark.IR.MCMem (MCMem) import Futhark.IR.Prop.Aliases import Futhark.IR.SeqMem (SeqMem) @@ -46,7 +46,7 @@ import qualified System.Info -- | Print the result to stdout. -printAction :: ASTLore lore => Action lore +printAction :: ASTRep rep => Action rep printAction = Action { actionName = "Prettyprint", @@ -55,7 +55,7 @@ } -- | Print the result to stdout, alias annotations. -printAliasesAction :: (ASTLore lore, CanBeAliased (Op lore)) => Action lore +printAliasesAction :: (ASTRep rep, CanBeAliased (Op rep)) => Action rep printAliasesAction = Action { actionName = "Prettyprint", @@ -64,7 +64,7 @@ } -- | Print metrics about AST node counts to stdout. -metricsAction :: OpMetrics (Op lore) => Action lore +metricsAction :: OpMetrics (Op rep) => Action rep metricsAction = Action { actionName = "Compute metrics", @@ -82,12 +82,12 @@ } -- | Convert the program to GPU ImpCode and print it to stdout. -kernelImpCodeGenAction :: Action KernelsMem +kernelImpCodeGenAction :: Action GPUMem kernelImpCodeGenAction = Action { actionName = "Compile imperative kernels", actionDescription = "Translate program into imperative IL with kernels and write it on standard output.", - actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenKernels.compileProgOpenCL + actionProcedure = liftIO . putStrLn . pretty . snd <=< ImpGenGPU.compileProgOpenCL } -- | Convert the program to CPU multicore ImpCode and print it to stdout. @@ -173,7 +173,7 @@ runCC cpath outpath ["-O3", "-std=c99"] ["-lm"] -- | The @futhark opencl@ action. -compileOpenCLAction :: FutharkConfig -> CompilerMode -> FilePath -> Action KernelsMem +compileOpenCLAction :: FutharkConfig -> CompilerMode -> FilePath -> Action GPUMem compileOpenCLAction fcfg mode outpath = Action { actionName = "Compile to OpenCL", @@ -206,7 +206,7 @@ runCC cpath outpath ["-O", "-std=c99"] ("-lm" : extra_options) -- | The @futhark cuda@ action. -compileCUDAAction :: FutharkConfig -> CompilerMode -> FilePath -> Action KernelsMem +compileCUDAAction :: FutharkConfig -> CompilerMode -> FilePath -> Action GPUMem compileCUDAAction fcfg mode outpath = Action { actionName = "Compile to CUDA", @@ -291,7 +291,7 @@ actionProcedure = pythonCommon SequentialPy.compileProg fcfg mode outpath } -compilePyOpenCLAction :: FutharkConfig -> CompilerMode -> FilePath -> Action KernelsMem +compilePyOpenCLAction :: FutharkConfig -> CompilerMode -> FilePath -> Action GPUMem compilePyOpenCLAction fcfg mode outpath = Action { actionName = "Compile to PyOpenCL", diff --git a/src/Futhark/Analysis/Alias.hs b/src/Futhark/Analysis/Alias.hs --- a/src/Futhark/Analysis/Alias.hs +++ b/src/Futhark/Analysis/Alias.hs @@ -1,7 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} -- | Alias analysis of a full Futhark program. Takes as input a --- program with an arbitrary lore and produces one with aliases. This +-- program with an arbitrary rep and produces one with aliases. This -- module does not implement the aliasing logic itself, and derives -- its information from definitions in -- "Futhark.IR.Prop.Aliases" and @@ -26,37 +26,37 @@ -- | Perform alias analysis on a Futhark program. aliasAnalysis :: - (ASTLore lore, CanBeAliased (Op lore)) => - Prog lore -> - Prog (Aliases lore) + (ASTRep rep, CanBeAliased (Op rep)) => + Prog rep -> + Prog (Aliases rep) aliasAnalysis (Prog consts funs) = Prog (fst (analyseStms mempty consts)) (map analyseFun funs) analyseFun :: - (ASTLore lore, CanBeAliased (Op lore)) => - FunDef lore -> - FunDef (Aliases lore) + (ASTRep rep, CanBeAliased (Op rep)) => + FunDef rep -> + FunDef (Aliases rep) analyseFun (FunDef entry attrs fname restype params body) = FunDef entry attrs fname restype params body' where body' = analyseBody mempty body analyseBody :: - ( ASTLore lore, - CanBeAliased (Op lore) + ( ASTRep rep, + CanBeAliased (Op rep) ) => AliasTable -> - Body lore -> - Body (Aliases lore) -analyseBody atable (Body lore stms result) = + Body rep -> + Body (Aliases rep) +analyseBody atable (Body rep stms result) = let (stms', _atable') = analyseStms atable stms - in mkAliasedBody lore stms' result + in mkAliasedBody rep stms' result analyseStms :: - (ASTLore lore, CanBeAliased (Op lore)) => + (ASTRep rep, CanBeAliased (Op rep)) => AliasTable -> - Stms lore -> - (Stms (Aliases lore), AliasesAndConsumed) + Stms rep -> + (Stms (Aliases rep), AliasesAndConsumed) analyseStms orig_aliases = foldl' f (mempty, (orig_aliases, mempty)) . stmsToList where @@ -66,21 +66,21 @@ in (stms <> oneStm stm', atable') analyseStm :: - (ASTLore lore, CanBeAliased (Op lore)) => + (ASTRep rep, CanBeAliased (Op rep)) => AliasTable -> - Stm lore -> - Stm (Aliases lore) + Stm rep -> + Stm (Aliases rep) analyseStm aliases (Let pat (StmAux cs attrs dec) e) = let e' = analyseExp aliases e pat' = addAliasesToPattern pat e' - lore' = (AliasDec $ consumedInExp e', dec) - in Let pat' (StmAux cs attrs lore') e' + rep' = (AliasDec $ consumedInExp e', dec) + in Let pat' (StmAux cs attrs rep') e' analyseExp :: - (ASTLore lore, CanBeAliased (Op lore)) => + (ASTRep rep, CanBeAliased (Op rep)) => AliasTable -> - Exp lore -> - Exp (Aliases lore) + Exp rep -> + Exp (Aliases rep) -- Would be better to put this in a BranchType annotation, but that -- requires a lot of other work. analyseExp aliases (If cond tb fb dec) = @@ -115,10 +115,10 @@ } analyseLambda :: - (ASTLore lore, CanBeAliased (Op lore)) => + (ASTRep rep, CanBeAliased (Op rep)) => AliasTable -> - Lambda lore -> - Lambda (Aliases lore) + Lambda rep -> + Lambda (Aliases rep) analyseLambda aliases lam = let body = analyseBody aliases $ lambdaBody lam in lam diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs --- a/src/Futhark/Analysis/DataDependencies.hs +++ b/src/Futhark/Analysis/DataDependencies.hs @@ -18,13 +18,13 @@ type Dependencies = M.Map VName Names -- | Compute the data dependencies for an entire body. -dataDependencies :: ASTLore lore => Body lore -> Dependencies +dataDependencies :: ASTRep rep => Body rep -> Dependencies dataDependencies = dataDependencies' M.empty dataDependencies' :: - ASTLore lore => + ASTRep rep => Dependencies -> - Body lore -> + Body rep -> Dependencies dataDependencies' startdeps = foldl grow startdeps . bodyStms where diff --git a/src/Futhark/Analysis/HORep/MapNest.hs b/src/Futhark/Analysis/HORep/MapNest.hs --- a/src/Futhark/Analysis/HORep/MapNest.hs +++ b/src/Futhark/Analysis/HORep/MapNest.hs @@ -24,7 +24,7 @@ import qualified Futhark.IR.SOACS.SOAC as Futhark import Futhark.Transform.Substitute -data Nesting lore = Nesting +data Nesting rep = Nesting { nestingParamNames :: [VName], nestingResult :: [VName], nestingReturnType :: [Type], @@ -32,25 +32,25 @@ } deriving (Eq, Ord, Show) -data MapNest lore = MapNest SubExp (Lambda lore) [Nesting lore] [SOAC.Input] +data MapNest rep = MapNest SubExp (Lambda rep) [Nesting rep] [SOAC.Input] deriving (Show) -typeOf :: MapNest lore -> [Type] +typeOf :: MapNest rep -> [Type] typeOf (MapNest w lam [] _) = map (`arrayOfRow` w) $ lambdaReturnType lam typeOf (MapNest w _ (nest : _) _) = map (`arrayOfRow` w) $ nestingReturnType nest -params :: MapNest lore -> [VName] +params :: MapNest rep -> [VName] params (MapNest _ lam [] _) = map paramName $ lambdaParams lam params (MapNest _ _ (nest : _) _) = nestingParamNames nest -inputs :: MapNest lore -> [SOAC.Input] +inputs :: MapNest rep -> [SOAC.Input] inputs (MapNest _ _ _ inps) = inps -setInputs :: [SOAC.Input] -> MapNest lore -> MapNest lore +setInputs :: [SOAC.Input] -> MapNest rep -> MapNest rep setInputs [] (MapNest w body ns _) = MapNest w body ns [] setInputs (inp : inps) (MapNest _ body ns _) = MapNest w body ns' (inp : inps) where @@ -60,24 +60,24 @@ setDepth n nw = n {nestingWidth = nw} fromSOAC :: - ( Bindable lore, + ( Bindable rep, MonadFreshNames m, - LocalScope lore m, - Op lore ~ Futhark.SOAC lore + LocalScope rep m, + Op rep ~ Futhark.SOAC rep ) => - SOAC lore -> - m (Maybe (MapNest lore)) + SOAC rep -> + m (Maybe (MapNest rep)) fromSOAC = fromSOAC' mempty fromSOAC' :: - ( Bindable lore, + ( Bindable rep, MonadFreshNames m, - LocalScope lore m, - Op lore ~ Futhark.SOAC lore + LocalScope rep m, + Op rep ~ Futhark.SOAC rep ) => [Ident] -> - SOAC lore -> - m (Maybe (MapNest lore)) + SOAC rep -> + m (Maybe (MapNest rep)) fromSOAC' bound (SOAC.Screma w (SOAC.ScremaForm [] [] lam) inps) = do maybenest <- case ( stmsToList $ bodyStms $ lambdaBody lam, bodyResult $ lambdaBody lam @@ -142,13 +142,13 @@ toSOAC :: ( MonadFreshNames m, - HasScope lore m, - Bindable lore, - BinderOps lore, - Op lore ~ Futhark.SOAC lore + HasScope rep m, + Bindable rep, + BinderOps rep, + Op rep ~ Futhark.SOAC rep ) => - MapNest lore -> - m (SOAC lore) + MapNest rep -> + m (SOAC rep) toSOAC (MapNest w lam [] inps) = return $ SOAC.Screma w (Futhark.mapSOAC lam) inps toSOAC (MapNest w lam (Nesting npnames nres nrettype nw : ns) inps) = do diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs --- a/src/Futhark/Analysis/HORep/SOAC.hs +++ b/src/Futhark/Analysis/HORep/SOAC.hs @@ -226,7 +226,7 @@ -- an input transformation of an array variable. If so, return the -- variable and the transformation. Only 'Rearrange' and 'Reshape' -- are possible to express this way. -transformFromExp :: Certificates -> Exp lore -> Maybe (VName, ArrayTransform) +transformFromExp :: Certificates -> Exp rep -> Maybe (VName, ArrayTransform) transformFromExp cs (BasicOp (Futhark.Rearrange perm v)) = Just (v, Rearrange cs perm) transformFromExp cs (BasicOp (Futhark.Reshape shape v)) = @@ -374,11 +374,11 @@ addTransform (Rearrange mempty $ transposeIndex k n [0 .. inputRank inp -1]) inp -- | A definite representation of a SOAC expression. -data SOAC lore - = Stream SubExp (StreamForm lore) (Lambda lore) [SubExp] [Input] - | Scatter SubExp (Lambda lore) [Input] [(Shape, Int, VName)] - | Screma SubExp (ScremaForm lore) [Input] - | Hist SubExp [HistOp lore] (Lambda lore) [Input] +data SOAC rep + = Stream SubExp (StreamForm rep) (Lambda rep) [SubExp] [Input] + | Scatter SubExp (Lambda rep) [Input] [(Shape, Int, VName)] + | Screma SubExp (ScremaForm rep) [Input] + | Hist SubExp [HistOp rep] (Lambda rep) [Input] deriving (Eq, Show) instance PP.Pretty Input where @@ -395,21 +395,21 @@ f e (Replicate cs ne) = text "replicate" <> ppr cs <> PP.apply [ppr ne, e] -instance PrettyLore lore => PP.Pretty (SOAC lore) where +instance PrettyRep rep => PP.Pretty (SOAC rep) where ppr (Screma w form arrs) = Futhark.ppScrema w arrs form ppr (Hist len ops bucket_fun imgs) = Futhark.ppHist len ops bucket_fun imgs ppr soac = text $ show soac -- | Returns the inputs used in a SOAC. -inputs :: SOAC lore -> [Input] +inputs :: SOAC rep -> [Input] inputs (Stream _ _ _ _ arrs) = arrs inputs (Scatter _len _lam ivs _as) = ivs inputs (Screma _ _ arrs) = arrs inputs (Hist _ _ _ inps) = inps -- | Set the inputs to a SOAC. -setInputs :: [Input] -> SOAC lore -> SOAC lore +setInputs :: [Input] -> SOAC rep -> SOAC rep setInputs arrs (Stream w form lam nes _) = Stream (newWidth arrs w) form lam nes arrs setInputs arrs (Scatter w lam _ivs as) = @@ -424,14 +424,14 @@ newWidth (inp : _) _ = arraySize 0 $ inputType inp -- | The lambda used in a given SOAC. -lambda :: SOAC lore -> Lambda lore +lambda :: SOAC rep -> Lambda rep lambda (Stream _ _ lam _ _) = lam lambda (Scatter _len lam _ivs _as) = lam lambda (Screma _ (ScremaForm _ _ lam) _) = lam lambda (Hist _ _ lam _) = lam -- | Set the lambda used in the SOAC. -setLambda :: Lambda lore -> SOAC lore -> SOAC lore +setLambda :: Lambda rep -> SOAC rep -> SOAC rep setLambda lam (Stream w form _ nes arrs) = Stream w form lam nes arrs setLambda lam (Scatter len _lam ivs as) = @@ -442,7 +442,7 @@ Hist w ops lam inps -- | The return type of a SOAC. -typeOf :: SOAC lore -> [Type] +typeOf :: SOAC rep -> [Type] typeOf (Stream w _ lam nes _) = let accrtps = take (length nes) $ lambdaReturnType lam arrtps = @@ -464,7 +464,7 @@ -- | The "width" of a SOAC is the expected outer size of its array -- inputs _after_ input-transforms have been carried out. -width :: SOAC lore -> SubExp +width :: SOAC rep -> SubExp width (Stream w _ _ _ _) = w width (Scatter len _lam _ivs _as) = len width (Screma w _ _) = w @@ -472,16 +472,16 @@ -- | Convert a SOAC to the corresponding expression. toExp :: - (MonadBinder m, Op (Lore m) ~ Futhark.SOAC (Lore m)) => - SOAC (Lore m) -> - m (Exp (Lore m)) + (MonadBinder m, Op (Rep m) ~ Futhark.SOAC (Rep m)) => + SOAC (Rep m) -> + m (Exp (Rep m)) toExp soac = Op <$> toSOAC soac -- | Convert a SOAC to a Futhark-level SOAC. toSOAC :: MonadBinder m => - SOAC (Lore m) -> - m (Futhark.SOAC (Lore m)) + SOAC (Rep m) -> + m (Futhark.SOAC (Rep m)) toSOAC (Stream w form lam nes inps) = Futhark.Stream w <$> inputsToSubExps inps <*> pure form <*> pure nes <*> pure lam toSOAC (Scatter len lam ivs dests) = do @@ -503,9 +503,9 @@ -- representation, or a reason why the expression does not have the -- valid form. fromExp :: - (Op lore ~ Futhark.SOAC lore, HasScope lore m) => - Exp lore -> - m (Either NotSOAC (SOAC lore)) + (Op rep ~ Futhark.SOAC rep, HasScope rep m) => + Exp rep -> + m (Either NotSOAC (SOAC rep)) fromExp (Op (Futhark.Stream w as form nes lam)) = Right . Stream w form lam nes <$> traverse varInput as fromExp (Op (Futhark.Scatter len lam ivs as)) = @@ -520,9 +520,9 @@ -- Returns the Stream SOAC and the -- extra-accumulator body-result ident if any. soacToStream :: - (MonadFreshNames m, Bindable lore, Op lore ~ Futhark.SOAC lore) => - SOAC lore -> - m (SOAC lore, [Ident]) + (MonadFreshNames m, Bindable rep, Op rep ~ Futhark.SOAC rep) => + SOAC rep -> + m (SOAC rep, [Ident]) soacToStream soac = do chunk_param <- newParam "chunk" $ Prim int64 let chvar = Futhark.Var $ paramName chunk_param @@ -690,10 +690,10 @@ _ -> return (soac, []) where mkMapPlusAccLam :: - (MonadFreshNames m, Bindable lore) => + (MonadFreshNames m, Bindable rep) => [SubExp] -> - Lambda lore -> - m (Lambda lore) + Lambda rep -> + m (Lambda rep) mkMapPlusAccLam accs plus = do let (accpars, rempars) = splitAt (length accs) $ lambdaParams plus parbnds = @@ -715,10 +715,10 @@ renameLambda $ Lambda rempars newlambdy $ lambdaReturnType plus mkPlusBnds :: - (MonadFreshNames m, Bindable lore) => - Lambda lore -> + (MonadFreshNames m, Bindable rep) => + Lambda rep -> [SubExp] -> - m (Body lore) + m (Body rep) mkPlusBnds plus accels = do plus' <- renameLambda plus let parbnds = diff --git a/src/Futhark/Analysis/Interference.hs b/src/Futhark/Analysis/Interference.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/Interference.hs @@ -0,0 +1,326 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Interference analysis for Futhark programs. +module Futhark.Analysis.Interference (Graph, analyseGPU) where + +import Control.Monad.Reader +import Data.Foldable (toList) +import Data.Function ((&)) +import Data.Functor ((<&>)) +import Data.Map (Map) +import qualified Data.Map as M +import Data.Maybe (catMaybes, fromMaybe) +import Data.Set (Set) +import qualified Data.Set as S +import Futhark.Analysis.LastUse (LastUseMap) +import Futhark.IR.GPUMem +import Futhark.Util (invertMap) + +-- | The set of `VName` currently in use. +type InUse = Names + +-- | The set of `VName` that are no longer in use. +type LastUsed = Names + +-- | An interference graph. An element `(x, y)` in the set means that there is +-- an undirected edge between `x` and `y`, and therefore the lifetimes of `x` +-- and `y` overlap and they "interfere" with each other. We assume that pairs +-- are always normalized, such that `x` < `y`, before inserting. This should +-- prevent any duplicates. We also don't allow any pairs where `x == y`. +type Graph a = Set (a, a) + +-- | Insert an edge between two values into the graph. +makeEdge :: Ord a => a -> a -> Graph a +makeEdge v1 v2 + | v1 == v2 = mempty + | otherwise = S.singleton (min v1 v2, max v1 v2) + +-- | Compute the cartesian product of two foldable collections, using the given +-- combinator function. +cartesian :: (Monoid m, Foldable t) => (a -> a -> m) -> t a -> t a -> m +cartesian f xs ys = + [(x, y) | x <- toList xs, y <- toList ys] + & foldMap (uncurry f) + +analyseStm :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + Stm GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseStm lumap inuse0 stm = + inScopeOf stm $ do + let pat_name = patElemName $ head $ patternValueElements $ stmPattern stm + + new_mems <- + stmPattern stm + & patternValueElements + & mapM (memInfo . patElemName) + <&> catMaybes + <&> namesFromList + + -- `new_mems` should interfere with any mems inside the statement expression + let inuse_outside = inuse0 <> new_mems + + -- `inuse` is the set of memory blocks that are inuse at the end of any code + -- bodies inside the expression. `lus` is the set of all memory blocks that + -- have reached their last use in any code bodies inside the + -- expression. `graph` is the interference graph computed for any code + -- bodies inside the expression. + (inuse, lus, graph) <- analyseExp lumap inuse_outside (stmExp stm) + + last_use_mems <- + M.lookup pat_name lumap + & fromMaybe mempty + & namesToList + & mapM memInfo + <&> catMaybes + <&> namesFromList + <&> namesIntersection inuse_outside + + return + ( (inuse_outside `namesSubtract` last_use_mems `namesSubtract` lus) + <> new_mems, + (lus <> last_use_mems) `namesSubtract` new_mems, + graph + <> cartesian + makeEdge + (namesToList inuse_outside) + (namesToList $ inuse_outside <> inuse <> lus <> last_use_mems) + ) + +analyseExp :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + Exp GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseExp lumap inuse_outside expr = + case expr of + If _ then_body else_body _ -> do + res1 <- analyseBody lumap inuse_outside then_body + res2 <- analyseBody lumap inuse_outside else_body + return $ res1 <> res2 + DoLoop _ _ _ body -> do + analyseBody lumap inuse_outside body + Op (Inner (SegOp segop)) -> do + analyseSegOp lumap inuse_outside segop + _ -> + return mempty + +analyseKernelBody :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + KernelBody GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseKernelBody lumap inuse body = analyseStms lumap inuse $ kernelBodyStms body + +analyseBody :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + Body GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseBody lumap inuse body = analyseStms lumap inuse $ bodyStms body + +analyseStms :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + Stms GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseStms lumap inuse0 stms = do + inScopeOf stms $ foldM helper (inuse0, mempty, mempty) $ stmsToList stms + where + helper (inuse, lus, graph) stm = do + (inuse', lus', graph') <- analyseStm lumap inuse stm + return (inuse', lus' <> lus, graph' <> graph) + +analyseSegOp :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + SegOp lvl GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseSegOp lumap inuse (SegMap _ _ _ body) = + analyseKernelBody lumap inuse body +analyseSegOp lumap inuse (SegRed _ _ binops _ body) = + segWithBinOps lumap inuse binops body +analyseSegOp lumap inuse (SegScan _ _ binops _ body) = do + segWithBinOps lumap inuse binops body +analyseSegOp lumap inuse (SegHist _ _ histops _ body) = do + (inuse', lus', graph) <- analyseKernelBody lumap inuse body + (inuse'', lus'', graph') <- mconcat <$> mapM (analyseHistOp lumap inuse') histops + return (inuse'', lus' <> lus'', graph <> graph') + +segWithBinOps :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + [SegBinOp GPUMem] -> + KernelBody GPUMem -> + m (InUse, LastUsed, Graph VName) +segWithBinOps lumap inuse binops body = do + (inuse', lus', graph) <- analyseKernelBody lumap inuse body + (inuse'', lus'', graph') <- + mconcat + <$> mapM + (analyseSegBinOp lumap inuse') + binops + return (inuse'', lus' <> lus'', graph <> graph') + +analyseSegBinOp :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + SegBinOp GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseSegBinOp lumap inuse (SegBinOp _ lambda _ _) = + analyseLambda lumap inuse lambda + +analyseHistOp :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + HistOp GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseHistOp lumap inuse histop = + analyseLambda lumap inuse (histOp histop) + +analyseLambda :: + LocalScope GPUMem m => + LastUseMap -> + InUse -> + Lambda GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseLambda lumap inuse (Lambda _ body _) = + analyseBody lumap inuse body + +-- | Perform interference analysis on the given statements. The result is a +-- triple of the names currently in use, names that hit their last use somewhere +-- within, and the resulting graph. +analyseGPU :: + LocalScope GPUMem m => + LastUseMap -> + Stms GPUMem -> + m (Graph VName) +analyseGPU lumap stms = do + (_, _, graph) <- analyseGPU' lumap stms + -- We need to insert edges between memory blocks which differ in size, if they + -- are in DefaultSpace. The problem is that during memory expansion, + -- DefaultSpace arrays in kernels are interleaved. If the element sizes of two + -- merged memory blocks are different, threads might try to read and write to + -- overlapping memory positions. More information here: + -- https://munksgaard.me/technical-diary/2020-12-30.html#org210775b + spaces <- M.filter (== DefaultSpace) <$> memSpaces stms + inv_size_map <- + memSizes stms + <&> flip M.restrictKeys (S.fromList $ M.keys spaces) + <&> invertMap + let new_edges = + cartesian + (\x y -> if x /= y then cartesian makeEdge x y else mempty) + inv_size_map + inv_size_map + return $ graph <> new_edges + +-- | Return a mapping from memory blocks to their element sizes in the given +-- statements. +memSizes :: LocalScope GPUMem m => Stms GPUMem -> m (Map VName Int) +memSizes stms = + inScopeOf stms $ fmap mconcat <$> mapM memSizesStm $ stmsToList stms + where + memSizesStm :: LocalScope GPUMem m => Stm GPUMem -> m (Map VName Int) + memSizesStm (Let pat _ e) = do + arraySizes <- fmap mconcat <$> mapM memElemSize $ patternNames pat + arraySizes' <- memSizesExp e + return $ arraySizes <> arraySizes' + memSizesExp :: LocalScope GPUMem m => Exp GPUMem -> m (Map VName Int) + memSizesExp (Op (Inner (SegOp segop))) = + let body = segBody segop + in inScopeOf (kernelBodyStms body) $ + fmap mconcat + <$> mapM memSizesStm + $ stmsToList $ kernelBodyStms body + memSizesExp (If _ then_body else_body _) = do + then_res <- memSizes $ bodyStms then_body + else_res <- memSizes $ bodyStms else_body + return $ then_res <> else_res + memSizesExp (DoLoop _ _ _ body) = + memSizes $ bodyStms body + memSizesExp _ = return mempty + +-- | Return a mapping from memory blocks to the space they are allocated in. +memSpaces :: LocalScope GPUMem m => Stms GPUMem -> m (Map VName Space) +memSpaces stms = + return $ foldMap getSpacesStm stms + where + getSpacesStm :: Stm GPUMem -> Map VName Space + getSpacesStm (Let (Pattern [] [PatElem name _]) _ (Op (Alloc _ sp))) = + M.singleton name sp + getSpacesStm (Let _ _ (Op (Alloc _ _))) = error "impossible" + getSpacesStm (Let _ _ (Op (Inner (SegOp segop)))) = + foldMap getSpacesStm $ kernelBodyStms $ segBody segop + getSpacesStm (Let _ _ (If _ then_body else_body _)) = + foldMap getSpacesStm (bodyStms then_body) + <> foldMap getSpacesStm (bodyStms else_body) + getSpacesStm (Let _ _ (DoLoop _ _ _ body)) = + foldMap getSpacesStm (bodyStms body) + getSpacesStm _ = mempty + +analyseGPU' :: + LocalScope GPUMem m => + LastUseMap -> + Stms GPUMem -> + m (InUse, LastUsed, Graph VName) +analyseGPU' lumap stms = + mconcat . toList <$> mapM helper stms + where + helper :: + LocalScope GPUMem m => + Stm GPUMem -> + m (InUse, LastUsed, Graph VName) + helper stm@Let {stmExp = Op (Inner (SegOp segop))} = + inScopeOf stm $ analyseSegOp lumap mempty segop + helper stm@Let {stmExp = If _ then_body else_body _} = + inScopeOf stm $ do + res1 <- analyseGPU' lumap (bodyStms then_body) + res2 <- analyseGPU' lumap (bodyStms else_body) + return (res1 <> res2) + helper stm@Let {stmExp = DoLoop _ _ _ body} = + inScopeOf stm $ + analyseGPU' lumap $ bodyStms body + helper stm = + inScopeOf stm $ return mempty + +nameInfoToMemInfo :: Mem rep => NameInfo rep -> MemBound NoUniqueness +nameInfoToMemInfo info = + case info of + FParamName summary -> noUniquenessReturns summary + LParamName summary -> summary + LetName summary -> summary + IndexName it -> MemPrim $ IntType it + +memInfo :: LocalScope GPUMem m => VName -> m (Maybe VName) +memInfo vname = do + summary <- asksScope (fmap nameInfoToMemInfo . M.lookup vname) + case summary of + Just (MemArray _ _ _ (ArrayIn mem _)) -> + return $ Just mem + _ -> + return Nothing + +-- | Returns a mapping from memory block to element size. The input is the +-- `VName` of a variable (supposedly an array), and the result is a mapping from +-- the memory block of that array to element size of the array. +memElemSize :: LocalScope GPUMem m => VName -> m (Map VName Int) +memElemSize vname = do + summary <- asksScope (fmap nameInfoToMemInfo . M.lookup vname) + case summary of + Just (MemArray pt _ _ (ArrayIn mem _)) -> + return $ M.singleton mem (primByteSize pt) + _ -> + return mempty diff --git a/src/Futhark/Analysis/LastUse.hs b/src/Futhark/Analysis/LastUse.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/LastUse.hs @@ -0,0 +1,161 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Provides last-use analysis for Futhark programs. +module Futhark.Analysis.LastUse (LastUseMap, analyseProg) where + +import Data.Bifunctor (first) +import Data.Foldable +import Data.Function ((&)) +import Data.Map (Map) +import qualified Data.Map as M +import Data.Tuple +import Futhark.Analysis.Alias (aliasAnalysis) +import Futhark.IR.Aliases +import Futhark.IR.GPUMem + +-- | `LastUseMap` tells which names were last used in a given statement. +-- Statements are uniquely identified by the `VName` of the first value +-- parameter in the statement pattern. `Names` is the set of names last used. +type LastUseMap = Map VName Names + +-- | `LastUse` is a mapping from a `VName` to the statement identifying it's +-- last use. `LastUseMap` is the inverse of `LastUse`. +type LastUse = Map VName VName + +-- | `Used` is the set of `VName` that were used somewhere in a statement, body +-- or otherwise. +type Used = Names + +-- | Analyses a program to return a last-use map, mapping each simple statement +-- in the program to the values that were last used within that statement, and +-- the set of all `VName` that were used inside. +analyseProg :: Prog GPUMem -> (LastUseMap, Used) +analyseProg prog = + let consts = + progConsts prog + & concatMap (toList . fmap patElemName . patternValueElements . stmPattern) + & namesFromList + funs = progFuns $ aliasAnalysis prog + (lus, used) = foldMap (analyseFun mempty consts) funs + in (flipMap lus, used) + +analyseFun :: LastUse -> Used -> FunDef (Aliases GPUMem) -> (LastUse, Used) +analyseFun lumap used fun = + let (lumap', used') = analyseBody lumap used $ funDefBody fun + in (lumap', used' <> freeIn (funDefParams fun)) + +analyseStms :: LastUse -> Used -> Stms (Aliases GPUMem) -> (LastUse, Used) +analyseStms lumap used stms = foldr analyseStm (lumap, used) $ stmsToList stms + +analyseStm :: Stm (Aliases GPUMem) -> (LastUse, Used) -> (LastUse, Used) +analyseStm (Let pat _ e) (lumap0, used0) = + let (lumap', used') = patternValueElements pat & foldl helper (lumap0, used0) + in analyseExp (lumap', used') e + where + helper (lumap_acc, used_acc) (PatElem name (aliases, _)) = + -- Any aliases of `name` should have the same last-use as `name` + ( case M.lookup name lumap_acc of + Just name' -> + insertNames name' (unAliases aliases) lumap_acc + Nothing -> lumap_acc, + used_acc <> unAliases aliases + ) + + pat_name = patElemName $ head $ patternValueElements pat + analyseExp :: (LastUse, Used) -> Exp (Aliases GPUMem) -> (LastUse, Used) + analyseExp (lumap, used) (BasicOp _) = + let nms = freeIn e `namesSubtract` used + in (insertNames pat_name nms lumap, used <> nms) + analyseExp (lumap, used) (Apply _ args _ _) = + let nms = freeIn $ map fst args + in (insertNames pat_name nms lumap, used <> nms) + analyseExp (lumap, used) (If cse then_body else_body dec) = + let (lumap_then, used_then) = analyseBody lumap used then_body + (lumap_else, used_else) = analyseBody lumap used else_body + used' = used_then <> used_else + nms = ((freeIn cse <> freeIn dec) `namesSubtract` used') + in (insertNames pat_name nms (lumap_then <> lumap_else), used' <> nms) + analyseExp (lumap, used) (DoLoop ctx vals form body) = + let (lumap', used') = analyseBody lumap used body + nms = (freeIn ctx <> freeIn vals <> freeIn form) `namesSubtract` used' + in (insertNames pat_name nms lumap', used' <> nms) + analyseExp (lumap, used) (Op (Alloc se sp)) = + let nms = (freeIn se <> freeIn sp) `namesSubtract` used + in (insertNames pat_name nms lumap, used <> nms) + analyseExp (lumap, used) (Op (Inner (SizeOp sop))) = + let nms = freeIn sop `namesSubtract` used + in (insertNames pat_name nms lumap, used <> nms) + analyseExp (lumap, used) (Op (Inner (OtherOp ()))) = + (lumap, used) + analyseExp (lumap, used) (Op (Inner (SegOp (SegMap lvl _ tps body)))) = + let (lumap', used') = analyseKernelBody (lumap, used) body + nms = (freeIn lvl <> freeIn tps) `namesSubtract` used' + in (insertNames pat_name nms lumap', used' <> nms) + analyseExp (lumap, used) (Op (Inner (SegOp (SegRed lvl _ binops tps body)))) = + segOpHelper lumap used lvl binops tps body + analyseExp (lumap, used) (Op (Inner (SegOp (SegScan lvl _ binops tps body)))) = + segOpHelper lumap used lvl binops tps body + analyseExp (lumap, used) (Op (Inner (SegOp (SegHist lvl _ binops tps body)))) = + let (lumap', used') = foldr analyseHistOp (lumap, used) binops + (lumap'', used'') = analyseKernelBody (lumap', used') body + nms = (freeIn lvl <> freeIn tps) `namesSubtract` used'' + in (insertNames pat_name nms lumap'', used'' <> nms) + analyseExp (lumap, used) (WithAcc _ l) = + analyseLambda (lumap, used) l + segOpHelper lumap used lvl binops tps body = + let (lumap', used') = foldr analyseSegBinOp (lumap, used) binops + (lumap'', used'') = analyseKernelBody (lumap', used') body + nms = (freeIn lvl <> freeIn tps) `namesSubtract` used'' + in (insertNames pat_name nms lumap'', used'' <> nms) + +analyseBody :: LastUse -> Used -> Body (Aliases GPUMem) -> (LastUse, Used) +analyseBody lumap used (Body _ stms result) = + let used' = used <> freeIn result + in analyseStms lumap used' stms + +analyseKernelBody :: + (LastUse, Used) -> + KernelBody (Aliases GPUMem) -> + (LastUse, Used) +analyseKernelBody (lumap, used) (KernelBody _ stms result) = + let used' = used <> freeIn result + in analyseStms lumap used' stms + +analyseSegBinOp :: + SegBinOp (Aliases GPUMem) -> + (LastUse, Used) -> + (LastUse, Used) +analyseSegBinOp (SegBinOp _ lambda neutral shp) (lumap, used) = + let (lumap', used') = analyseLambda (lumap, used) lambda + nms = (freeIn neutral <> freeIn shp) `namesSubtract` used' + in (lumap', used' <> nms) + +analyseHistOp :: + HistOp (Aliases GPUMem) -> + (LastUse, Used) -> + (LastUse, Used) +analyseHistOp (HistOp width race dest neutral shp lambda) (lumap, used) = + let (lumap', used') = analyseLambda (lumap, used) lambda + nms = + ( freeIn width <> freeIn race <> freeIn dest <> freeIn neutral + <> freeIn shp + ) + `namesSubtract` used' + in (lumap', used' <> nms) + +analyseLambda :: (LastUse, Used) -> Lambda (Aliases GPUMem) -> (LastUse, Used) +analyseLambda (lumap, used) (Lambda params body ret) = + let (lumap', used') = analyseBody lumap used body + used'' = used' <> freeIn params <> freeIn ret + in (lumap', used'') + +flipMap :: Map VName VName -> Map VName Names +flipMap m = + M.toList m + & fmap (swap . first oneName) + & foldr (uncurry $ M.insertWith (<>)) mempty + +insertNames :: VName -> Names -> LastUse -> LastUse +insertNames name names lumap = + foldr (flip (M.insertWith $ \_ x -> x) name) lumap $ namesToList names diff --git a/src/Futhark/Analysis/Metrics.hs b/src/Futhark/Analysis/Metrics.hs --- a/src/Futhark/Analysis/Metrics.hs +++ b/src/Futhark/Analysis/Metrics.hs @@ -80,7 +80,7 @@ addWhat' (ctx, k) = (what : ctx, k) -- | Compute the metrics for a program. -progMetrics :: OpMetrics (Op lore) => Prog lore -> AstMetrics +progMetrics :: OpMetrics (Op rep) => Prog rep -> AstMetrics progMetrics prog = actualMetrics $ execWriter $ @@ -88,17 +88,17 @@ mapM_ funDefMetrics $ progFuns prog mapM_ stmMetrics $ progConsts prog -funDefMetrics :: OpMetrics (Op lore) => FunDef lore -> MetricsM () +funDefMetrics :: OpMetrics (Op rep) => FunDef rep -> MetricsM () funDefMetrics = bodyMetrics . funDefBody -bodyMetrics :: OpMetrics (Op lore) => Body lore -> MetricsM () +bodyMetrics :: OpMetrics (Op rep) => Body rep -> MetricsM () bodyMetrics = mapM_ stmMetrics . bodyStms -- | Compute metrics for this statement. -stmMetrics :: OpMetrics (Op lore) => Stm lore -> MetricsM () +stmMetrics :: OpMetrics (Op rep) => Stm rep -> MetricsM () stmMetrics = expMetrics . stmExp -expMetrics :: OpMetrics (Op lore) => Exp lore -> MetricsM () +expMetrics :: OpMetrics (Op rep) => Exp rep -> MetricsM () expMetrics (BasicOp op) = seen "BasicOp" >> primOpMetrics op expMetrics (DoLoop _ _ ForLoop {} body) = @@ -139,5 +139,5 @@ primOpMetrics UpdateAcc {} = seen "UpdateAcc" -- | Compute metrics for this lambda. -lambdaMetrics :: OpMetrics (Op lore) => Lambda lore -> MetricsM () +lambdaMetrics :: OpMetrics (Op rep) => Lambda rep -> MetricsM () lambdaMetrics = bodyMetrics . lambdaBody diff --git a/src/Futhark/Analysis/PrimExp/Convert.hs b/src/Futhark/Analysis/PrimExp/Convert.hs --- a/src/Futhark/Analysis/PrimExp/Convert.hs +++ b/src/Futhark/Analysis/PrimExp/Convert.hs @@ -60,9 +60,9 @@ -- This includes constants and variable names, which are passed as -- t'SubExp's. primExpFromExp :: - (Fail.MonadFail m, Decorations lore) => + (Fail.MonadFail m, RepTypes rep) => (VName -> m (PrimExp v)) -> - Exp lore -> + Exp rep -> m (PrimExp v) primExpFromExp f (BasicOp (BinOp op x y)) = BinOpExp op <$> primExpFromSubExpM f x <*> primExpFromSubExpM f y diff --git a/src/Futhark/Analysis/PrimExp/Simplify.hs b/src/Futhark/Analysis/PrimExp/Simplify.hs --- a/src/Futhark/Analysis/PrimExp/Simplify.hs +++ b/src/Futhark/Analysis/PrimExp/Simplify.hs @@ -11,9 +11,9 @@ -- refers to a name that is a 'Constant', the node turns into a -- 'ValueExp'. simplifyPrimExp :: - SimplifiableLore lore => + SimplifiableRep rep => PrimExp VName -> - SimpleM lore (PrimExp VName) + SimpleM rep (PrimExp VName) simplifyPrimExp = simplifyAnyPrimExp onLeaf where onLeaf v pt = do @@ -24,9 +24,9 @@ -- | Like 'simplifyPrimExp', but where leaves may be 'Ext's. simplifyExtPrimExp :: - SimplifiableLore lore => + SimplifiableRep rep => PrimExp (Ext VName) -> - SimpleM lore (PrimExp (Ext VName)) + SimpleM rep (PrimExp (Ext VName)) simplifyExtPrimExp = simplifyAnyPrimExp onLeaf where onLeaf (Free v) pt = do @@ -37,10 +37,10 @@ onLeaf (Ext i) pt = return $ LeafExp (Ext i) pt simplifyAnyPrimExp :: - SimplifiableLore lore => - (a -> PrimType -> SimpleM lore (PrimExp a)) -> + SimplifiableRep rep => + (a -> PrimType -> SimpleM rep (PrimExp a)) -> PrimExp a -> - SimpleM lore (PrimExp a) + SimpleM rep (PrimExp a) simplifyAnyPrimExp f (LeafExp v pt) = f v pt simplifyAnyPrimExp _ (ValueExp pv) = return $ ValueExp pv diff --git a/src/Futhark/Analysis/Rephrase.hs b/src/Futhark/Analysis/Rephrase.hs --- a/src/Futhark/Analysis/Rephrase.hs +++ b/src/Futhark/Analysis/Rephrase.hs @@ -1,6 +1,6 @@ {-# LANGUAGE ConstraintKinds #-} --- | Facilities for changing the lore of some fragment, with no +-- | Facilities for changing the rep of some fragment, with no -- context. We call this "rephrasing", for no deep reason. module Futhark.Analysis.Rephrase ( rephraseProg, @@ -20,14 +20,14 @@ -- | A collection of functions that together allow us to rephrase some -- IR fragment, in some monad @m@. If we let @m@ be the 'Maybe' -- monad, we can conveniently do rephrasing that might fail. This is --- useful if you want to see if some IR in e.g. the @Kernels@ lore +-- useful if you want to see if some IR in e.g. the @Kernels@ rep -- actually uses any @Kernels@-specific operations. data Rephraser m from to = Rephraser - { rephraseExpLore :: ExpDec from -> m (ExpDec to), - rephraseLetBoundLore :: LetDec from -> m (LetDec to), - rephraseFParamLore :: FParamInfo from -> m (FParamInfo to), - rephraseLParamLore :: LParamInfo from -> m (LParamInfo to), - rephraseBodyLore :: BodyDec from -> m (BodyDec to), + { rephraseExpDec :: ExpDec from -> m (ExpDec to), + rephraseLetBoundDec :: LetDec from -> m (LetDec to), + rephraseFParamDec :: FParamInfo from -> m (FParamInfo to), + rephraseLParamDec :: LParamInfo from -> m (LParamInfo to), + rephraseBodyDec :: BodyDec from -> m (BodyDec to), rephraseRetType :: RetType from -> m (RetType to), rephraseBranchType :: BranchType from -> m (BranchType to), rephraseOp :: Op from -> m (Op to) @@ -44,7 +44,7 @@ rephraseFunDef :: Monad m => Rephraser m from to -> FunDef from -> m (FunDef to) rephraseFunDef rephraser fundec = do body' <- rephraseBody rephraser $ funDefBody fundec - params' <- mapM (rephraseParam $ rephraseFParamLore rephraser) $ funDefParams fundec + params' <- mapM (rephraseParam $ rephraseFParamDec rephraser) $ funDefParams fundec rettype' <- mapM (rephraseRetType rephraser) $ funDefRetType fundec return fundec {funDefBody = body', funDefParams = params', funDefRetType = rettype'} @@ -56,8 +56,8 @@ rephraseStm :: Monad m => Rephraser m from to -> Stm from -> m (Stm to) rephraseStm rephraser (Let pat (StmAux cs attrs dec) e) = Let - <$> rephrasePattern (rephraseLetBoundLore rephraser) pat - <*> (StmAux cs attrs <$> rephraseExpLore rephraser dec) + <$> rephrasePattern (rephraseLetBoundDec rephraser) pat + <*> (StmAux cs attrs <$> rephraseExpDec rephraser dec) <*> rephraseExp rephraser e -- | Rephrase a pattern. @@ -80,9 +80,9 @@ -- | Rephrase a body. rephraseBody :: Monad m => Rephraser m from to -> Body from -> m (Body to) -rephraseBody rephraser (Body lore bnds res) = +rephraseBody rephraser (Body rep bnds res) = Body - <$> rephraseBodyLore rephraser lore + <$> rephraseBodyDec rephraser rep <*> (stmsFromList <$> mapM (rephraseStm rephraser) (stmsToList bnds)) <*> pure res @@ -90,7 +90,7 @@ rephraseLambda :: Monad m => Rephraser m from to -> Lambda from -> m (Lambda to) rephraseLambda rephraser lam = do body' <- rephraseBody rephraser $ lambdaBody lam - params' <- mapM (rephraseParam $ rephraseLParamLore rephraser) $ lambdaParams lam + params' <- mapM (rephraseParam $ rephraseLParamDec rephraser) $ lambdaParams lam return lam {lambdaBody = body', lambdaParams = params'} mapper :: Monad m => Rephraser m from to -> Mapper from to m @@ -99,7 +99,7 @@ { mapOnBody = const $ rephraseBody rephraser, mapOnRetType = rephraseRetType rephraser, mapOnBranchType = rephraseBranchType rephraser, - mapOnFParam = rephraseParam (rephraseFParamLore rephraser), - mapOnLParam = rephraseParam (rephraseLParamLore rephraser), + mapOnFParam = rephraseParam (rephraseFParamDec rephraser), + mapOnLParam = rephraseParam (rephraseLParamDec rephraser), mapOnOp = rephraseOp rephraser } diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs --- a/src/Futhark/Analysis/SymbolTable.hs +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -61,9 +61,9 @@ import qualified Futhark.IR.Prop.Aliases as Aliases import Prelude hiding (elem, lookup) -data SymbolTable lore = SymbolTable +data SymbolTable rep = SymbolTable { loopDepth :: Int, - bindings :: M.Map VName (Entry lore), + bindings :: M.Map VName (Entry rep), -- | Which names are available just before the most enclosing -- loop? availableAtClosestLoop :: Names, @@ -73,7 +73,7 @@ simplifyMemory :: Bool } -instance Semigroup (SymbolTable lore) where +instance Semigroup (SymbolTable rep) where table1 <> table2 = SymbolTable { loopDepth = max (loopDepth table1) (loopDepth table2), @@ -84,21 +84,21 @@ simplifyMemory = simplifyMemory table1 || simplifyMemory table2 } -instance Monoid (SymbolTable lore) where +instance Monoid (SymbolTable rep) where mempty = empty -empty :: SymbolTable lore +empty :: SymbolTable rep empty = SymbolTable 0 M.empty mempty False -fromScope :: ASTLore lore => Scope lore -> SymbolTable lore +fromScope :: ASTRep rep => Scope rep -> SymbolTable rep fromScope = M.foldlWithKey' insertFreeVar' empty where insertFreeVar' m k dec = insertFreeVar k dec m -toScope :: SymbolTable lore -> Scope lore +toScope :: SymbolTable rep -> Scope rep toScope = M.map entryInfo . bindings -deepen :: SymbolTable lore -> SymbolTable lore +deepen :: SymbolTable rep -> SymbolTable rep deepen vtable = vtable { loopDepth = loopDepth vtable + 1, @@ -125,59 +125,59 @@ -- | Indexing a delayed array if possible. type IndexArray = [TPrimExp Int64 VName] -> Maybe Indexed -data Entry lore = Entry +data Entry rep = Entry { -- | True if consumed. entryConsumed :: Bool, entryDepth :: Int, -- | True if this name has been used as an array size, -- implying that it is non-negative. entryIsSize :: Bool, - entryType :: EntryType lore + entryType :: EntryType rep } -data EntryType lore - = LoopVar (LoopVarEntry lore) - | LetBound (LetBoundEntry lore) - | FParam (FParamEntry lore) - | LParam (LParamEntry lore) - | FreeVar (FreeVarEntry lore) +data EntryType rep + = LoopVar (LoopVarEntry rep) + | LetBound (LetBoundEntry rep) + | FParam (FParamEntry rep) + | LParam (LParamEntry rep) + | FreeVar (FreeVarEntry rep) -data LoopVarEntry lore = LoopVarEntry +data LoopVarEntry rep = LoopVarEntry { loopVarType :: IntType, loopVarBound :: SubExp } -data LetBoundEntry lore = LetBoundEntry - { letBoundDec :: LetDec lore, +data LetBoundEntry rep = LetBoundEntry + { letBoundDec :: LetDec rep, letBoundAliases :: Names, - letBoundStm :: Stm lore, + letBoundStm :: Stm rep, -- | Index a delayed array, if possible. letBoundIndex :: Int -> IndexArray } -data FParamEntry lore = FParamEntry - { fparamDec :: FParamInfo lore, +data FParamEntry rep = FParamEntry + { fparamDec :: FParamInfo rep, fparamAliases :: Names, -- | If a loop parameter, the initial value and the eventual -- result. The result need not be in scope in the symbol table. fparamMerge :: Maybe (SubExp, SubExp) } -data LParamEntry lore = LParamEntry - { lparamDec :: LParamInfo lore, +data LParamEntry rep = LParamEntry + { lparamDec :: LParamInfo rep, lparamIndex :: IndexArray } -data FreeVarEntry lore = FreeVarEntry - { freeVarDec :: NameInfo lore, +data FreeVarEntry rep = FreeVarEntry + { freeVarDec :: NameInfo rep, -- | Index a delayed array, if possible. freeVarIndex :: VName -> IndexArray } -instance ASTLore lore => Typed (Entry lore) where +instance ASTRep rep => Typed (Entry rep) where typeOf = typeOf . entryInfo -entryInfo :: Entry lore -> NameInfo lore +entryInfo :: Entry rep -> NameInfo rep entryInfo e = case entryType e of LetBound entry -> LetName $ letBoundDec entry LoopVar entry -> IndexName $ loopVarType entry @@ -185,54 +185,54 @@ LParam entry -> LParamName $ lparamDec entry FreeVar entry -> freeVarDec entry -isLetBound :: Entry lore -> Maybe (LetBoundEntry lore) +isLetBound :: Entry rep -> Maybe (LetBoundEntry rep) isLetBound e = case entryType e of LetBound entry -> Just entry _ -> Nothing -entryStm :: Entry lore -> Maybe (Stm lore) +entryStm :: Entry rep -> Maybe (Stm rep) entryStm = fmap letBoundStm . isLetBound -entryFParam :: Entry lore -> Maybe (FParamInfo lore) +entryFParam :: Entry rep -> Maybe (FParamInfo rep) entryFParam e = case entryType e of FParam e' -> Just $ fparamDec e' _ -> Nothing -entryLetBoundDec :: Entry lore -> Maybe (LetDec lore) +entryLetBoundDec :: Entry rep -> Maybe (LetDec rep) entryLetBoundDec = fmap letBoundDec . isLetBound -elem :: VName -> SymbolTable lore -> Bool +elem :: VName -> SymbolTable rep -> Bool elem name = isJust . lookup name -lookup :: VName -> SymbolTable lore -> Maybe (Entry lore) +lookup :: VName -> SymbolTable rep -> Maybe (Entry rep) lookup name = M.lookup name . bindings -lookupStm :: VName -> SymbolTable lore -> Maybe (Stm lore) +lookupStm :: VName -> SymbolTable rep -> Maybe (Stm rep) lookupStm name vtable = entryStm =<< lookup name vtable -lookupExp :: VName -> SymbolTable lore -> Maybe (Exp lore, Certificates) +lookupExp :: VName -> SymbolTable rep -> Maybe (Exp rep, Certificates) lookupExp name vtable = (stmExp &&& stmCerts) <$> lookupStm name vtable -lookupBasicOp :: VName -> SymbolTable lore -> Maybe (BasicOp, Certificates) +lookupBasicOp :: VName -> SymbolTable rep -> Maybe (BasicOp, Certificates) lookupBasicOp name vtable = case lookupExp name vtable of Just (BasicOp e, cs) -> Just (e, cs) _ -> Nothing -lookupType :: ASTLore lore => VName -> SymbolTable lore -> Maybe Type +lookupType :: ASTRep rep => VName -> SymbolTable rep -> Maybe Type lookupType name vtable = typeOf <$> lookup name vtable -lookupSubExpType :: ASTLore lore => SubExp -> SymbolTable lore -> Maybe Type +lookupSubExpType :: ASTRep rep => SubExp -> SymbolTable rep -> Maybe Type lookupSubExpType (Var v) = lookupType v lookupSubExpType (Constant v) = const $ Just $ Prim $ primValueType v -lookupSubExp :: VName -> SymbolTable lore -> Maybe (SubExp, Certificates) +lookupSubExp :: VName -> SymbolTable rep -> Maybe (SubExp, Certificates) lookupSubExp name vtable = do (e, cs) <- lookupExp name vtable case e of BasicOp (SubExp se) -> Just (se, cs) _ -> Nothing -lookupAliases :: VName -> SymbolTable lore -> Names +lookupAliases :: VName -> SymbolTable rep -> Names lookupAliases name vtable = case entryType <$> M.lookup name (bindings vtable) of Just (LetBound e) -> letBoundAliases e @@ -241,25 +241,25 @@ -- | If the given variable name is the name of a 'ForLoop' parameter, -- then return the bound of that loop. -lookupLoopVar :: VName -> SymbolTable lore -> Maybe SubExp +lookupLoopVar :: VName -> SymbolTable rep -> Maybe SubExp lookupLoopVar name vtable = do LoopVar e <- entryType <$> M.lookup name (bindings vtable) return $ loopVarBound e -lookupLoopParam :: VName -> SymbolTable lore -> Maybe (SubExp, SubExp) +lookupLoopParam :: VName -> SymbolTable rep -> Maybe (SubExp, SubExp) lookupLoopParam name vtable = do FParam e <- entryType <$> M.lookup name (bindings vtable) fparamMerge e -- | In symbol table and not consumed. -available :: VName -> SymbolTable lore -> Bool +available :: VName -> SymbolTable rep -> Bool available name = maybe False (not . entryConsumed) . M.lookup name . bindings index :: - ASTLore lore => + ASTRep rep => VName -> [SubExp] -> - SymbolTable lore -> + SymbolTable rep -> Maybe Indexed index name is table = do is' <- mapM asPrimExp is @@ -272,7 +272,7 @@ index' :: VName -> [TPrimExp Int64 VName] -> - SymbolTable lore -> + SymbolTable rep -> Maybe Indexed index' name is vtable = do entry <- lookup name vtable @@ -290,8 +290,8 @@ class IndexOp op where indexOp :: - (ASTLore lore, IndexOp (Op lore)) => - SymbolTable lore -> + (ASTRep rep, IndexOp (Op rep)) => + SymbolTable rep -> Int -> op -> [TPrimExp Int64 VName] -> @@ -301,9 +301,9 @@ instance IndexOp () indexExp :: - (IndexOp (Op lore), ASTLore lore) => - SymbolTable lore -> - Exp lore -> + (IndexOp (Op rep), ASTRep rep) => + SymbolTable rep -> + Exp rep -> Int -> IndexArray indexExp vtable (Op op) k is = @@ -345,12 +345,12 @@ indexExp _ _ _ _ = Nothing defBndEntry :: - (ASTLore lore, IndexOp (Op lore)) => - SymbolTable lore -> - PatElem lore -> + (ASTRep rep, IndexOp (Op rep)) => + SymbolTable rep -> + PatElem rep -> Names -> - Stm lore -> - LetBoundEntry lore + Stm rep -> + LetBoundEntry rep defBndEntry vtable patElem als bnd = LetBoundEntry { letBoundDec = patElemDec patElem, @@ -362,10 +362,10 @@ } bindingEntries :: - (ASTLore lore, Aliases.Aliased lore, IndexOp (Op lore)) => - Stm lore -> - SymbolTable lore -> - [LetBoundEntry lore] + (ASTRep rep, Aliases.Aliased rep, IndexOp (Op rep)) => + Stm rep -> + SymbolTable rep -> + [LetBoundEntry rep] bindingEntries bnd@(Let pat _ _) vtable = do pat_elem <- patternElements pat return $ defBndEntry vtable pat_elem (Aliases.aliasesOf pat_elem) bnd @@ -374,11 +374,11 @@ adjustSeveral f = flip $ foldl' $ flip $ M.adjust f insertEntry :: - ASTLore lore => + ASTRep rep => VName -> - EntryType lore -> - SymbolTable lore -> - SymbolTable lore + EntryType rep -> + SymbolTable rep -> + SymbolTable rep insertEntry name entry vtable = let entry' = Entry @@ -396,20 +396,20 @@ } insertEntries :: - ASTLore lore => - [(VName, EntryType lore)] -> - SymbolTable lore -> - SymbolTable lore + ASTRep rep => + [(VName, EntryType rep)] -> + SymbolTable rep -> + SymbolTable rep insertEntries entries vtable = foldl' add vtable entries where add vtable' (name, entry) = insertEntry name entry vtable' insertStm :: - (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) => - Stm lore -> - SymbolTable lore -> - SymbolTable lore + (ASTRep rep, IndexOp (Op rep), Aliases.Aliased rep) => + Stm rep -> + SymbolTable rep -> + SymbolTable rep insertStm stm vtable = flip (foldl' $ flip consume) (namesToList stm_consumed) $ flip (foldl' addRevAliases) (patternElements $ stmPattern stm) $ @@ -435,23 +435,23 @@ update' e = e insertStms :: - (ASTLore lore, IndexOp (Op lore), Aliases.Aliased lore) => - Stms lore -> - SymbolTable lore -> - SymbolTable lore + (ASTRep rep, IndexOp (Op rep), Aliases.Aliased rep) => + Stms rep -> + SymbolTable rep -> + SymbolTable rep insertStms stms vtable = foldl' (flip insertStm) vtable $ stmsToList stms -expandAliases :: Names -> SymbolTable lore -> Names +expandAliases :: Names -> SymbolTable rep -> Names expandAliases names vtable = names <> aliasesOfAliases where aliasesOfAliases = mconcat . map (`lookupAliases` vtable) . namesToList $ names insertFParam :: - ASTLore lore => - AST.FParam lore -> - SymbolTable lore -> - SymbolTable lore + ASTRep rep => + AST.FParam rep -> + SymbolTable rep -> + SymbolTable rep insertFParam fparam = insertEntry name entry where name = AST.paramName fparam @@ -464,13 +464,13 @@ } insertFParams :: - ASTLore lore => - [AST.FParam lore] -> - SymbolTable lore -> - SymbolTable lore + ASTRep rep => + [AST.FParam rep] -> + SymbolTable rep -> + SymbolTable rep insertFParams fparams symtable = foldl' (flip insertFParam) symtable fparams -insertLParam :: ASTLore lore => LParam lore -> SymbolTable lore -> SymbolTable lore +insertLParam :: ASTRep rep => LParam rep -> SymbolTable rep -> SymbolTable rep insertLParam param = insertEntry name bind where bind = @@ -489,10 +489,10 @@ -- used to help some loop optimisations detect invariant loop -- parameters. insertLoopMerge :: - ASTLore lore => - [(AST.FParam lore, SubExp, SubExp)] -> - SymbolTable lore -> - SymbolTable lore + ASTRep rep => + [(AST.FParam rep, SubExp, SubExp)] -> + SymbolTable rep -> + SymbolTable rep insertLoopMerge = flip $ foldl' $ flip bind where bind (p, initial, res) = @@ -504,7 +504,7 @@ fparamMerge = Just (initial, res) } -insertLoopVar :: ASTLore lore => VName -> IntType -> SubExp -> SymbolTable lore -> SymbolTable lore +insertLoopVar :: ASTRep rep => VName -> IntType -> SubExp -> SymbolTable rep -> SymbolTable rep insertLoopVar name it bound = insertEntry name bind where bind = @@ -514,7 +514,7 @@ loopVarBound = bound } -insertFreeVar :: ASTLore lore => VName -> NameInfo lore -> SymbolTable lore -> SymbolTable lore +insertFreeVar :: ASTRep rep => VName -> NameInfo rep -> SymbolTable rep -> SymbolTable rep insertFreeVar name dec = insertEntry name entry where entry = @@ -524,7 +524,7 @@ freeVarIndex = \_ _ -> Nothing } -consume :: VName -> SymbolTable lore -> SymbolTable lore +consume :: VName -> SymbolTable rep -> SymbolTable rep consume consumee vtable = foldl' consume' vtable $ namesToList $ @@ -535,7 +535,7 @@ consume'' e = e {entryConsumed = True} -- | Hide definitions of those entries that satisfy some predicate. -hideIf :: (Entry lore -> Bool) -> SymbolTable lore -> SymbolTable lore +hideIf :: (Entry rep -> Bool) -> SymbolTable rep -> SymbolTable rep hideIf hide vtable = vtable {bindings = M.map maybeHide $ bindings vtable} where maybeHide entry @@ -552,7 +552,7 @@ -- | Hide these definitions, if they are protected by certificates in -- the set of names. -hideCertified :: Names -> SymbolTable lore -> SymbolTable lore +hideCertified :: Names -> SymbolTable rep -> SymbolTable rep hideCertified to_hide = hideIf $ maybe False hide . entryStm where hide = any (`nameIn` to_hide) . unCertificates . stmCerts diff --git a/src/Futhark/Analysis/UsageTable.hs b/src/Futhark/Analysis/UsageTable.hs --- a/src/Futhark/Analysis/UsageTable.hs +++ b/src/Futhark/Analysis/UsageTable.hs @@ -141,16 +141,16 @@ withoutU :: Usages -> Usages -> Usages withoutU (Usages x) (Usages y) = Usages $ x .&. complement y -usageInBody :: Aliased lore => Body lore -> UsageTable +usageInBody :: Aliased rep => Body rep -> UsageTable usageInBody = foldMap consumedUsage . namesToList . consumedInBody -- | Produce a usage table reflecting the use of the free variables in -- a single statement. -usageInStm :: (ASTLore lore, Aliased lore) => Stm lore -> UsageTable -usageInStm (Let pat lore e) = +usageInStm :: (ASTRep rep, Aliased rep) => Stm rep -> UsageTable +usageInStm (Let pat rep e) = mconcat [ usageInPat, - usageInExpLore, + usageInExpDec, usageInExp e, usages (freeIn e) ] @@ -161,10 +161,10 @@ `namesSubtract` namesFromList (patternNames pat) ) <> sizeUsages (foldMap (freeIn . patElemType) (patternElements pat)) - usageInExpLore = - usages $ freeIn lore + usageInExpDec = + usages $ freeIn rep -usageInExp :: Aliased lore => Exp lore -> UsageTable +usageInExp :: Aliased rep => Exp rep -> UsageTable usageInExp (Apply _ args _ _) = mconcat [ mconcat $ diff --git a/src/Futhark/Bench.hs b/src/Futhark/Bench.hs --- a/src/Futhark/Bench.hs +++ b/src/Futhark/Bench.hs @@ -158,11 +158,13 @@ cmdMaybe . liftIO $ cmdClear server - valuesAsVars server (zip ins input_types) futhark dir input_spec - let freeOuts = cmdMaybe (cmdFree server outs) freeIns = cmdMaybe (cmdFree server ins) + loadInput = valuesAsVars server (zip ins $ map inputType input_types) futhark dir input_spec + reloadInput = freeIns >> loadInput + loadInput + let runtime l | Just l' <- T.stripPrefix "runtime: " l, [(x, "")] <- reads $ T.unpack l' = @@ -172,6 +174,7 @@ doRun = do call_lines <- cmdEither (cmdCall server entry outs ins) + when (any inputConsumed input_types) reloadInput case mapMaybe runtime call_lines of [call_runtime] -> do liftIO $ fromMaybe (const $ pure ()) (runResultAction opts) call_runtime @@ -197,7 +200,7 @@ report <- cmdEither $ cmdReport server - vs <- readResults server outs program <* freeOuts + vs <- readResults server outs <* freeOuts maybe_expected <- liftIO $ maybe (return Nothing) (fmap Just . getExpectedValues) expected_spec diff --git a/src/Futhark/Binder.hs b/src/Futhark/Binder.hs --- a/src/Futhark/Binder.hs +++ b/src/Futhark/Binder.hs @@ -42,44 +42,44 @@ import Futhark.IR -- | A 'BinderT' (and by extension, a 'Binder') is only an instance of --- 'MonadBinder' for lores that implement this type class, which --- contains methods for constructing statements. -class ASTLore lore => BinderOps lore where +-- 'MonadBinder' for representations that implement this type class, +-- which contains methods for constructing statements. +class ASTRep rep => BinderOps rep where mkExpDecB :: - (MonadBinder m, Lore m ~ lore) => - Pattern lore -> - Exp lore -> - m (ExpDec lore) + (MonadBinder m, Rep m ~ rep) => + Pattern rep -> + Exp rep -> + m (ExpDec rep) mkBodyB :: - (MonadBinder m, Lore m ~ lore) => - Stms lore -> + (MonadBinder m, Rep m ~ rep) => + Stms rep -> Result -> - m (Body lore) + m (Body rep) mkLetNamesB :: - (MonadBinder m, Lore m ~ lore) => + (MonadBinder m, Rep m ~ rep) => [VName] -> - Exp lore -> - m (Stm lore) + Exp rep -> + m (Stm rep) default mkExpDecB :: - (MonadBinder m, Bindable lore) => - Pattern lore -> - Exp lore -> - m (ExpDec lore) + (MonadBinder m, Bindable rep) => + Pattern rep -> + Exp rep -> + m (ExpDec rep) mkExpDecB pat e = return $ mkExpDec pat e default mkBodyB :: - (MonadBinder m, Bindable lore) => - Stms lore -> + (MonadBinder m, Bindable rep) => + Stms rep -> Result -> - m (Body lore) + m (Body rep) mkBodyB stms res = return $ mkBody stms res default mkLetNamesB :: - (MonadBinder m, Lore m ~ lore, Bindable lore) => + (MonadBinder m, Rep m ~ rep, Bindable rep) => [VName] -> - Exp lore -> - m (Stm lore) + Exp rep -> + m (Stm rep) mkLetNamesB = mkLetNames -- | A monad transformer that tracks statements and provides a @@ -88,22 +88,22 @@ -- constructing statements (possibly as part of a larger monad stack). -- If you find yourself needing to implement 'MonadBinder' from -- scratch, then it is likely that you are making a mistake. -newtype BinderT lore m a = BinderT (StateT (Stms lore, Scope lore) m a) +newtype BinderT rep m a = BinderT (StateT (Stms rep, Scope rep) m a) deriving (Functor, Monad, Applicative) -instance MonadTrans (BinderT lore) where +instance MonadTrans (BinderT rep) where lift = BinderT . lift -- | The most commonly used binder monad. -type Binder lore = BinderT lore (State VNameSource) +type Binder rep = BinderT rep (State VNameSource) -instance MonadFreshNames m => MonadFreshNames (BinderT lore m) where +instance MonadFreshNames m => MonadFreshNames (BinderT rep m) where getNameSource = lift getNameSource putNameSource = lift . putNameSource instance - (ASTLore lore, Monad m) => - HasScope lore (BinderT lore m) + (ASTRep rep, Monad m) => + HasScope rep (BinderT rep m) where lookupType name = do t <- BinderT $ gets $ M.lookup name . snd @@ -113,8 +113,8 @@ askScope = BinderT $ gets snd instance - (ASTLore lore, Monad m) => - LocalScope lore (BinderT lore m) + (ASTRep rep, Monad m) => + LocalScope rep (BinderT rep m) where localScope types (BinderT m) = BinderT $ do modify $ second (M.union types) @@ -123,10 +123,10 @@ return x instance - (ASTLore lore, MonadFreshNames m, BinderOps lore) => - MonadBinder (BinderT lore m) + (ASTRep rep, MonadFreshNames m, BinderOps rep) => + MonadBinder (BinderT rep m) where - type Lore (BinderT lore m) = lore + type Rep (BinderT rep m) = rep mkExpDecM = mkExpDecB mkBodyM = mkBodyB mkLetNamesM = mkLetNamesB @@ -148,9 +148,9 @@ -- the statements added ('addStm') during the action. runBinderT :: MonadFreshNames m => - BinderT lore m a -> - Scope lore -> - m (a, Stms lore) + BinderT rep m a -> + Scope rep -> + m (a, Stms rep) runBinderT (BinderT m) scope = do (x, (stms, _)) <- runStateT m (mempty, scope) return (x, stms) @@ -158,17 +158,17 @@ -- | Like 'runBinderT', but return only the statements. runBinderT_ :: MonadFreshNames m => - BinderT lore m () -> - Scope lore -> - m (Stms lore) + BinderT rep m () -> + Scope rep -> + m (Stms rep) runBinderT_ m = fmap snd . runBinderT m -- | Like 'runBinderT', but get the initial scope from the current -- monad. runBinderT' :: - (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => - BinderT lore m a -> - m (a, Stms lore) + (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) => + BinderT rep m a -> + m (a, Stms rep) runBinderT' m = do scope <- askScope runBinderT m $ castScope scope @@ -176,9 +176,9 @@ -- | Like 'runBinderT_', but get the initial scope from the current -- monad. runBinderT'_ :: - (MonadFreshNames m, HasScope somelore m, SameScope somelore lore) => - BinderT lore m a -> - m (Stms lore) + (MonadFreshNames m, HasScope somerep m, SameScope somerep rep) => + BinderT rep m a -> + m (Stms rep) runBinderT'_ = fmap snd . runBinderT' -- | Run a binder action, returning a value and the statements added @@ -186,11 +186,11 @@ -- provides initial scope and name source. runBinder :: ( MonadFreshNames m, - HasScope somelore m, - SameScope somelore lore + HasScope somerep m, + SameScope somerep rep ) => - Binder lore a -> - m (a, Stms lore) + Binder rep a -> + m (a, Stms rep) runBinder m = do types <- askScope modifyNameSource $ runState $ runBinderT m $ castScope types @@ -199,23 +199,23 @@ -- added statements. runBinder_ :: ( MonadFreshNames m, - HasScope somelore m, - SameScope somelore lore + HasScope somerep m, + SameScope somerep rep ) => - Binder lore a -> - m (Stms lore) + Binder rep a -> + m (Stms rep) runBinder_ = fmap snd . runBinder -- | Run a binder that produces a t'Body', and prefix that t'Body' by -- the statements produced during execution of the action. runBodyBinder :: - ( Bindable lore, + ( Bindable rep, MonadFreshNames m, - HasScope somelore m, - SameScope somelore lore + HasScope somerep m, + SameScope somerep rep ) => - Binder lore (Body lore) -> - m (Body lore) + Binder rep (Body rep) -> + m (Body rep) runBodyBinder = fmap (uncurry $ flip insertStms) . runBinder -- Utility instance defintions for MTL classes. These require @@ -223,26 +223,26 @@ mapInner :: Monad m => - ( m (a, (Stms lore, Scope lore)) -> - m (b, (Stms lore, Scope lore)) + ( m (a, (Stms rep, Scope rep)) -> + m (b, (Stms rep, Scope rep)) ) -> - BinderT lore m a -> - BinderT lore m b + BinderT rep m a -> + BinderT rep m b mapInner f (BinderT m) = BinderT $ do s <- get (x, s') <- lift $ f $ runStateT m s put s' return x -instance MonadReader r m => MonadReader r (BinderT lore m) where +instance MonadReader r m => MonadReader r (BinderT rep m) where ask = BinderT $ lift ask local f = mapInner $ local f -instance MonadState s m => MonadState s (BinderT lore m) where +instance MonadState s m => MonadState s (BinderT rep m) where get = BinderT $ lift get put = BinderT . lift . put -instance MonadWriter w m => MonadWriter w (BinderT lore m) where +instance MonadWriter w m => MonadWriter w (BinderT rep m) where tell = BinderT . lift . tell pass = mapInner $ \m -> pass $ do ((x, f), s) <- m @@ -251,7 +251,7 @@ ((x, s), y) <- listen m return ((x, y), s) -instance MonadError e m => MonadError e (BinderT lore m) where +instance MonadError e m => MonadError e (BinderT rep m) where throwError = lift . throwError catchError (BinderT m) f = BinderT $ catchError m $ unBinder . f diff --git a/src/Futhark/Binder/Class.hs b/src/Futhark/Binder/Class.hs --- a/src/Futhark/Binder/Class.hs +++ b/src/Futhark/Binder/Class.hs @@ -26,33 +26,33 @@ import Futhark.IR import Futhark.MonadFreshNames --- | The class of lores that can be constructed solely from an --- expression, within some monad. Very important: the methods should --- not have any significant side effects! They may be called more --- often than you think, and the results thrown away. If used +-- | The class of representations that can be constructed solely from +-- an expression, within some monad. Very important: the methods +-- should not have any significant side effects! They may be called +-- more often than you think, and the results thrown away. If used -- exclusively within a 'MonadBinder' instance, it is acceptable for -- them to create new bindings, however. class - ( ASTLore lore, - FParamInfo lore ~ DeclType, - LParamInfo lore ~ Type, - RetType lore ~ DeclExtType, - BranchType lore ~ ExtType, - SetType (LetDec lore) + ( ASTRep rep, + FParamInfo rep ~ DeclType, + LParamInfo rep ~ Type, + RetType rep ~ DeclExtType, + BranchType rep ~ ExtType, + SetType (LetDec rep) ) => - Bindable lore + Bindable rep where - mkExpPat :: [Ident] -> [Ident] -> Exp lore -> Pattern lore - mkExpDec :: Pattern lore -> Exp lore -> ExpDec lore - mkBody :: Stms lore -> Result -> Body lore + mkExpPat :: [Ident] -> [Ident] -> Exp rep -> Pattern rep + mkExpDec :: Pattern rep -> Exp rep -> ExpDec rep + mkBody :: Stms rep -> Result -> Body rep mkLetNames :: - (MonadFreshNames m, HasScope lore m) => + (MonadFreshNames m, HasScope rep m) => [VName] -> - Exp lore -> - m (Stm lore) + Exp rep -> + m (Stm rep) -- | A monad that supports the creation of bindings from expressions --- and bodies from bindings, with a specific lore. This is the main +-- and bodies from bindings, with a specific rep. This is the main -- typeclass that a monad must implement in order for it to be useful -- for generating or modifying Futhark code. Most importantly -- maintains a current state of 'Stms' (as well as a 'Scope') that @@ -63,29 +63,29 @@ -- results thrown away. It is acceptable for them to create new -- bindings, however. class - ( ASTLore (Lore m), + ( ASTRep (Rep m), MonadFreshNames m, Applicative m, Monad m, - LocalScope (Lore m) m + LocalScope (Rep m) m ) => MonadBinder m where - type Lore m :: Data.Kind.Type - mkExpDecM :: Pattern (Lore m) -> Exp (Lore m) -> m (ExpDec (Lore m)) - mkBodyM :: Stms (Lore m) -> Result -> m (Body (Lore m)) - mkLetNamesM :: [VName] -> Exp (Lore m) -> m (Stm (Lore m)) + type Rep m :: Data.Kind.Type + mkExpDecM :: Pattern (Rep m) -> Exp (Rep m) -> m (ExpDec (Rep m)) + mkBodyM :: Stms (Rep m) -> Result -> m (Body (Rep m)) + mkLetNamesM :: [VName] -> Exp (Rep m) -> m (Stm (Rep m)) -- | Add a statement to the 'Stms' under construction. - addStm :: Stm (Lore m) -> m () + addStm :: Stm (Rep m) -> m () addStm = addStms . oneStm -- | Add multiple statements to the 'Stms' under construction. - addStms :: Stms (Lore m) -> m () + addStms :: Stms (Rep m) -> m () -- | Obtain the statements constructed during a monadic action, -- instead of adding them to the state. - collectStms :: m a -> m (a, Stms (Lore m)) + collectStms :: m a -> m (a, Stms (Rep m)) -- | Add the provided certificates to any statements added during -- execution of the action. @@ -95,7 +95,7 @@ -- | Apply a function to the statements added by this action. censorStms :: MonadBinder m => - (Stms (Lore m) -> Stms (Lore m)) -> + (Stms (Rep m) -> Stms (Rep m)) -> m a -> m a censorStms f m = do @@ -112,7 +112,7 @@ -- | Add the certificates and attributes to any statements added by -- this action. -auxing :: MonadBinder m => StmAux anylore -> m a -> m a +auxing :: MonadBinder m => StmAux anyrep -> m a -> m a auxing (StmAux cs attrs _) = censorStms $ fmap onStm where onStm (Let pat aux e) = @@ -127,15 +127,15 @@ -- | Add a statement with the given pattern and expression. letBind :: MonadBinder m => - Pattern (Lore m) -> - Exp (Lore m) -> + Pattern (Rep m) -> + Exp (Rep m) -> m () letBind pat e = addStm =<< Let pat <$> (defAux <$> mkExpDecM pat e) <*> pure e -- | Construct a 'Stm' from identifiers for the context- and value -- part of the pattern, as well as the expression. -mkLet :: Bindable lore => [Ident] -> [Ident] -> Exp lore -> Stm lore +mkLet :: Bindable rep => [Ident] -> [Ident] -> Exp rep -> Stm rep mkLet ctx val e = let pat = mkExpPat ctx val e dec = mkExpDec pat e @@ -143,7 +143,7 @@ -- | Like mkLet, but also take attributes and certificates from the -- given 'StmAux'. -mkLet' :: Bindable lore => [Ident] -> [Ident] -> StmAux a -> Exp lore -> Stm lore +mkLet' :: Bindable rep => [Ident] -> [Ident] -> StmAux a -> Exp rep -> Stm rep mkLet' ctx val (StmAux cs attrs _) e = let pat = mkExpPat ctx val e dec = mkExpDec pat e @@ -151,23 +151,23 @@ -- | Add a statement with the given pattern element names and -- expression. -letBindNames :: MonadBinder m => [VName] -> Exp (Lore m) -> m () +letBindNames :: MonadBinder m => [VName] -> Exp (Rep m) -> m () letBindNames names e = addStm =<< mkLetNamesM names e -- | As 'collectStms', but throw away the ordinary result. -collectStms_ :: MonadBinder m => m a -> m (Stms (Lore m)) +collectStms_ :: MonadBinder m => m a -> m (Stms (Rep m)) collectStms_ = fmap snd . collectStms -- | Add the statements of the body, then return the body result. -bodyBind :: MonadBinder m => Body (Lore m) -> m [SubExp] +bodyBind :: MonadBinder m => Body (Rep m) -> m [SubExp] bodyBind (Body _ stms es) = do addStms stms return es -- | Add several bindings at the outermost level of a t'Body'. -insertStms :: Bindable lore => Stms lore -> Body lore -> Body lore +insertStms :: Bindable rep => Stms rep -> Body rep -> Body rep insertStms stms1 (Body _ stms2 res) = mkBody (stms1 <> stms2) res -- | Add a single binding at the outermost level of a t'Body'. -insertStm :: Bindable lore => Stm lore -> Body lore -> Body lore +insertStm :: Bindable rep => Stm rep -> Body rep -> Body rep insertStm = insertStms . oneStm diff --git a/src/Futhark/CLI/Autotune.hs b/src/Futhark/CLI/Autotune.hs --- a/src/Futhark/CLI/Autotune.hs +++ b/src/Futhark/CLI/Autotune.hs @@ -141,7 +141,7 @@ -- run, but unfortunately we can only set threshold parameters -- on startup. let progbin = "." dropExtension prog - withServer progbin (serverOptions path opts) $ \server -> + withServer (futharkServerCfg progbin (serverOptions path opts)) $ \server -> either (Left . T.unpack) (Right . bestRuntime) <$> benchmarkDataset server diff --git a/src/Futhark/CLI/Bench.hs b/src/Futhark/CLI/Bench.hs --- a/src/Futhark/CLI/Bench.hs +++ b/src/Futhark/CLI/Bench.hs @@ -177,7 +177,7 @@ | null runner = (binpath, extra_options) | otherwise = (runner, binpath : extra_options) - liftIO $ (Just <$> withServer to_run to_run_args f) `catch` onError + liftIO $ (Just <$> withServer (futharkServerCfg to_run to_run_args) f) `catch` onError where onError :: SomeException -> IO (Maybe a) onError e = do @@ -242,16 +242,30 @@ descString :: String -> Int -> String descString desc pad_to = desc ++ ": " ++ replicate (pad_to - length desc) ' ' +interimResult :: Int -> Int -> Int -> String +interimResult us_sum i runs = + printf "%10.0fμs " avg ++ progressBar i runs 10 + where + avg :: Double + avg = fromIntegral us_sum / fromIntegral i + mkProgressPrompt :: Int -> Int -> String -> IO (Maybe Int -> IO ()) mkProgressPrompt runs pad_to dataset_desc | fancyTerminal = do - count <- newIORef (0 :: Int) + count <- newIORef (0, 0) return $ \us -> do putStr "\r" -- Go to start of line. - i <- readIORef count - let i' = if isJust us then i + 1 else i - writeIORef count i' - putStr $ descString (atMostChars 40 dataset_desc) pad_to ++ progressBar i' runs 10 + let p s = + putStr $ + descString (atMostChars 40 dataset_desc) pad_to ++ s + (us_sum, i) <- readIORef count + case us of + Nothing -> p $ replicate 13 ' ' ++ progressBar i runs 10 + Just us' -> do + let us_sum' = us_sum + us' + i' = i + 1 + writeIORef count (us_sum', i') + p $ interimResult us_sum' i' runs putStr " " -- Just to move the cursor away from the progress bar. hFlush stdout | otherwise = do @@ -260,17 +274,19 @@ return $ const $ return () reportResult :: [RunResult] -> IO () -reportResult results = do - let runtimes = map (fromIntegral . runMicroseconds) results - avg = sum runtimes / fromIntegral (length runtimes) - rsd = stddevp runtimes / mean runtimes :: Double - putStrLn $ - printf - "%10.0fμs (RSD: %.3f; min: %3.0f%%; max: %+3.0f%%)" - avg - rsd - ((minimum runtimes / avg - 1) * 100) - ((maxinum runtimes / avg - 1) * 100) +reportResult = putStrLn . reportString + where + reportString results = + printf + "%10.0fμs (RSD: %.3f; min: %3.0f%%; max: %+3.0f%%)" + avg + rsd + ((minimum runtimes / avg - 1) * 100) + ((maxinum runtimes / avg - 1) * 100) + where + runtimes = map (fromIntegral . runMicroseconds) results + avg = sum runtimes / fromIntegral (length runtimes) + rsd = stddevp runtimes / mean runtimes :: Double runBenchmarkCase :: Server -> diff --git a/src/Futhark/CLI/Datacmp.hs b/src/Futhark/CLI/Datacmp.hs --- a/src/Futhark/CLI/Datacmp.hs +++ b/src/Futhark/CLI/Datacmp.hs @@ -5,7 +5,8 @@ import Control.Exception import qualified Data.ByteString.Lazy.Char8 as BS -import Futhark.Test.Values +import Futhark.Data.Compare +import Futhark.Data.Reader import Futhark.Util.Options import System.Exit import System.IO @@ -41,7 +42,7 @@ hPutStrLn stderr $ "Error reading values from " ++ file_b exitFailure (Just vs_a, Just vs_b) -> - case compareValues vs_a vs_b of + case compareSeveralValues (Tolerance 0.002) vs_a vs_b of [] -> return () es -> do mapM_ print es diff --git a/src/Futhark/CLI/Dataset.hs b/src/Futhark/CLI/Dataset.hs --- a/src/Futhark/CLI/Dataset.hs +++ b/src/Futhark/CLI/Dataset.hs @@ -8,17 +8,18 @@ import Control.Monad.ST import qualified Data.Binary as Bin import qualified Data.ByteString.Lazy.Char8 as BS -import qualified Data.Map.Strict as M import qualified Data.Text as T +import qualified Data.Text.IO as T import Data.Vector.Generic (freeze) import qualified Data.Vector.Storable as SVec import qualified Data.Vector.Storable.Mutable as USVec import Data.Word -import Futhark.Test.Values +import qualified Futhark.Data as V +import Futhark.Data.Reader (readValues) import Futhark.Util.Options import Language.Futhark.Parser import Language.Futhark.Pretty () -import Language.Futhark.Prop (UncheckedTypeExp, namesToPrimTypes) +import Language.Futhark.Prop (UncheckedTypeExp) import Language.Futhark.Syntax hiding ( FloatValue (..), IntValue (..), @@ -43,9 +44,9 @@ exitFailure Just vs -> case format config of - Text -> mapM_ (putStrLn . pretty) vs + Text -> mapM_ (T.putStrLn . V.valueText) vs Binary -> mapM_ (BS.putStr . Bin.encode) vs - Type -> mapM_ (putStrLn . pretty . valueType) vs + Type -> mapM_ (T.putStrLn . V.valueTypeText . V.valueType) vs | otherwise = Just $ zipWithM_ @@ -178,11 +179,11 @@ outValue fmt v where name = "option " ++ t - outValue Text = putStrLn . pretty + outValue Text = T.putStrLn . V.valueText outValue Binary = BS.putStr . Bin.encode - outValue Type = putStrLn . pretty . valueType + outValue Type = T.putStrLn . V.valueTypeText . V.valueType -toValueType :: UncheckedTypeExp -> Either String ValueType +toValueType :: UncheckedTypeExp -> Either String V.ValueType toValueType TETuple {} = Left "Cannot handle tuples yet." toValueType TERecord {} = Left "Cannot handle records yet." toValueType TEApply {} = Left "Cannot handle type applications yet." @@ -191,13 +192,16 @@ toValueType (TEUnique t _) = toValueType t toValueType (TEArray t d _) = do d' <- constantDim d - ValueType ds t' <- toValueType t - return $ ValueType (d' : ds) t' + V.ValueType ds t' <- toValueType t + return $ V.ValueType (d' : ds) t' where constantDim (DimExpConst k _) = Right k constantDim _ = Left "Array has non-constant dimension declaration." toValueType (TEVar (QualName [] v) _) - | Just t <- M.lookup v namesToPrimTypes = Right $ ValueType [] t + | Just t <- lookup v m = Right $ V.ValueType [] t + where + m = map f [minBound .. maxBound] + f t = (nameFromText (V.primTypeText t), t) toValueType (TEVar v _) = Left $ "Unknown type " ++ pretty v @@ -263,30 +267,30 @@ (0.0, 1.0) (0.0, 1.0) -randomValue :: RandomConfiguration -> ValueType -> Word64 -> Value -randomValue conf (ValueType ds t) seed = +randomValue :: RandomConfiguration -> V.ValueType -> Word64 -> V.Value +randomValue conf (V.ValueType ds t) seed = case t of - Signed Int8 -> gen i8Range Int8Value - Signed Int16 -> gen i16Range Int16Value - Signed Int32 -> gen i32Range Int32Value - Signed Int64 -> gen i64Range Int64Value - Unsigned Int8 -> gen u8Range Word8Value - Unsigned Int16 -> gen u16Range Word16Value - Unsigned Int32 -> gen u32Range Word32Value - Unsigned Int64 -> gen u64Range Word64Value - FloatType Float32 -> gen f32Range Float32Value - FloatType Float64 -> gen f64Range Float64Value - Bool -> gen (const (False, True)) BoolValue + V.I8 -> gen i8Range V.I8Value + V.I16 -> gen i16Range V.I16Value + V.I32 -> gen i32Range V.I32Value + V.I64 -> gen i64Range V.I64Value + V.U8 -> gen u8Range V.U8Value + V.U16 -> gen u16Range V.U16Value + V.U32 -> gen u32Range V.U32Value + V.U64 -> gen u64Range V.U64Value + V.F32 -> gen f32Range V.F32Value + V.F64 -> gen f64Range V.F64Value + V.Bool -> gen (const (False, True)) V.BoolValue where gen range final = randomVector (range conf) final ds seed randomVector :: (SVec.Storable v, Variate v) => Range v -> - (SVec.Vector Int -> SVec.Vector v -> Value) -> + (SVec.Vector Int -> SVec.Vector v -> V.Value) -> [Int] -> Word64 -> - Value + V.Value randomVector range final ds seed = runST $ do -- USe some nice impure computation where we can preallocate a -- vector of the desired size, populate it via the random number 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 @@ -12,11 +12,12 @@ import qualified Data.Text as T import qualified Data.Text.IO as T import Futhark.Actions +import qualified Futhark.Analysis.Alias as Alias import Futhark.Analysis.Metrics (OpMetrics) import Futhark.Compiler.CLI -import Futhark.IR (ASTLore, Op, Prog, pretty) -import qualified Futhark.IR.Kernels as Kernels -import qualified Futhark.IR.KernelsMem as KernelsMem +import Futhark.IR (ASTRep, Op, Prog, pretty) +import qualified Futhark.IR.GPU as GPU +import qualified Futhark.IR.GPUMem as GPUMem import qualified Futhark.IR.MC as MC import qualified Futhark.IR.MCMem as MCMem import Futhark.IR.Parse @@ -33,12 +34,13 @@ import Futhark.Optimise.Fusion import Futhark.Optimise.InPlaceLowering import Futhark.Optimise.InliningDeadFun +import qualified Futhark.Optimise.ReuseAllocations as ReuseAllocations import Futhark.Optimise.Sink import Futhark.Optimise.TileLoops import Futhark.Optimise.Unstream import Futhark.Pass import Futhark.Pass.ExpandAllocations -import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels +import qualified Futhark.Pass.ExplicitAllocations.GPU as GPU import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq import Futhark.Pass.ExtractKernels import Futhark.Pass.ExtractMulticore @@ -46,7 +48,7 @@ import Futhark.Pass.KernelBabysitting import Futhark.Pass.Simplify import Futhark.Passes -import Futhark.TypeCheck (Checkable) +import Futhark.TypeCheck (Checkable, checkProg) import Futhark.Util.Log import Futhark.Util.Options import qualified Futhark.Util.Pretty as PP @@ -95,10 +97,10 @@ data UntypedPassState = SOACS (Prog SOACS.SOACS) - | Kernels (Prog Kernels.Kernels) + | GPU (Prog GPU.GPU) | MC (Prog MC.MC) | Seq (Prog Seq.Seq) - | KernelsMem (Prog KernelsMem.KernelsMem) + | GPUMem (Prog GPUMem.GPUMem) | MCMem (Prog MCMem.MCMem) | SeqMem (Prog SeqMem.SeqMem) @@ -113,21 +115,21 @@ instance Representation UntypedPassState where representation (SOACS _) = "SOACS" - representation (Kernels _) = "Kernels" + representation (GPU _) = "GPU" representation (MC _) = "MC" representation (Seq _) = "Seq" - representation (KernelsMem _) = "KernelsMem" + representation (GPUMem _) = "GPUMem" representation (MCMem _) = "MCMem" representation (SeqMem _) = "SeqMEm" instance PP.Pretty UntypedPassState where ppr (SOACS prog) = PP.ppr prog - ppr (Kernels prog) = PP.ppr prog + ppr (GPU prog) = PP.ppr prog ppr (MC prog) = PP.ppr prog ppr (Seq prog) = PP.ppr prog ppr (SeqMem prog) = PP.ppr prog ppr (MCMem prog) = PP.ppr prog - ppr (KernelsMem prog) = PP.ppr prog + ppr (GPUMem prog) = PP.ppr prog newtype UntypedPass = UntypedPass @@ -138,31 +140,31 @@ data UntypedAction = SOACSAction (Action SOACS.SOACS) - | KernelsAction (Action Kernels.Kernels) - | KernelsMemAction (FilePath -> Action KernelsMem.KernelsMem) + | GPUAction (Action GPU.GPU) + | GPUMemAction (FilePath -> Action GPUMem.GPUMem) | MCMemAction (FilePath -> Action MCMem.MCMem) | SeqMemAction (FilePath -> Action SeqMem.SeqMem) | PolyAction - ( forall lore. - ( ASTLore lore, - (CanBeAliased (Op lore)), - (OpMetrics (Op lore)) + ( forall rep. + ( ASTRep rep, + (CanBeAliased (Op rep)), + (OpMetrics (Op rep)) ) => - Action lore + Action rep ) untypedActionName :: UntypedAction -> String untypedActionName (SOACSAction a) = actionName a -untypedActionName (KernelsAction a) = actionName a +untypedActionName (GPUAction a) = actionName a untypedActionName (SeqMemAction a) = actionName $ a "" -untypedActionName (KernelsMemAction a) = actionName $ a "" +untypedActionName (GPUMemAction a) = actionName $ a "" untypedActionName (MCMemAction a) = actionName $ a "" untypedActionName (PolyAction a) = actionName (a :: Action SOACS.SOACS) instance Representation UntypedAction where representation (SOACSAction _) = "SOACS" - representation (KernelsAction _) = "Kernels" - representation (KernelsMemAction _) = "KernelsMem" + representation (GPUAction _) = "GPU" + representation (GPUMemAction _) = "GPUMem" representation (MCMemAction _) = "MCMem" representation (SeqMemAction _) = "SeqMem" representation PolyAction {} = "" @@ -194,13 +196,13 @@ kernelsMemProg :: String -> UntypedPassState -> - FutharkM (Prog KernelsMem.KernelsMem) -kernelsMemProg _ (KernelsMem prog) = + FutharkM (Prog GPUMem.GPUMem) +kernelsMemProg _ (GPUMem prog) = return prog kernelsMemProg name rep = externalErrorS $ "Pass " ++ name - ++ " expects KernelsMem representation, but got " + ++ " expects GPUMem representation, but got " ++ representation rep soacsProg :: String -> UntypedPassState -> FutharkM (Prog SOACS.SOACS) @@ -212,18 +214,18 @@ ++ " expects SOACS representation, but got " ++ representation rep -kernelsProg :: String -> UntypedPassState -> FutharkM (Prog Kernels.Kernels) -kernelsProg _ (Kernels prog) = +kernelsProg :: String -> UntypedPassState -> FutharkM (Prog GPU.GPU) +kernelsProg _ (GPU prog) = return prog kernelsProg name rep = externalErrorS $ - "Pass " ++ name ++ " expects Kernels representation, but got " ++ representation rep + "Pass " ++ name ++ " expects GPU representation, but got " ++ representation rep typedPassOption :: - Checkable tolore => - (String -> UntypedPassState -> FutharkM (Prog fromlore)) -> - (Prog tolore -> UntypedPassState) -> - Pass fromlore tolore -> + Checkable torep => + (String -> UntypedPassState -> FutharkM (Prog fromrep)) -> + (Prog torep -> UntypedPassState) -> + Pass fromrep torep -> String -> FutharkOption typedPassOption getProg putProg pass short = @@ -240,18 +242,18 @@ typedPassOption soacsProg SOACS kernelsPassOption :: - Pass Kernels.Kernels Kernels.Kernels -> + Pass GPU.GPU GPU.GPU -> String -> FutharkOption kernelsPassOption = - typedPassOption kernelsProg Kernels + typedPassOption kernelsProg GPU kernelsMemPassOption :: - Pass KernelsMem.KernelsMem KernelsMem.KernelsMem -> + Pass GPUMem.GPUMem GPUMem.GPUMem -> String -> FutharkOption kernelsMemPassOption = - typedPassOption kernelsMemProg KernelsMem + typedPassOption kernelsMemProg GPUMem simplifyOption :: String -> FutharkOption simplifyOption short = @@ -259,16 +261,16 @@ where perform (SOACS prog) config = SOACS <$> runPipeline (onePass simplifySOACS) config prog - perform (Kernels prog) config = - Kernels <$> runPipeline (onePass simplifyKernels) config prog + perform (GPU prog) config = + GPU <$> runPipeline (onePass simplifyGPU) config prog perform (MC prog) config = MC <$> runPipeline (onePass simplifyMC) config prog perform (Seq prog) config = Seq <$> runPipeline (onePass simplifySeq) config prog perform (SeqMem prog) config = SeqMem <$> runPipeline (onePass simplifySeqMem) config prog - perform (KernelsMem prog) config = - KernelsMem <$> runPipeline (onePass simplifyKernelsMem) config prog + perform (GPUMem prog) config = + GPUMem <$> runPipeline (onePass simplifyGPUMem) config prog perform (MCMem prog) config = MCMem <$> runPipeline (onePass simplifyMCMem) config prog @@ -279,9 +281,9 @@ allocateOption short = passOption (passDescription pass) (UntypedPass perform) short long where - perform (Kernels prog) config = - KernelsMem - <$> runPipeline (onePass Kernels.explicitAllocations) config prog + perform (GPU prog) config = + GPUMem + <$> runPipeline (onePass GPU.explicitAllocations) config prog perform (Seq prog) config = SeqMem <$> runPipeline (onePass Seq.explicitAllocations) config prog @@ -296,9 +298,9 @@ iplOption short = passOption (passDescription pass) (UntypedPass perform) short long where - perform (Kernels prog) config = - Kernels - <$> runPipeline (onePass inPlaceLoweringKernels) config prog + perform (GPU prog) config = + GPU + <$> runPipeline (onePass inPlaceLoweringGPU) config prog perform (Seq prog) config = Seq <$> runPipeline (onePass inPlaceLoweringSeq) config prog @@ -315,16 +317,16 @@ where perform (SOACS prog) config = SOACS <$> runPipeline (onePass $ performCSE True) config prog - perform (Kernels prog) config = - Kernels <$> runPipeline (onePass $ performCSE True) config prog + perform (GPU prog) config = + GPU <$> runPipeline (onePass $ performCSE True) config prog perform (MC prog) config = MC <$> runPipeline (onePass $ performCSE True) config prog perform (Seq prog) config = Seq <$> runPipeline (onePass $ performCSE True) config prog perform (SeqMem prog) config = SeqMem <$> runPipeline (onePass $ performCSE False) config prog - perform (KernelsMem prog) config = - KernelsMem <$> runPipeline (onePass $ performCSE False) config prog + perform (GPUMem prog) config = + GPUMem <$> runPipeline (onePass $ performCSE False) config prog perform (MCMem prog) config = MCMem <$> runPipeline (onePass $ performCSE False) config prog @@ -332,11 +334,11 @@ pass = performCSE True :: Pass SOACS.SOACS SOACS.SOACS pipelineOption :: - (UntypedPassState -> Maybe (Prog fromlore)) -> + (UntypedPassState -> Maybe (Prog fromrep)) -> String -> - (Prog tolore -> UntypedPassState) -> + (Prog torep -> UntypedPassState) -> String -> - Pipeline fromlore tolore -> + Pipeline fromrep torep -> String -> [String] -> FutharkOption @@ -413,7 +415,7 @@ ["compile-imperative-kernels"] ( NoArg $ Right $ \opts -> - opts {futharkAction = KernelsMemAction $ const kernelImpCodeGenAction} + opts {futharkAction = GPUMemAction $ const kernelImpCodeGenAction} ) "Translate program into the imperative IL with kernels and write it on standard output.", Option @@ -429,7 +431,7 @@ ["compile-opencl"] ( NoArg $ Right $ \opts -> - opts {futharkAction = KernelsMemAction $ compileOpenCLAction newFutharkConfig ToExecutable} + opts {futharkAction = GPUMemAction $ compileOpenCLAction newFutharkConfig ToExecutable} ) "Compile the program using the OpenCL backend.", Option @@ -503,14 +505,15 @@ soacsPassOption inlineFunctions [], kernelsPassOption babysitKernels [], kernelsPassOption tileLoops [], - kernelsPassOption unstreamKernels [], - kernelsPassOption sinkKernels [], - typedPassOption soacsProg Kernels extractKernels [], + kernelsPassOption unstreamGPU [], + kernelsPassOption sinkGPU [], + typedPassOption soacsProg GPU extractKernels [], typedPassOption soacsProg MC extractMulticore [], iplOption [], allocateOption "a", - kernelsMemPassOption doubleBufferKernels [], + kernelsMemPassOption doubleBufferGPU [], kernelsMemPassOption expandAllocations [], + kernelsMemPassOption ReuseAllocations.optimise [], cseOption [], simplifyOption "e", soacsPipelineOption @@ -520,23 +523,23 @@ ["standard"], pipelineOption getSOACSProg - "Kernels" - Kernels + "GPU" + GPU "Run the default optimised kernels pipeline" kernelsPipeline [] ["kernels"], pipelineOption getSOACSProg - "KernelsMem" - KernelsMem + "GPUMem" + GPUMem "Run the full GPU compilation pipeline" gpuPipeline [] ["gpu"], pipelineOption getSOACSProg - "KernelsMem" + "GPUMem" SeqMem "Run the sequential CPU compilation pipeline" sequentialCpuPipeline @@ -637,7 +640,10 @@ input <- liftIO $ T.readFile file case parse file input of Left err -> externalErrorS $ T.unpack err - Right prog -> runPolyPasses config base $ construct prog + Right prog -> + case checkProg $ Alias.aliasAnalysis prog of + Left err -> externalErrorS $ show err + Right () -> runPolyPasses config base $ construct prog handlers = [ ( ".fut", @@ -648,8 +654,8 @@ (".fut_soacs", readCore parseSOACS SOACS), (".fut_seq", readCore parseSeq Seq), (".fut_seq_mem", readCore parseSeqMem SeqMem), - (".fut_kernels", readCore parseKernels Kernels), - (".fut_kernels_mem", readCore parseKernelsMem KernelsMem), + (".fut_kernels", readCore parseGPU GPU), + (".fut_kernels_mem", readCore parseGPUMem GPUMem), (".fut_mc", readCore parseMC MC), (".fut_mc_mem", readCore parseMCMem MCMem) ] @@ -675,23 +681,23 @@ case (end_prog, futharkAction config) of (SOACS prog, SOACSAction action) -> actionProcedure action prog - (Kernels prog, KernelsAction action) -> + (GPU prog, GPUAction action) -> actionProcedure action prog (SeqMem prog, SeqMemAction action) -> actionProcedure (action base) prog - (KernelsMem prog, KernelsMemAction action) -> + (GPUMem prog, GPUMemAction action) -> actionProcedure (action base) prog (MCMem prog, MCMemAction action) -> actionProcedure (action base) prog (SOACS soacs_prog, PolyAction acs) -> actionProcedure acs soacs_prog - (Kernels kernels_prog, PolyAction acs) -> + (GPU kernels_prog, PolyAction acs) -> actionProcedure acs kernels_prog (MC mc_prog, PolyAction acs) -> actionProcedure acs mc_prog (Seq seq_prog, PolyAction acs) -> actionProcedure acs seq_prog - (KernelsMem mem_prog, PolyAction acs) -> + (GPUMem mem_prog, PolyAction acs) -> actionProcedure acs mem_prog (SeqMem mem_prog, PolyAction acs) -> actionProcedure acs mem_prog diff --git a/src/Futhark/CLI/Literate.hs b/src/Futhark/CLI/Literate.hs --- a/src/Futhark/CLI/Literate.hs +++ b/src/Futhark/CLI/Literate.hs @@ -27,6 +27,7 @@ import qualified Data.Vector.Storable.ByteString as SVec import Data.Void import Data.Word (Word32, Word8) +import Futhark.Data import Futhark.Script import Futhark.Server import Futhark.Test @@ -345,19 +346,19 @@ in fromIntegral c :: Word8 valueToBMP :: Value -> Maybe LBS.ByteString -valueToBMP v@(Word32Value _ bytes) +valueToBMP v@(U32Value _ bytes) | [h, w] <- valueShape v = Just $ vecToBMP h w bytes -valueToBMP v@(Int32Value _ bytes) +valueToBMP v@(I32Value _ bytes) | [h, w] <- valueShape v = Just $ vecToBMP h w $ SVec.map fromIntegral bytes -valueToBMP v@(Float32Value _ bytes) +valueToBMP v@(F32Value _ bytes) | [h, w] <- valueShape v = Just $ vecToBMP h w $ greyFloatToImg bytes -valueToBMP v@(Word8Value _ bytes) +valueToBMP v@(U8Value _ bytes) | [h, w] <- valueShape v = Just $ vecToBMP h w $ greyByteToImg bytes -valueToBMP v@(Float64Value _ bytes) +valueToBMP v@(F64Value _ bytes) | [h, w] <- valueShape v = Just $ vecToBMP h w $ greyFloatToImg bytes valueToBMP v@(BoolValue _ bytes) @@ -454,7 +455,7 @@ b = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 2) a = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 3) in (a `shiftL` 24) .|. (r `shiftL` 16) .|. (g `shiftL` 8) .|. b - pure $ ValueAtom $ Word32Value shape $ SVec.generate (w * h) pix + pure $ ValueAtom $ U32Value shape $ SVec.generate (w * h) pix loadImage :: FilePath -> ScriptM (Compound Value) loadImage imgfile = @@ -878,8 +879,9 @@ let mdfile = fromMaybe (prog `replaceExtension` "md") $ scriptOutput opts imgdir = dropExtension mdfile <> "-img" run_options = scriptExtraOptions opts + cfg = futharkServerCfg ("." dropExtension prog) run_options - withScriptServer ("." dropExtension prog) run_options $ \server -> do + withScriptServer cfg $ \server -> do let env = Env { envServer = server, diff --git a/src/Futhark/CLI/Pkg.hs b/src/Futhark/CLI/Pkg.hs --- a/src/Futhark/CLI/Pkg.hs +++ b/src/Futhark/CLI/Pkg.hs @@ -344,7 +344,16 @@ [p] -> Just $ runPkgM cfg $ doCreate' $ T.pack p _ -> Nothing where + validPkgPath p = + not $ any (`elem` [".", ".."]) $ splitDirectories $ T.unpack p + doCreate' p = do + unless (validPkgPath p) . liftIO $ do + T.putStrLn $ "Not a valid package path: " <> p + T.putStrLn "Note: package paths are usually URIs." + T.putStrLn "Note: 'futhark init' is only needed when creating a package, not to use packages." + exitFailure + exists <- liftIO $ (||) <$> doesFileExist futharkPkg <*> doesDirectoryExist futharkPkg when exists $ liftIO $ do diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs --- a/src/Futhark/CLI/Test.hs +++ b/src/Futhark/CLI/Test.hs @@ -23,7 +23,6 @@ import Futhark.Util (atMostChars, fancyTerminal) import Futhark.Util.Console import Futhark.Util.Options -import Futhark.Util.Pretty (prettyText) import Futhark.Util.Table import System.Console.ANSI import qualified System.Console.Terminal.Size as Terminal @@ -97,7 +96,7 @@ "Running " <> T.pack (unwords $ binpath : extra_options) context prog_ctx $ - pureTestResults $ liftIO $ withServer to_run to_run_args f + pureTestResults $ liftIO $ withServer (futharkServerCfg to_run to_run_args) f data TestCase = TestCase { _testCaseMode :: TestMode, @@ -185,7 +184,7 @@ runInterpretedCase run@(TestRun _ inputValues _ index _) = unless (any (`elem` runTags run) ["compiled", "script"]) $ context ("Entry point: " <> entry <> "; dataset: " <> T.pack (runDescription run)) $ do - input <- T.unlines . map prettyText <$> getValues (FutharkExe futhark) dir inputValues + input <- T.unlines . map valueText <$> getValues (FutharkExe futhark) dir inputValues expectedResult' <- getExpectedResult (FutharkExe futhark) program entry run (code, output, err) <- liftIO $ @@ -271,12 +270,16 @@ runCompiledEntry :: FutharkExe -> Server -> FilePath -> InputOutputs -> IO [TestResult] runCompiledEntry futhark server program (InputOutputs entry run_cases) = do - Right output_types <- cmdOutputs server entry - Right input_types <- cmdInputs server entry - let outs = ["out" <> T.pack (show i) | i <- [0 .. length output_types -1]] - ins = ["in" <> T.pack (show i) | i <- [0 .. length input_types -1]] - onRes = either (Failure . pure) (const Success) - mapM (fmap onRes . runCompiledCase input_types outs ins) run_cases + output_types <- cmdOutputs server entry + input_types <- cmdInputs server entry + case (,) <$> output_types <*> input_types of + Left (CmdFailure _ err) -> + pure [Failure err] + Right (output_types', input_types') -> do + let outs = ["out" <> T.pack (show i) | i <- [0 .. length output_types' -1]] + ins = ["in" <> T.pack (show i) | i <- [0 .. length input_types' -1]] + onRes = either (Failure . pure) (const Success) + mapM (fmap onRes . runCompiledCase input_types' outs ins) run_cases where dir = takeDirectory program @@ -289,7 +292,7 @@ context1 case_ctx $ do expected <- getExpectedResult futhark program entry run - valuesAsVars server (zip ins input_types) futhark dir input_spec + valuesAsVars server (zip ins (map inputType input_types)) futhark dir input_spec call_r <- liftIO $ cmdCall server entry outs ins liftCommand $ cmdFree server ins @@ -299,7 +302,7 @@ pure $ ErrorResult $ T.unlines err Right _ -> SuccessResult - <$> readResults server outs program + <$> readResults server outs <* liftCommand (cmdFree server outs) compareResult entry index program expected res diff --git a/src/Futhark/CodeGen/Backends/CCUDA.hs b/src/Futhark/CodeGen/Backends/CCUDA.hs --- a/src/Futhark/CodeGen/Backends/CCUDA.hs +++ b/src/Futhark/CodeGen/Backends/CCUDA.hs @@ -20,7 +20,7 @@ import Futhark.CodeGen.Backends.GenericC.Options import Futhark.CodeGen.ImpCode.OpenCL import qualified Futhark.CodeGen.ImpGen.CUDA as ImpGen -import Futhark.IR.KernelsMem hiding +import Futhark.IR.GPUMem hiding ( CmpSizeLe, GetSize, GetSizeMax, @@ -29,7 +29,7 @@ import qualified Language.C.Quote.OpenCL as C -- | Compile the program to C with calls to CUDA. -compileProg :: MonadFreshNames m => Prog KernelsMem -> m (ImpGen.Warnings, GC.CParts) +compileProg :: MonadFreshNames m => Prog GPUMem -> m (ImpGen.Warnings, GC.CParts) compileProg prog = do (ws, Program cuda_code cuda_prelude kernels _ sizes failures prog') <- ImpGen.compileProg prog diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs --- a/src/Futhark/CodeGen/Backends/COpenCL.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL.hs @@ -19,7 +19,7 @@ import Futhark.CodeGen.Backends.GenericC.Options import Futhark.CodeGen.ImpCode.OpenCL import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen -import Futhark.IR.KernelsMem hiding +import Futhark.IR.GPUMem hiding ( CmpSizeLe, GetSize, GetSizeMax, @@ -29,7 +29,7 @@ import qualified Language.C.Syntax as C -- | Compile the program to C with calls to OpenCL. -compileProg :: MonadFreshNames m => Prog KernelsMem -> m (ImpGen.Warnings, GC.CParts) +compileProg :: MonadFreshNames m => Prog GPUMem -> m (ImpGen.Warnings, GC.CParts) compileProg prog = do ( ws, Program diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -99,9 +99,17 @@ import qualified Language.C.Quote.OpenCL as C import qualified Language.C.Syntax as C +-- How public an array type definition sould be. Public types show up +-- in the generated API, while private types are used only to +-- implement the members of opaques. +data Publicness = Private | Public + deriving (Eq, Ord, Show) + +type ArrayType = (Space, Signedness, PrimType, Int) + data CompilerState s = CompilerState - { compArrayStructs :: [((C.Type, Int), (C.Type, [C.Definition]))], - compOpaqueStructs :: [(String, (C.Type, [C.Definition]))], + { compArrayTypes :: M.Map ArrayType Publicness, + compOpaqueTypes :: M.Map String [ValueDesc], compEarlyDecls :: DL.DList C.Definition, compInit :: [C.Stm], compNameSrc :: VNameSource, @@ -118,8 +126,8 @@ newCompilerState :: VNameSource -> s -> CompilerState s newCompilerState src s = CompilerState - { compArrayStructs = [], - compOpaqueStructs = [], + { compArrayTypes = mempty, + compOpaqueTypes = mempty, compEarlyDecls = mempty, compInit = [], compNameSrc = src, @@ -318,10 +326,6 @@ envFatMemory :: CompilerEnv op s -> Bool envFatMemory = opsFatMemory . envOperations -arrayDefinitions, opaqueDefinitions :: CompilerState s -> [C.Definition] -arrayDefinitions = concatMap (snd . snd) . compArrayStructs -opaqueDefinitions = concatMap (snd . snd) . compOpaqueStructs - initDecls, arrayDecls, opaqueDecls, entryDecls, miscDecls :: CompilerState s -> [C.Definition] initDecls = concatMap (DL.toList . snd) . filter ((== InitDecl) . fst) . M.toList . compHeaderDecls arrayDecls = concatMap (DL.toList . snd) . filter (isArrayDecl . fst) . M.toList . compHeaderDecls @@ -795,14 +799,14 @@ |] arrayLibraryFunctions :: + Publicness -> Space -> PrimType -> Signedness -> - [DimSize] -> + Int -> CompilerM op s [C.Definition] -arrayLibraryFunctions space pt signed shape = do - let rank = length shape - pt' = signedPrimTypeToCType signed pt +arrayLibraryFunctions pub space pt signed rank = do + let pt' = signedPrimTypeToCType signed pt name = arrayName pt signed rank arr_name = "futhark_" ++ name array_type = [C.cty|struct $id:arr_name|] @@ -871,26 +875,23 @@ ctx_ty <- contextType ops <- asks envOperations - headerDecl - (ArrayDecl name) + let proto = case pub of + Public -> headerDecl (ArrayDecl name) + Private -> libDecl + + proto [C.cedecl|struct $id:arr_name;|] - headerDecl - (ArrayDecl name) + proto [C.cedecl|$ty:array_type* $id:new_array($ty:ctx_ty *ctx, const $ty:pt' *data, $params:shape_params);|] - headerDecl - (ArrayDecl name) + proto [C.cedecl|$ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, const $ty:memty data, int offset, $params:shape_params);|] - headerDecl - (ArrayDecl name) + proto [C.cedecl|int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - headerDecl - (ArrayDecl name) + proto [C.cedecl|int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data);|] - headerDecl - (ArrayDecl name) + proto [C.cedecl|$ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - headerDecl - (ArrayDecl name) + proto [C.cedecl|const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] return @@ -1033,6 +1034,9 @@ headerDecl (OpaqueDecl desc) + [C.cedecl|struct $id:name;|] + headerDecl + (OpaqueDecl desc) [C.cedecl|int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj);|] headerDecl (OpaqueDecl desc) @@ -1076,53 +1080,47 @@ } |] -valueDescToCType :: ValueDesc -> CompilerM op s C.Type -valueDescToCType (ScalarValue pt signed _) = +valueDescToCType :: Publicness -> ValueDesc -> CompilerM op s C.Type +valueDescToCType _ (ScalarValue pt signed _) = return $ signedPrimTypeToCType signed pt -valueDescToCType (ArrayValue mem space pt signed shape) = do - let pt' = signedPrimTypeToCType signed pt - rank = length shape - exists <- gets $ lookup (pt', rank) . compArrayStructs - case exists of - Just (cty, _) -> return cty - Nothing -> do - memty <- memToCType mem space - name <- publicName $ arrayName pt signed rank - let struct = [C.cedecl|struct $id:name { $ty:memty mem; typename int64_t shape[$int:rank]; };|] - stype = [C.cty|struct $id:name|] - library <- arrayLibraryFunctions space pt signed shape - modify $ \s -> - s - { compArrayStructs = - ((pt', rank), (stype, struct : library)) : compArrayStructs s - } - return stype +valueDescToCType pub (ArrayValue _ space pt signed shape) = do + let rank = length shape + name <- publicName $ arrayName pt signed rank + let add = M.insertWith max (space, signed, pt, rank) pub + modify $ \s -> s {compArrayTypes = add $ compArrayTypes s} + pure [C.cty|struct $id:name|] opaqueToCType :: String -> [ValueDesc] -> CompilerM op s C.Type opaqueToCType desc vds = do name <- publicName $ opaqueName desc vds - exists <- gets $ lookup name . compOpaqueStructs - case exists of - Just (ty, _) -> return ty - Nothing -> do - members <- zipWithM field vds [(0 :: Int) ..] - let struct = [C.cedecl|struct $id:name { $sdecls:members };|] - stype = [C.cty|struct $id:name|] - headerDecl (OpaqueDecl desc) [C.cedecl|struct $id:name;|] - library <- opaqueLibraryFunctions desc vds - modify $ \s -> - s - { compOpaqueStructs = - (name, (stype, struct : library)) : - compOpaqueStructs s - } - return stype + let add = M.insert desc vds + modify $ \s -> s {compOpaqueTypes = add $ compOpaqueTypes s} + -- Now ensure that the constituent array types will exist. + mapM_ (valueDescToCType Private) vds + pure [C.cty|struct $id:name|] + +generateAPITypes :: CompilerM op s () +generateAPITypes = do + mapM_ generateArray . M.toList =<< gets compArrayTypes + mapM_ generateOpaque . M.toList =<< gets compOpaqueTypes where + generateArray ((space, signed, pt, rank), pub) = do + name <- publicName $ arrayName pt signed rank + let memty = fatMemType space + libDecl [C.cedecl|struct $id:name { $ty:memty mem; typename int64_t shape[$int:rank]; };|] + mapM libDecl =<< arrayLibraryFunctions pub space pt signed rank + + generateOpaque (desc, vds) = do + name <- publicName $ opaqueName desc vds + members <- zipWithM field vds [(0 :: Int) ..] + libDecl [C.cedecl|struct $id:name { $sdecls:members };|] + mapM libDecl =<< opaqueLibraryFunctions desc vds + field vd@ScalarValue {} i = do - ct <- valueDescToCType vd + ct <- valueDescToCType Private vd return [C.csdecl|$ty:ct $id:(tupleField i);|] field vd i = do - ct <- valueDescToCType vd + ct <- valueDescToCType Private vd return [C.csdecl|$ty:ct *$id:(tupleField i);|] allTrue :: [C.Exp] -> C.Exp @@ -1136,35 +1134,35 @@ prepareEntryInputs args = collect' $ zipWithM prepare [(0 :: Int) ..] args where arg_names = namesFromList $ concatMap evNames args - evNames (OpaqueValue _ vds) = map vdName vds - evNames (TransparentValue vd) = [vdName vd] + evNames (OpaqueValue _ _ vds) = map vdName vds + evNames (TransparentValue _ vd) = [vdName vd] vdName (ArrayValue v _ _ _ _) = v vdName (ScalarValue _ _ v) = v - prepare pno (TransparentValue vd) = do + prepare pno (TransparentValue _ vd) = do let pname = "in" ++ show pno - (ty, check) <- prepareValue [C.cexp|$id:pname|] vd + (ty, check) <- prepareValue Public [C.cexp|$id:pname|] vd return ( [C.cparam|const $ty:ty $id:pname|], allTrue check ) - prepare pno (OpaqueValue desc vds) = do + prepare pno (OpaqueValue _ desc vds) = do ty <- opaqueToCType desc vds let pname = "in" ++ show pno field i ScalarValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] field i ArrayValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] - checks <- map snd <$> zipWithM prepareValue (zipWith field [0 ..] vds) vds + checks <- map snd <$> zipWithM (prepareValue Private) (zipWith field [0 ..] vds) vds return ( [C.cparam|const $ty:ty *$id:pname|], allTrue $ concat checks ) - prepareValue src (ScalarValue pt signed name) = do + prepareValue _ src (ScalarValue pt signed name) = do let pt' = signedPrimTypeToCType signed pt stm [C.cstm|$id:name = $exp:src;|] return (pt', []) - prepareValue src vd@(ArrayValue mem _ _ _ shape) = do - ty <- valueDescToCType vd + prepareValue pub src vd@(ArrayValue mem _ _ _ shape) = do + ty <- valueDescToCType pub vd stm [C.cstm|$exp:mem = $exp:src->mem;|] @@ -1188,9 +1186,9 @@ prepareEntryOutputs :: [ExternalValue] -> CompilerM op s ([C.Param], [C.BlockItem]) prepareEntryOutputs = collect' . zipWithM prepare [(0 :: Int) ..] where - prepare pno (TransparentValue vd) = do + prepare pno (TransparentValue _ vd) = do let pname = "out" ++ show pno - ty <- valueDescToCType vd + ty <- valueDescToCType Public vd case vd of ArrayValue {} -> do @@ -1200,10 +1198,10 @@ ScalarValue {} -> do prepareValue [C.cexp|*$id:pname|] vd return [C.cparam|$ty:ty *$id:pname|] - prepare pno (OpaqueValue desc vds) = do + prepare pno (OpaqueValue _ desc vds) = do let pname = "out" ++ show pno ty <- opaqueToCType desc vds - vd_ts <- mapM valueDescToCType vds + vd_ts <- mapM (valueDescToCType Private) vds stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] @@ -1338,7 +1336,12 @@ disableWarnings = pretty [C.cunit| -$esc:("#ifdef __GNUC__") +$esc:("#ifdef __clang__") +$esc:("#pragma clang diagnostic ignored \"-Wunused-function\"") +$esc:("#pragma clang diagnostic ignored \"-Wunused-variable\"") +$esc:("#pragma clang diagnostic ignored \"-Wparentheses\"") +$esc:("#pragma clang diagnostic ignored \"-Wunused-label\"") +$esc:("#elif __GNUC__") $esc:("#pragma GCC diagnostic ignored \"-Wunused-function\"") $esc:("#pragma GCC diagnostic ignored \"-Wunused-variable\"") $esc:("#pragma GCC diagnostic ignored \"-Wparentheses\"") @@ -1346,12 +1349,6 @@ $esc:("#pragma GCC diagnostic ignored \"-Wunused-but-set-variable\"") $esc:("#endif") -$esc:("#ifdef __clang__") -$esc:("#pragma clang diagnostic ignored \"-Wunused-function\"") -$esc:("#pragma clang diagnostic ignored \"-Wunused-variable\"") -$esc:("#pragma clang diagnostic ignored \"-Wparentheses\"") -$esc:("#pragma clang diagnostic ignored \"-Wunused-label\"") -$esc:("#endif") |] -- | Produce header and implementation files. @@ -1467,11 +1464,7 @@ $edecls:lib_decls -$edecls:(map funcToDef definitions) - -$edecls:(arrayDefinitions endstate) - -$edecls:(opaqueDefinitions endstate) +$edecls:definitions $edecls:entry_point_decls |] @@ -1495,7 +1488,7 @@ ctx_ty <- contextType - (prototypes, definitions) <- + (prototypes, functions) <- unzip <$> mapM (compileFun get_consts [[C.cparam|$ty:ctx_ty *ctx|]]) funs mapM_ earlyDecl memstructs @@ -1508,7 +1501,7 @@ commonLibFuns memreport - return (prototypes, definitions, entry_points) + return (prototypes, map funcToDef functions, entry_points) funcToDef func = C.FuncDef func loc where @@ -1527,6 +1520,7 @@ commonLibFuns :: [C.BlockItem] -> CompilerM op s () commonLibFuns memreport = do + generateAPITypes ctx <- contextType ops <- asks envOperations profilereport <- gets $ DL.toList . compProfileItems diff --git a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs --- a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs @@ -156,8 +156,8 @@ in [C.cty|struct $id:name|] externalValueToCType :: ExternalValue -> C.Type -externalValueToCType (TransparentValue vd) = valueDescToCType vd -externalValueToCType (OpaqueValue desc vds) = opaqueToCType desc vds +externalValueToCType (TransparentValue _ vd) = valueDescToCType vd +externalValueToCType (OpaqueValue _ desc vds) = opaqueToCType desc vds primTypeInfo :: PrimType -> Signedness -> C.Exp primTypeInfo (IntType it) t = case (it, t) of @@ -184,14 +184,14 @@ }|] readInput :: Int -> ExternalValue -> ([C.BlockItem], C.Stm, C.Stm, C.Stm, C.Exp) -readInput i (OpaqueValue desc _) = +readInput i (OpaqueValue _ desc _) = ( [C.citems|futhark_panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:desc);|], [C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|NULL|] ) -readInput i (TransparentValue (ScalarValue t ept _)) = +readInput i (TransparentValue _ (ScalarValue t ept _)) = let dest = "read_value_" ++ show i in ( [C.citems|$ty:(primTypeToCType t) $id:dest; $stm:(readPrimStm dest i t ept);|], @@ -200,7 +200,7 @@ [C.cstm|;|], [C.cexp|$id:dest|] ) -readInput i (TransparentValue (ArrayValue _ _ t ept dims)) = +readInput i (TransparentValue _ (ArrayValue _ _ t ept dims)) = let dest = "read_value_" ++ show i shape = "read_shape_" ++ show i arr = "read_arr_" ++ show i @@ -252,19 +252,19 @@ result = "result_" ++ show i case ev of - TransparentValue ScalarValue {} -> + TransparentValue _ ScalarValue {} -> ( [C.citem|$ty:ty $id:result;|], [C.cexp|$id:result|], [C.cstm|;|] ) - TransparentValue (ArrayValue _ _ t ept dims) -> + TransparentValue _ (ArrayValue _ _ t ept dims) -> let name = arrayName t ept $ length dims free_array = "futhark_free_" ++ name in ( [C.citem|$ty:ty *$id:result;|], [C.cexp|$id:result|], [C.cstm|assert($id:free_array(ctx, $id:result) == 0);|] ) - OpaqueValue desc vds -> + OpaqueValue _ desc vds -> let free_opaque = "futhark_free_" ++ opaqueName desc vds in ( [C.citem|$ty:ty *$id:result;|], [C.cexp|$id:result|], @@ -277,11 +277,11 @@ -- | Return a statement printing the given external value. printStm :: ExternalValue -> C.Exp -> C.Stm -printStm (OpaqueValue desc _) _ = +printStm (OpaqueValue _ desc _) _ = [C.cstm|printf("#", $string:desc);|] -printStm (TransparentValue (ScalarValue bt ept _)) e = +printStm (TransparentValue _ (ScalarValue bt ept _)) e = printPrimStm [C.cexp|stdout|] e bt ept -printStm (TransparentValue (ArrayValue _ _ bt ept shape)) e = +printStm (TransparentValue _ (ArrayValue _ _ bt ept shape)) e = let values_array = "futhark_values_" ++ name shape_array = "futhark_shape_" ++ name num_elems = cproduct [[C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0 .. rank -1]] diff --git a/src/Futhark/CodeGen/Backends/GenericC/Server.hs b/src/Futhark/CodeGen/Backends/GenericC/Server.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Server.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Server.hs @@ -101,22 +101,22 @@ ] typeStructName :: ExternalValue -> String -typeStructName (TransparentValue (ScalarValue pt signed _)) = +typeStructName (TransparentValue _ (ScalarValue pt signed _)) = let name = prettySigned (signed == TypeUnsigned) pt in "type_" ++ name -typeStructName (TransparentValue (ArrayValue _ _ pt signed shape)) = +typeStructName (TransparentValue _ (ArrayValue _ _ pt signed shape)) = let rank = length shape name = arrayName pt signed rank in "type_" ++ name -typeStructName (OpaqueValue name vds) = +typeStructName (OpaqueValue _ name vds) = "type_" ++ opaqueName name vds valueDescBoilerplate :: ExternalValue -> (String, (C.Initializer, [C.Definition])) -valueDescBoilerplate ev@(TransparentValue (ScalarValue pt signed _)) = +valueDescBoilerplate ev@(TransparentValue _ (ScalarValue pt signed _)) = let name = prettySigned (signed == TypeUnsigned) pt type_name = typeStructName ev in (name, ([C.cinit|&$id:type_name|], mempty)) -valueDescBoilerplate ev@(TransparentValue (ArrayValue _ _ pt signed shape)) = +valueDescBoilerplate ev@(TransparentValue _ (ArrayValue _ _ pt signed shape)) = let rank = length shape name = arrayName pt signed rank pt_name = prettySigned (signed == TypeUnsigned) pt @@ -156,7 +156,7 @@ };|] ) ) -valueDescBoilerplate ev@(OpaqueValue name vds) = +valueDescBoilerplate ev@(OpaqueValue _ name vds) = let type_name = typeStructName ev aux_name = type_name ++ "_aux" opaque_free = "futhark_free_" ++ opaqueName name vds @@ -199,6 +199,8 @@ in_types = functionArgs fun out_types_name = pretty name ++ "_out_types" in_types_name = pretty name ++ "_in_types" + out_unique_name = pretty name ++ "_out_unique" + in_unique_name = pretty name ++ "_in_unique" (out_items, out_args) | null out_types = ([C.citems|(void)outs;|], mempty) | otherwise = unzip $ zipWith loadOut [0 ..] out_types @@ -211,10 +213,16 @@ $inits:(map typeStructInit out_types), NULL }; + bool $id:out_unique_name[] = { + $inits:(map typeUniqueInit out_types) + }; struct type* $id:in_types_name[] = { $inits:(map typeStructInit in_types), NULL }; + bool $id:in_unique_name[] = { + $inits:(map typeUniqueInit in_types) + }; int $id:call_f(struct futhark_context *ctx, void **outs, void **ins) { $items:out_items $items:in_items @@ -225,11 +233,20 @@ .name = $string:(pretty ename), .f = $id:call_f, .in_types = $id:in_types_name, - .out_types = $id:out_types_name + .out_types = $id:out_types_name, + .in_unique = $id:in_unique_name, + .out_unique = $id:out_unique_name }|] ) where typeStructInit t = [C.cinit|&$id:(typeStructName t)|] + typeUniqueInit t = + case typeUnique t of + Unique -> [C.cinit|true|] + Nonunique -> [C.cinit|false|] + + typeUnique (TransparentValue u _) = u + typeUnique (OpaqueValue u _ _) = u loadOut i ev = let v = "out" ++ show (i :: Int) diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs --- a/src/Futhark/CodeGen/Backends/GenericPython.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython.hs @@ -59,6 +59,7 @@ import Futhark.IR.Syntax (Space (..)) import Futhark.MonadFreshNames import Futhark.Util (zEncodeString) +import Futhark.Util.Pretty (pretty) -- | A substitute expression compiler, tried before the main -- compilation function. @@ -330,7 +331,7 @@ filter (isJust . Imp.functionEntry . snd) funs where evd Imp.TransparentValue {} = mempty - evd (Imp.OpaqueValue name vds) = + evd (Imp.OpaqueValue _ name vds) = M.singleton name $ map (String . vd) vds vd (Imp.ScalarValue pt s _) = readTypeEnum pt s @@ -538,18 +539,18 @@ ] entryPointOutput :: Imp.ExternalValue -> CompilerM op s PyExp -entryPointOutput (Imp.OpaqueValue desc vs) = +entryPointOutput (Imp.OpaqueValue u desc vs) = simpleCall "opaque" . (String (pretty desc) :) - <$> mapM (entryPointOutput . Imp.TransparentValue) vs -entryPointOutput (Imp.TransparentValue (Imp.ScalarValue bt ept name)) = do + <$> mapM (entryPointOutput . Imp.TransparentValue u) vs +entryPointOutput (Imp.TransparentValue _ (Imp.ScalarValue bt ept name)) = do name' <- compileVar name return $ simpleCall tf [name'] where tf = compilePrimToExtNp bt ept -entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do +entryPointOutput (Imp.TransparentValue _ (Imp.ArrayValue mem (Imp.Space sid) bt ept dims)) = do pack_output <- asks envEntryOutput pack_output mem sid bt ept dims -entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ bt ept dims)) = do +entryPointOutput (Imp.TransparentValue _ (Imp.ArrayValue mem _ bt ept dims)) = do mem' <- compileVar mem let cast = Cast mem' (compilePrimTypeExt bt ept) dims' <- mapM compileDim dims @@ -615,14 +616,14 @@ declEntryPointInputSizes :: [Imp.ExternalValue] -> CompilerM op s () declEntryPointInputSizes = mapM_ onSize . concatMap sizes where - sizes (Imp.TransparentValue v) = valueSizes v - sizes (Imp.OpaqueValue _ vs) = concatMap valueSizes vs + sizes (Imp.TransparentValue _ v) = valueSizes v + sizes (Imp.OpaqueValue _ _ vs) = concatMap valueSizes vs valueSizes (Imp.ArrayValue _ _ _ _ dims) = subExpVars dims valueSizes Imp.ScalarValue {} = [] onSize v = stm $ Assign (Var (compileName v)) None entryPointInput :: (Int, Imp.ExternalValue, PyExp) -> CompilerM op s () -entryPointInput (i, Imp.OpaqueValue desc vs, e) = do +entryPointInput (i, Imp.OpaqueValue u desc vs, e) = do let type_is_ok = BinOp "and" @@ -630,9 +631,9 @@ (BinOp "==" (Field e "desc") (String desc)) stm $ If (UnOp "not" type_is_ok) [badInput i e desc] [] mapM_ entryPointInput $ - zip3 (repeat i) (map Imp.TransparentValue vs) $ + zip3 (repeat i) (map (Imp.TransparentValue u) vs) $ map (Index (Field e "data") . IdxExp . Integer) [0 ..] -entryPointInput (i, Imp.TransparentValue (Imp.ScalarValue bt s name), e) = do +entryPointInput (i, Imp.TransparentValue _ (Imp.ScalarValue bt s name), e) = do vname' <- compileVar name let -- HACK: A Numpy int64 will signal an OverflowError if we pass -- it a number bigger than 2**63. This does not happen if we @@ -650,7 +651,7 @@ (Tuple [Var "TypeError", Var "AssertionError"]) [badInput i e $ prettySigned (s == Imp.TypeUnsigned) bt] ] -entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do +entryPointInput (i, Imp.TransparentValue _ (Imp.ArrayValue mem (Imp.Space sid) bt ept dims), e) = do unpack_input <- asks envEntryInput mem' <- compileVar mem unpack <- collect $ unpack_input mem' sid bt ept dims e @@ -664,7 +665,7 @@ ++ prettySigned (ept == Imp.TypeUnsigned) bt ] ] -entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem _ t s dims), e) = do +entryPointInput (i, Imp.TransparentValue _ (Imp.ArrayValue mem _ t s dims), e) = do let type_is_wrong = UnOp "not" $ BinOp "in" (simpleCall "type" [e]) $ List [Var "np.ndarray"] let dtype_is_wrong = UnOp "not" $ BinOp "==" (Field e "dtype") $ Var $ compilePrimToExtNp t s let dim_is_wrong = UnOp "not" $ BinOp "==" (Field e "ndim") $ Integer $ toInteger $ length dims @@ -700,9 +701,9 @@ stm $ Assign dest unwrap_call extValueDescName :: Imp.ExternalValue -> String -extValueDescName (Imp.TransparentValue v) = extName $ valueDescName v -extValueDescName (Imp.OpaqueValue desc []) = extName $ zEncodeString desc -extValueDescName (Imp.OpaqueValue desc (v : _)) = +extValueDescName (Imp.TransparentValue _ v) = extName $ valueDescName v +extValueDescName (Imp.OpaqueValue _ desc []) = extName $ zEncodeString desc +extValueDescName (Imp.OpaqueValue _ desc (v : _)) = extName $ zEncodeString desc ++ "_" ++ pretty (baseTag (valueDescVName v)) extName :: String -> String @@ -731,15 +732,15 @@ readTypeEnum Unit _ = "bool" readInput :: Imp.ExternalValue -> PyStmt -readInput (Imp.OpaqueValue desc _) = +readInput (Imp.OpaqueValue _ desc _) = Raise $ simpleCall "Exception" [String $ "Cannot read argument of type " ++ desc ++ "."] -readInput decl@(Imp.TransparentValue (Imp.ScalarValue bt ept _)) = +readInput decl@(Imp.TransparentValue _ (Imp.ScalarValue bt ept _)) = let type_name = readTypeEnum bt ept in Assign (Var $ extValueDescName decl) $ simpleCall "read_value" [String type_name] -readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ bt ept dims)) = +readInput decl@(Imp.TransparentValue _ (Imp.ArrayValue _ _ bt ept dims)) = let type_name = readTypeEnum bt ept in Assign (Var $ extValueDescName decl) $ simpleCall @@ -754,17 +755,17 @@ -- that returns an equivalent Numpy array. This works for PyOpenCL, -- but we will probably need yet another plugin mechanism here in -- the future. - printValue' (Imp.OpaqueValue desc _) _ = + printValue' (Imp.OpaqueValue _ desc _) _ = return [ Exp $ simpleCall "sys.stdout.write" [String $ "#"] ] - printValue' (Imp.TransparentValue (Imp.ArrayValue mem (Space _) bt ept shape)) e = - printValue' (Imp.TransparentValue (Imp.ArrayValue mem DefaultSpace bt ept shape)) $ + printValue' (Imp.TransparentValue u (Imp.ArrayValue mem (Space _) bt ept shape)) e = + printValue' (Imp.TransparentValue u (Imp.ArrayValue mem DefaultSpace bt ept shape)) $ simpleCall (pretty e ++ ".get") [] - printValue' (Imp.TransparentValue _) e = + printValue' (Imp.TransparentValue _ _) e = return [ Exp $ Call @@ -919,10 +920,10 @@ map desc $ Imp.functionResult func ) where - desc (Imp.OpaqueValue d _) = d - desc (Imp.TransparentValue (Imp.ScalarValue pt s _)) = readTypeEnum pt s - desc (Imp.TransparentValue (Imp.ArrayValue _ _ pt s dims)) = - concat (replicate (length dims) "[]") ++ readTypeEnum pt s + desc (Imp.OpaqueValue u d _) = pretty u <> d + desc (Imp.TransparentValue u (Imp.ScalarValue pt s _)) = pretty u <> readTypeEnum pt s + desc (Imp.TransparentValue u (Imp.ArrayValue _ _ pt s dims)) = + pretty u <> concat (replicate (length dims) "[]") <> readTypeEnum pt s callEntryFun :: [PyStmt] -> diff --git a/src/Futhark/CodeGen/Backends/MulticoreC.hs b/src/Futhark/CodeGen/Backends/MulticoreC.hs --- a/src/Futhark/CodeGen/Backends/MulticoreC.hs +++ b/src/Futhark/CodeGen/Backends/MulticoreC.hs @@ -647,9 +647,9 @@ atomicOps aop doAtomic :: - (C.ToIdent a1, C.ToIdent a2) => + (C.ToIdent a1) => a1 -> - a2 -> + VName -> Count u (TExp Int32) -> Exp -> String -> @@ -658,15 +658,17 @@ doAtomic old arr ind val op ty = do ind' <- GC.compileExp $ untyped $ unCount ind val' <- GC.compileExp val - GC.stm [C.cstm|$id:old = $id:op(&(($ty:ty*)$id:arr.mem)[$exp:ind'], ($ty:ty) $exp:val', __ATOMIC_RELAXED);|] + arr' <- GC.rawMem arr + GC.stm [C.cstm|$id:old = $id:op(&(($ty:ty*)$exp:arr')[$exp:ind'], ($ty:ty) $exp:val', __ATOMIC_RELAXED);|] atomicOps :: AtomicOp -> GC.CompilerM op s () atomicOps (AtomicCmpXchg t old arr ind res val) = do ind' <- GC.compileExp $ untyped $ unCount ind new_val' <- GC.compileExp val let cast = [C.cty|$ty:(GC.primTypeToCType t)*|] + arr' <- GC.rawMem arr GC.stm - [C.cstm|$id:res = $id:op(&(($ty:cast)$id:arr.mem)[$exp:ind'], + [C.cstm|$id:res = $id:op(&(($ty:cast)$exp:arr')[$exp:ind'], ($ty:cast)&$id:old, $exp:new_val', 0, __ATOMIC_SEQ_CST, __ATOMIC_RELAXED);|] diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs --- a/src/Futhark/CodeGen/Backends/PyOpenCL.hs +++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs @@ -15,16 +15,17 @@ import Futhark.CodeGen.Backends.PyOpenCL.Boilerplate import qualified Futhark.CodeGen.ImpCode.OpenCL as Imp import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen -import Futhark.IR.KernelsMem (KernelsMem, Prog) +import Futhark.IR.GPUMem (GPUMem, Prog) import Futhark.MonadFreshNames import Futhark.Util (zEncodeString) +import Futhark.Util.Pretty (pretty) -- | Compile the program to Python with calls to OpenCL. compileProg :: MonadFreshNames m => Py.CompilerMode -> String -> - Prog KernelsMem -> + Prog GPUMem -> m (ImpGen.Warnings, String) compileProg mode class_name prog = do ( ws, diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs @@ -26,7 +26,7 @@ untyped, ) import Futhark.CodeGen.OpenCL.Heuristics -import Futhark.Util.Pretty (prettyText) +import Futhark.Util.Pretty (pretty, prettyText) import NeatInterpolation (text) errorMsgNumArgs :: ErrorMsg a -> Int diff --git a/src/Futhark/CodeGen/Backends/SimpleRep.hs b/src/Futhark/CodeGen/Backends/SimpleRep.hs --- a/src/Futhark/CodeGen/Backends/SimpleRep.hs +++ b/src/Futhark/CodeGen/Backends/SimpleRep.hs @@ -127,11 +127,11 @@ -- | The type used to expose a Futhark value in the C API. A pointer -- in the case of arrays and opaques. externalValueType :: ExternalValue -> C.Type -externalValueType (OpaqueValue desc vds) = +externalValueType (OpaqueValue _ desc vds) = [C.cty|struct $id:("futhark_" ++ opaqueName desc vds)*|] -externalValueType (TransparentValue (ArrayValue _ _ pt signed shape)) = +externalValueType (TransparentValue _ (ArrayValue _ _ pt signed shape)) = [C.cty|struct $id:("futhark_" ++ arrayName pt signed (length shape))*|] -externalValueType (TransparentValue (ScalarValue pt signed _)) = +externalValueType (TransparentValue _ (ScalarValue pt signed _)) = signedPrimTypeToCType signed pt -- | Return an expression multiplying together the given expressions. 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 @@ -52,10 +52,11 @@ withElemType, -- * Re-exports from other modules. + pretty, module Language.Futhark.Core, module Futhark.IR.Primitive, module Futhark.Analysis.PrimExp, - module Futhark.IR.Kernels.Sizes, + module Futhark.IR.GPU.Sizes, module Futhark.IR.Prop.Names, ) where @@ -65,7 +66,7 @@ import qualified Data.Set as S import Data.Traversable import Futhark.Analysis.PrimExp -import Futhark.IR.Kernels.Sizes (Count (..)) +import Futhark.IR.GPU.Sizes (Count (..)) import Futhark.IR.Pretty () import Futhark.IR.Primitive import Futhark.IR.Prop.Names @@ -131,7 +132,7 @@ data Signedness = TypeUnsigned | TypeDirect - deriving (Eq, Show) + deriving (Eq, Ord, Show) -- | A description of an externally meaningful value. data ValueDesc @@ -144,12 +145,14 @@ -- | ^ An externally visible value. This can be an opaque value -- (covering several physical internal values), or a single value that --- can be used externally. +-- can be used externally. We record the uniqueness because it is +-- important to the external interface as well. data ExternalValue - = -- | The string is a human-readable description - -- with no other semantics. - OpaqueValue String [ValueDesc] - | TransparentValue ValueDesc + = -- | The string is a human-readable description with no other + -- semantics. + -- not matter. + OpaqueValue Uniqueness String [ValueDesc] + | TransparentValue Uniqueness ValueDesc deriving (Show) -- | A imperative function, containing the body as well as its @@ -438,9 +441,9 @@ TypeDirect -> mempty instance Pretty ExternalValue where - ppr (TransparentValue v) = ppr v - ppr (OpaqueValue desc vs) = - text "opaque" <+> text desc + ppr (TransparentValue u v) = ppr u <> ppr v + ppr (OpaqueValue u desc vs) = + ppr u <> text "opaque" <+> text desc <+> nestedBlock "{" "}" (stack $ map ppr vs) instance Pretty ArrayContents where @@ -625,8 +628,8 @@ freeIn' ScalarValue {} = mempty instance FreeIn ExternalValue where - freeIn' (TransparentValue vd) = freeIn' vd - freeIn' (OpaqueValue _ vds) = foldMap freeIn' vds + freeIn' (TransparentValue _ vd) = freeIn' vd + freeIn' (OpaqueValue _ _ vds) = foldMap freeIn' vds instance FreeIn a => FreeIn (Code a) where freeIn' (x :>>: y) = diff --git a/src/Futhark/CodeGen/ImpCode/GPU.hs b/src/Futhark/CodeGen/ImpCode/GPU.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpCode/GPU.hs @@ -0,0 +1,266 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Variation of "Futhark.CodeGen.ImpCode" that contains the notion +-- of a kernel invocation. +module Futhark.CodeGen.ImpCode.GPU + ( Program, + Function, + FunctionT (Function), + Code, + KernelCode, + KernelConst (..), + KernelConstExp, + HostOp (..), + KernelOp (..), + Fence (..), + AtomicOp (..), + Kernel (..), + KernelUse (..), + module Futhark.CodeGen.ImpCode, + module Futhark.IR.GPU.Sizes, + ) +where + +import Futhark.CodeGen.ImpCode hiding (Code, Function) +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.IR.GPU.Sizes +import Futhark.IR.Pretty () +import Futhark.Util.Pretty + +-- | A program that calls kernels. +type Program = Imp.Definitions HostOp + +-- | A function that calls kernels. +type Function = Imp.Function HostOp + +-- | Host-level code that can call kernels. +type Code = Imp.Code HostOp + +-- | Code inside a kernel. +type KernelCode = Imp.Code KernelOp + +-- | A run-time constant related to kernels. +newtype KernelConst = SizeConst Name + deriving (Eq, Ord, Show) + +-- | An expression whose variables are kernel constants. +type KernelConstExp = PrimExp KernelConst + +-- | An operation that runs on the host (CPU). +data HostOp + = CallKernel Kernel + | GetSize VName Name SizeClass + | CmpSizeLe VName Name SizeClass Imp.Exp + | GetSizeMax VName SizeClass + deriving (Show) + +-- | A generic kernel containing arbitrary kernel code. +data Kernel = Kernel + { kernelBody :: Imp.Code KernelOp, + -- | The host variables referenced by the kernel. + kernelUses :: [KernelUse], + kernelNumGroups :: [Imp.Exp], + kernelGroupSize :: [Imp.Exp], + -- | A short descriptive and _unique_ name - should be + -- alphanumeric and without spaces. + kernelName :: Name, + -- | If true, this kernel does not need to check + -- whether we are in a failing state, as it can cope. + -- Intuitively, it means that the kernel does not + -- depend on any non-scalar parameters to make control + -- flow decisions. Replication, transpose, and copy + -- kernels are examples of this. + kernelFailureTolerant :: Bool + } + deriving (Show) + +-- | Information about a host-level variable that is used inside this +-- kernel. When generating the actual kernel code, this is used to +-- deduce which parameters are needed. +data KernelUse + = ScalarUse VName PrimType + | MemoryUse VName + | ConstUse VName KernelConstExp + deriving (Eq, Ord, Show) + +instance Pretty KernelConst where + ppr (SizeConst key) = text "get_size" <> parens (ppr key) + +instance Pretty KernelUse where + ppr (ScalarUse name t) = + oneLine $ text "scalar_copy" <> parens (commasep [ppr name, ppr t]) + ppr (MemoryUse name) = + oneLine $ text "mem_copy" <> parens (commasep [ppr name]) + ppr (ConstUse name e) = + oneLine $ text "const" <> parens (commasep [ppr name, ppr e]) + +instance Pretty HostOp where + ppr (GetSize dest key size_class) = + ppr dest <+> text "<-" + <+> text "get_size" <> parens (commasep [ppr key, ppr size_class]) + ppr (GetSizeMax dest size_class) = + ppr dest <+> text "<-" <+> text "get_size_max" <> parens (ppr size_class) + ppr (CmpSizeLe dest name size_class x) = + ppr dest <+> text "<-" + <+> text "get_size" <> parens (commasep [ppr name, ppr size_class]) + <+> text "<" + <+> ppr x + ppr (CallKernel c) = + ppr c + +instance FreeIn HostOp where + freeIn' (CallKernel c) = + freeIn' c + freeIn' (CmpSizeLe dest _ _ x) = + freeIn' dest <> freeIn' x + freeIn' (GetSizeMax dest _) = + freeIn' dest + freeIn' (GetSize dest _ _) = + freeIn' dest + +instance FreeIn Kernel where + freeIn' kernel = + freeIn' (kernelBody kernel) + <> freeIn' [kernelNumGroups kernel, kernelGroupSize kernel] + +instance Pretty Kernel where + ppr kernel = + text "kernel" + <+> brace + ( text "groups" <+> brace (ppr $ kernelNumGroups kernel) + text "group_size" <+> brace (ppr $ kernelGroupSize kernel) + text "uses" <+> brace (commasep $ map ppr $ kernelUses kernel) + text "failure_tolerant" <+> brace (ppr $ kernelFailureTolerant kernel) + text "body" <+> brace (ppr $ kernelBody kernel) + ) + +-- | When we do a barrier or fence, is it at the local or global +-- level? +data Fence = FenceLocal | FenceGlobal + deriving (Show) + +-- | An operation that occurs within a kernel body. +data KernelOp + = GetGroupId VName Int + | GetLocalId VName Int + | GetLocalSize VName Int + | GetGlobalSize VName Int + | GetGlobalId VName Int + | GetLockstepWidth VName + | Atomic Space AtomicOp + | Barrier Fence + | MemFence Fence + | LocalAlloc VName (Count Bytes (Imp.TExp Int64)) + | -- | Perform a barrier and also check whether any + -- threads have failed an assertion. Make sure all + -- threads would reach all 'ErrorSync's if any of them + -- do. A failing assertion will jump to the next + -- following 'ErrorSync', so make sure it's not inside + -- control flow or similar. + ErrorSync Fence + deriving (Show) + +-- | Atomic operations return the value stored before the update. +-- This old value is stored in the first 'VName'. The second 'VName' +-- is the memory block to update. The 'Exp' is the new value. +data AtomicOp + = AtomicAdd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicFAdd FloatType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicSMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicSMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicUMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicUMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int64)) Exp + | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp Exp + | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp + deriving (Show) + +instance FreeIn AtomicOp where + freeIn' (AtomicAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicFAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicSMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicSMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicUMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicUMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicAnd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicOr _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicXor _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + freeIn' (AtomicCmpXchg _ _ arr i x y) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' y + freeIn' (AtomicXchg _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x + +instance Pretty KernelOp where + ppr (GetGroupId dest i) = + ppr dest <+> "<-" + <+> "get_group_id" <> parens (ppr i) + ppr (GetLocalId dest i) = + ppr dest <+> "<-" + <+> "get_local_id" <> parens (ppr i) + ppr (GetLocalSize dest i) = + ppr dest <+> "<-" + <+> "get_local_size" <> parens (ppr i) + ppr (GetGlobalSize dest i) = + ppr dest <+> "<-" + <+> "get_global_size" <> parens (ppr i) + ppr (GetGlobalId dest i) = + ppr dest <+> "<-" + <+> "get_global_id" <> parens (ppr i) + ppr (GetLockstepWidth dest) = + ppr dest <+> "<-" + <+> "get_lockstep_width()" + ppr (Barrier FenceLocal) = + "local_barrier()" + ppr (Barrier FenceGlobal) = + "global_barrier()" + ppr (MemFence FenceLocal) = + "mem_fence_local()" + ppr (MemFence FenceGlobal) = + "mem_fence_global()" + ppr (LocalAlloc name size) = + ppr name <+> equals <+> "local_alloc" <> parens (ppr size) + ppr (ErrorSync FenceLocal) = + "error_sync_local()" + ppr (ErrorSync FenceGlobal) = + "error_sync_global()" + ppr (Atomic _ (AtomicAdd t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_add_" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicFAdd t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_fadd_" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicSMax t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_smax" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicSMin t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_smin" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicUMax t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_umax" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicUMin t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_umin" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicAnd t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_and" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicOr t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_or" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicXor t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_xor" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic _ (AtomicCmpXchg t old arr ind x y)) = + ppr old <+> "<-" <+> "atomic_cmp_xchg" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y]) + ppr (Atomic _ (AtomicXchg t old arr ind x)) = + ppr old <+> "<-" <+> "atomic_xchg" <> ppr t + <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + +instance FreeIn KernelOp where + freeIn' (Atomic _ op) = freeIn' op + freeIn' _ = mempty + +brace :: Doc -> Doc +brace body = " {" indent 2 body "}" diff --git a/src/Futhark/CodeGen/ImpCode/Kernels.hs b/src/Futhark/CodeGen/ImpCode/Kernels.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpCode/Kernels.hs +++ /dev/null @@ -1,266 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} - --- | Variation of "Futhark.CodeGen.ImpCode" that contains the notion --- of a kernel invocation. -module Futhark.CodeGen.ImpCode.Kernels - ( Program, - Function, - FunctionT (Function), - Code, - KernelCode, - KernelConst (..), - KernelConstExp, - HostOp (..), - KernelOp (..), - Fence (..), - AtomicOp (..), - Kernel (..), - KernelUse (..), - module Futhark.CodeGen.ImpCode, - module Futhark.IR.Kernels.Sizes, - ) -where - -import Futhark.CodeGen.ImpCode hiding (Code, Function) -import qualified Futhark.CodeGen.ImpCode as Imp -import Futhark.IR.Kernels.Sizes -import Futhark.IR.Pretty () -import Futhark.Util.Pretty - --- | A program that calls kernels. -type Program = Imp.Definitions HostOp - --- | A function that calls kernels. -type Function = Imp.Function HostOp - --- | Host-level code that can call kernels. -type Code = Imp.Code HostOp - --- | Code inside a kernel. -type KernelCode = Imp.Code KernelOp - --- | A run-time constant related to kernels. -newtype KernelConst = SizeConst Name - deriving (Eq, Ord, Show) - --- | An expression whose variables are kernel constants. -type KernelConstExp = PrimExp KernelConst - --- | An operation that runs on the host (CPU). -data HostOp - = CallKernel Kernel - | GetSize VName Name SizeClass - | CmpSizeLe VName Name SizeClass Imp.Exp - | GetSizeMax VName SizeClass - deriving (Show) - --- | A generic kernel containing arbitrary kernel code. -data Kernel = Kernel - { kernelBody :: Imp.Code KernelOp, - -- | The host variables referenced by the kernel. - kernelUses :: [KernelUse], - kernelNumGroups :: [Imp.Exp], - kernelGroupSize :: [Imp.Exp], - -- | A short descriptive and _unique_ name - should be - -- alphanumeric and without spaces. - kernelName :: Name, - -- | If true, this kernel does not need to check - -- whether we are in a failing state, as it can cope. - -- Intuitively, it means that the kernel does not - -- depend on any non-scalar parameters to make control - -- flow decisions. Replication, transpose, and copy - -- kernels are examples of this. - kernelFailureTolerant :: Bool - } - deriving (Show) - --- | Information about a host-level variable that is used inside this --- kernel. When generating the actual kernel code, this is used to --- deduce which parameters are needed. -data KernelUse - = ScalarUse VName PrimType - | MemoryUse VName - | ConstUse VName KernelConstExp - deriving (Eq, Ord, Show) - -instance Pretty KernelConst where - ppr (SizeConst key) = text "get_size" <> parens (ppr key) - -instance Pretty KernelUse where - ppr (ScalarUse name t) = - oneLine $ text "scalar_copy" <> parens (commasep [ppr name, ppr t]) - ppr (MemoryUse name) = - oneLine $ text "mem_copy" <> parens (commasep [ppr name]) - ppr (ConstUse name e) = - oneLine $ text "const" <> parens (commasep [ppr name, ppr e]) - -instance Pretty HostOp where - ppr (GetSize dest key size_class) = - ppr dest <+> text "<-" - <+> text "get_size" <> parens (commasep [ppr key, ppr size_class]) - ppr (GetSizeMax dest size_class) = - ppr dest <+> text "<-" <+> text "get_size_max" <> parens (ppr size_class) - ppr (CmpSizeLe dest name size_class x) = - ppr dest <+> text "<-" - <+> text "get_size" <> parens (commasep [ppr name, ppr size_class]) - <+> text "<" - <+> ppr x - ppr (CallKernel c) = - ppr c - -instance FreeIn HostOp where - freeIn' (CallKernel c) = - freeIn' c - freeIn' (CmpSizeLe dest _ _ x) = - freeIn' dest <> freeIn' x - freeIn' (GetSizeMax dest _) = - freeIn' dest - freeIn' (GetSize dest _ _) = - freeIn' dest - -instance FreeIn Kernel where - freeIn' kernel = - freeIn' (kernelBody kernel) - <> freeIn' [kernelNumGroups kernel, kernelGroupSize kernel] - -instance Pretty Kernel where - ppr kernel = - text "kernel" - <+> brace - ( text "groups" <+> brace (ppr $ kernelNumGroups kernel) - text "group_size" <+> brace (ppr $ kernelGroupSize kernel) - text "uses" <+> brace (commasep $ map ppr $ kernelUses kernel) - text "failure_tolerant" <+> brace (ppr $ kernelFailureTolerant kernel) - text "body" <+> brace (ppr $ kernelBody kernel) - ) - --- | When we do a barrier or fence, is it at the local or global --- level? -data Fence = FenceLocal | FenceGlobal - deriving (Show) - --- | An operation that occurs within a kernel body. -data KernelOp - = GetGroupId VName Int - | GetLocalId VName Int - | GetLocalSize VName Int - | GetGlobalSize VName Int - | GetGlobalId VName Int - | GetLockstepWidth VName - | Atomic Space AtomicOp - | Barrier Fence - | MemFence Fence - | LocalAlloc VName (Count Bytes (Imp.TExp Int64)) - | -- | Perform a barrier and also check whether any - -- threads have failed an assertion. Make sure all - -- threads would reach all 'ErrorSync's if any of them - -- do. A failing assertion will jump to the next - -- following 'ErrorSync', so make sure it's not inside - -- control flow or similar. - ErrorSync Fence - deriving (Show) - --- | Atomic operations return the value stored before the update. --- This old value is stored in the first 'VName'. The second 'VName' --- is the memory block to update. The 'Exp' is the new value. -data AtomicOp - = AtomicAdd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicFAdd FloatType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicSMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicSMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicUMax IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicUMin IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicAnd IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicOr IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicXor IntType VName VName (Count Elements (Imp.TExp Int64)) Exp - | AtomicCmpXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp Exp - | AtomicXchg PrimType VName VName (Count Elements (Imp.TExp Int64)) Exp - deriving (Show) - -instance FreeIn AtomicOp where - freeIn' (AtomicAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicFAdd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicSMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicSMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicUMax _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicUMin _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicAnd _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicOr _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicXor _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - freeIn' (AtomicCmpXchg _ _ arr i x y) = freeIn' arr <> freeIn' i <> freeIn' x <> freeIn' y - freeIn' (AtomicXchg _ _ arr i x) = freeIn' arr <> freeIn' i <> freeIn' x - -instance Pretty KernelOp where - ppr (GetGroupId dest i) = - ppr dest <+> "<-" - <+> "get_group_id" <> parens (ppr i) - ppr (GetLocalId dest i) = - ppr dest <+> "<-" - <+> "get_local_id" <> parens (ppr i) - ppr (GetLocalSize dest i) = - ppr dest <+> "<-" - <+> "get_local_size" <> parens (ppr i) - ppr (GetGlobalSize dest i) = - ppr dest <+> "<-" - <+> "get_global_size" <> parens (ppr i) - ppr (GetGlobalId dest i) = - ppr dest <+> "<-" - <+> "get_global_id" <> parens (ppr i) - ppr (GetLockstepWidth dest) = - ppr dest <+> "<-" - <+> "get_lockstep_width()" - ppr (Barrier FenceLocal) = - "local_barrier()" - ppr (Barrier FenceGlobal) = - "global_barrier()" - ppr (MemFence FenceLocal) = - "mem_fence_local()" - ppr (MemFence FenceGlobal) = - "mem_fence_global()" - ppr (LocalAlloc name size) = - ppr name <+> equals <+> "local_alloc" <> parens (ppr size) - ppr (ErrorSync FenceLocal) = - "error_sync_local()" - ppr (ErrorSync FenceGlobal) = - "error_sync_global()" - ppr (Atomic _ (AtomicAdd t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_add_" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicFAdd t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_fadd_" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicSMax t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_smax" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicSMin t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_smin" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicUMax t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_umax" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicUMin t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_umin" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicAnd t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_and" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicOr t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_or" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicXor t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_xor" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - ppr (Atomic _ (AtomicCmpXchg t old arr ind x y)) = - ppr old <+> "<-" <+> "atomic_cmp_xchg" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y]) - ppr (Atomic _ (AtomicXchg t old arr ind x)) = - ppr old <+> "<-" <+> "atomic_xchg" <> ppr t - <> parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) - -instance FreeIn KernelOp where - freeIn' (Atomic _ op) = freeIn' op - freeIn' _ = mempty - -brace :: Doc -> Doc -brace body = " {" indent 2 body "}" diff --git a/src/Futhark/CodeGen/ImpCode/OpenCL.hs b/src/Futhark/CodeGen/ImpCode/OpenCL.hs --- a/src/Futhark/CodeGen/ImpCode/OpenCL.hs +++ b/src/Futhark/CodeGen/ImpCode/OpenCL.hs @@ -19,14 +19,14 @@ KernelTarget (..), FailureMsg (..), module Futhark.CodeGen.ImpCode, - module Futhark.IR.Kernels.Sizes, + module Futhark.IR.GPU.Sizes, ) where import qualified Data.Map as M import Futhark.CodeGen.ImpCode hiding (Code, Function) import qualified Futhark.CodeGen.ImpCode as Imp -import Futhark.IR.Kernels.Sizes +import Futhark.IR.GPU.Sizes import Futhark.Util.Pretty -- | An program calling OpenCL kernels. diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs --- a/src/Futhark/CodeGen/ImpGen.hs +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -151,39 +151,39 @@ import Prelude hiding (quot) -- | How to compile an t'Op'. -type OpCompiler lore r op = Pattern lore -> Op lore -> ImpM lore r op () +type OpCompiler rep r op = Pattern rep -> Op rep -> ImpM rep r op () -- | How to compile some 'Stms'. -type StmsCompiler lore r op = Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op () +type StmsCompiler rep r op = Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op () -- | How to compile an 'Exp'. -type ExpCompiler lore r op = Pattern lore -> Exp lore -> ImpM lore r op () +type ExpCompiler rep r op = Pattern rep -> Exp rep -> ImpM rep r op () -type CopyCompiler lore r op = +type CopyCompiler rep r op = PrimType -> MemLocation -> Slice (Imp.TExp Int64) -> MemLocation -> Slice (Imp.TExp Int64) -> - ImpM lore r op () + ImpM rep r op () -- | An alternate way of compiling an allocation. -type AllocCompiler lore r op = VName -> Count Bytes (Imp.TExp Int64) -> ImpM lore r op () +type AllocCompiler rep r op = VName -> Count Bytes (Imp.TExp Int64) -> ImpM rep r op () -data Operations lore r op = Operations - { opsExpCompiler :: ExpCompiler lore r op, - opsOpCompiler :: OpCompiler lore r op, - opsStmsCompiler :: StmsCompiler lore r op, - opsCopyCompiler :: CopyCompiler lore r op, - opsAllocCompilers :: M.Map Space (AllocCompiler lore r op) +data Operations rep r op = Operations + { opsExpCompiler :: ExpCompiler rep r op, + opsOpCompiler :: OpCompiler rep r op, + opsStmsCompiler :: StmsCompiler rep r op, + opsCopyCompiler :: CopyCompiler rep r op, + opsAllocCompilers :: M.Map Space (AllocCompiler rep r op) } -- | An operations set for which the expression compiler always -- returns 'defCompileExp'. defaultOperations :: - (Mem lore, FreeIn op) => - OpCompiler lore r op -> - Operations lore r op + (Mem rep, FreeIn op) => + OpCompiler rep r op -> + Operations rep r op defaultOperations opc = Operations { opsExpCompiler = defCompileExp, @@ -219,11 +219,11 @@ deriving (Show) -- | Every non-scalar variable must be associated with an entry. -data VarEntry lore - = ArrayVar (Maybe (Exp lore)) ArrayEntry - | ScalarVar (Maybe (Exp lore)) ScalarEntry - | MemVar (Maybe (Exp lore)) MemEntry - | AccVar (Maybe (Exp lore)) (VName, Shape, [Type]) +data VarEntry rep + = ArrayVar (Maybe (Exp rep)) ArrayEntry + | ScalarVar (Maybe (Exp rep)) ScalarEntry + | MemVar (Maybe (Exp rep)) MemEntry + | AccVar (Maybe (Exp rep)) (VName, Shape, [Type]) deriving (Show) -- | When compiling an expression, this is a description of where the @@ -247,12 +247,12 @@ ArrayDestination (Maybe MemLocation) deriving (Show) -data Env lore r op = Env - { envExpCompiler :: ExpCompiler lore r op, - envStmsCompiler :: StmsCompiler lore r op, - envOpCompiler :: OpCompiler lore r op, - envCopyCompiler :: CopyCompiler lore r op, - envAllocCompilers :: M.Map Space (AllocCompiler lore r op), +data Env rep r op = Env + { envExpCompiler :: ExpCompiler rep r op, + envStmsCompiler :: StmsCompiler rep r op, + envOpCompiler :: OpCompiler rep r op, + envCopyCompiler :: CopyCompiler rep r op, + envAllocCompilers :: M.Map Space (AllocCompiler rep r op), envDefaultSpace :: Imp.Space, envVolatility :: Imp.Volatility, -- | User-extensible environment. @@ -264,7 +264,7 @@ envAttrs :: Attrs } -newEnv :: r -> Operations lore r op -> Imp.Space -> Env lore r op +newEnv :: r -> Operations rep r op -> Imp.Space -> Env rep r op newEnv r ops ds = Env { envExpCompiler = opsExpCompiler ops, @@ -280,10 +280,10 @@ } -- | The symbol table used during compilation. -type VTable lore = M.Map VName (VarEntry lore) +type VTable rep = M.Map VName (VarEntry rep) -data ImpState lore r op = ImpState - { stateVTable :: VTable lore, +data ImpState rep r op = ImpState + { stateVTable :: VTable rep, stateFunctions :: Imp.Functions op, stateCode :: Imp.Code op, stateWarnings :: Warnings, @@ -293,30 +293,30 @@ -- accumulator throughout its lifetime. If the arrays -- backing an accumulator is not in this mapping, the -- accumulator is scatter-like. - stateAccs :: M.Map VName ([VName], Maybe (Lambda lore, [SubExp])), + stateAccs :: M.Map VName ([VName], Maybe (Lambda rep, [SubExp])), stateNameSource :: VNameSource } -newState :: VNameSource -> ImpState lore r op +newState :: VNameSource -> ImpState rep r op newState = ImpState mempty mempty mempty mempty mempty -newtype ImpM lore r op a - = ImpM (ReaderT (Env lore r op) (State (ImpState lore r op)) a) +newtype ImpM rep r op a + = ImpM (ReaderT (Env rep r op) (State (ImpState rep r op)) a) deriving ( Functor, Applicative, Monad, - MonadState (ImpState lore r op), - MonadReader (Env lore r op) + MonadState (ImpState rep r op), + MonadReader (Env rep r op) ) -instance MonadFreshNames (ImpM lore r op) where +instance MonadFreshNames (ImpM rep r op) where getNameSource = gets stateNameSource putNameSource src = modify $ \s -> s {stateNameSource = src} -- Cannot be an KernelsMem scope because the index functions have -- the wrong leaves (VName instead of Imp.Exp). -instance HasScope SOACS (ImpM lore r op) where +instance HasScope SOACS (ImpM rep r op) where askScope = gets $ M.map (LetName . entryType) . stateVTable where entryType (MemVar _ memEntry) = @@ -332,26 +332,26 @@ Acc acc ispace ts NoUniqueness runImpM :: - ImpM lore r op a -> + ImpM rep r op a -> r -> - Operations lore r op -> + Operations rep r op -> Imp.Space -> - ImpState lore r op -> - (a, ImpState lore r op) + ImpState rep r op -> + (a, ImpState rep r op) runImpM (ImpM m) r ops space = runState (runReaderT m $ newEnv r ops space) subImpM_ :: r' -> - Operations lore r' op' -> - ImpM lore r' op' a -> - ImpM lore r op (Imp.Code op') + Operations rep r' op' -> + ImpM rep r' op' a -> + ImpM rep r op (Imp.Code op') subImpM_ r ops m = snd <$> subImpM r ops m subImpM :: r' -> - Operations lore r' op' -> - ImpM lore r' op' a -> - ImpM lore r op (a, Imp.Code op') + Operations rep r' op' -> + ImpM rep r' op' a -> + ImpM rep r op (a, Imp.Code op') subImpM r ops (ImpM m) = do env <- ask s <- get @@ -382,10 +382,10 @@ -- | Execute a code generation action, returning the code that was -- emitted. -collect :: ImpM lore r op () -> ImpM lore r op (Imp.Code op) +collect :: ImpM rep r op () -> ImpM rep r op (Imp.Code op) collect = fmap snd . collect' -collect' :: ImpM lore r op a -> ImpM lore r op (a, Imp.Code op) +collect' :: ImpM rep r op a -> ImpM rep r op (a, Imp.Code op) collect' m = do prev_code <- gets stateCode modify $ \s -> s {stateCode = mempty} @@ -396,36 +396,36 @@ -- | Execute a code generation action, wrapping the generated code -- within a 'Imp.Comment' with the given description. -comment :: String -> ImpM lore r op () -> ImpM lore r op () +comment :: String -> ImpM rep r op () -> ImpM rep r op () comment desc m = do code <- collect m emit $ Imp.Comment desc code -- | Emit some generated imperative code. -emit :: Imp.Code op -> ImpM lore r op () +emit :: Imp.Code op -> ImpM rep r op () emit code = modify $ \s -> s {stateCode = stateCode s <> code} -warnings :: Warnings -> ImpM lore r op () +warnings :: Warnings -> ImpM rep r op () warnings ws = modify $ \s -> s {stateWarnings = ws <> stateWarnings s} -- | Emit a warning about something the user should be aware of. -warn :: Located loc => loc -> [loc] -> String -> ImpM lore r op () +warn :: Located loc => loc -> [loc] -> String -> ImpM rep r op () warn loc locs problem = warnings $ singleWarning' (srclocOf loc) (map srclocOf locs) (fromString problem) -- | Emit a function in the generated code. -emitFunction :: Name -> Imp.Function op -> ImpM lore r op () +emitFunction :: Name -> Imp.Function op -> ImpM rep r op () emitFunction fname fun = do Imp.Functions fs <- gets stateFunctions modify $ \s -> s {stateFunctions = Imp.Functions $ (fname, fun) : fs} -- | Check if a function of a given name exists. -hasFunction :: Name -> ImpM lore r op Bool +hasFunction :: Name -> ImpM rep r op Bool hasFunction fname = gets $ \s -> let Imp.Functions fs = stateFunctions s in isJust $ lookup fname fs -constsVTable :: Mem lore => Stms lore -> VTable lore +constsVTable :: Mem rep => Stms rep -> VTable rep constsVTable = foldMap stmVtable where stmVtable (Let pat _ e) = @@ -434,11 +434,11 @@ M.singleton name $ memBoundToVarEntry (Just e) dec compileProg :: - (Mem lore, FreeIn op, MonadFreshNames m) => + (Mem rep, FreeIn op, MonadFreshNames m) => r -> - Operations lore r op -> + Operations rep r op -> Imp.Space -> - Prog lore -> + Prog rep -> m (Warnings, Imp.Definitions op) compileProg r ops space (Prog consts funs) = modifyNameSource $ \src -> @@ -473,7 +473,7 @@ mconcat $ map stateWarnings ss } -compileConsts :: Names -> Stms lore -> ImpM lore r op (Imp.Constants op) +compileConsts :: Names -> Stms rep -> ImpM rep r op (Imp.Constants op) compileConsts used_consts stms = do code <- collect $ compileStms used_consts stms $ pure () pure $ uncurry Imp.Constants $ first DL.toList $ extract code @@ -496,9 +496,9 @@ (mempty, s) compileInParam :: - Mem lore => - FParam lore -> - ImpM lore r op (Either Imp.Param ArrayDecl) + Mem rep => + FParam rep -> + ImpM rep r op (Either Imp.Param ArrayDecl) compileInParam fparam = case paramDec fparam of MemPrim bt -> return $ Left $ Imp.ScalarParam name bt @@ -517,10 +517,10 @@ data ArrayDecl = ArrayDecl VName PrimType MemLocation compileInParams :: - Mem lore => - [FParam lore] -> + Mem rep => + [FParam rep] -> [EntryPointType] -> - ImpM lore r op ([Imp.Param], [ArrayDecl], [Imp.ExternalValue]) + ImpM rep r op ([Imp.Param], [ArrayDecl], [Imp.ExternalValue]) compileInParams params orig_epts = do let (ctx_params, val_params) = splitAt (length params - sum (map entryPointSize orig_epts)) params @@ -548,17 +548,18 @@ _ -> Nothing - mkExts (TypeOpaque desc n : epts) fparams = + mkExts (TypeOpaque u desc n : epts) fparams = let (fparams', rest) = splitAt n fparams in Imp.OpaqueValue + u desc (mapMaybe (`mkValueDesc` Imp.TypeDirect) fparams') : mkExts epts rest - mkExts (TypeUnsigned : epts) (fparam : fparams) = - maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeUnsigned) + mkExts (TypeUnsigned u : epts) (fparam : fparams) = + maybeToList (Imp.TransparentValue u <$> mkValueDesc fparam Imp.TypeUnsigned) ++ mkExts epts fparams - mkExts (TypeDirect : epts) (fparam : fparams) = - maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeDirect) + mkExts (TypeDirect u : epts) (fparam : fparams) = + maybeToList (Imp.TransparentValue u <$> mkValueDesc fparam Imp.TypeDirect) ++ mkExts epts fparams mkExts _ _ = [] @@ -567,10 +568,10 @@ isArrayDecl x (ArrayDecl y _ _) = x == y compileOutParams :: - Mem lore => - [RetType lore] -> + Mem rep => + [RetType rep] -> [EntryPointType] -> - ImpM lore r op ([Imp.ExternalValue], [Imp.Param], Destination) + ImpM rep r op ([Imp.ExternalValue], [Imp.Param], Destination) compileOutParams orig_rts orig_epts = do ((extvs, dests), (outparams, ctx_dests)) <- runWriterT $ evalStateT (mkExts orig_epts orig_rts) (M.empty, M.empty) @@ -579,26 +580,26 @@ where imp = lift . lift - mkExts (TypeOpaque desc n : epts) rts = do + mkExts (TypeOpaque u desc n : epts) rts = do let (rts', rest) = splitAt n rts (evs, dests) <- unzip <$> zipWithM mkParam rts' (repeat Imp.TypeDirect) (more_values, more_dests) <- mkExts epts rest return - ( Imp.OpaqueValue desc evs : more_values, + ( Imp.OpaqueValue u desc evs : more_values, dests ++ more_dests ) - mkExts (TypeUnsigned : epts) (rt : rts) = do + mkExts (TypeUnsigned u : epts) (rt : rts) = do (ev, dest) <- mkParam rt Imp.TypeUnsigned (more_values, more_dests) <- mkExts epts rts return - ( Imp.TransparentValue ev : more_values, + ( Imp.TransparentValue u ev : more_values, dest : more_dests ) - mkExts (TypeDirect : epts) (rt : rts) = do + mkExts (TypeDirect u : epts) (rt : rts) = do (ev, dest) <- mkParam rt Imp.TypeDirect (more_values, more_dests) <- mkExts epts rts return - ( Imp.TransparentValue ev : more_values, + ( Imp.TransparentValue u ev : more_values, dest : more_dests ) mkExts _ _ = return ([], []) @@ -646,9 +647,9 @@ return se compileFunDef :: - Mem lore => - FunDef lore -> - ImpM lore r op () + Mem rep => + FunDef rep -> + ImpM rep r op () compileFunDef (FunDef entry _ fname rettype params body) = local (\env -> env {envFunction = name_entry `mplus` Just fname}) $ do ((outparams, inparams, results, args), body') <- collect' compile @@ -657,8 +658,8 @@ (name_entry, params_entry, ret_entry) = case entry of Nothing -> ( Nothing, - replicate (length params) TypeDirect, - replicate (length rettype) TypeDirect + replicate (length params) (TypeDirect mempty), + replicate (length rettype) (TypeDirect mempty) ) Just (x, y, z) -> (Just x, y, z) compile = do @@ -673,18 +674,18 @@ return (outparams, inparams, results, args) -compileBody :: (Mem lore) => Pattern lore -> Body lore -> ImpM lore r op () +compileBody :: (Mem rep) => Pattern rep -> Body rep -> ImpM rep r op () compileBody pat (Body _ bnds ses) = do Destination _ dests <- destinationFromPattern pat compileStms (freeIn ses) bnds $ forM_ (zip dests ses) $ \(d, se) -> copyDWIMDest d [] se [] -compileBody' :: [Param dec] -> Body lore -> ImpM lore r op () +compileBody' :: [Param dec] -> Body rep -> ImpM rep r op () compileBody' params (Body _ bnds ses) = compileStms (freeIn ses) bnds $ forM_ (zip params ses) $ \(param, se) -> copyDWIM (paramName param) [] se [] -compileLoopBody :: Typed dec => [Param dec] -> Body lore -> ImpM lore r op () +compileLoopBody :: Typed dec => [Param dec] -> Body rep -> ImpM rep r op () compileLoopBody mergeparams (Body _ bnds ses) = do -- We cannot write the results to the merge parameters immediately, -- as some of the results may actually *be* merge parameters, and @@ -707,17 +708,17 @@ _ -> return $ return () sequence_ copy_to_merge_params -compileStms :: Names -> Stms lore -> ImpM lore r op () -> ImpM lore r op () +compileStms :: Names -> Stms rep -> ImpM rep r op () -> ImpM rep r op () compileStms alive_after_stms all_stms m = do cb <- asks envStmsCompiler cb alive_after_stms all_stms m defCompileStms :: - (Mem lore, FreeIn op) => + (Mem rep, FreeIn op) => Names -> - Stms lore -> - ImpM lore r op () -> - ImpM lore r op () + Stms rep -> + ImpM rep r op () -> + ImpM rep r op () defCompileStms alive_after_stms all_stms m = -- We keep track of any memory blocks produced by the statements, -- and after the last time that memory block is used, we insert a @@ -752,16 +753,16 @@ Mem space -> Just (patElemName pe, space) _ -> Nothing -compileExp :: Pattern lore -> Exp lore -> ImpM lore r op () +compileExp :: Pattern rep -> Exp rep -> ImpM rep r op () compileExp pat e = do ec <- asks envExpCompiler ec pat e defCompileExp :: - (Mem lore) => - Pattern lore -> - Exp lore -> - ImpM lore r op () + (Mem rep) => + Pattern rep -> + Exp rep -> + ImpM rep r op () defCompileExp pat (If cond tbranch fbranch _) = sIf (toBoolExp cond) (compileBody pat tbranch) (compileBody pat fbranch) defCompileExp pat (Apply fname args _ _) = do @@ -827,10 +828,10 @@ opc pat op defCompileBasicOp :: - Mem lore => - Pattern lore -> + Mem rep => + Pattern rep -> BasicOp -> - ImpM lore r op () + ImpM rep r op () defCompileBasicOp (Pattern _ [pe]) (SubExp se) = copyDWIM (patElemName pe) [] se [] defCompileBasicOp (Pattern _ [pe]) (Opaque se) = @@ -967,7 +968,7 @@ ++ pretty e -- | Note: a hack to be used only for functions. -addArrays :: [ArrayDecl] -> ImpM lore r op () +addArrays :: [ArrayDecl] -> ImpM rep r op () addArrays = mapM_ addArray where addArray (ArrayDecl name bt location) = @@ -981,7 +982,7 @@ -- | Like 'dFParams', but does not create new declarations. -- Note: a hack to be used only for functions. -addFParams :: Mem lore => [FParam lore] -> ImpM lore r op () +addFParams :: Mem rep => [FParam rep] -> ImpM rep r op () addFParams = mapM_ addFParam where addFParam fparam = @@ -989,25 +990,25 @@ memBoundToVarEntry Nothing $ noUniquenessReturns $ paramDec fparam -- | Another hack. -addLoopVar :: VName -> IntType -> ImpM lore r op () +addLoopVar :: VName -> IntType -> ImpM rep r op () addLoopVar i it = addVar i $ ScalarVar Nothing $ ScalarEntry $ IntType it dVars :: - Mem lore => - Maybe (Exp lore) -> - [PatElem lore] -> - ImpM lore r op () + Mem rep => + Maybe (Exp rep) -> + [PatElem rep] -> + ImpM rep r op () dVars e = mapM_ dVar where dVar = dScope e . scopeOfPatElem -dFParams :: Mem lore => [FParam lore] -> ImpM lore r op () +dFParams :: Mem rep => [FParam rep] -> ImpM rep r op () dFParams = dScope Nothing . scopeOfFParams -dLParams :: Mem lore => [LParam lore] -> ImpM lore r op () +dLParams :: Mem rep => [LParam rep] -> ImpM rep r op () dLParams = dScope Nothing . scopeOfLParams -dPrimVol :: String -> PrimType -> Imp.TExp t -> ImpM lore r op (TV t) +dPrimVol :: String -> PrimType -> Imp.TExp t -> ImpM rep r op (TV t) dPrimVol name t e = do name' <- newVName name emit $ Imp.DeclareScalar name' Imp.Volatile t @@ -1015,7 +1016,7 @@ name' <~~ untyped e return $ TV name' t -dPrim_ :: VName -> PrimType -> ImpM lore r op () +dPrim_ :: VName -> PrimType -> ImpM rep r op () dPrim_ name t = do emit $ Imp.DeclareScalar name Imp.Nonvolatile t addVar name $ ScalarVar Nothing $ ScalarEntry t @@ -1023,35 +1024,35 @@ -- | The return type is polymorphic, so there is no guarantee it -- actually matches the 'PrimType', but at least we have to use it -- consistently. -dPrim :: String -> PrimType -> ImpM lore r op (TV t) +dPrim :: String -> PrimType -> ImpM rep r op (TV t) dPrim name t = do name' <- newVName name dPrim_ name' t return $ TV name' t -dPrimV_ :: VName -> Imp.TExp t -> ImpM lore r op () +dPrimV_ :: VName -> Imp.TExp t -> ImpM rep r op () dPrimV_ name e = do dPrim_ name t TV name t <-- e where t = primExpType $ untyped e -dPrimV :: String -> Imp.TExp t -> ImpM lore r op (TV t) +dPrimV :: String -> Imp.TExp t -> ImpM rep r op (TV t) dPrimV name e = do name' <- dPrim name $ primExpType $ untyped e name' <-- e return name' -dPrimVE :: String -> Imp.TExp t -> ImpM lore r op (Imp.TExp t) +dPrimVE :: String -> Imp.TExp t -> ImpM rep r op (Imp.TExp t) dPrimVE name e = do name' <- dPrim name $ primExpType $ untyped e name' <-- e return $ tvExp name' memBoundToVarEntry :: - Maybe (Exp lore) -> + Maybe (Exp rep) -> MemBound NoUniqueness -> - VarEntry lore + VarEntry rep memBoundToVarEntry e (MemPrim bt) = ScalarVar e ScalarEntry {entryScalarType = bt} memBoundToVarEntry e (MemMem space) = @@ -1068,8 +1069,8 @@ } infoDec :: - Mem lore => - NameInfo lore -> + Mem rep => + NameInfo rep -> MemInfo SubExp NoUniqueness MemBind infoDec (LetName dec) = dec infoDec (FParamName dec) = noUniquenessReturns dec @@ -1077,11 +1078,11 @@ infoDec (IndexName it) = MemPrim $ IntType it dInfo :: - Mem lore => - Maybe (Exp lore) -> + Mem rep => + Maybe (Exp rep) -> VName -> - NameInfo lore -> - ImpM lore r op () + NameInfo rep -> + ImpM rep r op () dInfo e name info = do let entry = memBoundToVarEntry e $ infoDec info case entry of @@ -1096,22 +1097,22 @@ addVar name entry dScope :: - Mem lore => - Maybe (Exp lore) -> - Scope lore -> - ImpM lore r op () + Mem rep => + Maybe (Exp rep) -> + Scope rep -> + ImpM rep r op () dScope e = mapM_ (uncurry $ dInfo e) . M.toList -dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op () +dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM rep r op () dArray name bt shape membind = addVar name $ memBoundToVarEntry Nothing $ MemArray bt shape NoUniqueness membind -everythingVolatile :: ImpM lore r op a -> ImpM lore r op a +everythingVolatile :: ImpM rep r op a -> ImpM rep r op a everythingVolatile = local $ \env -> env {envVolatility = Imp.Volatile} -- | Remove the array targets. -funcallTargets :: Destination -> ImpM lore r op [VName] +funcallTargets :: Destination -> ImpM rep r op [VName] funcallTargets (Destination _ dests) = concat <$> mapM funcallTarget dests where @@ -1151,7 +1152,7 @@ class ToExp a where -- | Compile to an 'Imp.Exp', where the type (must must still be a -- primitive) is deduced monadically. - toExp :: a -> ImpM lore r op Imp.Exp + toExp :: a -> ImpM rep r op Imp.Exp -- | Compile where we know the type in advance. toExp' :: PrimType -> a -> Imp.Exp @@ -1178,44 +1179,44 @@ toExp = pure . fmap Imp.ScalarVar toExp' _ = fmap Imp.ScalarVar -addVar :: VName -> VarEntry lore -> ImpM lore r op () +addVar :: VName -> VarEntry rep -> ImpM rep r op () addVar name entry = modify $ \s -> s {stateVTable = M.insert name entry $ stateVTable s} -localDefaultSpace :: Imp.Space -> ImpM lore r op a -> ImpM lore r op a +localDefaultSpace :: Imp.Space -> ImpM rep r op a -> ImpM rep r op a localDefaultSpace space = local (\env -> env {envDefaultSpace = space}) -askFunction :: ImpM lore r op (Maybe Name) +askFunction :: ImpM rep r op (Maybe Name) askFunction = asks envFunction -- | Generate a 'VName', prefixed with 'askFunction' if it exists. -newVNameForFun :: String -> ImpM lore r op VName +newVNameForFun :: String -> ImpM rep r op VName newVNameForFun s = do fname <- fmap nameToString <$> askFunction newVName $ maybe "" (++ ".") fname ++ s -- | Generate a 'Name', prefixed with 'askFunction' if it exists. -nameForFun :: String -> ImpM lore r op Name +nameForFun :: String -> ImpM rep r op Name nameForFun s = do fname <- askFunction return $ maybe "" (<> ".") fname <> nameFromString s -askEnv :: ImpM lore r op r +askEnv :: ImpM rep r op r askEnv = asks envEnv -localEnv :: (r -> r) -> ImpM lore r op a -> ImpM lore r op a +localEnv :: (r -> r) -> ImpM rep r op a -> ImpM rep r op a localEnv f = local $ \env -> env {envEnv = f $ envEnv env} -- | The active attributes, including those for the statement -- currently being compiled. -askAttrs :: ImpM lore r op Attrs +askAttrs :: ImpM rep r op Attrs askAttrs = asks envAttrs -- | Add more attributes to what is returning by 'askAttrs'. -localAttrs :: Attrs -> ImpM lore r op a -> ImpM lore r op a +localAttrs :: Attrs -> ImpM rep r op a -> ImpM rep r op a localAttrs attrs = local $ \env -> env {envAttrs = attrs <> envAttrs env} -localOps :: Operations lore r op -> ImpM lore r op a -> ImpM lore r op a +localOps :: Operations rep r op -> ImpM rep r op a -> ImpM rep r op a localOps ops = local $ \env -> env { envExpCompiler = opsExpCompiler ops, @@ -1226,15 +1227,15 @@ } -- | Get the current symbol table. -getVTable :: ImpM lore r op (VTable lore) +getVTable :: ImpM rep r op (VTable rep) getVTable = gets stateVTable -putVTable :: VTable lore -> ImpM lore r op () +putVTable :: VTable rep -> ImpM rep r op () putVTable vtable = modify $ \s -> s {stateVTable = vtable} -- | Run an action with a modified symbol table. All changes to the -- symbol table will be reverted once the action is done! -localVTable :: (VTable lore -> VTable lore) -> ImpM lore r op a -> ImpM lore r op a +localVTable :: (VTable rep -> VTable rep) -> ImpM rep r op a -> ImpM rep r op a localVTable f m = do old_vtable <- getVTable putVTable $ f old_vtable @@ -1242,28 +1243,28 @@ putVTable old_vtable return a -lookupVar :: VName -> ImpM lore r op (VarEntry lore) +lookupVar :: VName -> ImpM rep r op (VarEntry rep) lookupVar name = do res <- gets $ M.lookup name . stateVTable case res of Just entry -> return entry _ -> error $ "Unknown variable: " ++ pretty name -lookupArray :: VName -> ImpM lore r op ArrayEntry +lookupArray :: VName -> ImpM rep r op ArrayEntry lookupArray name = do res <- lookupVar name case res of ArrayVar _ entry -> return entry _ -> error $ "ImpGen.lookupArray: not an array: " ++ pretty name -lookupMemory :: VName -> ImpM lore r op MemEntry +lookupMemory :: VName -> ImpM rep r op MemEntry lookupMemory name = do res <- lookupVar name case res of MemVar _ entry -> return entry _ -> error $ "Unknown memory block: " ++ pretty name -lookupArraySpace :: VName -> ImpM lore r op Space +lookupArraySpace :: VName -> ImpM rep r op Space lookupArraySpace = fmap entryMemSpace . lookupMemory <=< fmap (memLocationName . entryArrayLocation) . lookupArray @@ -1273,7 +1274,7 @@ lookupAcc :: VName -> [Imp.TExp Int64] -> - ImpM lore r op (VName, Space, [VName], [Imp.TExp Int64], Maybe (Lambda lore)) + ImpM rep r op (VName, Space, [VName], [Imp.TExp Int64], Maybe (Lambda rep)) lookupAcc name is = do res <- lookupVar name case res of @@ -1300,7 +1301,7 @@ error $ "ImpGen.lookupAcc: unlisted accumulator: " ++ pretty name _ -> error $ "ImpGen.lookupAcc: not an accumulator: " ++ pretty name -destinationFromPattern :: Mem lore => Pattern lore -> ImpM lore r op Destination +destinationFromPattern :: Mem rep => Pattern rep -> ImpM rep r op Destination destinationFromPattern pat = fmap (Destination (baseTag <$> maybeHead (patternNames pat))) . mapM inspect $ patternElements pat @@ -1321,7 +1322,7 @@ fullyIndexArray :: VName -> [Imp.TExp Int64] -> - ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) + ImpM rep r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) fullyIndexArray name indices = do arr <- lookupArray name fullyIndexArray' (entryArrayLocation arr) indices @@ -1329,23 +1330,18 @@ fullyIndexArray' :: MemLocation -> [Imp.TExp Int64] -> - ImpM lore r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) + ImpM rep r op (VName, Imp.Space, Count Elements (Imp.TExp Int64)) fullyIndexArray' (MemLocation mem _ ixfun) indices = do space <- entryMemSpace <$> lookupMemory mem - let indices' = case space of - ScalarSpace ds _ -> - let (zero_is, is) = splitFromEnd (length ds) indices - in map (const 0) zero_is ++ is - _ -> indices return ( mem, space, - elements $ IxFun.index ixfun indices' + elements $ IxFun.index ixfun indices ) -- More complicated read/write operations that use index functions. -copy :: CopyCompiler lore r op +copy :: CopyCompiler rep r op copy bt dest destslice src srcslice = do cc <- asks envCopyCompiler cc bt dest destslice src srcslice @@ -1407,7 +1403,7 @@ mapTransposeName :: PrimType -> String mapTransposeName bt = "map_transpose_" ++ pretty bt -mapTransposeForType :: PrimType -> ImpM lore r op Name +mapTransposeForType :: PrimType -> ImpM rep r op Name mapTransposeForType bt = do let fname = nameFromString $ "builtin#" <> mapTransposeName bt @@ -1417,7 +1413,7 @@ return fname -- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'. -defaultCopy :: CopyCompiler lore r op +defaultCopy :: CopyCompiler rep r op defaultCopy pt dest destslice src srcslice | Just ( destoffset, @@ -1471,7 +1467,7 @@ isScalarSpace ScalarSpace {} = True isScalarSpace _ = False -copyElementWise :: CopyCompiler lore r op +copyElementWise :: CopyCompiler rep r op copyElementWise bt dest destslice src srcslice = do let bounds = sliceDims srcslice is <- replicateM (length bounds) (newVName "i") @@ -1494,7 +1490,7 @@ [DimIndex (Imp.TExp Int64)] -> MemLocation -> [DimIndex (Imp.TExp Int64)] -> - ImpM lore r op (Imp.Code op) + ImpM rep r op (Imp.Code op) copyArrayDWIM bt destlocation@(MemLocation _ destshape _) @@ -1544,7 +1540,7 @@ [DimIndex (Imp.TExp Int64)] -> SubExp -> [DimIndex (Imp.TExp Int64)] -> - ImpM lore r op () + ImpM rep r op () copyDWIMDest _ _ (Constant v) (_ : _) = error $ unwords ["copyDWIMDest: constant source", pretty v, "cannot be indexed."] @@ -1640,7 +1636,7 @@ [DimIndex (Imp.TExp Int64)] -> SubExp -> [DimIndex (Imp.TExp Int64)] -> - ImpM lore r op () + ImpM rep r op () copyDWIM dest dest_slice src src_slice = do dest_entry <- lookupVar dest let dest_target = @@ -1662,7 +1658,7 @@ [Imp.TExp Int64] -> SubExp -> [Imp.TExp Int64] -> - ImpM lore r op () + ImpM rep r op () copyDWIMFix dest dest_is src src_is = copyDWIM dest (map DimFix dest_is) src (map DimFix src_is) @@ -1670,11 +1666,11 @@ -- writing the result to @dest@, which must be a single -- 'MemoryDestination', compileAlloc :: - Mem lore => - Pattern lore -> + Mem rep => + Pattern rep -> SubExp -> Space -> - ImpM lore r op () + ImpM rep r op () compileAlloc (Pattern [] [mem]) e space = do let e' = Imp.bytes $ toInt64Exp e allocator <- asks $ M.lookup space . envAllocCompilers @@ -1703,7 +1699,7 @@ --- Building blocks for constructing code. -sFor' :: VName -> Imp.Exp -> ImpM lore r op () -> ImpM lore r op () +sFor' :: VName -> Imp.Exp -> ImpM rep r op () -> ImpM rep r op () sFor' i bound body = do let it = case primExpType bound of IntType bound_t -> bound_t @@ -1712,72 +1708,72 @@ body' <- collect body emit $ Imp.For i bound body' -sFor :: String -> Imp.TExp t -> (Imp.TExp t -> ImpM lore r op ()) -> ImpM lore r op () +sFor :: String -> Imp.TExp t -> (Imp.TExp t -> ImpM rep r op ()) -> ImpM rep r op () sFor i bound body = do i' <- newVName i sFor' i' (untyped bound) $ body $ TPrimExp $ Imp.var i' $ primExpType $ untyped bound -sWhile :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () +sWhile :: Imp.TExp Bool -> ImpM rep r op () -> ImpM rep r op () sWhile cond body = do body' <- collect body emit $ Imp.While cond body' -sComment :: String -> ImpM lore r op () -> ImpM lore r op () +sComment :: String -> ImpM rep r op () -> ImpM rep r op () sComment s code = do code' <- collect code emit $ Imp.Comment s code' -sIf :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () -> ImpM lore r op () +sIf :: Imp.TExp Bool -> ImpM rep r op () -> ImpM rep r op () -> ImpM rep r op () sIf cond tbranch fbranch = do tbranch' <- collect tbranch fbranch' <- collect fbranch emit $ Imp.If cond tbranch' fbranch' -sWhen :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () +sWhen :: Imp.TExp Bool -> ImpM rep r op () -> ImpM rep r op () sWhen cond tbranch = sIf cond tbranch (return ()) -sUnless :: Imp.TExp Bool -> ImpM lore r op () -> ImpM lore r op () +sUnless :: Imp.TExp Bool -> ImpM rep r op () -> ImpM rep r op () sUnless cond = sIf cond (return ()) -sOp :: op -> ImpM lore r op () +sOp :: op -> ImpM rep r op () sOp = emit . Imp.Op -sDeclareMem :: String -> Space -> ImpM lore r op VName +sDeclareMem :: String -> Space -> ImpM rep r op VName sDeclareMem name space = do name' <- newVName name emit $ Imp.DeclareMem name' space addVar name' $ MemVar Nothing $ MemEntry space return name' -sAlloc_ :: VName -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM lore r op () +sAlloc_ :: VName -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM rep r op () sAlloc_ name' size' space = do allocator <- asks $ M.lookup space . envAllocCompilers case allocator of Nothing -> emit $ Imp.Allocate name' size' space Just allocator' -> allocator' name' size' -sAlloc :: String -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM lore r op VName +sAlloc :: String -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM rep r op VName sAlloc name size space = do name' <- sDeclareMem name space sAlloc_ name' size space return name' -sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore r op VName +sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM rep r op VName sArray name bt shape membind = do name' <- newVName name dArray name' bt shape membind return name' -- | Declare an array in row-major order in the given memory block. -sArrayInMem :: String -> PrimType -> ShapeBase SubExp -> VName -> ImpM lore r op VName +sArrayInMem :: String -> PrimType -> ShapeBase SubExp -> VName -> ImpM rep r op VName sArrayInMem name pt shape mem = sArray name pt shape $ ArrayIn mem $ IxFun.iota $ map (isInt64 . primExpFromSubExp int64) $ shapeDims shape -- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified. -sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM lore r op VName +sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM rep r op VName sAllocArrayPerm name pt shape space perm = do let permuted_dims = rearrangeShape perm $ shapeDims shape mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space @@ -1786,12 +1782,12 @@ ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm -- | Uses linear/iota index function. -sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM lore r op VName +sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName sAllocArray name pt shape space = sAllocArrayPerm name pt shape space [0 .. shapeRank shape -1] -- | Uses linear/iota index function. -sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore r op VName +sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM rep r op VName sStaticArray name space pt vs = do let num_elems = case vs of Imp.ArrayValues vs' -> length vs' @@ -1802,19 +1798,19 @@ addVar mem $ MemVar Nothing $ MemEntry space sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems] -sWrite :: VName -> [Imp.TExp Int64] -> Imp.Exp -> ImpM lore r op () +sWrite :: VName -> [Imp.TExp Int64] -> Imp.Exp -> ImpM rep r op () sWrite arr is v = do (mem, space, offset) <- fullyIndexArray arr is vol <- asks envVolatility emit $ Imp.Write mem offset (primExpType v) space vol v -sUpdate :: VName -> Slice (Imp.TExp Int64) -> SubExp -> ImpM lore r op () +sUpdate :: VName -> Slice (Imp.TExp Int64) -> SubExp -> ImpM rep r op () sUpdate arr slice v = copyDWIM arr slice v [] sLoopNest :: Shape -> - ([Imp.TExp Int64] -> ImpM lore r op ()) -> - ImpM lore r op () + ([Imp.TExp Int64] -> ImpM rep r op ()) -> + ImpM rep r op () sLoopNest = sLoopNest' [] . shapeDims where sLoopNest' is [] f = f $ reverse is @@ -1822,13 +1818,13 @@ sFor "nest_i" (toInt64Exp d) $ \i -> sLoopNest' (i : is) ds f -- | Untyped assignment. -(<~~) :: VName -> Imp.Exp -> ImpM lore r op () +(<~~) :: VName -> Imp.Exp -> ImpM rep r op () x <~~ e = emit $ Imp.SetScalar x e infixl 3 <~~ -- | Typed assignment. -(<--) :: TV t -> Imp.TExp t -> ImpM lore r op () +(<--) :: TV t -> Imp.TExp t -> ImpM rep r op () TV x _ <-- e = emit $ Imp.SetScalar x $ untyped e infixl 3 <-- @@ -1839,8 +1835,8 @@ Name -> [Imp.Param] -> [Imp.Param] -> - ImpM lore r op () -> - ImpM lore r op () + ImpM rep r op () -> + ImpM rep r op () function fname outputs inputs m = local newFunction $ do body <- collect $ do mapM_ addParam $ outputs ++ inputs @@ -1853,7 +1849,7 @@ addVar name $ ScalarVar Nothing $ ScalarEntry bt newFunction env = env {envFunction = Just fname} -dSlices :: [Imp.TExp Int64] -> ImpM lore r op [Imp.TExp Int64] +dSlices :: [Imp.TExp Int64] -> ImpM rep r op [Imp.TExp Int64] dSlices = fmap (drop 1 . snd) . dSlices' where dSlices' [] = pure (1, [1]) @@ -1869,7 +1865,7 @@ dIndexSpace :: [(VName, Imp.TExp Int64)] -> Imp.TExp Int64 -> - ImpM lore r op () + ImpM rep r op () dIndexSpace vs_ds j = do slices <- dSlices (map snd vs_ds) loop (zip (map fst vs_ds) slices) j diff --git a/src/Futhark/CodeGen/ImpGen/CUDA.hs b/src/Futhark/CodeGen/ImpGen/CUDA.hs --- a/src/Futhark/CodeGen/ImpGen/CUDA.hs +++ b/src/Futhark/CodeGen/ImpGen/CUDA.hs @@ -7,11 +7,11 @@ import Data.Bifunctor (second) import Futhark.CodeGen.ImpCode.OpenCL -import Futhark.CodeGen.ImpGen.Kernels -import Futhark.CodeGen.ImpGen.Kernels.ToOpenCL -import Futhark.IR.KernelsMem +import Futhark.CodeGen.ImpGen.GPU +import Futhark.CodeGen.ImpGen.GPU.ToOpenCL +import Futhark.IR.GPUMem import Futhark.MonadFreshNames -- | Compile the program to ImpCode with CUDA kernels. -compileProg :: MonadFreshNames m => Prog KernelsMem -> m (Warnings, Program) +compileProg :: MonadFreshNames m => Prog GPUMem -> m (Warnings, Program) compileProg prog = second kernelsToCUDA <$> compileProgCUDA prog diff --git a/src/Futhark/CodeGen/ImpGen/GPU.hs b/src/Futhark/CodeGen/ImpGen/GPU.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU.hs @@ -0,0 +1,431 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Compile a 'GPUMem' program to imperative code with kernels. +-- This is mostly (but not entirely) the same process no matter if we +-- are targeting OpenCL or CUDA. The important distinctions (the host +-- level code) are introduced later. +module Futhark.CodeGen.ImpGen.GPU + ( compileProgOpenCL, + compileProgCUDA, + Warnings, + ) +where + +import Control.Monad.Except +import Data.Bifunctor (second) +import Data.List (foldl') +import qualified Data.Map as M +import Data.Maybe +import Futhark.CodeGen.ImpCode.GPU (bytes) +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen hiding (compileProg) +import qualified Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.CodeGen.ImpGen.GPU.SegHist +import Futhark.CodeGen.ImpGen.GPU.SegMap +import Futhark.CodeGen.ImpGen.GPU.SegRed +import Futhark.CodeGen.ImpGen.GPU.SegScan +import Futhark.CodeGen.ImpGen.GPU.Transpose +import Futhark.CodeGen.SetDefaultSpace +import Futhark.Error +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Util.IntegralExp (IntegralExp, divUp, quot, rem) +import Prelude hiding (quot, rem) + +callKernelOperations :: Operations GPUMem HostEnv Imp.HostOp +callKernelOperations = + Operations + { opsExpCompiler = expCompiler, + opsCopyCompiler = callKernelCopy, + opsOpCompiler = opCompiler, + opsStmsCompiler = defCompileStms, + opsAllocCompilers = mempty + } + +openclAtomics, cudaAtomics :: AtomicBinOp +(openclAtomics, cudaAtomics) = (flip lookup opencl, flip lookup cuda) + where + opencl64 = + [ (Add Int64 OverflowUndef, Imp.AtomicAdd Int64), + (SMax Int64, Imp.AtomicSMax Int64), + (SMin Int64, Imp.AtomicSMin Int64), + (UMax Int64, Imp.AtomicUMax Int64), + (UMin Int64, Imp.AtomicUMin Int64), + (And Int64, Imp.AtomicAnd Int64), + (Or Int64, Imp.AtomicOr Int64), + (Xor Int64, Imp.AtomicXor Int64) + ] + opencl32 = + [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32), + (SMax Int32, Imp.AtomicSMax Int32), + (SMin Int32, Imp.AtomicSMin Int32), + (UMax Int32, Imp.AtomicUMax Int32), + (UMin Int32, Imp.AtomicUMin Int32), + (And Int32, Imp.AtomicAnd Int32), + (Or Int32, Imp.AtomicOr Int32), + (Xor Int32, Imp.AtomicXor Int32) + ] + opencl = opencl32 ++ opencl64 + cuda = + opencl + ++ [ (FAdd Float32, Imp.AtomicFAdd Float32), + (FAdd Float64, Imp.AtomicFAdd Float64) + ] + +compileProg :: + MonadFreshNames m => + HostEnv -> + Prog GPUMem -> + m (Warnings, Imp.Program) +compileProg env prog = + second (setDefaultSpace (Imp.Space "device")) + <$> Futhark.CodeGen.ImpGen.compileProg env callKernelOperations (Imp.Space "device") prog + +-- | Compile a 'GPUMem' program to low-level parallel code, with +-- either CUDA or OpenCL characteristics. +compileProgOpenCL, + compileProgCUDA :: + MonadFreshNames m => Prog GPUMem -> m (Warnings, Imp.Program) +compileProgOpenCL = compileProg $ HostEnv openclAtomics OpenCL mempty +compileProgCUDA = compileProg $ HostEnv cudaAtomics CUDA mempty + +opCompiler :: + Pattern GPUMem -> + Op GPUMem -> + CallKernelGen () +opCompiler dest (Alloc e space) = + compileAlloc dest e space +opCompiler (Pattern _ [pe]) (Inner (SizeOp (GetSize key size_class))) = do + fname <- askFunction + sOp $ + Imp.GetSize (patElemName pe) (keyWithEntryPoint fname key) $ + sizeClassWithEntryPoint fname size_class +opCompiler (Pattern _ [pe]) (Inner (SizeOp (CmpSizeLe key size_class x))) = do + fname <- askFunction + let size_class' = sizeClassWithEntryPoint fname size_class + sOp . Imp.CmpSizeLe (patElemName pe) (keyWithEntryPoint fname key) size_class' + =<< toExp x +opCompiler (Pattern _ [pe]) (Inner (SizeOp (GetSizeMax size_class))) = + sOp $ Imp.GetSizeMax (patElemName pe) size_class +opCompiler (Pattern _ [pe]) (Inner (SizeOp (CalcNumGroups w64 max_num_groups_key group_size))) = do + fname <- askFunction + max_num_groups :: TV Int32 <- dPrim "max_num_groups" int32 + sOp $ + Imp.GetSize (tvVar max_num_groups) (keyWithEntryPoint fname max_num_groups_key) $ + sizeClassWithEntryPoint fname SizeNumGroups + + -- If 'w' is small, we launch fewer groups than we normally would. + -- We don't want any idle groups. + -- + -- The calculations are done with 64-bit integers to avoid overflow + -- issues. + let num_groups_maybe_zero = + sMin64 (toInt64Exp w64 `divUp` toInt64Exp group_size) $ + sExt64 (tvExp max_num_groups) + -- We also don't want zero groups. + let num_groups = sMax64 1 num_groups_maybe_zero + mkTV (patElemName pe) int32 <-- sExt32 num_groups +opCompiler dest (Inner (SegOp op)) = + segOpCompiler dest op +opCompiler pat e = + compilerBugS $ + "opCompiler: Invalid pattern\n " + ++ pretty pat + ++ "\nfor expression\n " + ++ pretty e + +sizeClassWithEntryPoint :: Maybe Name -> Imp.SizeClass -> Imp.SizeClass +sizeClassWithEntryPoint fname (Imp.SizeThreshold path def) = + Imp.SizeThreshold (map f path) def + where + f (name, x) = (keyWithEntryPoint fname name, x) +sizeClassWithEntryPoint _ size_class = size_class + +segOpCompiler :: + Pattern GPUMem -> + SegOp SegLevel GPUMem -> + CallKernelGen () +segOpCompiler pat (SegMap lvl space _ kbody) = + compileSegMap pat lvl space kbody +segOpCompiler pat (SegRed lvl@SegThread {} space reds _ kbody) = + compileSegRed pat lvl space reds kbody +segOpCompiler pat (SegScan lvl@SegThread {} space scans _ kbody) = + compileSegScan pat lvl space scans kbody +segOpCompiler pat (SegHist (SegThread num_groups group_size _) space ops _ kbody) = + compileSegHist pat num_groups group_size space ops kbody +segOpCompiler pat segop = + compilerBugS $ "segOpCompiler: unexpected " ++ pretty (segLevel segop) ++ " for rhs of pattern " ++ pretty pat + +-- Create boolean expression that checks whether all kernels in the +-- enclosed code do not use more local memory than we have available. +-- We look at *all* the kernels here, even those that might be +-- otherwise protected by their own multi-versioning branches deeper +-- down. Currently the compiler will not generate multi-versioning +-- that makes this a problem, but it might in the future. +checkLocalMemoryReqs :: Imp.Code -> CallKernelGen (Maybe (Imp.TExp Bool)) +checkLocalMemoryReqs code = do + scope <- askScope + let alloc_sizes = map (sum . map alignedSize . localAllocSizes . Imp.kernelBody) $ getGPU code + + -- If any of the sizes involve a variable that is not known at this + -- point, then we cannot check the requirements. + if any (`M.notMember` scope) (namesToList $ freeIn alloc_sizes) + then return Nothing + else do + local_memory_capacity :: TV Int32 <- dPrim "local_memory_capacity" int32 + sOp $ Imp.GetSizeMax (tvVar local_memory_capacity) SizeLocalMemory + + let local_memory_capacity_64 = + sExt64 $ tvExp local_memory_capacity + fits size = + unCount size .<=. local_memory_capacity_64 + return $ Just $ foldl' (.&&.) true (map fits alloc_sizes) + where + getGPU = foldMap getKernel + getKernel (Imp.CallKernel k) = [k] + getKernel _ = [] + + localAllocSizes = foldMap localAllocSize + localAllocSize (Imp.LocalAlloc _ size) = [size] + localAllocSize _ = [] + + -- These allocations will actually be padded to an 8-byte aligned + -- size, so we should take that into account when checking whether + -- they fit. + alignedSize x = x + ((8 - (x `rem` 8)) `rem` 8) + +withAcc :: + Pattern GPUMem -> + [(Shape, [VName], Maybe (Lambda GPUMem, [SubExp]))] -> + Lambda GPUMem -> + CallKernelGen () +withAcc pat inputs lam = do + atomics <- hostAtomics <$> askEnv + locksForInputs atomics $ zip accs inputs + where + accs = map paramName $ lambdaParams lam + locksForInputs _ [] = + defCompileExp pat $ WithAcc inputs lam + locksForInputs atomics ((c, (_, _, op)) : inputs') + | Just (op_lam, _) <- op, + AtomicLocking _ <- atomicUpdateLocking atomics op_lam = do + let num_locks = 100151 + locks_arr <- + sStaticArray "withacc_locks" (Space "device") int32 $ + Imp.ArrayZeros num_locks + let locks = Locks locks_arr num_locks + extend env = env {hostLocks = M.insert c locks $ hostLocks env} + localEnv extend $ locksForInputs atomics inputs' + | otherwise = + locksForInputs atomics inputs' + +expCompiler :: ExpCompiler GPUMem HostEnv Imp.HostOp +-- We generate a simple kernel for itoa and replicate. +expCompiler (Pattern _ [pe]) (BasicOp (Iota n x s et)) = do + x' <- toExp x + s' <- toExp s + + sIota (patElemName pe) (toInt64Exp n) x' s' et +expCompiler (Pattern _ [pe]) (BasicOp (Replicate _ se)) = + sReplicate (patElemName pe) se +-- Allocation in the "local" space is just a placeholder. +expCompiler _ (Op (Alloc _ (Space "local"))) = + return () +expCompiler pat (WithAcc inputs lam) = + withAcc pat inputs lam +-- This is a multi-versioning If created by incremental flattening. +-- We need to augment the conditional with a check that any local +-- memory requirements in tbranch are compatible with the hardware. +-- We do not check anything for fbranch, as we assume that it will +-- always be safe (and what would we do if none of the branches would +-- work?). +expCompiler dest (If cond tbranch fbranch (IfDec _ IfEquiv)) = do + tcode <- collect $ compileBody dest tbranch + fcode <- collect $ compileBody dest fbranch + check <- checkLocalMemoryReqs tcode + emit $ case check of + Nothing -> fcode + Just ok -> Imp.If (ok .&&. toBoolExp cond) tcode fcode +expCompiler dest e = + defCompileExp dest e + +callKernelCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp +callKernelCopy + bt + destloc@(MemLocation destmem _ destIxFun) + destslice + srcloc@(MemLocation srcmem srcshape srcIxFun) + srcslice + | Just + ( destoffset, + srcoffset, + num_arrays, + size_x, + size_y + ) <- + isMapTransposeCopy bt destloc destslice srcloc srcslice = do + fname <- mapTransposeForType bt + emit $ + Imp.Call + [] + fname + [ Imp.MemArg destmem, + Imp.ExpArg $ untyped destoffset, + Imp.MemArg srcmem, + Imp.ExpArg $ untyped srcoffset, + Imp.ExpArg $ untyped num_arrays, + Imp.ExpArg $ untyped size_x, + Imp.ExpArg $ untyped size_y + ] + | bt_size <- primByteSize bt, + Just destoffset <- + IxFun.linearWithOffset (IxFun.slice destIxFun destslice) bt_size, + Just srcoffset <- + IxFun.linearWithOffset (IxFun.slice srcIxFun srcslice) bt_size = do + let num_elems = Imp.elements $ product $ map toInt64Exp srcshape + srcspace <- entryMemSpace <$> lookupMemory srcmem + destspace <- entryMemSpace <$> lookupMemory destmem + emit $ + Imp.Copy + destmem + (bytes $ sExt64 destoffset) + destspace + srcmem + (bytes $ sExt64 srcoffset) + srcspace + $ num_elems `Imp.withElemType` bt + | otherwise = sCopy bt destloc destslice srcloc srcslice + +mapTransposeForType :: PrimType -> CallKernelGen Name +mapTransposeForType bt = do + let fname = nameFromString $ "builtin#" <> mapTransposeName bt + + exists <- hasFunction fname + unless exists $ emitFunction fname $ mapTransposeFunction bt + + return fname + +mapTransposeName :: PrimType -> String +mapTransposeName bt = "gpu_map_transpose_" ++ pretty bt + +mapTransposeFunction :: PrimType -> Imp.Function +mapTransposeFunction bt = + Imp.Function Nothing [] params transpose_code [] [] + where + params = + [ memparam destmem, + intparam destoffset, + memparam srcmem, + intparam srcoffset, + intparam num_arrays, + intparam x, + intparam y + ] + + space = Space "device" + memparam v = Imp.MemParam v space + intparam v = Imp.ScalarParam v $ IntType Int32 + + [ destmem, + destoffset, + srcmem, + srcoffset, + num_arrays, + x, + y, + mulx, + muly, + block + ] = + zipWith + (VName . nameFromString) + [ "destmem", + "destoffset", + "srcmem", + "srcoffset", + "num_arrays", + "x_elems", + "y_elems", + -- The following is only used for low width/height + -- transpose kernels + "mulx", + "muly", + "block" + ] + [0 ..] + + block_dim_int = 16 + + block_dim :: IntegralExp a => a + block_dim = 16 + + -- When an input array has either width==1 or height==1, performing a + -- transpose will be the same as performing a copy. + can_use_copy = + let onearr = Imp.vi32 num_arrays .==. 1 + height_is_one = Imp.vi32 y .==. 1 + width_is_one = Imp.vi32 x .==. 1 + in onearr .&&. (width_is_one .||. height_is_one) + + transpose_code = + Imp.If input_is_empty mempty $ + mconcat + [ Imp.DeclareScalar muly Imp.Nonvolatile (IntType Int32), + Imp.SetScalar muly $ untyped $ block_dim `quot` Imp.vi32 x, + Imp.DeclareScalar mulx Imp.Nonvolatile (IntType Int32), + Imp.SetScalar mulx $ untyped $ block_dim `quot` Imp.vi32 y, + Imp.If can_use_copy copy_code $ + Imp.If should_use_lowwidth (callTransposeKernel TransposeLowWidth) $ + Imp.If should_use_lowheight (callTransposeKernel TransposeLowHeight) $ + Imp.If should_use_small (callTransposeKernel TransposeSmall) $ + callTransposeKernel TransposeNormal + ] + + input_is_empty = + Imp.vi32 num_arrays .==. 0 .||. Imp.vi32 x .==. 0 .||. Imp.vi32 y .==. 0 + + should_use_small = + Imp.vi32 x .<=. (block_dim `quot` 2) + .&&. Imp.vi32 y .<=. (block_dim `quot` 2) + + should_use_lowwidth = + Imp.vi32 x .<=. (block_dim `quot` 2) + .&&. block_dim .<. Imp.vi32 y + + should_use_lowheight = + Imp.vi32 y .<=. (block_dim `quot` 2) + .&&. block_dim .<. Imp.vi32 x + + copy_code = + let num_bytes = sExt64 $ Imp.vi32 x * Imp.vi32 y * primByteSize bt + in Imp.Copy + destmem + (Imp.Count $ sExt64 $ Imp.vi32 destoffset) + space + srcmem + (Imp.Count $ sExt64 $ Imp.vi32 srcoffset) + space + (Imp.Count num_bytes) + + callTransposeKernel = + Imp.Op . Imp.CallKernel + . mapTransposeKernel + (mapTransposeName bt) + block_dim_int + ( destmem, + Imp.vi32 destoffset, + srcmem, + Imp.vi32 srcoffset, + Imp.vi32 x, + Imp.vi32 y, + Imp.vi32 mulx, + Imp.vi32 muly, + Imp.vi32 num_arrays, + block + ) + bt diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs @@ -0,0 +1,1785 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.CodeGen.ImpGen.GPU.Base + ( KernelConstants (..), + keyWithEntryPoint, + CallKernelGen, + InKernelGen, + Locks (..), + HostEnv (..), + Target (..), + KernelEnv (..), + computeThreadChunkSize, + groupReduce, + groupScan, + isActive, + sKernelThread, + sKernelGroup, + sReplicate, + sIota, + sCopy, + compileThreadResult, + compileGroupResult, + virtualiseGroups, + groupLoop, + kernelLoop, + groupCoverSpace, + precomputeSegOpIDs, + atomicUpdateLocking, + AtomicBinOp, + Locking (..), + AtomicUpdate (..), + DoAtomicUpdate, + ) +where + +import Control.Monad.Except +import Data.List (zip4) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.Error +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Transform.Rename +import Futhark.Util (chunks, dropLast, mapAccumLM, nubOrd, takeLast) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +-- | Which target are we ultimately generating code for? While most +-- of the kernels code is the same, there are some cases where we +-- generate special code based on the ultimate low-level API we are +-- targeting. +data Target = CUDA | OpenCL + +-- | Information about the locks available for accumulators. +data Locks = Locks + { locksArray :: VName, + locksCount :: Int + } + +data HostEnv = HostEnv + { hostAtomics :: AtomicBinOp, + hostTarget :: Target, + hostLocks :: M.Map VName Locks + } + +data KernelEnv = KernelEnv + { kernelAtomics :: AtomicBinOp, + kernelConstants :: KernelConstants, + kernelLocks :: M.Map VName Locks + } + +type CallKernelGen = ImpM GPUMem HostEnv Imp.HostOp + +type InKernelGen = ImpM GPUMem KernelEnv Imp.KernelOp + +data KernelConstants = KernelConstants + { kernelGlobalThreadId :: Imp.TExp Int32, + kernelLocalThreadId :: Imp.TExp Int32, + kernelGroupId :: Imp.TExp Int32, + kernelGlobalThreadIdVar :: VName, + kernelLocalThreadIdVar :: VName, + kernelGroupIdVar :: VName, + kernelNumGroups :: Imp.TExp Int64, + kernelGroupSize :: Imp.TExp Int64, + kernelNumThreads :: Imp.TExp Int32, + kernelWaveSize :: Imp.TExp Int32, + kernelThreadActive :: Imp.TExp Bool, + -- | A mapping from dimensions of nested SegOps to already + -- computed local thread IDs. + kernelLocalIdMap :: M.Map [SubExp] [Imp.TExp Int32] + } + +segOpSizes :: Stms GPUMem -> S.Set [SubExp] +segOpSizes = onStms + where + onStms = foldMap (onExp . stmExp) + onExp (Op (Inner (SegOp op))) = + S.singleton $ map snd $ unSegSpace $ segSpace op + onExp (If _ tbranch fbranch _) = + onStms (bodyStms tbranch) <> onStms (bodyStms fbranch) + onExp (DoLoop _ _ _ body) = + onStms (bodyStms body) + onExp _ = mempty + +precomputeSegOpIDs :: Stms GPUMem -> InKernelGen a -> InKernelGen a +precomputeSegOpIDs stms m = do + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (segOpSizes stms)) + let f env = + env + { kernelConstants = + (kernelConstants env) {kernelLocalIdMap = new_ids} + } + localEnv f m + where + mkMap ltid dims = do + let dims' = map (sExt32 . toInt64Exp) dims + ids' <- mapM (dPrimVE "ltid_pre") $ unflattenIndex dims' ltid + return (dims, ids') + +keyWithEntryPoint :: Maybe Name -> Name -> Name +keyWithEntryPoint fname key = + nameFromString $ maybe "" ((++ ".") . nameToString) fname ++ nameToString key + +allocLocal :: AllocCompiler GPUMem r Imp.KernelOp +allocLocal mem size = + sOp $ Imp.LocalAlloc mem size + +kernelAlloc :: + Pattern GPUMem -> + SubExp -> + Space -> + InKernelGen () +kernelAlloc (Pattern _ [_]) _ ScalarSpace {} = + -- Handled by the declaration of the memory block, which is then + -- translated to an actual scalar variable during C code generation. + return () +kernelAlloc (Pattern _ [mem]) size (Space "local") = + allocLocal (patElemName mem) $ Imp.bytes $ toInt64Exp size +kernelAlloc (Pattern _ [mem]) _ _ = + compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." +kernelAlloc dest _ _ = + error $ "Invalid target for in-kernel allocation: " ++ show dest + +splitSpace :: + (ToExp w, ToExp i, ToExp elems_per_thread) => + Pattern GPUMem -> + SplitOrdering -> + w -> + i -> + elems_per_thread -> + ImpM rep r op () +splitSpace (Pattern [] [size]) o w i elems_per_thread = do + num_elements <- Imp.elements . TPrimExp <$> toExp w + let i' = toInt64Exp i + elems_per_thread' <- Imp.elements . TPrimExp <$> toExp elems_per_thread + computeThreadChunkSize o i' elems_per_thread' num_elements (mkTV (patElemName size) int64) +splitSpace pat _ _ _ _ = + error $ "Invalid target for splitSpace: " ++ pretty pat + +updateAcc :: VName -> [SubExp] -> [SubExp] -> InKernelGen () +updateAcc acc is vs = sComment "UpdateAcc" $ do + -- See the ImpGen implementation of UpdateAcc for general notes. + let is' = map toInt64Exp is + (c, space, arrs, dims, op) <- lookupAcc acc is' + sWhen (inBounds (map DimFix is') dims) $ + case op of + Nothing -> + forM_ (zip arrs vs) $ \(arr, v) -> copyDWIMFix arr is' v [] + Just lam -> do + dLParams $ lambdaParams lam + let (_x_params, y_params) = + splitAt (length vs) $ map paramName $ lambdaParams lam + forM_ (zip y_params vs) $ \(yp, v) -> copyDWIM yp [] v [] + atomics <- kernelAtomics <$> askEnv + case atomicUpdateLocking atomics lam of + AtomicPrim f -> f space arrs is' + AtomicCAS f -> f space arrs is' + AtomicLocking f -> do + c_locks <- M.lookup c . kernelLocks <$> askEnv + case c_locks of + Just (Locks locks num_locks) -> do + let locking = + Locking locks 0 1 0 $ + pure . (`rem` fromIntegral num_locks) . flattenIndex dims + f locking space arrs is' + Nothing -> + error $ "Missing locks for " ++ pretty acc + +compileThreadExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp +compileThreadExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0 ..] es) $ \(i, e) -> + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] +compileThreadExp _ (BasicOp (UpdateAcc acc is vs)) = + updateAcc acc is vs +compileThreadExp dest e = + defCompileExp dest e + +-- | Assign iterations of a for-loop to all threads in the kernel. +-- The passed-in function is invoked with the (symbolic) iteration. +-- 'threadOperations' will be in effect in the body. For +-- multidimensional loops, use 'groupCoverSpace'. +kernelLoop :: + IntExp t => + Imp.TExp t -> + Imp.TExp t -> + Imp.TExp t -> + (Imp.TExp t -> InKernelGen ()) -> + InKernelGen () +kernelLoop tid num_threads n f = + localOps threadOperations $ + if n == num_threads + then f tid + else do + -- Compute how many elements this thread is responsible for. + -- Formula: (n - tid) / num_threads (rounded up). + let elems_for_this = (n - tid) `divUp` num_threads + + sFor "i" elems_for_this $ \i -> f $ i * num_threads + tid + +-- | Assign iterations of a for-loop to threads in the workgroup. The +-- passed-in function is invoked with the (symbolic) iteration. For +-- multidimensional loops, use 'groupCoverSpace'. +groupLoop :: + Imp.TExp Int64 -> + (Imp.TExp Int64 -> InKernelGen ()) -> + InKernelGen () +groupLoop n f = do + constants <- kernelConstants <$> askEnv + kernelLoop + (sExt64 $ kernelLocalThreadId constants) + (kernelGroupSize constants) + n + f + +-- | Iterate collectively though a multidimensional space, such that +-- all threads in the group participate. The passed-in function is +-- invoked with a (symbolic) point in the index space. +groupCoverSpace :: + [Imp.TExp Int64] -> + ([Imp.TExp Int64] -> InKernelGen ()) -> + InKernelGen () +groupCoverSpace ds f = + groupLoop (product ds) $ f . unflattenIndex ds + +compileGroupExp :: ExpCompiler GPUMem KernelEnv Imp.KernelOp +-- The static arrays stuff does not work inside kernels. +compileGroupExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0 ..] es) $ \(i, e) -> + copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] +compileGroupExp _ (BasicOp (UpdateAcc acc is vs)) = + updateAcc acc is vs +compileGroupExp (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do + let ds' = map toInt64Exp $ shapeDims ds + groupCoverSpace ds' $ \is -> + copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is) + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupExp (Pattern _ [dest]) (BasicOp (Iota n e s it)) = do + n' <- toExp n + e' <- toExp e + s' <- toExp s + groupLoop (TPrimExp n') $ \i' -> do + x <- + dPrimV "x" $ + TPrimExp $ + BinOpExp (Add it OverflowUndef) e' $ + BinOpExp (Mul it OverflowUndef) (untyped i') s' + copyDWIMFix (patElemName dest) [i'] (Var (tvVar x)) [] + sOp $ Imp.Barrier Imp.FenceLocal + +-- When generating code for a scalar in-place update, we must make +-- sure that only one thread performs the write. When writing an +-- array, the group-level copy code will take care of doing the right +-- thing. +compileGroupExp (Pattern _ [pe]) (BasicOp (Update _ slice se)) + | null $ sliceDims slice = do + sOp $ Imp.Barrier Imp.FenceLocal + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + sWhen (ltid .==. 0) $ + copyDWIM (patElemName pe) (map (fmap toInt64Exp) slice) se [] + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupExp dest e = + defCompileExp dest e + +sanityCheckLevel :: SegLevel -> InKernelGen () +sanityCheckLevel SegThread {} = return () +sanityCheckLevel SegGroup {} = + error "compileGroupOp: unexpected group-level SegOp." + +localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] +localThreadIDs dims = do + ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv + let dims' = map toInt64Exp dims + maybe (unflattenIndex dims' ltid) (map sExt64) + . M.lookup dims + . kernelLocalIdMap + . kernelConstants + <$> askEnv + +compileGroupSpace :: SegLevel -> SegSpace -> InKernelGen () +compileGroupSpace lvl space = do + sanityCheckLevel lvl + let (ltids, dims) = unzip $ unSegSpace space + zipWithM_ dPrimV_ ltids =<< localThreadIDs dims + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + dPrimV_ (segFlat space) ltid + +-- Construct the necessary lock arrays for an intra-group histogram. +prepareIntraGroupSegHist :: + Count GroupSize SubExp -> + [HistOp GPUMem] -> + InKernelGen [[Imp.TExp Int64] -> InKernelGen ()] +prepareIntraGroupSegHist group_size = + fmap snd . mapAccumLM onOp Nothing + where + onOp l op = do + constants <- kernelConstants <$> askEnv + atomicBinOp <- kernelAtomics <$> askEnv + + let local_subhistos = histDest op + + case (l, atomicUpdateLocking atomicBinOp $ histOp op) of + (_, AtomicPrim f) -> return (l, f (Space "local") local_subhistos) + (_, AtomicCAS f) -> return (l, f (Space "local") local_subhistos) + (Just l', AtomicLocking f) -> return (l, f l' (Space "local") local_subhistos) + (Nothing, AtomicLocking f) -> do + locks <- newVName "locks" + + let num_locks = toInt64Exp $ unCount group_size + dims = map toInt64Exp $ shapeDims (histShape op) ++ [histWidth op] + l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims) + locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness + + locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local" + dArray locks int32 (arrayShape locks_t) $ + ArrayIn locks_mem $ + IxFun.iota $ + map pe64 $ arrayDims locks_t + + sComment "All locks start out unlocked" $ + groupCoverSpace [kernelGroupSize constants] $ \is -> + copyDWIMFix locks is (intConst Int32 0) [] + + return (Just l', f l' (Space "local") local_subhistos) + +whenActive :: SegLevel -> SegSpace -> InKernelGen () -> InKernelGen () +whenActive lvl space m + | SegNoVirtFull <- segVirt lvl = m + | otherwise = do + group_size <- kernelGroupSize . kernelConstants <$> askEnv + -- XXX: the following check is too naive - we should also handle + -- the multi-dimensional case. + if [group_size] == map (toInt64Exp . snd) (unSegSpace space) + then m + else sWhen (isActive $ unSegSpace space) m + +compileGroupOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp +compileGroupOp pat (Alloc size space) = + kernelAlloc pat size space +compileGroupOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = + splitSpace pat o w i elems_per_thread +compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do + void $ compileGroupSpace lvl space + + whenActive lvl space $ + localOps threadOperations $ + compileStms mempty (kernelBodyStms body) $ + zipWithM_ (compileThreadResult space) (patternElements pat) $ + kernelBodyResult body + + sOp $ Imp.ErrorSync Imp.FenceLocal +compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do + compileGroupSpace lvl space + let (ltids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + + whenActive lvl space $ + compileStms mempty (kernelBodyStms body) $ + forM_ (zip (patternNames pat) $ kernelBodyResult body) $ \(dest, res) -> + copyDWIMFix + dest + (map Imp.vi64 ltids) + (kernelResultSubExp res) + [] + + sOp $ Imp.ErrorSync Imp.FenceLocal + + let segment_size = last dims' + crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) + + -- groupScan needs to treat the scan output as a one-dimensional + -- array of scan elements, so we invent some new flattened arrays + -- here. XXX: this assumes that the original index function is just + -- row-major, but does not actually verify it. + dims_flat <- dPrimV "dims_flat" $ product dims' + let flattened pe = do + MemLocation mem _ _ <- + entryArrayLocation <$> lookupArray (patElemName pe) + let pe_t = typeOf pe + arr_dims = Var (tvVar dims_flat) : drop (length dims') (arrayDims pe_t) + sArray + (baseString (patElemName pe) ++ "_flat") + (elemType pe_t) + (Shape arr_dims) + $ ArrayIn mem $ IxFun.iota $ map pe64 arr_dims + + num_scan_results = sum $ map (length . segBinOpNeutral) scans + + arrs_flat <- mapM flattened $ take num_scan_results $ patternElements pat + + forM_ scans $ \scan -> do + let scan_op = segBinOpLambda scan + groupScan (Just crossesSegment) (product dims') (product dims') scan_op arrs_flat +compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do + compileGroupSpace lvl space + + let (ltids, dims) = unzip $ unSegSpace space + (red_pes, map_pes) = + splitAt (segBinOpResults ops) $ patternElements pat + + dims' = map toInt64Exp dims + + mkTempArr t = + sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" + + tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops + let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs + + whenActive lvl space $ + compileStms mempty (kernelBodyStms body) $ do + let (red_res, map_res) = + splitAt (segBinOpResults ops) $ kernelBodyResult body + forM_ (zip tmp_arrs red_res) $ \(dest, res) -> + copyDWIMFix dest (map Imp.vi64 ltids) (kernelResultSubExp res) [] + zipWithM_ (compileThreadResult space) map_pes map_res + + sOp $ Imp.ErrorSync Imp.FenceLocal + + case dims' of + -- Nonsegmented case (or rather, a single segment) - this we can + -- handle directly with a group-level reduction. + [dim'] -> do + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> + groupReduce (sExt32 dim') (segBinOpLambda op) tmps + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> + copyDWIMFix (patElemName pe) [] (Var arr) [0] + _ -> do + -- Segmented intra-group reductions are turned into (regular) + -- segmented scans. It is possible that this can be done + -- better, but at least this approach is simple. + + -- groupScan operates on flattened arrays. This does not + -- involve copying anything; merely playing with the index + -- function. + dims_flat <- dPrimV "dims_flat" $ product dims' + let flatten arr = do + ArrayEntry arr_loc pt <- lookupArray arr + let flat_shape = + Shape $ + Var (tvVar dims_flat) : + drop (length ltids) (memLocationShape arr_loc) + sArray "red_arr_flat" pt flat_shape $ + ArrayIn (memLocationName arr_loc) $ + IxFun.iota $ map pe64 $ shapeDims flat_shape + + let segment_size = last dims' + crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size) + + forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do + tmps_flat <- mapM flatten tmps + groupScan + (Just crossesSegment) + (product dims') + (product dims') + (segBinOpLambda op) + tmps_flat + + sOp $ Imp.ErrorSync Imp.FenceLocal + + forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> + copyDWIM + (patElemName pe) + [] + (Var arr) + (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' -1]) + + sOp $ Imp.Barrier Imp.FenceLocal +compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do + compileGroupSpace lvl space + let ltids = map fst $ unSegSpace space + + -- We don't need the red_pes, because it is guaranteed by our type + -- rules that they occupy the same memory as the destinations for + -- the ops. + let num_red_res = length ops + sum (map (length . histNeutral) ops) + (_red_pes, map_pes) = + splitAt num_red_res $ patternElements pat + + ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops + + -- Ensure that all locks have been initialised. + sOp $ Imp.Barrier Imp.FenceLocal + + whenActive lvl space $ + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody + (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res + zipWithM_ (compileThreadResult space) map_pes map_res + + let vs_per_op = chunks (map (length . histDest) ops) red_vs + + forM_ (zip4 red_is vs_per_op ops' ops) $ + \(bin, op_vs, do_op, HistOp dest_w _ _ _ shape lam) -> do + let bin' = toInt64Exp bin + dest_w' = toInt64Exp dest_w + bin_in_bounds = 0 .<=. bin' .&&. bin' .<. dest_w' + bin_is = map Imp.vi64 (init ltids) ++ [bin'] + vs_params = takeLast (length op_vs) $ lambdaParams lam + + sComment "perform atomic updates" $ + sWhen bin_in_bounds $ do + dLParams $ lambdaParams lam + sLoopNest shape $ \is -> do + forM_ (zip vs_params op_vs) $ \(p, v) -> + copyDWIMFix (paramName p) [] v is + do_op (bin_is ++ is) + + sOp $ Imp.ErrorSync Imp.FenceLocal +compileGroupOp pat _ = + compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat + +compileThreadOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp +compileThreadOp pat (Alloc size space) = + kernelAlloc pat size space +compileThreadOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = + splitSpace pat o w i elems_per_thread +compileThreadOp pat _ = + compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat + +-- | Locking strategy used for an atomic update. +data Locking = Locking + { -- | Array containing the lock. + lockingArray :: VName, + -- | Value for us to consider the lock free. + lockingIsUnlocked :: Imp.TExp Int32, + -- | What to write when we lock it. + lockingToLock :: Imp.TExp Int32, + -- | What to write when we unlock it. + lockingToUnlock :: Imp.TExp Int32, + -- | A transformation from the logical lock index to the + -- physical position in the array. This can also be used + -- to make the lock array smaller. + lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64] + } + +-- | A function for generating code for an atomic update. Assumes +-- that the bucket is in-bounds. +type DoAtomicUpdate rep r = + Space -> [VName] -> [Imp.TExp Int64] -> ImpM rep r Imp.KernelOp () + +-- | The mechanism that will be used for performing the atomic update. +-- Approximates how efficient it will be. Ordered from most to least +-- efficient. +data AtomicUpdate rep r + = -- | Supported directly by primitive. + AtomicPrim (DoAtomicUpdate rep r) + | -- | Can be done by efficient swaps. + AtomicCAS (DoAtomicUpdate rep r) + | -- | Requires explicit locking. + AtomicLocking (Locking -> DoAtomicUpdate rep r) + +-- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? +type AtomicBinOp = + BinOp -> + Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int64) -> Imp.Exp -> Imp.AtomicOp) + +-- | Do an atomic update corresponding to a binary operator lambda. +atomicUpdateLocking :: + AtomicBinOp -> + Lambda GPUMem -> + AtomicUpdate GPUMem KernelEnv +atomicUpdateLocking atomicBinOp lam + | Just ops_and_ts <- lamIsBinOp lam, + all (\(_, t, _, _) -> primBitSize t `elem` [32, 64]) ops_and_ts = + primOrCas ops_and_ts $ \space arrs bucket -> + -- If the operator is a vectorised binary operator on 32/64-bit + -- values, we can use a particularly efficient + -- implementation. If the operator has an atomic implementation + -- we use that, otherwise it is still a binary operator which + -- can be implemented by atomic compare-and-swap if 32/64 bits. + forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do + -- Common variables. + old <- dPrim "old" t + + (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket + + case opHasAtomicSupport space (tvVar old) arr' bucket_offset op of + Just f -> sOp $ f $ Imp.var y t + Nothing -> + atomicUpdateCAS space t a (tvVar old) bucket x $ + x <~~ Imp.BinOpExp op (Imp.var x t) (Imp.var y t) + where + opHasAtomicSupport space old arr' bucket' bop = do + let atomic f = Imp.Atomic space . f old arr' bucket' + atomic <$> atomicBinOp bop + + primOrCas ops + | all isPrim ops = AtomicPrim + | otherwise = AtomicCAS + + isPrim (op, _, _, _) = isJust $ atomicBinOp op + +-- If the operator functions purely on single 32/64-bit values, we can +-- use an implementation based on CAS, no matter what the operator +-- does. +atomicUpdateLocking _ op + | [Prim t] <- lambdaReturnType op, + [xp, _] <- lambdaParams op, + primBitSize t `elem` [32, 64] = AtomicCAS $ \space [arr] bucket -> do + old <- dPrim "old" t + atomicUpdateCAS space t arr (tvVar old) bucket (paramName xp) $ + compileBody' [xp] $ lambdaBody op +atomicUpdateLocking _ op = AtomicLocking $ \locking space arrs bucket -> do + old <- dPrim "old" int32 + continue <- dPrimVol "continue" Bool true + + -- Correctly index into locks. + (locks', _locks_space, locks_offset) <- + fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket + + -- Critical section + let try_acquire_lock = + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + locks_offset + (untyped $ lockingIsUnlocked locking) + (untyped $ lockingToLock locking) + lock_acquired = tvExp old .==. lockingIsUnlocked locking + -- Even the releasing is done with an atomic rather than a + -- simple write, for memory coherency reasons. + release_lock = + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int32 + (tvVar old) + locks' + locks_offset + (untyped $ lockingToLock locking) + (untyped $ lockingToUnlock locking) + break_loop = continue <-- false + + -- Preparing parameters. It is assumed that the caller has already + -- filled the arr_params. We copy the current value to the + -- accumulator parameters. + -- + -- Note the use of 'everythingVolatile' when reading and writing the + -- buckets. This was necessary to ensure correct execution on a + -- newer NVIDIA GPU (RTX 2080). The 'volatile' modifiers likely + -- make the writes pass through the (SM-local) L1 cache, which is + -- necessary here, because we are really doing device-wide + -- synchronisation without atomics (naughty!). + let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op + bind_acc_params = + everythingVolatile $ + sComment "bind lhs" $ + forM_ (zip acc_params arrs) $ \(acc_p, arr) -> + copyDWIMFix (paramName acc_p) [] (Var arr) bucket + + let op_body = + sComment "execute operation" $ + compileBody' acc_params $ lambdaBody op + + do_hist = + everythingVolatile $ + sComment "update global result" $ + zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params + + fence = case space of + Space "local" -> sOp $ Imp.MemFence Imp.FenceLocal + _ -> sOp $ Imp.MemFence Imp.FenceGlobal + + -- While-loop: Try to insert your value + sWhile (tvExp continue) $ do + try_acquire_lock + sWhen lock_acquired $ do + dLParams acc_params + bind_acc_params + op_body + do_hist + fence + release_lock + break_loop + fence + where + writeArray bucket arr val = copyDWIMFix arr bucket val [] + +atomicUpdateCAS :: + Space -> + PrimType -> + VName -> + VName -> + [Imp.TExp Int64] -> + VName -> + InKernelGen () -> + InKernelGen () +atomicUpdateCAS space t arr old bucket x do_op = do + -- Code generation target: + -- + -- old = d_his[idx]; + -- do { + -- assumed = old; + -- x = do_op(assumed, y); + -- old = atomicCAS(&d_his[idx], assumed, tmp); + -- } while(assumed != old); + assumed <- tvVar <$> dPrim "assumed" t + run_loop <- dPrimV "run_loop" true + + -- XXX: CUDA may generate really bad code if this is not a volatile + -- read. Unclear why. The later reads are volatile, so maybe + -- that's it. + everythingVolatile $ copyDWIMFix old [] (Var arr) bucket + + (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket + + -- While-loop: Try to insert your value + let (toBits, fromBits) = + case t of + FloatType Float32 -> + ( \v -> Imp.FunExp "to_bits32" [v] int32, + \v -> Imp.FunExp "from_bits32" [v] t + ) + FloatType Float64 -> + ( \v -> Imp.FunExp "to_bits64" [v] int64, + \v -> Imp.FunExp "from_bits64" [v] t + ) + _ -> (id, id) + + int + | primBitSize t == 32 = int32 + | otherwise = int64 + + sWhile (tvExp run_loop) $ do + assumed <~~ Imp.var old t + x <~~ Imp.var assumed t + do_op + old_bits_v <- newVName "old_bits" + dPrim_ old_bits_v int + let old_bits = Imp.var old_bits_v int + sOp $ + Imp.Atomic space $ + Imp.AtomicCmpXchg + int + old_bits_v + arr' + bucket_offset + (toBits (Imp.var assumed t)) + (toBits (Imp.var x t)) + old <~~ fromBits old_bits + let won = CmpOpExp (CmpEq int) (toBits (Imp.var assumed t)) old_bits + sWhen (isBool won) (run_loop <-- false) + +computeKernelUses :: + FreeIn a => + a -> + [VName] -> + CallKernelGen [Imp.KernelUse] +computeKernelUses kernel_body bound_in_kernel = do + let actually_free = freeIn kernel_body `namesSubtract` namesFromList bound_in_kernel + -- Compute the variables that we need to pass to the kernel. + nubOrd <$> readsFromSet actually_free + +readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] +readsFromSet free = + fmap catMaybes $ + forM (namesToList free) $ \var -> do + t <- lookupType var + vtable <- getVTable + case t of + Array {} -> return Nothing + Acc {} -> return Nothing + Mem (Space "local") -> return Nothing + Mem {} -> return $ Just $ Imp.MemoryUse var + Prim bt -> + isConstExp vtable (Imp.var var bt) >>= \case + Just ce -> return $ Just $ Imp.ConstUse var ce + Nothing -> return $ Just $ Imp.ScalarUse var bt + +isConstExp :: + VTable GPUMem -> + Imp.Exp -> + ImpM rep r op (Maybe Imp.KernelConstExp) +isConstExp vtable size = do + fname <- askFunction + let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name + onLeaf Imp.Index {} _ = Nothing + lookupConstExp name = + constExp =<< hasExp =<< M.lookup name vtable + constExp (Op (Inner (SizeOp (GetSize key _)))) = + Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32 + constExp e = primExpFromExp lookupConstExp e + return $ replaceInPrimExpM onLeaf size + where + hasExp (ArrayVar e _) = e + hasExp (AccVar e _) = e + hasExp (ScalarVar e _) = e + hasExp (MemVar e _) = e + +computeThreadChunkSize :: + SplitOrdering -> + Imp.TExp Int64 -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + TV Int64 -> + ImpM rep r op () +computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = + chunk_var + <-- sMin64 + (Imp.unCount elements_per_thread) + ((Imp.unCount num_elements - thread_index) `divUp` toInt64Exp stride) +computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do + starting_point <- + dPrimV "starting_point" $ + thread_index * Imp.unCount elements_per_thread + remaining_elements <- + dPrimV "remaining_elements" $ + Imp.unCount num_elements - tvExp starting_point + + let no_remaining_elements = tvExp remaining_elements .<=. 0 + beyond_bounds = Imp.unCount num_elements .<=. tvExp starting_point + + sIf + (no_remaining_elements .||. beyond_bounds) + (chunk_var <-- 0) + ( sIf + is_last_thread + (chunk_var <-- Imp.unCount last_thread_elements) + (chunk_var <-- Imp.unCount elements_per_thread) + ) + where + last_thread_elements = + num_elements - Imp.elements thread_index * elements_per_thread + is_last_thread = + Imp.unCount num_elements + .<. (thread_index + 1) * Imp.unCount elements_per_thread + +kernelInitialisationSimple :: + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + CallKernelGen (KernelConstants, InKernelGen ()) +kernelInitialisationSimple (Count num_groups) (Count group_size) = do + global_tid <- newVName "global_tid" + local_tid <- newVName "local_tid" + group_id <- newVName "group_tid" + wave_size <- newVName "wave_size" + inner_group_size <- newVName "group_size" + let constants = + KernelConstants + (Imp.vi32 global_tid) + (Imp.vi32 local_tid) + (Imp.vi32 group_id) + global_tid + local_tid + group_id + num_groups + group_size + (sExt32 (group_size * num_groups)) + (Imp.vi32 wave_size) + true + mempty + + let set_constants = do + dPrim_ global_tid int32 + dPrim_ local_tid int32 + dPrim_ inner_group_size int64 + dPrim_ wave_size int32 + dPrim_ group_id int32 + + sOp (Imp.GetGlobalId global_tid 0) + sOp (Imp.GetLocalId local_tid 0) + sOp (Imp.GetLocalSize inner_group_size 0) + sOp (Imp.GetLockstepWidth wave_size) + sOp (Imp.GetGroupId group_id 0) + + return (constants, set_constants) + +isActive :: [(VName, SubExp)] -> Imp.TExp Bool +isActive limit = case actives of + [] -> true + x : xs -> foldl (.&&.) x xs + where + (is, ws) = unzip limit + actives = zipWith active is $ map toInt64Exp ws + active i = (Imp.vi64 i .<.) + +-- | Change every memory block to be in the global address space, +-- except those who are in the local memory space. This only affects +-- generated code - we still need to make sure that the memory is +-- actually present on the device (and dared as variables in the +-- kernel). +makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a +makeAllMemoryGlobal = + localDefaultSpace (Imp.Space "global") . localVTable (M.map globalMemory) + where + globalMemory (MemVar _ entry) + | entryMemSpace entry /= Space "local" = + MemVar Nothing entry {entryMemSpace = Imp.Space "global"} + globalMemory entry = + entry + +groupReduce :: + Imp.TExp Int32 -> + Lambda GPUMem -> + [VName] -> + InKernelGen () +groupReduce w lam arrs = do + offset <- dPrim "offset" int32 + groupReduceWithOffset offset w lam arrs + +groupReduceWithOffset :: + TV Int32 -> + Imp.TExp Int32 -> + Lambda GPUMem -> + [VName] -> + InKernelGen () +groupReduceWithOffset offset w lam arrs = do + constants <- kernelConstants <$> askEnv + + let local_tid = kernelLocalThreadId constants + global_tid = kernelGlobalThreadId constants + + barrier + | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal + | otherwise = sOp $ Imp.Barrier Imp.FenceGlobal + + readReduceArgument param arr + | Prim _ <- paramType param = do + let i = local_tid + tvExp offset + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] + | otherwise = do + let i = global_tid + tvExp offset + copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] + + writeReduceOpResult param arr + | Prim _ <- paramType param = + copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) [] + | otherwise = + return () + + let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam + + skip_waves <- dPrimV "skip_waves" (1 :: Imp.TExp Int32) + dLParams $ lambdaParams lam + + offset <-- (0 :: Imp.TExp Int32) + + comment "participating threads read initial accumulator" $ + sWhen (local_tid .<. w) $ + zipWithM_ readReduceArgument reduce_acc_params arrs + + let do_reduce = do + comment "read array element" $ + zipWithM_ readReduceArgument reduce_arr_params arrs + comment "apply reduction operation" $ + compileBody' reduce_acc_params $ lambdaBody lam + comment "write result of operation" $ + zipWithM_ writeReduceOpResult reduce_acc_params arrs + in_wave_reduce = everythingVolatile do_reduce + + wave_size = kernelWaveSize constants + group_size = kernelGroupSize constants + wave_id = local_tid `quot` wave_size + in_wave_id = local_tid - wave_id * wave_size + num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size + arg_in_bounds = local_tid + tvExp offset .<. w + + doing_in_wave_reductions = + tvExp offset .<. wave_size + apply_in_in_wave_iteration = + (in_wave_id .&. (2 * tvExp offset - 1)) .==. 0 + in_wave_reductions = do + offset <-- (1 :: Imp.TExp Int32) + sWhile doing_in_wave_reductions $ do + sWhen + (arg_in_bounds .&&. apply_in_in_wave_iteration) + in_wave_reduce + offset <-- tvExp offset * 2 + + doing_cross_wave_reductions = + tvExp skip_waves .<. num_waves + is_first_thread_in_wave = + in_wave_id .==. 0 + wave_not_skipped = + (wave_id .&. (2 * tvExp skip_waves - 1)) .==. 0 + apply_in_cross_wave_iteration = + arg_in_bounds .&&. is_first_thread_in_wave .&&. wave_not_skipped + cross_wave_reductions = + sWhile doing_cross_wave_reductions $ do + barrier + offset <-- tvExp skip_waves * wave_size + sWhen + apply_in_cross_wave_iteration + do_reduce + skip_waves <-- tvExp skip_waves * 2 + + in_wave_reductions + cross_wave_reductions + +groupScan :: + Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Lambda GPUMem -> + [VName] -> + InKernelGen () +groupScan seg_flag arrs_full_size w lam arrs = do + constants <- kernelConstants <$> askEnv + renamed_lam <- renameLambda lam + + let ltid32 = kernelLocalThreadId constants + ltid = sExt64 ltid32 + (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam + + dLParams (lambdaParams lam ++ lambdaParams renamed_lam) + + ltid_in_bounds <- dPrimVE "ltid_in_bounds" $ ltid .<. w + + -- The scan works by splitting the group into blocks, which are + -- scanned separately. Typically, these blocks are smaller than + -- the lockstep width, which enables barrier-free execution inside + -- them. + -- + -- We hardcode the block size here. The only requirement is that + -- it should not be less than the square root of the group size. + -- With 32, we will work on groups of size 1024 or smaller, which + -- fits every device Troels has seen. Still, it would be nicer if + -- it were a runtime parameter. Some day. + let block_size = 32 + simd_width = kernelWaveSize constants + block_id = ltid32 `quot` block_size + in_block_id = ltid32 - block_id * block_size + doInBlockScan seg_flag' active = + inBlockScan + constants + seg_flag' + arrs_full_size + simd_width + block_size + active + arrs + barrier + array_scan = not $ all primType $ lambdaReturnType lam + barrier + | array_scan = + sOp $ Imp.Barrier Imp.FenceGlobal + | otherwise = + sOp $ Imp.Barrier Imp.FenceLocal + + group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants + + writeBlockResult p arr + | primType $ paramType p = + copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) [] + | otherwise = + copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) [] + + readPrevBlockResult p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1] + + doInBlockScan seg_flag ltid_in_bounds lam + barrier + + let is_first_block = block_id .==. 0 + when array_scan $ do + sComment "save correct values for first block" $ + sWhen is_first_block $ + forM_ (zip x_params arrs) $ \(x, arr) -> + unless (primType $ paramType x) $ + copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) [] + + barrier + + let last_in_block = in_block_id .==. block_size - 1 + sComment "last thread of block 'i' writes its result to offset 'i'" $ + sWhen (last_in_block .&&. ltid_in_bounds) $ + everythingVolatile $ + zipWithM_ writeBlockResult x_params arrs + + barrier + + let first_block_seg_flag = do + flag_true <- seg_flag + Just $ \from to -> + flag_true (from * block_size + block_size -1) (to * block_size + block_size -1) + comment + "scan the first block, after which offset 'i' contains carry-in for block 'i+1'" + $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam + + barrier + + when array_scan $ do + sComment "move correct values for first block back a block" $ + sWhen is_first_block $ + forM_ (zip x_params arrs) $ \(x, arr) -> + unless (primType $ paramType x) $ + copyDWIM + arr + [DimFix $ arrs_full_size + group_offset + ltid] + (Var arr) + [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] + + barrier + + let read_carry_in = do + forM_ (zip x_params y_params) $ \(x, y) -> + copyDWIM (paramName y) [] (Var (paramName x)) [] + zipWithM_ readPrevBlockResult x_params arrs + + y_to_x = forM_ (zip x_params y_params) $ \(x, y) -> + when (primType (paramType x)) $ + copyDWIM (paramName x) [] (Var (paramName y)) [] + + op_to_x + | Nothing <- seg_flag = + compileBody' x_params $ lambdaBody lam + | Just flag_true <- seg_flag = do + inactive <- + dPrimVE "inactive" $ flag_true (block_id * block_size -1) ltid32 + sWhen inactive y_to_x + when array_scan barrier + sUnless inactive $ compileBody' x_params $ lambdaBody lam + + write_final_result = + forM_ (zip x_params arrs) $ \(p, arr) -> + when (primType $ paramType p) $ + copyDWIM arr [DimFix ltid] (Var $ paramName p) [] + + sComment "carry-in for every block except the first" $ + sUnless (is_first_block .||. bNot ltid_in_bounds) $ do + sComment "read operands" read_carry_in + sComment "perform operation" op_to_x + sComment "write final result" write_final_result + + barrier + + sComment "restore correct values for first block" $ + sWhen is_first_block $ + forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) -> + if primType (paramType y) + then copyDWIM arr [DimFix ltid] (Var $ paramName y) [] + else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid] + + barrier + +inBlockScan :: + KernelConstants -> + Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> + Imp.TExp Int64 -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + Imp.TExp Bool -> + [VName] -> + InKernelGen () -> + Lambda GPUMem -> + InKernelGen () +inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do + skip_threads <- dPrim "skip_threads" int32 + let in_block_thread_active = + tvExp skip_threads .<=. in_block_id + actual_params = lambdaParams scan_lam + (x_params, y_params) = + splitAt (length actual_params `div` 2) actual_params + y_to_x = + forM_ (zip x_params y_params) $ \(x, y) -> + when (primType (paramType x)) $ + copyDWIM (paramName x) [] (Var (paramName y)) [] + + -- Set initial y values + sComment "read input for in-block scan" $ + sWhen active $ do + zipWithM_ readInitial y_params arrs + -- Since the final result is expected to be in x_params, we may + -- need to copy it there for the first thread in the block. + sWhen (in_block_id .==. 0) y_to_x + + when array_scan barrier + + let op_to_x + | Nothing <- seg_flag = + compileBody' x_params $ lambdaBody scan_lam + | Just flag_true <- seg_flag = do + inactive <- + dPrimVE "inactive" $ + flag_true (ltid32 - tvExp skip_threads) ltid32 + sWhen inactive y_to_x + when array_scan barrier + sUnless inactive $ compileBody' x_params $ lambdaBody scan_lam + + maybeBarrier = + sWhen + (lockstep_width .<=. tvExp skip_threads) + barrier + + sComment "in-block scan (hopefully no barriers needed)" $ do + skip_threads <-- 1 + sWhile (tvExp skip_threads .<. block_size) $ do + sWhen (in_block_thread_active .&&. active) $ do + sComment "read operands" $ + zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs + sComment "perform operation" op_to_x + + maybeBarrier + + sWhen (in_block_thread_active .&&. active) $ + sComment "write result" $ + sequence_ $ zipWith3 writeResult x_params y_params arrs + + maybeBarrier + + skip_threads <-- tvExp skip_threads * 2 + where + block_id = ltid32 `quot` block_size + in_block_id = ltid32 - block_id * block_size + ltid32 = kernelLocalThreadId constants + ltid = sExt64 ltid32 + gtid = sExt64 $ kernelGlobalThreadId constants + array_scan = not $ all primType $ lambdaReturnType scan_lam + + readInitial p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix ltid] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix gtid] + + readParam behind p arr + | primType $ paramType p = + copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind] + | otherwise = + copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size] + + writeResult x y arr + | primType $ paramType x = do + copyDWIM arr [DimFix ltid] (Var $ paramName x) [] + copyDWIM (paramName y) [] (Var $ paramName x) [] + | otherwise = + copyDWIM (paramName y) [] (Var $ paramName x) [] + +computeMapKernelGroups :: Imp.TExp Int64 -> CallKernelGen (Imp.TExp Int64, Imp.TExp Int64) +computeMapKernelGroups kernel_size = do + group_size <- dPrim "group_size" int64 + fname <- askFunction + let group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty $ tvVar group_size + sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup + num_groups <- dPrimV "num_groups" $ kernel_size `divUp` tvExp group_size + return (tvExp num_groups, tvExp group_size) + +simpleKernelConstants :: + Imp.TExp Int64 -> + String -> + CallKernelGen (KernelConstants, InKernelGen ()) +simpleKernelConstants kernel_size desc = do + thread_gtid <- newVName $ desc ++ "_gtid" + thread_ltid <- newVName $ desc ++ "_ltid" + group_id <- newVName $ desc ++ "_gid" + (num_groups, group_size) <- computeMapKernelGroups kernel_size + let set_constants = do + dPrim_ thread_gtid int32 + dPrim_ thread_ltid int32 + dPrim_ group_id int32 + sOp (Imp.GetGlobalId thread_gtid 0) + sOp (Imp.GetLocalId thread_ltid 0) + sOp (Imp.GetGroupId group_id 0) + + return + ( KernelConstants + (Imp.vi32 thread_gtid) + (Imp.vi32 thread_ltid) + (Imp.vi32 group_id) + thread_gtid + thread_ltid + group_id + num_groups + group_size + (sExt32 (group_size * num_groups)) + 0 + (Imp.vi64 thread_gtid .<. kernel_size) + mempty, + set_constants + ) + +-- | For many kernels, we may not have enough physical groups to cover +-- the logical iteration space. Some groups thus have to perform +-- double duty; we put an outer loop to accomplish this. The +-- advantage over just launching a bazillion threads is that the cost +-- of memory expansion should be proportional to the number of +-- *physical* threads (hardware parallelism), not the amount of +-- application parallelism. +virtualiseGroups :: + SegVirt -> + Imp.TExp Int32 -> + (Imp.TExp Int32 -> InKernelGen ()) -> + InKernelGen () +virtualiseGroups SegVirt required_groups m = do + constants <- kernelConstants <$> askEnv + phys_group_id <- dPrim "phys_group_id" int32 + sOp $ Imp.GetGroupId (tvVar phys_group_id) 0 + let iterations = + (required_groups - tvExp phys_group_id) + `divUp` sExt32 (kernelNumGroups constants) + + sFor "i" iterations $ \i -> do + m . tvExp + =<< dPrimV + "virt_group_id" + (tvExp phys_group_id + i * sExt32 (kernelNumGroups constants)) + -- Make sure the virtual group is actually done before we let + -- another virtual group have its way with it. + sOp $ Imp.Barrier Imp.FenceGlobal +virtualiseGroups _ _ m = do + gid <- kernelGroupIdVar . kernelConstants <$> askEnv + m $ Imp.vi32 gid + +sKernelThread :: + String -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + VName -> + InKernelGen () -> + CallKernelGen () +sKernelThread = sKernel threadOperations kernelGlobalThreadId + +sKernelGroup :: + String -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + VName -> + InKernelGen () -> + CallKernelGen () +sKernelGroup = sKernel groupOperations kernelGroupId + +sKernelFailureTolerant :: + Bool -> + Operations GPUMem KernelEnv Imp.KernelOp -> + KernelConstants -> + Name -> + InKernelGen () -> + CallKernelGen () +sKernelFailureTolerant tol ops constants name m = do + HostEnv atomics _ locks <- askEnv + body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m + uses <- computeKernelUses body mempty + emit $ + Imp.Op $ + Imp.CallKernel + Imp.Kernel + { Imp.kernelBody = body, + Imp.kernelUses = uses, + Imp.kernelNumGroups = [untyped $ kernelNumGroups constants], + Imp.kernelGroupSize = [untyped $ kernelGroupSize constants], + Imp.kernelName = name, + Imp.kernelFailureTolerant = tol + } + +sKernel :: + Operations GPUMem KernelEnv Imp.KernelOp -> + (KernelConstants -> Imp.TExp Int32) -> + String -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + VName -> + InKernelGen () -> + CallKernelGen () +sKernel ops flatf name num_groups group_size v f = do + (constants, set_constants) <- kernelInitialisationSimple num_groups group_size + name' <- nameForFun $ name ++ "_" ++ show (baseTag v) + sKernelFailureTolerant False ops constants name' $ do + set_constants + dPrimV_ v $ flatf constants + f + +copyInGroup :: CopyCompiler GPUMem KernelEnv Imp.KernelOp +copyInGroup pt destloc destslice srcloc srcslice = do + dest_space <- entryMemSpace <$> lookupMemory (memLocationName destloc) + src_space <- entryMemSpace <$> lookupMemory (memLocationName srcloc) + + case (dest_space, src_space) of + (ScalarSpace destds _, ScalarSpace srcds _) -> do + let destslice' = + replicate (length destslice - length destds) (DimFix 0) + ++ takeLast (length destds) destslice + srcslice' = + replicate (length srcslice - length srcds) (DimFix 0) + ++ takeLast (length srcds) srcslice + copyElementWise pt destloc destslice' srcloc srcslice' + _ -> do + groupCoverSpace (sliceDims destslice) $ \is -> + copyElementWise + pt + destloc + (map DimFix $ fixSlice destslice is) + srcloc + (map DimFix $ fixSlice srcslice is) + sOp $ Imp.Barrier Imp.FenceLocal + +threadOperations, groupOperations :: Operations GPUMem KernelEnv Imp.KernelOp +threadOperations = + (defaultOperations compileThreadOp) + { opsCopyCompiler = copyElementWise, + opsExpCompiler = compileThreadExp, + opsStmsCompiler = \_ -> defCompileStms mempty, + opsAllocCompilers = + M.fromList [(Space "local", allocLocal)] + } +groupOperations = + (defaultOperations compileGroupOp) + { opsCopyCompiler = copyInGroup, + opsExpCompiler = compileGroupExp, + opsStmsCompiler = \_ -> defCompileStms mempty, + opsAllocCompilers = + M.fromList [(Space "local", allocLocal)] + } + +-- | Perform a Replicate with a kernel. +sReplicateKernel :: VName -> SubExp -> CallKernelGen () +sReplicateKernel arr se = do + t <- subExpType se + ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr + + let dims = map toInt64Exp $ ds ++ arrayDims t + (constants, set_constants) <- + simpleKernelConstants (product $ map sExt64 dims) "replicate" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + is' = unflattenIndex dims $ sExt64 $ kernelGlobalThreadId constants + + sKernelFailureTolerant True threadOperations constants name $ do + set_constants + sWhen (kernelThreadActive constants) $ + copyDWIMFix arr is' se $ drop (length ds) is' + +replicateName :: PrimType -> String +replicateName bt = "replicate_" ++ pretty bt + +replicateForType :: PrimType -> CallKernelGen Name +replicateForType bt = do + let fname = nameFromString $ "builtin#" <> replicateName bt + + exists <- hasFunction fname + unless exists $ do + mem <- newVName "mem" + num_elems <- newVName "num_elems" + val <- newVName "val" + + let params = + [ Imp.MemParam mem (Space "device"), + Imp.ScalarParam num_elems int32, + Imp.ScalarParam val bt + ] + shape = Shape [Var num_elems] + function fname [] params $ do + arr <- + sArray "arr" bt shape $ + ArrayIn mem $ + IxFun.iota $ + map pe64 $ shapeDims shape + sReplicateKernel arr $ Var val + + return fname + +replicateIsFill :: VName -> SubExp -> CallKernelGen (Maybe (CallKernelGen ())) +replicateIsFill arr v = do + ArrayEntry (MemLocation arr_mem arr_shape arr_ixfun) _ <- lookupArray arr + v_t <- subExpType v + case v_t of + Prim v_t' + | IxFun.isLinear arr_ixfun -> return $ + Just $ do + fname <- replicateForType v_t' + emit $ + Imp.Call + [] + fname + [ Imp.MemArg arr_mem, + Imp.ExpArg $ untyped $ product $ map toInt64Exp arr_shape, + Imp.ExpArg $ toExp' v_t' v + ] + _ -> return Nothing + +-- | Perform a Replicate with a kernel. +sReplicate :: VName -> SubExp -> CallKernelGen () +sReplicate arr se = do + -- If the replicate is of a particularly common and simple form + -- (morally a memset()/fill), then we use a common function. + is_fill <- replicateIsFill arr se + + case is_fill of + Just m -> m + Nothing -> sReplicateKernel arr se + +-- | Perform an Iota with a kernel. +sIotaKernel :: + VName -> + Imp.TExp Int64 -> + Imp.Exp -> + Imp.Exp -> + IntType -> + CallKernelGen () +sIotaKernel arr n x s et = do + destloc <- entryArrayLocation <$> lookupArray arr + (constants, set_constants) <- simpleKernelConstants n "iota" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "iota_" ++ pretty et ++ "_" + ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ do + set_constants + let gtid = sExt64 $ kernelGlobalThreadId constants + sWhen (kernelThreadActive constants) $ do + (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] + + emit $ + Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ + BinOpExp + (Add et OverflowWrap) + (BinOpExp (Mul et OverflowWrap) (Imp.sExt et $ untyped gtid) s) + x + +iotaName :: IntType -> String +iotaName bt = "iota_" ++ pretty bt + +iotaForType :: IntType -> CallKernelGen Name +iotaForType bt = do + let fname = nameFromString $ "builtin#" <> iotaName bt + + exists <- hasFunction fname + unless exists $ do + mem <- newVName "mem" + n <- newVName "n" + x <- newVName "x" + s <- newVName "s" + + let params = + [ Imp.MemParam mem (Space "device"), + Imp.ScalarParam n int32, + Imp.ScalarParam x $ IntType bt, + Imp.ScalarParam s $ IntType bt + ] + shape = Shape [Var n] + n' = Imp.vi64 n + x' = Imp.var x $ IntType bt + s' = Imp.var s $ IntType bt + + function fname [] params $ do + arr <- + sArray "arr" (IntType bt) shape $ + ArrayIn mem $ + IxFun.iota $ + map pe64 $ shapeDims shape + sIotaKernel arr (sExt64 n') x' s' bt + + return fname + +-- | Perform an Iota with a kernel. +sIota :: + VName -> + Imp.TExp Int64 -> + Imp.Exp -> + Imp.Exp -> + IntType -> + CallKernelGen () +sIota arr n x s et = do + ArrayEntry (MemLocation arr_mem _ arr_ixfun) _ <- lookupArray arr + if IxFun.isLinear arr_ixfun + then do + fname <- iotaForType et + emit $ + Imp.Call + [] + fname + [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s] + else sIotaKernel arr n x s et + +sCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp +sCopy + bt + destloc@(MemLocation destmem _ _) + destslice + srcloc@(MemLocation srcmem _ _) + srcslice = + do + -- Note that the shape of the destination and the source are + -- necessarily the same. + let shape = sliceDims srcslice + kernel_size = product shape + + (constants, set_constants) <- simpleKernelConstants kernel_size "copy" + + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ do + set_constants + + let gtid = sExt64 $ kernelGlobalThreadId constants + dest_is = unflattenIndex shape gtid + src_is = dest_is + + (_, destspace, destidx) <- + fullyIndexArray' destloc $ fixSlice destslice dest_is + (_, srcspace, srcidx) <- + fullyIndexArray' srcloc $ fixSlice srcslice src_is + + sWhen (gtid .<. kernel_size) $ + emit $ + Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ + Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile + +compileGroupResult :: + SegSpace -> + PatElem GPUMem -> + KernelResult -> + InKernelGen () +compileGroupResult _ pe (TileReturns [(w, per_group_elems)] what) = do + n <- toInt64Exp . arraySize 0 <$> lookupType what + + constants <- kernelConstants <$> askEnv + let ltid = sExt64 $ kernelLocalThreadId constants + offset = + toInt64Exp per_group_elems + * sExt64 (kernelGroupId constants) + + -- Avoid loop for the common case where each thread is statically + -- known to write at most one element. + localOps threadOperations $ + if toInt64Exp per_group_elems == kernelGroupSize constants + then + sWhen (ltid + offset .<. toInt64Exp w) $ + copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid] + else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do + j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid + sWhen (j + offset .<. toInt64Exp w) $ + copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] +compileGroupResult space pe (TileReturns dims what) = do + let gids = map fst $ unSegSpace space + out_tile_sizes = map (toInt64Exp . snd) dims + group_is = zipWith (*) (map Imp.vi64 gids) out_tile_sizes + local_is <- localThreadIDs $ map snd dims + is_for_thread <- + mapM (dPrimV "thread_out_index") $ + zipWith (+) group_is local_is + + localOps threadOperations $ + sWhen (isActive $ zip (map tvVar is_for_thread) $ map fst dims) $ + copyDWIMFix (patElemName pe) (map tvExp is_for_thread) (Var what) local_is +compileGroupResult space pe (RegTileReturns dims_n_tiles what) = do + constants <- kernelConstants <$> askEnv + + let gids = map fst $ unSegSpace space + (dims, group_tiles, reg_tiles) = unzip3 dims_n_tiles + group_tiles' = map toInt64Exp group_tiles + reg_tiles' = map toInt64Exp reg_tiles + + -- Which group tile is this group responsible for? + let group_tile_is = map Imp.vi64 gids + + -- Within the group tile, which register tile is this thread + -- responsible for? + reg_tile_is <- + mapM (dPrimVE "reg_tile_i") $ + unflattenIndex group_tiles' $ sExt64 $ kernelLocalThreadId constants + + -- Compute output array slice for the register tile belonging to + -- this thread. + let regTileSliceDim (group_tile, group_tile_i) (reg_tile, reg_tile_i) = do + tile_dim_start <- + dPrimVE "tile_dim_start" $ + reg_tile * (group_tile * group_tile_i + reg_tile_i) + return $ DimSlice tile_dim_start reg_tile 1 + reg_tile_slices <- + zipWithM + regTileSliceDim + (zip group_tiles' group_tile_is) + (zip reg_tiles' reg_tile_is) + + localOps threadOperations $ + sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do + let dest_is = fixSlice reg_tile_slices is_in_reg_tile + src_is = reg_tile_is ++ is_in_reg_tile + sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map toInt64Exp dims) $ + copyDWIMFix (patElemName pe) dest_is (Var what) src_is +compileGroupResult space pe (Returns _ what) = do + constants <- kernelConstants <$> askEnv + in_local_memory <- arrayInLocalMemory what + let gids = map (Imp.vi64 . fst) $ unSegSpace space + + if not in_local_memory + then + localOps threadOperations $ + sWhen (kernelLocalThreadId constants .==. 0) $ + copyDWIMFix (patElemName pe) gids what [] + else -- If the result of the group is an array in local memory, we + -- store it by collective copying among all the threads of the + -- group. TODO: also do this if the array is in global memory + -- (but this is a bit more tricky, synchronisation-wise). + copyDWIMFix (patElemName pe) gids what [] +compileGroupResult _ _ WriteReturns {} = + compilerLimitationS "compileGroupResult: WriteReturns not handled yet." +compileGroupResult _ _ ConcatReturns {} = + compilerLimitationS "compileGroupResult: ConcatReturns not handled yet." + +compileThreadResult :: + SegSpace -> + PatElem GPUMem -> + KernelResult -> + InKernelGen () +compileThreadResult _ _ RegTileReturns {} = + compilerLimitationS "compileThreadResult: RegTileReturns not yet handled." +compileThreadResult space pe (Returns _ what) = do + let is = map (Imp.vi64 . fst) $ unSegSpace space + copyDWIMFix (patElemName pe) is what [] +compileThreadResult _ pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do + constants <- kernelConstants <$> askEnv + let offset = + toInt64Exp per_thread_elems + * sExt64 (kernelGlobalThreadId constants) + n <- toInt64Exp . arraySize 0 <$> lookupType what + copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) [] +compileThreadResult _ pe (ConcatReturns (SplitStrided stride) _ _ what) = do + offset <- sExt64 . kernelGlobalThreadId . kernelConstants <$> askEnv + n <- toInt64Exp . arraySize 0 <$> lookupType what + copyDWIM (patElemName pe) [DimSlice offset n $ toInt64Exp stride] (Var what) [] +compileThreadResult _ pe (WriteReturns (Shape rws) _arr dests) = do + constants <- kernelConstants <$> askEnv + let rws' = map toInt64Exp rws + forM_ dests $ \(slice, e) -> do + let slice' = map (fmap toInt64Exp) slice + condInBounds (DimFix i) rw = + 0 .<=. i .&&. i .<. rw + condInBounds (DimSlice i n s) rw = + 0 .<=. i .&&. i + n * s .<. rw + write = + foldl (.&&.) (kernelThreadActive constants) $ + zipWith condInBounds slice' rws' + sWhen write $ copyDWIM (patElemName pe) slice' e [] +compileThreadResult _ _ TileReturns {} = + compilerBugS "compileThreadResult: TileReturns unhandled." + +arrayInLocalMemory :: SubExp -> InKernelGen Bool +arrayInLocalMemory (Var name) = do + res <- lookupVar name + case res of + ArrayVar _ entry -> + (Space "local" ==) . entryMemSpace + <$> lookupMemory (memLocationName (entryArrayLocation entry)) + _ -> return False +arrayInLocalMemory Constant {} = return False diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs @@ -0,0 +1,1151 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Our compilation strategy for 'SegHist' is based around avoiding +-- bin conflicts. We do this by splitting the input into chunks, and +-- for each chunk computing a single subhistogram. Then we combine +-- the subhistograms using an ordinary segmented reduction ('SegRed'). +-- +-- There are some branches around to efficiently handle the case where +-- we use only a single subhistogram (because it's large), so that we +-- respect the asymptotics, and do not copy the destination array. +-- +-- We also use a heuristic strategy for computing subhistograms in +-- local memory when possible. Given: +-- +-- H: total size of histograms in bytes, including any lock arrays. +-- +-- G: group size +-- +-- T: number of bytes of local memory each thread can be given without +-- impacting occupancy (determined experimentally, e.g. 32). +-- +-- LMAX: maximum amount of local memory per workgroup (hard limit). +-- +-- We wish to compute: +-- +-- COOP: cooperation level (number of threads per subhistogram) +-- +-- LH: number of local memory subhistograms +-- +-- We do this as: +-- +-- COOP = ceil(H / T) +-- LH = ceil((G*T)/H) +-- if COOP <= G && H <= LMAX then +-- use local memory +-- else +-- use global memory +module Futhark.CodeGen.ImpGen.GPU.SegHist (compileSegHist) where + +import Control.Monad.Except +import Data.List (foldl', genericLength, zip4, zip6) +import Data.Maybe +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.CodeGen.ImpGen.GPU.SegRed (compileSegRed') +import Futhark.Construct (fullSliceNum) +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Pass.ExplicitAllocations () +import Futhark.Util (chunks, mapAccumLM, maxinum, splitFromEnd, takeLast) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +data SubhistosInfo = SubhistosInfo + { subhistosArray :: VName, + subhistosAlloc :: CallKernelGen () + } + +data SegHistSlug = SegHistSlug + { slugOp :: HistOp GPUMem, + slugNumSubhistos :: TV Int64, + slugSubhistos :: [SubhistosInfo], + slugAtomicUpdate :: AtomicUpdate GPUMem KernelEnv + } + +histoSpaceUsage :: + HistOp GPUMem -> + Imp.Count Imp.Bytes (Imp.TExp Int64) +histoSpaceUsage op = + sum $ + map + ( typeSize + . (`arrayOfRow` histWidth op) + . (`arrayOfShape` histShape op) + ) + $ lambdaReturnType $ histOp op + +-- | Figure out how much memory is needed per histogram, both +-- segmented and unsegmented,, and compute some other auxiliary +-- information. +computeHistoUsage :: + SegSpace -> + HistOp GPUMem -> + CallKernelGen + ( Imp.Count Imp.Bytes (Imp.TExp Int64), + Imp.Count Imp.Bytes (Imp.TExp Int64), + SegHistSlug + ) +computeHistoUsage space op = do + let segment_dims = init $ unSegSpace space + num_segments = length segment_dims + + -- Create names for the intermediate array memory blocks, + -- memory block sizes, arrays, and number of subhistograms. + num_subhistos <- dPrim "num_subhistos" int32 + subhisto_infos <- forM (zip (histDest op) (histNeutral op)) $ \(dest, ne) -> do + dest_t <- lookupType dest + dest_mem <- entryArrayLocation <$> lookupArray dest + + subhistos_mem <- + sDeclareMem (baseString dest ++ "_subhistos_mem") (Space "device") + + let subhistos_shape = + Shape (map snd segment_dims ++ [tvSize num_subhistos]) + <> stripDims num_segments (arrayShape dest_t) + subhistos_membind = + ArrayIn subhistos_mem $ + IxFun.iota $ + map pe64 $ shapeDims subhistos_shape + subhistos <- + sArray + (baseString dest ++ "_subhistos") + (elemType dest_t) + subhistos_shape + subhistos_membind + + return $ + SubhistosInfo subhistos $ do + let unitHistoCase = + emit $ + Imp.SetMem subhistos_mem (memLocationName dest_mem) $ + Space "device" + + multiHistoCase = do + let num_elems = + foldl' (*) (sExt64 $ tvExp num_subhistos) $ + map toInt64Exp $ arrayDims dest_t + + let subhistos_mem_size = + Imp.bytes $ + Imp.unCount (Imp.elements num_elems `Imp.withElemType` elemType dest_t) + + sAlloc_ subhistos_mem subhistos_mem_size $ Space "device" + sReplicate subhistos ne + subhistos_t <- lookupType subhistos + let slice = + fullSliceNum (map toInt64Exp $ arrayDims subhistos_t) $ + map (unitSlice 0 . toInt64Exp . snd) segment_dims + ++ [DimFix 0] + sUpdate subhistos slice $ Var dest + + sIf (tvExp num_subhistos .==. 1) unitHistoCase multiHistoCase + + let h = histoSpaceUsage op + segmented_h = h * product (map (Imp.bytes . toInt64Exp) $ init $ segSpaceDims space) + + atomics <- hostAtomics <$> askEnv + + return + ( h, + segmented_h, + SegHistSlug op num_subhistos subhisto_infos $ + atomicUpdateLocking atomics $ histOp op + ) + +prepareAtomicUpdateGlobal :: + Maybe Locking -> + [VName] -> + SegHistSlug -> + CallKernelGen + ( Maybe Locking, + [Imp.TExp Int64] -> InKernelGen () + ) +prepareAtomicUpdateGlobal l dests slug = + -- We need a separate lock array if the operators are not all of a + -- particularly simple form that permits pure atomic operations. + case (l, slugAtomicUpdate slug) of + (_, AtomicPrim f) -> return (l, f (Space "global") dests) + (_, AtomicCAS f) -> return (l, f (Space "global") dests) + (Just l', AtomicLocking f) -> return (l, f l' (Space "global") dests) + (Nothing, AtomicLocking f) -> do + -- The number of locks used here is too low, but since we are + -- currently forced to inline a huge list, I'm keeping it down + -- for now. Some quick experiments suggested that it has little + -- impact anyway (maybe the locking case is just too slow). + -- + -- A fun solution would also be to use a simple hashing + -- algorithm to ensure good distribution of locks. + let num_locks = 100151 + dims = + map toInt64Exp $ + shapeDims (histShape (slugOp slug)) + ++ [ tvSize (slugNumSubhistos slug), + histWidth (slugOp slug) + ] + locks <- + sStaticArray "hist_locks" (Space "device") int32 $ + Imp.ArrayZeros num_locks + let l' = Locking locks 0 1 0 (pure . (`rem` fromIntegral num_locks) . flattenIndex dims) + return (Just l', f l' (Space "global") dests) + +-- | Some kernel bodies are not safe (or efficient) to execute +-- multiple times. +data Passage = MustBeSinglePass | MayBeMultiPass deriving (Eq, Ord) + +bodyPassage :: KernelBody GPUMem -> Passage +bodyPassage kbody + | mempty == consumedInKernelBody (aliasAnalyseKernelBody mempty kbody) = + MayBeMultiPass + | otherwise = + MustBeSinglePass + +prepareIntermediateArraysGlobal :: + Passage -> + Imp.TExp Int32 -> + Imp.TExp Int64 -> + [SegHistSlug] -> + CallKernelGen + ( Imp.TExp Int32, + [[Imp.TExp Int64] -> InKernelGen ()] + ) +prepareIntermediateArraysGlobal passage hist_T hist_N slugs = do + -- The paper formulae assume there is only one histogram, but in our + -- implementation there can be multiple that have been horisontally + -- fused. We do a bit of trickery with summings and averages to + -- pretend there is really only one. For the case of a single + -- histogram, the actual calculations should be the same as in the + -- paper. + + -- The sum of all Hs. + hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth . slugOp) slugs + + hist_RF <- + dPrimVE "hist_RF" $ + sum (map (r64 . toInt64Exp . histRaceFactor . slugOp) slugs) + / genericLength slugs + + hist_el_size <- dPrimVE "hist_el_size" $ sum $ map slugElAvgSize slugs + + hist_C_max <- + dPrimVE "hist_C_max" $ + fMin64 (r64 hist_T) $ r64 hist_H / hist_k_ct_min + + hist_M_min <- + dPrimVE "hist_M_min" $ + sMax32 1 $ sExt32 $ t64 $ r64 hist_T / hist_C_max + + -- Querying L2 cache size is not reliable. Instead we provide a + -- tunable knob with a hopefully sane default. + let hist_L2_def = 4 * 1024 * 1024 + hist_L2 <- dPrim "L2_size" int32 + entry <- askFunction + -- Equivalent to F_L2*L2 in paper. + sOp $ + Imp.GetSize + (tvVar hist_L2) + (keyWithEntryPoint entry $ nameFromString (pretty (tvVar hist_L2))) + $ Imp.SizeBespoke (nameFromString "L2_for_histogram") hist_L2_def + + let hist_L2_ln_sz = 16 * 4 -- L2 cache line size approximation + hist_RACE_exp <- + dPrimVE "hist_RACE_exp" $ + fMax64 1 $ + (hist_k_RF * hist_RF) + / (hist_L2_ln_sz / r64 hist_el_size) + + hist_S <- dPrim "hist_S" int32 + + -- For sparse histograms (H exceeds N) we only want a single chunk. + sIf + (hist_N .<. hist_H) + (hist_S <-- (1 :: Imp.TExp Int32)) + $ hist_S + <-- case passage of + MayBeMultiPass -> + sExt32 $ + (sExt64 hist_M_min * hist_H * sExt64 hist_el_size) + `divUp` t64 (hist_F_L2 * r64 (tvExp hist_L2) * hist_RACE_exp) + MustBeSinglePass -> + 1 + + emit $ Imp.DebugPrint "Race expansion factor (RACE^exp)" $ Just $ untyped hist_RACE_exp + emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ untyped $ tvExp hist_S + + histograms <- + snd + <$> mapAccumLM + (onOp (tvExp hist_L2) hist_M_min (tvExp hist_S) hist_RACE_exp) + Nothing + slugs + + return (tvExp hist_S, histograms) + where + hist_k_ct_min = 2 -- Chosen experimentally + hist_k_RF = 0.75 -- Chosen experimentally + hist_F_L2 = 0.4 -- Chosen experimentally + r64 = isF64 . ConvOpExp (SIToFP Int32 Float64) . untyped + t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped + + -- "Average element size" as computed by a formula that also takes + -- locking into account. + slugElAvgSize slug@(SegHistSlug op _ _ do_op) = + case do_op of + AtomicLocking {} -> + slugElSize slug `quot` (1 + genericLength (lambdaReturnType (histOp op))) + _ -> + slugElSize slug `quot` genericLength (lambdaReturnType (histOp op)) + + -- "Average element size" as computed by a formula that also takes + -- locking into account. + slugElSize (SegHistSlug op _ _ do_op) = + case do_op of + AtomicLocking {} -> + sExt32 $ + unCount $ + sum $ + map (typeSize . (`arrayOfShape` histShape op)) $ + Prim int32 : lambdaReturnType (histOp op) + _ -> + sExt32 $ + unCount $ + sum $ + map (typeSize . (`arrayOfShape` histShape op)) $ + lambdaReturnType (histOp op) + + onOp hist_L2 hist_M_min hist_S hist_RACE_exp l slug = do + let SegHistSlug op num_subhistos subhisto_info do_op = slug + hist_H = toInt64Exp $ histWidth op + + hist_H_chk <- dPrimVE "hist_H_chk" $ hist_H `divUp` sExt64 hist_S + + emit $ Imp.DebugPrint "Chunk size (H_chk)" $ Just $ untyped hist_H_chk + + hist_k_max <- + dPrimVE "hist_k_max" $ + fMin64 + (hist_F_L2 * (r64 hist_L2 / r64 (slugElSize slug)) * hist_RACE_exp) + (r64 hist_N) + / r64 hist_T + + hist_u <- dPrimVE "hist_u" $ + case do_op of + AtomicPrim {} -> 2 + _ -> 1 + + hist_C <- + dPrimVE "hist_C" $ + fMin64 (r64 hist_T) $ r64 (hist_u * hist_H_chk) / hist_k_max + + -- Number of subhistograms per result histogram. + hist_M <- dPrimVE "hist_M" $ + case slugAtomicUpdate slug of + AtomicPrim {} -> 1 + _ -> sMax32 hist_M_min $ sExt32 $ t64 $ r64 hist_T / hist_C + + emit $ Imp.DebugPrint "Elements/thread in L2 cache (k_max)" $ Just $ untyped hist_k_max + emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ untyped hist_M + emit $ Imp.DebugPrint "Cooperation level (C)" $ Just $ untyped hist_C + + -- num_subhistos is the variable we use to communicate back. + num_subhistos <-- sExt64 hist_M + + -- Initialise sub-histograms. + -- + -- If hist_M is 1, then we just reuse the original + -- destination. The idea is to avoid a copy if we are writing a + -- small number of values into a very large prior histogram. + dests <- forM (zip (histDest op) subhisto_info) $ \(dest, info) -> do + dest_mem <- entryArrayLocation <$> lookupArray dest + + sub_mem <- + fmap memLocationName $ + entryArrayLocation + <$> lookupArray (subhistosArray info) + + let unitHistoCase = + emit $ + Imp.SetMem sub_mem (memLocationName dest_mem) $ + Space "device" + + multiHistoCase = subhistosAlloc info + + sIf (hist_M .==. 1) unitHistoCase multiHistoCase + + return $ subhistosArray info + + (l', do_op') <- prepareAtomicUpdateGlobal l dests slug + + return (l', do_op') + +histKernelGlobalPass :: + [PatElem GPUMem] -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + SegSpace -> + [SegHistSlug] -> + KernelBody GPUMem -> + [[Imp.TExp Int64] -> InKernelGen ()] -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + CallKernelGen () +histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do + let (space_is, space_sizes) = unzip $ unSegSpace space + space_sizes_64 = map (sExt64 . toInt64Exp) space_sizes + total_w_64 = product space_sizes_64 + + hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> + dPrimVE "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S + + sKernelThread "seghist_global" num_groups group_size (segFlat space) $ do + constants <- kernelConstants <$> askEnv + + -- Compute subhistogram index for each thread, per histogram. + subhisto_inds <- forM slugs $ \slug -> + dPrimVE "subhisto_ind" $ + kernelGlobalThreadId constants + `quot` ( kernelNumThreads constants + `divUp` sExt32 (tvExp (slugNumSubhistos slug)) + ) + + -- Loop over flat offsets into the input and output. The + -- calculation is done with 64-bit integers to avoid overflow, + -- but the final unflattened segment indexes are 32 bit. + let gtid = sExt64 $ kernelGlobalThreadId constants + num_threads = sExt64 $ kernelNumThreads constants + kernelLoop gtid num_threads total_w_64 $ \offset -> do + -- Construct segment indices. + zipWithM_ dPrimV_ space_is $ + map sExt32 $ unflattenIndex space_sizes_64 offset + + -- We execute the bucket function once and update each histogram serially. + -- We apply the bucket function if j=offset+ltid is less than + -- num_elements. This also involves writing to the mapout + -- arrays. + let input_in_bounds = offset .<. total_w_64 + + sWhen input_in_bounds $ + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody + + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, res) -> + copyDWIMFix + (patElemName pe) + (map (Imp.vi64 . fst) $ unSegSpace space) + (kernelResultSubExp res) + [] + + let (buckets, vs) = splitAt (length slugs) red_res + perOp = chunks $ map (length . histDest . slugOp) slugs + + sComment "perform atomic updates" $ + forM_ (zip6 (map slugOp slugs) histograms buckets (perOp vs) subhisto_inds hist_H_chks) $ + \( HistOp dest_w _ _ _ shape lam, + do_op, + bucket, + vs', + subhisto_ind, + hist_H_chk + ) -> do + let chk_beg = sExt64 chk_i * hist_H_chk + bucket' = toInt64Exp $ kernelResultSubExp bucket + dest_w' = toInt64Exp dest_w + bucket_in_bounds = + chk_beg .<=. bucket' + .&&. bucket' .<. (chk_beg + hist_H_chk) + .&&. bucket' .<. dest_w' + vs_params = takeLast (length vs') $ lambdaParams lam + + sWhen bucket_in_bounds $ do + let bucket_is = + map Imp.vi64 (init space_is) + ++ [sExt64 subhisto_ind, bucket'] + dLParams $ lambdaParams lam + sLoopNest shape $ \is -> do + forM_ (zip vs_params vs') $ \(p, res) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp res) is + do_op (bucket_is ++ is) + +histKernelGlobal :: + [PatElem GPUMem] -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegHistSlug] -> + KernelBody GPUMem -> + CallKernelGen () +histKernelGlobal map_pes num_groups group_size space slugs kbody = do + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + let (_space_is, space_sizes) = unzip $ unSegSpace space + num_threads = sExt32 $ unCount num_groups' * unCount group_size' + + emit $ Imp.DebugPrint "## Using global memory" Nothing + + (hist_S, histograms) <- + prepareIntermediateArraysGlobal + (bodyPassage kbody) + num_threads + (toInt64Exp $ last space_sizes) + slugs + + sFor "chk_i" hist_S $ \chk_i -> + histKernelGlobalPass + map_pes + num_groups' + group_size' + space + slugs + kbody + histograms + hist_S + chk_i + +type InitLocalHistograms = + [ ( [VName], + SubExp -> + InKernelGen + ( [VName], + [Imp.TExp Int64] -> InKernelGen () + ) + ) + ] + +prepareIntermediateArraysLocal :: + TV Int32 -> + Count NumGroups (Imp.TExp Int64) -> + SegSpace -> + [SegHistSlug] -> + CallKernelGen InitLocalHistograms +prepareIntermediateArraysLocal num_subhistos_per_group groups_per_segment space slugs = do + num_segments <- + dPrimVE "num_segments" $ + product $ map (toInt64Exp . snd) $ init $ unSegSpace space + mapM (onOp num_segments) slugs + where + onOp num_segments (SegHistSlug op num_subhistos subhisto_info do_op) = do + num_subhistos <-- sExt64 (unCount groups_per_segment) * num_segments + + emit $ + Imp.DebugPrint "Number of subhistograms in global memory" $ + Just $ untyped $ tvExp num_subhistos + + mk_op <- + case do_op of + AtomicPrim f -> return $ const $ return f + AtomicCAS f -> return $ const $ return f + AtomicLocking f -> return $ \hist_H_chk -> do + let lock_shape = + Shape $ + tvSize num_subhistos_per_group : + shapeDims (histShape op) + ++ [hist_H_chk] + + let dims = map toInt64Exp $ shapeDims lock_shape + + locks <- sAllocArray "locks" int32 lock_shape $ Space "local" + + sComment "All locks start out unlocked" $ + groupCoverSpace dims $ \is -> + copyDWIMFix locks is (intConst Int32 0) [] + + return $ f $ Locking locks 0 1 0 id + + -- Initialise local-memory sub-histograms. These are + -- represented as two-dimensional arrays. + let init_local_subhistos hist_H_chk = do + local_subhistos <- + forM (histType op) $ \t -> do + let sub_local_shape = + Shape [tvSize num_subhistos_per_group] + <> (arrayShape t `setOuterDim` hist_H_chk) + sAllocArray + "subhistogram_local" + (elemType t) + sub_local_shape + (Space "local") + + do_op' <- mk_op hist_H_chk + + return (local_subhistos, do_op' (Space "local") local_subhistos) + + -- Initialise global-memory sub-histograms. + glob_subhistos <- forM subhisto_info $ \info -> do + subhistosAlloc info + return $ subhistosArray info + + return (glob_subhistos, init_local_subhistos) + +histKernelLocalPass :: + TV Int32 -> + Count NumGroups (Imp.TExp Int64) -> + [PatElem GPUMem] -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + SegSpace -> + [SegHistSlug] -> + KernelBody GPUMem -> + InitLocalHistograms -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + CallKernelGen () +histKernelLocalPass + num_subhistos_per_group_var + groups_per_segment + map_pes + num_groups + group_size + space + slugs + kbody + init_histograms + hist_S + chk_i = do + let (space_is, space_sizes) = unzip $ unSegSpace space + segment_is = init space_is + segment_dims = init space_sizes + (i_in_segment, segment_size) = last $ unSegSpace space + num_subhistos_per_group = tvExp num_subhistos_per_group_var + segment_size' = toInt64Exp segment_size + + num_segments <- + dPrimVE "num_segments" $ + product $ map toInt64Exp segment_dims + + hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> + dPrimV "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S + + histo_sizes <- forM (zip slugs hist_H_chks) $ \(slug, hist_H_chk) -> do + let histo_dims = + tvExp hist_H_chk : + map toInt64Exp (shapeDims (histShape (slugOp slug))) + histo_size <- + dPrimVE "histo_size" $ product histo_dims + let group_hists_size = + sExt64 num_subhistos_per_group * histo_size + init_per_thread <- + dPrimVE "init_per_thread" $ sExt32 $ group_hists_size `divUp` unCount group_size + return (histo_dims, histo_size, init_per_thread) + + sKernelThread "seghist_local" num_groups group_size (segFlat space) $ + virtualiseGroups SegVirt (sExt32 $ unCount groups_per_segment * num_segments) $ \group_id -> do + constants <- kernelConstants <$> askEnv + + flat_segment_id <- dPrimVE "flat_segment_id" $ group_id `quot` sExt32 (unCount groups_per_segment) + gid_in_segment <- dPrimVE "gid_in_segment" $ group_id `rem` sExt32 (unCount groups_per_segment) + -- This pgtid is kind of a "virtualised physical" gtid - not the + -- same thing as the gtid used for the SegHist itself. + pgtid_in_segment <- + dPrimVE "pgtid_in_segment" $ + gid_in_segment * sExt32 (kernelGroupSize constants) + + kernelLocalThreadId constants + threads_per_segment <- + dPrimVE "threads_per_segment" $ + sExt32 $ unCount groups_per_segment * kernelGroupSize constants + + -- Set segment indices. + zipWithM_ dPrimV_ segment_is $ + unflattenIndex (map toInt64Exp segment_dims) $ sExt64 flat_segment_id + + histograms <- forM (zip init_histograms hist_H_chks) $ + \((glob_subhistos, init_local_subhistos), hist_H_chk) -> do + (local_subhistos, do_op) <- init_local_subhistos $ Var $ tvVar hist_H_chk + return (zip glob_subhistos local_subhistos, hist_H_chk, do_op) + + -- Find index of local subhistograms updated by this thread. We + -- try to ensure, as much as possible, that threads in the same + -- warp use different subhistograms, to avoid conflicts. + thread_local_subhisto_i <- + dPrimVE "thread_local_subhisto_i" $ + kernelLocalThreadId constants `rem` num_subhistos_per_group + + let onSlugs f = + forM_ (zip3 slugs histograms histo_sizes) $ + \(slug, (dests, hist_H_chk, _), (histo_dims, histo_size, init_per_thread)) -> + f slug dests (tvExp hist_H_chk) histo_dims histo_size init_per_thread + + let onAllHistograms f = + onSlugs $ \slug dests hist_H_chk histo_dims histo_size init_per_thread -> do + let group_hists_size = num_subhistos_per_group * sExt32 histo_size + + forM_ (zip dests (histNeutral $ slugOp slug)) $ + \((dest_global, dest_local), ne) -> + sFor "local_i" init_per_thread $ \i -> do + j <- + dPrimVE "j" $ + i * sExt32 (kernelGroupSize constants) + + kernelLocalThreadId constants + j_offset <- + dPrimVE "j_offset" $ + num_subhistos_per_group * sExt32 histo_size * gid_in_segment + j + + local_subhisto_i <- dPrimVE "local_subhisto_i" $ j `quot` sExt32 histo_size + let local_bucket_is = unflattenIndex histo_dims $ sExt64 $ j `rem` sExt32 histo_size + global_bucket_is = + head local_bucket_is + sExt64 chk_i * hist_H_chk : + tail local_bucket_is + global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` sExt32 histo_size + + sWhen (j .<. group_hists_size) $ + f + dest_local + dest_global + (slugOp slug) + ne + local_subhisto_i + global_subhisto_i + local_bucket_is + global_bucket_is + + sComment "initialize histograms in local memory" $ + onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i local_bucket_is global_bucket_is -> + sComment "First subhistogram is initialised from global memory; others with neutral element." $ do + let global_is = map Imp.vi64 segment_is ++ [0] ++ global_bucket_is + local_is = sExt64 local_subhisto_i : local_bucket_is + sIf + (global_subhisto_i .==. 0) + (copyDWIMFix dest_local local_is (Var dest_global) global_is) + ( sLoopNest (histShape op) $ \is -> + copyDWIMFix dest_local (local_is ++ is) ne [] + ) + + sOp $ Imp.Barrier Imp.FenceLocal + + kernelLoop pgtid_in_segment threads_per_segment (sExt32 segment_size') $ \ie -> do + dPrimV_ i_in_segment ie + + -- We execute the bucket function once and update each histogram + -- serially. This also involves writing to the mapout arrays if + -- this is the first chunk. + + compileStms mempty (kernelBodyStms kbody) $ do + let (red_res, map_res) = + splitFromEnd (length map_pes) $ + map kernelResultSubExp $ kernelBodyResult kbody + (buckets, vs) = splitAt (length slugs) red_res + perOp = chunks $ map (length . histDest . slugOp) slugs + + sWhen (chk_i .==. 0) $ + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, se) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 space_is) + se + [] + + forM_ (zip4 (map slugOp slugs) histograms buckets (perOp vs)) $ + \( HistOp dest_w _ _ _ shape lam, + (_, hist_H_chk, do_op), + bucket, + vs' + ) -> do + let chk_beg = sExt64 chk_i * tvExp hist_H_chk + bucket' = toInt64Exp bucket + dest_w' = toInt64Exp dest_w + bucket_in_bounds = + bucket' .<. dest_w' + .&&. chk_beg .<=. bucket' + .&&. bucket' .<. (chk_beg + tvExp hist_H_chk) + bucket_is = [sExt64 thread_local_subhisto_i, bucket' - chk_beg] + vs_params = takeLast (length vs') $ lambdaParams lam + + sComment "perform atomic updates" $ + sWhen bucket_in_bounds $ do + dLParams $ lambdaParams lam + sLoopNest shape $ \is -> do + forM_ (zip vs_params vs') $ \(p, v) -> + copyDWIMFix (paramName p) [] v is + do_op (bucket_is ++ is) + + sOp $ Imp.ErrorSync Imp.FenceGlobal + + sComment "Compact the multiple local memory subhistograms to result in global memory" $ + onSlugs $ \slug dests hist_H_chk histo_dims _histo_size bins_per_thread -> do + trunc_H <- + dPrimV "trunc_H" $ + sMin64 hist_H_chk $ + toInt64Exp (histWidth (slugOp slug)) + - sExt64 chk_i * head histo_dims + let trunc_histo_dims = + tvExp trunc_H : + map toInt64Exp (shapeDims (histShape (slugOp slug))) + trunc_histo_size <- dPrimVE "histo_size" $ sExt32 $ product trunc_histo_dims + + sFor "local_i" bins_per_thread $ \i -> do + j <- + dPrimVE "j" $ + i * sExt32 (kernelGroupSize constants) + + kernelLocalThreadId constants + sWhen (j .<. trunc_histo_size) $ do + -- We are responsible for compacting the flat bin 'j', which + -- we immediately unflatten. + let local_bucket_is = unflattenIndex histo_dims $ sExt64 j + global_bucket_is = + head local_bucket_is + sExt64 chk_i * hist_H_chk : + tail local_bucket_is + dLParams $ lambdaParams $ histOp $ slugOp slug + let (global_dests, local_dests) = unzip dests + (xparams, yparams) = + splitAt (length local_dests) $ + lambdaParams $ histOp $ slugOp slug + + sComment "Read values from subhistogram 0." $ + forM_ (zip xparams local_dests) $ \(xp, subhisto) -> + copyDWIMFix + (paramName xp) + [] + (Var subhisto) + (0 : local_bucket_is) + + sComment "Accumulate based on values in other subhistograms." $ + sFor "subhisto_id" (num_subhistos_per_group - 1) $ \subhisto_id -> do + forM_ (zip yparams local_dests) $ \(yp, subhisto) -> + copyDWIMFix + (paramName yp) + [] + (Var subhisto) + (sExt64 subhisto_id + 1 : local_bucket_is) + compileBody' xparams $ lambdaBody $ histOp $ slugOp slug + + sComment "Put final bucket value in global memory." $ do + let global_is = + map Imp.vi64 segment_is + ++ [sExt64 group_id `rem` unCount groups_per_segment] + ++ global_bucket_is + forM_ (zip xparams global_dests) $ \(xp, global_dest) -> + copyDWIMFix global_dest global_is (Var $ paramName xp) [] + +histKernelLocal :: + TV Int32 -> + Count NumGroups (Imp.TExp Int64) -> + [PatElem GPUMem] -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + Imp.TExp Int32 -> + [SegHistSlug] -> + KernelBody GPUMem -> + CallKernelGen () +histKernelLocal num_subhistos_per_group_var groups_per_segment map_pes num_groups group_size space hist_S slugs kbody = do + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + num_subhistos_per_group = tvExp num_subhistos_per_group_var + + emit $ + Imp.DebugPrint "Number of local subhistograms per group" $ + Just $ untyped num_subhistos_per_group + + init_histograms <- + prepareIntermediateArraysLocal num_subhistos_per_group_var groups_per_segment space slugs + + sFor "chk_i" hist_S $ \chk_i -> + histKernelLocalPass + num_subhistos_per_group_var + groups_per_segment + map_pes + num_groups' + group_size' + space + slugs + kbody + init_histograms + hist_S + chk_i + +-- | The maximum number of passes we are willing to accept for this +-- kind of atomic update. +slugMaxLocalMemPasses :: SegHistSlug -> Int +slugMaxLocalMemPasses slug = + case slugAtomicUpdate slug of + AtomicPrim _ -> 3 + AtomicCAS _ -> 4 + AtomicLocking _ -> 6 + +localMemoryCase :: + [PatElem GPUMem] -> + Imp.TExp Int32 -> + SegSpace -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.TExp Int32 -> + [SegHistSlug] -> + KernelBody GPUMem -> + CallKernelGen (Imp.TExp Bool, CallKernelGen ()) +localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N _ slugs kbody = do + let space_sizes = segSpaceDims space + segment_dims = init space_sizes + segmented = not $ null segment_dims + + hist_L <- dPrim "hist_L" int32 + sOp $ Imp.GetSizeMax (tvVar hist_L) Imp.SizeLocalMemory + + max_group_size <- dPrim "max_group_size" int32 + sOp $ Imp.GetSizeMax (tvVar max_group_size) Imp.SizeGroup + let group_size = Imp.Count $ Var $ tvVar max_group_size + num_groups <- + fmap (Imp.Count . tvSize) $ + dPrimV "num_groups" $ + sExt64 hist_T `divUp` toInt64Exp (unCount group_size) + let num_groups' = toInt64Exp <$> num_groups + group_size' = toInt64Exp <$> group_size + + let r64 = isF64 . ConvOpExp (SIToFP Int64 Float64) . untyped + t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped + + -- M approximation. + hist_m' <- + dPrimVE "hist_m_prime" $ + r64 + ( sMin64 + (sExt64 (tvExp hist_L `quot` hist_el_size)) + (hist_N `divUp` sExt64 (unCount num_groups')) + ) + / r64 hist_H + + let hist_B = unCount group_size' + + -- M in the paper, but not adjusted for asymptotic efficiency. + hist_M0 <- + dPrimVE "hist_M0" $ + sMax64 1 $ sMin64 (t64 hist_m') hist_B + + -- Minimal sequential chunking factor. + let q_small = 2 + + -- The number of segments/histograms produced.. + hist_Nout <- dPrimVE "hist_Nout" $ product $ map toInt64Exp segment_dims + + hist_Nin <- dPrimVE "hist_Nin" $ toInt64Exp $ last space_sizes + + -- Maximum M for work efficiency. + work_asymp_M_max <- + if segmented + then do + hist_T_hist_min <- + dPrimVE "hist_T_hist_min" $ + sExt32 $ + sMin64 (sExt64 hist_Nin * sExt64 hist_Nout) (sExt64 hist_T) + `divUp` sExt64 hist_Nout + + -- Number of groups, rounded up. + let r = hist_T_hist_min `divUp` sExt32 hist_B + + dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (sExt64 r * hist_H) + else + dPrimVE "work_asymp_M_max" $ + (hist_Nout * hist_N) + `quot` ( (q_small * unCount num_groups' * hist_H) + `quot` genericLength slugs + ) + + -- Number of subhistograms per result histogram. + hist_M <- dPrimV "hist_M" $ sExt32 $ sMin64 hist_M0 work_asymp_M_max + + -- hist_M may be zero (which we'll check for below), but we need it + -- for some divisions first, so crudely make a nonzero form. + let hist_M_nonzero = sMax32 1 $ tvExp hist_M + + -- "Cooperation factor" - the number of threads cooperatively + -- working on the same (sub)histogram. + hist_C <- + dPrimVE "hist_C" $ + hist_B `divUp` sExt64 hist_M_nonzero + + emit $ Imp.DebugPrint "local hist_M0" $ Just $ untyped hist_M0 + emit $ Imp.DebugPrint "local work asymp M max" $ Just $ untyped work_asymp_M_max + emit $ Imp.DebugPrint "local C" $ Just $ untyped hist_C + emit $ Imp.DebugPrint "local B" $ Just $ untyped hist_B + emit $ Imp.DebugPrint "local M" $ Just $ untyped $ tvExp hist_M + emit $ + Imp.DebugPrint "local memory needed" $ + Just $ untyped $ hist_H * hist_el_size * sExt64 (tvExp hist_M) + + -- local_mem_needed is what we need to keep a single bucket in local + -- memory - this is an absolute minimum. We can fit anything else + -- by doing multiple passes, although more than a few is + -- (heuristically) not efficient. + local_mem_needed <- + dPrimVE "local_mem_needed" $ + hist_el_size * sExt64 (tvExp hist_M) + hist_S <- + dPrimVE "hist_S" $ + sExt32 $ + (hist_H * local_mem_needed) `divUp` tvExp hist_L + let max_S = case bodyPassage kbody of + MustBeSinglePass -> 1 + MayBeMultiPass -> fromIntegral $ maxinum $ map slugMaxLocalMemPasses slugs + + groups_per_segment <- + if segmented + then + fmap Count $ + dPrimVE "groups_per_segment" $ unCount num_groups' `divUp` hist_Nout + else pure num_groups' + + -- We only use local memory if the number of updates per histogram + -- at least matches the histogram size, as otherwise it is not + -- asymptotically efficient. This mostly matters for the segmented + -- case. + let pick_local = + hist_Nin .>=. hist_H + .&&. (local_mem_needed .<=. tvExp hist_L) + .&&. (hist_S .<=. max_S) + .&&. hist_C .<=. hist_B + .&&. tvExp hist_M .>. 0 + + run = do + emit $ Imp.DebugPrint "## Using local memory" Nothing + emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H + emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ untyped $ tvExp hist_M + emit $ Imp.DebugPrint "Cooperation level (C)" $ Just $ untyped hist_C + emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ untyped hist_S + when segmented $ + emit $ Imp.DebugPrint "Groups per segment" $ Just $ untyped $ unCount groups_per_segment + histKernelLocal + hist_M + groups_per_segment + map_pes + num_groups + group_size + space + hist_S + slugs + kbody + + return (pick_local, run) + +-- | Generate code for a segmented histogram called from the host. +compileSegHist :: + Pattern GPUMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [HistOp GPUMem] -> + KernelBody GPUMem -> + CallKernelGen () +compileSegHist (Pattern _ pes) num_groups group_size space ops kbody = do + -- Most of this function is not the histogram part itself, but + -- rather figuring out whether to use a local or global memory + -- strategy, as well as collapsing the subhistograms produced (which + -- are always in global memory, but their number may vary). + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + dims = map toInt64Exp $ segSpaceDims space + + num_red_res = length ops + sum (map (length . histNeutral) ops) + (all_red_pes, map_pes) = splitAt num_red_res pes + segment_size = last dims + + (op_hs, op_seg_hs, slugs) <- unzip3 <$> mapM (computeHistoUsage space) ops + h <- dPrimVE "h" $ Imp.unCount $ sum op_hs + seg_h <- dPrimVE "seg_h" $ Imp.unCount $ sum op_seg_hs + + -- Check for emptyness to avoid division-by-zero. + sUnless (seg_h .==. 0) $ do + -- Maximum group size (or actual, in this case). + let hist_B = unCount group_size' + + -- Size of a histogram. + hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth) ops + + -- Size of a single histogram element. Actually the weighted + -- average of histogram elements in cases where we have more than + -- one histogram operation, plus any locks. + let lockSize slug = case slugAtomicUpdate slug of + AtomicLocking {} -> Just $ primByteSize int32 + _ -> Nothing + hist_el_size <- + dPrimVE "hist_el_size" $ + foldl' (+) (h `divUp` hist_H) $ + mapMaybe lockSize slugs + + -- Input elements contributing to each histogram. + hist_N <- dPrimVE "hist_N" segment_size + + -- Compute RF as the average RF over all the histograms. + hist_RF <- + dPrimVE "hist_RF" $ + sExt32 $ + sum (map (toInt64Exp . histRaceFactor . slugOp) slugs) + `quot` genericLength slugs + + let hist_T = sExt32 $ unCount num_groups' * unCount group_size' + emit $ Imp.DebugPrint "\n# SegHist" Nothing + emit $ Imp.DebugPrint "Number of threads (T)" $ Just $ untyped hist_T + emit $ Imp.DebugPrint "Desired group size (B)" $ Just $ untyped hist_B + emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H + emit $ Imp.DebugPrint "Input elements per histogram (N)" $ Just $ untyped hist_N + emit $ + Imp.DebugPrint "Number of segments" $ + Just $ untyped $ product $ map (toInt64Exp . snd) segment_dims + emit $ Imp.DebugPrint "Histogram element size (el_size)" $ Just $ untyped hist_el_size + emit $ Imp.DebugPrint "Race factor (RF)" $ Just $ untyped hist_RF + emit $ Imp.DebugPrint "Memory per set of subhistograms per segment" $ Just $ untyped h + emit $ Imp.DebugPrint "Memory per set of subhistograms times segments" $ Just $ untyped seg_h + + (use_local_memory, run_in_local_memory) <- + localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N hist_RF slugs kbody + + sIf use_local_memory run_in_local_memory $ + histKernelGlobal map_pes num_groups group_size space slugs kbody + + let pes_per_op = chunks (map (length . histDest) ops) all_red_pes + + forM_ (zip3 slugs pes_per_op ops) $ \(slug, red_pes, op) -> do + let num_histos = slugNumSubhistos slug + subhistos = map subhistosArray $ slugSubhistos slug + + let unitHistoCase = + -- This is OK because the memory blocks are at least as + -- large as the ones we are supposed to use for the result. + forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do + pe_mem <- + memLocationName . entryArrayLocation + <$> lookupArray (patElemName pe) + subhisto_mem <- + memLocationName . entryArrayLocation + <$> lookupArray subhisto + emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device" + + sIf (tvExp num_histos .==. 1) unitHistoCase $ do + -- For the segmented reduction, we keep the segment dimensions + -- unchanged. To this, we add two dimensions: one over the number + -- of buckets, and one over the number of subhistograms. This + -- inner dimension is the one that is collapsed in the reduction. + let num_buckets = histWidth op + + bucket_id <- newVName "bucket_id" + subhistogram_id <- newVName "subhistogram_id" + vector_ids <- + mapM (const $ newVName "vector_id") $ + shapeDims $ histShape op + + flat_gtid <- newVName "flat_gtid" + + let lvl = SegThread num_groups group_size SegVirt + segred_space = + SegSpace flat_gtid $ + segment_dims + ++ [(bucket_id, num_buckets)] + ++ zip vector_ids (shapeDims $ histShape op) + ++ [(subhistogram_id, Var $ tvVar num_histos)] + + let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty + compileSegRed' (Pattern [] red_pes) lvl segred_space [segred_op] $ \red_cont -> + red_cont $ + flip map subhistos $ \subhisto -> + ( Var subhisto, + map Imp.vi64 $ + map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids + ) + + emit $ Imp.DebugPrint "" Nothing + where + segment_dims = init $ unSegSpace space diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs @@ -0,0 +1,61 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Code generation for 'SegMap' is quite straightforward. The only +-- trick is virtualisation in case the physical number of threads is +-- not sufficient to cover the logical thread space. This is handled +-- by having actual workgroups run a loop to imitate multiple workgroups. +module Futhark.CodeGen.ImpGen.GPU.SegMap (compileSegMap) where + +import Control.Monad.Except +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.IR.GPUMem +import Futhark.Util.IntegralExp (divUp) +import Prelude hiding (quot, rem) + +-- | Compile 'SegMap' instance code. +compileSegMap :: + Pattern GPUMem -> + SegLevel -> + SegSpace -> + KernelBody GPUMem -> + CallKernelGen () +compileSegMap pat lvl space kbody = do + let (is, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + num_groups' = toInt64Exp <$> segNumGroups lvl + group_size' = toInt64Exp <$> segGroupSize lvl + + emit $ Imp.DebugPrint "\n# SegMap" Nothing + case lvl of + SegThread {} -> do + let virt_num_groups = + sExt32 $ product dims' `divUp` unCount group_size' + sKernelThread "segmap" num_groups' group_size' (segFlat space) $ + virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do + local_tid <- kernelLocalThreadId . kernelConstants <$> askEnv + + global_tid <- + dPrimVE "global_tid" $ + sExt64 group_id * sExt64 (unCount group_size') + + sExt64 local_tid + + dIndexSpace (zip is dims') global_tid + + sWhen (isActive $ unSegSpace space) $ + compileStms mempty (kernelBodyStms kbody) $ + zipWithM_ (compileThreadResult space) (patternElements pat) $ + kernelBodyResult kbody + SegGroup {} -> + sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ do + let virt_num_groups = sExt32 $ product dims' + precomputeSegOpIDs (kernelBodyStms kbody) $ + virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do + dIndexSpace (zip is dims') $ sExt64 group_id + + compileStms mempty (kernelBodyStms kbody) $ + zipWithM_ (compileGroupResult space) (patternElements pat) $ + kernelBodyResult kbody + emit $ Imp.DebugPrint "" Nothing diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs @@ -0,0 +1,842 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | We generate code for non-segmented/single-segment SegRed using +-- the basic approach outlined in the paper "Design and GPGPU +-- Performance of Futhark’s Redomap Construct" (ARRAY '16). The main +-- deviations are: +-- +-- * While we still use two-phase reduction, we use only a single +-- kernel, with the final workgroup to write a result (tracked via +-- an atomic counter) performing the final reduction as well. +-- +-- * Instead of depending on storage layout transformations to handle +-- non-commutative reductions efficiently, we slide a +-- @groupsize@-sized window over the input, and perform a parallel +-- reduction for each window. This sacrifices the notion of +-- efficient sequentialisation, but is sometimes faster and +-- definitely simpler and more predictable (and uses less auxiliary +-- storage). +-- +-- For segmented reductions we use the approach from "Strategies for +-- Regular Segmented Reductions on GPU" (FHPC '17). This involves +-- having two different strategies, and dynamically deciding which one +-- to use based on the number of segments and segment size. We use the +-- (static) @group_size@ to decide which of the following two +-- strategies to choose: +-- +-- * Large: uses one or more groups to process a single segment. If +-- multiple groups are used per segment, the intermediate reduction +-- results must be recursively reduced, until there is only a single +-- value per segment. +-- +-- Each thread /can/ read multiple elements, which will greatly +-- increase performance; however, if the reduction is +-- non-commutative we will have to use a less efficient traversal +-- (with interim group-wide reductions) to enable coalesced memory +-- accesses, just as in the non-segmented case. +-- +-- * Small: is used to let each group process *multiple* segments +-- within a group. We will only use this approach when we can +-- process at least two segments within a single group. In those +-- cases, we would allocate a /whole/ group per segment with the +-- large strategy, but at most 50% of the threads in the group would +-- have any element to read, which becomes highly inefficient. +module Futhark.CodeGen.ImpGen.GPU.SegRed + ( compileSegRed, + compileSegRed', + DoSegBody, + ) +where + +import Control.Monad.Except +import Data.List (genericLength, zip7) +import Data.Maybe +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.Error +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.Transform.Rename +import Futhark.Util (chunks) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +-- | The maximum number of operators we support in a single SegRed. +-- This limit arises out of the static allocation of counters. +maxNumOps :: Int32 +maxNumOps = 10 + +-- | Code generation for the body of the SegRed, taking a continuation +-- for saving the results of the body. The results should be +-- represented as a pairing of a t'SubExp' along with a list of +-- indexes into that t'SubExp' for reading the result. +type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen () + +-- | Compile 'SegRed' instance to host-level code with calls to +-- various kernels. +compileSegRed :: + Pattern GPUMem -> + SegLevel -> + SegSpace -> + [SegBinOp GPUMem] -> + KernelBody GPUMem -> + CallKernelGen () +compileSegRed pat lvl space reds body = + compileSegRed' pat lvl space reds $ \red_cont -> + compileStms mempty (kernelBodyStms body) $ do + let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult body + + sComment "save map-out results" $ do + let map_arrs = drop (segBinOpResults reds) $ patternElements pat + zipWithM_ (compileThreadResult space) map_arrs map_res + + red_cont $ zip (map kernelResultSubExp red_res) $ repeat [] + +-- | Like 'compileSegRed', but where the body is a monadic action. +compileSegRed' :: + Pattern GPUMem -> + SegLevel -> + SegSpace -> + [SegBinOp GPUMem] -> + DoSegBody -> + CallKernelGen () +compileSegRed' pat lvl space reds body + | genericLength reds > maxNumOps = + compilerLimitationS $ + "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported." + | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space = + nonsegmentedReduction pat num_groups group_size space reds body + | otherwise = do + let group_size' = toInt64Exp $ unCount group_size + segment_size = toInt64Exp $ last $ segSpaceDims space + use_small_segments = segment_size * 2 .<. group_size' + sIf + use_small_segments + (smallSegmentsReduction pat num_groups group_size space reds body) + (largeSegmentsReduction pat num_groups group_size space reds body) + where + num_groups = segNumGroups lvl + group_size = segGroupSize lvl + +-- | Prepare intermediate arrays for the reduction. Prim-typed +-- arguments go in local memory (so we need to do the allocation of +-- those arrays inside the kernel), while array-typed arguments go in +-- global memory. Allocations for the former have already been +-- performed. This policy is baked into how the allocations are done +-- in ExplicitAllocations. +intermediateArrays :: + Count GroupSize SubExp -> + SubExp -> + SegBinOp GPUMem -> + InKernelGen [VName] +intermediateArrays (Count group_size) num_threads (SegBinOp _ red_op nes _) = do + let red_op_params = lambdaParams red_op + (red_acc_params, _) = splitAt (length nes) red_op_params + forM red_acc_params $ \p -> + case paramDec p of + MemArray pt shape _ (ArrayIn mem _) -> do + let shape' = Shape [num_threads] <> shape + sArray "red_arr" pt shape' $ + ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' + _ -> do + let pt = elemType $ paramType p + shape = Shape [group_size] + sAllocArray "red_arr" pt shape $ Space "local" + +-- | Arrays for storing group results. +-- +-- The group-result arrays have an extra dimension because they are +-- also used for keeping vectorised accumulators for first-stage +-- reduction, if necessary. If necessary, this dimension has size +-- group_size, and otherwise 1. When actually storing group results, +-- the first index is set to 0. +groupResultArrays :: + Count NumGroups SubExp -> + Count GroupSize SubExp -> + [SegBinOp GPUMem] -> + CallKernelGen [[VName]] +groupResultArrays (Count virt_num_groups) (Count group_size) reds = + forM reds $ \(SegBinOp _ lam _ shape) -> + forM (lambdaReturnType lam) $ \t -> do + let pt = elemType t + extra_dim + | primType t = intConst Int64 1 + | otherwise = group_size + full_shape = Shape [extra_dim, virt_num_groups] <> shape <> arrayShape t + -- Move the groupsize dimension last to ensure coalesced + -- memory access. + perm = [1 .. shapeRank full_shape -1] ++ [0] + sAllocArrayPerm "segred_tmp" pt full_shape (Space "device") perm + +nonsegmentedReduction :: + Pattern GPUMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp GPUMem] -> + DoSegBody -> + CallKernelGen () +nonsegmentedReduction segred_pat num_groups group_size space reds body = do + let (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + global_tid = Imp.vi64 $ segFlat space + w = last dims' + + counter <- + sStaticArray "counter" (Space "device") int32 $ + Imp.ArrayValues $ replicate (fromIntegral maxNumOps) $ IntValue $ Int32Value 0 + + reds_group_res_arrs <- groupResultArrays num_groups group_size reds + + num_threads <- + dPrimV "num_threads" $ + unCount num_groups' * unCount group_size' + + emit $ Imp.DebugPrint "\n# SegRed" Nothing + + sKernelThread "segred_nonseg" num_groups' group_size' (segFlat space) $ do + constants <- kernelConstants <$> askEnv + sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" + reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds + + -- Since this is the nonsegmented case, all outer segment IDs must + -- necessarily be 0. + forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64) + + let num_elements = Imp.elements w + elems_per_thread = + num_elements + `divUp` Imp.elements (sExt64 (kernelNumThreads constants)) + + slugs <- + mapM (segBinOpSlug (kernelLocalThreadId constants) (kernelGroupId constants)) $ + zip3 reds reds_arrs reds_group_res_arrs + reds_op_renamed <- + reductionStageOne + constants + (zip gtids dims') + num_elements + global_tid + elems_per_thread + (tvVar num_threads) + slugs + body + + let segred_pes = + chunks (map (length . segBinOpNeutral) reds) $ + patternElements segred_pat + forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $ + \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do + let (red_x_params, red_y_params) = splitAt (length nes) $ lambdaParams red_op + reductionStageTwo + constants + pes + (kernelGroupId constants) + 0 + [0] + 0 + (sExt64 $ kernelNumGroups constants) + slug + red_x_params + red_y_params + red_op_renamed + nes + 1 + counter + (fromInteger i) + sync_arr + group_res_arrs + red_arrs + + emit $ Imp.DebugPrint "" Nothing + +smallSegmentsReduction :: + Pattern GPUMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp GPUMem] -> + DoSegBody -> + CallKernelGen () +smallSegmentsReduction (Pattern _ segred_pes) num_groups group_size space reds body = do + let (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + segment_size = last dims' + + -- Careful to avoid division by zero now. + segment_size_nonzero <- + dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size + + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' + let num_segments = product $ init dims' + segments_per_group = unCount group_size' `quot` segment_size_nonzero + required_groups = sExt32 $ num_segments `divUp` segments_per_group + + emit $ Imp.DebugPrint "\n# SegRed-small" Nothing + emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments + emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size + emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group + emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups + + sKernelThread "segred_small" num_groups' group_size' (segFlat space) $ do + constants <- kernelConstants <$> askEnv + reds_arrs <- mapM (intermediateArrays group_size (Var $ tvVar num_threads)) reds + + -- We probably do not have enough actual workgroups to cover the + -- entire iteration space. Some groups thus have to perform double + -- duty; we put an outer loop to accomplish this. + virtualiseGroups SegVirt required_groups $ \group_id' -> do + -- Compute the 'n' input indices. The outer 'n-1' correspond to + -- the segment ID, and are computed from the group id. The inner + -- is computed from the local thread id, and may be out-of-bounds. + let ltid = sExt64 $ kernelLocalThreadId constants + segment_index = + (ltid `quot` segment_size_nonzero) + + (sExt64 group_id' * sExt64 segments_per_group) + index_within_segment = ltid `rem` segment_size + + dIndexSpace (zip (init gtids) (init dims')) segment_index + dPrimV_ (last gtids) index_within_segment + + let out_of_bounds = + forM_ (zip reds reds_arrs) $ \(SegBinOp _ _ nes _, red_arrs) -> + forM_ (zip red_arrs nes) $ \(arr, ne) -> + copyDWIMFix arr [ltid] ne [] + + in_bounds = + body $ \red_res -> + sComment "save results to be reduced" $ do + let red_dests = zip (concat reds_arrs) $ repeat [ltid] + forM_ (zip red_dests red_res) $ \((d, d_is), (res, res_is)) -> + copyDWIMFix d d_is res res_is + + sComment "apply map function if in bounds" $ + sIf + ( segment_size .>. 0 + .&&. isActive (init $ zip gtids dims) + .&&. ltid .<. segment_size * segments_per_group + ) + in_bounds + out_of_bounds + + sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. + let crossesSegment from to = + (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) + sWhen (segment_size .>. 0) $ + sComment "perform segmented scan to imitate reduction" $ + forM_ (zip reds reds_arrs) $ \(SegBinOp _ red_op _ _, red_arrs) -> + groupScan + (Just crossesSegment) + (sExt64 $ tvExp num_threads) + (segment_size * segments_per_group) + red_op + red_arrs + + sOp $ Imp.Barrier Imp.FenceLocal + + sComment "save final values of segments" $ + sWhen + ( sExt64 group_id' * segments_per_group + sExt64 ltid .<. num_segments + .&&. ltid .<. segments_per_group + ) + $ forM_ (zip segred_pes (concat reds_arrs)) $ \(pe, arr) -> do + -- Figure out which segment result this thread should write... + let flat_segment_index = + sExt64 group_id' * segments_per_group + sExt64 ltid + gtids' = + unflattenIndex (init dims') flat_segment_index + copyDWIMFix + (patElemName pe) + gtids' + (Var arr) + [(ltid + 1) * segment_size_nonzero - 1] + + -- Finally another barrier, because we will be writing to the + -- local memory array first thing in the next iteration. + sOp $ Imp.Barrier Imp.FenceLocal + + emit $ Imp.DebugPrint "" Nothing + +largeSegmentsReduction :: + Pattern GPUMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp GPUMem] -> + DoSegBody -> + CallKernelGen () +largeSegmentsReduction segred_pat num_groups group_size space reds body = do + let (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + num_segments = product $ init dims' + segment_size = last dims' + num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + + (groups_per_segment, elems_per_thread) <- + groupsPerSegmentAndElementsPerThread + segment_size + num_segments + num_groups' + group_size' + virt_num_groups <- + dPrimV "virt_num_groups" $ + groups_per_segment * num_segments + + num_threads <- + dPrimV "num_threads" $ + unCount num_groups' * unCount group_size' + + threads_per_segment <- + dPrimV "threads_per_segment" $ + groups_per_segment * unCount group_size' + + emit $ Imp.DebugPrint "\n# SegRed-large" Nothing + emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments + emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size + emit $ Imp.DebugPrint "virt_num_groups" $ Just $ untyped $ tvExp virt_num_groups + emit $ Imp.DebugPrint "num_groups" $ Just $ untyped $ Imp.unCount num_groups' + emit $ Imp.DebugPrint "group_size" $ Just $ untyped $ Imp.unCount group_size' + emit $ Imp.DebugPrint "elems_per_thread" $ Just $ untyped $ Imp.unCount elems_per_thread + emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment + + reds_group_res_arrs <- groupResultArrays (Count (tvSize virt_num_groups)) group_size reds + + -- In principle we should have a counter for every segment. Since + -- the number of segments is a dynamic quantity, we would have to + -- allocate and zero out an array here, which is expensive. + -- However, we exploit the fact that the number of segments being + -- reduced at any point in time is limited by the number of + -- workgroups. If we bound the number of workgroups, we can get away + -- with using that many counters. FIXME: Is this limit checked + -- anywhere? There are other places in the compiler that will fail + -- if the group count exceeds the maximum group size, which is at + -- most 1024 anyway. + let num_counters = fromIntegral maxNumOps * 1024 + counter <- + sStaticArray "counter" (Space "device") int32 $ + Imp.ArrayZeros num_counters + + sKernelThread "segred_large" num_groups' group_size' (segFlat space) $ do + constants <- kernelConstants <$> askEnv + reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds + sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" + + -- We probably do not have enough actual workgroups to cover the + -- entire iteration space. Some groups thus have to perform double + -- duty; we put an outer loop to accomplish this. + virtualiseGroups SegVirt (sExt32 (tvExp virt_num_groups)) $ \group_id -> do + let segment_gtids = init gtids + w = last dims + local_tid = kernelLocalThreadId constants + + flat_segment_id <- + dPrimVE "flat_segment_id" $ + group_id `quot` sExt32 groups_per_segment + + global_tid <- + dPrimVE "global_tid" $ + (sExt64 group_id * sExt64 (unCount group_size') + sExt64 local_tid) + `rem` (sExt64 (unCount group_size') * groups_per_segment) + + let first_group_for_segment = sExt64 flat_segment_id * groups_per_segment + dIndexSpace (zip segment_gtids (init dims')) $sExt64 flat_segment_id + dPrim_ (last gtids) int64 + let num_elements = Imp.elements $ toInt64Exp w + + slugs <- + mapM (segBinOpSlug local_tid group_id) $ + zip3 reds reds_arrs reds_group_res_arrs + reds_op_renamed <- + reductionStageOne + constants + (zip gtids dims') + num_elements + global_tid + elems_per_thread + (tvVar threads_per_segment) + slugs + body + + let segred_pes = + chunks (map (length . segBinOpNeutral) reds) $ + patternElements segred_pat + + multiple_groups_per_segment = + forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $ + \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do + let (red_x_params, red_y_params) = + splitAt (length nes) $ lambdaParams red_op + reductionStageTwo + constants + pes + group_id + flat_segment_id + (map Imp.vi64 segment_gtids) + (sExt64 first_group_for_segment) + groups_per_segment + slug + red_x_params + red_y_params + red_op_renamed + nes + (fromIntegral num_counters) + counter + (fromInteger i) + sync_arr + group_res_arrs + red_arrs + + one_group_per_segment = + comment "first thread in group saves final result to memory" $ + forM_ (zip slugs segred_pes) $ \(slug, pes) -> + sWhen (local_tid .==. 0) $ + forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) -> + copyDWIMFix (patElemName v) (map Imp.vi64 segment_gtids) (Var acc) acc_is + + sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment + + emit $ Imp.DebugPrint "" Nothing + +-- Careful to avoid division by zero here. We have at least one group +-- per segment. +groupsPerSegmentAndElementsPerThread :: + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Count NumGroups (Imp.TExp Int64) -> + Count GroupSize (Imp.TExp Int64) -> + CallKernelGen + ( Imp.TExp Int64, + Imp.Count Imp.Elements (Imp.TExp Int64) + ) +groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = do + groups_per_segment <- + dPrimVE "groups_per_segment" $ + unCount num_groups_hint `divUp` sMax64 1 num_segments + elements_per_thread <- + dPrimVE "elements_per_thread" $ + segment_size `divUp` (unCount group_size * groups_per_segment) + return (groups_per_segment, Imp.elements elements_per_thread) + +-- | A SegBinOp with auxiliary information. +data SegBinOpSlug = SegBinOpSlug + { slugOp :: SegBinOp GPUMem, + -- | The arrays used for computing the intra-group reduction + -- (either local or global memory). + slugArrs :: [VName], + -- | Places to store accumulator in stage 1 reduction. + slugAccs :: [(VName, [Imp.TExp Int64])] + } + +slugBody :: SegBinOpSlug -> Body GPUMem +slugBody = lambdaBody . segBinOpLambda . slugOp + +slugParams :: SegBinOpSlug -> [LParam GPUMem] +slugParams = lambdaParams . segBinOpLambda . slugOp + +slugNeutral :: SegBinOpSlug -> [SubExp] +slugNeutral = segBinOpNeutral . slugOp + +slugShape :: SegBinOpSlug -> Shape +slugShape = segBinOpShape . slugOp + +slugsComm :: [SegBinOpSlug] -> Commutativity +slugsComm = mconcat . map (segBinOpComm . slugOp) + +accParams, nextParams :: SegBinOpSlug -> [LParam GPUMem] +accParams slug = take (length (slugNeutral slug)) $ slugParams slug +nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug + +segBinOpSlug :: Imp.TExp Int32 -> Imp.TExp Int32 -> (SegBinOp GPUMem, [VName], [VName]) -> InKernelGen SegBinOpSlug +segBinOpSlug local_tid group_id (op, group_res_arrs, param_arrs) = + SegBinOpSlug op group_res_arrs + <$> zipWithM mkAcc (lambdaParams (segBinOpLambda op)) param_arrs + where + mkAcc p param_arr + | Prim t <- paramType p, + shapeRank (segBinOpShape op) == 0 = do + acc <- dPrim (baseString (paramName p) <> "_acc") t + return (tvVar acc, []) + | otherwise = + return (param_arr, [sExt64 local_tid, sExt64 group_id]) + +reductionStageZero :: + KernelConstants -> + [(VName, Imp.TExp Int64)] -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + Imp.TExp Int64 -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + VName -> + [SegBinOpSlug] -> + DoSegBody -> + InKernelGen ([Lambda GPUMem], InKernelGen ()) +reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do + let (gtids, _dims) = unzip ispace + gtid = mkTV (last gtids) int64 + local_tid = sExt64 $ kernelLocalThreadId constants + + -- Figure out how many elements this thread should process. + chunk_size <- dPrim "chunk_size" int64 + let ordering = case slugsComm slugs of + Commutative -> SplitStrided $ Var threads_per_segment + Noncommutative -> SplitContiguous + computeThreadChunkSize ordering (sExt64 global_tid) elems_per_thread num_elements chunk_size + + dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs + + sComment "neutral-initialise the accumulators" $ + forM_ slugs $ \slug -> + forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> + sLoopNest (slugShape slug) $ \vec_is -> + copyDWIMFix acc (acc_is ++ vec_is) ne [] + + slugs_op_renamed <- mapM (renameLambda . segBinOpLambda . slugOp) slugs + + let doTheReduction = + forM_ (zip slugs_op_renamed slugs) $ \(slug_op_renamed, slug) -> + sLoopNest (slugShape slug) $ \vec_is -> do + comment "to reduce current chunk, first store our result in memory" $ do + forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> + copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is) + + forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) -> + when (primType $ paramType p) $ + copyDWIMFix arr [local_tid] (Var $ paramName p) [] + + sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. + groupReduce (sExt32 (kernelGroupSize constants)) slug_op_renamed (slugArrs slug) + + sOp $ Imp.Barrier Imp.FenceLocal + + sComment "first thread saves the result in accumulator" $ + sWhen (local_tid .==. 0) $ + forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) -> + copyDWIMFix acc (acc_is ++ vec_is) (Var $ paramName p) [] + + -- If this is a non-commutative reduction, each thread must run the + -- loop the same number of iterations, because we will be performing + -- a group-wide reduction in there. + let comm = slugsComm slugs + (bound, check_bounds) = + case comm of + Commutative -> (tvExp chunk_size, id) + Noncommutative -> + ( Imp.unCount elems_per_thread, + sWhen (tvExp gtid .<. Imp.unCount num_elements) + ) + + sFor "i" bound $ \i -> do + gtid + <-- case comm of + Commutative -> + global_tid + Imp.vi64 threads_per_segment * i + Noncommutative -> + let index_in_segment = global_tid `quot` kernelGroupSize constants + in sExt64 local_tid + + (index_in_segment * Imp.unCount elems_per_thread + i) + * kernelGroupSize constants + + check_bounds $ + sComment "apply map function" $ + body $ \all_red_res -> do + let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res + + forM_ (zip slugs slugs_res) $ \(slug, red_res) -> + sLoopNest (slugShape slug) $ \vec_is -> do + sComment "load accumulator" $ + forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> + copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is) + sComment "load new values" $ + forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) -> + copyDWIMFix (paramName p) [] res (res_is ++ vec_is) + sComment "apply reduction operator" $ + compileStms mempty (bodyStms $ slugBody slug) $ + sComment "store in accumulator" $ + forM_ + ( zip + (slugAccs slug) + (bodyResult $ slugBody slug) + ) + $ \((acc, acc_is), se) -> + copyDWIMFix acc (acc_is ++ vec_is) se [] + + case comm of + Noncommutative -> do + doTheReduction + sComment "first thread keeps accumulator; others reset to neutral element" $ do + let reset_to_neutral = + forM_ slugs $ \slug -> + forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> + sLoopNest (slugShape slug) $ \vec_is -> + copyDWIMFix acc (acc_is ++ vec_is) ne [] + sUnless (local_tid .==. 0) reset_to_neutral + _ -> return () + + return (slugs_op_renamed, doTheReduction) + +reductionStageOne :: + KernelConstants -> + [(VName, Imp.TExp Int64)] -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + Imp.TExp Int64 -> + Imp.Count Imp.Elements (Imp.TExp Int64) -> + VName -> + [SegBinOpSlug] -> + DoSegBody -> + InKernelGen [Lambda GPUMem] +reductionStageOne constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do + (slugs_op_renamed, doTheReduction) <- + reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body + + case slugsComm slugs of + Noncommutative -> + forM_ slugs $ \slug -> + forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> + copyDWIMFix (paramName p) [] (Var acc) acc_is + _ -> doTheReduction + + return slugs_op_renamed + +reductionStageTwo :: + KernelConstants -> + [PatElem GPUMem] -> + Imp.TExp Int32 -> + Imp.TExp Int32 -> + [Imp.TExp Int64] -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + SegBinOpSlug -> + [LParam GPUMem] -> + [LParam GPUMem] -> + Lambda GPUMem -> + [SubExp] -> + Imp.TExp Int32 -> + VName -> + Imp.TExp Int32 -> + VName -> + [VName] -> + [VName] -> + InKernelGen () +reductionStageTwo + constants + segred_pes + group_id + flat_segment_id + segment_gtids + first_group_for_segment + groups_per_segment + slug + red_x_params + red_y_params + red_op_renamed + nes + num_counters + counter + counter_i + sync_arr + group_res_arrs + red_arrs = do + let local_tid = kernelLocalThreadId constants + group_size = kernelGroupSize constants + old_counter <- dPrim "old_counter" int32 + (counter_mem, _, counter_offset) <- + fullyIndexArray + counter + [ sExt64 $ + counter_i * num_counters + + flat_segment_id `rem` num_counters + ] + comment "first thread in group saves group result to global memory" $ + sWhen (local_tid .==. 0) $ do + forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) -> + copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is + sOp $ Imp.MemFence Imp.FenceGlobal + -- Increment the counter, thus stating that our result is + -- available. + sOp $ + Imp.Atomic DefaultSpace $ + Imp.AtomicAdd + Int32 + (tvVar old_counter) + counter_mem + counter_offset + $ untyped (1 :: Imp.TExp Int32) + -- Now check if we were the last group to write our result. If + -- so, it is our responsibility to produce the final result. + sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1 + + sOp $ Imp.Barrier Imp.FenceGlobal + + is_last_group <- dPrim "is_last_group" Bool + copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0] + sWhen (tvExp is_last_group) $ do + -- The final group has written its result (and it was + -- us!), so read in all the group results and perform the + -- final stage of the reduction. But first, we reset the + -- counter so it is ready for next time. This is done + -- with an atomic to avoid warnings about write/write + -- races in oclgrind. + sWhen (local_tid .==. 0) $ + sOp $ + Imp.Atomic DefaultSpace $ + Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $ + untyped $ negate groups_per_segment + + sLoopNest (slugShape slug) $ \vec_is -> do + -- 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 + -- way that may induce non-coalesced accesses, but the total + -- number of accesses should be tiny here. + comment "read in the per-group-results" $ do + read_per_thread <- + dPrimVE "read_per_thread" $ + groups_per_segment `divUp` sExt64 group_size + + forM_ (zip red_x_params nes) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] + + sFor "i" read_per_thread $ \i -> do + group_res_id <- + dPrimVE "group_res_id" $ + sExt64 local_tid * read_per_thread + i + index_of_group_res <- + dPrimVE "index_of_group_res" $ + first_group_for_segment + group_res_id + + sWhen (group_res_id .<. groups_per_segment) $ do + forM_ (zip red_y_params group_res_arrs) $ + \(p, group_res_arr) -> + copyDWIMFix + (paramName p) + [] + (Var group_res_arr) + ([0, index_of_group_res] ++ vec_is) + + compileStms mempty (bodyStms $ slugBody slug) $ + forM_ (zip red_x_params (bodyResult $ slugBody slug)) $ \(p, se) -> + copyDWIMFix (paramName p) [] se [] + + forM_ (zip red_x_params red_arrs) $ \(p, arr) -> + when (primType $ paramType p) $ + copyDWIMFix arr [sExt64 local_tid] (Var $ paramName p) [] + + sOp $ Imp.Barrier Imp.FenceLocal + + sComment "reduce the per-group results" $ do + groupReduce (sExt32 group_size) red_op_renamed red_arrs + + sComment "and back to memory with the final result" $ + sWhen (local_tid .==. 0) $ + forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) -> + copyDWIMFix + (patElemName pe) + (segment_gtids ++ vec_is) + (Var $ paramName p) + [] diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs @@ -0,0 +1,68 @@ +-- | Code generation for 'SegScan'. Dispatches to either a +-- single-pass or two-pass implementation, depending on the nature of +-- the scan and the chosen abckend. +module Futhark.CodeGen.ImpGen.GPU.SegScan (compileSegScan) where + +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen hiding (compileProg) +import Futhark.CodeGen.ImpGen.GPU.Base +import qualified Futhark.CodeGen.ImpGen.GPU.SegScan.SinglePass as SinglePass +import qualified Futhark.CodeGen.ImpGen.GPU.SegScan.TwoPass as TwoPass +import Futhark.IR.GPUMem + +-- The single-pass scan does not support multiple operators, so jam +-- them together here. +combineScans :: [SegBinOp GPUMem] -> SegBinOp GPUMem +combineScans ops = + SegBinOp + { segBinOpComm = mconcat (map segBinOpComm ops), + segBinOpLambda = lam', + segBinOpNeutral = concatMap segBinOpNeutral ops, + segBinOpShape = mempty -- Assumed + } + where + lams = map segBinOpLambda ops + xParams lam = take (length (lambdaReturnType lam)) (lambdaParams lam) + yParams lam = drop (length (lambdaReturnType lam)) (lambdaParams lam) + lam' = + Lambda + { lambdaParams = concatMap xParams lams ++ concatMap yParams lams, + lambdaReturnType = concatMap lambdaReturnType lams, + lambdaBody = + Body + () + (mconcat (map (bodyStms . lambdaBody) lams)) + (concatMap (bodyResult . lambdaBody) lams) + } + +canBeSinglePass :: SegSpace -> [SegBinOp GPUMem] -> Maybe (SegBinOp GPUMem) +canBeSinglePass space ops + | [_] <- unSegSpace space, + all ok ops = + Just $ combineScans ops + | otherwise = + Nothing + where + ok op = + segBinOpShape op == mempty + && all primType (lambdaReturnType (segBinOpLambda op)) + +-- | Compile 'SegScan' instance to host-level code with calls to +-- various kernels. +compileSegScan :: + Pattern GPUMem -> + SegLevel -> + SegSpace -> + [SegBinOp GPUMem] -> + KernelBody GPUMem -> + CallKernelGen () +compileSegScan pat lvl space scans kbody = sWhen (0 .<. n) $ do + emit $ Imp.DebugPrint "\n# SegScan" Nothing + target <- hostTarget <$> askEnv + case target of + CUDA + | Just scan' <- canBeSinglePass space scans -> + SinglePass.compileSegScan pat lvl space scan' kbody + _ -> TwoPass.compileSegScan pat lvl space scans kbody + where + n = product $ map toInt64Exp $ segSpaceDims space diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs @@ -0,0 +1,565 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Code generation for segmented and non-segmented scans. Uses a +-- fast single-pass algorithm, but which only works on NVIDIA GPUs and +-- with some constraints on the operator. We use this when we can. +module Futhark.CodeGen.ImpGen.GPU.SegScan.SinglePass (compileSegScan) where + +import Control.Monad.Except +import Data.List (zip4) +import Data.Maybe +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.Transform.Rename +import Futhark.Util (takeLast) +import Futhark.Util.IntegralExp (IntegralExp (mod, rem), divUp, quot) +import Prelude hiding (mod, quot, rem) + +xParams, yParams :: SegBinOp GPUMem -> [LParam GPUMem] +xParams scan = + take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) +yParams scan = + drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) + +alignTo :: IntegralExp a => a -> a -> a +alignTo x a = (x `divUp` a) * a + +createLocalArrays :: + Count GroupSize SubExp -> + SubExp -> + [PrimType] -> + InKernelGen (VName, [VName], [VName], VName, VName, [VName]) +createLocalArrays (Count groupSize) m types = do + let groupSizeE = toInt64Exp groupSize + workSize = toInt64Exp m * groupSizeE + prefixArraysSize = + foldl (\acc tySize -> alignTo acc tySize + tySize * groupSizeE) 0 $ + map primByteSize types + maxTransposedArraySize = + foldl1 sMax64 $ map (\ty -> workSize * primByteSize ty) types + + warpSize :: Num a => a + warpSize = 32 + maxWarpExchangeSize = + foldl (\acc tySize -> alignTo acc tySize + tySize * fromInteger warpSize) 0 $ + map primByteSize types + maxLookbackSize = maxWarpExchangeSize + warpSize + size = Imp.bytes $ maxLookbackSize `sMax64` prefixArraysSize `sMax64` maxTransposedArraySize + + varTE :: TV Int64 -> TPrimExp Int64 VName + varTE = le64 . tvVar + + byteOffsets <- + mapM (fmap varTE . dPrimV "byte_offsets") $ + scanl (\off tySize -> alignTo off tySize + toInt64Exp groupSize * tySize) 0 $ + map primByteSize types + + warpByteOffsets <- + mapM (fmap varTE . dPrimV "warp_byte_offset") $ + scanl (\off tySize -> alignTo off tySize + warpSize * tySize) warpSize $ + map primByteSize types + + sComment "Allocate reused shared memeory" $ return () + + localMem <- sAlloc "local_mem" size (Space "local") + transposeArrayLength <- dPrimV "trans_arr_len" workSize + + sharedId <- sArrayInMem "shared_id" int32 (Shape [constant (1 :: Int32)]) localMem + sharedReadOffset <- sArrayInMem "shared_read_offset" int32 (Shape [constant (1 :: Int32)]) localMem + + transposedArrays <- + forM types $ \ty -> + sArrayInMem + "local_transpose_arr" + ty + (Shape [tvSize transposeArrayLength]) + localMem + + prefixArrays <- + forM (zip byteOffsets types) $ \(off, ty) -> do + let off' = off `quot` primByteSize ty + sArray + "local_prefix_arr" + ty + (Shape [groupSize]) + $ ArrayIn localMem $ IxFun.iotaOffset off' [pe64 groupSize] + + warpscan <- sArrayInMem "warpscan" int8 (Shape [constant (warpSize :: Int64)]) localMem + warpExchanges <- + forM (zip warpByteOffsets types) $ \(off, ty) -> do + let off' = off `quot` primByteSize ty + sArray + "warp_exchange" + ty + (Shape [constant (warpSize :: Int64)]) + $ ArrayIn localMem $ IxFun.iotaOffset off' [warpSize] + + return (sharedId, transposedArrays, prefixArrays, sharedReadOffset, warpscan, warpExchanges) + +-- | Compile 'SegScan' instance to host-level code with calls to a +-- single-pass kernel. +compileSegScan :: + Pattern GPUMem -> + SegLevel -> + SegSpace -> + SegBinOp GPUMem -> + KernelBody GPUMem -> + CallKernelGen () +compileSegScan pat lvl space scanOp kbody = do + let Pattern _ all_pes = pat + group_size = toInt64Exp <$> segGroupSize lvl + n = product $ map toInt64Exp $ segSpaceDims space + num_groups = Count (n `divUp` (unCount group_size * m)) + num_threads = unCount num_groups * unCount group_size + (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + segmented = length dims' > 1 + not_segmented_e = if segmented then false else true + segment_size = last dims' + scanOpNe = segBinOpNeutral scanOp + tys = map (\(Prim pt) -> pt) $ lambdaReturnType $ segBinOpLambda scanOp + + statusX, statusA, statusP :: Num a => a + statusX = 0 + statusA = 1 + statusP = 2 + makeStatusUsed flag used = tvExp flag .|. (tvExp used .<<. 2) + unmakeStatusUsed :: TV Int8 -> TV Int8 -> TV Int8 -> InKernelGen () + unmakeStatusUsed flagUsed flag used = do + used <-- tvExp flagUsed .>>. 2 + flag <-- tvExp flagUsed .&. 3 + + sumT :: Integer + maxT :: Integer + sumT = foldl (\bytes typ -> bytes + primByteSize typ) 0 tys + primByteSize' = max 4 . primByteSize + sumT' = foldl (\bytes typ -> bytes + primByteSize' typ) 0 tys `div` 4 + maxT = maximum (map primByteSize tys) + -- TODO: Make these constants dynamic by querying device + -- RTX 2080 Ti constants (CC 7.5) + k_reg = 64 + k_mem = 48 --12*4 + mem_constraint = max k_mem sumT `div` maxT + reg_constraint = (k_reg -1 - sumT') `div` (2 * sumT' + 3) + m :: Num a => a + m = fromIntegral $ max 1 $ min mem_constraint reg_constraint + + -- Allocate the shared memory for output component + numThreads <- dPrimV "numThreads" num_threads + numGroups <- dPrimV "numGroups" $ unCount num_groups + + globalId <- sStaticArray "id_counter" (Space "device") int32 $ Imp.ArrayZeros 1 + statusFlags <- sAllocArray "status_flags" int8 (Shape [tvSize numGroups]) (Space "device") + (aggregateArrays, incprefixArrays) <- + fmap unzip $ + forM tys $ \ty -> + (,) <$> sAllocArray "aggregates" ty (Shape [tvSize numGroups]) (Space "device") + <*> sAllocArray "incprefixes" ty (Shape [tvSize numGroups]) (Space "device") + + sReplicate statusFlags $ intConst Int8 statusX + + sKernelThread "segscan" num_groups group_size (segFlat space) $ do + constants <- kernelConstants <$> askEnv + + (sharedId, transposedArrays, prefixArrays, sharedReadOffset, warpscan, exchanges) <- + createLocalArrays (segGroupSize lvl) (intConst Int64 m) tys + + dynamicId <- dPrim "dynamic_id" int32 + sWhen (kernelLocalThreadId constants .==. 0) $ do + (globalIdMem, _, globalIdOff) <- fullyIndexArray globalId [0] + sOp $ + Imp.Atomic DefaultSpace $ + Imp.AtomicAdd + Int32 + (tvVar dynamicId) + globalIdMem + (Count $ unCount globalIdOff) + (untyped (1 :: Imp.TExp Int32)) + copyDWIMFix sharedId [0] (tvSize dynamicId) [] + + let localBarrier = Imp.Barrier Imp.FenceLocal + localFence = Imp.MemFence Imp.FenceLocal + globalFence = Imp.MemFence Imp.FenceGlobal + + sOp localBarrier + copyDWIMFix (tvVar dynamicId) [] (Var sharedId) [0] + sOp localBarrier + + blockOff <- + dPrimV "blockOff" $ + sExt64 (tvExp dynamicId) * m * kernelGroupSize constants + sgmIdx <- dPrimVE "sgm_idx" $ tvExp blockOff `mod` segment_size + boundary <- + dPrimVE "boundary" $ + sExt32 $ sMin64 (m * unCount group_size) (segment_size - sgmIdx) + segsize_compact <- + dPrimVE "segsize_compact" $ + sExt32 $ sMin64 (m * unCount group_size) segment_size + privateArrays <- + forM tys $ \ty -> + sAllocArray + "private" + ty + (Shape [intConst Int64 m]) + (ScalarSpace [intConst Int64 m] ty) + + sComment "Load and map" $ + sFor "i" m $ \i -> do + -- The map's input index + phys_tid <- + dPrimVE "phys_tid" $ + tvExp blockOff + sExt64 (kernelLocalThreadId constants) + + i * kernelGroupSize constants + dIndexSpace (zip gtids dims') phys_tid + -- Perform the map + let in_bounds = + compileStms mempty (kernelBodyStms kbody) $ do + let (all_scan_res, map_res) = splitAt (segBinOpResults [scanOp]) $ kernelBodyResult kbody + + -- Write map results to their global memory destinations + forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(dest, src) -> + copyDWIMFix (patElemName dest) (map Imp.vi64 gtids) (kernelResultSubExp src) [] + + -- Write to-scan results to private memory. + forM_ (zip privateArrays $ map kernelResultSubExp all_scan_res) $ \(dest, src) -> + copyDWIMFix dest [i] src [] + + out_of_bounds = + forM_ (zip privateArrays scanOpNe) $ \(dest, ne) -> + copyDWIMFix dest [i] ne [] + + sIf (phys_tid .<. n) in_bounds out_of_bounds + + sComment "Transpose scan inputs" $ do + forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do + sOp localBarrier + sFor "i" m $ \i -> do + sharedIdx <- + dPrimVE "sharedIdx" $ + sExt64 (kernelLocalThreadId constants) + + i * kernelGroupSize constants + copyDWIMFix trans [sharedIdx] (Var priv) [i] + sOp localBarrier + sFor "i" m $ \i -> do + sharedIdx <- dPrimV "sharedIdx" $ kernelLocalThreadId constants * m + i + copyDWIMFix priv [sExt64 i] (Var trans) [sExt64 $ tvExp sharedIdx] + sOp localBarrier + + sComment "Per thread scan" $ do + -- We don't need to touch the first element, so only m-1 + -- iterations here. + globalIdx <- + dPrimVE "gidx" $ + (kernelLocalThreadId constants * m) + 1 + sFor "i" (m -1) $ \i -> do + let xs = map paramName $ xParams scanOp + ys = map paramName $ yParams scanOp + -- determine if start of segment + new_sgm <- + if segmented + then dPrimVE "new_sgm" $ (globalIdx + sExt32 i - boundary) `mod` segsize_compact .==. 0 + else pure false + -- skip scan of first element in segment + sUnless new_sgm $ do + forM_ (zip privateArrays $ zip3 xs ys tys) $ \(src, (x, y, ty)) -> do + dPrim_ x ty + dPrim_ y ty + copyDWIMFix x [] (Var src) [i] + copyDWIMFix y [] (Var src) [i + 1] + + compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scanOp) $ + forM_ (zip privateArrays $ bodyResult $ lambdaBody $ segBinOpLambda scanOp) $ \(dest, res) -> + copyDWIMFix dest [i + 1] res [] + + sComment "Publish results in shared memory" $ do + forM_ (zip prefixArrays privateArrays) $ \(dest, src) -> + copyDWIMFix dest [sExt64 $ kernelLocalThreadId constants] (Var src) [m - 1] + sOp localBarrier + + let crossesSegment = do + guard segmented + Just $ \from to -> + let from' = (from + 1) * m - 1 + to' = (to + 1) * m - 1 + in (to' - from') .>. (to' + segsize_compact - boundary) `mod` segsize_compact + + scanOp' <- renameLambda $ segBinOpLambda scanOp + + accs <- mapM (dPrim "acc") tys + sComment "Scan results (with warp scan)" $ do + groupScan + crossesSegment + (tvExp numThreads) + (kernelGroupSize constants) + scanOp' + prefixArrays + + sOp localBarrier + + let firstThread acc prefixes = + copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelGroupSize constants) - 1] + notFirstThread acc prefixes = + copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelLocalThreadId constants) - 1] + sIf + (kernelLocalThreadId constants .==. 0) + (zipWithM_ firstThread accs prefixArrays) + (zipWithM_ notFirstThread accs prefixArrays) + + sOp localBarrier + + prefixes <- forM (zip scanOpNe tys) $ \(ne, ty) -> + dPrimV "prefix" $ TPrimExp $ toExp' ty ne + blockNewSgm <- dPrimVE "block_new_sgm" $ sgmIdx .==. 0 + sComment "Perform lookback" $ do + sWhen (blockNewSgm .&&. kernelLocalThreadId constants .==. 0) $ do + everythingVolatile $ + forM_ (zip incprefixArrays accs) $ \(incprefixArray, acc) -> + copyDWIMFix incprefixArray [tvExp dynamicId] (tvSize acc) [] + sOp globalFence + everythingVolatile $ + copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusP) [] + forM_ (zip scanOpNe accs) $ \(ne, acc) -> + copyDWIMFix (tvVar acc) [] ne [] + -- end sWhen + + let warpSize = kernelWaveSize constants + sWhen (bNot blockNewSgm .&&. kernelLocalThreadId constants .<. warpSize) $ do + sWhen (kernelLocalThreadId constants .==. 0) $ do + sIf + (not_segmented_e .||. boundary .==. sExt32 (unCount group_size * m)) + ( do + everythingVolatile $ + forM_ (zip aggregateArrays accs) $ \(aggregateArray, acc) -> + copyDWIMFix aggregateArray [tvExp dynamicId] (tvSize acc) [] + sOp globalFence + everythingVolatile $ + copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusA) [] + ) + ( do + everythingVolatile $ + forM_ (zip incprefixArrays accs) $ \(incprefixArray, acc) -> + copyDWIMFix incprefixArray [tvExp dynamicId] (tvSize acc) [] + sOp globalFence + everythingVolatile $ + copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusP) [] + ) + copyDWIMFix warpscan [0] (Var statusFlags) [tvExp dynamicId - 1] + -- sWhen + sOp localFence + + status <- dPrim "status" int8 :: InKernelGen (TV Int8) + copyDWIMFix (tvVar status) [] (Var warpscan) [0] + + sIf + (tvExp status .==. statusP) + ( sWhen (kernelLocalThreadId constants .==. 0) $ + everythingVolatile $ + forM_ (zip prefixes incprefixArrays) $ \(prefix, incprefixArray) -> + copyDWIMFix (tvVar prefix) [] (Var incprefixArray) [tvExp dynamicId - 1] + ) + ( do + readOffset <- + dPrimV "readOffset" $ + sExt32 $ tvExp dynamicId - sExt64 (kernelWaveSize constants) + let loopStop = warpSize * (-1) + sameSegment readIdx + | segmented = + let startIdx = sExt64 (tvExp readIdx + 1) * kernelGroupSize constants * m - 1 + in tvExp blockOff - startIdx .<=. sgmIdx + | otherwise = true + sWhile (tvExp readOffset .>. loopStop) $ do + readI <- dPrimV "read_i" $ tvExp readOffset + kernelLocalThreadId constants + aggrs <- forM (zip scanOpNe tys) $ \(ne, ty) -> + dPrimV "aggr" $ TPrimExp $ toExp' ty ne + flag <- dPrimV "flag" statusX + used <- dPrimV "used" 0 + everythingVolatile . sWhen (tvExp readI .>=. 0) $ do + sIf + (sameSegment readI) + ( do + copyDWIMFix (tvVar flag) [] (Var statusFlags) [sExt64 $ tvExp readI] + sIf + (tvExp flag .==. statusP) + ( forM_ (zip incprefixArrays aggrs) $ \(incprefix, aggr) -> + copyDWIMFix (tvVar aggr) [] (Var incprefix) [sExt64 $ tvExp readI] + ) + ( sWhen (tvExp flag .==. statusA) $ do + forM_ (zip aggrs aggregateArrays) $ \(aggr, aggregate) -> + copyDWIMFix (tvVar aggr) [] (Var aggregate) [sExt64 $ tvExp readI] + used <-- 1 + ) + ) + (copyDWIMFix (tvVar flag) [] (intConst Int8 statusP) []) + -- end sIf + -- end sWhen + forM_ (zip exchanges aggrs) $ \(exchange, aggr) -> + copyDWIMFix exchange [sExt64 $ kernelLocalThreadId constants] (tvSize aggr) [] + tmp <- dPrimV "tmp" $ makeStatusUsed flag used + copyDWIMFix warpscan [sExt64 $ kernelLocalThreadId constants] (tvSize tmp) [] + sOp localFence + + (warpscanMem, warpscanSpace, warpscanOff) <- + fullyIndexArray warpscan [sExt64 warpSize - 1] + flag <-- TPrimExp (Imp.index warpscanMem warpscanOff int8 warpscanSpace Imp.Volatile) + sWhen (kernelLocalThreadId constants .==. 0) $ do + -- TODO: This is a single-threaded reduce + sIf + (bNot $ tvExp flag .==. statusP) + ( do + scanOp'' <- renameLambda scanOp' + let (agg1s, agg2s) = splitAt (length tys) $ map paramName $ lambdaParams scanOp'' + + forM_ (zip3 agg1s scanOpNe tys) $ \(agg1, ne, ty) -> + dPrimV_ agg1 $ TPrimExp $ toExp' ty ne + zipWithM_ dPrim_ agg2s tys + + flag1 <- dPrimV "flag1" statusX + flag2 <- dPrim "flag2" int8 + used1 <- dPrimV "used1" 0 + used2 <- dPrim "used2" int8 + sFor "i" warpSize $ \i -> do + copyDWIMFix (tvVar flag2) [] (Var warpscan) [sExt64 i] + unmakeStatusUsed flag2 flag2 used2 + forM_ (zip agg2s exchanges) $ \(agg2, exchange) -> + copyDWIMFix agg2 [] (Var exchange) [sExt64 i] + sIf + (bNot $ tvExp flag2 .==. statusA) + ( do + flag1 <-- tvExp flag2 + used1 <-- tvExp used2 + forM_ (zip3 agg1s tys agg2s) $ \(agg1, ty, agg2) -> + agg1 <~~ toExp' ty (Var agg2) + ) + ( do + used1 <-- tvExp used1 + tvExp used2 + compileStms mempty (bodyStms $ lambdaBody scanOp'') $ + forM_ (zip3 agg1s tys $ bodyResult $ lambdaBody scanOp'') $ + \(agg1, ty, res) -> agg1 <~~ toExp' ty res + ) + flag <-- tvExp flag1 + used <-- tvExp used1 + forM_ (zip3 aggrs tys agg1s) $ \(aggr, ty, agg1) -> + tvVar aggr <~~ toExp' ty (Var agg1) + ) + -- else + ( forM_ (zip aggrs exchanges) $ \(aggr, exchange) -> + copyDWIMFix (tvVar aggr) [] (Var exchange) [sExt64 warpSize - 1] + ) + -- end sIf + sIf + (tvExp flag .==. statusP) + (readOffset <-- loopStop) + (readOffset <-- tvExp readOffset - zExt32 (tvExp used)) + copyDWIMFix sharedReadOffset [0] (tvSize readOffset) [] + scanOp''' <- renameLambda scanOp' + let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams scanOp''' + forM_ (zip xs aggrs) $ \(x, aggr) -> dPrimV_ x (tvExp aggr) + forM_ (zip ys prefixes) $ \(y, prefix) -> dPrimV_ y (tvExp prefix) + compileStms mempty (bodyStms $ lambdaBody scanOp''') $ + forM_ (zip3 prefixes tys $ bodyResult $ lambdaBody scanOp''') $ + \(prefix, ty, res) -> prefix <-- TPrimExp (toExp' ty res) + -- end sWhen + sOp localFence + copyDWIMFix (tvVar readOffset) [] (Var sharedReadOffset) [0] + ) + -- end sWhile + -- end sIf + sWhen (kernelLocalThreadId constants .==. 0) $ do + scanOp'''' <- renameLambda scanOp' + let xs = map paramName $ take (length tys) $ lambdaParams scanOp'''' + ys = map paramName $ drop (length tys) $ lambdaParams scanOp'''' + sWhen (boundary .==. sExt32 (unCount group_size * m)) $ do + forM_ (zip xs prefixes) $ \(x, prefix) -> dPrimV_ x $ tvExp prefix + forM_ (zip ys accs) $ \(y, acc) -> dPrimV_ y $ tvExp acc + compileStms mempty (bodyStms $ lambdaBody scanOp'''') $ + everythingVolatile $ + forM_ (zip incprefixArrays $ bodyResult $ lambdaBody scanOp'''') $ + \(incprefixArray, res) -> copyDWIMFix incprefixArray [tvExp dynamicId] res [] + sOp globalFence + everythingVolatile $ copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusP) [] + forM_ (zip exchanges prefixes) $ \(exchange, prefix) -> + copyDWIMFix exchange [0] (tvSize prefix) [] + forM_ (zip3 accs tys scanOpNe) $ \(acc, ty, ne) -> + tvVar acc <~~ toExp' ty ne + -- end sWhen + -- end sWhen + + sWhen (bNot $ tvExp dynamicId .==. 0) $ do + sOp localBarrier + forM_ (zip exchanges prefixes) $ \(exchange, prefix) -> + copyDWIMFix (tvVar prefix) [] (Var exchange) [0] + sOp localBarrier + -- end sWhen + -- end sComment + + scanOp''''' <- renameLambda scanOp' + scanOp'''''' <- renameLambda scanOp' + + sComment "Distribute results" $ do + let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams scanOp''''' + (xs', ys') = splitAt (length tys) $ map paramName $ lambdaParams scanOp'''''' + + forM_ (zip4 (zip prefixes accs) (zip xs xs') (zip ys ys') tys) $ + \((prefix, acc), (x, x'), (y, y'), ty) -> do + dPrim_ x ty + dPrim_ y ty + dPrimV_ x' $ tvExp prefix + dPrimV_ y' $ tvExp acc + + sIf + (kernelLocalThreadId constants * m .<. boundary .&&. bNot blockNewSgm) + ( compileStms mempty (bodyStms $ lambdaBody scanOp'''''') $ + forM_ (zip3 xs tys $ bodyResult $ lambdaBody scanOp'''''') $ + \(x, ty, res) -> x <~~ toExp' ty res + ) + (forM_ (zip xs accs) $ \(x, acc) -> copyDWIMFix x [] (Var $ tvVar acc) []) + -- calculate where previous thread stopped, to determine number of + -- elements left before new segment. + stop <- + dPrimVE "stopping_point" $ + segsize_compact - (kernelLocalThreadId constants * m - 1 + segsize_compact - boundary) `rem` segsize_compact + sFor "i" m $ \i -> do + sWhen (sExt32 i .<. stop - 1) $ do + forM_ (zip privateArrays ys) $ \(src, y) -> + -- only include prefix for the first segment part per thread + copyDWIMFix y [] (Var src) [i] + compileStms mempty (bodyStms $ lambdaBody scanOp''''') $ + forM_ (zip privateArrays $ bodyResult $ lambdaBody scanOp''''') $ + \(dest, res) -> + copyDWIMFix dest [i] res [] + + sComment "Transpose scan output" $ do + forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do + sOp localBarrier + sFor "i" m $ \i -> do + sharedIdx <- + dPrimV "sharedIdx" $ + sExt64 (kernelLocalThreadId constants * m) + i + copyDWIMFix trans [tvExp sharedIdx] (Var priv) [i] + sOp localBarrier + sFor "i" m $ \i -> do + sharedIdx <- + dPrimV "sharedIdx" $ + kernelLocalThreadId constants + + sExt32 (kernelGroupSize constants * i) + copyDWIMFix priv [i] (Var trans) [sExt64 $ tvExp sharedIdx] + sOp localBarrier + + sComment "Write block scan results to global memory" $ + sFor "i" m $ \i -> do + flat_idx <- + dPrimVE "flat_idx" $ + tvExp blockOff + kernelGroupSize constants * i + + sExt64 (kernelLocalThreadId constants) + dIndexSpace (zip gtids dims') flat_idx + sWhen (flat_idx .<. n) $ do + forM_ (zip (map patElemName all_pes) privateArrays) $ \(dest, src) -> + copyDWIMFix dest (map Imp.vi64 gtids) (Var src) [i] + + sComment "If this is the last block, reset the dynamicId" $ + sWhen (tvExp dynamicId .==. unCount num_groups - 1) $ + copyDWIMFix globalId [0] (constant (0 :: Int32)) [] diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs @@ -0,0 +1,506 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Code generation for segmented and non-segmented scans. Uses a +-- fairly inefficient two-pass algorithm, but can handle anything. +module Futhark.CodeGen.ImpGen.GPU.SegScan.TwoPass (compileSegScan) where + +import Control.Monad.Except +import Control.Monad.State +import Data.List (delete, find, foldl', zip4) +import Data.Maybe +import qualified Futhark.CodeGen.ImpCode.GPU as Imp +import Futhark.CodeGen.ImpGen +import Futhark.CodeGen.ImpGen.GPU.Base +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.Transform.Rename +import Futhark.Util (takeLast) +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +-- Aggressively try to reuse memory for different SegBinOps, because +-- we will run them sequentially after another. +makeLocalArrays :: + Count GroupSize SubExp -> + SubExp -> + [SegBinOp GPUMem] -> + InKernelGen [[VName]] +makeLocalArrays (Count group_size) num_threads scans = do + (arrs, mems_and_sizes) <- runStateT (mapM onScan scans) mempty + let maxSize sizes = Imp.bytes $ foldl' sMax64 1 $ map Imp.unCount sizes + forM_ mems_and_sizes $ \(sizes, mem) -> + sAlloc_ mem (maxSize sizes) (Space "local") + return arrs + where + onScan (SegBinOp _ scan_op nes _) = do + let (scan_x_params, _scan_y_params) = + splitAt (length nes) $ lambdaParams scan_op + (arrs, used_mems) <- fmap unzip $ + forM scan_x_params $ \p -> + case paramDec p of + MemArray pt shape _ (ArrayIn mem _) -> do + let shape' = Shape [num_threads] <> shape + arr <- + lift $ + sArray "scan_arr" pt shape' $ + ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' + return (arr, []) + _ -> do + let pt = elemType $ paramType p + shape = Shape [group_size] + (sizes, mem') <- getMem pt shape + arr <- lift $ sArrayInMem "scan_arr" pt shape mem' + return (arr, [(sizes, mem')]) + modify (<> concat used_mems) + return arrs + + getMem pt shape = do + let size = typeSize $ Array pt shape NoUniqueness + mems <- get + case (find ((size `elem`) . fst) mems, mems) of + (Just mem, _) -> do + modify $ delete mem + return mem + (Nothing, (size', mem) : mems') -> do + put mems' + return (size : size', mem) + (Nothing, []) -> do + mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local" + return ([size], mem) + +type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool) + +localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int64 +localArrayIndex constants t = + if primType t + then sExt64 (kernelLocalThreadId constants) + else sExt64 (kernelGlobalThreadId constants) + +barrierFor :: Lambda GPUMem -> (Bool, Imp.Fence, InKernelGen ()) +barrierFor scan_op = (array_scan, fence, sOp $ Imp.Barrier fence) + where + array_scan = not $ all primType $ lambdaReturnType scan_op + fence + | array_scan = Imp.FenceGlobal + | otherwise = Imp.FenceLocal + +xParams, yParams :: SegBinOp GPUMem -> [LParam GPUMem] +xParams scan = + take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) +yParams scan = + drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) + +writeToScanValues :: + [VName] -> + ([PatElem GPUMem], SegBinOp GPUMem, [KernelResult]) -> + InKernelGen () +writeToScanValues gtids (pes, scan, scan_res) + | shapeRank (segBinOpShape scan) > 0 = + forM_ (zip pes scan_res) $ \(pe, res) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids) + (kernelResultSubExp res) + [] + | otherwise = + forM_ (zip (yParams scan) scan_res) $ \(p, res) -> + copyDWIMFix (paramName p) [] (kernelResultSubExp res) [] + +readToScanValues :: + [Imp.TExp Int64] -> + [PatElem GPUMem] -> + SegBinOp GPUMem -> + InKernelGen () +readToScanValues is pes scan + | shapeRank (segBinOpShape scan) > 0 = + forM_ (zip (yParams scan) pes) $ \(p, pe) -> + copyDWIMFix (paramName p) [] (Var (patElemName pe)) is + | otherwise = + return () + +readCarries :: + Imp.TExp Int64 -> + [Imp.TExp Int64] -> + [Imp.TExp Int64] -> + [PatElem GPUMem] -> + SegBinOp GPUMem -> + InKernelGen () +readCarries chunk_offset dims' vec_is pes scan + | shapeRank (segBinOpShape scan) > 0 = do + ltid <- kernelLocalThreadId . kernelConstants <$> askEnv + -- We may have to reload the carries from the output of the + -- previous chunk. + sIf + (chunk_offset .>. 0 .&&. ltid .==. 0) + ( do + let is = unflattenIndex dims' $ chunk_offset - 1 + forM_ (zip (xParams scan) pes) $ \(p, pe) -> + copyDWIMFix (paramName p) [] (Var (patElemName pe)) (is ++ vec_is) + ) + ( forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] + ) + | otherwise = + return () + +-- | Produce partially scanned intervals; one per workgroup. +scanStage1 :: + Pattern GPUMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + SegSpace -> + [SegBinOp GPUMem] -> + KernelBody GPUMem -> + CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment) +scanStage1 (Pattern _ all_pes) num_groups group_size space scans kbody = do + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size' + + let (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + let num_elements = product dims' + elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads) + elems_per_group = unCount group_size' * elems_per_thread + + let crossesSegment = + case reverse dims' of + segment_size : _ : _ -> Just $ \from to -> + (to - from) .>. (to `rem` segment_size) + _ -> Nothing + + sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ do + constants <- kernelConstants <$> askEnv + all_local_arrs <- makeLocalArrays group_size (tvSize num_threads) scans + + -- The variables from scan_op will be used for the carry and such + -- in the big chunking loop. + forM_ scans $ \scan -> do + dScope Nothing $ scopeOfLParams $ lambdaParams $ segBinOpLambda scan + forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] + + sFor "j" elems_per_thread $ \j -> do + chunk_offset <- + dPrimV "chunk_offset" $ + sExt64 (kernelGroupSize constants) * j + + sExt64 (kernelGroupId constants) * elems_per_group + flat_idx <- + dPrimV "flat_idx" $ + tvExp chunk_offset + sExt64 (kernelLocalThreadId constants) + -- Construct segment indices. + zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx + + let per_scan_pes = segBinOpChunks scans all_pes + + in_bounds = + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' + + when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do + let (all_scan_res, map_res) = + splitAt (segBinOpResults scans) $ kernelBodyResult kbody + per_scan_res = + segBinOpChunks scans all_scan_res + + sComment "write to-scan values to parameters" $ + mapM_ (writeToScanValues gtids) $ + zip3 per_scan_pes scans per_scan_res + + sComment "write mapped values results to global memory" $ + forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids) + (kernelResultSubExp se) + [] + + sComment "threads in bounds read input" $ + sWhen in_bounds when_in_bounds + + forM_ (zip3 per_scan_pes scans all_local_arrs) $ + \(pes, scan@(SegBinOp _ scan_op nes vec_shape), local_arrs) -> + sComment "do one intra-group scan operation" $ do + let rets = lambdaReturnType scan_op + scan_x_params = xParams scan + (array_scan, fence, barrier) = barrierFor scan_op + + when array_scan barrier + + sLoopNest vec_shape $ \vec_is -> do + sComment "maybe restore some to-scan values to parameters, or read neutral" $ + sIf + in_bounds + ( do + readToScanValues (map Imp.vi64 gtids ++ vec_is) pes scan + readCarries (tvExp chunk_offset) dims' vec_is pes scan + ) + ( forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> + copyDWIMFix (paramName p) [] ne [] + ) + + sComment "combine with carry and write to local memory" $ + compileStms mempty (bodyStms $ lambdaBody scan_op) $ + forM_ (zip3 rets local_arrs (bodyResult $ lambdaBody scan_op)) $ + \(t, arr, se) -> + copyDWIMFix arr [localArrayIndex constants t] se [] + + let crossesSegment' = do + f <- crossesSegment + Just $ \from to -> + let from' = sExt64 from + tvExp chunk_offset + to' = sExt64 to + tvExp chunk_offset + in f from' to' + + sOp $ Imp.ErrorSync fence + + -- We need to avoid parameter name clashes. + scan_op_renamed <- renameLambda scan_op + groupScan + crossesSegment' + (sExt64 $ tvExp num_threads) + (sExt64 $ kernelGroupSize constants) + scan_op_renamed + local_arrs + + sComment "threads in bounds write partial scan result" $ + sWhen in_bounds $ + forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids ++ vec_is) + (Var arr) + [localArrayIndex constants t] + + barrier + + let load_carry = + forM_ (zip local_arrs scan_x_params) $ \(arr, p) -> + copyDWIMFix + (paramName p) + [] + (Var arr) + [ if primType $ paramType p + then sExt64 (kernelGroupSize constants) - 1 + else + (sExt64 (kernelGroupId constants) + 1) + * sExt64 (kernelGroupSize constants) - 1 + ] + load_neutral = + forM_ (zip nes scan_x_params) $ \(ne, p) -> + copyDWIMFix (paramName p) [] ne [] + + sComment "first thread reads last element as carry-in for next iteration" $ do + crosses_segment <- dPrimVE "crosses_segment" $ + case crossesSegment of + Nothing -> false + Just f -> + f + ( tvExp chunk_offset + + sExt64 (kernelGroupSize constants) -1 + ) + ( tvExp chunk_offset + + sExt64 (kernelGroupSize constants) + ) + should_load_carry <- + dPrimVE "should_load_carry" $ + kernelLocalThreadId constants .==. 0 .&&. bNot crosses_segment + sWhen should_load_carry load_carry + when array_scan barrier + sUnless should_load_carry load_neutral + + barrier + + return (num_threads, elems_per_group, crossesSegment) + +scanStage2 :: + Pattern GPUMem -> + TV Int32 -> + Imp.TExp Int64 -> + Count NumGroups SubExp -> + CrossesSegment -> + SegSpace -> + [SegBinOp GPUMem] -> + CallKernelGen () +scanStage2 (Pattern _ all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do + let (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + + -- Our group size is the number of groups for the stage 1 kernel. + let group_size = Count $ unCount num_groups + group_size' = fmap toInt64Exp group_size + + let crossesSegment' = do + f <- crossesSegment + Just $ \from to -> + f + ((sExt64 from + 1) * elems_per_group - 1) + ((sExt64 to + 1) * elems_per_group - 1) + + sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ do + constants <- kernelConstants <$> askEnv + per_scan_local_arrs <- makeLocalArrays group_size (tvSize stage1_num_threads) scans + let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans + per_scan_pes = segBinOpChunks scans all_pes + + flat_idx <- + dPrimV "flat_idx" $ + (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1 + -- Construct segment indices. + zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx + + forM_ (zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $ + \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) -> + sLoopNest vec_shape $ \vec_is -> do + let glob_is = map Imp.vi64 gtids ++ vec_is + + in_bounds = + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' + + when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) -> + copyDWIMFix + arr + [localArrayIndex constants t] + (Var $ patElemName pe) + glob_is + + when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) -> + copyDWIMFix arr [localArrayIndex constants t] ne [] + (_, _, barrier) = + barrierFor scan_op + + sComment "threads in bound read carries; others get neutral element" $ + sIf in_bounds when_in_bounds when_out_of_bounds + + barrier + + groupScan + crossesSegment' + (sExt64 $ tvExp stage1_num_threads) + (sExt64 $ kernelGroupSize constants) + scan_op + local_arrs + + sComment "threads in bounds write scanned carries" $ + sWhen in_bounds $ + forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> + copyDWIMFix + (patElemName pe) + glob_is + (Var arr) + [localArrayIndex constants t] + +scanStage3 :: + Pattern GPUMem -> + Count NumGroups SubExp -> + Count GroupSize SubExp -> + Imp.TExp Int64 -> + CrossesSegment -> + SegSpace -> + [SegBinOp GPUMem] -> + CallKernelGen () +scanStage3 (Pattern _ all_pes) num_groups group_size elems_per_group crossesSegment space scans = do + let num_groups' = fmap toInt64Exp num_groups + group_size' = fmap toInt64Exp group_size + (gtids, dims) = unzip $ unSegSpace space + dims' = map toInt64Exp dims + required_groups <- + dPrimVE "required_groups" $ + sExt32 $ product dims' `divUp` sExt64 (unCount group_size') + + sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $ + virtualiseGroups SegVirt required_groups $ \virt_group_id -> do + constants <- kernelConstants <$> askEnv + + -- Compute our logical index. + flat_idx <- + dPrimVE "flat_idx" $ + sExt64 virt_group_id * sExt64 (unCount group_size') + + sExt64 (kernelLocalThreadId constants) + zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx + + -- Figure out which group this element was originally in. + orig_group <- dPrimV "orig_group" $ flat_idx `quot` elems_per_group + -- Then the index of the carry-in of the preceding group. + carry_in_flat_idx <- + dPrimV "carry_in_flat_idx" $ + tvExp orig_group * elems_per_group - 1 + -- Figure out the logical index of the carry-in. + let carry_in_idx = unflattenIndex dims' $ tvExp carry_in_flat_idx + + -- Apply the carry if we are not in the scan results for the first + -- group, and are not the last element in such a group (because + -- then the carry was updated in stage 2), and we are not crossing + -- a segment boundary. + let in_bounds = + foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' + crosses_segment = + fromMaybe false $ + crossesSegment + <*> pure (tvExp carry_in_flat_idx) + <*> pure flat_idx + is_a_carry = flat_idx .==. (tvExp orig_group + 1) * elems_per_group - 1 + no_carry_in = tvExp orig_group .==. 0 .||. is_a_carry .||. crosses_segment + + let per_scan_pes = segBinOpChunks scans all_pes + sWhen in_bounds $ + sUnless no_carry_in $ + forM_ (zip per_scan_pes scans) $ + \(pes, SegBinOp _ scan_op nes vec_shape) -> do + dScope Nothing $ scopeOfLParams $ lambdaParams scan_op + let (scan_x_params, scan_y_params) = + splitAt (length nes) $ lambdaParams scan_op + + sLoopNest vec_shape $ \vec_is -> do + forM_ (zip scan_x_params pes) $ \(p, pe) -> + copyDWIMFix + (paramName p) + [] + (Var $ patElemName pe) + (carry_in_idx ++ vec_is) + + forM_ (zip scan_y_params pes) $ \(p, pe) -> + copyDWIMFix + (paramName p) + [] + (Var $ patElemName pe) + (map Imp.vi64 gtids ++ vec_is) + + compileBody' scan_x_params $ lambdaBody scan_op + + forM_ (zip scan_x_params pes) $ \(p, pe) -> + copyDWIMFix + (patElemName pe) + (map Imp.vi64 gtids ++ vec_is) + (Var $ paramName p) + [] + +-- | Compile 'SegScan' instance to host-level code with calls to +-- various kernels. +compileSegScan :: + Pattern GPUMem -> + SegLevel -> + SegSpace -> + [SegBinOp GPUMem] -> + KernelBody GPUMem -> + CallKernelGen () +compileSegScan pat lvl space scans kbody = do + -- Since stage 2 involves a group size equal to the number of groups + -- used for stage 1, we have to cap this number to the maximum group + -- size. + stage1_max_num_groups <- dPrim "stage1_max_num_groups" int64 + sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup + + stage1_num_groups <- + fmap (Imp.Count . tvSize) $ + dPrimV "stage1_num_groups" $ + sMin64 (tvExp stage1_max_num_groups) $ + toInt64Exp $ Imp.unCount $ segNumGroups lvl + + (stage1_num_threads, elems_per_group, crossesSegment) <- + scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody + + emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group + + scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans + scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans diff --git a/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs @@ -0,0 +1,888 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} + +-- | This module defines a translation from imperative code with +-- kernels to imperative code with OpenCL calls. +module Futhark.CodeGen.ImpGen.GPU.ToOpenCL + ( kernelsToOpenCL, + kernelsToCUDA, + ) +where + +import Control.Monad.Identity +import Control.Monad.Reader +import Control.Monad.State +import Data.FileEmbed +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Set as S +import qualified Futhark.CodeGen.Backends.GenericC as GC +import Futhark.CodeGen.Backends.SimpleRep +import Futhark.CodeGen.ImpCode.GPU hiding (Program) +import qualified Futhark.CodeGen.ImpCode.GPU as ImpGPU +import Futhark.CodeGen.ImpCode.OpenCL hiding (Program) +import qualified Futhark.CodeGen.ImpCode.OpenCL as ImpOpenCL +import Futhark.Error (compilerLimitationS) +import Futhark.IR.Prop (isBuiltInFunction) +import Futhark.MonadFreshNames +import Futhark.Util (zEncodeString) +import Futhark.Util.Pretty (prettyOneLine) +import qualified Language.C.Quote.CUDA as CUDAC +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C + +kernelsToCUDA, kernelsToOpenCL :: ImpGPU.Program -> ImpOpenCL.Program +kernelsToCUDA = translateGPU TargetCUDA +kernelsToOpenCL = translateGPU TargetOpenCL + +-- | Translate a kernels-program to an OpenCL-program. +translateGPU :: + KernelTarget -> + ImpGPU.Program -> + ImpOpenCL.Program +translateGPU target prog = + let ( prog', + ToOpenCL kernels device_funs used_types sizes failures + ) = + (`runState` initialOpenCL) . (`runReaderT` defFuns prog) $ do + let ImpGPU.Definitions + (ImpGPU.Constants ps consts) + (ImpGPU.Functions funs) = prog + consts' <- traverse (onHostOp target) consts + funs' <- forM funs $ \(fname, fun) -> + (fname,) <$> traverse (onHostOp target) fun + + return $ + ImpOpenCL.Definitions + (ImpOpenCL.Constants ps consts') + (ImpOpenCL.Functions funs') + + (device_prototypes, device_defs) = unzip $ M.elems device_funs + kernels' = M.map fst kernels + opencl_code = openClCode $ map snd $ M.elems kernels + + opencl_prelude = + unlines + [ pretty $ genPrelude target used_types, + unlines $ map pretty device_prototypes, + unlines $ map pretty device_defs + ] + in ImpOpenCL.Program + opencl_code + opencl_prelude + kernels' + (S.toList used_types) + (cleanSizes sizes) + failures + prog' + where + genPrelude TargetOpenCL = genOpenClPrelude + genPrelude TargetCUDA = const genCUDAPrelude + +-- | Due to simplifications after kernel extraction, some threshold +-- parameters may contain KernelPaths that reference threshold +-- parameters that no longer exist. We remove these here. +cleanSizes :: M.Map Name SizeClass -> M.Map Name SizeClass +cleanSizes m = M.map clean m + where + known = M.keys m + clean (SizeThreshold path def) = + SizeThreshold (filter ((`elem` known) . fst) path) def + clean s = s + +pointerQuals :: Monad m => String -> m [C.TypeQual] +pointerQuals "global" = return [C.ctyquals|__global|] +pointerQuals "local" = return [C.ctyquals|__local|] +pointerQuals "private" = return [C.ctyquals|__private|] +pointerQuals "constant" = return [C.ctyquals|__constant|] +pointerQuals "write_only" = return [C.ctyquals|__write_only|] +pointerQuals "read_only" = return [C.ctyquals|__read_only|] +pointerQuals "kernel" = return [C.ctyquals|__kernel|] +pointerQuals s = error $ "'" ++ s ++ "' is not an OpenCL kernel address space." + +-- In-kernel name and per-workgroup size in bytes. +type LocalMemoryUse = (VName, Count Bytes Exp) + +data KernelState = KernelState + { kernelLocalMemory :: [LocalMemoryUse], + kernelFailures :: [FailureMsg], + kernelNextSync :: Int, + -- | Has a potential failure occurred sine the last + -- ErrorSync? + kernelSyncPending :: Bool, + kernelHasBarriers :: Bool + } + +newKernelState :: [FailureMsg] -> KernelState +newKernelState failures = KernelState mempty failures 0 False False + +errorLabel :: KernelState -> String +errorLabel = ("error_" ++) . show . kernelNextSync + +data ToOpenCL = ToOpenCL + { clGPU :: M.Map KernelName (KernelSafety, C.Func), + clDevFuns :: M.Map Name (C.Definition, C.Func), + clUsedTypes :: S.Set PrimType, + clSizes :: M.Map Name SizeClass, + clFailures :: [FailureMsg] + } + +initialOpenCL :: ToOpenCL +initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty + +type AllFunctions = ImpGPU.Functions ImpGPU.HostOp + +lookupFunction :: Name -> AllFunctions -> Maybe ImpGPU.Function +lookupFunction fname (ImpGPU.Functions fs) = lookup fname fs + +type OnKernelM = ReaderT AllFunctions (State ToOpenCL) + +addSize :: Name -> SizeClass -> OnKernelM () +addSize key sclass = + modify $ \s -> s {clSizes = M.insert key sclass $ clSizes s} + +onHostOp :: KernelTarget -> HostOp -> OnKernelM OpenCL +onHostOp target (CallKernel k) = onKernel target k +onHostOp _ (ImpGPU.GetSize v key size_class) = do + addSize key size_class + return $ ImpOpenCL.GetSize v key +onHostOp _ (ImpGPU.CmpSizeLe v key size_class x) = do + addSize key size_class + return $ ImpOpenCL.CmpSizeLe v key x +onHostOp _ (ImpGPU.GetSizeMax v size_class) = + return $ ImpOpenCL.GetSizeMax v size_class + +genGPUCode :: + OpsMode -> + KernelCode -> + [FailureMsg] -> + GC.CompilerM KernelOp KernelState a -> + (a, GC.CompilerState KernelState) +genGPUCode mode body failures = + GC.runCompilerM + (inKernelOperations mode body) + blankNameSource + (newKernelState failures) + +-- Compilation of a device function that is not not invoked from the +-- host, but is invoked by (perhaps multiple) kernels. +generateDeviceFun :: Name -> ImpGPU.Function -> OnKernelM () +generateDeviceFun fname 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 + when (any memParam $ functionInput host_func) bad + + 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 + + modify $ \s -> + s + { clUsedTypes = typesInCode (functionBody device_func) <> clUsedTypes s, + clDevFuns = M.insert fname func $ clDevFuns s, + clFailures = kernelFailures kstate + } + + -- Important to do this after the 'modify' call, so we propagate the + -- right clFailures. + void $ ensureDeviceFuns $ functionBody device_func + where + toDevice :: HostOp -> KernelOp + toDevice _ = bad + + memParam MemParam {} = True + memParam ScalarParam {} = False + + bad = compilerLimitationS "Cannot generate GPU functions that use arrays." + +-- Ensure that this device function is available, but don't regenerate +-- it if it already exists. +ensureDeviceFun :: Name -> ImpGPU.Function -> OnKernelM () +ensureDeviceFun fname host_func = do + exists <- gets $ M.member fname . clDevFuns + unless exists $ generateDeviceFun fname host_func + +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 func -> do + ensureDeviceFun fname func + return $ Just fname + Nothing -> return Nothing + +onKernel :: KernelTarget -> Kernel -> OnKernelM OpenCL +onKernel target kernel = do + called <- ensureDeviceFuns $ kernelBody kernel + + -- Crucial that this is done after 'ensureDeviceFuns', as the device + -- functions may themselves define failure points. + failures <- gets clFailures + + let (kernel_body, cstate) = + genGPUCode KernelMode (kernelBody kernel) failures $ + GC.blockScope $ GC.compileCode $ kernelBody kernel + kstate = GC.compUserState cstate + + use_params = mapMaybe useAsParam $ kernelUses kernel + + (local_memory_args, local_memory_params, local_memory_init) = + unzip3 $ + flip evalState (blankNameSource :: VNameSource) $ + mapM (prepareLocalMemory target) $ kernelLocalMemory kstate + + -- CUDA has very strict restrictions on the number of blocks + -- permitted along the 'y' and 'z' dimensions of the grid + -- (1<<16). To work around this, we are going to dynamically + -- permute the block dimensions to move the largest one to the + -- 'x' dimension, which has a higher limit (1<<31). This means + -- we need to extend the kernel with extra parameters that + -- contain information about this permutation, but we only do + -- this for multidimensional kernels (at the time of this + -- writing, only transposes). The corresponding arguments are + -- added automatically in CCUDA.hs. + (perm_params, block_dim_init) = + case (target, num_groups) of + (TargetCUDA, [_, _, _]) -> + ( [ [C.cparam|const int block_dim0|], + [C.cparam|const int block_dim1|], + [C.cparam|const int block_dim2|] + ], + mempty + ) + _ -> + ( mempty, + [ [C.citem|const int block_dim0 = 0;|], + [C.citem|const int block_dim1 = 1;|], + [C.citem|const int block_dim2 = 2;|] + ] + ) + + (const_defs, const_undefs) = unzip $ mapMaybe constDef $ kernelUses kernel + + let (safety, error_init) + -- We conservatively assume that any called function can fail. + | not $ null called = + (SafetyFull, []) + | length (kernelFailures kstate) == length failures = + if kernelFailureTolerant kernel + then (SafetyNone, []) + else -- No possible failures in this kernel, so if we make + -- it past an initial check, then we are good to go. + + ( SafetyCheap, + [C.citems|if (*global_failure >= 0) { return; }|] + ) + | otherwise = + if not (kernelHasBarriers kstate) + then + ( SafetyFull, + [C.citems|if (*global_failure >= 0) { return; }|] + ) + else + ( SafetyFull, + [C.citems| + volatile __local bool local_failure; + if (failure_is_an_option) { + int failed = *global_failure >= 0; + if (failed) { + return; + } + } + // All threads write this value - it looks like CUDA has a compiler bug otherwise. + local_failure = false; + barrier(CLK_LOCAL_MEM_FENCE); + |] + ) + + failure_params = + [ [C.cparam|__global int *global_failure|], + [C.cparam|int failure_is_an_option|], + [C.cparam|__global typename int64_t *global_failure_args|] + ] + + params = + perm_params + ++ take (numFailureParams safety) failure_params + ++ catMaybes local_memory_params + ++ use_params + + kernel_fun = + [C.cfun|__kernel void $id:name ($params:params) { + $items:const_defs + $items:block_dim_init + $items:local_memory_init + $items:error_init + $items:kernel_body + + $id:(errorLabel kstate): return; + + $items:const_undefs + }|] + modify $ \s -> + s + { clGPU = M.insert name (safety, kernel_fun) $ clGPU s, + clUsedTypes = typesInKernel kernel <> clUsedTypes s, + clFailures = kernelFailures kstate + } + + -- The argument corresponding to the global_failure parameters is + -- added automatically later. + let args = + catMaybes local_memory_args + ++ kernelArgs kernel + + return $ LaunchKernel safety name args num_groups group_size + where + name = kernelName kernel + num_groups = kernelNumGroups kernel + group_size = kernelGroupSize kernel + + prepareLocalMemory TargetOpenCL (mem, size) = do + mem_aligned <- newVName $ baseString mem ++ "_aligned" + return + ( Just $ SharedMemoryKArg size, + Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|], + [C.citem|__local volatile char* restrict $id:mem = (__local volatile char*)$id:mem_aligned;|] + ) + prepareLocalMemory TargetCUDA (mem, size) = do + param <- newVName $ baseString mem ++ "_offset" + return + ( Just $ SharedMemoryKArg size, + Just [C.cparam|uint $id:param|], + [C.citem|volatile char *$id:mem = &shared_mem[$id:param];|] + ) + +useAsParam :: KernelUse -> Maybe C.Param +useAsParam (ScalarUse name bt) = + let ctp = case bt of + -- OpenCL does not permit bool as a kernel parameter type. + Bool -> [C.cty|unsigned char|] + Unit -> [C.cty|unsigned char|] + _ -> GC.primTypeToCType bt + in Just [C.cparam|$ty:ctp $id:name|] +useAsParam (MemoryUse name) = + Just [C.cparam|__global unsigned char *$id:name|] +useAsParam ConstUse {} = + Nothing + +-- Constants are #defined as macros. Since a constant name in one +-- kernel might potentially (although unlikely) also be used for +-- something else in another kernel, we #undef them after the kernel. +constDef :: KernelUse -> Maybe (C.BlockItem, C.BlockItem) +constDef (ConstUse v e) = + Just + ( [C.citem|$escstm:def|], + [C.citem|$escstm:undef|] + ) + where + e' = compilePrimExp e + def = "#define " ++ pretty (C.toIdent v mempty) ++ " (" ++ prettyOneLine e' ++ ")" + undef = "#undef " ++ pretty (C.toIdent v mempty) +constDef _ = Nothing + +openClCode :: [C.Func] -> String +openClCode kernels = + pretty [C.cunit|$edecls:funcs|] + where + funcs = + [ [C.cedecl|$func:kernel_func|] + | kernel_func <- kernels + ] + +atomicsDefs :: String +atomicsDefs = $(embedStringFile "rts/c/atomics.h") + +genOpenClPrelude :: S.Set PrimType -> [C.Definition] +genOpenClPrelude ts = + -- Clang-based OpenCL implementations need this for 'static' to work. + [ [C.cedecl|$esc:("#ifdef cl_clang_storage_class_specifiers")|], + [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_clang_storage_class_specifiers : enable")|], + [C.cedecl|$esc:("#endif")|], + [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_byte_addressable_store : enable")|] + ] + ++ concat + [ [C.cunit|$esc:("#pragma OPENCL EXTENSION cl_khr_fp64 : enable") + $esc:("#define FUTHARK_F64_ENABLED")|] + | uses_float64 + ] + ++ [C.cunit| +/* Some OpenCL programs dislike empty progams, or programs with no kernels. + * Declare a dummy kernel to ensure they remain our friends. */ +__kernel void dummy_kernel(__global unsigned char *dummy, int n) +{ + const int thread_gid = get_global_id(0); + if (thread_gid >= n) return; +} + +$esc:("#pragma OPENCL EXTENSION cl_khr_int64_base_atomics : enable") +$esc:("#pragma OPENCL EXTENSION cl_khr_int64_extended_atomics : enable") + +typedef char int8_t; +typedef short int16_t; +typedef int int32_t; +typedef long int64_t; + +typedef uchar uint8_t; +typedef ushort uint16_t; +typedef uint uint32_t; +typedef ulong uint64_t; + +// NVIDIAs OpenCL does not create device-wide memory fences (see #734), so we +// use inline assembly if we detect we are on an NVIDIA GPU. +$esc:("#ifdef cl_nv_pragma_unroll") +static inline void mem_fence_global() { + asm("membar.gl;"); +} +$esc:("#else") +static inline void mem_fence_global() { + mem_fence(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE); +} +$esc:("#endif") +static inline void mem_fence_local() { + mem_fence(CLK_LOCAL_MEM_FENCE); +} +|] + ++ cIntOps + ++ cFloat32Ops + ++ cFloat32Funs + ++ (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else []) + ++ [[C.cedecl|$esc:atomicsDefs|]] + where + uses_float64 = FloatType Float64 `S.member` ts + +genCUDAPrelude :: [C.Definition] +genCUDAPrelude = + cudafy ++ ops + where + ops = + cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ cFloat64Ops + ++ cFloat64Funs + ++ cFloatConvOps + ++ [[C.cedecl|$esc:atomicsDefs|]] + cudafy = + [CUDAC.cunit| +$esc:("#define FUTHARK_CUDA") +$esc:("#define FUTHARK_F64_ENABLED") + +typedef char int8_t; +typedef short int16_t; +typedef int int32_t; +typedef long long int64_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef uint8_t uchar; +typedef uint16_t ushort; +typedef uint32_t uint; +typedef uint64_t ulong; +$esc:("#define __kernel extern \"C\" __global__ __launch_bounds__(MAX_THREADS_PER_BLOCK)") +$esc:("#define __global") +$esc:("#define __local") +$esc:("#define __private") +$esc:("#define __constant") +$esc:("#define __write_only") +$esc:("#define __read_only") + +static inline int get_group_id_fn(int block_dim0, int block_dim1, int block_dim2, int d) +{ + switch (d) { + case 0: d = block_dim0; break; + case 1: d = block_dim1; break; + case 2: d = block_dim2; break; + } + switch (d) { + case 0: return blockIdx.x; + case 1: return blockIdx.y; + case 2: return blockIdx.z; + default: return 0; + } +} +$esc:("#define get_group_id(d) get_group_id_fn(block_dim0, block_dim1, block_dim2, d)") + +static inline int get_num_groups_fn(int block_dim0, int block_dim1, int block_dim2, int d) +{ + switch (d) { + case 0: d = block_dim0; break; + case 1: d = block_dim1; break; + case 2: d = block_dim2; break; + } + switch(d) { + case 0: return gridDim.x; + case 1: return gridDim.y; + case 2: return gridDim.z; + default: return 0; + } +} +$esc:("#define get_num_groups(d) get_num_groups_fn(block_dim0, block_dim1, block_dim2, d)") + +static inline int get_local_id(int d) +{ + switch (d) { + case 0: return threadIdx.x; + case 1: return threadIdx.y; + case 2: return threadIdx.z; + default: return 0; + } +} + +static inline int get_local_size(int d) +{ + switch (d) { + case 0: return blockDim.x; + case 1: return blockDim.y; + case 2: return blockDim.z; + default: return 0; + } +} + +static inline int get_global_id_fn(int block_dim0, int block_dim1, int block_dim2, int d) +{ + return get_group_id(d) * get_local_size(d) + get_local_id(d); +} +$esc:("#define get_global_id(d) get_global_id_fn(block_dim0, block_dim1, block_dim2, d)") + +static inline int get_global_size(int block_dim0, int block_dim1, int block_dim2, int d) +{ + return get_num_groups(d) * get_local_size(d); +} + +$esc:("#define CLK_LOCAL_MEM_FENCE 1") +$esc:("#define CLK_GLOBAL_MEM_FENCE 2") +static inline void barrier(int x) +{ + __syncthreads(); +} +static inline void mem_fence_local() { + __threadfence_block(); +} +static inline void mem_fence_global() { + __threadfence(); +} + +$esc:("#define NAN (0.0/0.0)") +$esc:("#define INFINITY (1.0/0.0)") +extern volatile __shared__ char shared_mem[]; +|] + +compilePrimExp :: PrimExp KernelConst -> C.Exp +compilePrimExp e = runIdentity $ GC.compilePrimExp compileKernelConst e + where + compileKernelConst (SizeConst key) = + return [C.cexp|$id:(zEncodeString (pretty key))|] + +kernelArgs :: Kernel -> [KernelArg] +kernelArgs = mapMaybe useToArg . kernelUses + where + useToArg (MemoryUse mem) = Just $ MemKArg mem + useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt + useToArg ConstUse {} = Nothing + +nextErrorLabel :: GC.CompilerM KernelOp KernelState String +nextErrorLabel = + errorLabel <$> GC.getUserState + +incErrorLabel :: GC.CompilerM KernelOp KernelState () +incErrorLabel = + GC.modifyUserState $ \s -> s {kernelNextSync = kernelNextSync s + 1} + +pendingError :: Bool -> GC.CompilerM KernelOp KernelState () +pendingError b = + GC.modifyUserState $ \s -> s {kernelSyncPending = b} + +hasCommunication :: ImpGPU.KernelCode -> Bool +hasCommunication = any communicates + where + communicates ErrorSync {} = True + communicates Barrier {} = True + communicates _ = False + +-- Whether we are generating code for a kernel or a device function. +-- This has minor effects, such as exactly how failures are +-- propagated. +data OpsMode = KernelMode | FunMode deriving (Eq) + +inKernelOperations :: + OpsMode -> + ImpGPU.KernelCode -> + GC.Operations KernelOp KernelState +inKernelOperations mode body = + GC.Operations + { GC.opsCompiler = kernelOps, + GC.opsMemoryType = kernelMemoryType, + GC.opsWriteScalar = kernelWriteScalar, + GC.opsReadScalar = kernelReadScalar, + GC.opsAllocate = cannotAllocate, + GC.opsDeallocate = cannotDeallocate, + GC.opsCopy = copyInKernel, + GC.opsStaticArray = noStaticArrays, + GC.opsFatMemory = False, + GC.opsError = errorInKernel, + GC.opsCall = callInKernel, + GC.opsCritical = mempty + } + where + has_communication = hasCommunication body + + fence FenceLocal = [C.cexp|CLK_LOCAL_MEM_FENCE|] + fence FenceGlobal = [C.cexp|CLK_GLOBAL_MEM_FENCE | CLK_LOCAL_MEM_FENCE|] + + kernelOps :: GC.OpCompiler KernelOp KernelState + kernelOps (GetGroupId v i) = + GC.stm [C.cstm|$id:v = get_group_id($int:i);|] + kernelOps (GetLocalId v i) = + GC.stm [C.cstm|$id:v = get_local_id($int:i);|] + kernelOps (GetLocalSize v i) = + GC.stm [C.cstm|$id:v = get_local_size($int:i);|] + kernelOps (GetGlobalId v i) = + GC.stm [C.cstm|$id:v = get_global_id($int:i);|] + kernelOps (GetGlobalSize v i) = + GC.stm [C.cstm|$id:v = get_global_size($int:i);|] + kernelOps (GetLockstepWidth v) = + GC.stm [C.cstm|$id:v = LOCKSTEP_WIDTH;|] + kernelOps (Barrier f) = do + GC.stm [C.cstm|barrier($exp:(fence f));|] + GC.modifyUserState $ \s -> s {kernelHasBarriers = True} + kernelOps (MemFence FenceLocal) = + GC.stm [C.cstm|mem_fence_local();|] + kernelOps (MemFence FenceGlobal) = + GC.stm [C.cstm|mem_fence_global();|] + kernelOps (LocalAlloc name size) = do + name' <- newVName $ pretty name ++ "_backing" + GC.modifyUserState $ \s -> + s {kernelLocalMemory = (name', fmap untyped size) : kernelLocalMemory s} + GC.stm [C.cstm|$id:name = (__local char*) $id:name';|] + kernelOps (ErrorSync f) = do + label <- nextErrorLabel + pending <- kernelSyncPending <$> GC.getUserState + when pending $ do + pendingError False + GC.stm [C.cstm|$id:label: barrier($exp:(fence f));|] + GC.stm [C.cstm|if (local_failure) { return; }|] + GC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|] -- intentional + GC.modifyUserState $ \s -> s {kernelHasBarriers = True} + incErrorLabel + kernelOps (Atomic space aop) = atomicOps space aop + + atomicCast s t = do + let volatile = [C.ctyquals|volatile|] + quals <- case s of + Space sid -> pointerQuals sid + _ -> pointerQuals "global" + return [C.cty|$tyquals:(volatile++quals) $ty:t|] + + atomicSpace (Space sid) = sid + atomicSpace _ = "global" + + doAtomic s t old arr ind val op ty = do + ind' <- GC.compileExp $ untyped $ unCount ind + val' <- GC.compileExp val + cast <- atomicCast s ty + GC.stm [C.cstm|$id:old = $id:op'(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|] + where + op' = op ++ "_" ++ pretty t ++ "_" ++ atomicSpace s + + doAtomicCmpXchg s t old arr ind cmp val ty = do + ind' <- GC.compileExp $ untyped $ unCount ind + cmp' <- GC.compileExp cmp + val' <- GC.compileExp val + cast <- atomicCast s ty + GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|] + where + op = "atomic_cmpxchg_" ++ pretty t ++ "_" ++ atomicSpace s + doAtomicXchg s t old arr ind val ty = do + cast <- atomicCast s ty + ind' <- GC.compileExp $ untyped $ unCount ind + val' <- GC.compileExp val + GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|] + where + op = "atomic_chg_" ++ pretty t ++ "_" ++ atomicSpace s + -- First the 64-bit operations. + atomicOps s (AtomicAdd Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_add" [C.cty|typename int64_t|] + atomicOps s (AtomicFAdd Float64 old arr ind val) = + doAtomic s Float64 old arr ind val "atomic_fadd" [C.cty|double|] + atomicOps s (AtomicSMax Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_smax" [C.cty|typename int64_t|] + atomicOps s (AtomicSMin Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_smin" [C.cty|typename int64_t|] + atomicOps s (AtomicUMax Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_umax" [C.cty|unsigned int64_t|] + atomicOps s (AtomicUMin Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_umin" [C.cty|unsigned int64_t|] + atomicOps s (AtomicAnd Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_and" [C.cty|typename int64_t|] + atomicOps s (AtomicOr Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_or" [C.cty|typename int64_t|] + atomicOps s (AtomicXor Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_xor" [C.cty|typename int64_t|] + atomicOps s (AtomicCmpXchg (IntType Int64) old arr ind cmp val) = + doAtomicCmpXchg s (IntType Int64) old arr ind cmp val [C.cty|typename int64_t|] + atomicOps s (AtomicXchg (IntType Int64) old arr ind val) = + doAtomicXchg s (IntType Int64) old arr ind val [C.cty|typename int64_t|] + -- + atomicOps s (AtomicAdd t old arr ind val) = + doAtomic s t old arr ind val "atomic_add" [C.cty|int|] + atomicOps s (AtomicFAdd Float32 old arr ind val) = + doAtomic s Float32 old arr ind val "atomic_fadd" [C.cty|float|] + atomicOps s (AtomicSMax t old arr ind val) = + doAtomic s t old arr ind val "atomic_smax" [C.cty|int|] + atomicOps s (AtomicSMin t old arr ind val) = + doAtomic s t old arr ind val "atomic_smin" [C.cty|int|] + atomicOps s (AtomicUMax t old arr ind val) = + doAtomic s t old arr ind val "atomic_umax" [C.cty|unsigned int|] + atomicOps s (AtomicUMin t old arr ind val) = + doAtomic s t old arr ind val "atomic_umin" [C.cty|unsigned int|] + atomicOps s (AtomicAnd t old arr ind val) = + doAtomic s t old arr ind val "atomic_and" [C.cty|int|] + atomicOps s (AtomicOr t old arr ind val) = + doAtomic s t old arr ind val "atomic_or" [C.cty|int|] + atomicOps s (AtomicXor t old arr ind val) = + doAtomic s t old arr ind val "atomic_xor" [C.cty|int|] + atomicOps s (AtomicCmpXchg t old arr ind cmp val) = + doAtomicCmpXchg s t old arr ind cmp val [C.cty|int|] + atomicOps s (AtomicXchg t old arr ind val) = + doAtomicXchg s t old arr ind val [C.cty|int|] + + cannotAllocate :: GC.Allocate KernelOp KernelState + cannotAllocate _ = + error "Cannot allocate memory in kernel" + + cannotDeallocate :: GC.Deallocate KernelOp KernelState + cannotDeallocate _ _ = + error "Cannot deallocate memory in kernel" + + copyInKernel :: GC.Copy KernelOp KernelState + copyInKernel _ _ _ _ _ _ _ = + error "Cannot bulk copy in kernel." + + noStaticArrays :: GC.StaticArray KernelOp KernelState + noStaticArrays _ _ _ _ = + error "Cannot create static array in kernel." + + kernelMemoryType space = do + quals <- pointerQuals space + return [C.cty|$tyquals:quals $ty:defaultMemBlockType|] + + kernelWriteScalar = + GC.writeScalarPointerWithQuals pointerQuals + + kernelReadScalar = + GC.readScalarPointerWithQuals pointerQuals + + whatNext = do + label <- nextErrorLabel + pendingError True + return $ + if has_communication + then [C.citems|local_failure = true; goto $id:label;|] + else + if mode == FunMode + then [C.citems|return 1;|] + else [C.citems|return;|] + + callInKernel dests fname args + | isBuiltInFunction fname = + GC.opsCall GC.defaultOperations dests fname args + | otherwise = 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 + + GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|] + + errorInKernel msg@(ErrorMsg parts) backtrace = do + n <- length . kernelFailures <$> GC.getUserState + GC.modifyUserState $ \s -> + s {kernelFailures = kernelFailures s ++ [FailureMsg msg backtrace]} + let setArgs _ [] = return [] + setArgs i (ErrorString {} : parts') = setArgs i parts' + setArgs i (ErrorInt32 x : parts') = do + x' <- GC.compileExp x + stms <- setArgs (i + 1) parts' + return $ [C.cstm|global_failure_args[$int:i] = (typename int64_t)$exp:x';|] : stms + setArgs i (ErrorInt64 x : parts') = do + x' <- GC.compileExp x + stms <- setArgs (i + 1) parts' + return $ [C.cstm|global_failure_args[$int:i] = $exp:x';|] : stms + argstms <- setArgs (0 :: Int) parts + + what_next <- whatNext + + GC.stm + [C.cstm|{ if (atomic_cmpxchg_i32_global(global_failure, -1, $int:n) == -1) + { $stms:argstms; } + $items:what_next + }|] + +--- Checking requirements + +typesInKernel :: Kernel -> S.Set PrimType +typesInKernel kernel = typesInCode $ kernelBody kernel + +typesInCode :: ImpGPU.KernelCode -> S.Set PrimType +typesInCode Skip = mempty +typesInCode (c1 :>>: c2) = typesInCode c1 <> typesInCode c2 +typesInCode (For _ e c) = typesInExp e <> typesInCode c +typesInCode (While (TPrimExp e) c) = typesInExp e <> typesInCode c +typesInCode DeclareMem {} = mempty +typesInCode (DeclareScalar _ _ t) = S.singleton t +typesInCode (DeclareArray _ _ t _) = S.singleton t +typesInCode (Allocate _ (Count (TPrimExp e)) _) = typesInExp e +typesInCode Free {} = mempty +typesInCode + ( Copy + _ + (Count (TPrimExp e1)) + _ + _ + (Count (TPrimExp e2)) + _ + (Count (TPrimExp e3)) + ) = + typesInExp e1 <> typesInExp e2 <> typesInExp e3 +typesInCode (Write _ (Count (TPrimExp e1)) t _ _ e2) = + typesInExp e1 <> S.singleton t <> typesInExp e2 +typesInCode (SetScalar _ e) = typesInExp e +typesInCode SetMem {} = mempty +typesInCode (Call _ _ es) = mconcat $ map typesInArg es + where + typesInArg MemArg {} = mempty + typesInArg (ExpArg e) = typesInExp e +typesInCode (If (TPrimExp e) c1 c2) = + typesInExp e <> typesInCode c1 <> typesInCode c2 +typesInCode (Assert e _ _) = typesInExp e +typesInCode (Comment _ c) = typesInCode c +typesInCode (DebugPrint _ v) = maybe mempty typesInExp v +typesInCode Op {} = mempty + +typesInExp :: Exp -> S.Set PrimType +typesInExp (ValueExp v) = S.singleton $ primValueType v +typesInExp (BinOpExp _ e1 e2) = typesInExp e1 <> typesInExp e2 +typesInExp (CmpOpExp _ e1 e2) = typesInExp e1 <> typesInExp e2 +typesInExp (ConvOpExp op e) = S.fromList [from, to] <> typesInExp e + where + (from, to) = convOpType op +typesInExp (UnOpExp _ e) = typesInExp e +typesInExp (FunExp _ args t) = S.singleton t <> mconcat (map typesInExp args) +typesInExp (LeafExp (Index _ (Count (TPrimExp e)) t _ _) _) = S.singleton t <> typesInExp e +typesInExp (LeafExp ScalarVar {} _) = mempty diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs b/src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs @@ -0,0 +1,386 @@ +-- | Carefully optimised implementations of GPU transpositions. +-- Written in ImpCode so we can compile it to both CUDA and OpenCL. +module Futhark.CodeGen.ImpGen.GPU.Transpose + ( TransposeType (..), + TransposeArgs, + mapTransposeKernel, + ) +where + +import Futhark.CodeGen.ImpCode.GPU +import Futhark.IR.Prop.Types +import Futhark.Util.IntegralExp (divUp, quot, rem) +import Prelude hiding (quot, rem) + +-- | Which form of transposition to generate code for. +data TransposeType + = TransposeNormal + | TransposeLowWidth + | TransposeLowHeight + | -- | For small arrays that do not + -- benefit from coalescing. + TransposeSmall + deriving (Eq, Ord, Show) + +-- | The types of the arguments accepted by a transposition function. +type TransposeArgs = + ( VName, + TExp Int32, + VName, + TExp Int32, + TExp Int32, + TExp Int32, + TExp Int32, + TExp Int32, + TExp Int32, + VName + ) + +elemsPerThread :: TExp Int32 +elemsPerThread = 4 + +mapTranspose :: TExp Int32 -> TransposeArgs -> PrimType -> TransposeType -> KernelCode +mapTranspose block_dim args t kind = + case kind of + TransposeSmall -> + mconcat + [ get_ids, + dec our_array_offset $ vi32 get_global_id_0 `quot` (height * width) * (height * width), + dec x_index $ (vi32 get_global_id_0 `rem` (height * width)) `quot` height, + dec y_index $ vi32 get_global_id_0 `rem` height, + dec odata_offset $ + (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset, + dec idata_offset $ + (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset, + dec index_in $ vi32 y_index * width + vi32 x_index, + dec index_out $ vi32 x_index * height + vi32 y_index, + when + (vi32 get_global_id_0 .<. width * height * num_arrays) + ( Write odata (elements $ sExt64 $ vi32 odata_offset + vi32 index_out) t (Space "global") Nonvolatile $ + index idata (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) t (Space "global") Nonvolatile + ) + ] + TransposeLowWidth -> + mkTranspose $ + lowDimBody + (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_0 `quot` muly)) + ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_1 + + (vi32 get_local_id_0 `rem` muly) * block_dim + ) + ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_0 + + (vi32 get_local_id_1 `rem` muly) * block_dim + ) + (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_1 `quot` muly)) + TransposeLowHeight -> + mkTranspose $ + lowDimBody + ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_0 + + (vi32 get_local_id_1 `rem` mulx) * block_dim + ) + (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_1 `quot` mulx)) + (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_0 `quot` mulx)) + ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_1 + + (vi32 get_local_id_0 `rem` mulx) * block_dim + ) + TransposeNormal -> + mkTranspose $ + mconcat + [ dec x_index $ vi32 get_global_id_0, + dec y_index $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_1, + when (vi32 x_index .<. width) $ + For j (untyped elemsPerThread) $ + let i = vi32 j * (tile_dim `quot` elemsPerThread) + in mconcat + [ dec index_in $ (vi32 y_index + i) * width + vi32 x_index, + when (vi32 y_index + i .<. height) $ + Write + block + ( elements $ + sExt64 $ + (vi32 get_local_id_1 + i) * (tile_dim + 1) + + vi32 get_local_id_0 + ) + t + (Space "local") + Nonvolatile + $ index + idata + (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) + t + (Space "global") + Nonvolatile + ], + Op $ Barrier FenceLocal, + SetScalar x_index $ untyped $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_0, + SetScalar y_index $ untyped $ vi32 get_group_id_0 * tile_dim + vi32 get_local_id_1, + when (vi32 x_index .<. height) $ + For j (untyped elemsPerThread) $ + let i = vi32 j * (tile_dim `quot` elemsPerThread) + in mconcat + [ dec index_out $ (vi32 y_index + i) * height + vi32 x_index, + when (vi32 y_index + i .<. width) $ + Write + odata + (elements $ sExt64 $ vi32 odata_offset + vi32 index_out) + t + (Space "global") + Nonvolatile + $ index + block + ( elements $ + sExt64 $ + vi32 get_local_id_0 * (tile_dim + 1) + vi32 get_local_id_1 + i + ) + t + (Space "local") + Nonvolatile + ] + ] + where + dec v (TPrimExp e) = + DeclareScalar v Nonvolatile (primExpType e) <> SetScalar v e + tile_dim = 2 * block_dim + + when a b = If a b mempty + + ( odata, + basic_odata_offset, + idata, + basic_idata_offset, + width, + height, + mulx, + muly, + num_arrays, + block + ) = args + + -- Be extremely careful when editing this list to ensure that + -- the names match up. Also, be careful that the tags on + -- these names do not conflict with the tags of the + -- surrounding code. We accomplish the latter by using very + -- low tags (normal variables start at least in the low + -- hundreds). + [ our_array_offset, + x_index, + y_index, + odata_offset, + idata_offset, + index_in, + index_out, + get_global_id_0, + get_local_id_0, + get_local_id_1, + get_group_id_0, + get_group_id_1, + get_group_id_2, + j + ] = + zipWith (flip VName) [30 ..] $ + map + nameFromString + [ "our_array_offset", + "x_index", + "y_index", + "odata_offset", + "idata_offset", + "index_in", + "index_out", + "get_global_id_0", + "get_local_id_0", + "get_local_id_1", + "get_group_id_0", + "get_group_id_1", + "get_group_id_2", + "j" + ] + + get_ids = + mconcat + [ DeclareScalar get_global_id_0 Nonvolatile int32, + Op $ GetGlobalId get_global_id_0 0, + DeclareScalar get_local_id_0 Nonvolatile int32, + Op $ GetLocalId get_local_id_0 0, + DeclareScalar get_local_id_1 Nonvolatile int32, + Op $ GetLocalId get_local_id_1 1, + DeclareScalar get_group_id_0 Nonvolatile int32, + Op $ GetGroupId get_group_id_0 0, + DeclareScalar get_group_id_1 Nonvolatile int32, + Op $ GetGroupId get_group_id_1 1, + DeclareScalar get_group_id_2 Nonvolatile int32, + Op $ GetGroupId get_group_id_2 2 + ] + + mkTranspose body = + mconcat + [ get_ids, + dec our_array_offset $ vi32 get_group_id_2 * width * height, + dec odata_offset $ + (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset, + dec idata_offset $ + (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset, + body + ] + + lowDimBody x_in_index y_in_index x_out_index y_out_index = + mconcat + [ dec x_index x_in_index, + dec y_index y_in_index, + dec index_in $ vi32 y_index * width + vi32 x_index, + when (vi32 x_index .<. width .&&. vi32 y_index .<. height) $ + Write + block + (elements $ sExt64 $ vi32 get_local_id_1 * (block_dim + 1) + vi32 get_local_id_0) + t + (Space "local") + Nonvolatile + $ index + idata + (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) + t + (Space "global") + Nonvolatile, + Op $ Barrier FenceLocal, + SetScalar x_index $ untyped x_out_index, + SetScalar y_index $ untyped y_out_index, + dec index_out $ vi32 y_index * height + vi32 x_index, + when (vi32 x_index .<. height .&&. vi32 y_index .<. width) $ + Write + odata + (elements $ sExt64 (vi32 odata_offset + vi32 index_out)) + t + (Space "global") + Nonvolatile + $ index + block + (elements $ sExt64 $ vi32 get_local_id_0 * (block_dim + 1) + vi32 get_local_id_1) + t + (Space "local") + Nonvolatile + ] + +-- | Generate a transpose kernel. There is special support to handle +-- input arrays with low width, low height, or both. +-- +-- Normally when transposing a @[2][n]@ array we would use a @FUT_BLOCK_DIM x +-- FUT_BLOCK_DIM@ group to process a @[2][FUT_BLOCK_DIM]@ slice of the input +-- array. This would mean that many of the threads in a group would be inactive. +-- We try to remedy this by using a special kernel that will process a larger +-- part of the input, by using more complex indexing. In our example, we could +-- use all threads in a group if we are processing @(2/FUT_BLOCK_DIM)@ as large +-- a slice of each rows per group. The variable @mulx@ contains this factor for +-- the kernel to handle input arrays with low height. +-- +-- See issue #308 on GitHub for more details. +-- +-- These kernels are optimized to ensure all global reads and writes +-- are coalesced, and to avoid bank conflicts in shared memory. Each +-- thread group transposes a 2D tile of block_dim*2 by block_dim*2 +-- elements. The size of a thread group is block_dim/2 by +-- block_dim*2, meaning that each thread will process 4 elements in a +-- 2D tile. The shared memory array containing the 2D tile consists +-- of block_dim*2 by block_dim*2+1 elements. Padding each row with +-- an additional element prevents bank conflicts from occuring when +-- the tile is accessed column-wise. +mapTransposeKernel :: + String -> + Integer -> + TransposeArgs -> + PrimType -> + TransposeType -> + Kernel +mapTransposeKernel desc block_dim_int args t kind = + Kernel + { kernelBody = + DeclareMem block (Space "local") + <> Op (LocalAlloc block block_size) + <> mapTranspose block_dim args t kind, + kernelUses = uses, + kernelNumGroups = map untyped num_groups, + kernelGroupSize = map untyped group_size, + kernelName = nameFromString name, + kernelFailureTolerant = True + } + where + pad2DBytes k = k * (k + 1) * primByteSize t + block_size = + bytes $ + case kind of + TransposeSmall -> 1 :: TExp Int64 + -- Not used, but AMD's OpenCL + -- does not like zero-size + -- local memory. + TransposeNormal -> fromInteger $ pad2DBytes $ 2 * block_dim_int + TransposeLowWidth -> fromInteger $ pad2DBytes block_dim_int + TransposeLowHeight -> fromInteger $ pad2DBytes block_dim_int + block_dim = fromInteger block_dim_int :: TExp Int32 + + ( odata, + basic_odata_offset, + idata, + basic_idata_offset, + width, + height, + mulx, + muly, + num_arrays, + block + ) = args + + (num_groups, group_size) = + case kind of + TransposeSmall -> + ( [(num_arrays * width * height) `divUp` (block_dim * block_dim)], + [block_dim * block_dim] + ) + TransposeLowWidth -> + lowDimKernelAndGroupSize block_dim num_arrays width $ height `divUp` muly + TransposeLowHeight -> + lowDimKernelAndGroupSize block_dim num_arrays (width `divUp` mulx) height + TransposeNormal -> + let actual_dim = block_dim * 2 + in ( [ width `divUp` actual_dim, + height `divUp` actual_dim, + num_arrays + ], + [actual_dim, actual_dim `quot` elemsPerThread, 1] + ) + + uses = + map + (`ScalarUse` int32) + ( namesToList $ + mconcat $ + map + freeIn + [ basic_odata_offset, + basic_idata_offset, + num_arrays, + width, + height, + mulx, + muly + ] + ) + ++ map MemoryUse [odata, idata] + + name = + case kind of + TransposeSmall -> desc ++ "_small" + TransposeLowHeight -> desc ++ "_low_height" + TransposeLowWidth -> desc ++ "_low_width" + TransposeNormal -> desc + +lowDimKernelAndGroupSize :: + TExp Int32 -> + TExp Int32 -> + TExp Int32 -> + TExp Int32 -> + ([TExp Int32], [TExp Int32]) +lowDimKernelAndGroupSize block_dim num_arrays x_elems y_elems = + ( [ x_elems `divUp` block_dim, + y_elems `divUp` block_dim, + num_arrays + ], + [block_dim, block_dim, 1] + ) diff --git a/src/Futhark/CodeGen/ImpGen/Kernels.hs b/src/Futhark/CodeGen/ImpGen/Kernels.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels.hs +++ /dev/null @@ -1,431 +0,0 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} - --- | Compile a 'KernelsMem' program to imperative code with kernels. --- This is mostly (but not entirely) the same process no matter if we --- are targeting OpenCL or CUDA. The important distinctions (the host --- level code) are introduced later. -module Futhark.CodeGen.ImpGen.Kernels - ( compileProgOpenCL, - compileProgCUDA, - Warnings, - ) -where - -import Control.Monad.Except -import Data.Bifunctor (second) -import Data.List (foldl') -import qualified Data.Map as M -import Data.Maybe -import Futhark.CodeGen.ImpCode.Kernels (bytes) -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen hiding (compileProg) -import qualified Futhark.CodeGen.ImpGen -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.CodeGen.ImpGen.Kernels.SegHist -import Futhark.CodeGen.ImpGen.Kernels.SegMap -import Futhark.CodeGen.ImpGen.Kernels.SegRed -import Futhark.CodeGen.ImpGen.Kernels.SegScan -import Futhark.CodeGen.ImpGen.Kernels.Transpose -import Futhark.CodeGen.SetDefaultSpace -import Futhark.Error -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.MonadFreshNames -import Futhark.Util.IntegralExp (IntegralExp, divUp, quot, rem) -import Prelude hiding (quot, rem) - -callKernelOperations :: Operations KernelsMem HostEnv Imp.HostOp -callKernelOperations = - Operations - { opsExpCompiler = expCompiler, - opsCopyCompiler = callKernelCopy, - opsOpCompiler = opCompiler, - opsStmsCompiler = defCompileStms, - opsAllocCompilers = mempty - } - -openclAtomics, cudaAtomics :: AtomicBinOp -(openclAtomics, cudaAtomics) = (flip lookup opencl, flip lookup cuda) - where - opencl64 = - [ (Add Int64 OverflowUndef, Imp.AtomicAdd Int64), - (SMax Int64, Imp.AtomicSMax Int64), - (SMin Int64, Imp.AtomicSMin Int64), - (UMax Int64, Imp.AtomicUMax Int64), - (UMin Int64, Imp.AtomicUMin Int64), - (And Int64, Imp.AtomicAnd Int64), - (Or Int64, Imp.AtomicOr Int64), - (Xor Int64, Imp.AtomicXor Int64) - ] - opencl32 = - [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32), - (SMax Int32, Imp.AtomicSMax Int32), - (SMin Int32, Imp.AtomicSMin Int32), - (UMax Int32, Imp.AtomicUMax Int32), - (UMin Int32, Imp.AtomicUMin Int32), - (And Int32, Imp.AtomicAnd Int32), - (Or Int32, Imp.AtomicOr Int32), - (Xor Int32, Imp.AtomicXor Int32) - ] - opencl = opencl32 ++ opencl64 - cuda = - opencl - ++ [ (FAdd Float32, Imp.AtomicFAdd Float32), - (FAdd Float64, Imp.AtomicFAdd Float64) - ] - -compileProg :: - MonadFreshNames m => - HostEnv -> - Prog KernelsMem -> - m (Warnings, Imp.Program) -compileProg env prog = - second (setDefaultSpace (Imp.Space "device")) - <$> Futhark.CodeGen.ImpGen.compileProg env callKernelOperations (Imp.Space "device") prog - --- | Compile a 'KernelsMem' program to low-level parallel code, with --- either CUDA or OpenCL characteristics. -compileProgOpenCL, - compileProgCUDA :: - MonadFreshNames m => Prog KernelsMem -> m (Warnings, Imp.Program) -compileProgOpenCL = compileProg $ HostEnv openclAtomics OpenCL mempty -compileProgCUDA = compileProg $ HostEnv cudaAtomics CUDA mempty - -opCompiler :: - Pattern KernelsMem -> - Op KernelsMem -> - CallKernelGen () -opCompiler dest (Alloc e space) = - compileAlloc dest e space -opCompiler (Pattern _ [pe]) (Inner (SizeOp (GetSize key size_class))) = do - fname <- askFunction - sOp $ - Imp.GetSize (patElemName pe) (keyWithEntryPoint fname key) $ - sizeClassWithEntryPoint fname size_class -opCompiler (Pattern _ [pe]) (Inner (SizeOp (CmpSizeLe key size_class x))) = do - fname <- askFunction - let size_class' = sizeClassWithEntryPoint fname size_class - sOp . Imp.CmpSizeLe (patElemName pe) (keyWithEntryPoint fname key) size_class' - =<< toExp x -opCompiler (Pattern _ [pe]) (Inner (SizeOp (GetSizeMax size_class))) = - sOp $ Imp.GetSizeMax (patElemName pe) size_class -opCompiler (Pattern _ [pe]) (Inner (SizeOp (CalcNumGroups w64 max_num_groups_key group_size))) = do - fname <- askFunction - max_num_groups :: TV Int32 <- dPrim "max_num_groups" int32 - sOp $ - Imp.GetSize (tvVar max_num_groups) (keyWithEntryPoint fname max_num_groups_key) $ - sizeClassWithEntryPoint fname SizeNumGroups - - -- If 'w' is small, we launch fewer groups than we normally would. - -- We don't want any idle groups. - -- - -- The calculations are done with 64-bit integers to avoid overflow - -- issues. - let num_groups_maybe_zero = - sMin64 (toInt64Exp w64 `divUp` toInt64Exp group_size) $ - sExt64 (tvExp max_num_groups) - -- We also don't want zero groups. - let num_groups = sMax64 1 num_groups_maybe_zero - mkTV (patElemName pe) int32 <-- sExt32 num_groups -opCompiler dest (Inner (SegOp op)) = - segOpCompiler dest op -opCompiler pat e = - compilerBugS $ - "opCompiler: Invalid pattern\n " - ++ pretty pat - ++ "\nfor expression\n " - ++ pretty e - -sizeClassWithEntryPoint :: Maybe Name -> Imp.SizeClass -> Imp.SizeClass -sizeClassWithEntryPoint fname (Imp.SizeThreshold path def) = - Imp.SizeThreshold (map f path) def - where - f (name, x) = (keyWithEntryPoint fname name, x) -sizeClassWithEntryPoint _ size_class = size_class - -segOpCompiler :: - Pattern KernelsMem -> - SegOp SegLevel KernelsMem -> - CallKernelGen () -segOpCompiler pat (SegMap lvl space _ kbody) = - compileSegMap pat lvl space kbody -segOpCompiler pat (SegRed lvl@SegThread {} space reds _ kbody) = - compileSegRed pat lvl space reds kbody -segOpCompiler pat (SegScan lvl@SegThread {} space scans _ kbody) = - compileSegScan pat lvl space scans kbody -segOpCompiler pat (SegHist (SegThread num_groups group_size _) space ops _ kbody) = - compileSegHist pat num_groups group_size space ops kbody -segOpCompiler pat segop = - compilerBugS $ "segOpCompiler: unexpected " ++ pretty (segLevel segop) ++ " for rhs of pattern " ++ pretty pat - --- Create boolean expression that checks whether all kernels in the --- enclosed code do not use more local memory than we have available. --- We look at *all* the kernels here, even those that might be --- otherwise protected by their own multi-versioning branches deeper --- down. Currently the compiler will not generate multi-versioning --- that makes this a problem, but it might in the future. -checkLocalMemoryReqs :: Imp.Code -> CallKernelGen (Maybe (Imp.TExp Bool)) -checkLocalMemoryReqs code = do - scope <- askScope - let alloc_sizes = map (sum . map alignedSize . localAllocSizes . Imp.kernelBody) $ getKernels code - - -- If any of the sizes involve a variable that is not known at this - -- point, then we cannot check the requirements. - if any (`M.notMember` scope) (namesToList $ freeIn alloc_sizes) - then return Nothing - else do - local_memory_capacity :: TV Int32 <- dPrim "local_memory_capacity" int32 - sOp $ Imp.GetSizeMax (tvVar local_memory_capacity) SizeLocalMemory - - let local_memory_capacity_64 = - sExt64 $ tvExp local_memory_capacity - fits size = - unCount size .<=. local_memory_capacity_64 - return $ Just $ foldl' (.&&.) true (map fits alloc_sizes) - where - getKernels = foldMap getKernel - getKernel (Imp.CallKernel k) = [k] - getKernel _ = [] - - localAllocSizes = foldMap localAllocSize - localAllocSize (Imp.LocalAlloc _ size) = [size] - localAllocSize _ = [] - - -- These allocations will actually be padded to an 8-byte aligned - -- size, so we should take that into account when checking whether - -- they fit. - alignedSize x = x + ((8 - (x `rem` 8)) `rem` 8) - -withAcc :: - Pattern KernelsMem -> - [(Shape, [VName], Maybe (Lambda KernelsMem, [SubExp]))] -> - Lambda KernelsMem -> - CallKernelGen () -withAcc pat inputs lam = do - atomics <- hostAtomics <$> askEnv - locksForInputs atomics $ zip accs inputs - where - accs = map paramName $ lambdaParams lam - locksForInputs _ [] = - defCompileExp pat $ WithAcc inputs lam - locksForInputs atomics ((c, (_, _, op)) : inputs') - | Just (op_lam, _) <- op, - AtomicLocking _ <- atomicUpdateLocking atomics op_lam = do - let num_locks = 100151 - locks_arr <- - sStaticArray "withacc_locks" (Space "device") int32 $ - Imp.ArrayZeros num_locks - let locks = Locks locks_arr num_locks - extend env = env {hostLocks = M.insert c locks $ hostLocks env} - localEnv extend $ locksForInputs atomics inputs' - | otherwise = - locksForInputs atomics inputs' - -expCompiler :: ExpCompiler KernelsMem HostEnv Imp.HostOp --- We generate a simple kernel for itoa and replicate. -expCompiler (Pattern _ [pe]) (BasicOp (Iota n x s et)) = do - x' <- toExp x - s' <- toExp s - - sIota (patElemName pe) (toInt64Exp n) x' s' et -expCompiler (Pattern _ [pe]) (BasicOp (Replicate _ se)) = - sReplicate (patElemName pe) se --- Allocation in the "local" space is just a placeholder. -expCompiler _ (Op (Alloc _ (Space "local"))) = - return () -expCompiler pat (WithAcc inputs lam) = - withAcc pat inputs lam --- This is a multi-versioning If created by incremental flattening. --- We need to augment the conditional with a check that any local --- memory requirements in tbranch are compatible with the hardware. --- We do not check anything for fbranch, as we assume that it will --- always be safe (and what would we do if none of the branches would --- work?). -expCompiler dest (If cond tbranch fbranch (IfDec _ IfEquiv)) = do - tcode <- collect $ compileBody dest tbranch - fcode <- collect $ compileBody dest fbranch - check <- checkLocalMemoryReqs tcode - emit $ case check of - Nothing -> fcode - Just ok -> Imp.If (ok .&&. toBoolExp cond) tcode fcode -expCompiler dest e = - defCompileExp dest e - -callKernelCopy :: CopyCompiler KernelsMem HostEnv Imp.HostOp -callKernelCopy - bt - destloc@(MemLocation destmem _ destIxFun) - destslice - srcloc@(MemLocation srcmem srcshape srcIxFun) - srcslice - | Just - ( destoffset, - srcoffset, - num_arrays, - size_x, - size_y - ) <- - isMapTransposeCopy bt destloc destslice srcloc srcslice = do - fname <- mapTransposeForType bt - emit $ - Imp.Call - [] - fname - [ Imp.MemArg destmem, - Imp.ExpArg $ untyped destoffset, - Imp.MemArg srcmem, - Imp.ExpArg $ untyped srcoffset, - Imp.ExpArg $ untyped num_arrays, - Imp.ExpArg $ untyped size_x, - Imp.ExpArg $ untyped size_y - ] - | bt_size <- primByteSize bt, - Just destoffset <- - IxFun.linearWithOffset (IxFun.slice destIxFun destslice) bt_size, - Just srcoffset <- - IxFun.linearWithOffset (IxFun.slice srcIxFun srcslice) bt_size = do - let num_elems = Imp.elements $ product $ map toInt64Exp srcshape - srcspace <- entryMemSpace <$> lookupMemory srcmem - destspace <- entryMemSpace <$> lookupMemory destmem - emit $ - Imp.Copy - destmem - (bytes $ sExt64 destoffset) - destspace - srcmem - (bytes $ sExt64 srcoffset) - srcspace - $ num_elems `Imp.withElemType` bt - | otherwise = sCopy bt destloc destslice srcloc srcslice - -mapTransposeForType :: PrimType -> CallKernelGen Name -mapTransposeForType bt = do - let fname = nameFromString $ "builtin#" <> mapTransposeName bt - - exists <- hasFunction fname - unless exists $ emitFunction fname $ mapTransposeFunction bt - - return fname - -mapTransposeName :: PrimType -> String -mapTransposeName bt = "gpu_map_transpose_" ++ pretty bt - -mapTransposeFunction :: PrimType -> Imp.Function -mapTransposeFunction bt = - Imp.Function Nothing [] params transpose_code [] [] - where - params = - [ memparam destmem, - intparam destoffset, - memparam srcmem, - intparam srcoffset, - intparam num_arrays, - intparam x, - intparam y - ] - - space = Space "device" - memparam v = Imp.MemParam v space - intparam v = Imp.ScalarParam v $ IntType Int32 - - [ destmem, - destoffset, - srcmem, - srcoffset, - num_arrays, - x, - y, - mulx, - muly, - block - ] = - zipWith - (VName . nameFromString) - [ "destmem", - "destoffset", - "srcmem", - "srcoffset", - "num_arrays", - "x_elems", - "y_elems", - -- The following is only used for low width/height - -- transpose kernels - "mulx", - "muly", - "block" - ] - [0 ..] - - block_dim_int = 16 - - block_dim :: IntegralExp a => a - block_dim = 16 - - -- When an input array has either width==1 or height==1, performing a - -- transpose will be the same as performing a copy. - can_use_copy = - let onearr = Imp.vi32 num_arrays .==. 1 - height_is_one = Imp.vi32 y .==. 1 - width_is_one = Imp.vi32 x .==. 1 - in onearr .&&. (width_is_one .||. height_is_one) - - transpose_code = - Imp.If input_is_empty mempty $ - mconcat - [ Imp.DeclareScalar muly Imp.Nonvolatile (IntType Int32), - Imp.SetScalar muly $ untyped $ block_dim `quot` Imp.vi32 x, - Imp.DeclareScalar mulx Imp.Nonvolatile (IntType Int32), - Imp.SetScalar mulx $ untyped $ block_dim `quot` Imp.vi32 y, - Imp.If can_use_copy copy_code $ - Imp.If should_use_lowwidth (callTransposeKernel TransposeLowWidth) $ - Imp.If should_use_lowheight (callTransposeKernel TransposeLowHeight) $ - Imp.If should_use_small (callTransposeKernel TransposeSmall) $ - callTransposeKernel TransposeNormal - ] - - input_is_empty = - Imp.vi32 num_arrays .==. 0 .||. Imp.vi32 x .==. 0 .||. Imp.vi32 y .==. 0 - - should_use_small = - Imp.vi32 x .<=. (block_dim `quot` 2) - .&&. Imp.vi32 y .<=. (block_dim `quot` 2) - - should_use_lowwidth = - Imp.vi32 x .<=. (block_dim `quot` 2) - .&&. block_dim .<. Imp.vi32 y - - should_use_lowheight = - Imp.vi32 y .<=. (block_dim `quot` 2) - .&&. block_dim .<. Imp.vi32 x - - copy_code = - let num_bytes = sExt64 $ Imp.vi32 x * Imp.vi32 y * primByteSize bt - in Imp.Copy - destmem - (Imp.Count $ sExt64 $ Imp.vi32 destoffset) - space - srcmem - (Imp.Count $ sExt64 $ Imp.vi32 srcoffset) - space - (Imp.Count num_bytes) - - callTransposeKernel = - Imp.Op . Imp.CallKernel - . mapTransposeKernel - (mapTransposeName bt) - block_dim_int - ( destmem, - Imp.vi32 destoffset, - srcmem, - Imp.vi32 srcoffset, - Imp.vi32 x, - Imp.vi32 y, - Imp.vi32 mulx, - Imp.vi32 muly, - Imp.vi32 num_arrays, - block - ) - bt diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs +++ /dev/null @@ -1,1803 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE LambdaCase #-} -{-# LANGUAGE TypeFamilies #-} - -module Futhark.CodeGen.ImpGen.Kernels.Base - ( KernelConstants (..), - keyWithEntryPoint, - CallKernelGen, - InKernelGen, - Locks (..), - HostEnv (..), - Target (..), - KernelEnv (..), - computeThreadChunkSize, - groupReduce, - groupScan, - isActive, - sKernelThread, - sKernelGroup, - sReplicate, - sIota, - sCopy, - compileThreadResult, - compileGroupResult, - virtualiseGroups, - groupLoop, - kernelLoop, - groupCoverSpace, - precomputeSegOpIDs, - atomicUpdateLocking, - AtomicBinOp, - Locking (..), - AtomicUpdate (..), - DoAtomicUpdate, - ) -where - -import Control.Monad.Except -import Data.List (elemIndex, find, zip4) -import qualified Data.Map.Strict as M -import Data.Maybe -import qualified Data.Set as S -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.Error -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.MonadFreshNames -import Futhark.Transform.Rename -import Futhark.Util (chunks, dropLast, mapAccumLM, maybeNth, nubOrd, takeLast) -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Prelude hiding (quot, rem) - --- | Which target are we ultimately generating code for? While most --- of the kernels code is the same, there are some cases where we --- generate special code based on the ultimate low-level API we are --- targeting. -data Target = CUDA | OpenCL - --- | Information about the locks available for accumulators. -data Locks = Locks - { locksArray :: VName, - locksCount :: Int - } - -data HostEnv = HostEnv - { hostAtomics :: AtomicBinOp, - hostTarget :: Target, - hostLocks :: M.Map VName Locks - } - -data KernelEnv = KernelEnv - { kernelAtomics :: AtomicBinOp, - kernelConstants :: KernelConstants, - kernelLocks :: M.Map VName Locks - } - -type CallKernelGen = ImpM KernelsMem HostEnv Imp.HostOp - -type InKernelGen = ImpM KernelsMem KernelEnv Imp.KernelOp - -data KernelConstants = KernelConstants - { kernelGlobalThreadId :: Imp.TExp Int32, - kernelLocalThreadId :: Imp.TExp Int32, - kernelGroupId :: Imp.TExp Int32, - kernelGlobalThreadIdVar :: VName, - kernelLocalThreadIdVar :: VName, - kernelGroupIdVar :: VName, - kernelNumGroups :: Imp.TExp Int64, - kernelGroupSize :: Imp.TExp Int64, - kernelNumThreads :: Imp.TExp Int32, - kernelWaveSize :: Imp.TExp Int32, - kernelThreadActive :: Imp.TExp Bool, - -- | A mapping from dimensions of nested SegOps to already - -- computed local thread IDs. - kernelLocalIdMap :: M.Map [SubExp] [Imp.TExp Int32] - } - -segOpSizes :: Stms KernelsMem -> S.Set [SubExp] -segOpSizes = onStms - where - onStms = foldMap (onExp . stmExp) - onExp (Op (Inner (SegOp op))) = - S.singleton $ map snd $ unSegSpace $ segSpace op - onExp (If _ tbranch fbranch _) = - onStms (bodyStms tbranch) <> onStms (bodyStms fbranch) - onExp (DoLoop _ _ _ body) = - onStms (bodyStms body) - onExp _ = mempty - -precomputeSegOpIDs :: Stms KernelsMem -> InKernelGen a -> InKernelGen a -precomputeSegOpIDs stms m = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - new_ids <- M.fromList <$> mapM (mkMap ltid) (S.toList (segOpSizes stms)) - let f env = - env - { kernelConstants = - (kernelConstants env) {kernelLocalIdMap = new_ids} - } - localEnv f m - where - mkMap ltid dims = do - let dims' = map (sExt32 . toInt64Exp) dims - ids' <- mapM (dPrimVE "ltid_pre") $ unflattenIndex dims' ltid - return (dims, ids') - -keyWithEntryPoint :: Maybe Name -> Name -> Name -keyWithEntryPoint fname key = - nameFromString $ maybe "" ((++ ".") . nameToString) fname ++ nameToString key - -allocLocal :: AllocCompiler KernelsMem r Imp.KernelOp -allocLocal mem size = - sOp $ Imp.LocalAlloc mem size - -kernelAlloc :: - Pattern KernelsMem -> - SubExp -> - Space -> - InKernelGen () -kernelAlloc (Pattern _ [_]) _ ScalarSpace {} = - -- Handled by the declaration of the memory block, which is then - -- translated to an actual scalar variable during C code generation. - return () -kernelAlloc (Pattern _ [mem]) size (Space "local") = - allocLocal (patElemName mem) $ Imp.bytes $ toInt64Exp size -kernelAlloc (Pattern _ [mem]) _ _ = - compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." -kernelAlloc dest _ _ = - error $ "Invalid target for in-kernel allocation: " ++ show dest - -splitSpace :: - (ToExp w, ToExp i, ToExp elems_per_thread) => - Pattern KernelsMem -> - SplitOrdering -> - w -> - i -> - elems_per_thread -> - ImpM lore r op () -splitSpace (Pattern [] [size]) o w i elems_per_thread = do - num_elements <- Imp.elements . TPrimExp <$> toExp w - let i' = toInt64Exp i - elems_per_thread' <- Imp.elements . TPrimExp <$> toExp elems_per_thread - computeThreadChunkSize o i' elems_per_thread' num_elements (mkTV (patElemName size) int64) -splitSpace pat _ _ _ _ = - error $ "Invalid target for splitSpace: " ++ pretty pat - -updateAcc :: VName -> [SubExp] -> [SubExp] -> InKernelGen () -updateAcc acc is vs = sComment "UpdateAcc" $ do - -- See the ImpGen implementation of UpdateAcc for general notes. - let is' = map toInt64Exp is - (c, space, arrs, dims, op) <- lookupAcc acc is' - sWhen (inBounds (map DimFix is') dims) $ - case op of - Nothing -> - forM_ (zip arrs vs) $ \(arr, v) -> copyDWIMFix arr is' v [] - Just lam -> do - dLParams $ lambdaParams lam - let (_x_params, y_params) = - splitAt (length vs) $ map paramName $ lambdaParams lam - forM_ (zip y_params vs) $ \(yp, v) -> copyDWIM yp [] v [] - atomics <- kernelAtomics <$> askEnv - case atomicUpdateLocking atomics lam of - AtomicPrim f -> f space arrs is' - AtomicCAS f -> f space arrs is' - AtomicLocking f -> do - c_locks <- M.lookup c . kernelLocks <$> askEnv - case c_locks of - Just (Locks locks num_locks) -> do - let locking = - Locking locks 0 1 0 $ - pure . (`rem` fromIntegral num_locks) . flattenIndex dims - f locking space arrs is' - Nothing -> - error $ "Missing locks for " ++ pretty acc - -compileThreadExp :: ExpCompiler KernelsMem KernelEnv Imp.KernelOp -compileThreadExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = - forM_ (zip [0 ..] es) $ \(i, e) -> - copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] -compileThreadExp _ (BasicOp (UpdateAcc acc is vs)) = - updateAcc acc is vs -compileThreadExp dest e = - defCompileExp dest e - --- | Assign iterations of a for-loop to all threads in the kernel. --- The passed-in function is invoked with the (symbolic) iteration. --- 'threadOperations' will be in effect in the body. For --- multidimensional loops, use 'groupCoverSpace'. -kernelLoop :: - IntExp t => - Imp.TExp t -> - Imp.TExp t -> - Imp.TExp t -> - (Imp.TExp t -> InKernelGen ()) -> - InKernelGen () -kernelLoop tid num_threads n f = - localOps threadOperations $ - if n == num_threads - then f tid - else do - -- Compute how many elements this thread is responsible for. - -- Formula: (n - tid) / num_threads (rounded up). - let elems_for_this = (n - tid) `divUp` num_threads - - sFor "i" elems_for_this $ \i -> f $ i * num_threads + tid - --- | Assign iterations of a for-loop to threads in the workgroup. The --- passed-in function is invoked with the (symbolic) iteration. For --- multidimensional loops, use 'groupCoverSpace'. -groupLoop :: - Imp.TExp Int64 -> - (Imp.TExp Int64 -> InKernelGen ()) -> - InKernelGen () -groupLoop n f = do - constants <- kernelConstants <$> askEnv - kernelLoop - (sExt64 $ kernelLocalThreadId constants) - (kernelGroupSize constants) - n - f - --- | Iterate collectively though a multidimensional space, such that --- all threads in the group participate. The passed-in function is --- invoked with a (symbolic) point in the index space. -groupCoverSpace :: - [Imp.TExp Int64] -> - ([Imp.TExp Int64] -> InKernelGen ()) -> - InKernelGen () -groupCoverSpace ds f = - groupLoop (product ds) $ f . unflattenIndex ds - -compileGroupExp :: ExpCompiler KernelsMem KernelEnv Imp.KernelOp --- The static arrays stuff does not work inside kernels. -compileGroupExp (Pattern _ [dest]) (BasicOp (ArrayLit es _)) = - forM_ (zip [0 ..] es) $ \(i, e) -> - copyDWIMFix (patElemName dest) [fromIntegral (i :: Int64)] e [] -compileGroupExp _ (BasicOp (UpdateAcc acc is vs)) = - updateAcc acc is vs -compileGroupExp (Pattern _ [dest]) (BasicOp (Replicate ds se)) = do - let ds' = map toInt64Exp $ shapeDims ds - groupCoverSpace ds' $ \is -> - copyDWIMFix (patElemName dest) is se (drop (shapeRank ds) is) - sOp $ Imp.Barrier Imp.FenceLocal -compileGroupExp (Pattern _ [dest]) (BasicOp (Iota n e s it)) = do - n' <- toExp n - e' <- toExp e - s' <- toExp s - groupLoop (TPrimExp n') $ \i' -> do - x <- - dPrimV "x" $ - TPrimExp $ - BinOpExp (Add it OverflowUndef) e' $ - BinOpExp (Mul it OverflowUndef) (untyped i') s' - copyDWIMFix (patElemName dest) [i'] (Var (tvVar x)) [] - sOp $ Imp.Barrier Imp.FenceLocal - --- When generating code for a scalar in-place update, we must make --- sure that only one thread performs the write. When writing an --- array, the group-level copy code will take care of doing the right --- thing. -compileGroupExp (Pattern _ [pe]) (BasicOp (Update _ slice se)) - | null $ sliceDims slice = do - sOp $ Imp.Barrier Imp.FenceLocal - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - sWhen (ltid .==. 0) $ - copyDWIM (patElemName pe) (map (fmap toInt64Exp) slice) se [] - sOp $ Imp.Barrier Imp.FenceLocal -compileGroupExp dest e = - defCompileExp dest e - -sanityCheckLevel :: SegLevel -> InKernelGen () -sanityCheckLevel SegThread {} = return () -sanityCheckLevel SegGroup {} = - error "compileGroupOp: unexpected group-level SegOp." - -localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] -localThreadIDs dims = do - ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv - let dims' = map toInt64Exp dims - maybe (unflattenIndex dims' ltid) (map sExt64) - . M.lookup dims - . kernelLocalIdMap - . kernelConstants - <$> askEnv - -compileGroupSpace :: SegLevel -> SegSpace -> InKernelGen () -compileGroupSpace lvl space = do - sanityCheckLevel lvl - let (ltids, dims) = unzip $ unSegSpace space - zipWithM_ dPrimV_ ltids =<< localThreadIDs dims - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - dPrimV_ (segFlat space) ltid - --- Construct the necessary lock arrays for an intra-group histogram. -prepareIntraGroupSegHist :: - Count GroupSize SubExp -> - [HistOp KernelsMem] -> - InKernelGen [[Imp.TExp Int64] -> InKernelGen ()] -prepareIntraGroupSegHist group_size = - fmap snd . mapAccumLM onOp Nothing - where - onOp l op = do - constants <- kernelConstants <$> askEnv - atomicBinOp <- kernelAtomics <$> askEnv - - let local_subhistos = histDest op - - case (l, atomicUpdateLocking atomicBinOp $ histOp op) of - (_, AtomicPrim f) -> return (l, f (Space "local") local_subhistos) - (_, AtomicCAS f) -> return (l, f (Space "local") local_subhistos) - (Just l', AtomicLocking f) -> return (l, f l' (Space "local") local_subhistos) - (Nothing, AtomicLocking f) -> do - locks <- newVName "locks" - - let num_locks = toInt64Exp $ unCount group_size - dims = map toInt64Exp $ shapeDims (histShape op) ++ [histWidth op] - l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims) - locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness - - locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local" - dArray locks int32 (arrayShape locks_t) $ - ArrayIn locks_mem $ - IxFun.iota $ - map pe64 $ arrayDims locks_t - - sComment "All locks start out unlocked" $ - groupCoverSpace [kernelGroupSize constants] $ \is -> - copyDWIMFix locks is (intConst Int32 0) [] - - return (Just l', f l' (Space "local") local_subhistos) - -whenActive :: SegLevel -> SegSpace -> InKernelGen () -> InKernelGen () -whenActive lvl space m - | SegNoVirtFull <- segVirt lvl = m - | otherwise = do - group_size <- kernelGroupSize . kernelConstants <$> askEnv - -- XXX: the following check is too naive - we should also handle - -- the multi-dimensional case. - if [group_size] == map (toInt64Exp . snd) (unSegSpace space) - then m - else sWhen (isActive $ unSegSpace space) m - -compileGroupOp :: OpCompiler KernelsMem KernelEnv Imp.KernelOp -compileGroupOp pat (Alloc size space) = - kernelAlloc pat size space -compileGroupOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = - splitSpace pat o w i elems_per_thread -compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do - void $ compileGroupSpace lvl space - - whenActive lvl space $ - localOps threadOperations $ - compileStms mempty (kernelBodyStms body) $ - zipWithM_ (compileThreadResult space) (patternElements pat) $ - kernelBodyResult body - - sOp $ Imp.ErrorSync Imp.FenceLocal -compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do - compileGroupSpace lvl space - let (ltids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - - whenActive lvl space $ - compileStms mempty (kernelBodyStms body) $ - forM_ (zip (patternNames pat) $ kernelBodyResult body) $ \(dest, res) -> - copyDWIMFix - dest - (map Imp.vi64 ltids) - (kernelResultSubExp res) - [] - - sOp $ Imp.ErrorSync Imp.FenceLocal - - let segment_size = last dims' - crossesSegment from to = - (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) - - -- groupScan needs to treat the scan output as a one-dimensional - -- array of scan elements, so we invent some new flattened arrays - -- here. XXX: this assumes that the original index function is just - -- row-major, but does not actually verify it. - dims_flat <- dPrimV "dims_flat" $ product dims' - let flattened pe = do - MemLocation mem _ _ <- - entryArrayLocation <$> lookupArray (patElemName pe) - let pe_t = typeOf pe - arr_dims = Var (tvVar dims_flat) : drop (length dims') (arrayDims pe_t) - sArray - (baseString (patElemName pe) ++ "_flat") - (elemType pe_t) - (Shape arr_dims) - $ ArrayIn mem $ IxFun.iota $ map pe64 arr_dims - - num_scan_results = sum $ map (length . segBinOpNeutral) scans - - arrs_flat <- mapM flattened $ take num_scan_results $ patternElements pat - - forM_ scans $ \scan -> do - let scan_op = segBinOpLambda scan - groupScan (Just crossesSegment) (product dims') (product dims') scan_op arrs_flat -compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do - compileGroupSpace lvl space - - let (ltids, dims) = unzip $ unSegSpace space - (red_pes, map_pes) = - splitAt (segBinOpResults ops) $ patternElements pat - - dims' = map toInt64Exp dims - - mkTempArr t = - sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" - - tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops - let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs - - whenActive lvl space $ - compileStms mempty (kernelBodyStms body) $ do - let (red_res, map_res) = - splitAt (segBinOpResults ops) $ kernelBodyResult body - forM_ (zip tmp_arrs red_res) $ \(dest, res) -> - copyDWIMFix dest (map Imp.vi64 ltids) (kernelResultSubExp res) [] - zipWithM_ (compileThreadResult space) map_pes map_res - - sOp $ Imp.ErrorSync Imp.FenceLocal - - case dims' of - -- Nonsegmented case (or rather, a single segment) - this we can - -- handle directly with a group-level reduction. - [dim'] -> do - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> - groupReduce (sExt32 dim') (segBinOpLambda op) tmps - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> - copyDWIMFix (patElemName pe) [] (Var arr) [0] - _ -> do - -- Segmented intra-group reductions are turned into (regular) - -- segmented scans. It is possible that this can be done - -- better, but at least this approach is simple. - - -- groupScan operates on flattened arrays. This does not - -- involve copying anything; merely playing with the index - -- function. - dims_flat <- dPrimV "dims_flat" $ product dims' - let flatten arr = do - ArrayEntry arr_loc pt <- lookupArray arr - let flat_shape = - Shape $ - Var (tvVar dims_flat) : - drop (length ltids) (memLocationShape arr_loc) - sArray "red_arr_flat" pt flat_shape $ - ArrayIn (memLocationName arr_loc) $ - IxFun.iota $ map pe64 $ shapeDims flat_shape - - let segment_size = last dims' - crossesSegment from to = - (sExt64 to - sExt64 from) .>. (sExt64 to `rem` sExt64 segment_size) - - forM_ (zip ops tmps_for_ops) $ \(op, tmps) -> do - tmps_flat <- mapM flatten tmps - groupScan - (Just crossesSegment) - (product dims') - (product dims') - (segBinOpLambda op) - tmps_flat - - sOp $ Imp.ErrorSync Imp.FenceLocal - - forM_ (zip red_pes tmp_arrs) $ \(pe, arr) -> - copyDWIM - (patElemName pe) - [] - (Var arr) - (map (unitSlice 0) (init dims') ++ [DimFix $ last dims' -1]) - - sOp $ Imp.Barrier Imp.FenceLocal -compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do - compileGroupSpace lvl space - let ltids = map fst $ unSegSpace space - - -- We don't need the red_pes, because it is guaranteed by our type - -- rules that they occupy the same memory as the destinations for - -- the ops. - let num_red_res = length ops + sum (map (length . histNeutral) ops) - (_red_pes, map_pes) = - splitAt num_red_res $ patternElements pat - - ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops - - -- Ensure that all locks have been initialised. - sOp $ Imp.Barrier Imp.FenceLocal - - whenActive lvl space $ - compileStms mempty (kernelBodyStms kbody) $ do - let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody - (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res - zipWithM_ (compileThreadResult space) map_pes map_res - - let vs_per_op = chunks (map (length . histDest) ops) red_vs - - forM_ (zip4 red_is vs_per_op ops' ops) $ - \(bin, op_vs, do_op, HistOp dest_w _ _ _ shape lam) -> do - let bin' = toInt64Exp bin - dest_w' = toInt64Exp dest_w - bin_in_bounds = 0 .<=. bin' .&&. bin' .<. dest_w' - bin_is = map Imp.vi64 (init ltids) ++ [bin'] - vs_params = takeLast (length op_vs) $ lambdaParams lam - - sComment "perform atomic updates" $ - sWhen bin_in_bounds $ do - dLParams $ lambdaParams lam - sLoopNest shape $ \is -> do - forM_ (zip vs_params op_vs) $ \(p, v) -> - copyDWIMFix (paramName p) [] v is - do_op (bin_is ++ is) - - sOp $ Imp.ErrorSync Imp.FenceLocal -compileGroupOp pat _ = - compilerBugS $ "compileGroupOp: cannot compile rhs of binding " ++ pretty pat - -compileThreadOp :: OpCompiler KernelsMem KernelEnv Imp.KernelOp -compileThreadOp pat (Alloc size space) = - kernelAlloc pat size space -compileThreadOp pat (Inner (SizeOp (SplitSpace o w i elems_per_thread))) = - splitSpace pat o w i elems_per_thread -compileThreadOp pat _ = - compilerBugS $ "compileThreadOp: cannot compile rhs of binding " ++ pretty pat - --- | Locking strategy used for an atomic update. -data Locking = Locking - { -- | Array containing the lock. - lockingArray :: VName, - -- | Value for us to consider the lock free. - lockingIsUnlocked :: Imp.TExp Int32, - -- | What to write when we lock it. - lockingToLock :: Imp.TExp Int32, - -- | What to write when we unlock it. - lockingToUnlock :: Imp.TExp Int32, - -- | A transformation from the logical lock index to the - -- physical position in the array. This can also be used - -- to make the lock array smaller. - lockingMapping :: [Imp.TExp Int64] -> [Imp.TExp Int64] - } - --- | A function for generating code for an atomic update. Assumes --- that the bucket is in-bounds. -type DoAtomicUpdate lore r = - Space -> [VName] -> [Imp.TExp Int64] -> ImpM lore r Imp.KernelOp () - --- | The mechanism that will be used for performing the atomic update. --- Approximates how efficient it will be. Ordered from most to least --- efficient. -data AtomicUpdate lore r - = -- | Supported directly by primitive. - AtomicPrim (DoAtomicUpdate lore r) - | -- | Can be done by efficient swaps. - AtomicCAS (DoAtomicUpdate lore r) - | -- | Requires explicit locking. - AtomicLocking (Locking -> DoAtomicUpdate lore r) - --- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? -type AtomicBinOp = - BinOp -> - Maybe (VName -> VName -> Count Imp.Elements (Imp.TExp Int64) -> Imp.Exp -> Imp.AtomicOp) - --- | Do an atomic update corresponding to a binary operator lambda. -atomicUpdateLocking :: - AtomicBinOp -> - Lambda KernelsMem -> - AtomicUpdate KernelsMem KernelEnv -atomicUpdateLocking atomicBinOp lam - | Just ops_and_ts <- splitOp lam, - all (\(_, t, _, _) -> primBitSize t `elem` [32, 64]) ops_and_ts = - primOrCas ops_and_ts $ \space arrs bucket -> - -- If the operator is a vectorised binary operator on 32/64-bit - -- values, we can use a particularly efficient - -- implementation. If the operator has an atomic implementation - -- we use that, otherwise it is still a binary operator which - -- can be implemented by atomic compare-and-swap if 32/64 bits. - forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do - -- Common variables. - old <- dPrim "old" t - - (arr', _a_space, bucket_offset) <- fullyIndexArray a bucket - - case opHasAtomicSupport space (tvVar old) arr' bucket_offset op of - Just f -> sOp $ f $ Imp.var y t - Nothing -> - atomicUpdateCAS space t a (tvVar old) bucket x $ - x <~~ Imp.BinOpExp op (Imp.var x t) (Imp.var y t) - where - opHasAtomicSupport space old arr' bucket' bop = do - let atomic f = Imp.Atomic space . f old arr' bucket' - atomic <$> atomicBinOp bop - - primOrCas ops - | all isPrim ops = AtomicPrim - | otherwise = AtomicCAS - - isPrim (op, _, _, _) = isJust $ atomicBinOp op - --- If the operator functions purely on single 32/64-bit values, we can --- use an implementation based on CAS, no matter what the operator --- does. -atomicUpdateLocking _ op - | [Prim t] <- lambdaReturnType op, - [xp, _] <- lambdaParams op, - primBitSize t `elem` [32, 64] = AtomicCAS $ \space [arr] bucket -> do - old <- dPrim "old" t - atomicUpdateCAS space t arr (tvVar old) bucket (paramName xp) $ - compileBody' [xp] $ lambdaBody op -atomicUpdateLocking _ op = AtomicLocking $ \locking space arrs bucket -> do - old <- dPrim "old" int32 - continue <- dPrimVol "continue" Bool true - - -- Correctly index into locks. - (locks', _locks_space, locks_offset) <- - fullyIndexArray (lockingArray locking) $ lockingMapping locking bucket - - -- Critical section - let try_acquire_lock = - sOp $ - Imp.Atomic space $ - Imp.AtomicCmpXchg - int32 - (tvVar old) - locks' - locks_offset - (untyped $ lockingIsUnlocked locking) - (untyped $ lockingToLock locking) - lock_acquired = tvExp old .==. lockingIsUnlocked locking - -- Even the releasing is done with an atomic rather than a - -- simple write, for memory coherency reasons. - release_lock = - sOp $ - Imp.Atomic space $ - Imp.AtomicCmpXchg - int32 - (tvVar old) - locks' - locks_offset - (untyped $ lockingToLock locking) - (untyped $ lockingToUnlock locking) - break_loop = continue <-- false - - -- Preparing parameters. It is assumed that the caller has already - -- filled the arr_params. We copy the current value to the - -- accumulator parameters. - -- - -- Note the use of 'everythingVolatile' when reading and writing the - -- buckets. This was necessary to ensure correct execution on a - -- newer NVIDIA GPU (RTX 2080). The 'volatile' modifiers likely - -- make the writes pass through the (SM-local) L1 cache, which is - -- necessary here, because we are really doing device-wide - -- synchronisation without atomics (naughty!). - let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op - bind_acc_params = - everythingVolatile $ - sComment "bind lhs" $ - forM_ (zip acc_params arrs) $ \(acc_p, arr) -> - copyDWIMFix (paramName acc_p) [] (Var arr) bucket - - let op_body = - sComment "execute operation" $ - compileBody' acc_params $ lambdaBody op - - do_hist = - everythingVolatile $ - sComment "update global result" $ - zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params - - fence = case space of - Space "local" -> sOp $ Imp.MemFence Imp.FenceLocal - _ -> sOp $ Imp.MemFence Imp.FenceGlobal - - -- While-loop: Try to insert your value - sWhile (tvExp continue) $ do - try_acquire_lock - sWhen lock_acquired $ do - dLParams acc_params - bind_acc_params - op_body - do_hist - fence - release_lock - break_loop - fence - where - writeArray bucket arr val = copyDWIMFix arr bucket val [] - -atomicUpdateCAS :: - Space -> - PrimType -> - VName -> - VName -> - [Imp.TExp Int64] -> - VName -> - InKernelGen () -> - InKernelGen () -atomicUpdateCAS space t arr old bucket x do_op = do - -- Code generation target: - -- - -- old = d_his[idx]; - -- do { - -- assumed = old; - -- x = do_op(assumed, y); - -- old = atomicCAS(&d_his[idx], assumed, tmp); - -- } while(assumed != old); - assumed <- tvVar <$> dPrim "assumed" t - run_loop <- dPrimV "run_loop" true - - -- XXX: CUDA may generate really bad code if this is not a volatile - -- read. Unclear why. The later reads are volatile, so maybe - -- that's it. - everythingVolatile $ copyDWIMFix old [] (Var arr) bucket - - (arr', _a_space, bucket_offset) <- fullyIndexArray arr bucket - - -- While-loop: Try to insert your value - let (toBits, fromBits) = - case t of - FloatType Float32 -> - ( \v -> Imp.FunExp "to_bits32" [v] int32, - \v -> Imp.FunExp "from_bits32" [v] t - ) - FloatType Float64 -> - ( \v -> Imp.FunExp "to_bits64" [v] int64, - \v -> Imp.FunExp "from_bits64" [v] t - ) - _ -> (id, id) - - int - | primBitSize t == 32 = int32 - | otherwise = int64 - - sWhile (tvExp run_loop) $ do - assumed <~~ Imp.var old t - x <~~ Imp.var assumed t - do_op - old_bits_v <- newVName "old_bits" - dPrim_ old_bits_v int - let old_bits = Imp.var old_bits_v int - sOp $ - Imp.Atomic space $ - Imp.AtomicCmpXchg - int - old_bits_v - arr' - bucket_offset - (toBits (Imp.var assumed t)) - (toBits (Imp.var x t)) - old <~~ fromBits old_bits - let won = CmpOpExp (CmpEq int) (toBits (Imp.var assumed t)) old_bits - sWhen (isBool won) (run_loop <-- false) - --- | Horizontally fission a lambda that models a binary operator. -splitOp :: ASTLore lore => Lambda lore -> Maybe [(BinOp, PrimType, VName, VName)] -splitOp lam = mapM splitStm $ bodyResult $ lambdaBody lam - where - n = length $ lambdaReturnType lam - splitStm (Var res) = do - Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y))) <- - find (([res] ==) . patternNames . stmPattern) $ - stmsToList $ bodyStms $ lambdaBody lam - i <- Var res `elemIndex` bodyResult (lambdaBody lam) - xp <- maybeNth i $ lambdaParams lam - yp <- maybeNth (n + i) $ lambdaParams lam - guard $ paramName xp == x - guard $ paramName yp == y - Prim t <- Just $ patElemType pe - return (op, t, paramName xp, paramName yp) - splitStm _ = Nothing - -computeKernelUses :: - FreeIn a => - a -> - [VName] -> - CallKernelGen [Imp.KernelUse] -computeKernelUses kernel_body bound_in_kernel = do - let actually_free = freeIn kernel_body `namesSubtract` namesFromList bound_in_kernel - -- Compute the variables that we need to pass to the kernel. - nubOrd <$> readsFromSet actually_free - -readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] -readsFromSet free = - fmap catMaybes $ - forM (namesToList free) $ \var -> do - t <- lookupType var - vtable <- getVTable - case t of - Array {} -> return Nothing - Acc {} -> return Nothing - Mem (Space "local") -> return Nothing - Mem {} -> return $ Just $ Imp.MemoryUse var - Prim bt -> - isConstExp vtable (Imp.var var bt) >>= \case - Just ce -> return $ Just $ Imp.ConstUse var ce - Nothing -> return $ Just $ Imp.ScalarUse var bt - -isConstExp :: - VTable KernelsMem -> - Imp.Exp -> - ImpM lore r op (Maybe Imp.KernelConstExp) -isConstExp vtable size = do - fname <- askFunction - let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name - onLeaf Imp.Index {} _ = Nothing - lookupConstExp name = - constExp =<< hasExp =<< M.lookup name vtable - constExp (Op (Inner (SizeOp (GetSize key _)))) = - Just $ LeafExp (Imp.SizeConst $ keyWithEntryPoint fname key) int32 - constExp e = primExpFromExp lookupConstExp e - return $ replaceInPrimExpM onLeaf size - where - hasExp (ArrayVar e _) = e - hasExp (AccVar e _) = e - hasExp (ScalarVar e _) = e - hasExp (MemVar e _) = e - -computeThreadChunkSize :: - SplitOrdering -> - Imp.TExp Int64 -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - TV Int64 -> - ImpM lore r op () -computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = - chunk_var - <-- sMin64 - (Imp.unCount elements_per_thread) - ((Imp.unCount num_elements - thread_index) `divUp` toInt64Exp stride) -computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do - starting_point <- - dPrimV "starting_point" $ - thread_index * Imp.unCount elements_per_thread - remaining_elements <- - dPrimV "remaining_elements" $ - Imp.unCount num_elements - tvExp starting_point - - let no_remaining_elements = tvExp remaining_elements .<=. 0 - beyond_bounds = Imp.unCount num_elements .<=. tvExp starting_point - - sIf - (no_remaining_elements .||. beyond_bounds) - (chunk_var <-- 0) - ( sIf - is_last_thread - (chunk_var <-- Imp.unCount last_thread_elements) - (chunk_var <-- Imp.unCount elements_per_thread) - ) - where - last_thread_elements = - num_elements - Imp.elements thread_index * elements_per_thread - is_last_thread = - Imp.unCount num_elements - .<. (thread_index + 1) * Imp.unCount elements_per_thread - -kernelInitialisationSimple :: - Count NumGroups (Imp.TExp Int64) -> - Count GroupSize (Imp.TExp Int64) -> - CallKernelGen (KernelConstants, InKernelGen ()) -kernelInitialisationSimple (Count num_groups) (Count group_size) = do - global_tid <- newVName "global_tid" - local_tid <- newVName "local_tid" - group_id <- newVName "group_tid" - wave_size <- newVName "wave_size" - inner_group_size <- newVName "group_size" - let constants = - KernelConstants - (Imp.vi32 global_tid) - (Imp.vi32 local_tid) - (Imp.vi32 group_id) - global_tid - local_tid - group_id - num_groups - group_size - (sExt32 (group_size * num_groups)) - (Imp.vi32 wave_size) - true - mempty - - let set_constants = do - dPrim_ global_tid int32 - dPrim_ local_tid int32 - dPrim_ inner_group_size int64 - dPrim_ wave_size int32 - dPrim_ group_id int32 - - sOp (Imp.GetGlobalId global_tid 0) - sOp (Imp.GetLocalId local_tid 0) - sOp (Imp.GetLocalSize inner_group_size 0) - sOp (Imp.GetLockstepWidth wave_size) - sOp (Imp.GetGroupId group_id 0) - - return (constants, set_constants) - -isActive :: [(VName, SubExp)] -> Imp.TExp Bool -isActive limit = case actives of - [] -> true - x : xs -> foldl (.&&.) x xs - where - (is, ws) = unzip limit - actives = zipWith active is $ map toInt64Exp ws - active i = (Imp.vi64 i .<.) - --- | Change every memory block to be in the global address space, --- except those who are in the local memory space. This only affects --- generated code - we still need to make sure that the memory is --- actually present on the device (and dared as variables in the --- kernel). -makeAllMemoryGlobal :: CallKernelGen a -> CallKernelGen a -makeAllMemoryGlobal = - localDefaultSpace (Imp.Space "global") . localVTable (M.map globalMemory) - where - globalMemory (MemVar _ entry) - | entryMemSpace entry /= Space "local" = - MemVar Nothing entry {entryMemSpace = Imp.Space "global"} - globalMemory entry = - entry - -groupReduce :: - Imp.TExp Int32 -> - Lambda KernelsMem -> - [VName] -> - InKernelGen () -groupReduce w lam arrs = do - offset <- dPrim "offset" int32 - groupReduceWithOffset offset w lam arrs - -groupReduceWithOffset :: - TV Int32 -> - Imp.TExp Int32 -> - Lambda KernelsMem -> - [VName] -> - InKernelGen () -groupReduceWithOffset offset w lam arrs = do - constants <- kernelConstants <$> askEnv - - let local_tid = kernelLocalThreadId constants - global_tid = kernelGlobalThreadId constants - - barrier - | all primType $ lambdaReturnType lam = sOp $ Imp.Barrier Imp.FenceLocal - | otherwise = sOp $ Imp.Barrier Imp.FenceGlobal - - readReduceArgument param arr - | Prim _ <- paramType param = do - let i = local_tid + tvExp offset - copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] - | otherwise = do - let i = global_tid + tvExp offset - copyDWIMFix (paramName param) [] (Var arr) [sExt64 i] - - writeReduceOpResult param arr - | Prim _ <- paramType param = - copyDWIMFix arr [sExt64 local_tid] (Var $ paramName param) [] - | otherwise = - return () - - let (reduce_acc_params, reduce_arr_params) = splitAt (length arrs) $ lambdaParams lam - - skip_waves <- dPrimV "skip_waves" (1 :: Imp.TExp Int32) - dLParams $ lambdaParams lam - - offset <-- (0 :: Imp.TExp Int32) - - comment "participating threads read initial accumulator" $ - sWhen (local_tid .<. w) $ - zipWithM_ readReduceArgument reduce_acc_params arrs - - let do_reduce = do - comment "read array element" $ - zipWithM_ readReduceArgument reduce_arr_params arrs - comment "apply reduction operation" $ - compileBody' reduce_acc_params $ lambdaBody lam - comment "write result of operation" $ - zipWithM_ writeReduceOpResult reduce_acc_params arrs - in_wave_reduce = everythingVolatile do_reduce - - wave_size = kernelWaveSize constants - group_size = kernelGroupSize constants - wave_id = local_tid `quot` wave_size - in_wave_id = local_tid - wave_id * wave_size - num_waves = (sExt32 group_size + wave_size - 1) `quot` wave_size - arg_in_bounds = local_tid + tvExp offset .<. w - - doing_in_wave_reductions = - tvExp offset .<. wave_size - apply_in_in_wave_iteration = - (in_wave_id .&. (2 * tvExp offset - 1)) .==. 0 - in_wave_reductions = do - offset <-- (1 :: Imp.TExp Int32) - sWhile doing_in_wave_reductions $ do - sWhen - (arg_in_bounds .&&. apply_in_in_wave_iteration) - in_wave_reduce - offset <-- tvExp offset * 2 - - doing_cross_wave_reductions = - tvExp skip_waves .<. num_waves - is_first_thread_in_wave = - in_wave_id .==. 0 - wave_not_skipped = - (wave_id .&. (2 * tvExp skip_waves - 1)) .==. 0 - apply_in_cross_wave_iteration = - arg_in_bounds .&&. is_first_thread_in_wave .&&. wave_not_skipped - cross_wave_reductions = - sWhile doing_cross_wave_reductions $ do - barrier - offset <-- tvExp skip_waves * wave_size - sWhen - apply_in_cross_wave_iteration - do_reduce - skip_waves <-- tvExp skip_waves * 2 - - in_wave_reductions - cross_wave_reductions - -groupScan :: - Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> - Imp.TExp Int64 -> - Imp.TExp Int64 -> - Lambda KernelsMem -> - [VName] -> - InKernelGen () -groupScan seg_flag arrs_full_size w lam arrs = do - constants <- kernelConstants <$> askEnv - renamed_lam <- renameLambda lam - - let ltid32 = kernelLocalThreadId constants - ltid = sExt64 ltid32 - (x_params, y_params) = splitAt (length arrs) $ lambdaParams lam - - dLParams (lambdaParams lam ++ lambdaParams renamed_lam) - - ltid_in_bounds <- dPrimVE "ltid_in_bounds" $ ltid .<. w - - -- The scan works by splitting the group into blocks, which are - -- scanned separately. Typically, these blocks are smaller than - -- the lockstep width, which enables barrier-free execution inside - -- them. - -- - -- We hardcode the block size here. The only requirement is that - -- it should not be less than the square root of the group size. - -- With 32, we will work on groups of size 1024 or smaller, which - -- fits every device Troels has seen. Still, it would be nicer if - -- it were a runtime parameter. Some day. - let block_size = 32 - simd_width = kernelWaveSize constants - block_id = ltid32 `quot` block_size - in_block_id = ltid32 - block_id * block_size - doInBlockScan seg_flag' active = - inBlockScan - constants - seg_flag' - arrs_full_size - simd_width - block_size - active - arrs - barrier - array_scan = not $ all primType $ lambdaReturnType lam - barrier - | array_scan = - sOp $ Imp.Barrier Imp.FenceGlobal - | otherwise = - sOp $ Imp.Barrier Imp.FenceLocal - - group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants - - writeBlockResult p arr - | primType $ paramType p = - copyDWIM arr [DimFix $ sExt64 block_id] (Var $ paramName p) [] - | otherwise = - copyDWIM arr [DimFix $ group_offset + sExt64 block_id] (Var $ paramName p) [] - - readPrevBlockResult p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix $ sExt64 block_id - 1] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix $ group_offset + sExt64 block_id - 1] - - doInBlockScan seg_flag ltid_in_bounds lam - barrier - - let is_first_block = block_id .==. 0 - when array_scan $ do - sComment "save correct values for first block" $ - sWhen is_first_block $ - forM_ (zip x_params arrs) $ \(x, arr) -> - unless (primType $ paramType x) $ - copyDWIM arr [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] (Var $ paramName x) [] - - barrier - - let last_in_block = in_block_id .==. block_size - 1 - sComment "last thread of block 'i' writes its result to offset 'i'" $ - sWhen (last_in_block .&&. ltid_in_bounds) $ - everythingVolatile $ - zipWithM_ writeBlockResult x_params arrs - - barrier - - let first_block_seg_flag = do - flag_true <- seg_flag - Just $ \from to -> - flag_true (from * block_size + block_size -1) (to * block_size + block_size -1) - comment - "scan the first block, after which offset 'i' contains carry-in for block 'i+1'" - $ doInBlockScan first_block_seg_flag (is_first_block .&&. ltid_in_bounds) renamed_lam - - barrier - - when array_scan $ do - sComment "move correct values for first block back a block" $ - sWhen is_first_block $ - forM_ (zip x_params arrs) $ \(x, arr) -> - unless (primType $ paramType x) $ - copyDWIM - arr - [DimFix $ arrs_full_size + group_offset + ltid] - (Var arr) - [DimFix $ arrs_full_size + group_offset + sExt64 block_size + ltid] - - barrier - - let read_carry_in = do - forM_ (zip x_params y_params) $ \(x, y) -> - copyDWIM (paramName y) [] (Var (paramName x)) [] - zipWithM_ readPrevBlockResult x_params arrs - - y_to_x = forM_ (zip x_params y_params) $ \(x, y) -> - when (primType (paramType x)) $ - copyDWIM (paramName x) [] (Var (paramName y)) [] - - op_to_x - | Nothing <- seg_flag = - compileBody' x_params $ lambdaBody lam - | Just flag_true <- seg_flag = do - inactive <- - dPrimVE "inactive" $ flag_true (block_id * block_size -1) ltid32 - sWhen inactive y_to_x - when array_scan barrier - sUnless inactive $ compileBody' x_params $ lambdaBody lam - - write_final_result = - forM_ (zip x_params arrs) $ \(p, arr) -> - when (primType $ paramType p) $ - copyDWIM arr [DimFix ltid] (Var $ paramName p) [] - - sComment "carry-in for every block except the first" $ - sUnless (is_first_block .||. bNot ltid_in_bounds) $ do - sComment "read operands" read_carry_in - sComment "perform operation" op_to_x - sComment "write final result" write_final_result - - barrier - - sComment "restore correct values for first block" $ - sWhen is_first_block $ - forM_ (zip3 x_params y_params arrs) $ \(x, y, arr) -> - if primType (paramType y) - then copyDWIM arr [DimFix ltid] (Var $ paramName y) [] - else copyDWIM (paramName x) [] (Var arr) [DimFix $ arrs_full_size + group_offset + ltid] - - barrier - -inBlockScan :: - KernelConstants -> - Maybe (Imp.TExp Int32 -> Imp.TExp Int32 -> Imp.TExp Bool) -> - Imp.TExp Int64 -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> - Imp.TExp Bool -> - [VName] -> - InKernelGen () -> - Lambda KernelsMem -> - InKernelGen () -inBlockScan constants seg_flag arrs_full_size lockstep_width block_size active arrs barrier scan_lam = everythingVolatile $ do - skip_threads <- dPrim "skip_threads" int32 - let in_block_thread_active = - tvExp skip_threads .<=. in_block_id - actual_params = lambdaParams scan_lam - (x_params, y_params) = - splitAt (length actual_params `div` 2) actual_params - y_to_x = - forM_ (zip x_params y_params) $ \(x, y) -> - when (primType (paramType x)) $ - copyDWIM (paramName x) [] (Var (paramName y)) [] - - -- Set initial y values - sComment "read input for in-block scan" $ - sWhen active $ do - zipWithM_ readInitial y_params arrs - -- Since the final result is expected to be in x_params, we may - -- need to copy it there for the first thread in the block. - sWhen (in_block_id .==. 0) y_to_x - - when array_scan barrier - - let op_to_x - | Nothing <- seg_flag = - compileBody' x_params $ lambdaBody scan_lam - | Just flag_true <- seg_flag = do - inactive <- - dPrimVE "inactive" $ - flag_true (ltid32 - tvExp skip_threads) ltid32 - sWhen inactive y_to_x - when array_scan barrier - sUnless inactive $ compileBody' x_params $ lambdaBody scan_lam - - maybeBarrier = - sWhen - (lockstep_width .<=. tvExp skip_threads) - barrier - - sComment "in-block scan (hopefully no barriers needed)" $ do - skip_threads <-- 1 - sWhile (tvExp skip_threads .<. block_size) $ do - sWhen (in_block_thread_active .&&. active) $ do - sComment "read operands" $ - zipWithM_ (readParam (sExt64 $ tvExp skip_threads)) x_params arrs - sComment "perform operation" op_to_x - - maybeBarrier - - sWhen (in_block_thread_active .&&. active) $ - sComment "write result" $ - sequence_ $ zipWith3 writeResult x_params y_params arrs - - maybeBarrier - - skip_threads <-- tvExp skip_threads * 2 - where - block_id = ltid32 `quot` block_size - in_block_id = ltid32 - block_id * block_size - ltid32 = kernelLocalThreadId constants - ltid = sExt64 ltid32 - gtid = sExt64 $ kernelGlobalThreadId constants - array_scan = not $ all primType $ lambdaReturnType scan_lam - - readInitial p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix ltid] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix gtid] - - readParam behind p arr - | primType $ paramType p = - copyDWIM (paramName p) [] (Var arr) [DimFix $ ltid - behind] - | otherwise = - copyDWIM (paramName p) [] (Var arr) [DimFix $ gtid - behind + arrs_full_size] - - writeResult x y arr - | primType $ paramType x = do - copyDWIM arr [DimFix ltid] (Var $ paramName x) [] - copyDWIM (paramName y) [] (Var $ paramName x) [] - | otherwise = - copyDWIM (paramName y) [] (Var $ paramName x) [] - -computeMapKernelGroups :: Imp.TExp Int64 -> CallKernelGen (Imp.TExp Int64, Imp.TExp Int64) -computeMapKernelGroups kernel_size = do - group_size <- dPrim "group_size" int64 - fname <- askFunction - let group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty $ tvVar group_size - sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup - num_groups <- dPrimV "num_groups" $ kernel_size `divUp` tvExp group_size - return (tvExp num_groups, tvExp group_size) - -simpleKernelConstants :: - Imp.TExp Int64 -> - String -> - CallKernelGen (KernelConstants, InKernelGen ()) -simpleKernelConstants kernel_size desc = do - thread_gtid <- newVName $ desc ++ "_gtid" - thread_ltid <- newVName $ desc ++ "_ltid" - group_id <- newVName $ desc ++ "_gid" - (num_groups, group_size) <- computeMapKernelGroups kernel_size - let set_constants = do - dPrim_ thread_gtid int32 - dPrim_ thread_ltid int32 - dPrim_ group_id int32 - sOp (Imp.GetGlobalId thread_gtid 0) - sOp (Imp.GetLocalId thread_ltid 0) - sOp (Imp.GetGroupId group_id 0) - - return - ( KernelConstants - (Imp.vi32 thread_gtid) - (Imp.vi32 thread_ltid) - (Imp.vi32 group_id) - thread_gtid - thread_ltid - group_id - num_groups - group_size - (sExt32 (group_size * num_groups)) - 0 - (Imp.vi64 thread_gtid .<. kernel_size) - mempty, - set_constants - ) - --- | For many kernels, we may not have enough physical groups to cover --- the logical iteration space. Some groups thus have to perform --- double duty; we put an outer loop to accomplish this. The --- advantage over just launching a bazillion threads is that the cost --- of memory expansion should be proportional to the number of --- *physical* threads (hardware parallelism), not the amount of --- application parallelism. -virtualiseGroups :: - SegVirt -> - Imp.TExp Int32 -> - (Imp.TExp Int32 -> InKernelGen ()) -> - InKernelGen () -virtualiseGroups SegVirt required_groups m = do - constants <- kernelConstants <$> askEnv - phys_group_id <- dPrim "phys_group_id" int32 - sOp $ Imp.GetGroupId (tvVar phys_group_id) 0 - let iterations = - (required_groups - tvExp phys_group_id) - `divUp` sExt32 (kernelNumGroups constants) - - sFor "i" iterations $ \i -> do - m . tvExp - =<< dPrimV - "virt_group_id" - (tvExp phys_group_id + i * sExt32 (kernelNumGroups constants)) - -- Make sure the virtual group is actually done before we let - -- another virtual group have its way with it. - sOp $ Imp.Barrier Imp.FenceGlobal -virtualiseGroups _ _ m = do - gid <- kernelGroupIdVar . kernelConstants <$> askEnv - m $ Imp.vi32 gid - -sKernelThread :: - String -> - Count NumGroups (Imp.TExp Int64) -> - Count GroupSize (Imp.TExp Int64) -> - VName -> - InKernelGen () -> - CallKernelGen () -sKernelThread = sKernel threadOperations kernelGlobalThreadId - -sKernelGroup :: - String -> - Count NumGroups (Imp.TExp Int64) -> - Count GroupSize (Imp.TExp Int64) -> - VName -> - InKernelGen () -> - CallKernelGen () -sKernelGroup = sKernel groupOperations kernelGroupId - -sKernelFailureTolerant :: - Bool -> - Operations KernelsMem KernelEnv Imp.KernelOp -> - KernelConstants -> - Name -> - InKernelGen () -> - CallKernelGen () -sKernelFailureTolerant tol ops constants name m = do - HostEnv atomics _ locks <- askEnv - body <- makeAllMemoryGlobal $ subImpM_ (KernelEnv atomics constants locks) ops m - uses <- computeKernelUses body mempty - emit $ - Imp.Op $ - Imp.CallKernel - Imp.Kernel - { Imp.kernelBody = body, - Imp.kernelUses = uses, - Imp.kernelNumGroups = [untyped $ kernelNumGroups constants], - Imp.kernelGroupSize = [untyped $ kernelGroupSize constants], - Imp.kernelName = name, - Imp.kernelFailureTolerant = tol - } - -sKernel :: - Operations KernelsMem KernelEnv Imp.KernelOp -> - (KernelConstants -> Imp.TExp Int32) -> - String -> - Count NumGroups (Imp.TExp Int64) -> - Count GroupSize (Imp.TExp Int64) -> - VName -> - InKernelGen () -> - CallKernelGen () -sKernel ops flatf name num_groups group_size v f = do - (constants, set_constants) <- kernelInitialisationSimple num_groups group_size - name' <- nameForFun $ name ++ "_" ++ show (baseTag v) - sKernelFailureTolerant False ops constants name' $ do - set_constants - dPrimV_ v $ flatf constants - f - -copyInGroup :: CopyCompiler KernelsMem KernelEnv Imp.KernelOp -copyInGroup pt destloc destslice srcloc srcslice = do - dest_space <- entryMemSpace <$> lookupMemory (memLocationName destloc) - src_space <- entryMemSpace <$> lookupMemory (memLocationName srcloc) - - case (dest_space, src_space) of - (ScalarSpace destds _, ScalarSpace srcds _) -> do - let destslice' = - replicate (length destslice - length destds) (DimFix 0) - ++ takeLast (length destds) destslice - srcslice' = - replicate (length srcslice - length srcds) (DimFix 0) - ++ takeLast (length srcds) srcslice - copyElementWise pt destloc destslice' srcloc srcslice' - _ -> do - groupCoverSpace (sliceDims destslice) $ \is -> - copyElementWise - pt - destloc - (map DimFix $ fixSlice destslice is) - srcloc - (map DimFix $ fixSlice srcslice is) - sOp $ Imp.Barrier Imp.FenceLocal - -threadOperations, groupOperations :: Operations KernelsMem KernelEnv Imp.KernelOp -threadOperations = - (defaultOperations compileThreadOp) - { opsCopyCompiler = copyElementWise, - opsExpCompiler = compileThreadExp, - opsStmsCompiler = \_ -> defCompileStms mempty, - opsAllocCompilers = - M.fromList [(Space "local", allocLocal)] - } -groupOperations = - (defaultOperations compileGroupOp) - { opsCopyCompiler = copyInGroup, - opsExpCompiler = compileGroupExp, - opsStmsCompiler = \_ -> defCompileStms mempty, - opsAllocCompilers = - M.fromList [(Space "local", allocLocal)] - } - --- | Perform a Replicate with a kernel. -sReplicateKernel :: VName -> SubExp -> CallKernelGen () -sReplicateKernel arr se = do - t <- subExpType se - ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr - - let dims = map toInt64Exp $ ds ++ arrayDims t - (constants, set_constants) <- - simpleKernelConstants (product $ map sExt64 dims) "replicate" - - fname <- askFunction - let name = - keyWithEntryPoint fname $ - nameFromString $ - "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - is' = unflattenIndex dims $ sExt64 $ kernelGlobalThreadId constants - - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - sWhen (kernelThreadActive constants) $ - copyDWIMFix arr is' se $ drop (length ds) is' - -replicateName :: PrimType -> String -replicateName bt = "replicate_" ++ pretty bt - -replicateForType :: PrimType -> CallKernelGen Name -replicateForType bt = do - let fname = nameFromString $ "builtin#" <> replicateName bt - - exists <- hasFunction fname - unless exists $ do - mem <- newVName "mem" - num_elems <- newVName "num_elems" - val <- newVName "val" - - let params = - [ Imp.MemParam mem (Space "device"), - Imp.ScalarParam num_elems int32, - Imp.ScalarParam val bt - ] - shape = Shape [Var num_elems] - function fname [] params $ do - arr <- - sArray "arr" bt shape $ - ArrayIn mem $ - IxFun.iota $ - map pe64 $ shapeDims shape - sReplicateKernel arr $ Var val - - return fname - -replicateIsFill :: VName -> SubExp -> CallKernelGen (Maybe (CallKernelGen ())) -replicateIsFill arr v = do - ArrayEntry (MemLocation arr_mem arr_shape arr_ixfun) _ <- lookupArray arr - v_t <- subExpType v - case v_t of - Prim v_t' - | IxFun.isLinear arr_ixfun -> return $ - Just $ do - fname <- replicateForType v_t' - emit $ - Imp.Call - [] - fname - [ Imp.MemArg arr_mem, - Imp.ExpArg $ untyped $ product $ map toInt64Exp arr_shape, - Imp.ExpArg $ toExp' v_t' v - ] - _ -> return Nothing - --- | Perform a Replicate with a kernel. -sReplicate :: VName -> SubExp -> CallKernelGen () -sReplicate arr se = do - -- If the replicate is of a particularly common and simple form - -- (morally a memset()/fill), then we use a common function. - is_fill <- replicateIsFill arr se - - case is_fill of - Just m -> m - Nothing -> sReplicateKernel arr se - --- | Perform an Iota with a kernel. -sIotaKernel :: - VName -> - Imp.TExp Int64 -> - Imp.Exp -> - Imp.Exp -> - IntType -> - CallKernelGen () -sIotaKernel arr n x s et = do - destloc <- entryArrayLocation <$> lookupArray arr - (constants, set_constants) <- simpleKernelConstants n "iota" - - fname <- askFunction - let name = - keyWithEntryPoint fname $ - nameFromString $ - "iota_" ++ pretty et ++ "_" - ++ show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - let gtid = sExt64 $ kernelGlobalThreadId constants - sWhen (kernelThreadActive constants) $ do - (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] - - emit $ - Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ - BinOpExp - (Add et OverflowWrap) - (BinOpExp (Mul et OverflowWrap) (Imp.sExt et $ untyped gtid) s) - x - -iotaName :: IntType -> String -iotaName bt = "iota_" ++ pretty bt - -iotaForType :: IntType -> CallKernelGen Name -iotaForType bt = do - let fname = nameFromString $ "builtin#" <> iotaName bt - - exists <- hasFunction fname - unless exists $ do - mem <- newVName "mem" - n <- newVName "n" - x <- newVName "x" - s <- newVName "s" - - let params = - [ Imp.MemParam mem (Space "device"), - Imp.ScalarParam n int32, - Imp.ScalarParam x $ IntType bt, - Imp.ScalarParam s $ IntType bt - ] - shape = Shape [Var n] - n' = Imp.vi64 n - x' = Imp.var x $ IntType bt - s' = Imp.var s $ IntType bt - - function fname [] params $ do - arr <- - sArray "arr" (IntType bt) shape $ - ArrayIn mem $ - IxFun.iota $ - map pe64 $ shapeDims shape - sIotaKernel arr (sExt64 n') x' s' bt - - return fname - --- | Perform an Iota with a kernel. -sIota :: - VName -> - Imp.TExp Int64 -> - Imp.Exp -> - Imp.Exp -> - IntType -> - CallKernelGen () -sIota arr n x s et = do - ArrayEntry (MemLocation arr_mem _ arr_ixfun) _ <- lookupArray arr - if IxFun.isLinear arr_ixfun - then do - fname <- iotaForType et - emit $ - Imp.Call - [] - fname - [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s] - else sIotaKernel arr n x s et - -sCopy :: CopyCompiler KernelsMem HostEnv Imp.HostOp -sCopy - bt - destloc@(MemLocation destmem _ _) - destslice - srcloc@(MemLocation srcmem _ _) - srcslice = - do - -- Note that the shape of the destination and the source are - -- necessarily the same. - let shape = sliceDims srcslice - kernel_size = product shape - - (constants, set_constants) <- simpleKernelConstants kernel_size "copy" - - fname <- askFunction - let name = - keyWithEntryPoint fname $ - nameFromString $ - "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - - let gtid = sExt64 $ kernelGlobalThreadId constants - dest_is = unflattenIndex shape gtid - src_is = dest_is - - (_, destspace, destidx) <- - fullyIndexArray' destloc $ fixSlice destslice dest_is - (_, srcspace, srcidx) <- - fullyIndexArray' srcloc $ fixSlice srcslice src_is - - sWhen (gtid .<. kernel_size) $ - emit $ - Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ - Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile - -compileGroupResult :: - SegSpace -> - PatElem KernelsMem -> - KernelResult -> - InKernelGen () -compileGroupResult _ pe (TileReturns [(w, per_group_elems)] what) = do - n <- toInt64Exp . arraySize 0 <$> lookupType what - - constants <- kernelConstants <$> askEnv - let ltid = sExt64 $ kernelLocalThreadId constants - offset = - toInt64Exp per_group_elems - * sExt64 (kernelGroupId constants) - - -- Avoid loop for the common case where each thread is statically - -- known to write at most one element. - localOps threadOperations $ - if toInt64Exp per_group_elems == kernelGroupSize constants - then - sWhen (ltid + offset .<. toInt64Exp w) $ - copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid] - else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do - j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid - sWhen (j + offset .<. toInt64Exp w) $ - copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] -compileGroupResult space pe (TileReturns dims what) = do - let gids = map fst $ unSegSpace space - out_tile_sizes = map (toInt64Exp . snd) dims - group_is = zipWith (*) (map Imp.vi64 gids) out_tile_sizes - local_is <- localThreadIDs $ map snd dims - is_for_thread <- - mapM (dPrimV "thread_out_index") $ - zipWith (+) group_is local_is - - localOps threadOperations $ - sWhen (isActive $ zip (map tvVar is_for_thread) $ map fst dims) $ - copyDWIMFix (patElemName pe) (map tvExp is_for_thread) (Var what) local_is -compileGroupResult space pe (RegTileReturns dims_n_tiles what) = do - constants <- kernelConstants <$> askEnv - - let gids = map fst $ unSegSpace space - (dims, group_tiles, reg_tiles) = unzip3 dims_n_tiles - group_tiles' = map toInt64Exp group_tiles - reg_tiles' = map toInt64Exp reg_tiles - - -- Which group tile is this group responsible for? - let group_tile_is = map Imp.vi64 gids - - -- Within the group tile, which register tile is this thread - -- responsible for? - reg_tile_is <- - mapM (dPrimVE "reg_tile_i") $ - unflattenIndex group_tiles' $ sExt64 $ kernelLocalThreadId constants - - -- Compute output array slice for the register tile belonging to - -- this thread. - let regTileSliceDim (group_tile, group_tile_i) (reg_tile, reg_tile_i) = do - tile_dim_start <- - dPrimVE "tile_dim_start" $ - reg_tile * (group_tile * group_tile_i + reg_tile_i) - return $ DimSlice tile_dim_start reg_tile 1 - reg_tile_slices <- - zipWithM - regTileSliceDim - (zip group_tiles' group_tile_is) - (zip reg_tiles' reg_tile_is) - - localOps threadOperations $ - sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do - let dest_is = fixSlice reg_tile_slices is_in_reg_tile - src_is = reg_tile_is ++ is_in_reg_tile - sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map toInt64Exp dims) $ - copyDWIMFix (patElemName pe) dest_is (Var what) src_is -compileGroupResult space pe (Returns _ what) = do - constants <- kernelConstants <$> askEnv - in_local_memory <- arrayInLocalMemory what - let gids = map (Imp.vi64 . fst) $ unSegSpace space - - if not in_local_memory - then - localOps threadOperations $ - sWhen (kernelLocalThreadId constants .==. 0) $ - copyDWIMFix (patElemName pe) gids what [] - else -- If the result of the group is an array in local memory, we - -- store it by collective copying among all the threads of the - -- group. TODO: also do this if the array is in global memory - -- (but this is a bit more tricky, synchronisation-wise). - copyDWIMFix (patElemName pe) gids what [] -compileGroupResult _ _ WriteReturns {} = - compilerLimitationS "compileGroupResult: WriteReturns not handled yet." -compileGroupResult _ _ ConcatReturns {} = - compilerLimitationS "compileGroupResult: ConcatReturns not handled yet." - -compileThreadResult :: - SegSpace -> - PatElem KernelsMem -> - KernelResult -> - InKernelGen () -compileThreadResult _ _ RegTileReturns {} = - compilerLimitationS "compileThreadResult: RegTileReturns not yet handled." -compileThreadResult space pe (Returns _ what) = do - let is = map (Imp.vi64 . fst) $ unSegSpace space - copyDWIMFix (patElemName pe) is what [] -compileThreadResult _ pe (ConcatReturns SplitContiguous _ per_thread_elems what) = do - constants <- kernelConstants <$> askEnv - let offset = - toInt64Exp per_thread_elems - * sExt64 (kernelGlobalThreadId constants) - n <- toInt64Exp . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) [] -compileThreadResult _ pe (ConcatReturns (SplitStrided stride) _ _ what) = do - offset <- sExt64 . kernelGlobalThreadId . kernelConstants <$> askEnv - n <- toInt64Exp . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n $ toInt64Exp stride] (Var what) [] -compileThreadResult _ pe (WriteReturns (Shape rws) _arr dests) = do - constants <- kernelConstants <$> askEnv - let rws' = map toInt64Exp rws - forM_ dests $ \(slice, e) -> do - let slice' = map (fmap toInt64Exp) slice - condInBounds (DimFix i) rw = - 0 .<=. i .&&. i .<. rw - condInBounds (DimSlice i n s) rw = - 0 .<=. i .&&. i + n * s .<. rw - write = - foldl (.&&.) (kernelThreadActive constants) $ - zipWith condInBounds slice' rws' - sWhen write $ copyDWIM (patElemName pe) slice' e [] -compileThreadResult _ _ TileReturns {} = - compilerBugS "compileThreadResult: TileReturns unhandled." - -arrayInLocalMemory :: SubExp -> InKernelGen Bool -arrayInLocalMemory (Var name) = do - res <- lookupVar name - case res of - ArrayVar _ entry -> - (Space "local" ==) . entryMemSpace - <$> lookupMemory (memLocationName (entryArrayLocation entry)) - _ -> return False -arrayInLocalMemory Constant {} = return False diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs +++ /dev/null @@ -1,1151 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TypeFamilies #-} - --- | Our compilation strategy for 'SegHist' is based around avoiding --- bin conflicts. We do this by splitting the input into chunks, and --- for each chunk computing a single subhistogram. Then we combine --- the subhistograms using an ordinary segmented reduction ('SegRed'). --- --- There are some branches around to efficiently handle the case where --- we use only a single subhistogram (because it's large), so that we --- respect the asymptotics, and do not copy the destination array. --- --- We also use a heuristic strategy for computing subhistograms in --- local memory when possible. Given: --- --- H: total size of histograms in bytes, including any lock arrays. --- --- G: group size --- --- T: number of bytes of local memory each thread can be given without --- impacting occupancy (determined experimentally, e.g. 32). --- --- LMAX: maximum amount of local memory per workgroup (hard limit). --- --- We wish to compute: --- --- COOP: cooperation level (number of threads per subhistogram) --- --- LH: number of local memory subhistograms --- --- We do this as: --- --- COOP = ceil(H / T) --- LH = ceil((G*T)/H) --- if COOP <= G && H <= LMAX then --- use local memory --- else --- use global memory -module Futhark.CodeGen.ImpGen.Kernels.SegHist (compileSegHist) where - -import Control.Monad.Except -import Data.List (foldl', genericLength, zip4, zip6) -import Data.Maybe -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.CodeGen.ImpGen.Kernels.SegRed (compileSegRed') -import Futhark.Construct (fullSliceNum) -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.MonadFreshNames -import Futhark.Pass.ExplicitAllocations () -import Futhark.Util (chunks, mapAccumLM, maxinum, splitFromEnd, takeLast) -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Prelude hiding (quot, rem) - -data SubhistosInfo = SubhistosInfo - { subhistosArray :: VName, - subhistosAlloc :: CallKernelGen () - } - -data SegHistSlug = SegHistSlug - { slugOp :: HistOp KernelsMem, - slugNumSubhistos :: TV Int64, - slugSubhistos :: [SubhistosInfo], - slugAtomicUpdate :: AtomicUpdate KernelsMem KernelEnv - } - -histoSpaceUsage :: - HistOp KernelsMem -> - Imp.Count Imp.Bytes (Imp.TExp Int64) -histoSpaceUsage op = - sum $ - map - ( typeSize - . (`arrayOfRow` histWidth op) - . (`arrayOfShape` histShape op) - ) - $ lambdaReturnType $ histOp op - --- | Figure out how much memory is needed per histogram, both --- segmented and unsegmented,, and compute some other auxiliary --- information. -computeHistoUsage :: - SegSpace -> - HistOp KernelsMem -> - CallKernelGen - ( Imp.Count Imp.Bytes (Imp.TExp Int64), - Imp.Count Imp.Bytes (Imp.TExp Int64), - SegHistSlug - ) -computeHistoUsage space op = do - let segment_dims = init $ unSegSpace space - num_segments = length segment_dims - - -- Create names for the intermediate array memory blocks, - -- memory block sizes, arrays, and number of subhistograms. - num_subhistos <- dPrim "num_subhistos" int32 - subhisto_infos <- forM (zip (histDest op) (histNeutral op)) $ \(dest, ne) -> do - dest_t <- lookupType dest - dest_mem <- entryArrayLocation <$> lookupArray dest - - subhistos_mem <- - sDeclareMem (baseString dest ++ "_subhistos_mem") (Space "device") - - let subhistos_shape = - Shape (map snd segment_dims ++ [tvSize num_subhistos]) - <> stripDims num_segments (arrayShape dest_t) - subhistos_membind = - ArrayIn subhistos_mem $ - IxFun.iota $ - map pe64 $ shapeDims subhistos_shape - subhistos <- - sArray - (baseString dest ++ "_subhistos") - (elemType dest_t) - subhistos_shape - subhistos_membind - - return $ - SubhistosInfo subhistos $ do - let unitHistoCase = - emit $ - Imp.SetMem subhistos_mem (memLocationName dest_mem) $ - Space "device" - - multiHistoCase = do - let num_elems = - foldl' (*) (sExt64 $ tvExp num_subhistos) $ - map toInt64Exp $ arrayDims dest_t - - let subhistos_mem_size = - Imp.bytes $ - Imp.unCount (Imp.elements num_elems `Imp.withElemType` elemType dest_t) - - sAlloc_ subhistos_mem subhistos_mem_size $ Space "device" - sReplicate subhistos ne - subhistos_t <- lookupType subhistos - let slice = - fullSliceNum (map toInt64Exp $ arrayDims subhistos_t) $ - map (unitSlice 0 . toInt64Exp . snd) segment_dims - ++ [DimFix 0] - sUpdate subhistos slice $ Var dest - - sIf (tvExp num_subhistos .==. 1) unitHistoCase multiHistoCase - - let h = histoSpaceUsage op - segmented_h = h * product (map (Imp.bytes . toInt64Exp) $ init $ segSpaceDims space) - - atomics <- hostAtomics <$> askEnv - - return - ( h, - segmented_h, - SegHistSlug op num_subhistos subhisto_infos $ - atomicUpdateLocking atomics $ histOp op - ) - -prepareAtomicUpdateGlobal :: - Maybe Locking -> - [VName] -> - SegHistSlug -> - CallKernelGen - ( Maybe Locking, - [Imp.TExp Int64] -> InKernelGen () - ) -prepareAtomicUpdateGlobal l dests slug = - -- We need a separate lock array if the operators are not all of a - -- particularly simple form that permits pure atomic operations. - case (l, slugAtomicUpdate slug) of - (_, AtomicPrim f) -> return (l, f (Space "global") dests) - (_, AtomicCAS f) -> return (l, f (Space "global") dests) - (Just l', AtomicLocking f) -> return (l, f l' (Space "global") dests) - (Nothing, AtomicLocking f) -> do - -- The number of locks used here is too low, but since we are - -- currently forced to inline a huge list, I'm keeping it down - -- for now. Some quick experiments suggested that it has little - -- impact anyway (maybe the locking case is just too slow). - -- - -- A fun solution would also be to use a simple hashing - -- algorithm to ensure good distribution of locks. - let num_locks = 100151 - dims = - map toInt64Exp $ - shapeDims (histShape (slugOp slug)) - ++ [ tvSize (slugNumSubhistos slug), - histWidth (slugOp slug) - ] - locks <- - sStaticArray "hist_locks" (Space "device") int32 $ - Imp.ArrayZeros num_locks - let l' = Locking locks 0 1 0 (pure . (`rem` fromIntegral num_locks) . flattenIndex dims) - return (Just l', f l' (Space "global") dests) - --- | Some kernel bodies are not safe (or efficient) to execute --- multiple times. -data Passage = MustBeSinglePass | MayBeMultiPass deriving (Eq, Ord) - -bodyPassage :: KernelBody KernelsMem -> Passage -bodyPassage kbody - | mempty == consumedInKernelBody (aliasAnalyseKernelBody mempty kbody) = - MayBeMultiPass - | otherwise = - MustBeSinglePass - -prepareIntermediateArraysGlobal :: - Passage -> - Imp.TExp Int32 -> - Imp.TExp Int64 -> - [SegHistSlug] -> - CallKernelGen - ( Imp.TExp Int32, - [[Imp.TExp Int64] -> InKernelGen ()] - ) -prepareIntermediateArraysGlobal passage hist_T hist_N slugs = do - -- The paper formulae assume there is only one histogram, but in our - -- implementation there can be multiple that have been horisontally - -- fused. We do a bit of trickery with summings and averages to - -- pretend there is really only one. For the case of a single - -- histogram, the actual calculations should be the same as in the - -- paper. - - -- The sum of all Hs. - hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth . slugOp) slugs - - hist_RF <- - dPrimVE "hist_RF" $ - sum (map (r64 . toInt64Exp . histRaceFactor . slugOp) slugs) - / genericLength slugs - - hist_el_size <- dPrimVE "hist_el_size" $ sum $ map slugElAvgSize slugs - - hist_C_max <- - dPrimVE "hist_C_max" $ - fMin64 (r64 hist_T) $ r64 hist_H / hist_k_ct_min - - hist_M_min <- - dPrimVE "hist_M_min" $ - sMax32 1 $ sExt32 $ t64 $ r64 hist_T / hist_C_max - - -- Querying L2 cache size is not reliable. Instead we provide a - -- tunable knob with a hopefully sane default. - let hist_L2_def = 4 * 1024 * 1024 - hist_L2 <- dPrim "L2_size" int32 - entry <- askFunction - -- Equivalent to F_L2*L2 in paper. - sOp $ - Imp.GetSize - (tvVar hist_L2) - (keyWithEntryPoint entry $ nameFromString (pretty (tvVar hist_L2))) - $ Imp.SizeBespoke (nameFromString "L2_for_histogram") hist_L2_def - - let hist_L2_ln_sz = 16 * 4 -- L2 cache line size approximation - hist_RACE_exp <- - dPrimVE "hist_RACE_exp" $ - fMax64 1 $ - (hist_k_RF * hist_RF) - / (hist_L2_ln_sz / r64 hist_el_size) - - hist_S <- dPrim "hist_S" int32 - - -- For sparse histograms (H exceeds N) we only want a single chunk. - sIf - (hist_N .<. hist_H) - (hist_S <-- (1 :: Imp.TExp Int32)) - $ hist_S - <-- case passage of - MayBeMultiPass -> - sExt32 $ - (sExt64 hist_M_min * hist_H * sExt64 hist_el_size) - `divUp` t64 (hist_F_L2 * r64 (tvExp hist_L2) * hist_RACE_exp) - MustBeSinglePass -> - 1 - - emit $ Imp.DebugPrint "Race expansion factor (RACE^exp)" $ Just $ untyped hist_RACE_exp - emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ untyped $ tvExp hist_S - - histograms <- - snd - <$> mapAccumLM - (onOp (tvExp hist_L2) hist_M_min (tvExp hist_S) hist_RACE_exp) - Nothing - slugs - - return (tvExp hist_S, histograms) - where - hist_k_ct_min = 2 -- Chosen experimentally - hist_k_RF = 0.75 -- Chosen experimentally - hist_F_L2 = 0.4 -- Chosen experimentally - r64 = isF64 . ConvOpExp (SIToFP Int32 Float64) . untyped - t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped - - -- "Average element size" as computed by a formula that also takes - -- locking into account. - slugElAvgSize slug@(SegHistSlug op _ _ do_op) = - case do_op of - AtomicLocking {} -> - slugElSize slug `quot` (1 + genericLength (lambdaReturnType (histOp op))) - _ -> - slugElSize slug `quot` genericLength (lambdaReturnType (histOp op)) - - -- "Average element size" as computed by a formula that also takes - -- locking into account. - slugElSize (SegHistSlug op _ _ do_op) = - case do_op of - AtomicLocking {} -> - sExt32 $ - unCount $ - sum $ - map (typeSize . (`arrayOfShape` histShape op)) $ - Prim int32 : lambdaReturnType (histOp op) - _ -> - sExt32 $ - unCount $ - sum $ - map (typeSize . (`arrayOfShape` histShape op)) $ - lambdaReturnType (histOp op) - - onOp hist_L2 hist_M_min hist_S hist_RACE_exp l slug = do - let SegHistSlug op num_subhistos subhisto_info do_op = slug - hist_H = toInt64Exp $ histWidth op - - hist_H_chk <- dPrimVE "hist_H_chk" $ hist_H `divUp` sExt64 hist_S - - emit $ Imp.DebugPrint "Chunk size (H_chk)" $ Just $ untyped hist_H_chk - - hist_k_max <- - dPrimVE "hist_k_max" $ - fMin64 - (hist_F_L2 * (r64 hist_L2 / r64 (slugElSize slug)) * hist_RACE_exp) - (r64 hist_N) - / r64 hist_T - - hist_u <- dPrimVE "hist_u" $ - case do_op of - AtomicPrim {} -> 2 - _ -> 1 - - hist_C <- - dPrimVE "hist_C" $ - fMin64 (r64 hist_T) $ r64 (hist_u * hist_H_chk) / hist_k_max - - -- Number of subhistograms per result histogram. - hist_M <- dPrimVE "hist_M" $ - case slugAtomicUpdate slug of - AtomicPrim {} -> 1 - _ -> sMax32 hist_M_min $ sExt32 $ t64 $ r64 hist_T / hist_C - - emit $ Imp.DebugPrint "Elements/thread in L2 cache (k_max)" $ Just $ untyped hist_k_max - emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ untyped hist_M - emit $ Imp.DebugPrint "Cooperation level (C)" $ Just $ untyped hist_C - - -- num_subhistos is the variable we use to communicate back. - num_subhistos <-- sExt64 hist_M - - -- Initialise sub-histograms. - -- - -- If hist_M is 1, then we just reuse the original - -- destination. The idea is to avoid a copy if we are writing a - -- small number of values into a very large prior histogram. - dests <- forM (zip (histDest op) subhisto_info) $ \(dest, info) -> do - dest_mem <- entryArrayLocation <$> lookupArray dest - - sub_mem <- - fmap memLocationName $ - entryArrayLocation - <$> lookupArray (subhistosArray info) - - let unitHistoCase = - emit $ - Imp.SetMem sub_mem (memLocationName dest_mem) $ - Space "device" - - multiHistoCase = subhistosAlloc info - - sIf (hist_M .==. 1) unitHistoCase multiHistoCase - - return $ subhistosArray info - - (l', do_op') <- prepareAtomicUpdateGlobal l dests slug - - return (l', do_op') - -histKernelGlobalPass :: - [PatElem KernelsMem] -> - Count NumGroups (Imp.TExp Int64) -> - Count GroupSize (Imp.TExp Int64) -> - SegSpace -> - [SegHistSlug] -> - KernelBody KernelsMem -> - [[Imp.TExp Int64] -> InKernelGen ()] -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> - CallKernelGen () -histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do - let (space_is, space_sizes) = unzip $ unSegSpace space - space_sizes_64 = map (sExt64 . toInt64Exp) space_sizes - total_w_64 = product space_sizes_64 - - hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> - dPrimVE "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S - - sKernelThread "seghist_global" num_groups group_size (segFlat space) $ do - constants <- kernelConstants <$> askEnv - - -- Compute subhistogram index for each thread, per histogram. - subhisto_inds <- forM slugs $ \slug -> - dPrimVE "subhisto_ind" $ - kernelGlobalThreadId constants - `quot` ( kernelNumThreads constants - `divUp` sExt32 (tvExp (slugNumSubhistos slug)) - ) - - -- Loop over flat offsets into the input and output. The - -- calculation is done with 64-bit integers to avoid overflow, - -- but the final unflattened segment indexes are 32 bit. - let gtid = sExt64 $ kernelGlobalThreadId constants - num_threads = sExt64 $ kernelNumThreads constants - kernelLoop gtid num_threads total_w_64 $ \offset -> do - -- Construct segment indices. - zipWithM_ dPrimV_ space_is $ - map sExt32 $ unflattenIndex space_sizes_64 offset - - -- We execute the bucket function once and update each histogram serially. - -- We apply the bucket function if j=offset+ltid is less than - -- num_elements. This also involves writing to the mapout - -- arrays. - let input_in_bounds = offset .<. total_w_64 - - sWhen input_in_bounds $ - compileStms mempty (kernelBodyStms kbody) $ do - let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody - - sComment "save map-out results" $ - forM_ (zip map_pes map_res) $ \(pe, res) -> - copyDWIMFix - (patElemName pe) - (map (Imp.vi64 . fst) $ unSegSpace space) - (kernelResultSubExp res) - [] - - let (buckets, vs) = splitAt (length slugs) red_res - perOp = chunks $ map (length . histDest . slugOp) slugs - - sComment "perform atomic updates" $ - forM_ (zip6 (map slugOp slugs) histograms buckets (perOp vs) subhisto_inds hist_H_chks) $ - \( HistOp dest_w _ _ _ shape lam, - do_op, - bucket, - vs', - subhisto_ind, - hist_H_chk - ) -> do - let chk_beg = sExt64 chk_i * hist_H_chk - bucket' = toInt64Exp $ kernelResultSubExp bucket - dest_w' = toInt64Exp dest_w - bucket_in_bounds = - chk_beg .<=. bucket' - .&&. bucket' .<. (chk_beg + hist_H_chk) - .&&. bucket' .<. dest_w' - vs_params = takeLast (length vs') $ lambdaParams lam - - sWhen bucket_in_bounds $ do - let bucket_is = - map Imp.vi64 (init space_is) - ++ [sExt64 subhisto_ind, bucket'] - dLParams $ lambdaParams lam - sLoopNest shape $ \is -> do - forM_ (zip vs_params vs') $ \(p, res) -> - copyDWIMFix (paramName p) [] (kernelResultSubExp res) is - do_op (bucket_is ++ is) - -histKernelGlobal :: - [PatElem KernelsMem] -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - SegSpace -> - [SegHistSlug] -> - KernelBody KernelsMem -> - CallKernelGen () -histKernelGlobal map_pes num_groups group_size space slugs kbody = do - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - let (_space_is, space_sizes) = unzip $ unSegSpace space - num_threads = sExt32 $ unCount num_groups' * unCount group_size' - - emit $ Imp.DebugPrint "## Using global memory" Nothing - - (hist_S, histograms) <- - prepareIntermediateArraysGlobal - (bodyPassage kbody) - num_threads - (toInt64Exp $ last space_sizes) - slugs - - sFor "chk_i" hist_S $ \chk_i -> - histKernelGlobalPass - map_pes - num_groups' - group_size' - space - slugs - kbody - histograms - hist_S - chk_i - -type InitLocalHistograms = - [ ( [VName], - SubExp -> - InKernelGen - ( [VName], - [Imp.TExp Int64] -> InKernelGen () - ) - ) - ] - -prepareIntermediateArraysLocal :: - TV Int32 -> - Count NumGroups (Imp.TExp Int64) -> - SegSpace -> - [SegHistSlug] -> - CallKernelGen InitLocalHistograms -prepareIntermediateArraysLocal num_subhistos_per_group groups_per_segment space slugs = do - num_segments <- - dPrimVE "num_segments" $ - product $ map (toInt64Exp . snd) $ init $ unSegSpace space - mapM (onOp num_segments) slugs - where - onOp num_segments (SegHistSlug op num_subhistos subhisto_info do_op) = do - num_subhistos <-- sExt64 (unCount groups_per_segment) * num_segments - - emit $ - Imp.DebugPrint "Number of subhistograms in global memory" $ - Just $ untyped $ tvExp num_subhistos - - mk_op <- - case do_op of - AtomicPrim f -> return $ const $ return f - AtomicCAS f -> return $ const $ return f - AtomicLocking f -> return $ \hist_H_chk -> do - let lock_shape = - Shape $ - tvSize num_subhistos_per_group : - shapeDims (histShape op) - ++ [hist_H_chk] - - let dims = map toInt64Exp $ shapeDims lock_shape - - locks <- sAllocArray "locks" int32 lock_shape $ Space "local" - - sComment "All locks start out unlocked" $ - groupCoverSpace dims $ \is -> - copyDWIMFix locks is (intConst Int32 0) [] - - return $ f $ Locking locks 0 1 0 id - - -- Initialise local-memory sub-histograms. These are - -- represented as two-dimensional arrays. - let init_local_subhistos hist_H_chk = do - local_subhistos <- - forM (histType op) $ \t -> do - let sub_local_shape = - Shape [tvSize num_subhistos_per_group] - <> (arrayShape t `setOuterDim` hist_H_chk) - sAllocArray - "subhistogram_local" - (elemType t) - sub_local_shape - (Space "local") - - do_op' <- mk_op hist_H_chk - - return (local_subhistos, do_op' (Space "local") local_subhistos) - - -- Initialise global-memory sub-histograms. - glob_subhistos <- forM subhisto_info $ \info -> do - subhistosAlloc info - return $ subhistosArray info - - return (glob_subhistos, init_local_subhistos) - -histKernelLocalPass :: - TV Int32 -> - Count NumGroups (Imp.TExp Int64) -> - [PatElem KernelsMem] -> - Count NumGroups (Imp.TExp Int64) -> - Count GroupSize (Imp.TExp Int64) -> - SegSpace -> - [SegHistSlug] -> - KernelBody KernelsMem -> - InitLocalHistograms -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> - CallKernelGen () -histKernelLocalPass - num_subhistos_per_group_var - groups_per_segment - map_pes - num_groups - group_size - space - slugs - kbody - init_histograms - hist_S - chk_i = do - let (space_is, space_sizes) = unzip $ unSegSpace space - segment_is = init space_is - segment_dims = init space_sizes - (i_in_segment, segment_size) = last $ unSegSpace space - num_subhistos_per_group = tvExp num_subhistos_per_group_var - segment_size' = toInt64Exp segment_size - - num_segments <- - dPrimVE "num_segments" $ - product $ map toInt64Exp segment_dims - - hist_H_chks <- forM (map (histWidth . slugOp) slugs) $ \w -> - dPrimV "hist_H_chk" $ toInt64Exp w `divUp` sExt64 hist_S - - histo_sizes <- forM (zip slugs hist_H_chks) $ \(slug, hist_H_chk) -> do - let histo_dims = - tvExp hist_H_chk : - map toInt64Exp (shapeDims (histShape (slugOp slug))) - histo_size <- - dPrimVE "histo_size" $ product histo_dims - let group_hists_size = - sExt64 num_subhistos_per_group * histo_size - init_per_thread <- - dPrimVE "init_per_thread" $ sExt32 $ group_hists_size `divUp` unCount group_size - return (histo_dims, histo_size, init_per_thread) - - sKernelThread "seghist_local" num_groups group_size (segFlat space) $ - virtualiseGroups SegVirt (sExt32 $ unCount groups_per_segment * num_segments) $ \group_id -> do - constants <- kernelConstants <$> askEnv - - flat_segment_id <- dPrimVE "flat_segment_id" $ group_id `quot` sExt32 (unCount groups_per_segment) - gid_in_segment <- dPrimVE "gid_in_segment" $ group_id `rem` sExt32 (unCount groups_per_segment) - -- This pgtid is kind of a "virtualised physical" gtid - not the - -- same thing as the gtid used for the SegHist itself. - pgtid_in_segment <- - dPrimVE "pgtid_in_segment" $ - gid_in_segment * sExt32 (kernelGroupSize constants) - + kernelLocalThreadId constants - threads_per_segment <- - dPrimVE "threads_per_segment" $ - sExt32 $ unCount groups_per_segment * kernelGroupSize constants - - -- Set segment indices. - zipWithM_ dPrimV_ segment_is $ - unflattenIndex (map toInt64Exp segment_dims) $ sExt64 flat_segment_id - - histograms <- forM (zip init_histograms hist_H_chks) $ - \((glob_subhistos, init_local_subhistos), hist_H_chk) -> do - (local_subhistos, do_op) <- init_local_subhistos $ Var $ tvVar hist_H_chk - return (zip glob_subhistos local_subhistos, hist_H_chk, do_op) - - -- Find index of local subhistograms updated by this thread. We - -- try to ensure, as much as possible, that threads in the same - -- warp use different subhistograms, to avoid conflicts. - thread_local_subhisto_i <- - dPrimVE "thread_local_subhisto_i" $ - kernelLocalThreadId constants `rem` num_subhistos_per_group - - let onSlugs f = - forM_ (zip3 slugs histograms histo_sizes) $ - \(slug, (dests, hist_H_chk, _), (histo_dims, histo_size, init_per_thread)) -> - f slug dests (tvExp hist_H_chk) histo_dims histo_size init_per_thread - - let onAllHistograms f = - onSlugs $ \slug dests hist_H_chk histo_dims histo_size init_per_thread -> do - let group_hists_size = num_subhistos_per_group * sExt32 histo_size - - forM_ (zip dests (histNeutral $ slugOp slug)) $ - \((dest_global, dest_local), ne) -> - sFor "local_i" init_per_thread $ \i -> do - j <- - dPrimVE "j" $ - i * sExt32 (kernelGroupSize constants) - + kernelLocalThreadId constants - j_offset <- - dPrimVE "j_offset" $ - num_subhistos_per_group * sExt32 histo_size * gid_in_segment + j - - local_subhisto_i <- dPrimVE "local_subhisto_i" $ j `quot` sExt32 histo_size - let local_bucket_is = unflattenIndex histo_dims $ sExt64 $ j `rem` sExt32 histo_size - global_bucket_is = - head local_bucket_is + sExt64 chk_i * hist_H_chk : - tail local_bucket_is - global_subhisto_i <- dPrimVE "global_subhisto_i" $ j_offset `quot` sExt32 histo_size - - sWhen (j .<. group_hists_size) $ - f - dest_local - dest_global - (slugOp slug) - ne - local_subhisto_i - global_subhisto_i - local_bucket_is - global_bucket_is - - sComment "initialize histograms in local memory" $ - onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i local_bucket_is global_bucket_is -> - sComment "First subhistogram is initialised from global memory; others with neutral element." $ do - let global_is = map Imp.vi64 segment_is ++ [0] ++ global_bucket_is - local_is = sExt64 local_subhisto_i : local_bucket_is - sIf - (global_subhisto_i .==. 0) - (copyDWIMFix dest_local local_is (Var dest_global) global_is) - ( sLoopNest (histShape op) $ \is -> - copyDWIMFix dest_local (local_is ++ is) ne [] - ) - - sOp $ Imp.Barrier Imp.FenceLocal - - kernelLoop pgtid_in_segment threads_per_segment (sExt32 segment_size') $ \ie -> do - dPrimV_ i_in_segment ie - - -- We execute the bucket function once and update each histogram - -- serially. This also involves writing to the mapout arrays if - -- this is the first chunk. - - compileStms mempty (kernelBodyStms kbody) $ do - let (red_res, map_res) = - splitFromEnd (length map_pes) $ - map kernelResultSubExp $ kernelBodyResult kbody - (buckets, vs) = splitAt (length slugs) red_res - perOp = chunks $ map (length . histDest . slugOp) slugs - - sWhen (chk_i .==. 0) $ - sComment "save map-out results" $ - forM_ (zip map_pes map_res) $ \(pe, se) -> - copyDWIMFix - (patElemName pe) - (map Imp.vi64 space_is) - se - [] - - forM_ (zip4 (map slugOp slugs) histograms buckets (perOp vs)) $ - \( HistOp dest_w _ _ _ shape lam, - (_, hist_H_chk, do_op), - bucket, - vs' - ) -> do - let chk_beg = sExt64 chk_i * tvExp hist_H_chk - bucket' = toInt64Exp bucket - dest_w' = toInt64Exp dest_w - bucket_in_bounds = - bucket' .<. dest_w' - .&&. chk_beg .<=. bucket' - .&&. bucket' .<. (chk_beg + tvExp hist_H_chk) - bucket_is = [sExt64 thread_local_subhisto_i, bucket' - chk_beg] - vs_params = takeLast (length vs') $ lambdaParams lam - - sComment "perform atomic updates" $ - sWhen bucket_in_bounds $ do - dLParams $ lambdaParams lam - sLoopNest shape $ \is -> do - forM_ (zip vs_params vs') $ \(p, v) -> - copyDWIMFix (paramName p) [] v is - do_op (bucket_is ++ is) - - sOp $ Imp.ErrorSync Imp.FenceGlobal - - sComment "Compact the multiple local memory subhistograms to result in global memory" $ - onSlugs $ \slug dests hist_H_chk histo_dims _histo_size bins_per_thread -> do - trunc_H <- - dPrimV "trunc_H" $ - sMin64 hist_H_chk $ - toInt64Exp (histWidth (slugOp slug)) - - sExt64 chk_i * head histo_dims - let trunc_histo_dims = - tvExp trunc_H : - map toInt64Exp (shapeDims (histShape (slugOp slug))) - trunc_histo_size <- dPrimVE "histo_size" $ sExt32 $ product trunc_histo_dims - - sFor "local_i" bins_per_thread $ \i -> do - j <- - dPrimVE "j" $ - i * sExt32 (kernelGroupSize constants) - + kernelLocalThreadId constants - sWhen (j .<. trunc_histo_size) $ do - -- We are responsible for compacting the flat bin 'j', which - -- we immediately unflatten. - let local_bucket_is = unflattenIndex histo_dims $ sExt64 j - global_bucket_is = - head local_bucket_is + sExt64 chk_i * hist_H_chk : - tail local_bucket_is - dLParams $ lambdaParams $ histOp $ slugOp slug - let (global_dests, local_dests) = unzip dests - (xparams, yparams) = - splitAt (length local_dests) $ - lambdaParams $ histOp $ slugOp slug - - sComment "Read values from subhistogram 0." $ - forM_ (zip xparams local_dests) $ \(xp, subhisto) -> - copyDWIMFix - (paramName xp) - [] - (Var subhisto) - (0 : local_bucket_is) - - sComment "Accumulate based on values in other subhistograms." $ - sFor "subhisto_id" (num_subhistos_per_group - 1) $ \subhisto_id -> do - forM_ (zip yparams local_dests) $ \(yp, subhisto) -> - copyDWIMFix - (paramName yp) - [] - (Var subhisto) - (sExt64 subhisto_id + 1 : local_bucket_is) - compileBody' xparams $ lambdaBody $ histOp $ slugOp slug - - sComment "Put final bucket value in global memory." $ do - let global_is = - map Imp.vi64 segment_is - ++ [sExt64 group_id `rem` unCount groups_per_segment] - ++ global_bucket_is - forM_ (zip xparams global_dests) $ \(xp, global_dest) -> - copyDWIMFix global_dest global_is (Var $ paramName xp) [] - -histKernelLocal :: - TV Int32 -> - Count NumGroups (Imp.TExp Int64) -> - [PatElem KernelsMem] -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - SegSpace -> - Imp.TExp Int32 -> - [SegHistSlug] -> - KernelBody KernelsMem -> - CallKernelGen () -histKernelLocal num_subhistos_per_group_var groups_per_segment map_pes num_groups group_size space hist_S slugs kbody = do - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - num_subhistos_per_group = tvExp num_subhistos_per_group_var - - emit $ - Imp.DebugPrint "Number of local subhistograms per group" $ - Just $ untyped num_subhistos_per_group - - init_histograms <- - prepareIntermediateArraysLocal num_subhistos_per_group_var groups_per_segment space slugs - - sFor "chk_i" hist_S $ \chk_i -> - histKernelLocalPass - num_subhistos_per_group_var - groups_per_segment - map_pes - num_groups' - group_size' - space - slugs - kbody - init_histograms - hist_S - chk_i - --- | The maximum number of passes we are willing to accept for this --- kind of atomic update. -slugMaxLocalMemPasses :: SegHistSlug -> Int -slugMaxLocalMemPasses slug = - case slugAtomicUpdate slug of - AtomicPrim _ -> 3 - AtomicCAS _ -> 4 - AtomicLocking _ -> 6 - -localMemoryCase :: - [PatElem KernelsMem] -> - Imp.TExp Int32 -> - SegSpace -> - Imp.TExp Int64 -> - Imp.TExp Int64 -> - Imp.TExp Int64 -> - Imp.TExp Int32 -> - [SegHistSlug] -> - KernelBody KernelsMem -> - CallKernelGen (Imp.TExp Bool, CallKernelGen ()) -localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N _ slugs kbody = do - let space_sizes = segSpaceDims space - segment_dims = init space_sizes - segmented = not $ null segment_dims - - hist_L <- dPrim "hist_L" int32 - sOp $ Imp.GetSizeMax (tvVar hist_L) Imp.SizeLocalMemory - - max_group_size <- dPrim "max_group_size" int32 - sOp $ Imp.GetSizeMax (tvVar max_group_size) Imp.SizeGroup - let group_size = Imp.Count $ Var $ tvVar max_group_size - num_groups <- - fmap (Imp.Count . tvSize) $ - dPrimV "num_groups" $ - sExt64 hist_T `divUp` toInt64Exp (unCount group_size) - let num_groups' = toInt64Exp <$> num_groups - group_size' = toInt64Exp <$> group_size - - let r64 = isF64 . ConvOpExp (SIToFP Int64 Float64) . untyped - t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped - - -- M approximation. - hist_m' <- - dPrimVE "hist_m_prime" $ - r64 - ( sMin64 - (sExt64 (tvExp hist_L `quot` hist_el_size)) - (hist_N `divUp` sExt64 (unCount num_groups')) - ) - / r64 hist_H - - let hist_B = unCount group_size' - - -- M in the paper, but not adjusted for asymptotic efficiency. - hist_M0 <- - dPrimVE "hist_M0" $ - sMax64 1 $ sMin64 (t64 hist_m') hist_B - - -- Minimal sequential chunking factor. - let q_small = 2 - - -- The number of segments/histograms produced.. - hist_Nout <- dPrimVE "hist_Nout" $ product $ map toInt64Exp segment_dims - - hist_Nin <- dPrimVE "hist_Nin" $ toInt64Exp $ last space_sizes - - -- Maximum M for work efficiency. - work_asymp_M_max <- - if segmented - then do - hist_T_hist_min <- - dPrimVE "hist_T_hist_min" $ - sExt32 $ - sMin64 (sExt64 hist_Nin * sExt64 hist_Nout) (sExt64 hist_T) - `divUp` sExt64 hist_Nout - - -- Number of groups, rounded up. - let r = hist_T_hist_min `divUp` sExt32 hist_B - - dPrimVE "work_asymp_M_max" $ hist_Nin `quot` (sExt64 r * hist_H) - else - dPrimVE "work_asymp_M_max" $ - (hist_Nout * hist_N) - `quot` ( (q_small * unCount num_groups' * hist_H) - `quot` genericLength slugs - ) - - -- Number of subhistograms per result histogram. - hist_M <- dPrimV "hist_M" $ sExt32 $ sMin64 hist_M0 work_asymp_M_max - - -- hist_M may be zero (which we'll check for below), but we need it - -- for some divisions first, so crudely make a nonzero form. - let hist_M_nonzero = sMax32 1 $ tvExp hist_M - - -- "Cooperation factor" - the number of threads cooperatively - -- working on the same (sub)histogram. - hist_C <- - dPrimVE "hist_C" $ - hist_B `divUp` sExt64 hist_M_nonzero - - emit $ Imp.DebugPrint "local hist_M0" $ Just $ untyped hist_M0 - emit $ Imp.DebugPrint "local work asymp M max" $ Just $ untyped work_asymp_M_max - emit $ Imp.DebugPrint "local C" $ Just $ untyped hist_C - emit $ Imp.DebugPrint "local B" $ Just $ untyped hist_B - emit $ Imp.DebugPrint "local M" $ Just $ untyped $ tvExp hist_M - emit $ - Imp.DebugPrint "local memory needed" $ - Just $ untyped $ hist_H * hist_el_size * sExt64 (tvExp hist_M) - - -- local_mem_needed is what we need to keep a single bucket in local - -- memory - this is an absolute minimum. We can fit anything else - -- by doing multiple passes, although more than a few is - -- (heuristically) not efficient. - local_mem_needed <- - dPrimVE "local_mem_needed" $ - hist_el_size * sExt64 (tvExp hist_M) - hist_S <- - dPrimVE "hist_S" $ - sExt32 $ - (hist_H * local_mem_needed) `divUp` tvExp hist_L - let max_S = case bodyPassage kbody of - MustBeSinglePass -> 1 - MayBeMultiPass -> fromIntegral $ maxinum $ map slugMaxLocalMemPasses slugs - - groups_per_segment <- - if segmented - then - fmap Count $ - dPrimVE "groups_per_segment" $ unCount num_groups' `divUp` hist_Nout - else pure num_groups' - - -- We only use local memory if the number of updates per histogram - -- at least matches the histogram size, as otherwise it is not - -- asymptotically efficient. This mostly matters for the segmented - -- case. - let pick_local = - hist_Nin .>=. hist_H - .&&. (local_mem_needed .<=. tvExp hist_L) - .&&. (hist_S .<=. max_S) - .&&. hist_C .<=. hist_B - .&&. tvExp hist_M .>. 0 - - run = do - emit $ Imp.DebugPrint "## Using local memory" Nothing - emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H - emit $ Imp.DebugPrint "Multiplication degree (M)" $ Just $ untyped $ tvExp hist_M - emit $ Imp.DebugPrint "Cooperation level (C)" $ Just $ untyped hist_C - emit $ Imp.DebugPrint "Number of chunks (S)" $ Just $ untyped hist_S - when segmented $ - emit $ Imp.DebugPrint "Groups per segment" $ Just $ untyped $ unCount groups_per_segment - histKernelLocal - hist_M - groups_per_segment - map_pes - num_groups - group_size - space - hist_S - slugs - kbody - - return (pick_local, run) - --- | Generate code for a segmented histogram called from the host. -compileSegHist :: - Pattern KernelsMem -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - SegSpace -> - [HistOp KernelsMem] -> - KernelBody KernelsMem -> - CallKernelGen () -compileSegHist (Pattern _ pes) num_groups group_size space ops kbody = do - -- Most of this function is not the histogram part itself, but - -- rather figuring out whether to use a local or global memory - -- strategy, as well as collapsing the subhistograms produced (which - -- are always in global memory, but their number may vary). - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - dims = map toInt64Exp $ segSpaceDims space - - num_red_res = length ops + sum (map (length . histNeutral) ops) - (all_red_pes, map_pes) = splitAt num_red_res pes - segment_size = last dims - - (op_hs, op_seg_hs, slugs) <- unzip3 <$> mapM (computeHistoUsage space) ops - h <- dPrimVE "h" $ Imp.unCount $ sum op_hs - seg_h <- dPrimVE "seg_h" $ Imp.unCount $ sum op_seg_hs - - -- Check for emptyness to avoid division-by-zero. - sUnless (seg_h .==. 0) $ do - -- Maximum group size (or actual, in this case). - let hist_B = unCount group_size' - - -- Size of a histogram. - hist_H <- dPrimVE "hist_H" $ sum $ map (toInt64Exp . histWidth) ops - - -- Size of a single histogram element. Actually the weighted - -- average of histogram elements in cases where we have more than - -- one histogram operation, plus any locks. - let lockSize slug = case slugAtomicUpdate slug of - AtomicLocking {} -> Just $ primByteSize int32 - _ -> Nothing - hist_el_size <- - dPrimVE "hist_el_size" $ - foldl' (+) (h `divUp` hist_H) $ - mapMaybe lockSize slugs - - -- Input elements contributing to each histogram. - hist_N <- dPrimVE "hist_N" segment_size - - -- Compute RF as the average RF over all the histograms. - hist_RF <- - dPrimVE "hist_RF" $ - sExt32 $ - sum (map (toInt64Exp . histRaceFactor . slugOp) slugs) - `quot` genericLength slugs - - let hist_T = sExt32 $ unCount num_groups' * unCount group_size' - emit $ Imp.DebugPrint "\n# SegHist" Nothing - emit $ Imp.DebugPrint "Number of threads (T)" $ Just $ untyped hist_T - emit $ Imp.DebugPrint "Desired group size (B)" $ Just $ untyped hist_B - emit $ Imp.DebugPrint "Histogram size (H)" $ Just $ untyped hist_H - emit $ Imp.DebugPrint "Input elements per histogram (N)" $ Just $ untyped hist_N - emit $ - Imp.DebugPrint "Number of segments" $ - Just $ untyped $ product $ map (toInt64Exp . snd) segment_dims - emit $ Imp.DebugPrint "Histogram element size (el_size)" $ Just $ untyped hist_el_size - emit $ Imp.DebugPrint "Race factor (RF)" $ Just $ untyped hist_RF - emit $ Imp.DebugPrint "Memory per set of subhistograms per segment" $ Just $ untyped h - emit $ Imp.DebugPrint "Memory per set of subhistograms times segments" $ Just $ untyped seg_h - - (use_local_memory, run_in_local_memory) <- - localMemoryCase map_pes hist_T space hist_H hist_el_size hist_N hist_RF slugs kbody - - sIf use_local_memory run_in_local_memory $ - histKernelGlobal map_pes num_groups group_size space slugs kbody - - let pes_per_op = chunks (map (length . histDest) ops) all_red_pes - - forM_ (zip3 slugs pes_per_op ops) $ \(slug, red_pes, op) -> do - let num_histos = slugNumSubhistos slug - subhistos = map subhistosArray $ slugSubhistos slug - - let unitHistoCase = - -- This is OK because the memory blocks are at least as - -- large as the ones we are supposed to use for the result. - forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do - pe_mem <- - memLocationName . entryArrayLocation - <$> lookupArray (patElemName pe) - subhisto_mem <- - memLocationName . entryArrayLocation - <$> lookupArray subhisto - emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device" - - sIf (tvExp num_histos .==. 1) unitHistoCase $ do - -- For the segmented reduction, we keep the segment dimensions - -- unchanged. To this, we add two dimensions: one over the number - -- of buckets, and one over the number of subhistograms. This - -- inner dimension is the one that is collapsed in the reduction. - let num_buckets = histWidth op - - bucket_id <- newVName "bucket_id" - subhistogram_id <- newVName "subhistogram_id" - vector_ids <- - mapM (const $ newVName "vector_id") $ - shapeDims $ histShape op - - flat_gtid <- newVName "flat_gtid" - - let lvl = SegThread num_groups group_size SegVirt - segred_space = - SegSpace flat_gtid $ - segment_dims - ++ [(bucket_id, num_buckets)] - ++ zip vector_ids (shapeDims $ histShape op) - ++ [(subhistogram_id, Var $ tvVar num_histos)] - - let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty - compileSegRed' (Pattern [] red_pes) lvl segred_space [segred_op] $ \red_cont -> - red_cont $ - flip map subhistos $ \subhisto -> - ( Var subhisto, - map Imp.vi64 $ - map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids - ) - - emit $ Imp.DebugPrint "" Nothing - where - segment_dims = init $ unSegSpace space diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegMap.hs +++ /dev/null @@ -1,61 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} - --- | Code generation for 'SegMap' is quite straightforward. The only --- trick is virtualisation in case the physical number of threads is --- not sufficient to cover the logical thread space. This is handled --- by having actual workgroups run a loop to imitate multiple workgroups. -module Futhark.CodeGen.ImpGen.Kernels.SegMap (compileSegMap) where - -import Control.Monad.Except -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.IR.KernelsMem -import Futhark.Util.IntegralExp (divUp) -import Prelude hiding (quot, rem) - --- | Compile 'SegMap' instance code. -compileSegMap :: - Pattern KernelsMem -> - SegLevel -> - SegSpace -> - KernelBody KernelsMem -> - CallKernelGen () -compileSegMap pat lvl space kbody = do - let (is, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - num_groups' = toInt64Exp <$> segNumGroups lvl - group_size' = toInt64Exp <$> segGroupSize lvl - - emit $ Imp.DebugPrint "\n# SegMap" Nothing - case lvl of - SegThread {} -> do - let virt_num_groups = - sExt32 $ product dims' `divUp` unCount group_size' - sKernelThread "segmap" num_groups' group_size' (segFlat space) $ - virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do - local_tid <- kernelLocalThreadId . kernelConstants <$> askEnv - - global_tid <- - dPrimVE "global_tid" $ - sExt64 group_id * sExt64 (unCount group_size') - + sExt64 local_tid - - dIndexSpace (zip is dims') global_tid - - sWhen (isActive $ unSegSpace space) $ - compileStms mempty (kernelBodyStms kbody) $ - zipWithM_ (compileThreadResult space) (patternElements pat) $ - kernelBodyResult kbody - SegGroup {} -> - sKernelGroup "segmap_intragroup" num_groups' group_size' (segFlat space) $ do - let virt_num_groups = sExt32 $ product dims' - precomputeSegOpIDs (kernelBodyStms kbody) $ - virtualiseGroups (segVirt lvl) virt_num_groups $ \group_id -> do - dIndexSpace (zip is dims') $ sExt64 group_id - - compileStms mempty (kernelBodyStms kbody) $ - zipWithM_ (compileGroupResult space) (patternElements pat) $ - kernelBodyResult kbody - emit $ Imp.DebugPrint "" Nothing diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs +++ /dev/null @@ -1,838 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} - --- | We generate code for non-segmented/single-segment SegRed using --- the basic approach outlined in the paper "Design and GPGPU --- Performance of Futhark’s Redomap Construct" (ARRAY '16). The main --- deviations are: --- --- * While we still use two-phase reduction, we use only a single --- kernel, with the final workgroup to write a result (tracked via --- an atomic counter) performing the final reduction as well. --- --- * Instead of depending on storage layout transformations to handle --- non-commutative reductions efficiently, we slide a --- @groupsize@-sized window over the input, and perform a parallel --- reduction for each window. This sacrifices the notion of --- efficient sequentialisation, but is sometimes faster and --- definitely simpler and more predictable (and uses less auxiliary --- storage). --- --- For segmented reductions we use the approach from "Strategies for --- Regular Segmented Reductions on GPU" (FHPC '17). This involves --- having two different strategies, and dynamically deciding which one --- to use based on the number of segments and segment size. We use the --- (static) @group_size@ to decide which of the following two --- strategies to choose: --- --- * Large: uses one or more groups to process a single segment. If --- multiple groups are used per segment, the intermediate reduction --- results must be recursively reduced, until there is only a single --- value per segment. --- --- Each thread /can/ read multiple elements, which will greatly --- increase performance; however, if the reduction is --- non-commutative we will have to use a less efficient traversal --- (with interim group-wide reductions) to enable coalesced memory --- accesses, just as in the non-segmented case. --- --- * Small: is used to let each group process *multiple* segments --- within a group. We will only use this approach when we can --- process at least two segments within a single group. In those --- cases, we would allocate a /whole/ group per segment with the --- large strategy, but at most 50% of the threads in the group would --- have any element to read, which becomes highly inefficient. -module Futhark.CodeGen.ImpGen.Kernels.SegRed - ( compileSegRed, - compileSegRed', - DoSegBody, - ) -where - -import Control.Monad.Except -import Data.List (genericLength, zip7) -import Data.Maybe -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.Error -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Transform.Rename -import Futhark.Util (chunks) -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Prelude hiding (quot, rem) - --- | The maximum number of operators we support in a single SegRed. --- This limit arises out of the static allocation of counters. -maxNumOps :: Int32 -maxNumOps = 10 - --- | Code generation for the body of the SegRed, taking a continuation --- for saving the results of the body. The results should be --- represented as a pairing of a t'SubExp' along with a list of --- indexes into that t'SubExp' for reading the result. -type DoSegBody = ([(SubExp, [Imp.TExp Int64])] -> InKernelGen ()) -> InKernelGen () - --- | Compile 'SegRed' instance to host-level code with calls to --- various kernels. -compileSegRed :: - Pattern KernelsMem -> - SegLevel -> - SegSpace -> - [SegBinOp KernelsMem] -> - KernelBody KernelsMem -> - CallKernelGen () -compileSegRed pat lvl space reds body = - compileSegRed' pat lvl space reds $ \red_cont -> - compileStms mempty (kernelBodyStms body) $ do - let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult body - - sComment "save map-out results" $ do - let map_arrs = drop (segBinOpResults reds) $ patternElements pat - zipWithM_ (compileThreadResult space) map_arrs map_res - - red_cont $ zip (map kernelResultSubExp red_res) $ repeat [] - --- | Like 'compileSegRed', but where the body is a monadic action. -compileSegRed' :: - Pattern KernelsMem -> - SegLevel -> - SegSpace -> - [SegBinOp KernelsMem] -> - DoSegBody -> - CallKernelGen () -compileSegRed' pat lvl space reds body - | genericLength reds > maxNumOps = - compilerLimitationS $ - "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported." - | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space = - nonsegmentedReduction pat num_groups group_size space reds body - | otherwise = do - let group_size' = toInt64Exp $ unCount group_size - segment_size = toInt64Exp $ last $ segSpaceDims space - use_small_segments = segment_size * 2 .<. group_size' - sIf - use_small_segments - (smallSegmentsReduction pat num_groups group_size space reds body) - (largeSegmentsReduction pat num_groups group_size space reds body) - where - num_groups = segNumGroups lvl - group_size = segGroupSize lvl - --- | Prepare intermediate arrays for the reduction. Prim-typed --- arguments go in local memory (so we need to do the allocation of --- those arrays inside the kernel), while array-typed arguments go in --- global memory. Allocations for the former have already been --- performed. This policy is baked into how the allocations are done --- in ExplicitAllocations. -intermediateArrays :: - Count GroupSize SubExp -> - SubExp -> - SegBinOp KernelsMem -> - InKernelGen [VName] -intermediateArrays (Count group_size) num_threads (SegBinOp _ red_op nes _) = do - let red_op_params = lambdaParams red_op - (red_acc_params, _) = splitAt (length nes) red_op_params - forM red_acc_params $ \p -> - case paramDec p of - MemArray pt shape _ (ArrayIn mem _) -> do - let shape' = Shape [num_threads] <> shape - sArray "red_arr" pt shape' $ - ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' - _ -> do - let pt = elemType $ paramType p - shape = Shape [group_size] - sAllocArray "red_arr" pt shape $ Space "local" - --- | Arrays for storing group results. --- --- The group-result arrays have an extra dimension (of size groupsize) --- because they are also used for keeping vectorised accumulators for --- first-stage reduction, if necessary. When actually storing group --- results, the first index is set to 0. -groupResultArrays :: - Count NumGroups SubExp -> - Count GroupSize SubExp -> - [SegBinOp KernelsMem] -> - CallKernelGen [[VName]] -groupResultArrays (Count virt_num_groups) (Count group_size) reds = - forM reds $ \(SegBinOp _ lam _ shape) -> - forM (lambdaReturnType lam) $ \t -> do - let pt = elemType t - full_shape = Shape [group_size, virt_num_groups] <> shape <> arrayShape t - -- Move the groupsize dimension last to ensure coalesced - -- memory access. - perm = [1 .. shapeRank full_shape -1] ++ [0] - sAllocArrayPerm "group_res_arr" pt full_shape (Space "device") perm - -nonsegmentedReduction :: - Pattern KernelsMem -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - SegSpace -> - [SegBinOp KernelsMem] -> - DoSegBody -> - CallKernelGen () -nonsegmentedReduction segred_pat num_groups group_size space reds body = do - let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - global_tid = Imp.vi64 $ segFlat space - w = last dims' - - counter <- - sStaticArray "counter" (Space "device") int32 $ - Imp.ArrayValues $ replicate (fromIntegral maxNumOps) $ IntValue $ Int32Value 0 - - reds_group_res_arrs <- groupResultArrays num_groups group_size reds - - num_threads <- - dPrimV "num_threads" $ - unCount num_groups' * unCount group_size' - - emit $ Imp.DebugPrint "\n# SegRed" Nothing - - sKernelThread "segred_nonseg" num_groups' group_size' (segFlat space) $ do - constants <- kernelConstants <$> askEnv - sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" - reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds - - -- Since this is the nonsegmented case, all outer segment IDs must - -- necessarily be 0. - forM_ gtids $ \v -> dPrimV_ v (0 :: Imp.TExp Int64) - - let num_elements = Imp.elements w - elems_per_thread = - num_elements - `divUp` Imp.elements (sExt64 (kernelNumThreads constants)) - - slugs <- - mapM (segBinOpSlug (kernelLocalThreadId constants) (kernelGroupId constants)) $ - zip3 reds reds_arrs reds_group_res_arrs - reds_op_renamed <- - reductionStageOne - constants - (zip gtids dims') - num_elements - global_tid - elems_per_thread - (tvVar num_threads) - slugs - body - - let segred_pes = - chunks (map (length . segBinOpNeutral) reds) $ - patternElements segred_pat - forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $ - \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do - let (red_x_params, red_y_params) = splitAt (length nes) $ lambdaParams red_op - reductionStageTwo - constants - pes - (kernelGroupId constants) - 0 - [0] - 0 - (sExt64 $ kernelNumGroups constants) - slug - red_x_params - red_y_params - red_op_renamed - nes - 1 - counter - (fromInteger i) - sync_arr - group_res_arrs - red_arrs - - emit $ Imp.DebugPrint "" Nothing - -smallSegmentsReduction :: - Pattern KernelsMem -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - SegSpace -> - [SegBinOp KernelsMem] -> - DoSegBody -> - CallKernelGen () -smallSegmentsReduction (Pattern _ segred_pes) num_groups group_size space reds body = do - let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - segment_size = last dims' - - -- Careful to avoid division by zero now. - segment_size_nonzero <- - dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size - - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' - let num_segments = product $ init dims' - segments_per_group = unCount group_size' `quot` segment_size_nonzero - required_groups = sExt32 $ num_segments `divUp` segments_per_group - - emit $ Imp.DebugPrint "\n# SegRed-small" Nothing - emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments - emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size - emit $ Imp.DebugPrint "segments_per_group" $ Just $ untyped segments_per_group - emit $ Imp.DebugPrint "required_groups" $ Just $ untyped required_groups - - sKernelThread "segred_small" num_groups' group_size' (segFlat space) $ do - constants <- kernelConstants <$> askEnv - reds_arrs <- mapM (intermediateArrays group_size (Var $ tvVar num_threads)) reds - - -- We probably do not have enough actual workgroups to cover the - -- entire iteration space. Some groups thus have to perform double - -- duty; we put an outer loop to accomplish this. - virtualiseGroups SegVirt required_groups $ \group_id' -> do - -- Compute the 'n' input indices. The outer 'n-1' correspond to - -- the segment ID, and are computed from the group id. The inner - -- is computed from the local thread id, and may be out-of-bounds. - let ltid = sExt64 $ kernelLocalThreadId constants - segment_index = - (ltid `quot` segment_size_nonzero) - + (sExt64 group_id' * sExt64 segments_per_group) - index_within_segment = ltid `rem` segment_size - - dIndexSpace (zip (init gtids) (init dims')) segment_index - dPrimV_ (last gtids) index_within_segment - - let out_of_bounds = - forM_ (zip reds reds_arrs) $ \(SegBinOp _ _ nes _, red_arrs) -> - forM_ (zip red_arrs nes) $ \(arr, ne) -> - copyDWIMFix arr [ltid] ne [] - - in_bounds = - body $ \red_res -> - sComment "save results to be reduced" $ do - let red_dests = zip (concat reds_arrs) $ repeat [ltid] - forM_ (zip red_dests red_res) $ \((d, d_is), (res, res_is)) -> - copyDWIMFix d d_is res res_is - - sComment "apply map function if in bounds" $ - sIf - ( segment_size .>. 0 - .&&. isActive (init $ zip gtids dims) - .&&. ltid .<. segment_size * segments_per_group - ) - in_bounds - out_of_bounds - - sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. - let crossesSegment from to = - (sExt64 to - sExt64 from) .>. (sExt64 to `rem` segment_size) - sWhen (segment_size .>. 0) $ - sComment "perform segmented scan to imitate reduction" $ - forM_ (zip reds reds_arrs) $ \(SegBinOp _ red_op _ _, red_arrs) -> - groupScan - (Just crossesSegment) - (sExt64 $ tvExp num_threads) - (segment_size * segments_per_group) - red_op - red_arrs - - sOp $ Imp.Barrier Imp.FenceLocal - - sComment "save final values of segments" $ - sWhen - ( sExt64 group_id' * segments_per_group + sExt64 ltid .<. num_segments - .&&. ltid .<. segments_per_group - ) - $ forM_ (zip segred_pes (concat reds_arrs)) $ \(pe, arr) -> do - -- Figure out which segment result this thread should write... - let flat_segment_index = - sExt64 group_id' * segments_per_group + sExt64 ltid - gtids' = - unflattenIndex (init dims') flat_segment_index - copyDWIMFix - (patElemName pe) - gtids' - (Var arr) - [(ltid + 1) * segment_size_nonzero - 1] - - -- Finally another barrier, because we will be writing to the - -- local memory array first thing in the next iteration. - sOp $ Imp.Barrier Imp.FenceLocal - - emit $ Imp.DebugPrint "" Nothing - -largeSegmentsReduction :: - Pattern KernelsMem -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - SegSpace -> - [SegBinOp KernelsMem] -> - DoSegBody -> - CallKernelGen () -largeSegmentsReduction segred_pat num_groups group_size space reds body = do - let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - num_segments = product $ init dims' - segment_size = last dims' - num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - - (groups_per_segment, elems_per_thread) <- - groupsPerSegmentAndElementsPerThread - segment_size - num_segments - num_groups' - group_size' - virt_num_groups <- - dPrimV "virt_num_groups" $ - groups_per_segment * num_segments - - num_threads <- - dPrimV "num_threads" $ - unCount num_groups' * unCount group_size' - - threads_per_segment <- - dPrimV "threads_per_segment" $ - groups_per_segment * unCount group_size' - - emit $ Imp.DebugPrint "\n# SegRed-large" Nothing - emit $ Imp.DebugPrint "num_segments" $ Just $ untyped num_segments - emit $ Imp.DebugPrint "segment_size" $ Just $ untyped segment_size - emit $ Imp.DebugPrint "virt_num_groups" $ Just $ untyped $ tvExp virt_num_groups - emit $ Imp.DebugPrint "num_groups" $ Just $ untyped $ Imp.unCount num_groups' - emit $ Imp.DebugPrint "group_size" $ Just $ untyped $ Imp.unCount group_size' - emit $ Imp.DebugPrint "elems_per_thread" $ Just $ untyped $ Imp.unCount elems_per_thread - emit $ Imp.DebugPrint "groups_per_segment" $ Just $ untyped groups_per_segment - - reds_group_res_arrs <- groupResultArrays (Count (tvSize virt_num_groups)) group_size reds - - -- In principle we should have a counter for every segment. Since - -- the number of segments is a dynamic quantity, we would have to - -- allocate and zero out an array here, which is expensive. - -- However, we exploit the fact that the number of segments being - -- reduced at any point in time is limited by the number of - -- workgroups. If we bound the number of workgroups, we can get away - -- with using that many counters. FIXME: Is this limit checked - -- anywhere? There are other places in the compiler that will fail - -- if the group count exceeds the maximum group size, which is at - -- most 1024 anyway. - let num_counters = fromIntegral maxNumOps * 1024 - counter <- - sStaticArray "counter" (Space "device") int32 $ - Imp.ArrayZeros num_counters - - sKernelThread "segred_large" num_groups' group_size' (segFlat space) $ do - constants <- kernelConstants <$> askEnv - reds_arrs <- mapM (intermediateArrays group_size (tvSize num_threads)) reds - sync_arr <- sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" - - -- We probably do not have enough actual workgroups to cover the - -- entire iteration space. Some groups thus have to perform double - -- duty; we put an outer loop to accomplish this. - virtualiseGroups SegVirt (sExt32 (tvExp virt_num_groups)) $ \group_id -> do - let segment_gtids = init gtids - w = last dims - local_tid = kernelLocalThreadId constants - - flat_segment_id <- - dPrimVE "flat_segment_id" $ - group_id `quot` sExt32 groups_per_segment - - global_tid <- - dPrimVE "global_tid" $ - (sExt64 group_id * sExt64 (unCount group_size') + sExt64 local_tid) - `rem` (sExt64 (unCount group_size') * groups_per_segment) - - let first_group_for_segment = sExt64 flat_segment_id * groups_per_segment - dIndexSpace (zip segment_gtids (init dims')) $sExt64 flat_segment_id - dPrim_ (last gtids) int64 - let num_elements = Imp.elements $ toInt64Exp w - - slugs <- - mapM (segBinOpSlug local_tid group_id) $ - zip3 reds reds_arrs reds_group_res_arrs - reds_op_renamed <- - reductionStageOne - constants - (zip gtids dims') - num_elements - global_tid - elems_per_thread - (tvVar threads_per_segment) - slugs - body - - let segred_pes = - chunks (map (length . segBinOpNeutral) reds) $ - patternElements segred_pat - - multiple_groups_per_segment = - forM_ (zip7 reds reds_arrs reds_group_res_arrs segred_pes slugs reds_op_renamed [0 ..]) $ - \(SegBinOp _ red_op nes _, red_arrs, group_res_arrs, pes, slug, red_op_renamed, i) -> do - let (red_x_params, red_y_params) = - splitAt (length nes) $ lambdaParams red_op - reductionStageTwo - constants - pes - group_id - flat_segment_id - (map Imp.vi64 segment_gtids) - (sExt64 first_group_for_segment) - groups_per_segment - slug - red_x_params - red_y_params - red_op_renamed - nes - (fromIntegral num_counters) - counter - (fromInteger i) - sync_arr - group_res_arrs - red_arrs - - one_group_per_segment = - comment "first thread in group saves final result to memory" $ - forM_ (zip slugs segred_pes) $ \(slug, pes) -> - sWhen (local_tid .==. 0) $ - forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) -> - copyDWIMFix (patElemName v) (map Imp.vi64 segment_gtids) (Var acc) acc_is - - sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment - - emit $ Imp.DebugPrint "" Nothing - --- Careful to avoid division by zero here. We have at least one group --- per segment. -groupsPerSegmentAndElementsPerThread :: - Imp.TExp Int64 -> - Imp.TExp Int64 -> - Count NumGroups (Imp.TExp Int64) -> - Count GroupSize (Imp.TExp Int64) -> - CallKernelGen - ( Imp.TExp Int64, - Imp.Count Imp.Elements (Imp.TExp Int64) - ) -groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = do - groups_per_segment <- - dPrimVE "groups_per_segment" $ - unCount num_groups_hint `divUp` sMax64 1 num_segments - elements_per_thread <- - dPrimVE "elements_per_thread" $ - segment_size `divUp` (unCount group_size * groups_per_segment) - return (groups_per_segment, Imp.elements elements_per_thread) - --- | A SegBinOp with auxiliary information. -data SegBinOpSlug = SegBinOpSlug - { slugOp :: SegBinOp KernelsMem, - -- | The arrays used for computing the intra-group reduction - -- (either local or global memory). - slugArrs :: [VName], - -- | Places to store accumulator in stage 1 reduction. - slugAccs :: [(VName, [Imp.TExp Int64])] - } - -slugBody :: SegBinOpSlug -> Body KernelsMem -slugBody = lambdaBody . segBinOpLambda . slugOp - -slugParams :: SegBinOpSlug -> [LParam KernelsMem] -slugParams = lambdaParams . segBinOpLambda . slugOp - -slugNeutral :: SegBinOpSlug -> [SubExp] -slugNeutral = segBinOpNeutral . slugOp - -slugShape :: SegBinOpSlug -> Shape -slugShape = segBinOpShape . slugOp - -slugsComm :: [SegBinOpSlug] -> Commutativity -slugsComm = mconcat . map (segBinOpComm . slugOp) - -accParams, nextParams :: SegBinOpSlug -> [LParam KernelsMem] -accParams slug = take (length (slugNeutral slug)) $ slugParams slug -nextParams slug = drop (length (slugNeutral slug)) $ slugParams slug - -segBinOpSlug :: Imp.TExp Int32 -> Imp.TExp Int32 -> (SegBinOp KernelsMem, [VName], [VName]) -> InKernelGen SegBinOpSlug -segBinOpSlug local_tid group_id (op, group_res_arrs, param_arrs) = - SegBinOpSlug op group_res_arrs - <$> zipWithM mkAcc (lambdaParams (segBinOpLambda op)) param_arrs - where - mkAcc p param_arr - | Prim t <- paramType p, - shapeRank (segBinOpShape op) == 0 = do - acc <- dPrim (baseString (paramName p) <> "_acc") t - return (tvVar acc, []) - | otherwise = - return (param_arr, [sExt64 local_tid, sExt64 group_id]) - -reductionStageZero :: - KernelConstants -> - [(VName, Imp.TExp Int64)] -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - Imp.TExp Int64 -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - VName -> - [SegBinOpSlug] -> - DoSegBody -> - InKernelGen ([Lambda KernelsMem], InKernelGen ()) -reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do - let (gtids, _dims) = unzip ispace - gtid = mkTV (last gtids) int64 - local_tid = sExt64 $ kernelLocalThreadId constants - - -- Figure out how many elements this thread should process. - chunk_size <- dPrim "chunk_size" int64 - let ordering = case slugsComm slugs of - Commutative -> SplitStrided $ Var threads_per_segment - Noncommutative -> SplitContiguous - computeThreadChunkSize ordering (sExt64 global_tid) elems_per_thread num_elements chunk_size - - dScope Nothing $ scopeOfLParams $ concatMap slugParams slugs - - sComment "neutral-initialise the accumulators" $ - forM_ slugs $ \slug -> - forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> - sLoopNest (slugShape slug) $ \vec_is -> - copyDWIMFix acc (acc_is ++ vec_is) ne [] - - slugs_op_renamed <- mapM (renameLambda . segBinOpLambda . slugOp) slugs - - let doTheReduction = - forM_ (zip slugs_op_renamed slugs) $ \(slug_op_renamed, slug) -> - sLoopNest (slugShape slug) $ \vec_is -> do - comment "to reduce current chunk, first store our result in memory" $ do - forM_ (zip (slugParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> - copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is) - - forM_ (zip (slugArrs slug) (slugParams slug)) $ \(arr, p) -> - when (primType $ paramType p) $ - copyDWIMFix arr [local_tid] (Var $ paramName p) [] - - sOp $ Imp.ErrorSync Imp.FenceLocal -- Also implicitly barrier. - groupReduce (sExt32 (kernelGroupSize constants)) slug_op_renamed (slugArrs slug) - - sOp $ Imp.Barrier Imp.FenceLocal - - sComment "first thread saves the result in accumulator" $ - sWhen (local_tid .==. 0) $ - forM_ (zip (slugAccs slug) (lambdaParams slug_op_renamed)) $ \((acc, acc_is), p) -> - copyDWIMFix acc (acc_is ++ vec_is) (Var $ paramName p) [] - - -- If this is a non-commutative reduction, each thread must run the - -- loop the same number of iterations, because we will be performing - -- a group-wide reduction in there. - let comm = slugsComm slugs - (bound, check_bounds) = - case comm of - Commutative -> (tvExp chunk_size, id) - Noncommutative -> - ( Imp.unCount elems_per_thread, - sWhen (tvExp gtid .<. Imp.unCount num_elements) - ) - - sFor "i" bound $ \i -> do - gtid - <-- case comm of - Commutative -> - global_tid + Imp.vi64 threads_per_segment * i - Noncommutative -> - let index_in_segment = global_tid `quot` kernelGroupSize constants - in sExt64 local_tid - + (index_in_segment * Imp.unCount elems_per_thread + i) - * kernelGroupSize constants - - check_bounds $ - sComment "apply map function" $ - body $ \all_red_res -> do - let slugs_res = chunks (map (length . slugNeutral) slugs) all_red_res - - forM_ (zip slugs slugs_res) $ \(slug, red_res) -> - sLoopNest (slugShape slug) $ \vec_is -> do - sComment "load accumulator" $ - forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> - copyDWIMFix (paramName p) [] (Var acc) (acc_is ++ vec_is) - sComment "load new values" $ - forM_ (zip (nextParams slug) red_res) $ \(p, (res, res_is)) -> - copyDWIMFix (paramName p) [] res (res_is ++ vec_is) - sComment "apply reduction operator" $ - compileStms mempty (bodyStms $ slugBody slug) $ - sComment "store in accumulator" $ - forM_ - ( zip - (slugAccs slug) - (bodyResult $ slugBody slug) - ) - $ \((acc, acc_is), se) -> - copyDWIMFix acc (acc_is ++ vec_is) se [] - - case comm of - Noncommutative -> do - doTheReduction - sComment "first thread keeps accumulator; others reset to neutral element" $ do - let reset_to_neutral = - forM_ slugs $ \slug -> - forM_ (zip (slugAccs slug) (slugNeutral slug)) $ \((acc, acc_is), ne) -> - sLoopNest (slugShape slug) $ \vec_is -> - copyDWIMFix acc (acc_is ++ vec_is) ne [] - sUnless (local_tid .==. 0) reset_to_neutral - _ -> return () - - return (slugs_op_renamed, doTheReduction) - -reductionStageOne :: - KernelConstants -> - [(VName, Imp.TExp Int64)] -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - Imp.TExp Int64 -> - Imp.Count Imp.Elements (Imp.TExp Int64) -> - VName -> - [SegBinOpSlug] -> - DoSegBody -> - InKernelGen [Lambda KernelsMem] -reductionStageOne constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body = do - (slugs_op_renamed, doTheReduction) <- - reductionStageZero constants ispace num_elements global_tid elems_per_thread threads_per_segment slugs body - - case slugsComm slugs of - Noncommutative -> - forM_ slugs $ \slug -> - forM_ (zip (accParams slug) (slugAccs slug)) $ \(p, (acc, acc_is)) -> - copyDWIMFix (paramName p) [] (Var acc) acc_is - _ -> doTheReduction - - return slugs_op_renamed - -reductionStageTwo :: - KernelConstants -> - [PatElem KernelsMem] -> - Imp.TExp Int32 -> - Imp.TExp Int32 -> - [Imp.TExp Int64] -> - Imp.TExp Int64 -> - Imp.TExp Int64 -> - SegBinOpSlug -> - [LParam KernelsMem] -> - [LParam KernelsMem] -> - Lambda KernelsMem -> - [SubExp] -> - Imp.TExp Int32 -> - VName -> - Imp.TExp Int32 -> - VName -> - [VName] -> - [VName] -> - InKernelGen () -reductionStageTwo - constants - segred_pes - group_id - flat_segment_id - segment_gtids - first_group_for_segment - groups_per_segment - slug - red_x_params - red_y_params - red_op_renamed - nes - num_counters - counter - counter_i - sync_arr - group_res_arrs - red_arrs = do - let local_tid = kernelLocalThreadId constants - group_size = kernelGroupSize constants - old_counter <- dPrim "old_counter" int32 - (counter_mem, _, counter_offset) <- - fullyIndexArray - counter - [ sExt64 $ - counter_i * num_counters - + flat_segment_id `rem` num_counters - ] - comment "first thread in group saves group result to global memory" $ - sWhen (local_tid .==. 0) $ do - forM_ (take (length nes) $ zip group_res_arrs (slugAccs slug)) $ \(v, (acc, acc_is)) -> - copyDWIMFix v [0, sExt64 group_id] (Var acc) acc_is - sOp $ Imp.MemFence Imp.FenceGlobal - -- Increment the counter, thus stating that our result is - -- available. - sOp $ - Imp.Atomic DefaultSpace $ - Imp.AtomicAdd - Int32 - (tvVar old_counter) - counter_mem - counter_offset - $ untyped (1 :: Imp.TExp Int32) - -- Now check if we were the last group to write our result. If - -- so, it is our responsibility to produce the final result. - sWrite sync_arr [0] $ untyped $ tvExp old_counter .==. groups_per_segment - 1 - - sOp $ Imp.Barrier Imp.FenceGlobal - - is_last_group <- dPrim "is_last_group" Bool - copyDWIMFix (tvVar is_last_group) [] (Var sync_arr) [0] - sWhen (tvExp is_last_group) $ do - -- The final group has written its result (and it was - -- us!), so read in all the group results and perform the - -- final stage of the reduction. But first, we reset the - -- counter so it is ready for next time. This is done - -- with an atomic to avoid warnings about write/write - -- races in oclgrind. - sWhen (local_tid .==. 0) $ - sOp $ - Imp.Atomic DefaultSpace $ - Imp.AtomicAdd Int32 (tvVar old_counter) counter_mem counter_offset $ - untyped $ negate groups_per_segment - - sLoopNest (slugShape slug) $ \vec_is -> do - -- 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 - -- way that may induce non-coalesced accesses, but the total - -- number of accesses should be tiny here. - comment "read in the per-group-results" $ do - read_per_thread <- - dPrimVE "read_per_thread" $ - groups_per_segment `divUp` sExt64 group_size - - forM_ (zip red_x_params nes) $ \(p, ne) -> - copyDWIMFix (paramName p) [] ne [] - - sFor "i" read_per_thread $ \i -> do - group_res_id <- - dPrimVE "group_res_id" $ - sExt64 local_tid * read_per_thread + i - index_of_group_res <- - dPrimVE "index_of_group_res" $ - first_group_for_segment + group_res_id - - sWhen (group_res_id .<. groups_per_segment) $ do - forM_ (zip red_y_params group_res_arrs) $ - \(p, group_res_arr) -> - copyDWIMFix - (paramName p) - [] - (Var group_res_arr) - ([0, index_of_group_res] ++ vec_is) - - compileStms mempty (bodyStms $ slugBody slug) $ - forM_ (zip red_x_params (bodyResult $ slugBody slug)) $ \(p, se) -> - copyDWIMFix (paramName p) [] se [] - - forM_ (zip red_x_params red_arrs) $ \(p, arr) -> - when (primType $ paramType p) $ - copyDWIMFix arr [sExt64 local_tid] (Var $ paramName p) [] - - sOp $ Imp.Barrier Imp.FenceLocal - - sComment "reduce the per-group results" $ do - groupReduce (sExt32 group_size) red_op_renamed red_arrs - - sComment "and back to memory with the final result" $ - sWhen (local_tid .==. 0) $ - forM_ (zip segred_pes $ lambdaParams red_op_renamed) $ \(pe, p) -> - copyDWIMFix - (patElemName pe) - (segment_gtids ++ vec_is) - (Var $ paramName p) - [] diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan.hs +++ /dev/null @@ -1,68 +0,0 @@ --- | Code generation for 'SegScan'. Dispatches to either a --- single-pass or two-pass implementation, depending on the nature of --- the scan and the chosen abckend. -module Futhark.CodeGen.ImpGen.Kernels.SegScan (compileSegScan) where - -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen hiding (compileProg) -import Futhark.CodeGen.ImpGen.Kernels.Base -import qualified Futhark.CodeGen.ImpGen.Kernels.SegScan.SinglePass as SinglePass -import qualified Futhark.CodeGen.ImpGen.Kernels.SegScan.TwoPass as TwoPass -import Futhark.IR.KernelsMem - --- The single-pass scan does not support multiple operators, so jam --- them together here. -combineScans :: [SegBinOp KernelsMem] -> SegBinOp KernelsMem -combineScans ops = - SegBinOp - { segBinOpComm = mconcat (map segBinOpComm ops), - segBinOpLambda = lam', - segBinOpNeutral = concatMap segBinOpNeutral ops, - segBinOpShape = mempty -- Assumed - } - where - lams = map segBinOpLambda ops - xParams lam = take (length (lambdaReturnType lam)) (lambdaParams lam) - yParams lam = drop (length (lambdaReturnType lam)) (lambdaParams lam) - lam' = - Lambda - { lambdaParams = concatMap xParams lams ++ concatMap yParams lams, - lambdaReturnType = concatMap lambdaReturnType lams, - lambdaBody = - Body - () - (mconcat (map (bodyStms . lambdaBody) lams)) - (concatMap (bodyResult . lambdaBody) lams) - } - -canBeSinglePass :: SegSpace -> [SegBinOp KernelsMem] -> Maybe (SegBinOp KernelsMem) -canBeSinglePass space ops - | [_] <- unSegSpace space, - all ok ops = - Just $ combineScans ops - | otherwise = - Nothing - where - ok op = - segBinOpShape op == mempty - && all primType (lambdaReturnType (segBinOpLambda op)) - --- | Compile 'SegScan' instance to host-level code with calls to --- various kernels. -compileSegScan :: - Pattern KernelsMem -> - SegLevel -> - SegSpace -> - [SegBinOp KernelsMem] -> - KernelBody KernelsMem -> - CallKernelGen () -compileSegScan pat lvl space scans kbody = sWhen (0 .<. n) $ do - emit $ Imp.DebugPrint "\n# SegScan" Nothing - target <- hostTarget <$> askEnv - case target of - CUDA - | Just scan' <- canBeSinglePass space scans -> - SinglePass.compileSegScan pat lvl space scan' kbody - _ -> TwoPass.compileSegScan pat lvl space scans kbody - where - n = product $ map toInt64Exp $ segSpaceDims space diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan/SinglePass.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan/SinglePass.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan/SinglePass.hs +++ /dev/null @@ -1,489 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} - --- | Code generation for segmented and non-segmented scans. Uses a --- fast single-pass algorithm, but which only works on NVIDIA GPUs and --- with some constraints on the operator. We use this when we can. -module Futhark.CodeGen.ImpGen.Kernels.SegScan.SinglePass (compileSegScan) where - -import Control.Monad.Except -import Data.List (zip4) -import Data.Maybe -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Transform.Rename -import Futhark.Util (takeLast) -import Futhark.Util.IntegralExp (IntegralExp, divUp, quot) -import Prelude hiding (quot) - -xParams, yParams :: SegBinOp KernelsMem -> [LParam KernelsMem] -xParams scan = - take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) -yParams scan = - drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) - -alignTo :: IntegralExp a => a -> a -> a -alignTo x a = (x `divUp` a) * a - -createLocalArrays :: - Count GroupSize SubExp -> - SubExp -> - [PrimType] -> - InKernelGen (VName, [VName], [VName], VName, VName, [VName]) -createLocalArrays (Count groupSize) m types = do - let groupSizeE = toInt64Exp groupSize - workSize = toInt64Exp m * groupSizeE - prefixArraysSize = - foldl (\acc tySize -> alignTo acc tySize + tySize * groupSizeE) 0 $ - map primByteSize types - maxTransposedArraySize = - foldl1 sMax64 $ map (\ty -> workSize * primByteSize ty) types - - warpSize :: Num a => a - warpSize = 32 - maxWarpExchangeSize = - foldl (\acc tySize -> alignTo acc tySize + tySize * fromInteger warpSize) 0 $ - map primByteSize types - maxLookbackSize = maxWarpExchangeSize + warpSize - size = Imp.bytes $ maxLookbackSize `sMax64` prefixArraysSize `sMax64` maxTransposedArraySize - - varTE :: TV Int64 -> TPrimExp Int64 VName - varTE = le64 . tvVar - - byteOffsets <- - mapM (fmap varTE . dPrimV "byte_offsets") $ - scanl (\off tySize -> alignTo off tySize + toInt64Exp groupSize * tySize) 0 $ - map primByteSize types - - warpByteOffsets <- - mapM (fmap varTE . dPrimV "warp_byte_offset") $ - scanl (\off tySize -> alignTo off tySize + warpSize * tySize) warpSize $ - map primByteSize types - - sComment "Allocate reused shared memeory" $ return () - - localMem <- sAlloc "local_mem" size (Space "local") - transposeArrayLength <- dPrimV "trans_arr_len" workSize - - sharedId <- sArrayInMem "shared_id" int32 (Shape [constant (1 :: Int32)]) localMem - sharedReadOffset <- sArrayInMem "shared_read_offset" int32 (Shape [constant (1 :: Int32)]) localMem - - transposedArrays <- - forM types $ \ty -> - sArrayInMem - "local_transpose_arr" - ty - (Shape [tvSize transposeArrayLength]) - localMem - - prefixArrays <- - forM (zip byteOffsets types) $ \(off, ty) -> do - let off' = off `quot` primByteSize ty - sArray - "local_prefix_arr" - ty - (Shape [groupSize]) - $ ArrayIn localMem $ IxFun.iotaOffset off' [pe64 groupSize] - - warpscan <- sArrayInMem "warpscan" int8 (Shape [constant (warpSize :: Int64)]) localMem - warpExchanges <- - forM (zip warpByteOffsets types) $ \(off, ty) -> do - let off' = off `quot` primByteSize ty - sArray - "warp_exchange" - ty - (Shape [constant (warpSize :: Int64)]) - $ ArrayIn localMem $ IxFun.iotaOffset off' [warpSize] - - return (sharedId, transposedArrays, prefixArrays, sharedReadOffset, warpscan, warpExchanges) - --- | Compile 'SegScan' instance to host-level code with calls to a --- single-pass kernel. -compileSegScan :: - Pattern KernelsMem -> - SegLevel -> - SegSpace -> - SegBinOp KernelsMem -> - KernelBody KernelsMem -> - CallKernelGen () -compileSegScan pat lvl space scanOp kbody = do - let Pattern _ all_pes = pat - group_size = toInt64Exp <$> segGroupSize lvl - n = product $ map toInt64Exp $ segSpaceDims space - m :: Num a => a - m = 9 - num_groups = Count (n `divUp` (unCount group_size * m)) - num_threads = unCount num_groups * unCount group_size - (mapIdx, _) = head $ unSegSpace space - scanOpNe = segBinOpNeutral scanOp - tys = map (\(Prim pt) -> pt) $ lambdaReturnType $ segBinOpLambda scanOp - statusX, statusA, statusP :: Num a => a - statusX = 0 - statusA = 1 - statusP = 2 - makeStatusUsed flag used = tvExp flag .|. (tvExp used .<<. 2) - unmakeStatusUsed :: TV Int8 -> TV Int8 -> TV Int8 -> InKernelGen () - unmakeStatusUsed flagUsed flag used = do - used <-- tvExp flagUsed .>>. 2 - flag <-- tvExp flagUsed .&. 3 - - -- Allocate the shared memory for output component - numThreads <- dPrimV "numThreads" num_threads - numGroups <- dPrimV "numGroups" $ unCount num_groups - - globalId <- sStaticArray "id_counter" (Space "device") int32 $ Imp.ArrayZeros 1 - statusFlags <- sAllocArray "status_flags" int8 (Shape [tvSize numGroups]) (Space "device") - (aggregateArrays, incprefixArrays) <- - fmap unzip $ - forM tys $ \ty -> - (,) <$> sAllocArray "aggregates" ty (Shape [tvSize numGroups]) (Space "device") - <*> sAllocArray "incprefixes" ty (Shape [tvSize numGroups]) (Space "device") - - sReplicate statusFlags $ intConst Int8 statusX - - sKernelThread "segscan" num_groups group_size (segFlat space) $ do - constants <- kernelConstants <$> askEnv - - (sharedId, transposedArrays, prefixArrays, sharedReadOffset, warpscan, exchanges) <- - createLocalArrays (segGroupSize lvl) (intConst Int64 m) tys - - dynamicId <- dPrim "dynamic_id" int32 - sWhen (kernelLocalThreadId constants .==. 0) $ do - (globalIdMem, _, globalIdOff) <- fullyIndexArray globalId [0] - sOp $ - Imp.Atomic DefaultSpace $ - Imp.AtomicAdd - Int32 - (tvVar dynamicId) - globalIdMem - (Count $ unCount globalIdOff) - (untyped (1 :: Imp.TExp Int32)) - copyDWIMFix sharedId [0] (tvSize dynamicId) [] - - let localBarrier = Imp.Barrier Imp.FenceLocal - localFence = Imp.MemFence Imp.FenceLocal - globalFence = Imp.MemFence Imp.FenceGlobal - - sOp localBarrier - copyDWIMFix (tvVar dynamicId) [] (Var sharedId) [0] - sOp localBarrier - - blockOff <- - dPrimV "blockOff" $ - sExt64 (tvExp dynamicId) * m * kernelGroupSize constants - - privateArrays <- - forM tys $ \ty -> - sAllocArray - "private" - ty - (Shape [intConst Int64 m]) - (ScalarSpace [intConst Int64 m] ty) - - sComment "Load and map" $ - sFor "i" m $ \i -> do - -- The map's input index - dPrimV_ mapIdx $ - tvExp blockOff + sExt64 (kernelLocalThreadId constants) - + i * kernelGroupSize constants - -- Perform the map - let in_bounds = - compileStms mempty (kernelBodyStms kbody) $ do - let (all_scan_res, map_res) = splitAt (segBinOpResults [scanOp]) $ kernelBodyResult kbody - - -- Write map results to their global memory destinations - forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(dest, src) -> - copyDWIMFix (patElemName dest) [Imp.vi64 mapIdx] (kernelResultSubExp src) [] - - -- Write to-scan results to private memory. - forM_ (zip privateArrays $ map kernelResultSubExp all_scan_res) $ \(dest, src) -> - copyDWIMFix dest [i] src [] - - out_of_bounds = - forM_ (zip privateArrays scanOpNe) $ \(dest, ne) -> - copyDWIMFix dest [i] ne [] - - sIf (Imp.vi64 mapIdx .<. n) in_bounds out_of_bounds - - sComment "Transpose scan inputs" $ do - forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do - sOp localBarrier - sFor "i" m $ \i -> do - sharedIdx <- - dPrimVE "sharedIdx" $ - sExt64 (kernelLocalThreadId constants) - + i * kernelGroupSize constants - copyDWIMFix trans [sharedIdx] (Var priv) [i] - sOp localBarrier - sFor "i" m $ \i -> do - sharedIdx <- dPrimV "sharedIdx" $ kernelLocalThreadId constants * m + i - copyDWIMFix priv [sExt64 i] (Var trans) [sExt64 $ tvExp sharedIdx] - sOp localBarrier - - sComment "Per thread scan" $ - -- We don't need to touch the first element, so only m-1 - -- iterations here. - sFor "i" (m -1) $ \i -> do - let xs = map paramName $ xParams scanOp - ys = map paramName $ yParams scanOp - - forM_ (zip privateArrays $ zip3 xs ys tys) $ \(src, (x, y, ty)) -> do - dPrim_ x ty - dPrim_ y ty - copyDWIMFix x [] (Var src) [i] - copyDWIMFix y [] (Var src) [i + 1] - - compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scanOp) $ - forM_ (zip privateArrays $ bodyResult $ lambdaBody $ segBinOpLambda scanOp) $ \(dest, res) -> - copyDWIMFix dest [i + 1] res [] - - sComment "Publish results in shared memory" $ do - forM_ (zip prefixArrays privateArrays) $ \(dest, src) -> - copyDWIMFix dest [sExt64 $ kernelLocalThreadId constants] (Var src) [m - 1] - sOp localBarrier - - scanOp' <- renameLambda $ segBinOpLambda scanOp - - accs <- mapM (dPrim "acc") tys - sComment "Scan results (with warp scan)" $ do - groupScan - Nothing -- TODO - (tvExp numThreads) - (kernelGroupSize constants) - scanOp' - prefixArrays - - sOp localBarrier - - let firstThread acc prefixes = - copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelGroupSize constants) - 1] - notFirstThread acc prefixes = - copyDWIMFix (tvVar acc) [] (Var prefixes) [sExt64 (kernelLocalThreadId constants) - 1] - sIf - (kernelLocalThreadId constants .==. 0) - (zipWithM_ firstThread accs prefixArrays) - (zipWithM_ notFirstThread accs prefixArrays) - - sOp localBarrier - - prefixes <- forM (zip scanOpNe tys) $ \(ne, ty) -> - dPrimV "prefix" $ TPrimExp $ toExp' ty ne - sComment "Perform lookback" $ do - sWhen (tvExp dynamicId .==. 0 .&&. kernelLocalThreadId constants .==. 0) $ do - everythingVolatile $ - forM_ (zip incprefixArrays accs) $ \(incprefixArray, acc) -> - copyDWIMFix incprefixArray [tvExp dynamicId] (tvSize acc) [] - sOp globalFence - everythingVolatile $ - copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusP) [] - forM_ (zip scanOpNe accs) $ \(ne, acc) -> - copyDWIMFix (tvVar acc) [] ne [] - -- end sWhen - - let warpSize = kernelWaveSize constants - sWhen (bNot (tvExp dynamicId .==. 0) .&&. kernelLocalThreadId constants .<. warpSize) $ do - sWhen (kernelLocalThreadId constants .==. 0) $ do - everythingVolatile $ - forM_ (zip aggregateArrays accs) $ \(aggregateArray, acc) -> - copyDWIMFix aggregateArray [tvExp dynamicId] (tvSize acc) [] - sOp globalFence - everythingVolatile $ - copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusA) [] - copyDWIMFix warpscan [0] (Var statusFlags) [tvExp dynamicId - 1] - -- sWhen - sOp localFence - - status <- dPrim "status" int8 :: InKernelGen (TV Int8) - copyDWIMFix (tvVar status) [] (Var warpscan) [0] - - sIf - (tvExp status .==. statusP) - ( sWhen (kernelLocalThreadId constants .==. 0) $ - everythingVolatile $ - forM_ (zip prefixes incprefixArrays) $ \(prefix, incprefixArray) -> - copyDWIMFix (tvVar prefix) [] (Var incprefixArray) [tvExp dynamicId - 1] - ) - ( do - readOffset <- - dPrimV "readOffset" $ - sExt32 $ tvExp dynamicId - sExt64 (kernelWaveSize constants) - let loopStop = warpSize * (-1) - sWhile (tvExp readOffset .>. loopStop) $ do - readI <- dPrimV "read_i" $ tvExp readOffset + kernelLocalThreadId constants - aggrs <- forM (zip scanOpNe tys) $ \(ne, ty) -> - dPrimV "aggr" $ TPrimExp $ toExp' ty ne - flag <- dPrimV "flag" statusX - used <- dPrimV "used" (0 :: Imp.TExp Int8) - everythingVolatile $ - sWhen (tvExp readI .>=. 0) $ do - copyDWIMFix (tvVar flag) [] (Var statusFlags) [sExt64 $ tvExp readI] - sIf - (tvExp flag .==. statusP) - ( forM_ (zip incprefixArrays aggrs) $ \(incprefix, aggr) -> - copyDWIMFix (tvVar aggr) [] (Var incprefix) [sExt64 $ tvExp readI] - ) - ( sWhen (tvExp flag .==. statusA) $ do - forM_ (zip aggrs aggregateArrays) $ \(aggr, aggregate) -> - copyDWIMFix (tvVar aggr) [] (Var aggregate) [sExt64 $ tvExp readI] - used <-- (1 :: Imp.TExp Int8) - ) - -- end sIf - -- end sWhen - forM_ (zip exchanges aggrs) $ \(exchange, aggr) -> - copyDWIMFix exchange [sExt64 $ kernelLocalThreadId constants] (tvSize aggr) [] - tmp <- dPrimV "tmp" $ makeStatusUsed flag used - copyDWIMFix warpscan [sExt64 $ kernelLocalThreadId constants] (tvSize tmp) [] - sOp localFence - - (warpscanMem, warpscanSpace, warpscanOff) <- - fullyIndexArray warpscan [sExt64 warpSize - 1] - flag <-- TPrimExp (Imp.index warpscanMem warpscanOff int8 warpscanSpace Imp.Volatile) - sWhen (kernelLocalThreadId constants .==. 0) $ do - -- TODO: This is a single-threaded reduce - sIf - (bNot $ tvExp flag .==. statusP) - ( do - scanOp'' <- renameLambda scanOp' - let (agg1s, agg2s) = splitAt (length tys) $ map paramName $ lambdaParams scanOp'' - - forM_ (zip3 agg1s scanOpNe tys) $ \(agg1, ne, ty) -> - dPrimV_ agg1 $ TPrimExp $ toExp' ty ne - zipWithM_ dPrim_ agg2s tys - - flag1 <- dPrimV "flag1" statusX - flag2 <- dPrim "flag2" int8 - used1 <- dPrimV "used1" (0 :: Imp.TExp Int8) - used2 <- dPrim "used2" int8 - sFor "i" warpSize $ \i -> do - copyDWIMFix (tvVar flag2) [] (Var warpscan) [sExt64 i] - unmakeStatusUsed flag2 flag2 used2 - forM_ (zip agg2s exchanges) $ \(agg2, exchange) -> - copyDWIMFix agg2 [] (Var exchange) [sExt64 i] - sIf - (bNot $ tvExp flag2 .==. statusA) - ( do - flag1 <-- tvExp flag2 - used1 <-- tvExp used2 - forM_ (zip3 agg1s tys agg2s) $ \(agg1, ty, agg2) -> - agg1 <~~ toExp' ty (Var agg2) - ) - ( do - used1 <-- tvExp used1 + tvExp used2 - compileStms mempty (bodyStms $ lambdaBody scanOp'') $ - forM_ (zip3 agg1s tys $ bodyResult $ lambdaBody scanOp'') $ - \(agg1, ty, res) -> agg1 <~~ toExp' ty res - ) - flag <-- tvExp flag1 - used <-- tvExp used1 - forM_ (zip3 aggrs tys agg1s) $ \(aggr, ty, agg1) -> - tvVar aggr <~~ toExp' ty (Var agg1) - ) - -- else - ( forM_ (zip aggrs exchanges) $ \(aggr, exchange) -> - copyDWIMFix (tvVar aggr) [] (Var exchange) [sExt64 warpSize - 1] - ) - -- end sIf - sIf - (tvExp flag .==. statusP) - (readOffset <-- loopStop) - (readOffset <-- tvExp readOffset - zExt32 (tvExp used)) - copyDWIMFix sharedReadOffset [0] (tvSize readOffset) [] - scanOp''' <- renameLambda scanOp' - let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams scanOp''' - forM_ (zip xs aggrs) $ \(x, aggr) -> dPrimV_ x (tvExp aggr) - forM_ (zip ys prefixes) $ \(y, prefix) -> dPrimV_ y (tvExp prefix) - compileStms mempty (bodyStms $ lambdaBody scanOp''') $ - forM_ (zip3 prefixes tys $ bodyResult $ lambdaBody scanOp''') $ - \(prefix, ty, res) -> prefix <-- TPrimExp (toExp' ty res) - -- end sWhen - sOp localFence - copyDWIMFix (tvVar readOffset) [] (Var sharedReadOffset) [0] - ) - -- end sWhile - -- end sIf - sWhen (kernelLocalThreadId constants .==. 0) $ do - scanOp'''' <- renameLambda scanOp' - let xs = map paramName $ take (length tys) $ lambdaParams scanOp'''' - ys = map paramName $ drop (length tys) $ lambdaParams scanOp'''' - forM_ (zip xs prefixes) $ \(x, prefix) -> dPrimV_ x $ tvExp prefix - forM_ (zip ys accs) $ \(y, acc) -> dPrimV_ y $ tvExp acc - compileStms mempty (bodyStms $ lambdaBody scanOp'''') $ - everythingVolatile $ - forM_ (zip incprefixArrays $ bodyResult $ lambdaBody scanOp'''') $ - \(incprefixArray, res) -> copyDWIMFix incprefixArray [tvExp dynamicId] res [] - sOp globalFence - everythingVolatile $ copyDWIMFix statusFlags [tvExp dynamicId] (intConst Int8 statusP) [] - forM_ (zip exchanges prefixes) $ \(exchange, prefix) -> - copyDWIMFix exchange [0] (tvSize prefix) [] - forM_ (zip3 accs tys scanOpNe) $ \(acc, ty, ne) -> - tvVar acc <~~ toExp' ty ne - -- end sWhen - -- end sWhen - - sWhen (bNot $ tvExp dynamicId .==. 0) $ do - sOp localBarrier - forM_ (zip exchanges prefixes) $ \(exchange, prefix) -> - copyDWIMFix (tvVar prefix) [] (Var exchange) [0] - sOp localBarrier - -- end sWhen - -- end sComment - - scanOp''''' <- renameLambda scanOp' - scanOp'''''' <- renameLambda scanOp' - - sComment "Distribute results" $ do - let (xs, ys) = splitAt (length tys) $ map paramName $ lambdaParams scanOp''''' - (xs', ys') = splitAt (length tys) $ map paramName $ lambdaParams scanOp'''''' - - forM_ (zip4 (zip prefixes accs) (zip xs xs') (zip ys ys') tys) $ - \((prefix, acc), (x, x'), (y, y'), ty) -> do - dPrim_ x ty - dPrim_ y ty - dPrimV_ x' $ tvExp prefix - dPrimV_ y' $ tvExp acc - - compileStms mempty (bodyStms $ lambdaBody scanOp'''''') $ - forM_ (zip3 xs tys $ bodyResult $ lambdaBody scanOp'''''') $ - \(x, ty, res) -> x <~~ toExp' ty res - - sFor "i" m $ \i -> do - forM_ (zip privateArrays ys) $ \(src, y) -> - copyDWIMFix y [] (Var src) [i] - - compileStms mempty (bodyStms $ lambdaBody scanOp''''') $ - forM_ (zip privateArrays $ bodyResult $ lambdaBody scanOp''''') $ - \(dest, res) -> - copyDWIMFix dest [i] res [] - - sComment "Transpose scan output" $ do - forM_ (zip transposedArrays privateArrays) $ \(trans, priv) -> do - sOp localBarrier - sFor "i" m $ \i -> do - sharedIdx <- - dPrimV "sharedIdx" $ - sExt64 (kernelLocalThreadId constants * m) + i - copyDWIMFix trans [tvExp sharedIdx] (Var priv) [i] - sOp localBarrier - sFor "i" m $ \i -> do - sharedIdx <- - dPrimV "sharedIdx" $ - kernelLocalThreadId constants - + sExt32 (kernelGroupSize constants * i) - copyDWIMFix priv [i] (Var trans) [sExt64 $ tvExp sharedIdx] - sOp localBarrier - - sComment "Write block scan results to global memory" $ - forM_ (zip (map patElemName all_pes) privateArrays) $ \(dest, src) -> - sFor "i" m $ \i -> do - dPrimV_ mapIdx $ - tvExp blockOff + kernelGroupSize constants * i - + sExt64 (kernelLocalThreadId constants) - sWhen (Imp.vi64 mapIdx .<. n) $ - copyDWIMFix dest [Imp.vi64 mapIdx] (Var src) [i] - - sComment "If this is the last block, reset the dynamicId" $ - sWhen (tvExp dynamicId .==. unCount num_groups - 1) $ - copyDWIMFix globalId [0] (constant (0 :: Int32)) [] diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan/TwoPass.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegScan/TwoPass.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegScan/TwoPass.hs +++ /dev/null @@ -1,506 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} - --- | Code generation for segmented and non-segmented scans. Uses a --- fairly inefficient two-pass algorithm, but can handle anything. -module Futhark.CodeGen.ImpGen.Kernels.SegScan.TwoPass (compileSegScan) where - -import Control.Monad.Except -import Control.Monad.State -import Data.List (delete, find, foldl', zip4) -import Data.Maybe -import qualified Futhark.CodeGen.ImpCode.Kernels as Imp -import Futhark.CodeGen.ImpGen -import Futhark.CodeGen.ImpGen.Kernels.Base -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Transform.Rename -import Futhark.Util (takeLast) -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Prelude hiding (quot, rem) - --- Aggressively try to reuse memory for different SegBinOps, because --- we will run them sequentially after another. -makeLocalArrays :: - Count GroupSize SubExp -> - SubExp -> - [SegBinOp KernelsMem] -> - InKernelGen [[VName]] -makeLocalArrays (Count group_size) num_threads scans = do - (arrs, mems_and_sizes) <- runStateT (mapM onScan scans) mempty - let maxSize sizes = Imp.bytes $ foldl' sMax64 1 $ map Imp.unCount sizes - forM_ mems_and_sizes $ \(sizes, mem) -> - sAlloc_ mem (maxSize sizes) (Space "local") - return arrs - where - onScan (SegBinOp _ scan_op nes _) = do - let (scan_x_params, _scan_y_params) = - splitAt (length nes) $ lambdaParams scan_op - (arrs, used_mems) <- fmap unzip $ - forM scan_x_params $ \p -> - case paramDec p of - MemArray pt shape _ (ArrayIn mem _) -> do - let shape' = Shape [num_threads] <> shape - arr <- - lift $ - sArray "scan_arr" pt shape' $ - ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape' - return (arr, []) - _ -> do - let pt = elemType $ paramType p - shape = Shape [group_size] - (sizes, mem') <- getMem pt shape - arr <- lift $ sArrayInMem "scan_arr" pt shape mem' - return (arr, [(sizes, mem')]) - modify (<> concat used_mems) - return arrs - - getMem pt shape = do - let size = typeSize $ Array pt shape NoUniqueness - mems <- get - case (find ((size `elem`) . fst) mems, mems) of - (Just mem, _) -> do - modify $ delete mem - return mem - (Nothing, (size', mem) : mems') -> do - put mems' - return (size : size', mem) - (Nothing, []) -> do - mem <- lift $ sDeclareMem "scan_arr_mem" $ Space "local" - return ([size], mem) - -type CrossesSegment = Maybe (Imp.TExp Int64 -> Imp.TExp Int64 -> Imp.TExp Bool) - -localArrayIndex :: KernelConstants -> Type -> Imp.TExp Int64 -localArrayIndex constants t = - if primType t - then sExt64 (kernelLocalThreadId constants) - else sExt64 (kernelGlobalThreadId constants) - -barrierFor :: Lambda KernelsMem -> (Bool, Imp.Fence, InKernelGen ()) -barrierFor scan_op = (array_scan, fence, sOp $ Imp.Barrier fence) - where - array_scan = not $ all primType $ lambdaReturnType scan_op - fence - | array_scan = Imp.FenceGlobal - | otherwise = Imp.FenceLocal - -xParams, yParams :: SegBinOp KernelsMem -> [LParam KernelsMem] -xParams scan = - take (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) -yParams scan = - drop (length (segBinOpNeutral scan)) (lambdaParams (segBinOpLambda scan)) - -writeToScanValues :: - [VName] -> - ([PatElem KernelsMem], SegBinOp KernelsMem, [KernelResult]) -> - InKernelGen () -writeToScanValues gtids (pes, scan, scan_res) - | shapeRank (segBinOpShape scan) > 0 = - forM_ (zip pes scan_res) $ \(pe, res) -> - copyDWIMFix - (patElemName pe) - (map Imp.vi64 gtids) - (kernelResultSubExp res) - [] - | otherwise = - forM_ (zip (yParams scan) scan_res) $ \(p, res) -> - copyDWIMFix (paramName p) [] (kernelResultSubExp res) [] - -readToScanValues :: - [Imp.TExp Int64] -> - [PatElem KernelsMem] -> - SegBinOp KernelsMem -> - InKernelGen () -readToScanValues is pes scan - | shapeRank (segBinOpShape scan) > 0 = - forM_ (zip (yParams scan) pes) $ \(p, pe) -> - copyDWIMFix (paramName p) [] (Var (patElemName pe)) is - | otherwise = - return () - -readCarries :: - Imp.TExp Int64 -> - [Imp.TExp Int64] -> - [Imp.TExp Int64] -> - [PatElem KernelsMem] -> - SegBinOp KernelsMem -> - InKernelGen () -readCarries chunk_offset dims' vec_is pes scan - | shapeRank (segBinOpShape scan) > 0 = do - ltid <- kernelLocalThreadId . kernelConstants <$> askEnv - -- We may have to reload the carries from the output of the - -- previous chunk. - sIf - (chunk_offset .>. 0 .&&. ltid .==. 0) - ( do - let is = unflattenIndex dims' $ chunk_offset - 1 - forM_ (zip (xParams scan) pes) $ \(p, pe) -> - copyDWIMFix (paramName p) [] (Var (patElemName pe)) (is ++ vec_is) - ) - ( forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> - copyDWIMFix (paramName p) [] ne [] - ) - | otherwise = - return () - --- | Produce partially scanned intervals; one per workgroup. -scanStage1 :: - Pattern KernelsMem -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - SegSpace -> - [SegBinOp KernelsMem] -> - KernelBody KernelsMem -> - CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment) -scanStage1 (Pattern _ all_pes) num_groups group_size space scans kbody = do - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size' - - let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - let num_elements = product dims' - elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads) - elems_per_group = unCount group_size' * elems_per_thread - - let crossesSegment = - case reverse dims' of - segment_size : _ : _ -> Just $ \from to -> - (to - from) .>. (to `rem` segment_size) - _ -> Nothing - - sKernelThread "scan_stage1" num_groups' group_size' (segFlat space) $ do - constants <- kernelConstants <$> askEnv - all_local_arrs <- makeLocalArrays group_size (tvSize num_threads) scans - - -- The variables from scan_op will be used for the carry and such - -- in the big chunking loop. - forM_ scans $ \scan -> do - dScope Nothing $ scopeOfLParams $ lambdaParams $ segBinOpLambda scan - forM_ (zip (xParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> - copyDWIMFix (paramName p) [] ne [] - - sFor "j" elems_per_thread $ \j -> do - chunk_offset <- - dPrimV "chunk_offset" $ - sExt64 (kernelGroupSize constants) * j - + sExt64 (kernelGroupId constants) * elems_per_group - flat_idx <- - dPrimV "flat_idx" $ - tvExp chunk_offset + sExt64 (kernelLocalThreadId constants) - -- Construct segment indices. - zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx - - let per_scan_pes = segBinOpChunks scans all_pes - - in_bounds = - foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' - - when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do - let (all_scan_res, map_res) = - splitAt (segBinOpResults scans) $ kernelBodyResult kbody - per_scan_res = - segBinOpChunks scans all_scan_res - - sComment "write to-scan values to parameters" $ - mapM_ (writeToScanValues gtids) $ - zip3 per_scan_pes scans per_scan_res - - sComment "write mapped values results to global memory" $ - forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) -> - copyDWIMFix - (patElemName pe) - (map Imp.vi64 gtids) - (kernelResultSubExp se) - [] - - sComment "threads in bounds read input" $ - sWhen in_bounds when_in_bounds - - forM_ (zip3 per_scan_pes scans all_local_arrs) $ - \(pes, scan@(SegBinOp _ scan_op nes vec_shape), local_arrs) -> - sComment "do one intra-group scan operation" $ do - let rets = lambdaReturnType scan_op - scan_x_params = xParams scan - (array_scan, fence, barrier) = barrierFor scan_op - - when array_scan barrier - - sLoopNest vec_shape $ \vec_is -> do - sComment "maybe restore some to-scan values to parameters, or read neutral" $ - sIf - in_bounds - ( do - readToScanValues (map Imp.vi64 gtids ++ vec_is) pes scan - readCarries (tvExp chunk_offset) dims' vec_is pes scan - ) - ( forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) -> - copyDWIMFix (paramName p) [] ne [] - ) - - sComment "combine with carry and write to local memory" $ - compileStms mempty (bodyStms $ lambdaBody scan_op) $ - forM_ (zip3 rets local_arrs (bodyResult $ lambdaBody scan_op)) $ - \(t, arr, se) -> - copyDWIMFix arr [localArrayIndex constants t] se [] - - let crossesSegment' = do - f <- crossesSegment - Just $ \from to -> - let from' = sExt64 from + tvExp chunk_offset - to' = sExt64 to + tvExp chunk_offset - in f from' to' - - sOp $ Imp.ErrorSync fence - - -- We need to avoid parameter name clashes. - scan_op_renamed <- renameLambda scan_op - groupScan - crossesSegment' - (sExt64 $ tvExp num_threads) - (sExt64 $ kernelGroupSize constants) - scan_op_renamed - local_arrs - - sComment "threads in bounds write partial scan result" $ - sWhen in_bounds $ - forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> - copyDWIMFix - (patElemName pe) - (map Imp.vi64 gtids ++ vec_is) - (Var arr) - [localArrayIndex constants t] - - barrier - - let load_carry = - forM_ (zip local_arrs scan_x_params) $ \(arr, p) -> - copyDWIMFix - (paramName p) - [] - (Var arr) - [ if primType $ paramType p - then sExt64 (kernelGroupSize constants) - 1 - else - (sExt64 (kernelGroupId constants) + 1) - * sExt64 (kernelGroupSize constants) - 1 - ] - load_neutral = - forM_ (zip nes scan_x_params) $ \(ne, p) -> - copyDWIMFix (paramName p) [] ne [] - - sComment "first thread reads last element as carry-in for next iteration" $ do - crosses_segment <- dPrimVE "crosses_segment" $ - case crossesSegment of - Nothing -> false - Just f -> - f - ( tvExp chunk_offset - + sExt64 (kernelGroupSize constants) -1 - ) - ( tvExp chunk_offset - + sExt64 (kernelGroupSize constants) - ) - should_load_carry <- - dPrimVE "should_load_carry" $ - kernelLocalThreadId constants .==. 0 .&&. bNot crosses_segment - sWhen should_load_carry load_carry - when array_scan barrier - sUnless should_load_carry load_neutral - - barrier - - return (num_threads, elems_per_group, crossesSegment) - -scanStage2 :: - Pattern KernelsMem -> - TV Int32 -> - Imp.TExp Int64 -> - Count NumGroups SubExp -> - CrossesSegment -> - SegSpace -> - [SegBinOp KernelsMem] -> - CallKernelGen () -scanStage2 (Pattern _ all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do - let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - - -- Our group size is the number of groups for the stage 1 kernel. - let group_size = Count $ unCount num_groups - group_size' = fmap toInt64Exp group_size - - let crossesSegment' = do - f <- crossesSegment - Just $ \from to -> - f - ((sExt64 from + 1) * elems_per_group - 1) - ((sExt64 to + 1) * elems_per_group - 1) - - sKernelThread "scan_stage2" 1 group_size' (segFlat space) $ do - constants <- kernelConstants <$> askEnv - per_scan_local_arrs <- makeLocalArrays group_size (tvSize stage1_num_threads) scans - let per_scan_rets = map (lambdaReturnType . segBinOpLambda) scans - per_scan_pes = segBinOpChunks scans all_pes - - flat_idx <- - dPrimV "flat_idx" $ - (sExt64 (kernelLocalThreadId constants) + 1) * elems_per_group - 1 - -- Construct segment indices. - zipWithM_ dPrimV_ gtids $ unflattenIndex dims' $ tvExp flat_idx - - forM_ (zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $ - \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) -> - sLoopNest vec_shape $ \vec_is -> do - let glob_is = map Imp.vi64 gtids ++ vec_is - - in_bounds = - foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' - - when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) -> - copyDWIMFix - arr - [localArrayIndex constants t] - (Var $ patElemName pe) - glob_is - - when_out_of_bounds = forM_ (zip3 rets local_arrs nes) $ \(t, arr, ne) -> - copyDWIMFix arr [localArrayIndex constants t] ne [] - (_, _, barrier) = - barrierFor scan_op - - sComment "threads in bound read carries; others get neutral element" $ - sIf in_bounds when_in_bounds when_out_of_bounds - - barrier - - groupScan - crossesSegment' - (sExt64 $ tvExp stage1_num_threads) - (sExt64 $ kernelGroupSize constants) - scan_op - local_arrs - - sComment "threads in bounds write scanned carries" $ - sWhen in_bounds $ - forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) -> - copyDWIMFix - (patElemName pe) - glob_is - (Var arr) - [localArrayIndex constants t] - -scanStage3 :: - Pattern KernelsMem -> - Count NumGroups SubExp -> - Count GroupSize SubExp -> - Imp.TExp Int64 -> - CrossesSegment -> - SegSpace -> - [SegBinOp KernelsMem] -> - CallKernelGen () -scanStage3 (Pattern _ all_pes) num_groups group_size elems_per_group crossesSegment space scans = do - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - required_groups <- - dPrimVE "required_groups" $ - sExt32 $ product dims' `divUp` sExt64 (unCount group_size') - - sKernelThread "scan_stage3" num_groups' group_size' (segFlat space) $ - virtualiseGroups SegVirt required_groups $ \virt_group_id -> do - constants <- kernelConstants <$> askEnv - - -- Compute our logical index. - flat_idx <- - dPrimVE "flat_idx" $ - sExt64 virt_group_id * sExt64 (unCount group_size') - + sExt64 (kernelLocalThreadId constants) - zipWithM_ dPrimV_ gtids $ unflattenIndex dims' flat_idx - - -- Figure out which group this element was originally in. - orig_group <- dPrimV "orig_group" $ flat_idx `quot` elems_per_group - -- Then the index of the carry-in of the preceding group. - carry_in_flat_idx <- - dPrimV "carry_in_flat_idx" $ - tvExp orig_group * elems_per_group - 1 - -- Figure out the logical index of the carry-in. - let carry_in_idx = unflattenIndex dims' $ tvExp carry_in_flat_idx - - -- Apply the carry if we are not in the scan results for the first - -- group, and are not the last element in such a group (because - -- then the carry was updated in stage 2), and we are not crossing - -- a segment boundary. - let in_bounds = - foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims' - crosses_segment = - fromMaybe false $ - crossesSegment - <*> pure (tvExp carry_in_flat_idx) - <*> pure flat_idx - is_a_carry = flat_idx .==. (tvExp orig_group + 1) * elems_per_group - 1 - no_carry_in = tvExp orig_group .==. 0 .||. is_a_carry .||. crosses_segment - - let per_scan_pes = segBinOpChunks scans all_pes - sWhen in_bounds $ - sUnless no_carry_in $ - forM_ (zip per_scan_pes scans) $ - \(pes, SegBinOp _ scan_op nes vec_shape) -> do - dScope Nothing $ scopeOfLParams $ lambdaParams scan_op - let (scan_x_params, scan_y_params) = - splitAt (length nes) $ lambdaParams scan_op - - sLoopNest vec_shape $ \vec_is -> do - forM_ (zip scan_x_params pes) $ \(p, pe) -> - copyDWIMFix - (paramName p) - [] - (Var $ patElemName pe) - (carry_in_idx ++ vec_is) - - forM_ (zip scan_y_params pes) $ \(p, pe) -> - copyDWIMFix - (paramName p) - [] - (Var $ patElemName pe) - (map Imp.vi64 gtids ++ vec_is) - - compileBody' scan_x_params $ lambdaBody scan_op - - forM_ (zip scan_x_params pes) $ \(p, pe) -> - copyDWIMFix - (patElemName pe) - (map Imp.vi64 gtids ++ vec_is) - (Var $ paramName p) - [] - --- | Compile 'SegScan' instance to host-level code with calls to --- various kernels. -compileSegScan :: - Pattern KernelsMem -> - SegLevel -> - SegSpace -> - [SegBinOp KernelsMem] -> - KernelBody KernelsMem -> - CallKernelGen () -compileSegScan pat lvl space scans kbody = do - -- Since stage 2 involves a group size equal to the number of groups - -- used for stage 1, we have to cap this number to the maximum group - -- size. - stage1_max_num_groups <- dPrim "stage1_max_num_groups" int64 - sOp $ Imp.GetSizeMax (tvVar stage1_max_num_groups) SizeGroup - - stage1_num_groups <- - fmap (Imp.Count . tvSize) $ - dPrimV "stage1_num_groups" $ - sMin64 (tvExp stage1_max_num_groups) $ - toInt64Exp $ Imp.unCount $ segNumGroups lvl - - (stage1_num_threads, elems_per_group, crossesSegment) <- - scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody - - emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group - - scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans - scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs +++ /dev/null @@ -1,888 +0,0 @@ -{-# LANGUAGE QuasiQuotes #-} -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE TupleSections #-} - --- | This module defines a translation from imperative code with --- kernels to imperative code with OpenCL calls. -module Futhark.CodeGen.ImpGen.Kernels.ToOpenCL - ( kernelsToOpenCL, - kernelsToCUDA, - ) -where - -import Control.Monad.Identity -import Control.Monad.Reader -import Control.Monad.State -import Data.FileEmbed -import qualified Data.Map.Strict as M -import Data.Maybe -import qualified Data.Set as S -import qualified Futhark.CodeGen.Backends.GenericC as GC -import Futhark.CodeGen.Backends.SimpleRep -import Futhark.CodeGen.ImpCode.Kernels hiding (Program) -import qualified Futhark.CodeGen.ImpCode.Kernels as ImpKernels -import Futhark.CodeGen.ImpCode.OpenCL hiding (Program) -import qualified Futhark.CodeGen.ImpCode.OpenCL as ImpOpenCL -import Futhark.Error (compilerLimitationS) -import Futhark.IR.Prop (isBuiltInFunction) -import Futhark.MonadFreshNames -import Futhark.Util (zEncodeString) -import Futhark.Util.Pretty (prettyOneLine) -import qualified Language.C.Quote.CUDA as CUDAC -import qualified Language.C.Quote.OpenCL as C -import qualified Language.C.Syntax as C - -kernelsToCUDA, kernelsToOpenCL :: ImpKernels.Program -> ImpOpenCL.Program -kernelsToCUDA = translateKernels TargetCUDA -kernelsToOpenCL = translateKernels TargetOpenCL - --- | Translate a kernels-program to an OpenCL-program. -translateKernels :: - KernelTarget -> - ImpKernels.Program -> - ImpOpenCL.Program -translateKernels target prog = - let ( prog', - ToOpenCL kernels device_funs used_types sizes failures - ) = - (`runState` initialOpenCL) . (`runReaderT` defFuns prog) $ do - let ImpKernels.Definitions - (ImpKernels.Constants ps consts) - (ImpKernels.Functions funs) = prog - consts' <- traverse (onHostOp target) consts - funs' <- forM funs $ \(fname, fun) -> - (fname,) <$> traverse (onHostOp target) fun - - return $ - ImpOpenCL.Definitions - (ImpOpenCL.Constants ps consts') - (ImpOpenCL.Functions funs') - - (device_prototypes, device_defs) = unzip $ M.elems device_funs - kernels' = M.map fst kernels - opencl_code = openClCode $ map snd $ M.elems kernels - - opencl_prelude = - unlines - [ pretty $ genPrelude target used_types, - unlines $ map pretty device_prototypes, - unlines $ map pretty device_defs - ] - in ImpOpenCL.Program - opencl_code - opencl_prelude - kernels' - (S.toList used_types) - (cleanSizes sizes) - failures - prog' - where - genPrelude TargetOpenCL = genOpenClPrelude - genPrelude TargetCUDA = const genCUDAPrelude - --- | Due to simplifications after kernel extraction, some threshold --- parameters may contain KernelPaths that reference threshold --- parameters that no longer exist. We remove these here. -cleanSizes :: M.Map Name SizeClass -> M.Map Name SizeClass -cleanSizes m = M.map clean m - where - known = M.keys m - clean (SizeThreshold path def) = - SizeThreshold (filter ((`elem` known) . fst) path) def - clean s = s - -pointerQuals :: Monad m => String -> m [C.TypeQual] -pointerQuals "global" = return [C.ctyquals|__global|] -pointerQuals "local" = return [C.ctyquals|__local|] -pointerQuals "private" = return [C.ctyquals|__private|] -pointerQuals "constant" = return [C.ctyquals|__constant|] -pointerQuals "write_only" = return [C.ctyquals|__write_only|] -pointerQuals "read_only" = return [C.ctyquals|__read_only|] -pointerQuals "kernel" = return [C.ctyquals|__kernel|] -pointerQuals s = error $ "'" ++ s ++ "' is not an OpenCL kernel address space." - --- In-kernel name and per-workgroup size in bytes. -type LocalMemoryUse = (VName, Count Bytes Exp) - -data KernelState = KernelState - { kernelLocalMemory :: [LocalMemoryUse], - kernelFailures :: [FailureMsg], - kernelNextSync :: Int, - -- | Has a potential failure occurred sine the last - -- ErrorSync? - kernelSyncPending :: Bool, - kernelHasBarriers :: Bool - } - -newKernelState :: [FailureMsg] -> KernelState -newKernelState failures = KernelState mempty failures 0 False False - -errorLabel :: KernelState -> String -errorLabel = ("error_" ++) . show . kernelNextSync - -data ToOpenCL = ToOpenCL - { clKernels :: M.Map KernelName (KernelSafety, C.Func), - clDevFuns :: M.Map Name (C.Definition, C.Func), - clUsedTypes :: S.Set PrimType, - clSizes :: M.Map Name SizeClass, - clFailures :: [FailureMsg] - } - -initialOpenCL :: ToOpenCL -initialOpenCL = ToOpenCL mempty mempty mempty mempty mempty - -type AllFunctions = ImpKernels.Functions ImpKernels.HostOp - -lookupFunction :: Name -> AllFunctions -> Maybe ImpKernels.Function -lookupFunction fname (ImpKernels.Functions fs) = lookup fname fs - -type OnKernelM = ReaderT AllFunctions (State ToOpenCL) - -addSize :: Name -> SizeClass -> OnKernelM () -addSize key sclass = - modify $ \s -> s {clSizes = M.insert key sclass $ clSizes s} - -onHostOp :: KernelTarget -> HostOp -> OnKernelM OpenCL -onHostOp target (CallKernel k) = onKernel target k -onHostOp _ (ImpKernels.GetSize v key size_class) = do - addSize key size_class - return $ ImpOpenCL.GetSize v key -onHostOp _ (ImpKernels.CmpSizeLe v key size_class x) = do - addSize key size_class - return $ ImpOpenCL.CmpSizeLe v key x -onHostOp _ (ImpKernels.GetSizeMax v size_class) = - return $ ImpOpenCL.GetSizeMax v size_class - -genGPUCode :: - OpsMode -> - KernelCode -> - [FailureMsg] -> - GC.CompilerM KernelOp KernelState a -> - (a, GC.CompilerState KernelState) -genGPUCode mode body failures = - GC.runCompilerM - (inKernelOperations mode body) - blankNameSource - (newKernelState failures) - --- Compilation of a device function that is not not invoked from the --- host, but is invoked by (perhaps multiple) kernels. -generateDeviceFun :: Name -> ImpKernels.Function -> OnKernelM () -generateDeviceFun fname 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 - when (any memParam $ functionInput host_func) bad - - 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 - - modify $ \s -> - s - { clUsedTypes = typesInCode (functionBody device_func) <> clUsedTypes s, - clDevFuns = M.insert fname func $ clDevFuns s, - clFailures = kernelFailures kstate - } - - -- Important to do this after the 'modify' call, so we propagate the - -- right clFailures. - void $ ensureDeviceFuns $ functionBody device_func - where - toDevice :: HostOp -> KernelOp - toDevice _ = bad - - memParam MemParam {} = True - memParam ScalarParam {} = False - - bad = compilerLimitationS "Cannot generate GPU functions that use arrays." - --- Ensure that this device function is available, but don't regenerate --- it if it already exists. -ensureDeviceFun :: Name -> ImpKernels.Function -> OnKernelM () -ensureDeviceFun fname host_func = do - exists <- gets $ M.member fname . clDevFuns - unless exists $ generateDeviceFun fname host_func - -ensureDeviceFuns :: ImpKernels.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 func -> do - ensureDeviceFun fname func - return $ Just fname - Nothing -> return Nothing - -onKernel :: KernelTarget -> Kernel -> OnKernelM OpenCL -onKernel target kernel = do - called <- ensureDeviceFuns $ kernelBody kernel - - -- Crucial that this is done after 'ensureDeviceFuns', as the device - -- functions may themselves define failure points. - failures <- gets clFailures - - let (kernel_body, cstate) = - genGPUCode KernelMode (kernelBody kernel) failures $ - GC.blockScope $ GC.compileCode $ kernelBody kernel - kstate = GC.compUserState cstate - - use_params = mapMaybe useAsParam $ kernelUses kernel - - (local_memory_args, local_memory_params, local_memory_init) = - unzip3 $ - flip evalState (blankNameSource :: VNameSource) $ - mapM (prepareLocalMemory target) $ kernelLocalMemory kstate - - -- CUDA has very strict restrictions on the number of blocks - -- permitted along the 'y' and 'z' dimensions of the grid - -- (1<<16). To work around this, we are going to dynamically - -- permute the block dimensions to move the largest one to the - -- 'x' dimension, which has a higher limit (1<<31). This means - -- we need to extend the kernel with extra parameters that - -- contain information about this permutation, but we only do - -- this for multidimensional kernels (at the time of this - -- writing, only transposes). The corresponding arguments are - -- added automatically in CCUDA.hs. - (perm_params, block_dim_init) = - case (target, num_groups) of - (TargetCUDA, [_, _, _]) -> - ( [ [C.cparam|const int block_dim0|], - [C.cparam|const int block_dim1|], - [C.cparam|const int block_dim2|] - ], - mempty - ) - _ -> - ( mempty, - [ [C.citem|const int block_dim0 = 0;|], - [C.citem|const int block_dim1 = 1;|], - [C.citem|const int block_dim2 = 2;|] - ] - ) - - (const_defs, const_undefs) = unzip $ mapMaybe constDef $ kernelUses kernel - - let (safety, error_init) - -- We conservatively assume that any called function can fail. - | not $ null called = - (SafetyFull, []) - | length (kernelFailures kstate) == length failures = - if kernelFailureTolerant kernel - then (SafetyNone, []) - else -- No possible failures in this kernel, so if we make - -- it past an initial check, then we are good to go. - - ( SafetyCheap, - [C.citems|if (*global_failure >= 0) { return; }|] - ) - | otherwise = - if not (kernelHasBarriers kstate) - then - ( SafetyFull, - [C.citems|if (*global_failure >= 0) { return; }|] - ) - else - ( SafetyFull, - [C.citems| - volatile __local bool local_failure; - if (failure_is_an_option) { - int failed = *global_failure >= 0; - if (failed) { - return; - } - } - // All threads write this value - it looks like CUDA has a compiler bug otherwise. - local_failure = false; - barrier(CLK_LOCAL_MEM_FENCE); - |] - ) - - failure_params = - [ [C.cparam|__global int *global_failure|], - [C.cparam|int failure_is_an_option|], - [C.cparam|__global typename int64_t *global_failure_args|] - ] - - params = - perm_params - ++ take (numFailureParams safety) failure_params - ++ catMaybes local_memory_params - ++ use_params - - kernel_fun = - [C.cfun|__kernel void $id:name ($params:params) { - $items:const_defs - $items:block_dim_init - $items:local_memory_init - $items:error_init - $items:kernel_body - - $id:(errorLabel kstate): return; - - $items:const_undefs - }|] - modify $ \s -> - s - { clKernels = M.insert name (safety, kernel_fun) $ clKernels s, - clUsedTypes = typesInKernel kernel <> clUsedTypes s, - clFailures = kernelFailures kstate - } - - -- The argument corresponding to the global_failure parameters is - -- added automatically later. - let args = - catMaybes local_memory_args - ++ kernelArgs kernel - - return $ LaunchKernel safety name args num_groups group_size - where - name = kernelName kernel - num_groups = kernelNumGroups kernel - group_size = kernelGroupSize kernel - - prepareLocalMemory TargetOpenCL (mem, size) = do - mem_aligned <- newVName $ baseString mem ++ "_aligned" - return - ( Just $ SharedMemoryKArg size, - Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|], - [C.citem|__local volatile char* restrict $id:mem = (__local volatile char*)$id:mem_aligned;|] - ) - prepareLocalMemory TargetCUDA (mem, size) = do - param <- newVName $ baseString mem ++ "_offset" - return - ( Just $ SharedMemoryKArg size, - Just [C.cparam|uint $id:param|], - [C.citem|volatile char *$id:mem = &shared_mem[$id:param];|] - ) - -useAsParam :: KernelUse -> Maybe C.Param -useAsParam (ScalarUse name bt) = - let ctp = case bt of - -- OpenCL does not permit bool as a kernel parameter type. - Bool -> [C.cty|unsigned char|] - Unit -> [C.cty|unsigned char|] - _ -> GC.primTypeToCType bt - in Just [C.cparam|$ty:ctp $id:name|] -useAsParam (MemoryUse name) = - Just [C.cparam|__global unsigned char *$id:name|] -useAsParam ConstUse {} = - Nothing - --- Constants are #defined as macros. Since a constant name in one --- kernel might potentially (although unlikely) also be used for --- something else in another kernel, we #undef them after the kernel. -constDef :: KernelUse -> Maybe (C.BlockItem, C.BlockItem) -constDef (ConstUse v e) = - Just - ( [C.citem|$escstm:def|], - [C.citem|$escstm:undef|] - ) - where - e' = compilePrimExp e - def = "#define " ++ pretty (C.toIdent v mempty) ++ " (" ++ prettyOneLine e' ++ ")" - undef = "#undef " ++ pretty (C.toIdent v mempty) -constDef _ = Nothing - -openClCode :: [C.Func] -> String -openClCode kernels = - pretty [C.cunit|$edecls:funcs|] - where - funcs = - [ [C.cedecl|$func:kernel_func|] - | kernel_func <- kernels - ] - -atomicsDefs :: String -atomicsDefs = $(embedStringFile "rts/c/atomics.h") - -genOpenClPrelude :: S.Set PrimType -> [C.Definition] -genOpenClPrelude ts = - -- Clang-based OpenCL implementations need this for 'static' to work. - [ [C.cedecl|$esc:("#ifdef cl_clang_storage_class_specifiers")|], - [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_clang_storage_class_specifiers : enable")|], - [C.cedecl|$esc:("#endif")|], - [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_byte_addressable_store : enable")|] - ] - ++ concat - [ [C.cunit|$esc:("#pragma OPENCL EXTENSION cl_khr_fp64 : enable") - $esc:("#define FUTHARK_F64_ENABLED")|] - | uses_float64 - ] - ++ [C.cunit| -/* Some OpenCL programs dislike empty progams, or programs with no kernels. - * Declare a dummy kernel to ensure they remain our friends. */ -__kernel void dummy_kernel(__global unsigned char *dummy, int n) -{ - const int thread_gid = get_global_id(0); - if (thread_gid >= n) return; -} - -$esc:("#pragma OPENCL EXTENSION cl_khr_int64_base_atomics : enable") -$esc:("#pragma OPENCL EXTENSION cl_khr_int64_extended_atomics : enable") - -typedef char int8_t; -typedef short int16_t; -typedef int int32_t; -typedef long int64_t; - -typedef uchar uint8_t; -typedef ushort uint16_t; -typedef uint uint32_t; -typedef ulong uint64_t; - -// NVIDIAs OpenCL does not create device-wide memory fences (see #734), so we -// use inline assembly if we detect we are on an NVIDIA GPU. -$esc:("#ifdef cl_nv_pragma_unroll") -static inline void mem_fence_global() { - asm("membar.gl;"); -} -$esc:("#else") -static inline void mem_fence_global() { - mem_fence(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE); -} -$esc:("#endif") -static inline void mem_fence_local() { - mem_fence(CLK_LOCAL_MEM_FENCE); -} -|] - ++ cIntOps - ++ cFloat32Ops - ++ cFloat32Funs - ++ (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else []) - ++ [[C.cedecl|$esc:atomicsDefs|]] - where - uses_float64 = FloatType Float64 `S.member` ts - -genCUDAPrelude :: [C.Definition] -genCUDAPrelude = - cudafy ++ ops - where - ops = - cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ cFloat64Ops - ++ cFloat64Funs - ++ cFloatConvOps - ++ [[C.cedecl|$esc:atomicsDefs|]] - cudafy = - [CUDAC.cunit| -$esc:("#define FUTHARK_CUDA") -$esc:("#define FUTHARK_F64_ENABLED") - -typedef char int8_t; -typedef short int16_t; -typedef int int32_t; -typedef long long int64_t; -typedef unsigned char uint8_t; -typedef unsigned short uint16_t; -typedef unsigned int uint32_t; -typedef unsigned long long uint64_t; -typedef uint8_t uchar; -typedef uint16_t ushort; -typedef uint32_t uint; -typedef uint64_t ulong; -$esc:("#define __kernel extern \"C\" __global__ __launch_bounds__(MAX_THREADS_PER_BLOCK)") -$esc:("#define __global") -$esc:("#define __local") -$esc:("#define __private") -$esc:("#define __constant") -$esc:("#define __write_only") -$esc:("#define __read_only") - -static inline int get_group_id_fn(int block_dim0, int block_dim1, int block_dim2, int d) -{ - switch (d) { - case 0: d = block_dim0; break; - case 1: d = block_dim1; break; - case 2: d = block_dim2; break; - } - switch (d) { - case 0: return blockIdx.x; - case 1: return blockIdx.y; - case 2: return blockIdx.z; - default: return 0; - } -} -$esc:("#define get_group_id(d) get_group_id_fn(block_dim0, block_dim1, block_dim2, d)") - -static inline int get_num_groups_fn(int block_dim0, int block_dim1, int block_dim2, int d) -{ - switch (d) { - case 0: d = block_dim0; break; - case 1: d = block_dim1; break; - case 2: d = block_dim2; break; - } - switch(d) { - case 0: return gridDim.x; - case 1: return gridDim.y; - case 2: return gridDim.z; - default: return 0; - } -} -$esc:("#define get_num_groups(d) get_num_groups_fn(block_dim0, block_dim1, block_dim2, d)") - -static inline int get_local_id(int d) -{ - switch (d) { - case 0: return threadIdx.x; - case 1: return threadIdx.y; - case 2: return threadIdx.z; - default: return 0; - } -} - -static inline int get_local_size(int d) -{ - switch (d) { - case 0: return blockDim.x; - case 1: return blockDim.y; - case 2: return blockDim.z; - default: return 0; - } -} - -static inline int get_global_id_fn(int block_dim0, int block_dim1, int block_dim2, int d) -{ - return get_group_id(d) * get_local_size(d) + get_local_id(d); -} -$esc:("#define get_global_id(d) get_global_id_fn(block_dim0, block_dim1, block_dim2, d)") - -static inline int get_global_size(int block_dim0, int block_dim1, int block_dim2, int d) -{ - return get_num_groups(d) * get_local_size(d); -} - -$esc:("#define CLK_LOCAL_MEM_FENCE 1") -$esc:("#define CLK_GLOBAL_MEM_FENCE 2") -static inline void barrier(int x) -{ - __syncthreads(); -} -static inline void mem_fence_local() { - __threadfence_block(); -} -static inline void mem_fence_global() { - __threadfence(); -} - -$esc:("#define NAN (0.0/0.0)") -$esc:("#define INFINITY (1.0/0.0)") -extern volatile __shared__ char shared_mem[]; -|] - -compilePrimExp :: PrimExp KernelConst -> C.Exp -compilePrimExp e = runIdentity $ GC.compilePrimExp compileKernelConst e - where - compileKernelConst (SizeConst key) = - return [C.cexp|$id:(zEncodeString (pretty key))|] - -kernelArgs :: Kernel -> [KernelArg] -kernelArgs = mapMaybe useToArg . kernelUses - where - useToArg (MemoryUse mem) = Just $ MemKArg mem - useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt - useToArg ConstUse {} = Nothing - -nextErrorLabel :: GC.CompilerM KernelOp KernelState String -nextErrorLabel = - errorLabel <$> GC.getUserState - -incErrorLabel :: GC.CompilerM KernelOp KernelState () -incErrorLabel = - GC.modifyUserState $ \s -> s {kernelNextSync = kernelNextSync s + 1} - -pendingError :: Bool -> GC.CompilerM KernelOp KernelState () -pendingError b = - GC.modifyUserState $ \s -> s {kernelSyncPending = b} - -hasCommunication :: ImpKernels.KernelCode -> Bool -hasCommunication = any communicates - where - communicates ErrorSync {} = True - communicates Barrier {} = True - communicates _ = False - --- Whether we are generating code for a kernel or a device function. --- This has minor effects, such as exactly how failures are --- propagated. -data OpsMode = KernelMode | FunMode deriving (Eq) - -inKernelOperations :: - OpsMode -> - ImpKernels.KernelCode -> - GC.Operations KernelOp KernelState -inKernelOperations mode body = - GC.Operations - { GC.opsCompiler = kernelOps, - GC.opsMemoryType = kernelMemoryType, - GC.opsWriteScalar = kernelWriteScalar, - GC.opsReadScalar = kernelReadScalar, - GC.opsAllocate = cannotAllocate, - GC.opsDeallocate = cannotDeallocate, - GC.opsCopy = copyInKernel, - GC.opsStaticArray = noStaticArrays, - GC.opsFatMemory = False, - GC.opsError = errorInKernel, - GC.opsCall = callInKernel, - GC.opsCritical = mempty - } - where - has_communication = hasCommunication body - - fence FenceLocal = [C.cexp|CLK_LOCAL_MEM_FENCE|] - fence FenceGlobal = [C.cexp|CLK_GLOBAL_MEM_FENCE | CLK_LOCAL_MEM_FENCE|] - - kernelOps :: GC.OpCompiler KernelOp KernelState - kernelOps (GetGroupId v i) = - GC.stm [C.cstm|$id:v = get_group_id($int:i);|] - kernelOps (GetLocalId v i) = - GC.stm [C.cstm|$id:v = get_local_id($int:i);|] - kernelOps (GetLocalSize v i) = - GC.stm [C.cstm|$id:v = get_local_size($int:i);|] - kernelOps (GetGlobalId v i) = - GC.stm [C.cstm|$id:v = get_global_id($int:i);|] - kernelOps (GetGlobalSize v i) = - GC.stm [C.cstm|$id:v = get_global_size($int:i);|] - kernelOps (GetLockstepWidth v) = - GC.stm [C.cstm|$id:v = LOCKSTEP_WIDTH;|] - kernelOps (Barrier f) = do - GC.stm [C.cstm|barrier($exp:(fence f));|] - GC.modifyUserState $ \s -> s {kernelHasBarriers = True} - kernelOps (MemFence FenceLocal) = - GC.stm [C.cstm|mem_fence_local();|] - kernelOps (MemFence FenceGlobal) = - GC.stm [C.cstm|mem_fence_global();|] - kernelOps (LocalAlloc name size) = do - name' <- newVName $ pretty name ++ "_backing" - GC.modifyUserState $ \s -> - s {kernelLocalMemory = (name', fmap untyped size) : kernelLocalMemory s} - GC.stm [C.cstm|$id:name = (__local char*) $id:name';|] - kernelOps (ErrorSync f) = do - label <- nextErrorLabel - pending <- kernelSyncPending <$> GC.getUserState - when pending $ do - pendingError False - GC.stm [C.cstm|$id:label: barrier($exp:(fence f));|] - GC.stm [C.cstm|if (local_failure) { return; }|] - GC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|] -- intentional - GC.modifyUserState $ \s -> s {kernelHasBarriers = True} - incErrorLabel - kernelOps (Atomic space aop) = atomicOps space aop - - atomicCast s t = do - let volatile = [C.ctyquals|volatile|] - quals <- case s of - Space sid -> pointerQuals sid - _ -> pointerQuals "global" - return [C.cty|$tyquals:(volatile++quals) $ty:t|] - - atomicSpace (Space sid) = sid - atomicSpace _ = "global" - - doAtomic s t old arr ind val op ty = do - ind' <- GC.compileExp $ untyped $ unCount ind - val' <- GC.compileExp val - cast <- atomicCast s ty - GC.stm [C.cstm|$id:old = $id:op'(&(($ty:cast *)$id:arr)[$exp:ind'], ($ty:ty) $exp:val');|] - where - op' = op ++ "_" ++ pretty t ++ "_" ++ atomicSpace s - - doAtomicCmpXchg s t old arr ind cmp val ty = do - ind' <- GC.compileExp $ untyped $ unCount ind - cmp' <- GC.compileExp cmp - val' <- GC.compileExp val - cast <- atomicCast s ty - GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|] - where - op = "atomic_cmpxchg_" ++ pretty t ++ "_" ++ atomicSpace s - doAtomicXchg s t old arr ind val ty = do - cast <- atomicCast s ty - ind' <- GC.compileExp $ untyped $ unCount ind - val' <- GC.compileExp val - GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|] - where - op = "atomic_chg_" ++ pretty t ++ "_" ++ atomicSpace s - -- First the 64-bit operations. - atomicOps s (AtomicAdd Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_add" [C.cty|typename int64_t|] - atomicOps s (AtomicFAdd Float64 old arr ind val) = - doAtomic s Float64 old arr ind val "atomic_fadd" [C.cty|double|] - atomicOps s (AtomicSMax Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_smax" [C.cty|typename int64_t|] - atomicOps s (AtomicSMin Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_smin" [C.cty|typename int64_t|] - atomicOps s (AtomicUMax Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_umax" [C.cty|unsigned int64_t|] - atomicOps s (AtomicUMin Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_umin" [C.cty|unsigned int64_t|] - atomicOps s (AtomicAnd Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_and" [C.cty|typename int64_t|] - atomicOps s (AtomicOr Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_or" [C.cty|typename int64_t|] - atomicOps s (AtomicXor Int64 old arr ind val) = - doAtomic s Int64 old arr ind val "atomic_xor" [C.cty|typename int64_t|] - atomicOps s (AtomicCmpXchg (IntType Int64) old arr ind cmp val) = - doAtomicCmpXchg s (IntType Int64) old arr ind cmp val [C.cty|typename int64_t|] - atomicOps s (AtomicXchg (IntType Int64) old arr ind val) = - doAtomicXchg s (IntType Int64) old arr ind val [C.cty|typename int64_t|] - -- - atomicOps s (AtomicAdd t old arr ind val) = - doAtomic s t old arr ind val "atomic_add" [C.cty|int|] - atomicOps s (AtomicFAdd Float32 old arr ind val) = - doAtomic s Float32 old arr ind val "atomic_fadd" [C.cty|float|] - atomicOps s (AtomicSMax t old arr ind val) = - doAtomic s t old arr ind val "atomic_smax" [C.cty|int|] - atomicOps s (AtomicSMin t old arr ind val) = - doAtomic s t old arr ind val "atomic_smin" [C.cty|int|] - atomicOps s (AtomicUMax t old arr ind val) = - doAtomic s t old arr ind val "atomic_umax" [C.cty|unsigned int|] - atomicOps s (AtomicUMin t old arr ind val) = - doAtomic s t old arr ind val "atomic_umin" [C.cty|unsigned int|] - atomicOps s (AtomicAnd t old arr ind val) = - doAtomic s t old arr ind val "atomic_and" [C.cty|int|] - atomicOps s (AtomicOr t old arr ind val) = - doAtomic s t old arr ind val "atomic_or" [C.cty|int|] - atomicOps s (AtomicXor t old arr ind val) = - doAtomic s t old arr ind val "atomic_xor" [C.cty|int|] - atomicOps s (AtomicCmpXchg t old arr ind cmp val) = - doAtomicCmpXchg s t old arr ind cmp val [C.cty|int|] - atomicOps s (AtomicXchg t old arr ind val) = - doAtomicXchg s t old arr ind val [C.cty|int|] - - cannotAllocate :: GC.Allocate KernelOp KernelState - cannotAllocate _ = - error "Cannot allocate memory in kernel" - - cannotDeallocate :: GC.Deallocate KernelOp KernelState - cannotDeallocate _ _ = - error "Cannot deallocate memory in kernel" - - copyInKernel :: GC.Copy KernelOp KernelState - copyInKernel _ _ _ _ _ _ _ = - error "Cannot bulk copy in kernel." - - noStaticArrays :: GC.StaticArray KernelOp KernelState - noStaticArrays _ _ _ _ = - error "Cannot create static array in kernel." - - kernelMemoryType space = do - quals <- pointerQuals space - return [C.cty|$tyquals:quals $ty:defaultMemBlockType|] - - kernelWriteScalar = - GC.writeScalarPointerWithQuals pointerQuals - - kernelReadScalar = - GC.readScalarPointerWithQuals pointerQuals - - whatNext = do - label <- nextErrorLabel - pendingError True - return $ - if has_communication - then [C.citems|local_failure = true; goto $id:label;|] - else - if mode == FunMode - then [C.citems|return 1;|] - else [C.citems|return;|] - - callInKernel dests fname args - | isBuiltInFunction fname = - GC.opsCall GC.defaultOperations dests fname args - | otherwise = 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 - - GC.item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:what_next; }|] - - errorInKernel msg@(ErrorMsg parts) backtrace = do - n <- length . kernelFailures <$> GC.getUserState - GC.modifyUserState $ \s -> - s {kernelFailures = kernelFailures s ++ [FailureMsg msg backtrace]} - let setArgs _ [] = return [] - setArgs i (ErrorString {} : parts') = setArgs i parts' - setArgs i (ErrorInt32 x : parts') = do - x' <- GC.compileExp x - stms <- setArgs (i + 1) parts' - return $ [C.cstm|global_failure_args[$int:i] = (typename int64_t)$exp:x';|] : stms - setArgs i (ErrorInt64 x : parts') = do - x' <- GC.compileExp x - stms <- setArgs (i + 1) parts' - return $ [C.cstm|global_failure_args[$int:i] = $exp:x';|] : stms - argstms <- setArgs (0 :: Int) parts - - what_next <- whatNext - - GC.stm - [C.cstm|{ if (atomic_cmpxchg_i32_global(global_failure, -1, $int:n) == -1) - { $stms:argstms; } - $items:what_next - }|] - ---- Checking requirements - -typesInKernel :: Kernel -> S.Set PrimType -typesInKernel kernel = typesInCode $ kernelBody kernel - -typesInCode :: ImpKernels.KernelCode -> S.Set PrimType -typesInCode Skip = mempty -typesInCode (c1 :>>: c2) = typesInCode c1 <> typesInCode c2 -typesInCode (For _ e c) = typesInExp e <> typesInCode c -typesInCode (While (TPrimExp e) c) = typesInExp e <> typesInCode c -typesInCode DeclareMem {} = mempty -typesInCode (DeclareScalar _ _ t) = S.singleton t -typesInCode (DeclareArray _ _ t _) = S.singleton t -typesInCode (Allocate _ (Count (TPrimExp e)) _) = typesInExp e -typesInCode Free {} = mempty -typesInCode - ( Copy - _ - (Count (TPrimExp e1)) - _ - _ - (Count (TPrimExp e2)) - _ - (Count (TPrimExp e3)) - ) = - typesInExp e1 <> typesInExp e2 <> typesInExp e3 -typesInCode (Write _ (Count (TPrimExp e1)) t _ _ e2) = - typesInExp e1 <> S.singleton t <> typesInExp e2 -typesInCode (SetScalar _ e) = typesInExp e -typesInCode SetMem {} = mempty -typesInCode (Call _ _ es) = mconcat $ map typesInArg es - where - typesInArg MemArg {} = mempty - typesInArg (ExpArg e) = typesInExp e -typesInCode (If (TPrimExp e) c1 c2) = - typesInExp e <> typesInCode c1 <> typesInCode c2 -typesInCode (Assert e _ _) = typesInExp e -typesInCode (Comment _ c) = typesInCode c -typesInCode (DebugPrint _ v) = maybe mempty typesInExp v -typesInCode Op {} = mempty - -typesInExp :: Exp -> S.Set PrimType -typesInExp (ValueExp v) = S.singleton $ primValueType v -typesInExp (BinOpExp _ e1 e2) = typesInExp e1 <> typesInExp e2 -typesInExp (CmpOpExp _ e1 e2) = typesInExp e1 <> typesInExp e2 -typesInExp (ConvOpExp op e) = S.fromList [from, to] <> typesInExp e - where - (from, to) = convOpType op -typesInExp (UnOpExp _ e) = typesInExp e -typesInExp (FunExp _ args t) = S.singleton t <> mconcat (map typesInExp args) -typesInExp (LeafExp (Index _ (Count (TPrimExp e)) t _ _) _) = S.singleton t <> typesInExp e -typesInExp (LeafExp ScalarVar {} _) = mempty diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs deleted file mode 100644 --- a/src/Futhark/CodeGen/ImpGen/Kernels/Transpose.hs +++ /dev/null @@ -1,386 +0,0 @@ --- | Carefully optimised implementations of GPU transpositions. --- Written in ImpCode so we can compile it to both CUDA and OpenCL. -module Futhark.CodeGen.ImpGen.Kernels.Transpose - ( TransposeType (..), - TransposeArgs, - mapTransposeKernel, - ) -where - -import Futhark.CodeGen.ImpCode.Kernels -import Futhark.IR.Prop.Types -import Futhark.Util.IntegralExp (divUp, quot, rem) -import Prelude hiding (quot, rem) - --- | Which form of transposition to generate code for. -data TransposeType - = TransposeNormal - | TransposeLowWidth - | TransposeLowHeight - | -- | For small arrays that do not - -- benefit from coalescing. - TransposeSmall - deriving (Eq, Ord, Show) - --- | The types of the arguments accepted by a transposition function. -type TransposeArgs = - ( VName, - TExp Int32, - VName, - TExp Int32, - TExp Int32, - TExp Int32, - TExp Int32, - TExp Int32, - TExp Int32, - VName - ) - -elemsPerThread :: TExp Int32 -elemsPerThread = 4 - -mapTranspose :: TExp Int32 -> TransposeArgs -> PrimType -> TransposeType -> KernelCode -mapTranspose block_dim args t kind = - case kind of - TransposeSmall -> - mconcat - [ get_ids, - dec our_array_offset $ vi32 get_global_id_0 `quot` (height * width) * (height * width), - dec x_index $ (vi32 get_global_id_0 `rem` (height * width)) `quot` height, - dec y_index $ vi32 get_global_id_0 `rem` height, - dec odata_offset $ - (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset, - dec idata_offset $ - (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset, - dec index_in $ vi32 y_index * width + vi32 x_index, - dec index_out $ vi32 x_index * height + vi32 y_index, - when - (vi32 get_global_id_0 .<. width * height * num_arrays) - ( Write odata (elements $ sExt64 $ vi32 odata_offset + vi32 index_out) t (Space "global") Nonvolatile $ - index idata (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) t (Space "global") Nonvolatile - ) - ] - TransposeLowWidth -> - mkTranspose $ - lowDimBody - (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_0 `quot` muly)) - ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_1 - + (vi32 get_local_id_0 `rem` muly) * block_dim - ) - ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_0 - + (vi32 get_local_id_1 `rem` muly) * block_dim - ) - (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_1 `quot` muly)) - TransposeLowHeight -> - mkTranspose $ - lowDimBody - ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_0 - + (vi32 get_local_id_1 `rem` mulx) * block_dim - ) - (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_1 `quot` mulx)) - (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_0 `quot` mulx)) - ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_1 - + (vi32 get_local_id_0 `rem` mulx) * block_dim - ) - TransposeNormal -> - mkTranspose $ - mconcat - [ dec x_index $ vi32 get_global_id_0, - dec y_index $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_1, - when (vi32 x_index .<. width) $ - For j (untyped elemsPerThread) $ - let i = vi32 j * (tile_dim `quot` elemsPerThread) - in mconcat - [ dec index_in $ (vi32 y_index + i) * width + vi32 x_index, - when (vi32 y_index + i .<. height) $ - Write - block - ( elements $ - sExt64 $ - (vi32 get_local_id_1 + i) * (tile_dim + 1) - + vi32 get_local_id_0 - ) - t - (Space "local") - Nonvolatile - $ index - idata - (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) - t - (Space "global") - Nonvolatile - ], - Op $ Barrier FenceLocal, - SetScalar x_index $ untyped $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_0, - SetScalar y_index $ untyped $ vi32 get_group_id_0 * tile_dim + vi32 get_local_id_1, - when (vi32 x_index .<. height) $ - For j (untyped elemsPerThread) $ - let i = vi32 j * (tile_dim `quot` elemsPerThread) - in mconcat - [ dec index_out $ (vi32 y_index + i) * height + vi32 x_index, - when (vi32 y_index + i .<. width) $ - Write - odata - (elements $ sExt64 $ vi32 odata_offset + vi32 index_out) - t - (Space "global") - Nonvolatile - $ index - block - ( elements $ - sExt64 $ - vi32 get_local_id_0 * (tile_dim + 1) + vi32 get_local_id_1 + i - ) - t - (Space "local") - Nonvolatile - ] - ] - where - dec v (TPrimExp e) = - DeclareScalar v Nonvolatile (primExpType e) <> SetScalar v e - tile_dim = 2 * block_dim - - when a b = If a b mempty - - ( odata, - basic_odata_offset, - idata, - basic_idata_offset, - width, - height, - mulx, - muly, - num_arrays, - block - ) = args - - -- Be extremely careful when editing this list to ensure that - -- the names match up. Also, be careful that the tags on - -- these names do not conflict with the tags of the - -- surrounding code. We accomplish the latter by using very - -- low tags (normal variables start at least in the low - -- hundreds). - [ our_array_offset, - x_index, - y_index, - odata_offset, - idata_offset, - index_in, - index_out, - get_global_id_0, - get_local_id_0, - get_local_id_1, - get_group_id_0, - get_group_id_1, - get_group_id_2, - j - ] = - zipWith (flip VName) [30 ..] $ - map - nameFromString - [ "our_array_offset", - "x_index", - "y_index", - "odata_offset", - "idata_offset", - "index_in", - "index_out", - "get_global_id_0", - "get_local_id_0", - "get_local_id_1", - "get_group_id_0", - "get_group_id_1", - "get_group_id_2", - "j" - ] - - get_ids = - mconcat - [ DeclareScalar get_global_id_0 Nonvolatile int32, - Op $ GetGlobalId get_global_id_0 0, - DeclareScalar get_local_id_0 Nonvolatile int32, - Op $ GetLocalId get_local_id_0 0, - DeclareScalar get_local_id_1 Nonvolatile int32, - Op $ GetLocalId get_local_id_1 1, - DeclareScalar get_group_id_0 Nonvolatile int32, - Op $ GetGroupId get_group_id_0 0, - DeclareScalar get_group_id_1 Nonvolatile int32, - Op $ GetGroupId get_group_id_1 1, - DeclareScalar get_group_id_2 Nonvolatile int32, - Op $ GetGroupId get_group_id_2 2 - ] - - mkTranspose body = - mconcat - [ get_ids, - dec our_array_offset $ vi32 get_group_id_2 * width * height, - dec odata_offset $ - (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset, - dec idata_offset $ - (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset, - body - ] - - lowDimBody x_in_index y_in_index x_out_index y_out_index = - mconcat - [ dec x_index x_in_index, - dec y_index y_in_index, - dec index_in $ vi32 y_index * width + vi32 x_index, - when (vi32 x_index .<. width .&&. vi32 y_index .<. height) $ - Write - block - (elements $ sExt64 $ vi32 get_local_id_1 * (block_dim + 1) + vi32 get_local_id_0) - t - (Space "local") - Nonvolatile - $ index - idata - (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) - t - (Space "global") - Nonvolatile, - Op $ Barrier FenceLocal, - SetScalar x_index $ untyped x_out_index, - SetScalar y_index $ untyped y_out_index, - dec index_out $ vi32 y_index * height + vi32 x_index, - when (vi32 x_index .<. height .&&. vi32 y_index .<. width) $ - Write - odata - (elements $ sExt64 (vi32 odata_offset + vi32 index_out)) - t - (Space "global") - Nonvolatile - $ index - block - (elements $ sExt64 $ vi32 get_local_id_0 * (block_dim + 1) + vi32 get_local_id_1) - t - (Space "local") - Nonvolatile - ] - --- | Generate a transpose kernel. There is special support to handle --- input arrays with low width, low height, or both. --- --- Normally when transposing a @[2][n]@ array we would use a @FUT_BLOCK_DIM x --- FUT_BLOCK_DIM@ group to process a @[2][FUT_BLOCK_DIM]@ slice of the input --- array. This would mean that many of the threads in a group would be inactive. --- We try to remedy this by using a special kernel that will process a larger --- part of the input, by using more complex indexing. In our example, we could --- use all threads in a group if we are processing @(2/FUT_BLOCK_DIM)@ as large --- a slice of each rows per group. The variable @mulx@ contains this factor for --- the kernel to handle input arrays with low height. --- --- See issue #308 on GitHub for more details. --- --- These kernels are optimized to ensure all global reads and writes --- are coalesced, and to avoid bank conflicts in shared memory. Each --- thread group transposes a 2D tile of block_dim*2 by block_dim*2 --- elements. The size of a thread group is block_dim/2 by --- block_dim*2, meaning that each thread will process 4 elements in a --- 2D tile. The shared memory array containing the 2D tile consists --- of block_dim*2 by block_dim*2+1 elements. Padding each row with --- an additional element prevents bank conflicts from occuring when --- the tile is accessed column-wise. -mapTransposeKernel :: - String -> - Integer -> - TransposeArgs -> - PrimType -> - TransposeType -> - Kernel -mapTransposeKernel desc block_dim_int args t kind = - Kernel - { kernelBody = - DeclareMem block (Space "local") - <> Op (LocalAlloc block block_size) - <> mapTranspose block_dim args t kind, - kernelUses = uses, - kernelNumGroups = map untyped num_groups, - kernelGroupSize = map untyped group_size, - kernelName = nameFromString name, - kernelFailureTolerant = True - } - where - pad2DBytes k = k * (k + 1) * primByteSize t - block_size = - bytes $ - case kind of - TransposeSmall -> 1 :: TExp Int64 - -- Not used, but AMD's OpenCL - -- does not like zero-size - -- local memory. - TransposeNormal -> fromInteger $ pad2DBytes $ 2 * block_dim_int - TransposeLowWidth -> fromInteger $ pad2DBytes block_dim_int - TransposeLowHeight -> fromInteger $ pad2DBytes block_dim_int - block_dim = fromInteger block_dim_int :: TExp Int32 - - ( odata, - basic_odata_offset, - idata, - basic_idata_offset, - width, - height, - mulx, - muly, - num_arrays, - block - ) = args - - (num_groups, group_size) = - case kind of - TransposeSmall -> - ( [(num_arrays * width * height) `divUp` (block_dim * block_dim)], - [block_dim * block_dim] - ) - TransposeLowWidth -> - lowDimKernelAndGroupSize block_dim num_arrays width $ height `divUp` muly - TransposeLowHeight -> - lowDimKernelAndGroupSize block_dim num_arrays (width `divUp` mulx) height - TransposeNormal -> - let actual_dim = block_dim * 2 - in ( [ width `divUp` actual_dim, - height `divUp` actual_dim, - num_arrays - ], - [actual_dim, actual_dim `quot` elemsPerThread, 1] - ) - - uses = - map - (`ScalarUse` int32) - ( namesToList $ - mconcat $ - map - freeIn - [ basic_odata_offset, - basic_idata_offset, - num_arrays, - width, - height, - mulx, - muly - ] - ) - ++ map MemoryUse [odata, idata] - - name = - case kind of - TransposeSmall -> desc ++ "_small" - TransposeLowHeight -> desc ++ "_low_height" - TransposeLowWidth -> desc ++ "_low_width" - TransposeNormal -> desc - -lowDimKernelAndGroupSize :: - TExp Int32 -> - TExp Int32 -> - TExp Int32 -> - TExp Int32 -> - ([TExp Int32], [TExp Int32]) -lowDimKernelAndGroupSize block_dim num_arrays x_elems y_elems = - ( [ x_elems `divUp` block_dim, - y_elems `divUp` block_dim, - num_arrays - ], - [block_dim, block_dim, 1] - ) diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs @@ -23,7 +23,6 @@ import Control.Monad import Data.Bifunctor -import Data.List (elemIndex, find) import qualified Data.Map as M import Data.Maybe import qualified Futhark.CodeGen.ImpCode.Multicore as Imp @@ -31,7 +30,6 @@ import Futhark.Error import Futhark.IR.MCMem import Futhark.Transform.Rename -import Futhark.Util (maybeNth) import Prelude hiding (quot, rem) -- | Is there an atomic t'BinOp' corresponding to this t'BinOp'? @@ -70,7 +68,7 @@ toParam name (Prim pt) = return [Imp.ScalarParam name pt] toParam name (Mem space) = return [Imp.MemParam name space] toParam name Array {} = pure <$> arrParam name -toParam name Acc {} = error $ "toParam Acc: " ++ pretty name +toParam _name Acc {} = pure [] -- FIXME? Are we sure this works? getSpace :: SegOp () MCMem -> SegSpace getSpace (SegHist _ space _ _ _) = space @@ -156,10 +154,10 @@ isLoadBalanced (Imp.Op (Imp.ParLoop _ _ _ code _ _ _)) = isLoadBalanced code isLoadBalanced _ = True -segBinOpComm' :: [SegBinOp lore] -> Commutativity +segBinOpComm' :: [SegBinOp rep] -> Commutativity segBinOpComm' = mconcat . map segBinOpComm -decideScheduling' :: SegOp () lore -> Imp.Code -> Imp.Scheduling +decideScheduling' :: SegOp () rep -> Imp.Code -> Imp.Scheduling decideScheduling' SegHist {} _ = Imp.Static decideScheduling' SegScan {} _ = Imp.Static decideScheduling' (SegRed _ _ reds _ _) code = @@ -244,25 +242,25 @@ -- | A function for generating code for an atomic update. Assumes -- that the bucket is in-bounds. -type DoAtomicUpdate lore r = +type DoAtomicUpdate rep r = [VName] -> [Imp.TExp Int64] -> MulticoreGen () -- | The mechanism that will be used for performing the atomic update. -- Approximates how efficient it will be. Ordered from most to least -- efficient. -data AtomicUpdate lore r - = AtomicPrim (DoAtomicUpdate lore r) +data AtomicUpdate rep r + = AtomicPrim (DoAtomicUpdate rep r) | -- | Can be done by efficient swaps. - AtomicCAS (DoAtomicUpdate lore r) + AtomicCAS (DoAtomicUpdate rep r) | -- | Requires explicit locking. - AtomicLocking (Locking -> DoAtomicUpdate lore r) + AtomicLocking (Locking -> DoAtomicUpdate rep r) atomicUpdateLocking :: AtomicBinOp -> Lambda MCMem -> AtomicUpdate MCMem () atomicUpdateLocking atomicBinOp lam - | Just ops_and_ts <- splitOp lam, + | Just ops_and_ts <- lamIsBinOp lam, all (\(_, t, _, _) -> supportedPrims $ primBitSize t) ops_and_ts = primOrCas ops_and_ts $ \arrs bucket -> -- If the operator is a vectorised binary operator on 32-bit values, @@ -408,24 +406,6 @@ (sExt32 <$> bucket_offset) (tvVar run_loop) (toBits (Imp.var x t)) - --- | Horizontally fission a lambda that models a binary operator. -splitOp :: ASTLore lore => Lambda lore -> Maybe [(BinOp, PrimType, VName, VName)] -splitOp lam = mapM splitStm $ bodyResult $ lambdaBody lam - where - n = length $ lambdaReturnType lam - splitStm (Var res) = do - Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y))) <- - find (([res] ==) . patternNames . stmPattern) $ - stmsToList $ bodyStms $ lambdaBody lam - i <- Var res `elemIndex` bodyResult (lambdaBody lam) - xp <- maybeNth i $ lambdaParams lam - yp <- maybeNth (n + i) $ lambdaParams lam - guard $ paramName xp == x - guard $ paramName yp == y - Prim t <- Just $ patElemType pe - return (op, t, paramName xp, paramName yp) - splitStm _ = Nothing -- TODO for supporting 8 and 16 bits (and 128) -- we need a functions for converting to and from bits diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs @@ -30,7 +30,7 @@ -- | Split some list into chunks equal to the number of values -- returned by each 'SegBinOp' -segHistOpChunks :: [HistOp lore] -> [a] -> [[a]] +segHistOpChunks :: [HistOp rep] -> [a] -> [[a]] segHistOpChunks = chunks . map (length . histNeutral) nonsegmentedHist :: diff --git a/src/Futhark/CodeGen/ImpGen/OpenCL.hs b/src/Futhark/CodeGen/ImpGen/OpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/OpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/OpenCL.hs @@ -7,11 +7,11 @@ import Data.Bifunctor (second) import qualified Futhark.CodeGen.ImpCode.OpenCL as OpenCL -import Futhark.CodeGen.ImpGen.Kernels -import Futhark.CodeGen.ImpGen.Kernels.ToOpenCL -import Futhark.IR.KernelsMem +import Futhark.CodeGen.ImpGen.GPU +import Futhark.CodeGen.ImpGen.GPU.ToOpenCL +import Futhark.IR.GPUMem import Futhark.MonadFreshNames -- | Compile the program to ImpCode with OpenCL kernels. -compileProg :: MonadFreshNames m => Prog KernelsMem -> m (Warnings, OpenCL.Program) +compileProg :: MonadFreshNames m => Prog GPUMem -> m (Warnings, OpenCL.Program) compileProg prog = second kernelsToOpenCL <$> compileProgOpenCL prog diff --git a/src/Futhark/CodeGen/SetDefaultSpace.hs b/src/Futhark/CodeGen/SetDefaultSpace.hs --- a/src/Futhark/CodeGen/SetDefaultSpace.hs +++ b/src/Futhark/CodeGen/SetDefaultSpace.hs @@ -38,10 +38,10 @@ param setExtValueSpace :: Space -> ExternalValue -> ExternalValue -setExtValueSpace space (OpaqueValue desc vs) = - OpaqueValue desc $ map (setValueSpace space) vs -setExtValueSpace space (TransparentValue v) = - TransparentValue $ setValueSpace space v +setExtValueSpace space (OpaqueValue u desc vs) = + OpaqueValue u desc $ map (setValueSpace space) vs +setExtValueSpace space (TransparentValue u v) = + TransparentValue u $ setValueSpace space v setValueSpace :: Space -> ValueDesc -> ValueDesc setValueSpace space (ArrayValue mem _ bt ept shape) = diff --git a/src/Futhark/Compiler.hs b/src/Futhark/Compiler.hs --- a/src/Futhark/Compiler.hs +++ b/src/Futhark/Compiler.hs @@ -70,8 +70,8 @@ -- 'Pipeline', and finish up with the given 'Action'. runCompilerOnProgram :: FutharkConfig -> - Pipeline I.SOACS lore -> - Action lore -> + Pipeline I.SOACS rep -> + Action rep -> FilePath -> IO () runCompilerOnProgram config pipeline action file = do @@ -95,9 +95,9 @@ -- 'Pipeline', and return it. runPipelineOnProgram :: FutharkConfig -> - Pipeline I.SOACS tolore -> + Pipeline I.SOACS torep -> FilePath -> - FutharkM (Prog tolore) + FutharkM (Prog torep) runPipelineOnProgram config pipeline file = do when (pipelineVerbose pipeline_config) $ logMsg ("Reading and type-checking source program" :: String) diff --git a/src/Futhark/Compiler/CLI.hs b/src/Futhark/Compiler/CLI.hs --- a/src/Futhark/Compiler/CLI.hs +++ b/src/Futhark/Compiler/CLI.hs @@ -32,13 +32,13 @@ -- | The longer action description. String -> -- | The pipeline to use. - Pipeline SOACS lore -> + Pipeline SOACS rep -> -- | The action to take on the result of the pipeline. ( FutharkConfig -> cfg -> CompilerMode -> FilePath -> - Prog lore -> + Prog rep -> FutharkM () ) -> -- | Program name diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -28,8 +28,8 @@ -- A monad that implements 'MonadBinder' tracks the statements added -- so far, the current names in scope, and allows you to add -- additional statements with 'addStm'. Any monad that implements --- 'MonadBinder' also implements the t'Lore' type family, which --- indicates which lore it works with. Inside a 'MonadBinder' we can +-- 'MonadBinder' also implements the t'REp' type family, which +-- indicates which rep it works with. Inside a 'MonadBinder' we can -- use 'collectStms' to gather up the 'Stms' added with 'addStm' in -- some nested computation. -- @@ -119,7 +119,7 @@ letSubExp :: MonadBinder m => String -> - Exp (Lore m) -> + Exp (Rep m) -> m SubExp letSubExp _ (BasicOp (SubExp se)) = return se letSubExp desc e = Var <$> letExp desc e @@ -127,7 +127,7 @@ letExp :: MonadBinder m => String -> - Exp (Lore m) -> + Exp (Rep m) -> m VName letExp _ (BasicOp (SubExp (Var v))) = return v @@ -144,7 +144,7 @@ String -> VName -> Slice SubExp -> - Exp (Lore m) -> + Exp (Rep m) -> m VName letInPlace desc src slice e = do tmp <- letSubExp (desc ++ "_tmp") e @@ -153,14 +153,14 @@ letSubExps :: MonadBinder m => String -> - [Exp (Lore m)] -> + [Exp (Rep m)] -> m [SubExp] letSubExps desc = mapM $ letSubExp desc letTupExp :: (MonadBinder m) => String -> - Exp (Lore m) -> + Exp (Rep m) -> m [VName] letTupExp _ (BasicOp (SubExp (Var v))) = return [v] @@ -173,7 +173,7 @@ letTupExp' :: (MonadBinder m) => String -> - Exp (Lore m) -> + Exp (Rep m) -> m [SubExp] letTupExp' _ (BasicOp (SubExp se)) = return [se] letTupExp' name ses = map Var <$> letTupExp name ses @@ -181,25 +181,25 @@ eSubExp :: MonadBinder m => SubExp -> - m (Exp (Lore m)) + m (Exp (Rep m)) eSubExp = pure . BasicOp . SubExp eIf :: - (MonadBinder m, BranchType (Lore m) ~ ExtType) => - m (Exp (Lore m)) -> - m (Body (Lore m)) -> - m (Body (Lore m)) -> - m (Exp (Lore m)) + (MonadBinder m, BranchType (Rep m) ~ ExtType) => + m (Exp (Rep m)) -> + m (Body (Rep m)) -> + m (Body (Rep m)) -> + m (Exp (Rep m)) eIf ce te fe = eIf' ce te fe IfNormal -- | As 'eIf', but an 'IfSort' can be given. eIf' :: - (MonadBinder m, BranchType (Lore m) ~ ExtType) => - m (Exp (Lore m)) -> - m (Body (Lore m)) -> - m (Body (Lore m)) -> + (MonadBinder m, BranchType (Rep m) ~ ExtType) => + m (Exp (Rep m)) -> + m (Body (Rep m)) -> + m (Body (Rep m)) -> IfSort -> - m (Exp (Lore m)) + m (Exp (Rep m)) eIf' ce te fe if_sort = do ce' <- letSubExp "cond" =<< ce te' <- insertStmsM te @@ -222,7 +222,7 @@ -- The type of a body. Watch out: this only works for the degenerate -- case where the body does not already return its context. -bodyExtType :: (HasScope lore m, Monad m) => Body lore -> m [ExtType] +bodyExtType :: (HasScope rep m, Monad m) => Body rep -> m [ExtType] bodyExtType (Body _ stms res) = existentialiseExtTypes (M.keys stmsscope) . staticShapes <$> extendedScope (traverse subExpType res) stmsscope @@ -232,9 +232,9 @@ eBinOp :: MonadBinder m => BinOp -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) + m (Exp (Rep m)) -> + m (Exp (Rep m)) -> + m (Exp (Rep m)) eBinOp op x y = do x' <- letSubExp "x" =<< x y' <- letSubExp "y" =<< y @@ -243,9 +243,9 @@ eCmpOp :: MonadBinder m => CmpOp -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) + m (Exp (Rep m)) -> + m (Exp (Rep m)) -> + m (Exp (Rep m)) eCmpOp op x y = do x' <- letSubExp "x" =<< x y' <- letSubExp "y" =<< y @@ -254,16 +254,16 @@ eConvOp :: MonadBinder m => ConvOp -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) + m (Exp (Rep m)) -> + m (Exp (Rep m)) eConvOp op x = do x' <- letSubExp "x" =<< x return $ BasicOp $ ConvOp op x' eSignum :: MonadBinder m => - m (Exp (Lore m)) -> - m (Exp (Lore m)) + m (Exp (Rep m)) -> + m (Exp (Rep m)) eSignum em = do e <- em e' <- letSubExp "signum_arg" e @@ -276,14 +276,14 @@ eCopy :: MonadBinder m => - m (Exp (Lore m)) -> - m (Exp (Lore m)) + m (Exp (Rep m)) -> + m (Exp (Rep m)) eCopy e = BasicOp . Copy <$> (letExp "copy_arg" =<< e) eBody :: (MonadBinder m) => - [m (Exp (Lore m))] -> - m (Body (Lore m)) + [m (Exp (Rep m))] -> + m (Body (Rep m)) eBody es = buildBody_ $ do es' <- sequence es xs <- mapM (letTupExp "x") es' @@ -291,8 +291,8 @@ eLambda :: MonadBinder m => - Lambda (Lore m) -> - [m (Exp (Lore m))] -> + Lambda (Rep m) -> + [m (Exp (Rep m))] -> m [SubExp] eLambda lam args = do zipWithM_ bindParam (lambdaParams lam) args @@ -303,9 +303,9 @@ eRoundToMultipleOf :: MonadBinder m => IntType -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) + m (Exp (Rep m)) -> + m (Exp (Rep m)) -> + m (Exp (Rep m)) eRoundToMultipleOf t x d = ePlus x (eMod (eMinus d (eMod x d)) d) where @@ -318,9 +318,9 @@ MonadBinder m => Int -> VName -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) + m (Exp (Rep m)) -> + m (Exp (Rep m)) -> + m (Exp (Rep m)) eSliceArray d arr i n = do arr_t <- lookupType arr let skips = map (slice (constant (0 :: Int64))) $ take d $ arrayDims arr_t @@ -334,8 +334,8 @@ eOutOfBounds :: MonadBinder m => VName -> - [m (Exp (Lore m))] -> - m (Exp (Lore m)) + [m (Exp (Rep m))] -> + m (Exp (Rep m)) eOutOfBounds arr is = do arr_t <- lookupType arr let ws = arrayDims arr_t @@ -354,11 +354,11 @@ -- | Write to an index of the array, if within bounds. Otherwise, -- nothing. Produces the updated array. eWriteArray :: - (MonadBinder m, BranchType (Lore m) ~ ExtType) => + (MonadBinder m, BranchType (Rep m) ~ ExtType) => VName -> - [m (Exp (Lore m))] -> - m (Exp (Lore m)) -> - m (Exp (Lore m)) + [m (Exp (Rep m))] -> + m (Exp (Rep m)) -> + m (Exp (Rep m)) eWriteArray arr is v = do arr_t <- lookupType arr is' <- mapM (letSubExp "write_i") =<< sequence is @@ -381,7 +381,7 @@ ifCommon [arr_t] -- | Construct an unspecified value of the given type. -eBlank :: MonadBinder m => Type -> m (Exp (Lore m)) +eBlank :: MonadBinder m => Type -> m (Exp (Rep m)) eBlank (Prim t) = return $ BasicOp $ SubExp $ Constant $ blankPrimValue t eBlank (Array t shape _) = return $ BasicOp $ Scratch t $ shapeDims shape eBlank Acc {} = error "eBlank: cannot create blank accumulator" @@ -419,14 +419,14 @@ BinOp -> SubExp -> [SubExp] -> - m (Exp (Lore m)) + m (Exp (Rep m)) foldBinOp _ ne [] = return $ BasicOp $ SubExp ne foldBinOp bop ne (e : es) = eBinOp bop (pure $ BasicOp $ SubExp e) (foldBinOp bop ne es) -- | True if all operands are true. -eAll :: MonadBinder m => [SubExp] -> m (Exp (Lore m)) +eAll :: MonadBinder m => [SubExp] -> m (Exp (Rep m)) eAll [] = pure $ BasicOp $ SubExp $ constant True eAll (x : xs) = foldBinOp LogAnd x xs @@ -435,25 +435,25 @@ -- result types are the same. (This assumption should be fixed at -- some point.) binOpLambda :: - (MonadBinder m, Bindable (Lore m)) => + (MonadBinder m, Bindable (Rep m)) => BinOp -> PrimType -> - m (Lambda (Lore m)) + m (Lambda (Rep m)) binOpLambda bop t = binLambda (BinOp bop) t t -- | As 'binOpLambda', but for t'CmpOp's. cmpOpLambda :: - (MonadBinder m, Bindable (Lore m)) => + (MonadBinder m, Bindable (Rep m)) => CmpOp -> - m (Lambda (Lore m)) + m (Lambda (Rep m)) cmpOpLambda cop = binLambda (CmpOp cop) (cmpOpType cop) Bool binLambda :: - (MonadBinder m, Bindable (Lore m)) => + (MonadBinder m, Bindable (Rep m)) => (SubExp -> SubExp -> BasicOp) -> PrimType -> PrimType -> - m (Lambda (Lore m)) + m (Lambda (Rep m)) binLambda bop arg_t ret_t = do x <- newVName "x" y <- newVName "y" @@ -473,9 +473,9 @@ -- | Easily construct a 'Lambda' within a 'MonadBinder'. mkLambda :: MonadBinder m => - [LParam (Lore m)] -> + [LParam (Rep m)] -> m Result -> - m (Lambda (Lore m)) + m (Lambda (Rep m)) mkLambda params m = do (body, ret) <- buildBody . localScope (scopeOfLParams params) $ do res <- m @@ -521,7 +521,7 @@ ifCommon ts = IfDec (staticShapes ts) IfNormal -- | Conveniently construct a body that contains no bindings. -resultBody :: Bindable lore => [SubExp] -> Body lore +resultBody :: Bindable rep => [SubExp] -> Body rep resultBody = mkBody mempty -- | Conveniently construct a body that contains no bindings - but @@ -529,15 +529,15 @@ resultBodyM :: MonadBinder m => [SubExp] -> - m (Body (Lore m)) + m (Body (Rep m)) resultBodyM = mkBodyM mempty -- | Evaluate the action, producing a body, then wrap it in all the -- bindings it created using 'addStm'. insertStmsM :: (MonadBinder m) => - m (Body (Lore m)) -> - m (Body (Lore m)) + m (Body (Rep m)) -> + m (Body (Rep m)) insertStmsM m = do (Body _ bnds res, otherbnds) <- collectStms m mkBodyM (otherbnds <> bnds) res @@ -548,7 +548,7 @@ buildBody :: MonadBinder m => m (Result, a) -> - m (Body (Lore m), a) + m (Body (Rep m), a) buildBody m = do ((res, v), stms) <- collectStms m body <- mkBodyM stms res @@ -558,16 +558,16 @@ buildBody_ :: MonadBinder m => m Result -> - m (Body (Lore m)) + m (Body (Rep m)) buildBody_ m = fst <$> buildBody ((,()) <$> m) -- | Change that result where evaluation of the body would stop. Also -- change type annotations at branches. mapResult :: - Bindable lore => - (Result -> Body lore) -> - Body lore -> - Body lore + Bindable rep => + (Result -> Body rep) -> + Body rep -> + Body rep mapResult f (Body _ bnds res) = let Body _ bnds2 newres = f res in mkBody (bnds <> bnds2) newres @@ -625,15 +625,15 @@ -- | Can be used as the definition of 'mkLetNames' for a 'Bindable' -- instance for simple representations. simpleMkLetNames :: - ( ExpDec lore ~ (), - LetDec lore ~ Type, + ( ExpDec rep ~ (), + LetDec rep ~ Type, MonadFreshNames m, - TypedOp (Op lore), - HasScope lore m + TypedOp (Op rep), + HasScope rep m ) => [VName] -> - Exp lore -> - m (Stm lore) + Exp rep -> + m (Stm rep) simpleMkLetNames names e = do et <- expExtType e (ts, shapes) <- instantiateShapes' et @@ -644,7 +644,7 @@ -- | Instances of this class can be converted to Futhark expressions -- within a 'MonadBinder'. class ToExp a where - toExp :: MonadBinder m => a -> m (Exp (Lore m)) + toExp :: MonadBinder m => a -> m (Exp (Rep m)) instance ToExp SubExp where toExp = return . BasicOp . SubExp diff --git a/src/Futhark/IR.hs b/src/Futhark/IR.hs --- a/src/Futhark/IR.hs +++ b/src/Futhark/IR.hs @@ -1,7 +1,4 @@ --- | A convenient re-export of basic AST modules. Note that --- "Futhark.IR.Lore" is not exported, as this would --- cause name clashes. You are advised to use a qualified import of --- the lore module, if you need it. +-- | A convenient re-export of basic AST modules. module Futhark.IR ( module Futhark.IR.Prop, module Futhark.IR.Traversals, diff --git a/src/Futhark/IR/Aliases.hs b/src/Futhark/IR/Aliases.hs --- a/src/Futhark/IR/Aliases.hs +++ b/src/Futhark/IR/Aliases.hs @@ -8,7 +8,7 @@ -- | A representation where all bindings are annotated with aliasing -- information. module Futhark.IR.Aliases - ( -- * The Lore definition + ( -- * The representation definition Aliases, AliasDec (..), VarAliases, @@ -60,8 +60,8 @@ import Futhark.Transform.Substitute import qualified Futhark.Util.Pretty as PP --- | The lore for the basic representation. -data Aliases lore +-- | The rep for the basic representation. +data Aliases rep -- | A wrapper around 'AliasDec' to get around the fact that we need an -- 'Ord' instance, which 'AliasDec does not have. @@ -102,18 +102,15 @@ -- consumed inside of it. type BodyAliasing = ([VarAliases], ConsumedInExp) -instance - (Decorations lore, CanBeAliased (Op lore)) => - Decorations (Aliases lore) - where - type LetDec (Aliases lore) = (VarAliases, LetDec lore) - type ExpDec (Aliases lore) = (ConsumedInExp, ExpDec lore) - type BodyDec (Aliases lore) = (BodyAliasing, BodyDec lore) - type FParamInfo (Aliases lore) = FParamInfo lore - type LParamInfo (Aliases lore) = LParamInfo lore - type RetType (Aliases lore) = RetType lore - type BranchType (Aliases lore) = BranchType lore - type Op (Aliases lore) = OpWithAliases (Op lore) +instance (RepTypes rep, CanBeAliased (Op rep)) => RepTypes (Aliases rep) where + type LetDec (Aliases rep) = (VarAliases, LetDec rep) + type ExpDec (Aliases rep) = (ConsumedInExp, ExpDec rep) + type BodyDec (Aliases rep) = (BodyAliasing, BodyDec rep) + type FParamInfo (Aliases rep) = FParamInfo rep + type LParamInfo (Aliases rep) = LParamInfo rep + type RetType (Aliases rep) = RetType rep + type BranchType (Aliases rep) = BranchType rep + type Op (Aliases rep) = OpWithAliases (Op rep) instance AliasesOf (VarAliases, dec) where aliasesOf = unAliases . fst @@ -121,28 +118,28 @@ instance FreeDec AliasDec withoutAliases :: - (HasScope (Aliases lore) m, Monad m) => - ReaderT (Scope lore) m a -> + (HasScope (Aliases rep) m, Monad m) => + ReaderT (Scope rep) m a -> m a withoutAliases m = do scope <- asksScope removeScopeAliases runReaderT m scope -instance (ASTLore lore, CanBeAliased (Op lore)) => ASTLore (Aliases lore) where +instance (ASTRep rep, CanBeAliased (Op rep)) => ASTRep (Aliases rep) where expTypesFromPattern = withoutAliases . expTypesFromPattern . removePatternAliases -instance (ASTLore lore, CanBeAliased (Op lore)) => Aliased (Aliases lore) where +instance (ASTRep rep, CanBeAliased (Op rep)) => Aliased (Aliases rep) where bodyAliases = map unAliases . fst . fst . bodyDec consumedInBody = unAliases . snd . fst . bodyDec -instance (ASTLore lore, CanBeAliased (Op lore)) => PrettyLore (Aliases lore) where - ppExpLore (consumed, inner) e = +instance (ASTRep rep, CanBeAliased (Op rep)) => PrettyRep (Aliases rep) where + ppExpDec (consumed, inner) e = maybeComment $ catMaybes [ exp_dec, merge_dec, - ppExpLore inner $ removeExpAliases e + ppExpDec inner $ removeExpAliases e ] where merge_dec = @@ -180,20 +177,20 @@ PP.text "-- Result of " <> PP.ppr name <> PP.text " aliases " <> PP.commasep (map PP.ppr als') -removeAliases :: CanBeAliased (Op lore) => Rephraser Identity (Aliases lore) lore +removeAliases :: CanBeAliased (Op rep) => Rephraser Identity (Aliases rep) rep removeAliases = Rephraser - { rephraseExpLore = return . snd, - rephraseLetBoundLore = return . snd, - rephraseBodyLore = return . snd, - rephraseFParamLore = return, - rephraseLParamLore = return, + { rephraseExpDec = return . snd, + rephraseLetBoundDec = return . snd, + rephraseBodyDec = return . snd, + rephraseFParamDec = return, + rephraseLParamDec = return, rephraseRetType = return, rephraseBranchType = return, rephraseOp = return . removeOpAliases } -removeScopeAliases :: Scope (Aliases lore) -> Scope lore +removeScopeAliases :: Scope (Aliases rep) -> Scope rep removeScopeAliases = M.map unAlias where unAlias (LetName (_, dec)) = LetName dec @@ -202,33 +199,33 @@ unAlias (IndexName it) = IndexName it removeProgAliases :: - CanBeAliased (Op lore) => - Prog (Aliases lore) -> - Prog lore + CanBeAliased (Op rep) => + Prog (Aliases rep) -> + Prog rep removeProgAliases = runIdentity . rephraseProg removeAliases removeFunDefAliases :: - CanBeAliased (Op lore) => - FunDef (Aliases lore) -> - FunDef lore + CanBeAliased (Op rep) => + FunDef (Aliases rep) -> + FunDef rep removeFunDefAliases = runIdentity . rephraseFunDef removeAliases removeExpAliases :: - CanBeAliased (Op lore) => - Exp (Aliases lore) -> - Exp lore + CanBeAliased (Op rep) => + Exp (Aliases rep) -> + Exp rep removeExpAliases = runIdentity . rephraseExp removeAliases removeStmAliases :: - CanBeAliased (Op lore) => - Stm (Aliases lore) -> - Stm lore + CanBeAliased (Op rep) => + Stm (Aliases rep) -> + Stm rep removeStmAliases = runIdentity . rephraseStm removeAliases removeLambdaAliases :: - CanBeAliased (Op lore) => - Lambda (Aliases lore) -> - Lambda lore + CanBeAliased (Op rep) => + Lambda (Aliases rep) -> + Lambda rep removeLambdaAliases = runIdentity . rephraseLambda removeAliases removePatternAliases :: @@ -237,26 +234,26 @@ removePatternAliases = runIdentity . rephrasePattern (return . snd) addAliasesToPattern :: - (ASTLore lore, CanBeAliased (Op lore), Typed dec) => + (ASTRep rep, CanBeAliased (Op rep), Typed dec) => PatternT dec -> - Exp (Aliases lore) -> + Exp (Aliases rep) -> PatternT (VarAliases, dec) addAliasesToPattern pat e = uncurry Pattern $ mkPatternAliases pat e mkAliasedBody :: - (ASTLore lore, CanBeAliased (Op lore)) => - BodyDec lore -> - Stms (Aliases lore) -> + (ASTRep rep, CanBeAliased (Op rep)) => + BodyDec rep -> + Stms (Aliases rep) -> Result -> - Body (Aliases lore) -mkAliasedBody innerlore bnds res = - Body (mkBodyAliases bnds res, innerlore) bnds res + Body (Aliases rep) +mkAliasedBody dec bnds res = + Body (mkBodyAliases bnds res, dec) bnds res mkPatternAliases :: - (Aliased lore, Typed dec) => + (Aliased rep, Typed dec) => PatternT dec -> - Exp lore -> + Exp rep -> ( [PatElemT (VarAliases, dec)], [PatElemT (VarAliases, dec)] ) @@ -273,7 +270,7 @@ ) where annotateBindee bindee names = - bindee `setPatElemLore` (AliasDec names', patElemDec bindee) + bindee `setPatElemDec` (AliasDec names', patElemDec bindee) where names' = case patElemType bindee of @@ -282,9 +279,9 @@ _ -> mempty mkContextAliases :: - Aliased lore => + Aliased rep => PatternT dec -> - Exp lore -> + Exp rep -> [Names] mkContextAliases pat (DoLoop ctxmerge valmerge _ body) = let ctx = map fst ctxmerge @@ -307,8 +304,8 @@ replicate (length $ patternContextElements pat) mempty mkBodyAliases :: - Aliased lore => - Stms lore -> + Aliased rep => + Stms rep -> Result -> BodyAliasing mkBodyAliases bnds res = @@ -326,8 +323,8 @@ -- | The aliases of the result and everything consumed in the given -- statements. mkStmsAliases :: - Aliased lore => - Stms lore -> + Aliased rep => + Stms rep -> [SubExp] -> ([Names], Names) mkStmsAliases bnds res = delve mempty $ stmsToList bnds @@ -349,9 +346,9 @@ ) trackAliases :: - Aliased lore => + Aliased rep => AliasesAndConsumed -> - Stm lore -> + Stm rep -> AliasesAndConsumed trackAliases (aliasmap, consumed) stm = let pat = stmPattern stm @@ -371,18 +368,18 @@ look k = M.findWithDefault mempty k aliasmap mkAliasedLetStm :: - (ASTLore lore, CanBeAliased (Op lore)) => - Pattern lore -> - StmAux (ExpDec lore) -> - Exp (Aliases lore) -> - Stm (Aliases lore) + (ASTRep rep, CanBeAliased (Op rep)) => + Pattern rep -> + StmAux (ExpDec rep) -> + Exp (Aliases rep) -> + Stm (Aliases rep) mkAliasedLetStm pat (StmAux cs attrs dec) e = Let (addAliasesToPattern pat e) (StmAux cs attrs (AliasDec $ consumedInExp e, dec)) e -instance (Bindable lore, CanBeAliased (Op lore)) => Bindable (Aliases lore) where +instance (Bindable rep, CanBeAliased (Op rep)) => Bindable (Aliases rep) where mkExpDec pat e = let dec = mkExpDec (removePatternAliases pat) $ removeExpAliases e in (AliasDec $ consumedInExp e, dec) @@ -397,7 +394,7 @@ return $ mkAliasedLetStm pat dec e mkBody bnds res = - let Body bodylore _ _ = mkBody (fmap removeStmAliases bnds) res - in mkAliasedBody bodylore bnds res + let Body bodyrep _ _ = mkBody (fmap removeStmAliases bnds) res + in mkAliasedBody bodyrep bnds res -instance (ASTLore (Aliases lore), Bindable (Aliases lore)) => BinderOps (Aliases lore) +instance (ASTRep (Aliases rep), Bindable (Aliases rep)) => BinderOps (Aliases rep) diff --git a/src/Futhark/IR/Decorations.hs b/src/Futhark/IR/Decorations.hs deleted file mode 100644 --- a/src/Futhark/IR/Decorations.hs +++ /dev/null @@ -1,92 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} - --- | The core Futhark AST is parameterised by a @lore@ type parameter, --- which is then used to invoke the type families defined here. -module Futhark.IR.Decorations - ( Decorations (..), - module Futhark.IR.RetType, - ) -where - -import qualified Data.Kind -import Futhark.IR.Prop.Types -import Futhark.IR.RetType -import Futhark.IR.Syntax.Core - --- | A collection of type families, along with constraints specifying --- that the types they map to should satisfy some minimal --- requirements. -class - ( Show (LetDec l), - Show (ExpDec l), - Show (BodyDec l), - Show (FParamInfo l), - Show (LParamInfo l), - Show (RetType l), - Show (BranchType l), - Show (Op l), - Eq (LetDec l), - Eq (ExpDec l), - Eq (BodyDec l), - Eq (FParamInfo l), - Eq (LParamInfo l), - Eq (RetType l), - Eq (BranchType l), - Eq (Op l), - Ord (LetDec l), - Ord (ExpDec l), - Ord (BodyDec l), - Ord (FParamInfo l), - Ord (LParamInfo l), - Ord (RetType l), - Ord (BranchType l), - Ord (Op l), - IsRetType (RetType l), - IsBodyType (BranchType l), - Typed (FParamInfo l), - Typed (LParamInfo l), - Typed (LetDec l), - DeclTyped (FParamInfo l) - ) => - Decorations l - where - -- | Decoration for every let-pattern element. - type LetDec l :: Data.Kind.Type - - type LetDec l = Type - - -- | Decoration for every expression. - type ExpDec l :: Data.Kind.Type - - type ExpDec l = () - - -- | Decoration for every body. - type BodyDec l :: Data.Kind.Type - - type BodyDec l = () - - -- | Decoration for every (non-lambda) function parameter. - type FParamInfo l :: Data.Kind.Type - - type FParamInfo l = DeclType - - -- | Decoration for every lambda function parameter. - type LParamInfo l :: Data.Kind.Type - - type LParamInfo l = Type - - -- | The return type decoration of function calls. - type RetType l :: Data.Kind.Type - - type RetType l = DeclExtType - - -- | The return type decoration of branches. - type BranchType l :: Data.Kind.Type - - type BranchType l = ExtType - - -- | Extensible operation. - type Op l :: Data.Kind.Type - - type Op l = () diff --git a/src/Futhark/IR/GPU.hs b/src/Futhark/IR/GPU.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/GPU.hs @@ -0,0 +1,61 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} + +-- | A representation with flat parallelism via GPU-oriented kernels. +module Futhark.IR.GPU + ( GPU, + + -- * Module re-exports + module Futhark.IR.Prop, + module Futhark.IR.Traversals, + module Futhark.IR.Pretty, + module Futhark.IR.Syntax, + module Futhark.IR.GPU.Kernel, + module Futhark.IR.GPU.Sizes, + module Futhark.IR.SOACS.SOAC, + ) +where + +import Futhark.Binder +import Futhark.Construct +import Futhark.IR.GPU.Kernel +import Futhark.IR.GPU.Sizes +import Futhark.IR.Pretty +import Futhark.IR.Prop +import Futhark.IR.SOACS.SOAC hiding (HistOp (..)) +import Futhark.IR.Syntax +import Futhark.IR.Traversals +import qualified Futhark.TypeCheck as TypeCheck + +-- | The phantom data type for the kernels representation. +data GPU + +instance RepTypes GPU where + type Op GPU = HostOp GPU (SOAC GPU) + +instance ASTRep GPU where + expTypesFromPattern = return . expExtTypesFromPattern + +instance TypeCheck.CheckableOp GPU where + checkOp = typeCheckGPUOp Nothing + where + typeCheckGPUOp lvl = + typeCheckHostOp (typeCheckGPUOp . Just) lvl typeCheckSOAC + +instance TypeCheck.Checkable GPU + +instance Bindable GPU where + mkBody = Body () + mkExpPat ctx val _ = basicPattern ctx val + mkExpDec _ _ = () + mkLetNames = simpleMkLetNames + +instance BinderOps GPU + +instance PrettyRep GPU + +instance HasSegOp GPU where + type SegOpLevel GPU = SegLevel + asSegOp (SegOp op) = Just op + asSegOp _ = Nothing + segOp = SegOp diff --git a/src/Futhark/IR/GPU/Kernel.hs b/src/Futhark/IR/GPU/Kernel.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/GPU/Kernel.hs @@ -0,0 +1,346 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Futhark.IR.GPU.Kernel + ( -- * Size operations + SizeOp (..), + + -- * Host operations + HostOp (..), + typeCheckHostOp, + + -- * SegOp refinements + SegLevel (..), + + -- * Reexports + module Futhark.IR.GPU.Sizes, + module Futhark.IR.SegOp, + ) +where + +import Futhark.Analysis.Metrics +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.IR +import Futhark.IR.Aliases (Aliases) +import Futhark.IR.GPU.Sizes +import Futhark.IR.Prop.Aliases +import Futhark.IR.SegOp +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Rep +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import qualified Futhark.TypeCheck as TC +import Futhark.Util.Pretty + ( commasep, + parens, + ppr, + text, + (<+>), + ) +import qualified Futhark.Util.Pretty as PP +import Prelude hiding (id, (.)) + +-- | At which level the *body* of a t'SegOp' executes. +data SegLevel + = SegThread + { segNumGroups :: Count NumGroups SubExp, + segGroupSize :: Count GroupSize SubExp, + segVirt :: SegVirt + } + | SegGroup + { segNumGroups :: Count NumGroups SubExp, + segGroupSize :: Count GroupSize SubExp, + segVirt :: SegVirt + } + deriving (Eq, Ord, Show) + +instance PP.Pretty SegLevel where + ppr lvl = + PP.parens + ( lvl' <> PP.semi + <+> text "#groups=" <> ppr (segNumGroups lvl) <> PP.semi + <+> text "groupsize=" <> ppr (segGroupSize lvl) <> virt + ) + where + lvl' = case lvl of + SegThread {} -> "thread" + SegGroup {} -> "group" + virt = case segVirt lvl of + SegNoVirt -> mempty + SegNoVirtFull -> PP.semi <+> text "full" + SegVirt -> PP.semi <+> text "virtualise" + +instance Engine.Simplifiable SegLevel where + simplify (SegThread num_groups group_size virt) = + SegThread <$> traverse Engine.simplify num_groups + <*> traverse Engine.simplify group_size + <*> pure virt + simplify (SegGroup num_groups group_size virt) = + SegGroup <$> traverse Engine.simplify num_groups + <*> traverse Engine.simplify group_size + <*> pure virt + +instance Substitute SegLevel where + substituteNames substs (SegThread num_groups group_size virt) = + SegThread + (substituteNames substs num_groups) + (substituteNames substs group_size) + virt + substituteNames substs (SegGroup num_groups group_size virt) = + SegGroup + (substituteNames substs num_groups) + (substituteNames substs group_size) + virt + +instance Rename SegLevel where + rename = substituteRename + +instance FreeIn SegLevel where + freeIn' (SegThread num_groups group_size _) = + freeIn' num_groups <> freeIn' group_size + freeIn' (SegGroup num_groups group_size _) = + freeIn' num_groups <> freeIn' group_size + +-- | A simple size-level query or computation. +data SizeOp + = -- | @SplitSpace o w i elems_per_thread@. + -- + -- Computes how to divide array elements to + -- threads in a kernel. Returns the number of + -- elements in the chunk that the current thread + -- should take. + -- + -- @w@ is the length of the outer dimension in + -- the array. @i@ is the current thread + -- index. Each thread takes at most + -- @elems_per_thread@ elements. + -- + -- If the order @o@ is 'SplitContiguous', thread with index @i@ + -- should receive elements + -- @i*elems_per_tread, i*elems_per_thread + 1, + -- ..., i*elems_per_thread + (elems_per_thread-1)@. + -- + -- If the order @o@ is @'SplitStrided' stride@, + -- the thread will receive elements @i, + -- i+stride, i+2*stride, ..., + -- i+(elems_per_thread-1)*stride@. + SplitSpace SplitOrdering SubExp SubExp SubExp + | -- | Produce some runtime-configurable size. + GetSize Name SizeClass + | -- | The maximum size of some class. + GetSizeMax SizeClass + | -- | Compare size (likely a threshold) with some integer value. + CmpSizeLe Name SizeClass SubExp + | -- | @CalcNumGroups w max_num_groups group_size@ calculates the + -- number of GPU workgroups to use for an input of the given size. + -- The @Name@ is a size name. Note that @w@ is an i64 to avoid + -- overflow issues. + CalcNumGroups SubExp Name SubExp + deriving (Eq, Ord, Show) + +instance Substitute SizeOp where + substituteNames subst (SplitSpace o w i elems_per_thread) = + SplitSpace + (substituteNames subst o) + (substituteNames subst w) + (substituteNames subst i) + (substituteNames subst elems_per_thread) + substituteNames substs (CmpSizeLe name sclass x) = + CmpSizeLe name sclass (substituteNames substs x) + substituteNames substs (CalcNumGroups w max_num_groups group_size) = + CalcNumGroups + (substituteNames substs w) + max_num_groups + (substituteNames substs group_size) + substituteNames _ op = op + +instance Rename SizeOp where + rename (SplitSpace o w i elems_per_thread) = + SplitSpace + <$> rename o + <*> rename w + <*> rename i + <*> rename elems_per_thread + rename (CmpSizeLe name sclass x) = + CmpSizeLe name sclass <$> rename x + rename (CalcNumGroups w max_num_groups group_size) = + CalcNumGroups <$> rename w <*> pure max_num_groups <*> rename group_size + rename x = pure x + +instance IsOp SizeOp where + safeOp _ = True + cheapOp _ = True + +instance TypedOp SizeOp where + opType SplitSpace {} = pure [Prim int64] + opType (GetSize _ _) = pure [Prim int64] + opType (GetSizeMax _) = pure [Prim int64] + opType CmpSizeLe {} = pure [Prim Bool] + opType CalcNumGroups {} = pure [Prim int64] + +instance AliasedOp SizeOp where + opAliases _ = [mempty] + consumedInOp _ = mempty + +instance FreeIn SizeOp where + freeIn' (SplitSpace o w i elems_per_thread) = + freeIn' o <> freeIn' [w, i, elems_per_thread] + freeIn' (CmpSizeLe _ _ x) = freeIn' x + freeIn' (CalcNumGroups w _ group_size) = freeIn' w <> freeIn' group_size + freeIn' _ = mempty + +instance PP.Pretty SizeOp where + ppr (SplitSpace SplitContiguous w i elems_per_thread) = + text "split_space" + <> parens (commasep [ppr w, ppr i, ppr elems_per_thread]) + ppr (SplitSpace (SplitStrided stride) w i elems_per_thread) = + text "split_space_strided" + <> parens (commasep [ppr stride, ppr w, ppr i, ppr elems_per_thread]) + ppr (GetSize name size_class) = + text "get_size" <> parens (commasep [ppr name, ppr size_class]) + ppr (GetSizeMax size_class) = + text "get_size_max" <> parens (commasep [ppr size_class]) + ppr (CmpSizeLe name size_class x) = + text "cmp_size" <> parens (commasep [ppr name, ppr size_class]) + <+> text "<=" + <+> ppr x + ppr (CalcNumGroups w max_num_groups group_size) = + text "calc_num_groups" <> parens (commasep [ppr w, ppr max_num_groups, ppr group_size]) + +instance OpMetrics SizeOp where + opMetrics SplitSpace {} = seen "SplitSpace" + opMetrics GetSize {} = seen "GetSize" + opMetrics GetSizeMax {} = seen "GetSizeMax" + opMetrics CmpSizeLe {} = seen "CmpSizeLe" + opMetrics CalcNumGroups {} = seen "CalcNumGroups" + +typeCheckSizeOp :: TC.Checkable rep => SizeOp -> TC.TypeM rep () +typeCheckSizeOp (SplitSpace o w i elems_per_thread) = do + case o of + SplitContiguous -> return () + SplitStrided stride -> TC.require [Prim int64] stride + mapM_ (TC.require [Prim int64]) [w, i, elems_per_thread] +typeCheckSizeOp GetSize {} = return () +typeCheckSizeOp GetSizeMax {} = return () +typeCheckSizeOp (CmpSizeLe _ _ x) = TC.require [Prim int64] x +typeCheckSizeOp (CalcNumGroups w _ group_size) = do + TC.require [Prim int64] w + TC.require [Prim int64] group_size + +-- | A host-level operation; parameterised by what else it can do. +data HostOp rep op + = -- | A segmented operation. + SegOp (SegOp SegLevel rep) + | SizeOp SizeOp + | OtherOp op + deriving (Eq, Ord, Show) + +instance (ASTRep rep, Substitute op) => Substitute (HostOp rep op) where + substituteNames substs (SegOp op) = + SegOp $ substituteNames substs op + substituteNames substs (OtherOp op) = + OtherOp $ substituteNames substs op + substituteNames substs (SizeOp op) = + SizeOp $ substituteNames substs op + +instance (ASTRep rep, Rename op) => Rename (HostOp rep op) where + rename (SegOp op) = SegOp <$> rename op + rename (OtherOp op) = OtherOp <$> rename op + rename (SizeOp op) = SizeOp <$> rename op + +instance (ASTRep rep, IsOp op) => IsOp (HostOp rep op) where + safeOp (SegOp op) = safeOp op + safeOp (OtherOp op) = safeOp op + safeOp (SizeOp op) = safeOp op + + cheapOp (SegOp op) = cheapOp op + cheapOp (OtherOp op) = cheapOp op + cheapOp (SizeOp op) = cheapOp op + +instance TypedOp op => TypedOp (HostOp rep op) where + opType (SegOp op) = opType op + opType (OtherOp op) = opType op + opType (SizeOp op) = opType op + +instance (Aliased rep, AliasedOp op, ASTRep rep) => AliasedOp (HostOp rep op) where + opAliases (SegOp op) = opAliases op + opAliases (OtherOp op) = opAliases op + opAliases (SizeOp op) = opAliases op + + consumedInOp (SegOp op) = consumedInOp op + consumedInOp (OtherOp op) = consumedInOp op + consumedInOp (SizeOp op) = consumedInOp op + +instance (ASTRep rep, FreeIn op) => FreeIn (HostOp rep op) where + freeIn' (SegOp op) = freeIn' op + freeIn' (OtherOp op) = freeIn' op + freeIn' (SizeOp op) = freeIn' op + +instance (CanBeAliased (Op rep), CanBeAliased op, ASTRep rep) => CanBeAliased (HostOp rep op) where + type OpWithAliases (HostOp rep op) = HostOp (Aliases rep) (OpWithAliases op) + + addOpAliases aliases (SegOp op) = SegOp $ addOpAliases aliases op + addOpAliases aliases (OtherOp op) = OtherOp $ addOpAliases aliases op + addOpAliases _ (SizeOp op) = SizeOp op + + removeOpAliases (SegOp op) = SegOp $ removeOpAliases op + removeOpAliases (OtherOp op) = OtherOp $ removeOpAliases op + removeOpAliases (SizeOp op) = SizeOp op + +instance (CanBeWise (Op rep), CanBeWise op, ASTRep rep) => CanBeWise (HostOp rep op) where + type OpWithWisdom (HostOp rep op) = HostOp (Wise rep) (OpWithWisdom op) + + removeOpWisdom (SegOp op) = SegOp $ removeOpWisdom op + removeOpWisdom (OtherOp op) = OtherOp $ removeOpWisdom op + removeOpWisdom (SizeOp op) = SizeOp op + +instance (ASTRep rep, ST.IndexOp op) => ST.IndexOp (HostOp rep op) where + indexOp vtable k (SegOp op) is = ST.indexOp vtable k op is + indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is + indexOp _ _ _ _ = Nothing + +instance (PrettyRep rep, PP.Pretty op) => PP.Pretty (HostOp rep op) where + ppr (SegOp op) = ppr op + ppr (OtherOp op) = ppr op + ppr (SizeOp op) = ppr op + +instance (OpMetrics (Op rep), OpMetrics op) => OpMetrics (HostOp rep op) where + opMetrics (SegOp op) = opMetrics op + opMetrics (OtherOp op) = opMetrics op + opMetrics (SizeOp op) = opMetrics op + +checkSegLevel :: + TC.Checkable rep => + Maybe SegLevel -> + SegLevel -> + TC.TypeM rep () +checkSegLevel Nothing lvl = do + TC.require [Prim int64] $ unCount $ segNumGroups lvl + TC.require [Prim int64] $ unCount $ segGroupSize lvl +checkSegLevel (Just SegThread {}) _ = + TC.bad $ TC.TypeError "SegOps cannot occur when already at thread level." +checkSegLevel (Just x) y + | x == y = TC.bad $ TC.TypeError $ "Already at at level " ++ pretty x + | segNumGroups x /= segNumGroups y || segGroupSize x /= segGroupSize y = + TC.bad $ TC.TypeError "Physical layout for SegLevel does not match parent SegLevel." + | otherwise = + return () + +typeCheckHostOp :: + TC.Checkable rep => + (SegLevel -> OpWithAliases (Op rep) -> TC.TypeM rep ()) -> + Maybe SegLevel -> + (op -> TC.TypeM rep ()) -> + HostOp (Aliases rep) op -> + TC.TypeM rep () +typeCheckHostOp checker lvl _ (SegOp op) = + TC.checkOpWith (checker $ segLevel op) $ + typeCheckSegOp (checkSegLevel lvl) op +typeCheckHostOp _ _ f (OtherOp op) = f op +typeCheckHostOp _ _ _ (SizeOp op) = typeCheckSizeOp op diff --git a/src/Futhark/IR/GPU/Simplify.hs b/src/Futhark/IR/GPU/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/GPU/Simplify.hs @@ -0,0 +1,113 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Futhark.IR.GPU.Simplify + ( simplifyGPU, + simplifyLambda, + GPU, + + -- * Building blocks + simplifyKernelOp, + ) +where + +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.IR.GPU +import qualified Futhark.IR.SOACS.Simplify as SOAC +import Futhark.MonadFreshNames +import qualified Futhark.Optimise.Simplify as Simplify +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Rep +import Futhark.Optimise.Simplify.Rule +import Futhark.Optimise.Simplify.Rules +import Futhark.Pass +import Futhark.Tools +import qualified Futhark.Transform.FirstOrderTransform as FOT + +simpleGPU :: Simplify.SimpleOps GPU +simpleGPU = Simplify.bindableSimpleOps $ simplifyKernelOp SOAC.simplifySOAC + +simplifyGPU :: Prog GPU -> PassM (Prog GPU) +simplifyGPU = + Simplify.simplifyProg simpleGPU kernelRules Simplify.noExtraHoistBlockers + +simplifyLambda :: + (HasScope GPU m, MonadFreshNames m) => + Lambda GPU -> + m (Lambda GPU) +simplifyLambda = + Simplify.simplifyLambda simpleGPU kernelRules Engine.noExtraHoistBlockers + +simplifyKernelOp :: + ( Engine.SimplifiableRep rep, + BodyDec rep ~ () + ) => + Simplify.SimplifyOp rep op -> + HostOp rep op -> + Engine.SimpleM rep (HostOp (Wise rep) (OpWithWisdom op), Stms (Wise rep)) +simplifyKernelOp f (OtherOp op) = do + (op', stms) <- f op + return (OtherOp op', stms) +simplifyKernelOp _ (SegOp op) = do + (op', hoisted) <- simplifySegOp op + return (SegOp op', hoisted) +simplifyKernelOp _ (SizeOp (SplitSpace o w i elems_per_thread)) = + (,) + <$> ( SizeOp + <$> ( SplitSpace <$> Engine.simplify o <*> Engine.simplify w + <*> Engine.simplify i + <*> Engine.simplify elems_per_thread + ) + ) + <*> pure mempty +simplifyKernelOp _ (SizeOp (GetSize key size_class)) = + return (SizeOp $ GetSize key size_class, mempty) +simplifyKernelOp _ (SizeOp (GetSizeMax size_class)) = + return (SizeOp $ GetSizeMax size_class, mempty) +simplifyKernelOp _ (SizeOp (CmpSizeLe key size_class x)) = do + x' <- Engine.simplify x + return (SizeOp $ CmpSizeLe key size_class x', mempty) +simplifyKernelOp _ (SizeOp (CalcNumGroups w max_num_groups group_size)) = do + w' <- Engine.simplify w + return (SizeOp $ CalcNumGroups w' max_num_groups group_size, mempty) + +instance BinderOps (Wise GPU) + +instance HasSegOp (Wise GPU) where + type SegOpLevel (Wise GPU) = SegLevel + asSegOp (SegOp op) = Just op + asSegOp _ = Nothing + segOp = SegOp + +instance SOAC.HasSOAC (Wise GPU) where + asSOAC (OtherOp soac) = Just soac + asSOAC _ = Nothing + soacOp = OtherOp + +kernelRules :: RuleBook (Wise GPU) +kernelRules = + standardRules <> segOpRules + <> ruleBook + [ RuleOp redomapIotaToLoop, + RuleOp SOAC.simplifyKnownIterationSOAC, + RuleOp SOAC.removeReplicateMapping, + RuleOp SOAC.liftIdentityMapping + ] + [ RuleBasicOp removeUnnecessaryCopy + ] + +-- We turn reductions over (solely) iotas into do-loops, because there +-- is no useful structure here anyway. This is mostly a hack to work +-- around the fact that loop tiling would otherwise pointlessly tile +-- them. +redomapIotaToLoop :: TopDownRuleOp (Wise GPU) +redomapIotaToLoop vtable pat aux (OtherOp soac@(Screma _ [arr] form)) + | Just _ <- isRedomapSOAC form, + Just (Iota {}, _) <- ST.lookupBasicOp arr vtable = + Simplify $ certifying (stmAuxCerts aux) $ FOT.transformSOAC pat soac +redomapIotaToLoop _ _ _ _ = + Skip diff --git a/src/Futhark/IR/GPU/Sizes.hs b/src/Futhark/IR/GPU/Sizes.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/GPU/Sizes.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE Trustworthy #-} + +-- | In the context of this module, a "size" is any kind of tunable +-- (run-time) constant. +module Futhark.IR.GPU.Sizes + ( SizeClass (..), + sizeDefault, + KernelPath, + Count (..), + NumGroups, + GroupSize, + NumThreads, + ) +where + +import Data.Int (Int64) +import Data.Traversable +import Futhark.IR.Prop.Names (FreeIn) +import Futhark.Transform.Substitute +import Futhark.Util.IntegralExp (IntegralExp) +import Futhark.Util.Pretty +import Language.Futhark.Core (Name) +import Prelude hiding (id, (.)) + +-- | An indication of which comparisons have been performed to get to +-- this point, as well as the result of each comparison. +type KernelPath = [(Name, Bool)] + +-- | The class of some kind of configurable size. Each class may +-- impose constraints on the valid values. +data SizeClass + = -- | A threshold with an optional default. + SizeThreshold KernelPath (Maybe Int64) + | SizeGroup + | SizeNumGroups + | SizeTile + | SizeRegTile + | -- | Likely not useful on its own, but querying the + -- maximum can be handy. + SizeLocalMemory + | -- | A bespoke size with a default. + SizeBespoke Name Int64 + deriving (Eq, Ord, Show) + +instance Pretty SizeClass where + ppr (SizeThreshold path def) = + "threshold" <> parens (def' <> comma <+> spread (map pStep path)) + where + pStep (v, True) = ppr v + pStep (v, False) = "!" <> ppr v + def' = maybe "def" ppr def + ppr SizeGroup = text "group_size" + ppr SizeNumGroups = text "num_groups" + ppr SizeTile = text "tile_size" + ppr SizeRegTile = text "reg_tile_size" + ppr SizeLocalMemory = text "local_memory" + ppr (SizeBespoke k def) = + text "bespoke" <> parens (ppr k <> comma <+> ppr def) + +-- | The default value for the size. If 'Nothing', that means the backend gets to decide. +sizeDefault :: SizeClass -> Maybe Int64 +sizeDefault (SizeThreshold _ x) = x +sizeDefault (SizeBespoke _ x) = Just x +sizeDefault _ = Nothing + +-- | A wrapper supporting a phantom type for indicating what we are counting. +newtype Count u e = Count {unCount :: e} + deriving (Eq, Ord, Show, Num, IntegralExp, FreeIn, Pretty, Substitute) + +instance Functor (Count u) where + fmap = fmapDefault + +instance Foldable (Count u) where + foldMap = foldMapDefault + +instance Traversable (Count u) where + traverse f (Count x) = Count <$> f x + +-- | Phantom type for the number of groups of some kernel. +data NumGroups + +-- | Phantom type for the group size of some kernel. +data GroupSize + +-- | Phantom type for number of threads. +data NumThreads diff --git a/src/Futhark/IR/GPUMem.hs b/src/Futhark/IR/GPUMem.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/GPUMem.hs @@ -0,0 +1,111 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.IR.GPUMem + ( GPUMem, + + -- * Simplification + simplifyProg, + simplifyStms, + simpleGPUMem, + + -- * Module re-exports + module Futhark.IR.Mem, + module Futhark.IR.GPU.Kernel, + ) +where + +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.IR.GPU.Kernel +import Futhark.IR.GPU.Simplify (simplifyKernelOp) +import Futhark.IR.Mem +import Futhark.IR.Mem.Simplify +import Futhark.MonadFreshNames +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Pass +import Futhark.Pass.ExplicitAllocations (BinderOps (..), mkLetNamesB', mkLetNamesB'') +import qualified Futhark.TypeCheck as TC + +data GPUMem + +instance RepTypes GPUMem where + type LetDec GPUMem = LetDecMem + type FParamInfo GPUMem = FParamMem + type LParamInfo GPUMem = LParamMem + type RetType GPUMem = RetTypeMem + type BranchType GPUMem = BranchTypeMem + type Op GPUMem = MemOp (HostOp GPUMem ()) + +instance ASTRep GPUMem where + expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern + +instance OpReturns GPUMem where + opReturns (Alloc _ space) = + return [MemMem space] + opReturns (Inner (SegOp op)) = segOpReturns op + opReturns k = extReturns <$> opType k + +instance PrettyRep GPUMem + +instance TC.CheckableOp GPUMem where + checkOp = typeCheckMemoryOp Nothing + where + typeCheckMemoryOp _ (Alloc size _) = + TC.require [Prim int64] size + typeCheckMemoryOp lvl (Inner op) = + typeCheckHostOp (typeCheckMemoryOp . Just) lvl (const $ return ()) op + +instance TC.Checkable GPUMem where + checkFParamDec = checkMemInfo + checkLParamDec = checkMemInfo + checkLetBoundDec = checkMemInfo + checkRetType = mapM_ $ TC.checkExtType . declExtTypeOf + primFParam name t = return $ Param name (MemPrim t) + matchPattern = matchPatternToExp + matchReturnType = matchFunctionReturnType + matchBranchType = matchBranchReturnType + matchLoopResult = matchLoopResultMem + +instance BinderOps GPUMem where + mkExpDecB _ _ = return () + mkBodyB stms res = return $ Body () stms res + mkLetNamesB = mkLetNamesB' () + +instance BinderOps (Engine.Wise GPUMem) where + mkExpDecB pat e = return $ Engine.mkWiseExpDec pat () e + mkBodyB stms res = return $ Engine.mkWiseBody () stms res + mkLetNamesB = mkLetNamesB'' + +simplifyProg :: Prog GPUMem -> PassM (Prog GPUMem) +simplifyProg = simplifyProgGeneric simpleGPUMem + +simplifyStms :: + (HasScope GPUMem m, MonadFreshNames m) => + Stms GPUMem -> + m + ( Engine.SymbolTable (Engine.Wise GPUMem), + Stms GPUMem + ) +simplifyStms = simplifyStmsGeneric simpleGPUMem + +simpleGPUMem :: Engine.SimpleOps GPUMem +simpleGPUMem = + simpleGeneric usage $ simplifyKernelOp $ const $ return ((), mempty) + where + -- Slightly hackily, we look at the inside of SegGroup operations + -- to figure out the sizes of local memory allocations, and add + -- usages for those sizes. This is necessary so the simplifier + -- will hoist those sizes out as far as possible (most + -- importantly, past the versioning If). + usage (SegOp (SegMap SegGroup {} _ _ kbody)) = localAllocs kbody + usage _ = mempty + localAllocs = foldMap stmLocalAlloc . kernelBodyStms + stmLocalAlloc = expLocalAlloc . stmExp + expLocalAlloc (Op (Alloc (Var v) (Space "local"))) = + UT.sizeUsage v + expLocalAlloc _ = + mempty diff --git a/src/Futhark/IR/Kernels.hs b/src/Futhark/IR/Kernels.hs deleted file mode 100644 --- a/src/Futhark/IR/Kernels.hs +++ /dev/null @@ -1,62 +0,0 @@ -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} - --- | A representation with flat parallelism via GPU-oriented kernels. -module Futhark.IR.Kernels - ( -- * The Lore definition - Kernels, - - -- * Module re-exports - module Futhark.IR.Prop, - module Futhark.IR.Traversals, - module Futhark.IR.Pretty, - module Futhark.IR.Syntax, - module Futhark.IR.Kernels.Kernel, - module Futhark.IR.Kernels.Sizes, - module Futhark.IR.SOACS.SOAC, - ) -where - -import Futhark.Binder -import Futhark.Construct -import Futhark.IR.Kernels.Kernel -import Futhark.IR.Kernels.Sizes -import Futhark.IR.Pretty -import Futhark.IR.Prop -import Futhark.IR.SOACS.SOAC hiding (HistOp (..)) -import Futhark.IR.Syntax -import Futhark.IR.Traversals -import qualified Futhark.TypeCheck as TypeCheck - --- | The phantom data type for the kernels representation. -data Kernels - -instance Decorations Kernels where - type Op Kernels = HostOp Kernels (SOAC Kernels) - -instance ASTLore Kernels where - expTypesFromPattern = return . expExtTypesFromPattern - -instance TypeCheck.CheckableOp Kernels where - checkOp = typeCheckKernelsOp Nothing - where - typeCheckKernelsOp lvl = - typeCheckHostOp (typeCheckKernelsOp . Just) lvl typeCheckSOAC - -instance TypeCheck.Checkable Kernels - -instance Bindable Kernels where - mkBody = Body () - mkExpPat ctx val _ = basicPattern ctx val - mkExpDec _ _ = () - mkLetNames = simpleMkLetNames - -instance BinderOps Kernels - -instance PrettyLore Kernels - -instance HasSegOp Kernels where - type SegOpLevel Kernels = SegLevel - asSegOp (SegOp op) = Just op - asSegOp _ = Nothing - segOp = SegOp diff --git a/src/Futhark/IR/Kernels/Kernel.hs b/src/Futhark/IR/Kernels/Kernel.hs deleted file mode 100644 --- a/src/Futhark/IR/Kernels/Kernel.hs +++ /dev/null @@ -1,346 +0,0 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} - -module Futhark.IR.Kernels.Kernel - ( -- * Size operations - SizeOp (..), - - -- * Host operations - HostOp (..), - typeCheckHostOp, - - -- * SegOp refinements - SegLevel (..), - - -- * Reexports - module Futhark.IR.Kernels.Sizes, - module Futhark.IR.SegOp, - ) -where - -import Futhark.Analysis.Metrics -import qualified Futhark.Analysis.SymbolTable as ST -import Futhark.IR -import Futhark.IR.Aliases (Aliases) -import Futhark.IR.Kernels.Sizes -import Futhark.IR.Prop.Aliases -import Futhark.IR.SegOp -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore -import Futhark.Transform.Rename -import Futhark.Transform.Substitute -import qualified Futhark.TypeCheck as TC -import Futhark.Util.Pretty - ( commasep, - parens, - ppr, - text, - (<+>), - ) -import qualified Futhark.Util.Pretty as PP -import Prelude hiding (id, (.)) - --- | At which level the *body* of a t'SegOp' executes. -data SegLevel - = SegThread - { segNumGroups :: Count NumGroups SubExp, - segGroupSize :: Count GroupSize SubExp, - segVirt :: SegVirt - } - | SegGroup - { segNumGroups :: Count NumGroups SubExp, - segGroupSize :: Count GroupSize SubExp, - segVirt :: SegVirt - } - deriving (Eq, Ord, Show) - -instance PP.Pretty SegLevel where - ppr lvl = - PP.parens - ( lvl' <> PP.semi - <+> text "#groups=" <> ppr (segNumGroups lvl) <> PP.semi - <+> text "groupsize=" <> ppr (segGroupSize lvl) <> virt - ) - where - lvl' = case lvl of - SegThread {} -> "thread" - SegGroup {} -> "group" - virt = case segVirt lvl of - SegNoVirt -> mempty - SegNoVirtFull -> PP.semi <+> text "full" - SegVirt -> PP.semi <+> text "virtualise" - -instance Engine.Simplifiable SegLevel where - simplify (SegThread num_groups group_size virt) = - SegThread <$> traverse Engine.simplify num_groups - <*> traverse Engine.simplify group_size - <*> pure virt - simplify (SegGroup num_groups group_size virt) = - SegGroup <$> traverse Engine.simplify num_groups - <*> traverse Engine.simplify group_size - <*> pure virt - -instance Substitute SegLevel where - substituteNames substs (SegThread num_groups group_size virt) = - SegThread - (substituteNames substs num_groups) - (substituteNames substs group_size) - virt - substituteNames substs (SegGroup num_groups group_size virt) = - SegGroup - (substituteNames substs num_groups) - (substituteNames substs group_size) - virt - -instance Rename SegLevel where - rename = substituteRename - -instance FreeIn SegLevel where - freeIn' (SegThread num_groups group_size _) = - freeIn' num_groups <> freeIn' group_size - freeIn' (SegGroup num_groups group_size _) = - freeIn' num_groups <> freeIn' group_size - --- | A simple size-level query or computation. -data SizeOp - = -- | @SplitSpace o w i elems_per_thread@. - -- - -- Computes how to divide array elements to - -- threads in a kernel. Returns the number of - -- elements in the chunk that the current thread - -- should take. - -- - -- @w@ is the length of the outer dimension in - -- the array. @i@ is the current thread - -- index. Each thread takes at most - -- @elems_per_thread@ elements. - -- - -- If the order @o@ is 'SplitContiguous', thread with index @i@ - -- should receive elements - -- @i*elems_per_tread, i*elems_per_thread + 1, - -- ..., i*elems_per_thread + (elems_per_thread-1)@. - -- - -- If the order @o@ is @'SplitStrided' stride@, - -- the thread will receive elements @i, - -- i+stride, i+2*stride, ..., - -- i+(elems_per_thread-1)*stride@. - SplitSpace SplitOrdering SubExp SubExp SubExp - | -- | Produce some runtime-configurable size. - GetSize Name SizeClass - | -- | The maximum size of some class. - GetSizeMax SizeClass - | -- | Compare size (likely a threshold) with some integer value. - CmpSizeLe Name SizeClass SubExp - | -- | @CalcNumGroups w max_num_groups group_size@ calculates the - -- number of GPU workgroups to use for an input of the given size. - -- The @Name@ is a size name. Note that @w@ is an i64 to avoid - -- overflow issues. - CalcNumGroups SubExp Name SubExp - deriving (Eq, Ord, Show) - -instance Substitute SizeOp where - substituteNames subst (SplitSpace o w i elems_per_thread) = - SplitSpace - (substituteNames subst o) - (substituteNames subst w) - (substituteNames subst i) - (substituteNames subst elems_per_thread) - substituteNames substs (CmpSizeLe name sclass x) = - CmpSizeLe name sclass (substituteNames substs x) - substituteNames substs (CalcNumGroups w max_num_groups group_size) = - CalcNumGroups - (substituteNames substs w) - max_num_groups - (substituteNames substs group_size) - substituteNames _ op = op - -instance Rename SizeOp where - rename (SplitSpace o w i elems_per_thread) = - SplitSpace - <$> rename o - <*> rename w - <*> rename i - <*> rename elems_per_thread - rename (CmpSizeLe name sclass x) = - CmpSizeLe name sclass <$> rename x - rename (CalcNumGroups w max_num_groups group_size) = - CalcNumGroups <$> rename w <*> pure max_num_groups <*> rename group_size - rename x = pure x - -instance IsOp SizeOp where - safeOp _ = True - cheapOp _ = True - -instance TypedOp SizeOp where - opType SplitSpace {} = pure [Prim int64] - opType (GetSize _ _) = pure [Prim int64] - opType (GetSizeMax _) = pure [Prim int64] - opType CmpSizeLe {} = pure [Prim Bool] - opType CalcNumGroups {} = pure [Prim int64] - -instance AliasedOp SizeOp where - opAliases _ = [mempty] - consumedInOp _ = mempty - -instance FreeIn SizeOp where - freeIn' (SplitSpace o w i elems_per_thread) = - freeIn' o <> freeIn' [w, i, elems_per_thread] - freeIn' (CmpSizeLe _ _ x) = freeIn' x - freeIn' (CalcNumGroups w _ group_size) = freeIn' w <> freeIn' group_size - freeIn' _ = mempty - -instance PP.Pretty SizeOp where - ppr (SplitSpace SplitContiguous w i elems_per_thread) = - text "split_space" - <> parens (commasep [ppr w, ppr i, ppr elems_per_thread]) - ppr (SplitSpace (SplitStrided stride) w i elems_per_thread) = - text "split_space_strided" - <> parens (commasep [ppr stride, ppr w, ppr i, ppr elems_per_thread]) - ppr (GetSize name size_class) = - text "get_size" <> parens (commasep [ppr name, ppr size_class]) - ppr (GetSizeMax size_class) = - text "get_size_max" <> parens (commasep [ppr size_class]) - ppr (CmpSizeLe name size_class x) = - text "cmp_size" <> parens (commasep [ppr name, ppr size_class]) - <+> text "<=" - <+> ppr x - ppr (CalcNumGroups w max_num_groups group_size) = - text "calc_num_groups" <> parens (commasep [ppr w, ppr max_num_groups, ppr group_size]) - -instance OpMetrics SizeOp where - opMetrics SplitSpace {} = seen "SplitSpace" - opMetrics GetSize {} = seen "GetSize" - opMetrics GetSizeMax {} = seen "GetSizeMax" - opMetrics CmpSizeLe {} = seen "CmpSizeLe" - opMetrics CalcNumGroups {} = seen "CalcNumGroups" - -typeCheckSizeOp :: TC.Checkable lore => SizeOp -> TC.TypeM lore () -typeCheckSizeOp (SplitSpace o w i elems_per_thread) = do - case o of - SplitContiguous -> return () - SplitStrided stride -> TC.require [Prim int64] stride - mapM_ (TC.require [Prim int64]) [w, i, elems_per_thread] -typeCheckSizeOp GetSize {} = return () -typeCheckSizeOp GetSizeMax {} = return () -typeCheckSizeOp (CmpSizeLe _ _ x) = TC.require [Prim int64] x -typeCheckSizeOp (CalcNumGroups w _ group_size) = do - TC.require [Prim int64] w - TC.require [Prim int64] group_size - --- | A host-level operation; parameterised by what else it can do. -data HostOp lore op - = -- | A segmented operation. - SegOp (SegOp SegLevel lore) - | SizeOp SizeOp - | OtherOp op - deriving (Eq, Ord, Show) - -instance (ASTLore lore, Substitute op) => Substitute (HostOp lore op) where - substituteNames substs (SegOp op) = - SegOp $ substituteNames substs op - substituteNames substs (OtherOp op) = - OtherOp $ substituteNames substs op - substituteNames substs (SizeOp op) = - SizeOp $ substituteNames substs op - -instance (ASTLore lore, Rename op) => Rename (HostOp lore op) where - rename (SegOp op) = SegOp <$> rename op - rename (OtherOp op) = OtherOp <$> rename op - rename (SizeOp op) = SizeOp <$> rename op - -instance (ASTLore lore, IsOp op) => IsOp (HostOp lore op) where - safeOp (SegOp op) = safeOp op - safeOp (OtherOp op) = safeOp op - safeOp (SizeOp op) = safeOp op - - cheapOp (SegOp op) = cheapOp op - cheapOp (OtherOp op) = cheapOp op - cheapOp (SizeOp op) = cheapOp op - -instance TypedOp op => TypedOp (HostOp lore op) where - opType (SegOp op) = opType op - opType (OtherOp op) = opType op - opType (SizeOp op) = opType op - -instance (Aliased lore, AliasedOp op, ASTLore lore) => AliasedOp (HostOp lore op) where - opAliases (SegOp op) = opAliases op - opAliases (OtherOp op) = opAliases op - opAliases (SizeOp op) = opAliases op - - consumedInOp (SegOp op) = consumedInOp op - consumedInOp (OtherOp op) = consumedInOp op - consumedInOp (SizeOp op) = consumedInOp op - -instance (ASTLore lore, FreeIn op) => FreeIn (HostOp lore op) where - freeIn' (SegOp op) = freeIn' op - freeIn' (OtherOp op) = freeIn' op - freeIn' (SizeOp op) = freeIn' op - -instance (CanBeAliased (Op lore), CanBeAliased op, ASTLore lore) => CanBeAliased (HostOp lore op) where - type OpWithAliases (HostOp lore op) = HostOp (Aliases lore) (OpWithAliases op) - - addOpAliases aliases (SegOp op) = SegOp $ addOpAliases aliases op - addOpAliases aliases (OtherOp op) = OtherOp $ addOpAliases aliases op - addOpAliases _ (SizeOp op) = SizeOp op - - removeOpAliases (SegOp op) = SegOp $ removeOpAliases op - removeOpAliases (OtherOp op) = OtherOp $ removeOpAliases op - removeOpAliases (SizeOp op) = SizeOp op - -instance (CanBeWise (Op lore), CanBeWise op, ASTLore lore) => CanBeWise (HostOp lore op) where - type OpWithWisdom (HostOp lore op) = HostOp (Wise lore) (OpWithWisdom op) - - removeOpWisdom (SegOp op) = SegOp $ removeOpWisdom op - removeOpWisdom (OtherOp op) = OtherOp $ removeOpWisdom op - removeOpWisdom (SizeOp op) = SizeOp op - -instance (ASTLore lore, ST.IndexOp op) => ST.IndexOp (HostOp lore op) where - indexOp vtable k (SegOp op) is = ST.indexOp vtable k op is - indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is - indexOp _ _ _ _ = Nothing - -instance (PrettyLore lore, PP.Pretty op) => PP.Pretty (HostOp lore op) where - ppr (SegOp op) = ppr op - ppr (OtherOp op) = ppr op - ppr (SizeOp op) = ppr op - -instance (OpMetrics (Op lore), OpMetrics op) => OpMetrics (HostOp lore op) where - opMetrics (SegOp op) = opMetrics op - opMetrics (OtherOp op) = opMetrics op - opMetrics (SizeOp op) = opMetrics op - -checkSegLevel :: - TC.Checkable lore => - Maybe SegLevel -> - SegLevel -> - TC.TypeM lore () -checkSegLevel Nothing lvl = do - TC.require [Prim int64] $ unCount $ segNumGroups lvl - TC.require [Prim int64] $ unCount $ segGroupSize lvl -checkSegLevel (Just SegThread {}) _ = - TC.bad $ TC.TypeError "SegOps cannot occur when already at thread level." -checkSegLevel (Just x) y - | x == y = TC.bad $ TC.TypeError $ "Already at at level " ++ pretty x - | segNumGroups x /= segNumGroups y || segGroupSize x /= segGroupSize y = - TC.bad $ TC.TypeError "Physical layout for SegLevel does not match parent SegLevel." - | otherwise = - return () - -typeCheckHostOp :: - TC.Checkable lore => - (SegLevel -> OpWithAliases (Op lore) -> TC.TypeM lore ()) -> - Maybe SegLevel -> - (op -> TC.TypeM lore ()) -> - HostOp (Aliases lore) op -> - TC.TypeM lore () -typeCheckHostOp checker lvl _ (SegOp op) = - TC.checkOpWith (checker $ segLevel op) $ - typeCheckSegOp (checkSegLevel lvl) op -typeCheckHostOp _ _ f (OtherOp op) = f op -typeCheckHostOp _ _ _ (SizeOp op) = typeCheckSizeOp op diff --git a/src/Futhark/IR/Kernels/Simplify.hs b/src/Futhark/IR/Kernels/Simplify.hs deleted file mode 100644 --- a/src/Futhark/IR/Kernels/Simplify.hs +++ /dev/null @@ -1,113 +0,0 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} - -module Futhark.IR.Kernels.Simplify - ( simplifyKernels, - simplifyLambda, - Kernels, - - -- * Building blocks - simplifyKernelOp, - ) -where - -import qualified Futhark.Analysis.SymbolTable as ST -import Futhark.IR.Kernels -import qualified Futhark.IR.SOACS.Simplify as SOAC -import Futhark.MonadFreshNames -import qualified Futhark.Optimise.Simplify as Simplify -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore -import Futhark.Optimise.Simplify.Rule -import Futhark.Optimise.Simplify.Rules -import Futhark.Pass -import Futhark.Tools -import qualified Futhark.Transform.FirstOrderTransform as FOT - -simpleKernels :: Simplify.SimpleOps Kernels -simpleKernels = Simplify.bindableSimpleOps $ simplifyKernelOp SOAC.simplifySOAC - -simplifyKernels :: Prog Kernels -> PassM (Prog Kernels) -simplifyKernels = - Simplify.simplifyProg simpleKernels kernelRules Simplify.noExtraHoistBlockers - -simplifyLambda :: - (HasScope Kernels m, MonadFreshNames m) => - Lambda Kernels -> - m (Lambda Kernels) -simplifyLambda = - Simplify.simplifyLambda simpleKernels kernelRules Engine.noExtraHoistBlockers - -simplifyKernelOp :: - ( Engine.SimplifiableLore lore, - BodyDec lore ~ () - ) => - Simplify.SimplifyOp lore op -> - HostOp lore op -> - Engine.SimpleM lore (HostOp (Wise lore) (OpWithWisdom op), Stms (Wise lore)) -simplifyKernelOp f (OtherOp op) = do - (op', stms) <- f op - return (OtherOp op', stms) -simplifyKernelOp _ (SegOp op) = do - (op', hoisted) <- simplifySegOp op - return (SegOp op', hoisted) -simplifyKernelOp _ (SizeOp (SplitSpace o w i elems_per_thread)) = - (,) - <$> ( SizeOp - <$> ( SplitSpace <$> Engine.simplify o <*> Engine.simplify w - <*> Engine.simplify i - <*> Engine.simplify elems_per_thread - ) - ) - <*> pure mempty -simplifyKernelOp _ (SizeOp (GetSize key size_class)) = - return (SizeOp $ GetSize key size_class, mempty) -simplifyKernelOp _ (SizeOp (GetSizeMax size_class)) = - return (SizeOp $ GetSizeMax size_class, mempty) -simplifyKernelOp _ (SizeOp (CmpSizeLe key size_class x)) = do - x' <- Engine.simplify x - return (SizeOp $ CmpSizeLe key size_class x', mempty) -simplifyKernelOp _ (SizeOp (CalcNumGroups w max_num_groups group_size)) = do - w' <- Engine.simplify w - return (SizeOp $ CalcNumGroups w' max_num_groups group_size, mempty) - -instance BinderOps (Wise Kernels) - -instance HasSegOp (Wise Kernels) where - type SegOpLevel (Wise Kernels) = SegLevel - asSegOp (SegOp op) = Just op - asSegOp _ = Nothing - segOp = SegOp - -instance SOAC.HasSOAC (Wise Kernels) where - asSOAC (OtherOp soac) = Just soac - asSOAC _ = Nothing - soacOp = OtherOp - -kernelRules :: RuleBook (Wise Kernels) -kernelRules = - standardRules <> segOpRules - <> ruleBook - [ RuleOp redomapIotaToLoop, - RuleOp SOAC.simplifyKnownIterationSOAC, - RuleOp SOAC.removeReplicateMapping, - RuleOp SOAC.liftIdentityMapping - ] - [ RuleBasicOp removeUnnecessaryCopy - ] - --- We turn reductions over (solely) iotas into do-loops, because there --- is no useful structure here anyway. This is mostly a hack to work --- around the fact that loop tiling would otherwise pointlessly tile --- them. -redomapIotaToLoop :: TopDownRuleOp (Wise Kernels) -redomapIotaToLoop vtable pat aux (OtherOp soac@(Screma _ [arr] form)) - | Just _ <- isRedomapSOAC form, - Just (Iota {}, _) <- ST.lookupBasicOp arr vtable = - Simplify $ certifying (stmAuxCerts aux) $ FOT.transformSOAC pat soac -redomapIotaToLoop _ _ _ _ = - Skip diff --git a/src/Futhark/IR/Kernels/Sizes.hs b/src/Futhark/IR/Kernels/Sizes.hs deleted file mode 100644 --- a/src/Futhark/IR/Kernels/Sizes.hs +++ /dev/null @@ -1,88 +0,0 @@ -{-# LANGUAGE GeneralizedNewtypeDeriving #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE Trustworthy #-} - --- | In the context of this module, a "size" is any kind of tunable --- (run-time) constant. -module Futhark.IR.Kernels.Sizes - ( SizeClass (..), - sizeDefault, - KernelPath, - Count (..), - NumGroups, - GroupSize, - NumThreads, - ) -where - -import Data.Int (Int64) -import Data.Traversable -import Futhark.IR.Prop.Names (FreeIn) -import Futhark.Transform.Substitute -import Futhark.Util.IntegralExp (IntegralExp) -import Futhark.Util.Pretty -import Language.Futhark.Core (Name) -import Prelude hiding (id, (.)) - --- | An indication of which comparisons have been performed to get to --- this point, as well as the result of each comparison. -type KernelPath = [(Name, Bool)] - --- | The class of some kind of configurable size. Each class may --- impose constraints on the valid values. -data SizeClass - = -- | A threshold with an optional default. - SizeThreshold KernelPath (Maybe Int64) - | SizeGroup - | SizeNumGroups - | SizeTile - | SizeRegTile - | -- | Likely not useful on its own, but querying the - -- maximum can be handy. - SizeLocalMemory - | -- | A bespoke size with a default. - SizeBespoke Name Int64 - deriving (Eq, Ord, Show) - -instance Pretty SizeClass where - ppr (SizeThreshold path def) = - "threshold" <> parens (def' <> comma <+> spread (map pStep path)) - where - pStep (v, True) = ppr v - pStep (v, False) = "!" <> ppr v - def' = maybe "def" ppr def - ppr SizeGroup = text "group_size" - ppr SizeNumGroups = text "num_groups" - ppr SizeTile = text "tile_size" - ppr SizeRegTile = text "reg_tile_size" - ppr SizeLocalMemory = text "local_memory" - ppr (SizeBespoke k def) = - text "bespoke" <> parens (ppr k <> comma <+> ppr def) - --- | The default value for the size. If 'Nothing', that means the backend gets to decide. -sizeDefault :: SizeClass -> Maybe Int64 -sizeDefault (SizeThreshold _ x) = x -sizeDefault (SizeBespoke _ x) = Just x -sizeDefault _ = Nothing - --- | A wrapper supporting a phantom type for indicating what we are counting. -newtype Count u e = Count {unCount :: e} - deriving (Eq, Ord, Show, Num, IntegralExp, FreeIn, Pretty, Substitute) - -instance Functor (Count u) where - fmap = fmapDefault - -instance Foldable (Count u) where - foldMap = foldMapDefault - -instance Traversable (Count u) where - traverse f (Count x) = Count <$> f x - --- | Phantom type for the number of groups of some kernel. -data NumGroups - --- | Phantom type for the group size of some kernel. -data GroupSize - --- | Phantom type for number of threads. -data NumThreads diff --git a/src/Futhark/IR/KernelsMem.hs b/src/Futhark/IR/KernelsMem.hs deleted file mode 100644 --- a/src/Futhark/IR/KernelsMem.hs +++ /dev/null @@ -1,111 +0,0 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE TypeFamilies #-} - -module Futhark.IR.KernelsMem - ( KernelsMem, - - -- * Simplification - simplifyProg, - simplifyStms, - simpleKernelsMem, - - -- * Module re-exports - module Futhark.IR.Mem, - module Futhark.IR.Kernels.Kernel, - ) -where - -import Futhark.Analysis.PrimExp.Convert -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.IR.Kernels.Kernel -import Futhark.IR.Kernels.Simplify (simplifyKernelOp) -import Futhark.IR.Mem -import Futhark.IR.Mem.Simplify -import Futhark.MonadFreshNames -import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Pass -import Futhark.Pass.ExplicitAllocations (BinderOps (..), mkLetNamesB', mkLetNamesB'') -import qualified Futhark.TypeCheck as TC - -data KernelsMem - -instance Decorations KernelsMem where - type LetDec KernelsMem = LetDecMem - type FParamInfo KernelsMem = FParamMem - type LParamInfo KernelsMem = LParamMem - type RetType KernelsMem = RetTypeMem - type BranchType KernelsMem = BranchTypeMem - type Op KernelsMem = MemOp (HostOp KernelsMem ()) - -instance ASTLore KernelsMem where - expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern - -instance OpReturns KernelsMem where - opReturns (Alloc _ space) = - return [MemMem space] - opReturns (Inner (SegOp op)) = segOpReturns op - opReturns k = extReturns <$> opType k - -instance PrettyLore KernelsMem - -instance TC.CheckableOp KernelsMem where - checkOp = typeCheckMemoryOp Nothing - where - typeCheckMemoryOp _ (Alloc size _) = - TC.require [Prim int64] size - typeCheckMemoryOp lvl (Inner op) = - typeCheckHostOp (typeCheckMemoryOp . Just) lvl (const $ return ()) op - -instance TC.Checkable KernelsMem where - checkFParamLore = checkMemInfo - checkLParamLore = checkMemInfo - checkLetBoundLore = checkMemInfo - checkRetType = mapM_ $ TC.checkExtType . declExtTypeOf - primFParam name t = return $ Param name (MemPrim t) - matchPattern = matchPatternToExp - matchReturnType = matchFunctionReturnType - matchBranchType = matchBranchReturnType - matchLoopResult = matchLoopResultMem - -instance BinderOps KernelsMem where - mkExpDecB _ _ = return () - mkBodyB stms res = return $ Body () stms res - mkLetNamesB = mkLetNamesB' () - -instance BinderOps (Engine.Wise KernelsMem) where - mkExpDecB pat e = return $ Engine.mkWiseExpDec pat () e - mkBodyB stms res = return $ Engine.mkWiseBody () stms res - mkLetNamesB = mkLetNamesB'' - -simplifyProg :: Prog KernelsMem -> PassM (Prog KernelsMem) -simplifyProg = simplifyProgGeneric simpleKernelsMem - -simplifyStms :: - (HasScope KernelsMem m, MonadFreshNames m) => - Stms KernelsMem -> - m - ( Engine.SymbolTable (Engine.Wise KernelsMem), - Stms KernelsMem - ) -simplifyStms = simplifyStmsGeneric simpleKernelsMem - -simpleKernelsMem :: Engine.SimpleOps KernelsMem -simpleKernelsMem = - simpleGeneric usage $ simplifyKernelOp $ const $ return ((), mempty) - where - -- Slightly hackily, we look at the inside of SegGroup operations - -- to figure out the sizes of local memory allocations, and add - -- usages for those sizes. This is necessary so the simplifier - -- will hoist those sizes out as far as possible (most - -- importantly, past the versioning If). - usage (SegOp (SegMap SegGroup {} _ _ kbody)) = localAllocs kbody - usage _ = mempty - localAllocs = foldMap stmLocalAlloc . kernelBodyStms - stmLocalAlloc = expLocalAlloc . stmExp - expLocalAlloc (Op (Alloc (Var v) (Space "local"))) = - UT.sizeUsage v - expLocalAlloc _ = - mempty diff --git a/src/Futhark/IR/MC.hs b/src/Futhark/IR/MC.hs --- a/src/Futhark/IR/MC.hs +++ b/src/Futhark/IR/MC.hs @@ -4,8 +4,7 @@ -- | A representation for multicore CPU parallelism. module Futhark.IR.MC - ( -- * The Lore definition - MC, + ( MC, -- * Simplification simplifyProg, @@ -39,10 +38,10 @@ data MC -instance Decorations MC where +instance RepTypes MC where type Op MC = MCOp MC (SOAC MC) -instance ASTLore MC where +instance ASTRep MC where expTypesFromPattern = return . expExtTypesFromPattern instance TypeCheck.CheckableOp MC where @@ -60,7 +59,7 @@ instance BinderOps (Engine.Wise MC) -instance PrettyLore MC +instance PrettyRep MC simpleMC :: Simplify.SimpleOps MC simpleMC = Simplify.bindableSimpleOps $ simplifyMCOp SOAC.simplifySOAC diff --git a/src/Futhark/IR/MC/Op.hs b/src/Futhark/IR/MC/Op.hs --- a/src/Futhark/IR/MC/Op.hs +++ b/src/Futhark/IR/MC/Op.hs @@ -24,7 +24,7 @@ import Futhark.IR.SegOp import qualified Futhark.Optimise.Simplify as Simplify import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Transform.Rename import Futhark.Transform.Substitute import qualified Futhark.TypeCheck as TC @@ -40,45 +40,45 @@ -- | An operation for the multicore representation. Feel free to -- extend this on an ad hoc basis as needed. Parameterised with some -- other operation. -data MCOp lore op +data MCOp rep op = -- | The first 'SegOp' (if it exists) contains nested parallelism, -- while the second one has a fully sequential body. They are -- semantically fully equivalent. ParOp - (Maybe (SegOp () lore)) - (SegOp () lore) + (Maybe (SegOp () rep)) + (SegOp () rep) | -- | Something else (in practice often a SOAC). OtherOp op deriving (Eq, Ord, Show) -instance (ASTLore lore, Substitute op) => Substitute (MCOp lore op) where +instance (ASTRep rep, Substitute op) => Substitute (MCOp rep op) where substituteNames substs (ParOp par_op op) = ParOp (substituteNames substs <$> par_op) (substituteNames substs op) substituteNames substs (OtherOp op) = OtherOp $ substituteNames substs op -instance (ASTLore lore, Rename op) => Rename (MCOp lore op) where +instance (ASTRep rep, Rename op) => Rename (MCOp rep op) where rename (ParOp par_op op) = ParOp <$> rename par_op <*> rename op rename (OtherOp op) = OtherOp <$> rename op -instance (ASTLore lore, FreeIn op) => FreeIn (MCOp lore op) where +instance (ASTRep rep, FreeIn op) => FreeIn (MCOp rep op) where freeIn' (ParOp par_op op) = freeIn' par_op <> freeIn' op freeIn' (OtherOp op) = freeIn' op -instance (ASTLore lore, IsOp op) => IsOp (MCOp lore op) where +instance (ASTRep rep, IsOp op) => IsOp (MCOp rep op) where safeOp (ParOp _ op) = safeOp op safeOp (OtherOp op) = safeOp op cheapOp (ParOp _ op) = cheapOp op cheapOp (OtherOp op) = cheapOp op -instance TypedOp op => TypedOp (MCOp lore op) where +instance TypedOp op => TypedOp (MCOp rep op) where opType (ParOp _ op) = opType op opType (OtherOp op) = opType op instance - (Aliased lore, AliasedOp op, ASTLore lore) => - AliasedOp (MCOp lore op) + (Aliased rep, AliasedOp op, ASTRep rep) => + AliasedOp (MCOp rep op) where opAliases (ParOp _ op) = opAliases op opAliases (OtherOp op) = opAliases op @@ -87,10 +87,10 @@ consumedInOp (OtherOp op) = consumedInOp op instance - (CanBeAliased (Op lore), CanBeAliased op, ASTLore lore) => - CanBeAliased (MCOp lore op) + (CanBeAliased (Op rep), CanBeAliased op, ASTRep rep) => + CanBeAliased (MCOp rep op) where - type OpWithAliases (MCOp lore op) = MCOp (Aliases lore) (OpWithAliases op) + type OpWithAliases (MCOp rep op) = MCOp (Aliases rep) (OpWithAliases op) addOpAliases aliases (ParOp par_op op) = ParOp (addOpAliases aliases <$> par_op) (addOpAliases aliases op) @@ -103,36 +103,36 @@ OtherOp $ removeOpAliases op instance - (CanBeWise (Op lore), CanBeWise op, ASTLore lore) => - CanBeWise (MCOp lore op) + (CanBeWise (Op rep), CanBeWise op, ASTRep rep) => + CanBeWise (MCOp rep op) where - type OpWithWisdom (MCOp lore op) = MCOp (Wise lore) (OpWithWisdom op) + type OpWithWisdom (MCOp rep op) = MCOp (Wise rep) (OpWithWisdom op) removeOpWisdom (ParOp par_op op) = ParOp (removeOpWisdom <$> par_op) (removeOpWisdom op) removeOpWisdom (OtherOp op) = OtherOp $ removeOpWisdom op -instance (ASTLore lore, ST.IndexOp op) => ST.IndexOp (MCOp lore op) where +instance (ASTRep rep, ST.IndexOp op) => ST.IndexOp (MCOp rep op) where indexOp vtable k (ParOp _ op) is = ST.indexOp vtable k op is indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is -instance (PrettyLore lore, Pretty op) => Pretty (MCOp lore op) where +instance (PrettyRep rep, Pretty op) => Pretty (MCOp rep op) where ppr (ParOp Nothing op) = ppr op ppr (ParOp (Just par_op) op) = "par" <+> nestedBlock "{" "}" (ppr par_op) "seq" <+> nestedBlock "{" "}" (ppr op) ppr (OtherOp op) = ppr op -instance (OpMetrics (Op lore), OpMetrics op) => OpMetrics (MCOp lore op) where +instance (OpMetrics (Op rep), OpMetrics op) => OpMetrics (MCOp rep op) where opMetrics (ParOp par_op op) = opMetrics par_op >> opMetrics op opMetrics (OtherOp op) = opMetrics op typeCheckMCOp :: - TC.Checkable lore => - (op -> TC.TypeM lore ()) -> - MCOp (Aliases lore) op -> - TC.TypeM lore () + TC.Checkable rep => + (op -> TC.TypeM rep ()) -> + MCOp (Aliases rep) op -> + TC.TypeM rep () typeCheckMCOp _ (ParOp (Just par_op) op) = do -- It is valid for the same array to be consumed in both par_op and op. _ <- typeCheckSegOp return par_op `TC.alternative` typeCheckSegOp return op @@ -142,12 +142,12 @@ typeCheckMCOp f (OtherOp op) = f op simplifyMCOp :: - ( Engine.SimplifiableLore lore, - BodyDec lore ~ () + ( Engine.SimplifiableRep rep, + BodyDec rep ~ () ) => - Simplify.SimplifyOp lore op -> - MCOp lore op -> - Engine.SimpleM lore (MCOp (Wise lore) (OpWithWisdom op), Stms (Wise lore)) + Simplify.SimplifyOp rep op -> + MCOp rep op -> + Engine.SimpleM rep (MCOp (Wise rep) (OpWithWisdom op), Stms (Wise rep)) simplifyMCOp f (OtherOp op) = do (op', stms) <- f op return (OtherOp op', stms) diff --git a/src/Futhark/IR/MCMem.hs b/src/Futhark/IR/MCMem.hs --- a/src/Futhark/IR/MCMem.hs +++ b/src/Futhark/IR/MCMem.hs @@ -29,7 +29,7 @@ data MCMem -instance Decorations MCMem where +instance RepTypes MCMem where type LetDec MCMem = LetDecMem type FParamInfo MCMem = FParamMem type LParamInfo MCMem = LParamMem @@ -37,7 +37,7 @@ type BranchType MCMem = BranchTypeMem type Op MCMem = MemOp (MCOp MCMem ()) -instance ASTLore MCMem where +instance ASTRep MCMem where expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern instance OpReturns MCMem where @@ -45,7 +45,7 @@ opReturns (Inner (ParOp _ op)) = segOpReturns op opReturns (Inner (OtherOp ())) = pure [] -instance PrettyLore MCMem +instance PrettyRep MCMem instance TC.CheckableOp MCMem where checkOp = typeCheckMemoryOp @@ -56,9 +56,9 @@ typeCheckMCOp pure op instance TC.Checkable MCMem where - checkFParamLore = checkMemInfo - checkLParamLore = checkMemInfo - checkLetBoundLore = checkMemInfo + checkFParamDec = checkMemInfo + checkLParamDec = checkMemInfo + checkLetBoundDec = checkMemInfo checkRetType = mapM_ (TC.checkExtType . declExtTypeOf) primFParam name t = return $ Param name (MemPrim t) matchPattern = matchPatternToExp diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs --- a/src/Futhark/IR/Mem.hs +++ b/src/Futhark/IR/Mem.hs @@ -128,7 +128,7 @@ import Futhark.IR.Syntax import Futhark.IR.Traversals import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Transform.Rename import Futhark.Transform.Substitute import qualified Futhark.TypeCheck as TC @@ -151,16 +151,15 @@ class AllocOp op where allocOp :: SubExp -> Space -> op -type Mem lore = - ( AllocOp (Op lore), - FParamInfo lore ~ FParamMem, - LParamInfo lore ~ LParamMem, - LetDec lore ~ LetDecMem, - RetType lore ~ RetTypeMem, - BranchType lore ~ BranchTypeMem, - ASTLore lore, - Decorations lore, - OpReturns lore +type Mem rep = + ( AllocOp (Op rep), + FParamInfo rep ~ FParamMem, + LParamInfo rep ~ LParamMem, + LetDec rep ~ LetDecMem, + RetType rep ~ RetTypeMem, + BranchType rep ~ BranchTypeMem, + ASTRep rep, + OpReturns rep ) instance IsRetType FunReturns where @@ -324,15 +323,15 @@ rename = substituteRename simplifyIxFun :: - Engine.SimplifiableLore lore => + Engine.SimplifiableRep rep => IxFun -> - Engine.SimpleM lore IxFun + Engine.SimpleM rep IxFun simplifyIxFun = traverse $ fmap isInt64 . simplifyPrimExp . untyped simplifyExtIxFun :: - Engine.SimplifiableLore lore => + Engine.SimplifiableRep rep => ExtIxFun -> - Engine.SimpleM lore ExtIxFun + Engine.SimpleM rep ExtIxFun simplifyExtIxFun = traverse $ fmap isInt64 . simplifyExtPrimExp . untyped isStaticIxFun :: ExtIxFun -> Maybe IxFun @@ -533,20 +532,20 @@ bodyReturnsToExpReturns = noUniquenessReturns . maybeReturns matchRetTypeToResult :: - (Mem lore, TC.Checkable lore) => + (Mem rep, TC.Checkable rep) => [FunReturns] -> Result -> - TC.TypeM lore () + TC.TypeM rep () matchRetTypeToResult rettype result = do scope <- askScope result_ts <- runReaderT (mapM subExpMemInfo result) $ removeScopeAliases scope matchReturnType rettype result result_ts matchFunctionReturnType :: - (Mem lore, TC.Checkable lore) => + (Mem rep, TC.Checkable rep) => [FunReturns] -> Result -> - TC.TypeM lore () + TC.TypeM rep () matchFunctionReturnType rettype result = do matchRetTypeToResult rettype result mapM_ checkResultSubExp result @@ -570,11 +569,11 @@ ++ pretty ixfun matchLoopResultMem :: - (Mem lore, TC.Checkable lore) => - [FParam (Aliases lore)] -> - [FParam (Aliases lore)] -> + (Mem rep, TC.Checkable rep) => + [FParam (Aliases rep)] -> + [FParam (Aliases rep)] -> [SubExp] -> - TC.TypeM lore () + TC.TypeM rep () matchLoopResultMem ctx val = matchRetTypeToResult rettype where ctx_names = map paramName ctx @@ -607,10 +606,10 @@ ixfun' = existentialiseIxFun ctx_names ixfun matchBranchReturnType :: - (Mem lore, TC.Checkable lore) => + (Mem rep, TC.Checkable rep) => [BodyReturns] -> - Body (Aliases lore) -> - TC.TypeM lore () + Body (Aliases rep) -> + TC.TypeM rep () matchBranchReturnType rettype (Body _ stms res) = do scope <- askScope ts <- runReaderT (mapM subExpMemInfo res) $ removeScopeAliases (scope <> scopeOf stms) @@ -649,7 +648,7 @@ [MemInfo ExtSize u MemReturn] -> [SubExp] -> [MemInfo SubExp NoUniqueness MemBind] -> - TC.TypeM lore () + TC.TypeM rep () matchReturnType rettype res ts = do let (ctx_ts, val_ts) = splitFromEnd (length rettype) ts (ctx_res, _val_res) = splitFromEnd (length rettype) res @@ -769,7 +768,7 @@ pretty y ] - bad :: String -> TC.TypeM lore a + bad :: String -> TC.TypeM rep a bad s = TC.bad $ TC.TypeError $ @@ -792,10 +791,10 @@ either bad return =<< runExceptT (zipWithM_ checkReturn rettype val_ts) matchPatternToExp :: - (Mem lore, TC.Checkable lore) => - Pattern (Aliases lore) -> - Exp (Aliases lore) -> - TC.TypeM lore () + (Mem rep, TC.Checkable rep) => + Pattern (Aliases rep) -> + Exp (Aliases rep) -> + TC.TypeM rep () matchPatternToExp pat e = do scope <- asksScope removeScopeAliases rt <- runReaderT (expReturns $ removeExpAliases e) scope @@ -866,9 +865,9 @@ extInIxFn ixfun = S.fromList $ concatMap (mapMaybe isExt . toList) ixfun varMemInfo :: - Mem lore => + Mem rep => VName -> - TC.TypeM lore (MemInfo SubExp NoUniqueness MemBind) + TC.TypeM rep (MemInfo SubExp NoUniqueness MemBind) varMemInfo name = do dec <- TC.lookupVar name @@ -878,7 +877,7 @@ LParamName summary -> return summary IndexName it -> return $ MemPrim $ IntType it -nameInfoToMemInfo :: Mem lore => NameInfo lore -> MemBound NoUniqueness +nameInfoToMemInfo :: Mem rep => NameInfo rep -> MemBound NoUniqueness nameInfoToMemInfo info = case info of FParamName summary -> noUniquenessReturns summary @@ -887,20 +886,20 @@ IndexName it -> MemPrim $ IntType it lookupMemInfo :: - (HasScope lore m, Mem lore) => + (HasScope rep m, Mem rep) => VName -> m (MemInfo SubExp NoUniqueness MemBind) lookupMemInfo = fmap nameInfoToMemInfo . lookupInfo subExpMemInfo :: - (HasScope lore m, Monad m, Mem lore) => + (HasScope rep m, Monad m, Mem rep) => SubExp -> m (MemInfo SubExp NoUniqueness MemBind) subExpMemInfo (Var v) = lookupMemInfo v subExpMemInfo (Constant v) = return $ MemPrim $ primValueType v lookupArraySummary :: - (Mem lore, HasScope lore m, Monad m) => + (Mem rep, HasScope rep m, Monad m) => VName -> m (VName, IxFun.IxFun (TPrimExp Int64 VName)) lookupArraySummary name = do @@ -912,10 +911,10 @@ error $ "Variable " ++ pretty name ++ " does not look like an array." checkMemInfo :: - TC.Checkable lore => + TC.Checkable rep => VName -> MemInfo SubExp u MemBind -> - TC.TypeM lore () + TC.TypeM rep () checkMemInfo _ (MemPrim _) = return () checkMemInfo _ (MemMem (ScalarSpace d _)) = mapM_ (TC.require [Prim int64]) d checkMemInfo _ (MemMem _) = return () @@ -1002,7 +1001,7 @@ convert (Free v) = Free <$> pe64 v arrayVarReturns :: - (HasScope lore m, Monad m, Mem lore) => + (HasScope rep m, Monad m, Mem rep) => VName -> m (PrimType, Shape, VName, IxFun) arrayVarReturns v = do @@ -1014,7 +1013,7 @@ error $ "arrayVarReturns: " ++ pretty v ++ " is not an array." varReturns :: - (HasScope lore m, Monad m, Mem lore) => + (HasScope rep m, Monad m, Mem rep) => VName -> m ExpReturns varReturns v = do @@ -1031,7 +1030,7 @@ MemAcc acc ispace ts u -> return $ MemAcc acc ispace ts u -subExpReturns :: (HasScope lore m, Monad m, Mem lore) => SubExp -> m ExpReturns +subExpReturns :: (HasScope rep m, Monad m, Mem rep) => SubExp -> m ExpReturns subExpReturns (Var v) = varReturns v subExpReturns (Constant v) = @@ -1041,10 +1040,10 @@ -- "return type with memory annotations" of the expression. expReturns :: ( Monad m, - LocalScope lore m, - Mem lore + LocalScope rep m, + Mem rep ) => - Exp lore -> + Exp rep -> m [ExpReturns] expReturns (BasicOp (SubExp se)) = pure <$> subExpReturns se @@ -1137,7 +1136,7 @@ num_accs = length inputs sliceInfo :: - (Monad m, HasScope lore m, Mem lore) => + (Monad m, HasScope rep m, Mem rep) => VName -> Slice SubExp -> m (MemInfo SubExp NoUniqueness MemBind) @@ -1153,10 +1152,10 @@ ixfun (map (fmap (isInt64 . primExpFromSubExp int64)) slice) -class TypedOp (Op lore) => OpReturns lore where +class TypedOp (Op rep) => OpReturns rep where opReturns :: - (Monad m, HasScope lore m) => - Op lore -> + (Monad m, HasScope rep m) => + Op rep -> m [ExpReturns] opReturns op = extReturns <$> opType op diff --git a/src/Futhark/IR/Mem/Simplify.hs b/src/Futhark/IR/Mem/Simplify.hs --- a/src/Futhark/IR/Mem/Simplify.hs +++ b/src/Futhark/IR/Mem/Simplify.hs @@ -22,7 +22,7 @@ import qualified Futhark.IR.Syntax as AST import qualified Futhark.Optimise.Simplify as Simplify import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Optimise.Simplify.Rule import Futhark.Optimise.Simplify.Rules import Futhark.Pass @@ -30,17 +30,17 @@ import Futhark.Util simpleGeneric :: - (SimplifyMemory lore, Op lore ~ MemOp inner) => + (SimplifyMemory rep, Op rep ~ MemOp inner) => (OpWithWisdom inner -> UT.UsageTable) -> - Simplify.SimplifyOp lore inner -> - Simplify.SimpleOps lore + Simplify.SimplifyOp rep inner -> + Simplify.SimpleOps rep simpleGeneric = simplifiable simplifyProgGeneric :: - (SimplifyMemory lore, Op lore ~ MemOp inner) => - Simplify.SimpleOps lore -> - Prog lore -> - PassM (Prog lore) + (SimplifyMemory rep, Op rep ~ MemOp inner) => + Simplify.SimpleOps rep -> + Prog rep -> + PassM (Prog rep) simplifyProgGeneric ops = Simplify.simplifyProg ops @@ -59,14 +59,14 @@ not $ all primType $ patternTypes pat simplifyStmsGeneric :: - ( HasScope lore m, + ( HasScope rep m, MonadFreshNames m, - SimplifyMemory lore, - Op lore ~ MemOp inner + SimplifyMemory rep, + Op rep ~ MemOp inner ) => - Simplify.SimpleOps lore -> - Stms lore -> - m (ST.SymbolTable (Wise lore), Stms lore) + Simplify.SimpleOps rep -> + Stms rep -> + m (ST.SymbolTable (Wise rep), Stms rep) simplifyStmsGeneric ops stms = do scope <- askScope Simplify.simplifyStms @@ -76,18 +76,18 @@ scope stms -isResultAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore +isResultAlloc :: Op rep ~ MemOp op => Engine.BlockPred rep isResultAlloc _ usage (Let (AST.Pattern [] [bindee]) _ (Op Alloc {})) = UT.isInResult (patElemName bindee) usage isResultAlloc _ _ _ = False -isAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore +isAlloc :: Op rep ~ MemOp op => Engine.BlockPred rep isAlloc _ _ (Let _ _ (Op Alloc {})) = True isAlloc _ _ _ = False blockers :: - (Op lore ~ MemOp inner) => - Simplify.HoistBlockers lore + (Op rep ~ MemOp inner) => + Simplify.HoistBlockers rep blockers = Engine.noExtraHoistBlockers { Engine.blockHoistPar = isAlloc, @@ -96,17 +96,17 @@ } -- | Some constraints that must hold for the simplification rules to work. -type SimplifyMemory lore = - ( Simplify.SimplifiableLore lore, - ExpDec lore ~ (), - BodyDec lore ~ (), - AllocOp (Op (Wise lore)), - CanBeWise (Op lore), - BinderOps (Wise lore), - Mem lore +type SimplifyMemory rep = + ( Simplify.SimplifiableRep rep, + ExpDec rep ~ (), + BodyDec rep ~ (), + AllocOp (Op (Wise rep)), + CanBeWise (Op rep), + BinderOps (Wise rep), + Mem rep ) -callKernelRules :: SimplifyMemory lore => RuleBook (Wise lore) +callKernelRules :: SimplifyMemory rep => RuleBook (Wise rep) callKernelRules = standardRules <> ruleBook @@ -121,7 +121,7 @@ -- the array is not existential, and the index function of the array -- does not refer to any names in the pattern, then we can create a -- block of the proper size and always return there. -unExistentialiseMemory :: SimplifyMemory lore => TopDownRuleIf (Wise lore) +unExistentialiseMemory :: SimplifyMemory rep => TopDownRuleIf (Wise rep) unExistentialiseMemory vtable pat _ (cond, tbranch, fbranch, ifdec) | ST.simplifyMemory vtable, fixable <- foldl hasConcretisableMemory mempty $ patternElements pat, @@ -192,10 +192,10 @@ -- | If we are copying something that is itself a copy, just copy the -- original one instead. copyCopyToCopy :: - ( BinderOps lore, - LetDec lore ~ (VarWisdom, MemBound u) + ( BinderOps rep, + LetDec rep ~ (VarWisdom, MemBound u) ) => - TopDownRuleBasicOp lore + TopDownRuleBasicOp rep copyCopyToCopy vtable pat@(Pattern [] [pat_elem]) _ (Copy v1) | Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable, Just (_, MemArray _ _ _ (ArrayIn srcmem src_ixfun)) <- @@ -218,10 +218,10 @@ -- | If the destination of a copy is the same as the source, just -- remove it. removeIdentityCopy :: - ( BinderOps lore, - LetDec lore ~ (VarWisdom, MemBound u) + ( BinderOps rep, + LetDec rep ~ (VarWisdom, MemBound u) ) => - TopDownRuleBasicOp lore + TopDownRuleBasicOp rep removeIdentityCopy vtable pat@(Pattern [] [pe]) _ (Copy v) | (_, MemArray _ _ _ (ArrayIn dest_mem dest_ixfun)) <- patElemDec pe, Just (_, MemArray _ _ _ (ArrayIn src_mem src_ixfun)) <- @@ -234,7 +234,7 @@ -- If an allocation is statically known to be safe, then we can remove -- the certificates on it. This can help hoist things that would -- otherwise be stuck inside loops or branches. -decertifySafeAlloc :: SimplifyMemory lore => TopDownRuleOp (Wise lore) +decertifySafeAlloc :: SimplifyMemory rep => TopDownRuleOp (Wise rep) decertifySafeAlloc _ pat (StmAux cs attrs _) op | cs /= mempty, [Mem _] <- patternTypes pat, diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -5,8 +5,8 @@ -- | Parser for the Futhark core language. module Futhark.IR.Parse ( parseSOACS, - parseKernels, - parseKernelsMem, + parseGPU, + parseGPUMem, parseMC, parseMCMem, parseSeq, @@ -23,9 +23,9 @@ import Data.Void import Futhark.Analysis.PrimExp.Parse import Futhark.IR -import Futhark.IR.Kernels (Kernels) -import qualified Futhark.IR.Kernels.Kernel as Kernel -import Futhark.IR.KernelsMem (KernelsMem) +import Futhark.IR.GPU (GPU) +import qualified Futhark.IR.GPU.Kernel as Kernel +import Futhark.IR.GPUMem (GPUMem) import Futhark.IR.MC (MC) import qualified Futhark.IR.MC.Op as MC import Futhark.IR.MCMem (MCMem) @@ -335,46 +335,46 @@ -- bits. Essentially a manually passed-around type class dictionary, -- because ambiguities make it impossible to write this with actual -- type classes. -data PR lore = PR - { pRetType :: Parser (RetType lore), - pBranchType :: Parser (BranchType lore), - pFParamInfo :: Parser (FParamInfo lore), - pLParamInfo :: Parser (LParamInfo lore), - pLetDec :: Parser (LetDec lore), - pOp :: Parser (Op lore), - pBodyDec :: BodyDec lore, - pExpDec :: ExpDec lore +data PR rep = PR + { pRetType :: Parser (RetType rep), + pBranchType :: Parser (BranchType rep), + pFParamInfo :: Parser (FParamInfo rep), + pLParamInfo :: Parser (LParamInfo rep), + pLetDec :: Parser (LetDec rep), + pOp :: Parser (Op rep), + pBodyDec :: BodyDec rep, + pExpDec :: ExpDec rep } -pRetTypes :: PR lore -> Parser [RetType lore] +pRetTypes :: PR rep -> Parser [RetType rep] pRetTypes pr = braces $ pRetType pr `sepBy` pComma -pBranchTypes :: PR lore -> Parser [BranchType lore] +pBranchTypes :: PR rep -> Parser [BranchType rep] pBranchTypes pr = braces $ pBranchType pr `sepBy` pComma pParam :: Parser t -> Parser (Param t) pParam p = Param <$> pVName <*> (pColon *> p) -pFParam :: PR lore -> Parser (FParam lore) +pFParam :: PR rep -> Parser (FParam rep) pFParam = pParam . pFParamInfo -pFParams :: PR lore -> Parser [FParam lore] +pFParams :: PR rep -> Parser [FParam rep] pFParams pr = parens $ pFParam pr `sepBy` pComma -pLParam :: PR lore -> Parser (LParam lore) +pLParam :: PR rep -> Parser (LParam rep) pLParam = pParam . pLParamInfo -pLParams :: PR lore -> Parser [LParam lore] +pLParams :: PR rep -> Parser [LParam rep] pLParams pr = braces $ pLParam pr `sepBy` pComma -pPatElem :: PR lore -> Parser (PatElem lore) +pPatElem :: PR rep -> Parser (PatElem rep) pPatElem pr = (PatElem <$> pVName <*> (pColon *> pLetDec pr)) "pattern element" -pPattern :: PR lore -> Parser (Pattern lore) +pPattern :: PR rep -> Parser (Pattern rep) pPattern pr = uncurry Pattern <$> pPatternLike (pPatElem pr) -pIf :: PR lore -> Parser (Exp lore) +pIf :: PR rep -> Parser (Exp rep) pIf pr = keyword "if" $> f <*> pSort <*> pSubExp <*> (keyword "then" *> pBranchBody) @@ -395,22 +395,24 @@ braces (pBody pr) ] -pApply :: PR lore -> Parser (Exp lore) +pApply :: PR rep -> Parser (Exp rep) pApply pr = - keyword "apply" - $> Apply - <*> pName - <*> parens (pArg `sepBy` pComma) <* pColon - <*> pRetTypes pr - <*> pure (Safe, mempty, mempty) + keyword "apply" *> (p =<< choice [lexeme "" $> Unsafe, pure Safe]) where + p safety = + Apply + <$> pName + <*> parens (pArg `sepBy` pComma) <* pColon + <*> pRetTypes pr + <*> pure (safety, mempty, mempty) + pArg = choice [ lexeme "*" $> (,Consume) <*> pSubExp, (,Observe) <$> pSubExp ] -pLoop :: PR lore -> Parser (Exp lore) +pLoop :: PR rep -> Parser (Exp rep) pLoop pr = keyword "loop" $> uncurry DoLoop <*> pLoopParams @@ -434,7 +436,7 @@ keyword "while" $> WhileLoop <*> pVName ] -pLambda :: PR lore -> Parser (Lambda lore) +pLambda :: PR rep -> Parser (Lambda rep) pLambda pr = choice [ lexeme "\\" @@ -447,20 +449,20 @@ where lam params ret body = Lambda params body ret -pReduce :: PR lore -> Parser (SOAC.Reduce lore) +pReduce :: PR rep -> Parser (SOAC.Reduce rep) pReduce pr = SOAC.Reduce <$> pComm <*> pLambda pr <* pComma <*> braces (pSubExp `sepBy` pComma) -pScan :: PR lore -> Parser (SOAC.Scan lore) +pScan :: PR rep -> Parser (SOAC.Scan rep) pScan pr = SOAC.Scan <$> pLambda pr <* pComma <*> braces (pSubExp `sepBy` pComma) -pWithAcc :: PR lore -> Parser (Exp lore) +pWithAcc :: PR rep -> Parser (Exp rep) pWithAcc pr = keyword "with_acc" *> parens (WithAcc <$> braces (pInput `sepBy` pComma) <* pComma <*> pLambda pr) @@ -474,7 +476,7 @@ ) pCombFun = parens ((,) <$> pLambda pr <* pComma <*> pSubExps) -pExp :: PR lore -> Parser (Exp lore) +pExp :: PR rep -> Parser (Exp rep) pExp pr = choice [ pIf pr, @@ -485,7 +487,7 @@ BasicOp <$> pBasicOp ] -pStm :: PR lore -> Parser (Stm lore) +pStm :: PR rep -> Parser (Stm rep) pStm pr = keyword "let" $> Let <*> pPattern pr <* pEqual <*> pStmAux <*> pExp pr where @@ -497,10 +499,10 @@ pure mempty ] -pStms :: PR lore -> Parser (Stms lore) +pStms :: PR rep -> Parser (Stms rep) pStms pr = stmsFromList <$> many (pStm pr) -pBody :: PR lore -> Parser (Body lore) +pBody :: PR rep -> Parser (Body rep) pBody pr = choice [ Body (pBodyDec pr) <$> pStms pr <* keyword "in" <*> pResult, @@ -517,14 +519,15 @@ <* pComma <*> pEntryPointTypes where pEntryPointTypes = braces (pEntryPointType `sepBy` pComma) - pEntryPointType = + pEntryPointType = do + u <- pUniqueness choice - [ "direct" $> TypeDirect, - "unsigned" $> TypeUnsigned, - "opaque" *> parens (TypeOpaque <$> pStringLiteral <* pComma <*> pInt) + [ "direct" $> TypeDirect u, + "unsigned" $> TypeUnsigned u, + "opaque" *> parens (TypeOpaque u <$> pStringLiteral <* pComma <*> pInt) ] -pFunDef :: PR lore -> Parser (FunDef lore) +pFunDef :: PR rep -> Parser (FunDef rep) pFunDef pr = do attrs <- pAttrs entry <- @@ -538,10 +541,10 @@ FunDef entry attrs fname ret fparams <$> (pEqual *> braces (pBody pr)) -pProg :: PR lore -> Parser (Prog lore) +pProg :: PR rep -> Parser (Prog rep) pProg pr = Prog <$> pStms pr <*> many (pFunDef pr) -pSOAC :: PR lore -> Parser (SOAC.SOAC lore) +pSOAC :: PR rep -> Parser (SOAC.SOAC rep) pSOAC pr = choice [ keyword "map" *> pScrema pMapForm, @@ -742,13 +745,13 @@ pure (dim, blk_tile, reg_tile) pWrite = (,) <$> pSlice <* pEqual <*> pSubExp -pKernelBody :: PR lore -> Parser (SegOp.KernelBody lore) +pKernelBody :: PR rep -> Parser (SegOp.KernelBody rep) pKernelBody pr = SegOp.KernelBody (pBodyDec pr) <$> pStms pr <* keyword "return" <*> braces (pKernelResult `sepBy` pComma) -pSegOp :: PR lore -> Parser lvl -> Parser (SegOp.SegOp lvl lore) +pSegOp :: PR rep -> Parser lvl -> Parser (SegOp.SegOp lvl rep) pSegOp pr pLvl = choice [ keyword "segmap" *> pSegMap, @@ -805,7 +808,7 @@ pure SegOp.SegNoVirt ] -pHostOp :: PR lore -> Parser op -> Parser (Kernel.HostOp lore op) +pHostOp :: PR rep -> Parser op -> Parser (Kernel.HostOp rep op) pHostOp pr pOther = choice [ Kernel.SegOp <$> pSegOp pr pSegLevel, @@ -813,7 +816,7 @@ Kernel.OtherOp <$> pOther ] -pMCOp :: PR lore -> Parser op -> Parser (MC.MCOp lore op) +pMCOp :: PR rep -> Parser op -> Parser (MC.MCOp rep op) pMCOp pr pOther = choice [ MC.ParOp . Just @@ -944,17 +947,17 @@ where op = pMemOp empty -prKernels :: PR Kernels -prKernels = +prGPU :: PR GPU +prGPU = PR pDeclExtType pExtType pDeclType pType pType op () () where - op = pHostOp prKernels (pSOAC prKernels) + op = pHostOp prGPU (pSOAC prGPU) -prKernelsMem :: PR KernelsMem -prKernelsMem = +prGPUMem :: PR GPUMem +prGPUMem = PR pRetTypeMem pBranchTypeMem pFParamMem pLParamMem pLetDecMem op () () where - op = pMemOp $ pHostOp prKernelsMem empty + op = pMemOp $ pHostOp prGPUMem empty prMC :: PR MC prMC = @@ -968,28 +971,28 @@ where op = pMemOp $ pMCOp prMCMem empty -parseLore :: PR lore -> FilePath -> T.Text -> Either T.Text (Prog lore) -parseLore pr fname s = +parseRep :: PR rep -> FilePath -> T.Text -> Either T.Text (Prog rep) +parseRep pr fname s = either (Left . T.pack . errorBundlePretty) Right $ parse (whitespace *> pProg pr <* eof) fname s parseSOACS :: FilePath -> T.Text -> Either T.Text (Prog SOACS) -parseSOACS = parseLore prSOACS +parseSOACS = parseRep prSOACS parseSeq :: FilePath -> T.Text -> Either T.Text (Prog Seq) -parseSeq = parseLore prSeq +parseSeq = parseRep prSeq parseSeqMem :: FilePath -> T.Text -> Either T.Text (Prog SeqMem) -parseSeqMem = parseLore prSeqMem +parseSeqMem = parseRep prSeqMem -parseKernels :: FilePath -> T.Text -> Either T.Text (Prog Kernels) -parseKernels = parseLore prKernels +parseGPU :: FilePath -> T.Text -> Either T.Text (Prog GPU) +parseGPU = parseRep prGPU -parseKernelsMem :: FilePath -> T.Text -> Either T.Text (Prog KernelsMem) -parseKernelsMem = parseLore prKernelsMem +parseGPUMem :: FilePath -> T.Text -> Either T.Text (Prog GPUMem) +parseGPUMem = parseRep prGPUMem parseMC :: FilePath -> T.Text -> Either T.Text (Prog MC) -parseMC = parseLore prMC +parseMC = parseRep prMC parseMCMem :: FilePath -> T.Text -> Either T.Text (Prog MCMem) -parseMCMem = parseLore prMCMem +parseMCMem = parseRep prMCMem diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs --- a/src/Futhark/IR/Pretty.hs +++ b/src/Futhark/IR/Pretty.hs @@ -11,7 +11,7 @@ module Futhark.IR.Pretty ( prettyTuple, pretty, - PrettyLore (..), + PrettyRep (..), ppTuple', ) where @@ -21,31 +21,31 @@ import Futhark.IR.Syntax import Futhark.Util.Pretty --- | The class of lores whose annotations can be prettyprinted. +-- | The class of representations whose annotations can be prettyprinted. class - ( Decorations lore, - Pretty (RetType lore), - Pretty (BranchType lore), - Pretty (FParamInfo lore), - Pretty (LParamInfo lore), - Pretty (LetDec lore), - Pretty (Op lore) + ( RepTypes rep, + Pretty (RetType rep), + Pretty (BranchType rep), + Pretty (FParamInfo rep), + Pretty (LParamInfo rep), + Pretty (LetDec rep), + Pretty (Op rep) ) => - PrettyLore lore + PrettyRep rep where - ppExpLore :: ExpDec lore -> Exp lore -> Maybe Doc - ppExpLore _ _ = Nothing + ppExpDec :: ExpDec rep -> Exp rep -> Maybe Doc + ppExpDec _ _ = Nothing instance Pretty VName where ppr (VName vn i) = ppr vn <> text "_" <> text (show i) -instance Pretty NoUniqueness where - ppr _ = mempty - instance Pretty Commutativity where ppr Commutative = text "commutative" ppr Noncommutative = text "noncommutative" +instance Pretty NoUniqueness where + ppr _ = mempty + instance Pretty Shape where ppr = mconcat . map (brackets . ppr) . shapeDims @@ -96,10 +96,10 @@ ppr (Certificates []) = empty ppr (Certificates cs) = text "#" <> braces (commasep (map ppr cs)) -instance PrettyLore lore => Pretty (Stms lore) where +instance PrettyRep rep => Pretty (Stms rep) where ppr = stack . map ppr . stmsToList -instance PrettyLore lore => Pretty (Body lore) where +instance PrettyRep rep => Pretty (Body rep) where ppr (Body _ stms res) | null stms = braces (commasep $ map ppr res) | otherwise = @@ -115,7 +115,7 @@ where f v = text "#[" <> ppr v <> text "]" -stmAttrAnnots :: Stm lore -> [Doc] +stmAttrAnnots :: Stm rep -> [Doc] stmAttrAnnots = attrAnnots . stmAuxAttrs . stmAux certAnnots :: Certificates -> [Doc] @@ -123,7 +123,7 @@ | cs == mempty = [] | otherwise = [ppr cs] -stmCertAnnots :: Stm lore -> [Doc] +stmCertAnnots :: Stm rep -> [Doc] stmCertAnnots = certAnnots . stmAuxCerts . stmAux instance Pretty (PatElemT dec) => Pretty (PatternT dec) where @@ -135,7 +135,7 @@ instance Pretty t => Pretty (Param t) where ppr (Param name t) = ppr name <+> colon <+> align (ppr t) -instance PrettyLore lore => Pretty (Stm lore) where +instance PrettyRep rep => Pretty (Stm rep) where ppr bnd@(Let pat aux e) = align . hang 2 $ text "let" <+> align (ppr pat) @@ -154,7 +154,7 @@ stmannot = concat - [ maybeToList (ppExpLore (stmAuxDec aux) e), + [ maybeToList (ppExpDec (stmAuxDec aux) e), stmAttrAnnots bnd, stmCertAnnots bnd ] @@ -210,7 +210,7 @@ p (ErrorInt32 x) = ppr x <+> colon <+> text "i32" p (ErrorInt64 x) = ppr x <+> colon <+> text "i64" -instance PrettyLore lore => Pretty (Exp lore) where +instance PrettyRep rep => Pretty (Exp rep) where ppr (If c t f (IfDec ret ifsort)) = text "if" <+> info' <+> ppr c text "then" @@ -229,18 +229,17 @@ | otherwise = nestedBlock "{" "}" $ ppr b ppr (BasicOp op) = ppr op ppr (Apply fname args ret (safety, _, _)) = - text "apply" + applykw <+> text (nameToString fname) - <> safety' <> apply (map (align . pprArg) args) colon <+> braces (commasep $ map ppr ret) where pprArg (arg, Consume) = text "*" <> ppr arg pprArg (arg, _) = ppr arg - safety' = case safety of - Unsafe -> text "" - Safe -> mempty + applykw = case safety of + Unsafe -> text "apply " + Safe -> text "apply" ppr (Op op) = ppr op ppr (DoLoop ctx val form loopbody) = text "loop" <+> ppPattern ctxparams valparams @@ -273,7 +272,7 @@ pprLoopVar (p, a) = ppr p <+> text "in" <+> ppr a ppr (WithAcc inputs lam) = text "with_acc" - <> parens (braces (commasep $ map ppInput inputs) <> comma ppr lam) + <> parens (braces (commastack $ map ppInput inputs) <> comma ppr lam) where ppInput (shape, arrs, op) = parens @@ -284,7 +283,7 @@ comma parens (ppr op' <> comma ppTuple' (map ppr nes)) ) -instance PrettyLore lore => Pretty (Lambda lore) where +instance PrettyRep rep => Pretty (Lambda rep) where ppr (Lambda [] (Body _ stms []) []) | stms == mempty = text "nilFn" ppr (Lambda params body rettype) = text "\\" <+> ppTuple' params @@ -292,14 +291,15 @@ indent 2 (ppr body) instance Pretty EntryPointType where - ppr TypeDirect = "direct" - ppr TypeUnsigned = "unsigned" - ppr (TypeOpaque desc n) = "opaque" <> apply [ppr (show desc), ppr n] + ppr (TypeDirect u) = ppr u <> "direct" + ppr (TypeUnsigned u) = ppr u <> "unsigned" + ppr (TypeOpaque u desc n) = ppr u <> "opaque" <> apply [ppr (show desc), ppr n] -instance PrettyLore lore => Pretty (FunDef lore) where +instance PrettyRep rep => Pretty (FunDef rep) where ppr (FunDef entry attrs name rettype fparams body) = annot (attrAnnots attrs) $ - fun <+> text (nameToString name) + fun + indent 2 (text (nameToString name)) <+> apply (map ppr fparams) indent 2 (colon <+> align (ppTuple' rettype)) <+> equals @@ -315,7 +315,7 @@ ppTuple' ret_entry ) -instance PrettyLore lore => Pretty (Prog lore) where +instance PrettyRep rep => Pretty (Prog rep) where ppr (Prog consts funs) = stack $ punctuate line $ ppr consts : map ppr funs diff --git a/src/Futhark/IR/Prop.hs b/src/Futhark/IR/Prop.hs --- a/src/Futhark/IR/Prop.hs +++ b/src/Futhark/IR/Prop.hs @@ -35,13 +35,15 @@ certify, expExtTypesFromPattern, attrsForAssert, + lamIsBinOp, ASTConstraints, IsOp (..), - ASTLore (..), + ASTRep (..), ) where -import Data.List (find) +import Control.Monad +import Data.List (elemIndex, find) import qualified Data.Map.Strict as M import Data.Maybe (isJust, mapMaybe) import qualified Data.Set as S @@ -57,6 +59,7 @@ import Futhark.IR.Syntax import Futhark.Transform.Rename (Rename, Renameable) import Futhark.Transform.Substitute (Substitutable, Substitute) +import Futhark.Util (maybeNth) import Futhark.Util.Pretty -- | @isBuiltInFunction k@ is 'True' if @k@ is an element of 'builtInFunctions'. @@ -70,7 +73,7 @@ namify (k, (paramts, ret, _)) = (nameFromString k, (ret, paramts)) -- | If the expression is a t'BasicOp', return it, otherwise 'Nothing'. -asBasicOp :: Exp lore -> Maybe BasicOp +asBasicOp :: Exp rep -> Maybe BasicOp asBasicOp (BasicOp op) = Just op asBasicOp _ = Nothing @@ -78,7 +81,7 @@ -- any required certificates have been checked) in any context. For -- example, array indexing is not safe, as the index may be out of -- bounds. On the other hand, adding two numbers cannot fail. -safeExp :: IsOp (Op lore) => Exp lore -> Bool +safeExp :: IsOp (Op rep) => Exp rep -> Bool safeExp (BasicOp op) = safeBasicOp op where safeBasicOp (BinOp (SDiv _ Safe) _ _) = True @@ -131,7 +134,7 @@ safeExp WithAcc {} = True -- Although unlikely to matter. safeExp (Op op) = safeOp op -safeBody :: IsOp (Op lore) => Body lore -> Bool +safeBody :: IsOp (Op rep) => Body rep -> Bool safeBody = all (safeExp . stmExp) . bodyStms -- | Return the variable names used in 'Var' subexpressions. May contain @@ -148,7 +151,7 @@ -- Based on pattern matching and checking whether the lambda -- represents a known arithmetic operator; don't expect anything -- clever here. -commutativeLambda :: Lambda lore -> Bool +commutativeLambda :: Lambda rep -> Bool commutativeLambda lam = let body = lambdaBody lam n2 = length (lambdaParams lam) `div` 2 @@ -171,20 +174,20 @@ -- the parameters of the original function (they must all come at the -- end). entryPointSize :: EntryPointType -> Int -entryPointSize (TypeOpaque _ x) = x -entryPointSize TypeUnsigned = 1 -entryPointSize TypeDirect = 1 +entryPointSize (TypeOpaque _ _ x) = x +entryPointSize (TypeUnsigned _) = 1 +entryPointSize (TypeDirect _) = 1 -- | A 'StmAux' with empty 'Certificates'. defAux :: dec -> StmAux dec defAux = StmAux mempty mempty -- | The certificates associated with a statement. -stmCerts :: Stm lore -> Certificates +stmCerts :: Stm rep -> Certificates stmCerts = stmAuxCerts . stmAux -- | Add certificates to a statement. -certify :: Certificates -> Stm lore -> Stm lore +certify :: Certificates -> Stm rep -> Stm rep certify cs1 (Let pat (StmAux cs2 attrs dec) e) = Let pat (StmAux (cs2 <> cs1) attrs dec) e @@ -205,31 +208,31 @@ safeOp () = True cheapOp () = True --- | Lore-specific attributes; also means the lore supports some basic --- facilities. +-- | Representation-specific attributes; also means the rep supports +-- some basic facilities. class - ( Decorations lore, - PrettyLore lore, - Renameable lore, - Substitutable lore, - FreeDec (ExpDec lore), - FreeIn (LetDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - IsOp (Op lore) + ( RepTypes rep, + PrettyRep rep, + Renameable rep, + Substitutable rep, + FreeDec (ExpDec rep), + FreeIn (LetDec rep), + FreeDec (BodyDec rep), + FreeIn (FParamInfo rep), + FreeIn (LParamInfo rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + IsOp (Op rep) ) => - ASTLore lore + ASTRep rep where -- | Given a pattern, construct the type of a body that would match - -- it. An implementation for many lores would be + -- it. An implementation for many representations would be -- 'expExtTypesFromPattern'. expTypesFromPattern :: - (HasScope lore m, Monad m) => - Pattern lore -> - m [BranchType lore] + (HasScope rep m, Monad m) => + Pattern rep -> + m [BranchType rep] -- | Construct the type of an expression that would match the pattern. expExtTypesFromPattern :: Typed dec => PatternT dec -> [ExtType] @@ -244,3 +247,21 @@ Attrs $ S.filter attrForAssert attrs where attrForAssert = (== AttrComp "warn" ["safety_checks"]) + +-- | Horizontally fission a lambda that models a binary operator. +lamIsBinOp :: ASTRep rep => Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)] +lamIsBinOp lam = mapM splitStm $ bodyResult $ lambdaBody lam + where + n = length $ lambdaReturnType lam + splitStm (Var res) = do + Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y))) <- + find (([res] ==) . patternNames . stmPattern) $ + stmsToList $ bodyStms $ lambdaBody lam + i <- Var res `elemIndex` bodyResult (lambdaBody lam) + xp <- maybeNth i $ lambdaParams lam + yp <- maybeNth (n + i) $ lambdaParams lam + guard $ paramName xp == x + guard $ paramName yp == y + Prim t <- Just $ patElemType pe + return (op, t, paramName xp, paramName yp) + splitStm _ = Nothing diff --git a/src/Futhark/IR/Prop/Aliases.hs b/src/Futhark/IR/Prop/Aliases.hs --- a/src/Futhark/IR/Prop/Aliases.hs +++ b/src/Futhark/IR/Prop/Aliases.hs @@ -40,19 +40,13 @@ import Futhark.IR.Prop.Types import Futhark.IR.Syntax --- | The class of lores that contain aliasing information. -class - ( Decorations lore, - AliasedOp (Op lore), - AliasesOf (LetDec lore) - ) => - Aliased lore - where +-- | The class of representations that contain aliasing information. +class (RepTypes rep, AliasedOp (Op rep), AliasesOf (LetDec rep)) => Aliased rep where -- | The aliases of the body results. - bodyAliases :: Body lore -> [Names] + bodyAliases :: Body rep -> [Names] -- | The variables consumed in the body. - consumedInBody :: Body lore -> Names + consumedInBody :: Body rep -> Names vnameAliases :: VName -> Names vnameAliases = oneName @@ -95,7 +89,7 @@ returnAliases t [(subExpAliases se, d) | (se, d) <- args] -- | The aliases of an expression, one per non-context value returned. -expAliases :: (Aliased lore) => Exp lore -> [Names] +expAliases :: (Aliased rep) => Exp rep -> [Names] expAliases (If _ tb fb dec) = drop (length all_aliases - length ts) all_aliases where @@ -140,11 +134,11 @@ maskAliases als Observe = als -- | The variables consumed in this statement. -consumedInStm :: Aliased lore => Stm lore -> Names +consumedInStm :: Aliased rep => Stm rep -> Names consumedInStm = consumedInExp . stmExp -- | The variables consumed in this expression. -consumedInExp :: (Aliased lore) => Exp lore -> Names +consumedInExp :: (Aliased rep) => Exp rep -> Names consumedInExp (Apply _ args _ _) = mconcat (map (consumeArg . first subExpAliases) args) where @@ -178,7 +172,7 @@ consumedInExp (Op op) = consumedInOp op -- | The variables consumed by this lambda. -consumedByLambda :: Aliased lore => Lambda lore -> Names +consumedByLambda :: Aliased rep => Lambda rep -> Names consumedByLambda = consumedInBody . lambdaBody -- | The aliases of each pattern element (including the context). @@ -197,7 +191,7 @@ aliasesOf = aliasesOf . patElemDec -- | Also includes the name itself. -lookupAliases :: AliasesOf (LetDec lore) => VName -> Scope lore -> Names +lookupAliases :: AliasesOf (LetDec rep) => VName -> Scope rep -> Names lookupAliases v scope = case M.lookup v scope of Just (LetName dec) -> @@ -220,7 +214,7 @@ -- | The class of operations that can be given aliasing information. -- This is a somewhat subtle concept that is only used in the --- simplifier and when using "lore adapters". +-- simplifier and when using "rep adapters". class AliasedOp (OpWithAliases op) => CanBeAliased op where -- | The op that results when we add aliases to this op. type OpWithAliases op :: Data.Kind.Type diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs --- a/src/Futhark/IR/Prop/Names.hs +++ b/src/Futhark/IR/Prop/Names.hs @@ -132,16 +132,16 @@ fvNames = FV freeWalker :: - ( FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - FreeIn (Op lore) + ( FreeDec (ExpDec rep), + FreeDec (BodyDec rep), + FreeIn (FParamInfo rep), + FreeIn (LParamInfo rep), + FreeIn (LetDec rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + FreeIn (Op rep) ) => - Walker lore (State FV) + Walker rep (State FV) freeWalker = Walker { walkOnSubExp = modify . (<>) . freeIn', @@ -160,16 +160,16 @@ -- statements and result. Filters away the names that are bound by -- the statements. freeInStmsAndRes :: - ( FreeIn (Op lore), - FreeIn (LetDec lore), - FreeIn (LParamInfo lore), - FreeIn (FParamInfo lore), - FreeDec (BodyDec lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - FreeDec (ExpDec lore) + ( FreeIn (Op rep), + FreeIn (LetDec rep), + FreeIn (LParamInfo rep), + FreeIn (FParamInfo rep), + FreeDec (BodyDec rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + FreeDec (ExpDec rep) ) => - Stms lore -> + Stms rep -> Result -> FV freeInStmsAndRes stms res = @@ -207,63 +207,63 @@ freeIn' = foldMap freeIn' instance - ( FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - FreeIn (Op lore) + ( FreeDec (ExpDec rep), + FreeDec (BodyDec rep), + FreeIn (FParamInfo rep), + FreeIn (LParamInfo rep), + FreeIn (LetDec rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + FreeIn (Op rep) ) => - FreeIn (FunDef lore) + FreeIn (FunDef rep) where freeIn' (FunDef _ _ _ rettype params body) = fvBind (namesFromList $ map paramName params) $ freeIn' rettype <> freeIn' params <> freeIn' body instance - ( FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - FreeIn (Op lore) + ( FreeDec (ExpDec rep), + FreeDec (BodyDec rep), + FreeIn (FParamInfo rep), + FreeIn (LParamInfo rep), + FreeIn (LetDec rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + FreeIn (Op rep) ) => - FreeIn (Lambda lore) + FreeIn (Lambda rep) where freeIn' (Lambda params body rettype) = fvBind (namesFromList $ map paramName params) $ freeIn' rettype <> freeIn' params <> freeIn' body instance - ( FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - FreeIn (Op lore) + ( FreeDec (ExpDec rep), + FreeDec (BodyDec rep), + FreeIn (FParamInfo rep), + FreeIn (LParamInfo rep), + FreeIn (LetDec rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + FreeIn (Op rep) ) => - FreeIn (Body lore) + FreeIn (Body rep) where freeIn' (Body dec stms res) = precomputed dec $ freeIn' dec <> freeInStmsAndRes stms res instance - ( FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - FreeIn (Op lore) + ( FreeDec (ExpDec rep), + FreeDec (BodyDec rep), + FreeIn (FParamInfo rep), + FreeIn (LParamInfo rep), + FreeIn (LetDec rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + FreeIn (Op rep) ) => - FreeIn (Exp lore) + FreeIn (Exp rep) where freeIn' (DoLoop ctxmerge valmerge form loopbody) = let (ctxparams, ctxinits) = unzip ctxmerge @@ -282,22 +282,22 @@ freeIn' e = execState (walkExpM freeWalker e) mempty instance - ( FreeDec (ExpDec lore), - FreeDec (BodyDec lore), - FreeIn (FParamInfo lore), - FreeIn (LParamInfo lore), - FreeIn (LetDec lore), - FreeIn (RetType lore), - FreeIn (BranchType lore), - FreeIn (Op lore) + ( FreeDec (ExpDec rep), + FreeDec (BodyDec rep), + FreeIn (FParamInfo rep), + FreeIn (LParamInfo rep), + FreeIn (LetDec rep), + FreeIn (RetType rep), + FreeIn (BranchType rep), + FreeIn (Op rep) ) => - FreeIn (Stm lore) + FreeIn (Stm rep) where freeIn' (Let pat (StmAux cs attrs dec) e) = freeIn' cs <> freeIn' attrs <> precomputed dec (freeIn' dec <> freeIn' e <> freeIn' pat) -instance FreeIn (Stm lore) => FreeIn (Stms lore) where +instance FreeIn (Stm rep) => FreeIn (Stms rep) where freeIn' = foldMap freeIn' instance FreeIn Names where @@ -346,7 +346,7 @@ instance FreeIn dec => FreeIn (PatElemT dec) where freeIn' (PatElem _ dec) = freeIn' dec -instance FreeIn (LParamInfo lore) => FreeIn (LoopForm lore) where +instance FreeIn (LParamInfo rep) => FreeIn (LoopForm rep) where freeIn' (ForLoop _ _ bound loop_vars) = freeIn' bound <> freeIn' loop_vars freeIn' (WhileLoop cond) = freeIn' cond @@ -398,17 +398,17 @@ precomputed _ fv = fv -- | The names bound by the bindings immediately in a t'Body'. -boundInBody :: Body lore -> Names +boundInBody :: Body rep -> Names boundInBody = boundByStms . bodyStms -- | The names bound by a binding. -boundByStm :: Stm lore -> Names +boundByStm :: Stm rep -> Names boundByStm = namesFromList . patternNames . stmPattern -- | The names bound by the bindings. -boundByStms :: Stms lore -> Names +boundByStms :: Stms rep -> Names boundByStms = foldMap boundByStm -- | The names of the lambda parameters plus the index parameter. -boundByLambda :: Lambda lore -> [VName] +boundByLambda :: Lambda rep -> [VName] boundByLambda lam = map paramName (lambdaParams lam) diff --git a/src/Futhark/IR/Prop/Patterns.hs b/src/Futhark/IR/Prop/Patterns.hs --- a/src/Futhark/IR/Prop/Patterns.hs +++ b/src/Futhark/IR/Prop/Patterns.hs @@ -12,7 +12,7 @@ -- * Pattern elements patElemIdent, patElemType, - setPatElemLore, + setPatElemDec, patternElements, patternIdents, patternContextIdents, @@ -52,9 +52,9 @@ patElemType :: Typed dec => PatElemT dec -> Type patElemType = typeOf --- | Set the lore of a t'PatElem'. -setPatElemLore :: PatElemT oldattr -> newattr -> PatElemT newattr -setPatElemLore pe x = fmap (const x) pe +-- | Set the rep of a t'PatElem'. +setPatElemDec :: PatElemT oldattr -> newattr -> PatElemT newattr +setPatElemDec pe x = fmap (const x) pe -- | All pattern elements in the pattern - context first, then values. patternElements :: PatternT dec -> [PatElemT dec] diff --git a/src/Futhark/IR/Prop/Reshape.hs b/src/Futhark/IR/Prop/Reshape.hs --- a/src/Futhark/IR/Prop/Reshape.hs +++ b/src/Futhark/IR/Prop/Reshape.hs @@ -48,7 +48,7 @@ -- | Construct a 'Reshape' where all dimension changes are -- 'DimCoercion's. -shapeCoerce :: [SubExp] -> VName -> Exp lore +shapeCoerce :: [SubExp] -> VName -> Exp rep shapeCoerce newdims arr = BasicOp $ Reshape (map DimCoercion newdims) arr diff --git a/src/Futhark/IR/Prop/Scope.hs b/src/Futhark/IR/Prop/Scope.hs --- a/src/Futhark/IR/Prop/Scope.hs +++ b/src/Futhark/IR/Prop/Scope.hs @@ -41,22 +41,22 @@ import qualified Control.Monad.RWS.Strict import Control.Monad.Reader import qualified Data.Map.Strict as M -import Futhark.IR.Decorations import Futhark.IR.Pretty () import Futhark.IR.Prop.Patterns import Futhark.IR.Prop.Types +import Futhark.IR.Rep import Futhark.IR.Syntax -- | How some name in scope was bound. -data NameInfo lore - = LetName (LetDec lore) - | FParamName (FParamInfo lore) - | LParamName (LParamInfo lore) +data NameInfo rep + = LetName (LetDec rep) + | FParamName (FParamInfo rep) + | LParamName (LParamInfo rep) | IndexName IntType -deriving instance Decorations lore => Show (NameInfo lore) +deriving instance RepTypes rep => Show (NameInfo rep) -instance Decorations lore => Typed (NameInfo lore) where +instance RepTypes rep => Typed (NameInfo rep) where typeOf (LetName dec) = typeOf dec typeOf (FParamName dec) = typeOf dec typeOf (LParamName dec) = typeOf dec @@ -64,14 +64,14 @@ -- | A scope is a mapping from variable names to information about -- that name. -type Scope lore = M.Map VName (NameInfo lore) +type Scope rep = M.Map VName (NameInfo rep) -- | The class of applicative functors (or more common in practice: -- monads) that permit the lookup of variable types. A default method -- for 'lookupType' exists, which is sufficient (if not always -- maximally efficient, and using 'error' to fail) when 'askScope' -- is defined. -class (Applicative m, Decorations lore) => HasScope lore m | m -> lore where +class (Applicative m, RepTypes rep) => HasScope rep m | m -> rep where -- | Return the type of the given variable, or fail if it is not in -- the type environment. lookupType :: VName -> m Type @@ -79,7 +79,7 @@ -- | Return the info of the given variable, or fail if it is not in -- the type environment. - lookupInfo :: VName -> m (NameInfo lore) + lookupInfo :: VName -> m (NameInfo rep) lookupInfo name = asksScope (M.findWithDefault notFound name) where @@ -90,144 +90,144 @@ -- | Return the type environment contained in the applicative -- functor. - askScope :: m (Scope lore) + askScope :: m (Scope rep) -- | Return the result of applying some function to the type -- environment. - asksScope :: (Scope lore -> a) -> m a + asksScope :: (Scope rep -> a) -> m a asksScope f = f <$> askScope instance - (Applicative m, Monad m, Decorations lore) => - HasScope lore (ReaderT (Scope lore) m) + (Applicative m, Monad m, RepTypes rep) => + HasScope rep (ReaderT (Scope rep) m) where askScope = ask -instance (Monad m, HasScope lore m) => HasScope lore (ExceptT e m) where +instance (Monad m, HasScope rep m) => HasScope rep (ExceptT e m) where askScope = lift askScope instance - (Applicative m, Monad m, Monoid w, Decorations lore) => - HasScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) + (Applicative m, Monad m, Monoid w, RepTypes rep) => + HasScope rep (Control.Monad.RWS.Strict.RWST (Scope rep) w s m) where askScope = ask instance - (Applicative m, Monad m, Monoid w, Decorations lore) => - HasScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m) + (Applicative m, Monad m, Monoid w, RepTypes rep) => + HasScope rep (Control.Monad.RWS.Lazy.RWST (Scope rep) w s m) where askScope = ask -- | The class of monads that not only provide a 'Scope', but also -- the ability to locally extend it. A 'Reader' containing a -- 'Scope' is the prototypical example of such a monad. -class (HasScope lore m, Monad m) => LocalScope lore m where +class (HasScope rep m, Monad m) => LocalScope rep m where -- | Run a computation with an extended type environment. Note that -- this is intended to *add* to the current type environment, it -- does not replace it. - localScope :: Scope lore -> m a -> m a + localScope :: Scope rep -> m a -> m a -instance (Monad m, LocalScope lore m) => LocalScope lore (ExceptT e m) where +instance (Monad m, LocalScope rep m) => LocalScope rep (ExceptT e m) where localScope = mapExceptT . localScope instance - (Applicative m, Monad m, Decorations lore) => - LocalScope lore (ReaderT (Scope lore) m) + (Applicative m, Monad m, RepTypes rep) => + LocalScope rep (ReaderT (Scope rep) m) where localScope = local . M.union instance - (Applicative m, Monad m, Monoid w, Decorations lore) => - LocalScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) + (Applicative m, Monad m, Monoid w, RepTypes rep) => + LocalScope rep (Control.Monad.RWS.Strict.RWST (Scope rep) w s m) where localScope = local . M.union instance - (Applicative m, Monad m, Monoid w, Decorations lore) => - LocalScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) w s m) + (Applicative m, Monad m, Monoid w, RepTypes rep) => + LocalScope rep (Control.Monad.RWS.Lazy.RWST (Scope rep) w s m) where localScope = local . M.union -- | The class of things that can provide a scope. There is no -- overarching rule for what this means. For a 'Stm', it is the -- corresponding pattern. For a t'Lambda', is is the parameters. -class Scoped lore a | a -> lore where - scopeOf :: a -> Scope lore +class Scoped rep a | a -> rep where + scopeOf :: a -> Scope rep -- | Extend the monadic scope with the 'scopeOf' the given value. -inScopeOf :: (Scoped lore a, LocalScope lore m) => a -> m b -> m b +inScopeOf :: (Scoped rep a, LocalScope rep m) => a -> m b -> m b inScopeOf = localScope . scopeOf -instance Scoped lore a => Scoped lore [a] where +instance Scoped rep a => Scoped rep [a] where scopeOf = mconcat . map scopeOf -instance Scoped lore (Stms lore) where +instance Scoped rep (Stms rep) where scopeOf = foldMap scopeOf -instance Scoped lore (Stm lore) where +instance Scoped rep (Stm rep) where scopeOf = scopeOfPattern . stmPattern -instance Scoped lore (FunDef lore) where +instance Scoped rep (FunDef rep) where scopeOf = scopeOfFParams . funDefParams -instance Scoped lore (VName, NameInfo lore) where +instance Scoped rep (VName, NameInfo rep) where scopeOf = uncurry M.singleton -instance Scoped lore (LoopForm lore) where +instance Scoped rep (LoopForm rep) where scopeOf (WhileLoop _) = mempty scopeOf (ForLoop i it _ xs) = M.insert i (IndexName it) $ scopeOfLParams (map fst xs) -- | The scope of a pattern. -scopeOfPattern :: LetDec lore ~ dec => PatternT dec -> Scope lore +scopeOfPattern :: LetDec rep ~ dec => PatternT dec -> Scope rep scopeOfPattern = mconcat . map scopeOfPatElem . patternElements -- | The scope of a pattern element. -scopeOfPatElem :: LetDec lore ~ dec => PatElemT dec -> Scope lore +scopeOfPatElem :: LetDec rep ~ dec => PatElemT dec -> Scope rep scopeOfPatElem (PatElem name dec) = M.singleton name $ LetName dec -- | The scope of some lambda parameters. scopeOfLParams :: - LParamInfo lore ~ dec => + LParamInfo rep ~ dec => [Param dec] -> - Scope lore + Scope rep scopeOfLParams = M.fromList . map f where f param = (paramName param, LParamName $ paramDec param) -- | The scope of some function or loop parameters. scopeOfFParams :: - FParamInfo lore ~ dec => + FParamInfo rep ~ dec => [Param dec] -> - Scope lore + Scope rep scopeOfFParams = M.fromList . map f where f param = (paramName param, FParamName $ paramDec param) -instance Scoped lore (Lambda lore) where +instance Scoped rep (Lambda rep) where scopeOf lam = scopeOfLParams $ lambdaParams lam --- | A constraint that indicates two lores have the same 'NameInfo' +-- | A constraint that indicates two representations have the same 'NameInfo' -- representation. -type SameScope lore1 lore2 = - ( LetDec lore1 ~ LetDec lore2, - FParamInfo lore1 ~ FParamInfo lore2, - LParamInfo lore1 ~ LParamInfo lore2 +type SameScope rep1 rep2 = + ( LetDec rep1 ~ LetDec rep2, + FParamInfo rep1 ~ FParamInfo rep2, + LParamInfo rep1 ~ LParamInfo rep2 ) -- | If two scopes are really the same, then you can convert one to -- the other. castScope :: - SameScope fromlore tolore => - Scope fromlore -> - Scope tolore + SameScope fromrep torep => + Scope fromrep -> + Scope torep castScope = M.map castNameInfo castNameInfo :: - SameScope fromlore tolore => - NameInfo fromlore -> - NameInfo tolore + SameScope fromrep torep => + NameInfo fromrep -> + NameInfo torep castNameInfo (LetName dec) = LetName dec castNameInfo (FParamName dec) = FParamName dec castNameInfo (LParamName dec) = LParamName dec @@ -236,17 +236,17 @@ -- | A monad transformer that carries around an extended 'Scope'. -- Its 'lookupType' method will first look in the extended 'Scope', -- and then use the 'lookupType' method of the underlying monad. -newtype ExtendedScope lore m a = ExtendedScope (ReaderT (Scope lore) m a) +newtype ExtendedScope rep m a = ExtendedScope (ReaderT (Scope rep) m a) deriving ( Functor, Applicative, Monad, - MonadReader (Scope lore) + MonadReader (Scope rep) ) instance - (HasScope lore m, Monad m) => - HasScope lore (ExtendedScope lore m) + (HasScope rep m, Monad m) => + HasScope rep (ExtendedScope rep m) where lookupType name = do res <- asks $ fmap typeOf . M.lookup name @@ -255,7 +255,7 @@ -- | Run a computation in the extended type environment. extendedScope :: - ExtendedScope lore m a -> - Scope lore -> + ExtendedScope rep m a -> + Scope rep -> m a extendedScope (ExtendedScope m) = runReaderT m diff --git a/src/Futhark/IR/Prop/TypeOf.hs b/src/Futhark/IR/Prop/TypeOf.hs --- a/src/Futhark/IR/Prop/TypeOf.hs +++ b/src/Futhark/IR/Prop/TypeOf.hs @@ -51,14 +51,14 @@ -- | @mapType f arrts@ wraps each element in the return type of @f@ in -- an array with size equal to the outermost dimension of the first -- element of @arrts@. -mapType :: SubExp -> Lambda lore -> [Type] +mapType :: SubExp -> Lambda rep -> [Type] mapType outersize f = [ arrayOf t (Shape [outersize]) NoUniqueness | t <- lambdaReturnType f ] -- | The type of a primitive operation. -primOpType :: HasScope lore m => BasicOp -> m [Type] +primOpType :: HasScope rep m => BasicOp -> m [Type] primOpType (SubExp se) = pure <$> subExpType se primOpType (Opaque se) = @@ -121,8 +121,8 @@ -- | The type of an expression. expExtType :: - (HasScope lore m, TypedOp (Op lore)) => - Exp lore -> + (HasScope rep m, TypedOp (Op rep)) => + Exp rep -> m [ExtType] expExtType (Apply _ _ rt _) = pure $ map (fromDecl . declExtTypeOf) rt expExtType (If _ _ _ rt) = pure $ map extTypeOf $ ifReturns rt @@ -141,22 +141,22 @@ -- | The number of values returned by an expression. expExtTypeSize :: - (Decorations lore, TypedOp (Op lore)) => - Exp lore -> + (RepTypes rep, TypedOp (Op rep)) => + Exp rep -> Int expExtTypeSize = length . feelBad . expExtType -- FIXME, this is a horrible quick hack. -newtype FeelBad lore a = FeelBad {feelBad :: a} +newtype FeelBad rep a = FeelBad {feelBad :: a} -instance Functor (FeelBad lore) where +instance Functor (FeelBad rep) where fmap f = FeelBad . f . feelBad -instance Applicative (FeelBad lore) where +instance Applicative (FeelBad rep) where pure = FeelBad f <*> x = FeelBad $ feelBad f $ feelBad x -instance Decorations lore => HasScope lore (FeelBad lore) where +instance RepTypes rep => HasScope rep (FeelBad rep) where lookupType = const $ pure $ Prim $ IntType Int64 askScope = pure mempty diff --git a/src/Futhark/IR/Rep.hs b/src/Futhark/IR/Rep.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/IR/Rep.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | The core Futhark AST is parameterised by a @rep@ type parameter, +-- which is then used to invoke the type families defined here. +module Futhark.IR.Rep + ( RepTypes (..), + module Futhark.IR.RetType, + ) +where + +import qualified Data.Kind +import Futhark.IR.Prop.Types +import Futhark.IR.RetType +import Futhark.IR.Syntax.Core + +-- | A collection of type families giving various common types for a +-- representation, along with constraints specifying that the types +-- they map to should satisfy some minimal requirements. +class + ( Show (LetDec l), + Show (ExpDec l), + Show (BodyDec l), + Show (FParamInfo l), + Show (LParamInfo l), + Show (RetType l), + Show (BranchType l), + Show (Op l), + Eq (LetDec l), + Eq (ExpDec l), + Eq (BodyDec l), + Eq (FParamInfo l), + Eq (LParamInfo l), + Eq (RetType l), + Eq (BranchType l), + Eq (Op l), + Ord (LetDec l), + Ord (ExpDec l), + Ord (BodyDec l), + Ord (FParamInfo l), + Ord (LParamInfo l), + Ord (RetType l), + Ord (BranchType l), + Ord (Op l), + IsRetType (RetType l), + IsBodyType (BranchType l), + Typed (FParamInfo l), + Typed (LParamInfo l), + Typed (LetDec l), + DeclTyped (FParamInfo l) + ) => + RepTypes l + where + -- | Decoration for every let-pattern element. + type LetDec l :: Data.Kind.Type + + type LetDec l = Type + + -- | Decoration for every expression. + type ExpDec l :: Data.Kind.Type + + type ExpDec l = () + + -- | Decoration for every body. + type BodyDec l :: Data.Kind.Type + + type BodyDec l = () + + -- | Decoration for every (non-lambda) function parameter. + type FParamInfo l :: Data.Kind.Type + + type FParamInfo l = DeclType + + -- | Decoration for every lambda function parameter. + type LParamInfo l :: Data.Kind.Type + + type LParamInfo l = Type + + -- | The return type decoration of function calls. + type RetType l :: Data.Kind.Type + + type RetType l = DeclExtType + + -- | The return type decoration of branches. + type BranchType l :: Data.Kind.Type + + type BranchType l = ExtType + + -- | Extensible operation. + type Op l :: Data.Kind.Type + + type Op l = () diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -3,8 +3,7 @@ -- | A simple representation with SOACs and nested parallelism. module Futhark.IR.SOACS - ( -- * The Lore definition - SOACS, + ( SOACS, -- * Syntax types Body, @@ -54,13 +53,13 @@ -- like Standard ML. Instead, we have to abuse the namespace/module -- system. --- | The lore for the basic representation. +-- | The rep for the basic representation. data SOACS -instance Decorations SOACS where +instance RepTypes SOACS where type Op SOACS = SOAC SOACS -instance ASTLore SOACS where +instance ASTRep SOACS where expTypesFromPattern = return . expExtTypesFromPattern type Exp = AST.Exp SOACS @@ -94,4 +93,4 @@ instance BinderOps SOACS -instance PrettyLore SOACS +instance PrettyRep SOACS diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs --- a/src/Futhark/IR/SOACS/SOAC.hs +++ b/src/Futhark/IR/SOACS/SOAC.hs @@ -66,7 +66,7 @@ import Futhark.IR import Futhark.IR.Aliases (Aliases, removeLambdaAliases) import Futhark.IR.Prop.Aliases -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Transform.Rename import Futhark.Transform.Substitute import qualified Futhark.TypeCheck as TC @@ -76,8 +76,8 @@ import Prelude hiding (id, (.)) -- | A second-order array combinator (SOAC). -data SOAC lore - = Stream SubExp [VName] (StreamForm lore) [SubExp] (Lambda lore) +data SOAC rep + = Stream SubExp [VName] (StreamForm rep) [SubExp] (Lambda rep) | -- | @Scatter @ -- -- Scatter maps values from a set of input arrays to indices and values of a @@ -115,27 +115,27 @@ -- will correspond to the first two output values, and so on. For this -- example, should return a total of 11 values, 8 index values and -- 3 output values. - Scatter SubExp (Lambda lore) [VName] [(Shape, Int, VName)] + Scatter SubExp (Lambda rep) [VName] [(Shape, Int, VName)] | -- | @Hist @ -- -- The first SubExp is the length of the input arrays. The first -- list describes the operations to perform. The t'Lambda' is the -- bucket function. Finally comes the input images. - Hist SubExp [HistOp lore] (Lambda lore) [VName] + Hist SubExp [HistOp rep] (Lambda rep) [VName] | -- | A combination of scan, reduction, and map. The first -- t'SubExp' is the size of the input arrays. - Screma SubExp [VName] (ScremaForm lore) + Screma SubExp [VName] (ScremaForm rep) deriving (Eq, Ord, Show) -- | Information about computing a single histogram. -data HistOp lore = HistOp +data HistOp rep = HistOp { histWidth :: SubExp, -- | Race factor @RF@ means that only @1/RF@ -- bins are used. histRaceFactor :: SubExp, histDest :: [VName], histNeutral :: [SubExp], - histOp :: Lambda lore + histOp :: Lambda rep } deriving (Eq, Ord, Show) @@ -147,21 +147,21 @@ deriving (Eq, Ord, Show) -- | What kind of stream is this? -data StreamForm lore - = Parallel StreamOrd Commutativity (Lambda lore) +data StreamForm rep + = Parallel StreamOrd Commutativity (Lambda rep) | Sequential deriving (Eq, Ord, Show) -- | The essential parts of a 'Screma' factored out (everything -- except the input arrays). -data ScremaForm lore +data ScremaForm rep = ScremaForm - [Scan lore] - [Reduce lore] - (Lambda lore) + [Scan rep] + [Reduce rep] + (Lambda rep) deriving (Eq, Ord, Show) -singleBinOp :: Bindable lore => [Lambda lore] -> Lambda lore +singleBinOp :: Bindable rep => [Lambda rep] -> Lambda rep singleBinOp lams = Lambda { lambdaParams = concatMap xParams lams ++ concatMap yParams lams, @@ -176,37 +176,37 @@ yParams lam = drop (length (lambdaReturnType lam)) (lambdaParams lam) -- | How to compute a single scan result. -data Scan lore = Scan - { scanLambda :: Lambda lore, +data Scan rep = Scan + { scanLambda :: Lambda rep, scanNeutral :: [SubExp] } deriving (Eq, Ord, Show) -- | How many reduction results are produced by these 'Scan's? -scanResults :: [Scan lore] -> Int +scanResults :: [Scan rep] -> Int scanResults = sum . map (length . scanNeutral) -- | Combine multiple scan operators to a single operator. -singleScan :: Bindable lore => [Scan lore] -> Scan lore +singleScan :: Bindable rep => [Scan rep] -> Scan rep singleScan scans = let scan_nes = concatMap scanNeutral scans scan_lam = singleBinOp $ map scanLambda scans in Scan scan_lam scan_nes -- | How to compute a single reduction result. -data Reduce lore = Reduce +data Reduce rep = Reduce { redComm :: Commutativity, - redLambda :: Lambda lore, + redLambda :: Lambda rep, redNeutral :: [SubExp] } deriving (Eq, Ord, Show) -- | How many reduction results are produced by these 'Reduce's? -redResults :: [Reduce lore] -> Int +redResults :: [Reduce rep] -> Int redResults = sum . map (length . redNeutral) -- | Combine multiple reduction operators to a single operator. -singleReduce :: Bindable lore => [Reduce lore] -> Reduce lore +singleReduce :: Bindable rep => [Reduce rep] -> Reduce rep singleReduce reds = let red_nes = concatMap redNeutral reds red_lam = singleBinOp $ map redLambda reds @@ -214,7 +214,7 @@ -- | The types produced by a single 'Screma', given the size of the -- input array. -scremaType :: SubExp -> ScremaForm lore -> [Type] +scremaType :: SubExp -> ScremaForm rep -> [Type] scremaType w (ScremaForm scans reds map_lam) = scan_tps ++ red_tps ++ map (`arrayOfRow` w) map_tps where @@ -227,9 +227,9 @@ -- | Construct a lambda that takes parameters of the given types and -- simply returns them unchanged. mkIdentityLambda :: - (Bindable lore, MonadFreshNames m) => + (Bindable rep, MonadFreshNames m) => [Type] -> - m (Lambda lore) + m (Lambda rep) mkIdentityLambda ts = do params <- mapM (newParam "x") ts return @@ -240,31 +240,31 @@ } -- | Is the given lambda an identity lambda? -isIdentityLambda :: Lambda lore -> Bool +isIdentityLambda :: Lambda rep -> Bool isIdentityLambda lam = bodyResult (lambdaBody lam) == map (Var . paramName) (lambdaParams lam) -- | A lambda with no parameters that returns no values. -nilFn :: Bindable lore => Lambda lore +nilFn :: Bindable rep => Lambda rep nilFn = Lambda mempty (mkBody mempty mempty) mempty -- | Construct a Screma with possibly multiple scans, and -- the given map function. -scanomapSOAC :: [Scan lore] -> Lambda lore -> ScremaForm lore +scanomapSOAC :: [Scan rep] -> Lambda rep -> ScremaForm rep scanomapSOAC scans = ScremaForm scans [] -- | Construct a Screma with possibly multiple reductions, and -- the given map function. -redomapSOAC :: [Reduce lore] -> Lambda lore -> ScremaForm lore +redomapSOAC :: [Reduce rep] -> Lambda rep -> ScremaForm rep redomapSOAC = ScremaForm [] -- | Construct a Screma with possibly multiple scans, and identity map -- function. scanSOAC :: - (Bindable lore, MonadFreshNames m) => - [Scan lore] -> - m (ScremaForm lore) + (Bindable rep, MonadFreshNames m) => + [Scan rep] -> + m (ScremaForm rep) scanSOAC scans = scanomapSOAC scans <$> mkIdentityLambda ts where ts = concatMap (lambdaReturnType . scanLambda) scans @@ -272,40 +272,40 @@ -- | Construct a Screma with possibly multiple reductions, and -- identity map function. reduceSOAC :: - (Bindable lore, MonadFreshNames m) => - [Reduce lore] -> - m (ScremaForm lore) + (Bindable rep, MonadFreshNames m) => + [Reduce rep] -> + m (ScremaForm rep) reduceSOAC reds = redomapSOAC reds <$> mkIdentityLambda ts where ts = concatMap (lambdaReturnType . redLambda) reds -- | Construct a Screma corresponding to a map. -mapSOAC :: Lambda lore -> ScremaForm lore +mapSOAC :: Lambda rep -> ScremaForm rep mapSOAC = ScremaForm [] [] -- | Does this Screma correspond to a scan-map composition? -isScanomapSOAC :: ScremaForm lore -> Maybe ([Scan lore], Lambda lore) +isScanomapSOAC :: ScremaForm rep -> Maybe ([Scan rep], Lambda rep) isScanomapSOAC (ScremaForm scans reds map_lam) = do guard $ null reds guard $ not $ null scans return (scans, map_lam) -- | Does this Screma correspond to pure scan? -isScanSOAC :: ScremaForm lore -> Maybe [Scan lore] +isScanSOAC :: ScremaForm rep -> Maybe [Scan rep] isScanSOAC form = do (scans, map_lam) <- isScanomapSOAC form guard $ isIdentityLambda map_lam return scans -- | Does this Screma correspond to a reduce-map composition? -isRedomapSOAC :: ScremaForm lore -> Maybe ([Reduce lore], Lambda lore) +isRedomapSOAC :: ScremaForm rep -> Maybe ([Reduce rep], Lambda rep) isRedomapSOAC (ScremaForm scans reds map_lam) = do guard $ null scans guard $ not $ null reds return (reds, map_lam) -- | Does this Screma correspond to a pure reduce? -isReduceSOAC :: ScremaForm lore -> Maybe [Reduce lore] +isReduceSOAC :: ScremaForm rep -> Maybe [Reduce rep] isReduceSOAC form = do (reds, map_lam) <- isRedomapSOAC form guard $ isIdentityLambda map_lam @@ -313,7 +313,7 @@ -- | Does this Screma correspond to a simple map, without any -- reduction or scan results? -isMapSOAC :: ScremaForm lore -> Maybe (Lambda lore) +isMapSOAC :: ScremaForm rep -> Maybe (Lambda rep) isMapSOAC (ScremaForm scans reds map_lam) = do guard $ null scans guard $ null reds @@ -369,14 +369,14 @@ in splitAt num_indices results -- | Like 'Mapper', but just for 'SOAC's. -data SOACMapper flore tlore m = SOACMapper +data SOACMapper frep trep m = SOACMapper { mapOnSOACSubExp :: SubExp -> m SubExp, - mapOnSOACLambda :: Lambda flore -> m (Lambda tlore), + mapOnSOACLambda :: Lambda frep -> m (Lambda trep), mapOnSOACVName :: VName -> m VName } -- | A mapper that simply returns the SOAC verbatim. -identitySOACMapper :: Monad m => SOACMapper lore lore m +identitySOACMapper :: Monad m => SOACMapper rep rep m identitySOACMapper = SOACMapper { mapOnSOACSubExp = return, @@ -389,9 +389,9 @@ -- and is done left-to-right. mapSOACM :: (Applicative m, Monad m) => - SOACMapper flore tlore m -> - SOAC flore -> - m (SOAC tlore) + SOACMapper frep trep m -> + SOAC frep -> + m (SOAC trep) mapSOACM tv (Stream size arrs form accs lam) = Stream <$> mapOnSOACSubExp tv size <*> mapM (mapOnSOACVName tv) arrs @@ -448,7 +448,7 @@ <*> mapOnSOACLambda tv map_lam ) -instance ASTLore lore => FreeIn (SOAC lore) where +instance ASTRep rep => FreeIn (SOAC rep) where freeIn' = flip execState mempty . mapSOACM free where walk f x = modify (<> f x) >> return x @@ -459,7 +459,7 @@ mapOnSOACVName = walk freeIn' } -instance ASTLore lore => Substitute (SOAC lore) where +instance ASTRep rep => Substitute (SOAC rep) where substituteNames subst = runIdentity . mapSOACM substitute where @@ -470,13 +470,13 @@ mapOnSOACVName = return . substituteNames subst } -instance ASTLore lore => Rename (SOAC lore) where +instance ASTRep rep => Rename (SOAC rep) where rename = mapSOACM renamer where renamer = SOACMapper rename rename rename -- | The type of a SOAC. -soacType :: SOAC lore -> [Type] +soacType :: SOAC rep -> [Type] soacType (Stream outersize _ _ accs lam) = map (substNamesInType substs) rtp where @@ -495,10 +495,10 @@ soacType (Screma w _arrs form) = scremaType w form -instance TypedOp (SOAC lore) where +instance TypedOp (SOAC rep) where opType = pure . staticShapes . soacType -instance (ASTLore lore, Aliased lore) => AliasedOp (SOAC lore) where +instance (ASTRep rep, Aliased rep) => AliasedOp (SOAC rep) where opAliases = map (const mempty) . soacType -- Only map functions can consume anything. The operands to scan @@ -527,20 +527,20 @@ namesFromList $ concatMap histDest ops mapHistOp :: - (Lambda flore -> Lambda tlore) -> - HistOp flore -> - HistOp tlore + (Lambda frep -> Lambda trep) -> + HistOp frep -> + HistOp trep mapHistOp f (HistOp w rf dests nes lam) = HistOp w rf dests nes $ f lam instance - ( ASTLore lore, - ASTLore (Aliases lore), - CanBeAliased (Op lore) + ( ASTRep rep, + ASTRep (Aliases rep), + CanBeAliased (Op rep) ) => - CanBeAliased (SOAC lore) + CanBeAliased (SOAC rep) where - type OpWithAliases (SOAC lore) = SOAC (Aliases lore) + type OpWithAliases (SOAC rep) = SOAC (Aliases rep) addOpAliases aliases (Stream size arr form accs lam) = Stream size arr (analyseStreamForm form) accs $ @@ -571,7 +571,7 @@ where remove = SOACMapper return (return . removeLambdaAliases) return -instance ASTLore lore => IsOp (SOAC lore) where +instance ASTRep rep => IsOp (SOAC rep) where safeOp _ = False cheapOp _ = True @@ -588,14 +588,14 @@ substNamesInSubExp subs (Var idd) = M.findWithDefault (Var idd) idd subs -instance (ASTLore lore, CanBeWise (Op lore)) => CanBeWise (SOAC lore) where - type OpWithWisdom (SOAC lore) = SOAC (Wise lore) +instance (ASTRep rep, CanBeWise (Op rep)) => CanBeWise (SOAC rep) where + type OpWithWisdom (SOAC rep) = SOAC (Wise rep) removeOpWisdom = runIdentity . mapSOACM remove where remove = SOACMapper return (return . removeLambdaWisdom) return -instance Decorations lore => ST.IndexOp (SOAC lore) where +instance RepTypes rep => ST.IndexOp (SOAC rep) where indexOp vtable k soac [i] = do (lam, se, arr_params, arrs) <- lambdaAndSubExp soac let arr_indexes = M.fromList $ catMaybes $ zipWith arrIndex arr_params arrs @@ -634,7 +634,7 @@ indexOp _ _ _ _ = Nothing -- | Type-check a SOAC. -typeCheckSOAC :: TC.Checkable lore => SOAC (Aliases lore) -> TC.TypeM lore () +typeCheckSOAC :: TC.Checkable rep => SOAC (Aliases rep) -> TC.TypeM rep () typeCheckSOAC (Stream size arrexps form accexps lam) = do TC.require [Prim int64] size accargs <- mapM TC.checkArg accexps @@ -804,7 +804,7 @@ "Map function return type " ++ prettyTuple map_lam_ts ++ " wrong for given scan and reduction functions." -instance OpMetrics (Op lore) => OpMetrics (SOAC lore) where +instance OpMetrics (Op rep) => OpMetrics (SOAC rep) where opMetrics (Stream _ _ _ _ lam) = inside "Stream" $ lambdaMetrics lam opMetrics (Scatter _len lam _ivs _as) = @@ -817,7 +817,7 @@ mapM_ (lambdaMetrics . redLambda) reds lambdaMetrics map_lam -instance PrettyLore lore => PP.Pretty (SOAC lore) where +instance PrettyRep rep => PP.Pretty (SOAC rep) where ppr (Stream size arrs form acc lam) = case form of Parallel o comm lam0 -> @@ -879,7 +879,7 @@ -- | Prettyprint the given Screma. ppScrema :: - (PrettyLore lore, Pretty inp) => SubExp -> [inp] -> ScremaForm lore -> Doc + (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> ScremaForm rep -> Doc ppScrema w arrs (ScremaForm scans reds map_lam) = text "screma" <> parens @@ -890,7 +890,7 @@ ppr map_lam ) -instance PrettyLore lore => Pretty (Scan lore) where +instance PrettyRep rep => Pretty (Scan rep) where ppr (Scan scan_lam scan_nes) = ppr scan_lam <> comma PP.braces (commasep $ map ppr scan_nes) @@ -898,17 +898,17 @@ ppComm Noncommutative = mempty ppComm Commutative = text "commutative " -instance PrettyLore lore => Pretty (Reduce lore) where +instance PrettyRep rep => Pretty (Reduce rep) where ppr (Reduce comm red_lam red_nes) = ppComm comm <> ppr red_lam <> comma PP.braces (commasep $ map ppr red_nes) -- | Prettyprint the given histogram operation. ppHist :: - (PrettyLore lore, Pretty inp) => + (PrettyRep rep, Pretty inp) => SubExp -> - [HistOp lore] -> - Lambda lore -> + [HistOp rep] -> + Lambda rep -> [inp] -> Doc ppHist len ops bucket_fun imgs = 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 @@ -41,7 +41,7 @@ import Futhark.MonadFreshNames import qualified Futhark.Optimise.Simplify as Simplify import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Optimise.Simplify.Rule import Futhark.Optimise.Simplify.Rules import Futhark.Optimise.Simplify.Rules.ClosedForm @@ -93,8 +93,8 @@ Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers mempty simplifySOAC :: - Simplify.SimplifiableLore lore => - Simplify.SimplifyOp lore (SOAC lore) + Simplify.SimplifiableRep rep => + Simplify.SimplifyOp rep (SOAC rep) simplifySOAC (Stream outerdim arr form nes lam) = do outerdim' <- Engine.simplify outerdim (form', form_hoisted) <- simplifyStreamForm form @@ -155,10 +155,10 @@ instance BinderOps (Wise SOACS) fixLambdaParams :: - (MonadBinder m, Bindable (Lore m), BinderOps (Lore m)) => - AST.Lambda (Lore m) -> + (MonadBinder m, Bindable (Rep m), BinderOps (Rep m)) => + AST.Lambda (Rep m) -> [Maybe SubExp] -> - m (AST.Lambda (Lore m)) + m (AST.Lambda (Rep m)) fixLambdaParams lam fixes = do body <- runBodyBinder $ localScope (scopeOfLParams $ lambdaParams lam) $ do @@ -177,7 +177,7 @@ maybeFix p (Just x) = letBindNames [paramName p] $ BasicOp $ SubExp x maybeFix _ Nothing = return () -removeLambdaResults :: [Bool] -> AST.Lambda lore -> AST.Lambda lore +removeLambdaResults :: [Bool] -> AST.Lambda rep -> AST.Lambda rep removeLambdaResults keep lam = lam { lambdaBody = lam_body', @@ -193,11 +193,11 @@ soacRules :: RuleBook (Wise SOACS) soacRules = standardRules <> ruleBook topDownRules bottomUpRules --- | Does this lore contain 'SOAC's in its t'Op's? A lore must be an +-- | Does this rep contain 'SOAC's in its t'Op's? A rep must be an -- instance of this class for the simplification rules to work. -class HasSOAC lore where - asSOAC :: Op lore -> Maybe (SOAC lore) - soacOp :: SOAC lore -> Op lore +class HasSOAC rep where + asSOAC :: Op rep -> Maybe (SOAC rep) + soacOp :: SOAC rep -> Op rep instance HasSOAC (Wise SOACS) where asSOAC = Just @@ -256,11 +256,11 @@ Skip liftIdentityMapping :: - forall lore. - (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => - TopDownRuleOp (Wise lore) + forall rep. + (Bindable rep, Simplify.SimplifiableRep rep, HasSOAC (Wise rep)) => + TopDownRuleOp (Wise rep) liftIdentityMapping _ pat aux op - | Just (Screma w arrs form :: SOAC (Wise lore)) <- asSOAC op, + | Just (Screma w arrs form :: SOAC (Wise rep)) <- asSOAC op, Just fun <- isMapSOAC form = do let inputMap = M.fromList $ zip (map paramName $ lambdaParams fun) arrs free = freeIn $ lambdaBody fun @@ -339,8 +339,8 @@ -- | Remove all arguments to the map that are simply replicates. -- These can be turned into free variables instead. removeReplicateMapping :: - (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => - TopDownRuleOp (Wise lore) + (Bindable rep, Simplify.SimplifiableRep rep, HasSOAC (Wise rep)) => + TopDownRuleOp (Wise rep) removeReplicateMapping vtable pat aux op | Just (Screma w arrs form) <- asSOAC op, Just fun <- isMapSOAC form, @@ -358,13 +358,13 @@ removeReplicateWrite _ _ _ _ = Skip removeReplicateInput :: - Aliased lore => - ST.SymbolTable lore -> - AST.Lambda lore -> + Aliased rep => + ST.SymbolTable rep -> + AST.Lambda rep -> [VName] -> Maybe - ( [([VName], Certificates, AST.Exp lore)], - AST.Lambda lore, + ( [([VName], Certificates, AST.Exp rep)], + AST.Lambda rep, [VName] ) removeReplicateInput vtable fun arrs @@ -659,8 +659,8 @@ -- For now we just remove singleton SOACs. simplifyKnownIterationSOAC :: - (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => - TopDownRuleOp (Wise lore) + (Bindable rep, Simplify.SimplifiableRep rep, HasSOAC (Wise rep)) => + TopDownRuleOp (Wise rep) simplifyKnownIterationSOAC _ pat _ op | Just (Screma (Constant k) arrs (ScremaForm scans reds map_lam)) <- asSOAC op, oneIsh k = Simplify $ do diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs --- a/src/Futhark/IR/SegOp.hs +++ b/src/Futhark/IR/SegOp.hs @@ -14,6 +14,7 @@ ( SegOp (..), SegVirt (..), segLevel, + segBody, segSpace, typeCheckSegOp, SegSpace (..), @@ -79,7 +80,7 @@ import Futhark.IR.Mem import Futhark.IR.Prop.Aliases import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Optimise.Simplify.Rule import Futhark.Tools import Futhark.Transform.Rename @@ -121,7 +122,7 @@ SplitStrided <$> rename stride -- | An operator for 'SegHist'. -data HistOp lore = HistOp +data HistOp rep = HistOp { histWidth :: SubExp, histRaceFactor :: SubExp, histDest :: [VName], @@ -132,14 +133,14 @@ -- "dimensions". This is used to generate more efficient -- code. histShape :: Shape, - histOp :: Lambda lore + histOp :: Lambda rep } deriving (Eq, Ord, Show) -- | The type of a histogram produced by a 'HistOp'. This can be -- different from the type of the 'histDest's in case we are -- dealing with a segmented histogram. -histType :: HistOp lore -> [Type] +histType :: HistOp rep -> [Type] histType op = map ( (`arrayOfRow` histWidth op) @@ -148,9 +149,9 @@ $ lambdaReturnType $ histOp op -- | An operator for 'SegScan' and 'SegRed'. -data SegBinOp lore = SegBinOp +data SegBinOp rep = SegBinOp { segBinOpComm :: Commutativity, - segBinOpLambda :: Lambda lore, + segBinOpLambda :: Lambda rep, segBinOpNeutral :: [SubExp], -- | In case this operator is semantically a vectorised -- operator (corresponding to a perfect map nest in the @@ -162,26 +163,26 @@ deriving (Eq, Ord, Show) -- | How many reduction results are produced by these 'SegBinOp's? -segBinOpResults :: [SegBinOp lore] -> Int +segBinOpResults :: [SegBinOp rep] -> Int segBinOpResults = sum . map (length . segBinOpNeutral) -- | Split some list into chunks equal to the number of values -- returned by each 'SegBinOp' -segBinOpChunks :: [SegBinOp lore] -> [a] -> [[a]] +segBinOpChunks :: [SegBinOp rep] -> [a] -> [[a]] segBinOpChunks = chunks . map (length . segBinOpNeutral) -- | The body of a 'SegOp'. -data KernelBody lore = KernelBody - { kernelBodyLore :: BodyDec lore, - kernelBodyStms :: Stms lore, +data KernelBody rep = KernelBody + { kernelBodyDec :: BodyDec rep, + kernelBodyStms :: Stms rep, kernelBodyResult :: [KernelResult] } -deriving instance Decorations lore => Ord (KernelBody lore) +deriving instance RepTypes rep => Ord (KernelBody rep) -deriving instance Decorations lore => Show (KernelBody lore) +deriving instance RepTypes rep => Show (KernelBody rep) -deriving instance Decorations lore => Eq (KernelBody lore) +deriving instance RepTypes rep => Eq (KernelBody rep) -- | Metadata about whether there is a subtle point to this -- 'KernelResult'. This is used to protect things like tiling, which @@ -249,13 +250,13 @@ freeIn' (RegTileReturns dims_n_tiles v) = freeIn' dims_n_tiles <> freeIn' v -instance ASTLore lore => FreeIn (KernelBody lore) where +instance ASTRep rep => FreeIn (KernelBody rep) where freeIn' (KernelBody dec stms res) = fvBind bound_in_stms $ freeIn' dec <> freeIn' stms <> freeIn' res where bound_in_stms = foldMap boundByStm stms -instance ASTLore lore => Substitute (KernelBody lore) where +instance ASTRep rep => Substitute (KernelBody rep) where substituteNames subst (KernelBody dec stms res) = KernelBody (substituteNames subst dec) @@ -283,7 +284,7 @@ (substituteNames subst dims_n_tiles) (substituteNames subst v) -instance ASTLore lore => Rename (KernelBody lore) where +instance ASTRep rep => Rename (KernelBody rep) where rename (KernelBody dec stms res) = do dec' <- rename dec renamingStms stms $ \stms' -> @@ -294,35 +295,35 @@ -- | Perform alias analysis on a 'KernelBody'. aliasAnalyseKernelBody :: - ( ASTLore lore, - CanBeAliased (Op lore) + ( ASTRep rep, + CanBeAliased (Op rep) ) => AliasTable -> - KernelBody lore -> - KernelBody (Aliases lore) + KernelBody rep -> + KernelBody (Aliases rep) aliasAnalyseKernelBody aliases (KernelBody dec stms res) = let Body dec' stms' _ = Alias.analyseBody aliases $ Body dec stms [] in KernelBody dec' stms' res removeKernelBodyAliases :: - CanBeAliased (Op lore) => - KernelBody (Aliases lore) -> - KernelBody lore + CanBeAliased (Op rep) => + KernelBody (Aliases rep) -> + KernelBody rep removeKernelBodyAliases (KernelBody (_, dec) stms res) = KernelBody dec (fmap removeStmAliases stms) res removeKernelBodyWisdom :: - CanBeWise (Op lore) => - KernelBody (Wise lore) -> - KernelBody lore + CanBeWise (Op rep) => + KernelBody (Wise rep) -> + KernelBody rep removeKernelBodyWisdom (KernelBody dec stms res) = let Body dec' stms' _ = removeBodyWisdom $ Body dec stms [] in KernelBody dec' stms' res -- | The variables consumed in the kernel body. consumedInKernelBody :: - Aliased lore => - KernelBody lore -> + Aliased rep => + KernelBody rep -> Names consumedInKernelBody (KernelBody dec stms res) = consumedInBody (Body dec stms []) <> mconcat (map consumedByReturn res) @@ -331,12 +332,12 @@ consumedByReturn _ = mempty checkKernelBody :: - TC.Checkable lore => + TC.Checkable rep => [Type] -> - KernelBody (Aliases lore) -> - TC.TypeM lore () + KernelBody (Aliases rep) -> + TC.TypeM rep () checkKernelBody ts (KernelBody (_, dec) stms kres) = do - TC.checkBodyLore dec + TC.checkBodyDec dec -- We consume the kernel results (when applicable) before -- type-checking the stms, so we will get an error if a statement -- uses an array that is written to in a result. @@ -408,10 +409,10 @@ (dims, blk_tiles, reg_tiles) = unzip3 dims_n_tiles expected = t `arrayOfShape` Shape (blk_tiles ++ reg_tiles) -kernelBodyMetrics :: OpMetrics (Op lore) => KernelBody lore -> MetricsM () +kernelBodyMetrics :: OpMetrics (Op rep) => KernelBody rep -> MetricsM () kernelBodyMetrics = mapM_ stmMetrics . kernelBodyStms -instance PrettyLore lore => Pretty (KernelBody lore) where +instance PrettyRep rep => Pretty (KernelBody rep) where ppr (KernelBody _ stms res) = PP.stack (map ppr (stmsToList stms)) text "return" <+> PP.braces (PP.commasep $ map ppr res) @@ -477,11 +478,11 @@ -- | A 'Scope' containing all the identifiers brought into scope by -- this 'SegSpace'. -scopeOfSegSpace :: SegSpace -> Scope lore +scopeOfSegSpace :: SegSpace -> Scope rep scopeOfSegSpace (SegSpace phys space) = M.fromList $ zip (phys : map fst space) $ repeat $ IndexName Int64 -checkSegSpace :: TC.Checkable lore => SegSpace -> TC.TypeM lore () +checkSegSpace :: TC.Checkable rep => SegSpace -> TC.TypeM rep () checkSegSpace (SegSpace _ dims) = mapM_ (TC.require [Prim int64] . snd) dims @@ -495,29 +496,38 @@ -- of information. For example, in GPU backends, it is used to -- indicate whether the 'SegOp' is expected to run at the thread-level -- or the group-level. -data SegOp lvl lore - = SegMap lvl SegSpace [Type] (KernelBody lore) +data SegOp lvl rep + = SegMap lvl SegSpace [Type] (KernelBody rep) | -- | The KernelSpace must always have at least two dimensions, -- implying that the result of a SegRed is always an array. - SegRed lvl SegSpace [SegBinOp lore] [Type] (KernelBody lore) - | SegScan lvl SegSpace [SegBinOp lore] [Type] (KernelBody lore) - | SegHist lvl SegSpace [HistOp lore] [Type] (KernelBody lore) + SegRed lvl SegSpace [SegBinOp rep] [Type] (KernelBody rep) + | SegScan lvl SegSpace [SegBinOp rep] [Type] (KernelBody rep) + | SegHist lvl SegSpace [HistOp rep] [Type] (KernelBody rep) deriving (Eq, Ord, Show) -- | The level of a 'SegOp'. -segLevel :: SegOp lvl lore -> lvl +segLevel :: SegOp lvl rep -> lvl segLevel (SegMap lvl _ _ _) = lvl segLevel (SegRed lvl _ _ _ _) = lvl segLevel (SegScan lvl _ _ _ _) = lvl segLevel (SegHist lvl _ _ _ _) = lvl -- | The space of a 'SegOp'. -segSpace :: SegOp lvl lore -> SegSpace +segSpace :: SegOp lvl rep -> SegSpace segSpace (SegMap _ lvl _ _) = lvl segSpace (SegRed _ lvl _ _ _) = lvl segSpace (SegScan _ lvl _ _ _) = lvl segSpace (SegHist _ lvl _ _ _) = lvl +-- | The body of a 'SegOp'. +segBody :: SegOp lvl rep -> KernelBody rep +segBody segop = + case segop of + SegMap _ _ _ body -> body + SegRed _ _ _ _ body -> body + SegScan _ _ _ _ body -> body + SegHist _ _ _ _ body -> body + segResultShape :: SegSpace -> Type -> KernelResult -> Type segResultShape _ t (WriteReturns shape _ _) = t `arrayOfShape` shape @@ -531,7 +541,7 @@ t `arrayOfShape` Shape (map (\(dim, _, _) -> dim) dims_n_tiles) -- | The return type of a 'SegOp'. -segOpType :: SegOp lvl lore -> [Type] +segOpType :: SegOp lvl rep -> [Type] segOpType (SegMap _ space ts kbody) = zipWith (segResultShape space) ts $ kernelBodyResult kbody segOpType (SegRed _ space reds ts kbody) = @@ -567,12 +577,12 @@ dims = segSpaceDims space segment_dims = init dims -instance TypedOp (SegOp lvl lore) where +instance TypedOp (SegOp lvl rep) where opType = pure . staticShapes . segOpType instance - (ASTLore lore, Aliased lore, ASTConstraints lvl) => - AliasedOp (SegOp lvl lore) + (ASTRep rep, Aliased rep, ASTConstraints lvl) => + AliasedOp (SegOp lvl rep) where opAliases = map (const mempty) . segOpType @@ -587,10 +597,10 @@ -- | Type check a 'SegOp', given a checker for its level. typeCheckSegOp :: - TC.Checkable lore => - (lvl -> TC.TypeM lore ()) -> - SegOp lvl (Aliases lore) -> - TC.TypeM lore () + TC.Checkable rep => + (lvl -> TC.TypeM rep ()) -> + SegOp lvl (Aliases rep) -> + TC.TypeM rep () typeCheckSegOp checkLvl (SegMap lvl space ts kbody) = do checkLvl lvl checkScanRed space [] ts kbody @@ -660,12 +670,12 @@ segment_dims = init $ segSpaceDims space checkScanRed :: - TC.Checkable lore => + TC.Checkable rep => SegSpace -> - [(Lambda (Aliases lore), [SubExp], Shape)] -> + [(Lambda (Aliases rep), [SubExp], Shape)] -> [Type] -> - KernelBody (Aliases lore) -> - TC.TypeM lore () + KernelBody (Aliases rep) -> + TC.TypeM rep () checkScanRed space ops ts kbody = do checkSegSpace space mapM_ TC.checkType ts @@ -698,16 +708,16 @@ checkKernelBody ts kbody -- | Like 'Mapper', but just for 'SegOp's. -data SegOpMapper lvl flore tlore m = SegOpMapper +data SegOpMapper lvl frep trep m = SegOpMapper { mapOnSegOpSubExp :: SubExp -> m SubExp, - mapOnSegOpLambda :: Lambda flore -> m (Lambda tlore), - mapOnSegOpBody :: KernelBody flore -> m (KernelBody tlore), + mapOnSegOpLambda :: Lambda frep -> m (Lambda trep), + mapOnSegOpBody :: KernelBody frep -> m (KernelBody trep), mapOnSegOpVName :: VName -> m VName, mapOnSegOpLevel :: lvl -> m lvl } -- | A mapper that simply returns the 'SegOp' verbatim. -identitySegOpMapper :: Monad m => SegOpMapper lvl lore lore m +identitySegOpMapper :: Monad m => SegOpMapper lvl rep rep m identitySegOpMapper = SegOpMapper { mapOnSegOpSubExp = return, @@ -719,7 +729,7 @@ mapOnSegSpace :: Monad f => - SegOpMapper lvl flore tlore f -> + SegOpMapper lvl frep trep f -> SegSpace -> f SegSpace mapOnSegSpace tv (SegSpace phys dims) = @@ -727,9 +737,9 @@ mapSegBinOp :: Monad m => - SegOpMapper lvl flore tlore m -> - SegBinOp flore -> - m (SegBinOp tlore) + SegOpMapper lvl frep trep m -> + SegBinOp frep -> + m (SegBinOp trep) mapSegBinOp tv (SegBinOp comm red_op nes shape) = SegBinOp comm <$> mapOnSegOpLambda tv red_op @@ -739,9 +749,9 @@ -- | Apply a 'SegOpMapper' to the given 'SegOp'. mapSegOpM :: (Applicative m, Monad m) => - SegOpMapper lvl flore tlore m -> - SegOp lvl flore -> - m (SegOp lvl tlore) + SegOpMapper lvl frep trep m -> + SegOp lvl frep -> + m (SegOp lvl trep) mapSegOpM tv (SegMap lvl space ts body) = SegMap <$> mapOnSegOpLevel tv lvl @@ -780,7 +790,7 @@ mapOnSegOpType :: Monad m => - SegOpMapper lvl flore tlore m -> + SegOpMapper lvl frep trep m -> Type -> m Type mapOnSegOpType _tv t@Prim {} = pure t @@ -795,8 +805,8 @@ mapOnSegOpType _tv (Mem s) = pure $ Mem s instance - (ASTLore lore, Substitute lvl) => - Substitute (SegOp lvl lore) + (ASTRep rep, Substitute lvl) => + Substitute (SegOp lvl rep) where substituteNames subst = runIdentity . mapSegOpM substitute where @@ -810,16 +820,16 @@ } instance - (ASTLore lore, ASTConstraints lvl) => - Rename (SegOp lvl lore) + (ASTRep rep, ASTConstraints lvl) => + Rename (SegOp lvl rep) where rename = mapSegOpM renamer where renamer = SegOpMapper rename rename rename rename rename instance - (ASTLore lore, FreeIn (LParamInfo lore), FreeIn lvl) => - FreeIn (SegOp lvl lore) + (ASTRep rep, FreeIn (LParamInfo rep), FreeIn lvl) => + FreeIn (SegOp lvl rep) where freeIn' e = flip execState mempty $ mapSegOpM free e where @@ -833,7 +843,7 @@ mapOnSegOpLevel = walk freeIn' } -instance OpMetrics (Op lore) => OpMetrics (SegOp lvl lore) where +instance OpMetrics (Op rep) => OpMetrics (SegOp lvl rep) where opMetrics (SegMap _ _ _ body) = inside "SegMap" $ kernelBodyMetrics body opMetrics (SegRed _ _ reds _ body) = @@ -858,7 +868,7 @@ ) <+> parens (text "~" <> ppr phys) -instance PrettyLore lore => Pretty (SegBinOp lore) where +instance PrettyRep rep => Pretty (SegBinOp rep) where ppr (SegBinOp comm lam nes shape) = PP.braces (PP.commasep $ map ppr nes) <> PP.comma ppr shape <> PP.comma @@ -868,7 +878,7 @@ Commutative -> text "commutative " Noncommutative -> mempty -instance (PrettyLore lore, PP.Pretty lvl) => PP.Pretty (SegOp lvl lore) where +instance (PrettyRep rep, PP.Pretty lvl) => PP.Pretty (SegOp lvl rep) where ppr (SegMap lvl space ts body) = text "segmap" <> ppr lvl PP.align (ppr space) @@ -905,14 +915,14 @@ ppr op instance - ( ASTLore lore, - ASTLore (Aliases lore), - CanBeAliased (Op lore), + ( ASTRep rep, + ASTRep (Aliases rep), + CanBeAliased (Op rep), ASTConstraints lvl ) => - CanBeAliased (SegOp lvl lore) + CanBeAliased (SegOp lvl rep) where - type OpWithAliases (SegOp lvl lore) = SegOp lvl (Aliases lore) + type OpWithAliases (SegOp lvl rep) = SegOp lvl (Aliases rep) addOpAliases aliases = runIdentity . mapSegOpM alias where @@ -935,10 +945,10 @@ return instance - (CanBeWise (Op lore), ASTLore lore, ASTConstraints lvl) => - CanBeWise (SegOp lvl lore) + (CanBeWise (Op rep), ASTRep rep, ASTConstraints lvl) => + CanBeWise (SegOp lvl rep) where - type OpWithWisdom (SegOp lvl lore) = SegOp lvl (Wise lore) + type OpWithWisdom (SegOp lvl rep) = SegOp lvl (Wise rep) removeOpWisdom = runIdentity . mapSegOpM remove where @@ -950,7 +960,7 @@ return return -instance ASTLore lore => ST.IndexOp (SegOp lvl lore) where +instance ASTRep rep => ST.IndexOp (SegOp lvl rep) where indexOp vtable k (SegMap _ space _ kbody) is = do Returns ResultMaySimplify se <- maybeNth k $ kernelBodyResult kbody guard $ length gtids <= length is @@ -995,8 +1005,8 @@ indexOp _ _ _ _ = Nothing instance - (ASTLore lore, ASTConstraints lvl) => - IsOp (SegOp lvl lore) + (ASTRep rep, ASTConstraints lvl) => + IsOp (SegOp lvl rep) where cheapOp _ = False safeOp _ = True @@ -1032,11 +1042,11 @@ <*> Engine.simplify what mkWiseKernelBody :: - (ASTLore lore, CanBeWise (Op lore)) => - BodyDec lore -> - Stms (Wise lore) -> + (ASTRep rep, CanBeWise (Op rep)) => + BodyDec rep -> + Stms (Wise rep) -> [KernelResult] -> - KernelBody (Wise lore) + KernelBody (Wise rep) mkWiseKernelBody dec bnds res = let Body dec' _ _ = mkWiseBody dec bnds res_vs in KernelBody dec' bnds res @@ -1045,9 +1055,9 @@ mkKernelBodyM :: MonadBinder m => - Stms (Lore m) -> + Stms (Rep m) -> [KernelResult] -> - m (KernelBody (Lore m)) + m (KernelBody (Rep m)) mkKernelBodyM stms kres = do Body dec' _ _ <- mkBodyM stms res_ses return $ KernelBody dec' stms kres @@ -1055,10 +1065,10 @@ res_ses = map kernelResultSubExp kres simplifyKernelBody :: - (Engine.SimplifiableLore lore, BodyDec lore ~ ()) => + (Engine.SimplifiableRep rep, BodyDec rep ~ ()) => SegSpace -> - KernelBody lore -> - Engine.SimpleM lore (KernelBody (Wise lore), Stms (Wise lore)) + KernelBody rep -> + Engine.SimpleM rep (KernelBody (Wise rep), Stms (Wise rep)) simplifyKernelBody space (KernelBody _ stms res) = do par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers @@ -1090,16 +1100,16 @@ consumedInResult _ = [] -segSpaceSymbolTable :: ASTLore lore => SegSpace -> ST.SymbolTable lore +segSpaceSymbolTable :: ASTRep rep => SegSpace -> ST.SymbolTable rep segSpaceSymbolTable (SegSpace flat gtids_and_dims) = foldl' f (ST.fromScope $ M.singleton flat $ IndexName Int64) gtids_and_dims where f vtable (gtid, dim) = ST.insertLoopVar gtid Int64 dim vtable simplifySegBinOp :: - Engine.SimplifiableLore lore => - SegBinOp lore -> - Engine.SimpleM lore (SegBinOp (Wise lore), Stms (Wise lore)) + Engine.SimplifiableRep rep => + SegBinOp rep -> + Engine.SimpleM rep (SegBinOp (Wise rep), Stms (Wise rep)) simplifySegBinOp (SegBinOp comm lam nes shape) = do (lam', hoisted) <- Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $ @@ -1110,12 +1120,12 @@ -- | Simplify the given 'SegOp'. simplifySegOp :: - ( Engine.SimplifiableLore lore, - BodyDec lore ~ (), + ( Engine.SimplifiableRep rep, + BodyDec rep ~ (), Engine.Simplifiable lvl ) => - SegOp lvl lore -> - Engine.SimpleM lore (SegOp lvl (Wise lore), Stms (Wise lore)) + SegOp lvl rep -> + Engine.SimpleM rep (SegOp lvl (Wise rep), Stms (Wise rep)) simplifySegOp (SegMap lvl space ts kbody) = do (lvl', space', ts') <- Engine.simplify (lvl, space, ts) (kbody', body_hoisted) <- simplifyKernelBody space kbody @@ -1181,23 +1191,23 @@ scope = scopeOfSegSpace space scope_vtable = ST.fromScope scope --- | Does this lore contain 'SegOp's in its t'Op's? A lore must be an +-- | Does this rep contain 'SegOp's in its t'Op's? A rep must be an -- instance of this class for the simplification rules to work. -class HasSegOp lore where - type SegOpLevel lore - asSegOp :: Op lore -> Maybe (SegOp (SegOpLevel lore) lore) - segOp :: SegOp (SegOpLevel lore) lore -> Op lore +class HasSegOp rep where + type SegOpLevel rep + asSegOp :: Op rep -> Maybe (SegOp (SegOpLevel rep) rep) + segOp :: SegOp (SegOpLevel rep) rep -> Op rep -- | Simplification rules for simplifying 'SegOp's. segOpRules :: - (HasSegOp lore, BinderOps lore, Bindable lore) => - RuleBook lore + (HasSegOp rep, BinderOps rep, Bindable rep) => + RuleBook rep segOpRules = ruleBook [RuleOp segOpRuleTopDown] [RuleOp segOpRuleBottomUp] segOpRuleTopDown :: - (HasSegOp lore, BinderOps lore, Bindable lore) => - TopDownRuleOp lore + (HasSegOp rep, BinderOps rep, Bindable rep) => + TopDownRuleOp rep segOpRuleTopDown vtable pat dec op | Just op' <- asSegOp op = topDownSegOp vtable pat dec op' @@ -1205,8 +1215,8 @@ Skip segOpRuleBottomUp :: - (HasSegOp lore, BinderOps lore) => - BottomUpRuleOp lore + (HasSegOp rep, BinderOps rep) => + BottomUpRuleOp rep segOpRuleBottomUp vtable pat dec op | Just op' <- asSegOp op = bottomUpSegOp vtable pat dec op' @@ -1214,12 +1224,12 @@ Skip topDownSegOp :: - (HasSegOp lore, BinderOps lore, Bindable lore) => - ST.SymbolTable lore -> - Pattern lore -> - StmAux (ExpDec lore) -> - SegOp (SegOpLevel lore) lore -> - Rule lore + (HasSegOp rep, BinderOps rep, Bindable rep) => + ST.SymbolTable rep -> + Pattern rep -> + StmAux (ExpDec rep) -> + SegOp (SegOpLevel rep) rep -> + Rule rep -- If a SegOp produces something invariant to the SegOp, turn it -- into a replicate. topDownSegOp vtable (Pattern [] kpes) dec (SegMap lvl space ts (KernelBody _ kstms kres)) = Simplify $ do @@ -1308,11 +1318,11 @@ -- A convenient way of operating on the type and body of a SegOp, -- without worrying about exactly what kind it is. segOpGuts :: - SegOp (SegOpLevel lore) lore -> + SegOp (SegOpLevel rep) rep -> ( [Type], - KernelBody lore, + KernelBody rep, Int, - [Type] -> KernelBody lore -> SegOp (SegOpLevel lore) lore + [Type] -> KernelBody rep -> SegOp (SegOpLevel rep) rep ) segOpGuts (SegMap lvl space kts body) = (kts, body, 0, SegMap lvl space) @@ -1324,12 +1334,12 @@ (kts, body, sum $ map (length . histDest) ops, SegHist lvl space ops) bottomUpSegOp :: - (HasSegOp lore, BinderOps lore) => - (ST.SymbolTable lore, UT.UsageTable) -> - Pattern lore -> - StmAux (ExpDec lore) -> - SegOp (SegOpLevel lore) lore -> - Rule lore + (HasSegOp rep, BinderOps rep) => + (ST.SymbolTable rep, UT.UsageTable) -> + Pattern rep -> + StmAux (ExpDec rep) -> + SegOp (SegOpLevel rep) rep -> + Rule rep -- Some SegOp results can be moved outside the SegOp, which can -- simplify further analysis. bottomUpSegOp (vtable, used) (Pattern [] kpes) dec segop = Simplify $ do @@ -1419,8 +1429,8 @@ --- Memory kernelBodyReturns :: - (Mem lore, HasScope lore m, Monad m) => - KernelBody lore -> + (Mem rep, HasScope rep m, Monad m) => + KernelBody rep -> [ExpReturns] -> m [ExpReturns] kernelBodyReturns = zipWithM correct . kernelBodyResult @@ -1430,8 +1440,8 @@ -- | Like 'segOpType', but for memory representations. segOpReturns :: - (Mem lore, Monad m, HasScope lore m) => - SegOp lvl lore -> + (Mem rep, Monad m, HasScope rep m) => + SegOp lvl rep -> m [ExpReturns] segOpReturns k@(SegMap _ _ _ kbody) = kernelBodyReturns kbody . extReturns =<< opType k diff --git a/src/Futhark/IR/Seq.hs b/src/Futhark/IR/Seq.hs --- a/src/Futhark/IR/Seq.hs +++ b/src/Futhark/IR/Seq.hs @@ -3,8 +3,7 @@ -- | A sequential representation. module Futhark.IR.Seq - ( -- * The Lore definition - Seq, + ( Seq, -- * Simplification simplifyProg, @@ -32,10 +31,10 @@ -- | The phantom type for the Seq representation. data Seq -instance Decorations Seq where +instance RepTypes Seq where type Op Seq = () -instance ASTLore Seq where +instance ASTRep Seq where expTypesFromPattern = return . expExtTypesFromPattern instance TypeCheck.CheckableOp Seq where @@ -51,7 +50,7 @@ instance BinderOps Seq -instance PrettyLore Seq +instance PrettyRep Seq instance BinderOps (Engine.Wise Seq) diff --git a/src/Futhark/IR/SeqMem.hs b/src/Futhark/IR/SeqMem.hs --- a/src/Futhark/IR/SeqMem.hs +++ b/src/Futhark/IR/SeqMem.hs @@ -13,12 +13,10 @@ -- * Module re-exports module Futhark.IR.Mem, - module Futhark.IR.Kernels.Kernel, ) where import Futhark.Analysis.PrimExp.Convert -import Futhark.IR.Kernels.Kernel import Futhark.IR.Mem import Futhark.IR.Mem.Simplify import qualified Futhark.Optimise.Simplify.Engine as Engine @@ -28,7 +26,7 @@ data SeqMem -instance Decorations SeqMem where +instance RepTypes SeqMem where type LetDec SeqMem = LetDecMem type FParamInfo SeqMem = FParamMem type LParamInfo SeqMem = LParamMem @@ -36,14 +34,14 @@ type BranchType SeqMem = BranchTypeMem type Op SeqMem = MemOp () -instance ASTLore SeqMem where +instance ASTRep SeqMem where expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern instance OpReturns SeqMem where opReturns (Alloc _ space) = return [MemMem space] opReturns (Inner ()) = pure [] -instance PrettyLore SeqMem +instance PrettyRep SeqMem instance TC.CheckableOp SeqMem where checkOp (Alloc size _) = @@ -52,9 +50,9 @@ pure () instance TC.Checkable SeqMem where - checkFParamLore = checkMemInfo - checkLParamLore = checkMemInfo - checkLetBoundLore = checkMemInfo + checkFParamDec = checkMemInfo + checkLParamDec = checkMemInfo + checkLetBoundDec = checkMemInfo checkRetType = mapM_ (TC.checkExtType . declExtTypeOf) primFParam name t = return $ Param name (MemPrim t) matchPattern = matchPatternToExp diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs --- a/src/Futhark/IR/Syntax.hs +++ b/src/Futhark/IR/Syntax.hs @@ -71,37 +71,37 @@ -- in {b_13} -- @ -- --- == Lores +-- == Representations -- -- Most AST types ('Stm', 'ExpT', t'Prog', etc) are parameterised by a --- type parameter with the somewhat silly name @lore@. The lore --- specifies how to fill out various polymorphic parts of the AST. --- For example, 'ExpT' has a constructor v'Op' whose payload depends --- on @lore@, via the use of a type family called t'Op' (a kind of --- type-level function) which is applied to the @lore@. The SOACS --- representation ("Futhark.IR.SOACS") thus uses a lore --- called @SOACS@, and defines that @Op SOACS@ is a SOAC, while the --- Kernels representation ("Futhark.IR.Kernels") defines --- @Op Kernels@ as some kind of kernel construct. Similarly, various --- other decorations (e.g. what information we store in a t'PatElemT') --- are also type families. +-- type parameter @rep@. The representation specifies how to fill out +-- various polymorphic parts of the AST. For example, 'ExpT' has a +-- constructor v'Op' whose payload depends on @rep@, via the use of a +-- type family called t'Op' (a kind of type-level function) which is +-- applied to the @rep@. The SOACS representation +-- ("Futhark.IR.SOACS") thus uses a rep called @SOACS@, and defines +-- that @Op SOACS@ is a SOAC, while the Kernels representation +-- ("Futhark.IR.Kernels") defines @Op Kernels@ as some kind of kernel +-- construct. Similarly, various other decorations (e.g. what +-- information we store in a t'PatElemT') are also type families. -- -- The full list of possible decorations is defined as part of the --- type class 'Decorations' (although other type families are also +-- type class 'RepTypes' (although other type families are also -- used elsewhere in the compiler on an ad hoc basis). -- --- Essentially, the @lore@ type parameter functions as a kind of +-- Essentially, the @rep@ type parameter functions as a kind of -- proxy, saving us from having to parameterise the AST type with all -- the different forms of decorations that we desire (it would easily -- become a type with a dozen type parameters). -- --- Defining a new representation (or /lore/) thus requires you to +-- Defining a new representation (or /rep/) thus requires you to -- define an empty datatype and implement a handful of type class -- instances for it. See the source of "Futhark.IR.Seq" -- for what is likely the simplest example. module Futhark.IR.Syntax ( module Language.Futhark.Core, - module Futhark.IR.Decorations, + pretty, + module Futhark.IR.Rep, module Futhark.IR.Syntax.Core, -- * Types @@ -172,8 +172,9 @@ import qualified Data.Set as S import Data.String import Data.Traversable (fmapDefault, foldMapDefault) -import Futhark.IR.Decorations +import Futhark.IR.Rep import Futhark.IR.Syntax.Core +import Futhark.Util.Pretty (pretty) import Language.Futhark.Core import Prelude hiding (id, (.)) @@ -204,7 +205,7 @@ withoutAttrs (Attrs x) (Attrs y) = Attrs $ x `S.difference` y -- | A type alias for namespace control. -type PatElem lore = PatElemT (LetDec lore) +type PatElem rep = PatElemT (LetDec rep) -- | A pattern is conceptually just a list of names and their types. data PatternT dec = Pattern @@ -232,7 +233,7 @@ Pattern <$> traverse (traverse f) ctx <*> traverse (traverse f) vals -- | A type alias for namespace control. -type Pattern lore = PatternT (LetDec lore) +type Pattern rep = PatternT (LetDec rep) -- | Auxilliary Information associated with a statement. data StmAux dec = StmAux @@ -247,38 +248,38 @@ StmAux (cs1 <> cs2) (attrs1 <> attrs2) (dec1 <> dec2) -- | A local variable binding. -data Stm lore = Let +data Stm rep = Let { -- | Pattern. - stmPattern :: Pattern lore, + stmPattern :: Pattern rep, -- | Auxiliary information statement. - stmAux :: StmAux (ExpDec lore), + stmAux :: StmAux (ExpDec rep), -- | Expression. - stmExp :: Exp lore + stmExp :: Exp rep } -deriving instance Decorations lore => Ord (Stm lore) +deriving instance RepTypes rep => Ord (Stm rep) -deriving instance Decorations lore => Show (Stm lore) +deriving instance RepTypes rep => Show (Stm rep) -deriving instance Decorations lore => Eq (Stm lore) +deriving instance RepTypes rep => Eq (Stm rep) -- | A sequence of statements. -type Stms lore = Seq.Seq (Stm lore) +type Stms rep = Seq.Seq (Stm rep) -- | A single statement. -oneStm :: Stm lore -> Stms lore +oneStm :: Stm rep -> Stms rep oneStm = Seq.singleton -- | Convert a statement list to a statement sequence. -stmsFromList :: [Stm lore] -> Stms lore +stmsFromList :: [Stm rep] -> Stms rep stmsFromList = Seq.fromList -- | Convert a statement sequence to a statement list. -stmsToList :: Stms lore -> [Stm lore] +stmsToList :: Stms rep -> [Stm rep] stmsToList = toList -- | The first statement in the sequence, if any. -stmsHead :: Stms lore -> Maybe (Stm lore, Stms lore) +stmsHead :: Stms rep -> Maybe (Stm rep, Stms rep) stmsHead stms = case Seq.viewl stms of stm Seq.:< stms' -> Just (stm, stms') Seq.EmptyL -> Nothing @@ -288,17 +289,17 @@ -- | A body consists of a number of bindings, terminating in a result -- (essentially a tuple literal). -data BodyT lore = Body - { bodyDec :: BodyDec lore, - bodyStms :: Stms lore, +data BodyT rep = Body + { bodyDec :: BodyDec rep, + bodyStms :: Stms rep, bodyResult :: Result } -deriving instance Decorations lore => Ord (BodyT lore) +deriving instance RepTypes rep => Ord (BodyT rep) -deriving instance Decorations lore => Show (BodyT lore) +deriving instance RepTypes rep => Show (BodyT rep) -deriving instance Decorations lore => Eq (BodyT lore) +deriving instance RepTypes rep => Eq (BodyT rep) -- | Type alias for namespace reasons. type Body = BodyT @@ -405,42 +406,42 @@ deriving (Eq, Ord, Show) -- | The root Futhark expression type. The v'Op' constructor contains --- a lore-specific operation. Do-loops, branches and function calls +-- a rep-specific operation. Do-loops, branches and function calls -- are special. Everything else is a simple t'BasicOp'. -data ExpT lore +data ExpT rep = -- | A simple (non-recursive) operation. BasicOp BasicOp - | Apply Name [(SubExp, Diet)] [RetType lore] (Safety, SrcLoc, [SrcLoc]) - | If SubExp (BodyT lore) (BodyT lore) (IfDec (BranchType lore)) + | Apply Name [(SubExp, Diet)] [RetType rep] (Safety, SrcLoc, [SrcLoc]) + | If SubExp (BodyT rep) (BodyT rep) (IfDec (BranchType rep)) | -- | @loop {a} = {v} (for i < n|while b) do b@. The merge -- parameters are divided into context and value part. - DoLoop [(FParam lore, SubExp)] [(FParam lore, SubExp)] (LoopForm lore) (BodyT lore) + DoLoop [(FParam rep, SubExp)] [(FParam rep, SubExp)] (LoopForm rep) (BodyT rep) | -- | Create accumulators backed by the given arrays (which are -- consumed) and pass them to the lambda, which must return the -- updated accumulators and possibly some extra values. The -- accumulators are turned back into arrays. The 'Shape' is the -- write index space. The corresponding arrays must all have this -- shape outermost. This construct is not part of 'BasicOp' - -- because we need the @lore@ parameter. - WithAcc [(Shape, [VName], Maybe (Lambda lore, [SubExp]))] (Lambda lore) - | Op (Op lore) + -- because we need the @rep@ parameter. + WithAcc [(Shape, [VName], Maybe (Lambda rep, [SubExp]))] (Lambda rep) + | Op (Op rep) -deriving instance Decorations lore => Eq (ExpT lore) +deriving instance RepTypes rep => Eq (ExpT rep) -deriving instance Decorations lore => Show (ExpT lore) +deriving instance RepTypes rep => Show (ExpT rep) -deriving instance Decorations lore => Ord (ExpT lore) +deriving instance RepTypes rep => Ord (ExpT rep) -- | For-loop or while-loop? -data LoopForm lore - = ForLoop VName IntType SubExp [(LParam lore, VName)] +data LoopForm rep + = ForLoop VName IntType SubExp [(LParam rep, VName)] | WhileLoop VName -deriving instance Decorations lore => Eq (LoopForm lore) +deriving instance RepTypes rep => Eq (LoopForm rep) -deriving instance Decorations lore => Show (LoopForm lore) +deriving instance RepTypes rep => Show (LoopForm rep) -deriving instance Decorations lore => Ord (LoopForm lore) +deriving instance RepTypes rep => Ord (LoopForm rep) -- | Data associated with a branch. data IfDec rt = IfDec @@ -472,44 +473,44 @@ type Exp = ExpT -- | Anonymous function for use in a SOAC. -data LambdaT lore = Lambda - { lambdaParams :: [LParam lore], - lambdaBody :: BodyT lore, +data LambdaT rep = Lambda + { lambdaParams :: [LParam rep], + lambdaBody :: BodyT rep, lambdaReturnType :: [Type] } -deriving instance Decorations lore => Eq (LambdaT lore) +deriving instance RepTypes rep => Eq (LambdaT rep) -deriving instance Decorations lore => Show (LambdaT lore) +deriving instance RepTypes rep => Show (LambdaT rep) -deriving instance Decorations lore => Ord (LambdaT lore) +deriving instance RepTypes rep => Ord (LambdaT rep) -- | Type alias for namespacing reasons. type Lambda = LambdaT -- | A function and loop parameter. -type FParam lore = Param (FParamInfo lore) +type FParam rep = Param (FParamInfo rep) -- | A lambda parameter. -type LParam lore = Param (LParamInfo lore) +type LParam rep = Param (LParamInfo rep) -- | Function Declarations -data FunDef lore = FunDef +data FunDef rep = FunDef { -- | Contains a value if this function is -- an entry point. funDefEntryPoint :: Maybe EntryPoint, funDefAttrs :: Attrs, funDefName :: Name, - funDefRetType :: [RetType lore], - funDefParams :: [FParam lore], - funDefBody :: BodyT lore + funDefRetType :: [RetType rep], + funDefParams :: [FParam rep], + funDefBody :: BodyT rep } -deriving instance Decorations lore => Eq (FunDef lore) +deriving instance RepTypes rep => Eq (FunDef rep) -deriving instance Decorations lore => Show (FunDef lore) +deriving instance RepTypes rep => Show (FunDef rep) -deriving instance Decorations lore => Ord (FunDef lore) +deriving instance RepTypes rep => Ord (FunDef rep) -- | Information about the parameters and return value of an entry -- point. The first element is for parameters, the second for return @@ -521,24 +522,24 @@ data EntryPointType = -- | Is an unsigned integer or array of unsigned -- integers. - TypeUnsigned + TypeUnsigned Uniqueness | -- | A black box type comprising this many core -- values. The string is a human-readable -- description with no other semantics. - TypeOpaque String Int + TypeOpaque Uniqueness String Int | -- | Maps directly. - TypeDirect + TypeDirect Uniqueness deriving (Eq, Show, Ord) -- | An entire Futhark program. -data Prog lore = Prog +data Prog rep = Prog { -- | Top-level constants that are computed at program startup, and -- which are in scope inside all functions. - progConsts :: Stms lore, + progConsts :: Stms rep, -- | The functions comprising the program. All funtions are also -- available in scope in the definitions of the constants, so be -- careful not to introduce circular dependencies (not currently -- checked). - progFuns :: [FunDef lore] + progFuns :: [FunDef rep] } deriving (Eq, Ord, Show) diff --git a/src/Futhark/IR/Syntax/Core.hs b/src/Futhark/IR/Syntax/Core.hs --- a/src/Futhark/IR/Syntax/Core.hs +++ b/src/Futhark/IR/Syntax/Core.hs @@ -5,7 +5,7 @@ -- | The most primitive ("core") aspects of the AST. Split out of -- "Futhark.IR.Syntax" in order for --- "Futhark.IR.Decorations" to use these definitions. This +-- "Futhark.IR.Rep" to use these definitions. This -- module is re-exported from "Futhark.IR.Syntax" and -- there should be no reason to include it explicitly. module Futhark.IR.Syntax.Core @@ -13,6 +13,7 @@ module Futhark.IR.Primitive, -- * Types + Commutativity (..), Uniqueness (..), NoUniqueness (..), ShapeBase (..), @@ -66,6 +67,19 @@ import Futhark.IR.Primitive import Language.Futhark.Core import Prelude hiding (id, (.)) + +-- | Whether some operator is commutative or not. The 'Monoid' +-- instance returns the least commutative of its arguments. +data Commutativity + = Noncommutative + | Commutative + deriving (Eq, Ord, Show) + +instance Semigroup Commutativity where + (<>) = min + +instance Monoid Commutativity where + mempty = Commutative -- | The size of an array type as a list of its dimension sizes, with -- the type of sizes being parametric. diff --git a/src/Futhark/IR/Traversals.hs b/src/Futhark/IR/Traversals.hs --- a/src/Futhark/IR/Traversals.hs +++ b/src/Futhark/IR/Traversals.hs @@ -46,21 +46,21 @@ -- | Express a monad mapping operation on a syntax node. Each element -- of this structure expresses the operation to be performed on a -- given child. -data Mapper flore tlore m = Mapper +data Mapper frep trep m = Mapper { mapOnSubExp :: SubExp -> m SubExp, -- | Most bodies are enclosed in a scope, which is passed along -- for convenience. - mapOnBody :: Scope tlore -> Body flore -> m (Body tlore), + mapOnBody :: Scope trep -> Body frep -> m (Body trep), mapOnVName :: VName -> m VName, - mapOnRetType :: RetType flore -> m (RetType tlore), - mapOnBranchType :: BranchType flore -> m (BranchType tlore), - mapOnFParam :: FParam flore -> m (FParam tlore), - mapOnLParam :: LParam flore -> m (LParam tlore), - mapOnOp :: Op flore -> m (Op tlore) + mapOnRetType :: RetType frep -> m (RetType trep), + mapOnBranchType :: BranchType frep -> m (BranchType trep), + mapOnFParam :: FParam frep -> m (FParam trep), + mapOnLParam :: LParam frep -> m (LParam trep), + mapOnOp :: Op frep -> m (Op trep) } -- | A mapper that simply returns the tree verbatim. -identityMapper :: Monad m => Mapper lore lore m +identityMapper :: Monad m => Mapper rep rep m identityMapper = Mapper { mapOnSubExp = return, @@ -78,9 +78,9 @@ -- into subexpressions. The mapping is done left-to-right. mapExpM :: (Applicative m, Monad m) => - Mapper flore tlore m -> - Exp flore -> - m (Exp tlore) + Mapper frep trep m -> + Exp frep -> + m (Exp trep) mapExpM tv (BasicOp (SubExp se)) = BasicOp <$> (SubExp <$> mapOnSubExp tv se) mapExpM tv (BasicOp (ArrayLit els rowt)) = @@ -174,14 +174,14 @@ mapExpM tv (Op op) = Op <$> mapOnOp tv op -mapOnShape :: Monad m => Mapper flore tlore m -> Shape -> m Shape +mapOnShape :: Monad m => Mapper frep trep m -> Shape -> m Shape mapOnShape tv (Shape ds) = Shape <$> mapM (mapOnSubExp tv) ds mapOnLoopForm :: Monad m => - Mapper flore tlore m -> - LoopForm flore -> - m (LoopForm tlore) + Mapper frep trep m -> + LoopForm frep -> + m (LoopForm trep) mapOnLoopForm tv (ForLoop i it bound loop_vars) = ForLoop <$> mapOnVName tv i <*> pure it <*> mapOnSubExp tv bound <*> (zip <$> mapM (mapOnLParam tv) loop_lparams <*> mapM (mapOnVName tv) loop_arrs) @@ -192,9 +192,9 @@ mapOnLambda :: Monad m => - Mapper flore tlore m -> - Lambda flore -> - m (Lambda tlore) + Mapper frep trep m -> + Lambda frep -> + m (Lambda trep) mapOnLambda tv (Lambda params body ret) = do params' <- mapM (mapOnLParam tv) params Lambda params' @@ -202,25 +202,25 @@ <*> mapM (mapOnType (mapOnSubExp tv)) ret -- | Like 'mapExpM', but in the 'Identity' monad. -mapExp :: Mapper flore tlore Identity -> Exp flore -> Exp tlore +mapExp :: Mapper frep trep Identity -> Exp frep -> Exp trep mapExp m = runIdentity . mapExpM m -- | Express a monad expression on a syntax node. Each element of -- this structure expresses the action to be performed on a given -- child. -data Walker lore m = Walker +data Walker rep m = Walker { walkOnSubExp :: SubExp -> m (), - walkOnBody :: Scope lore -> Body lore -> m (), + walkOnBody :: Scope rep -> Body rep -> m (), walkOnVName :: VName -> m (), - walkOnRetType :: RetType lore -> m (), - walkOnBranchType :: BranchType lore -> m (), - walkOnFParam :: FParam lore -> m (), - walkOnLParam :: LParam lore -> m (), - walkOnOp :: Op lore -> m () + walkOnRetType :: RetType rep -> m (), + walkOnBranchType :: BranchType rep -> m (), + walkOnFParam :: FParam rep -> m (), + walkOnLParam :: LParam rep -> m (), + walkOnOp :: Op rep -> m () } -- | A no-op traversal. -identityWalker :: Monad m => Walker lore m +identityWalker :: Monad m => Walker rep m identityWalker = Walker { walkOnSubExp = const $ return (), @@ -233,10 +233,10 @@ walkOnOp = const $ return () } -walkOnShape :: Monad m => Walker lore m -> Shape -> m () +walkOnShape :: Monad m => Walker rep m -> Shape -> m () walkOnShape tv (Shape ds) = mapM_ (walkOnSubExp tv) ds -walkOnType :: Monad m => Walker lore m -> Type -> m () +walkOnType :: Monad m => Walker rep m -> Type -> m () walkOnType _ Prim {} = return () walkOnType tv (Acc acc ispace ts _) = do walkOnVName tv acc @@ -245,7 +245,7 @@ walkOnType _ Mem {} = return () walkOnType tv (Array _ shape _) = walkOnShape tv shape -walkOnLoopForm :: Monad m => Walker lore m -> LoopForm lore -> m () +walkOnLoopForm :: Monad m => Walker rep m -> LoopForm rep -> m () walkOnLoopForm tv (ForLoop i _ bound loop_vars) = walkOnVName tv i >> walkOnSubExp tv bound >> mapM_ (walkOnLParam tv) loop_lparams @@ -255,14 +255,14 @@ walkOnLoopForm tv (WhileLoop cond) = walkOnVName tv cond -walkOnLambda :: Monad m => Walker lore m -> Lambda lore -> m () +walkOnLambda :: Monad m => Walker rep m -> Lambda rep -> m () walkOnLambda tv (Lambda params body ret) = do mapM_ (walkOnLParam tv) params walkOnBody tv (scopeOfLParams params) body mapM_ (walkOnType tv) ret -- | As 'mapExpM', but do not construct a result AST. -walkExpM :: Monad m => Walker lore m -> Exp lore -> m () +walkExpM :: Monad m => Walker rep m -> Exp rep -> m () walkExpM tv (BasicOp (SubExp se)) = walkOnSubExp tv se walkExpM tv (BasicOp (ArrayLit els rowt)) = 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 @@ -1091,11 +1091,9 @@ typeFromSV (Dynamic tp) = tp typeFromSV (LambdaSV _ _ _ env) = - Scalar $ - Record $ - M.fromList $ - map (bimap (nameFromString . pretty) (typeFromSV . bindingSV)) $ - M.toList env + Scalar . Record . M.fromList $ + map (bimap (nameFromString . pretty) (typeFromSV . bindingSV)) $ + M.toList env typeFromSV (RecordSV ls) = let ts = map (fmap typeFromSV) ls in Scalar $ Record $ M.fromList ts @@ -1174,13 +1172,7 @@ | ps' <- sortOn fst ps, svs' <- sortOn fst svs = RecordPattern - ( zipWith - ( \(n, p) (_, sv) -> - (n, updatePattern p sv) - ) - ps' - svs' - ) + (zipWith (\(n, p) (_, sv) -> (n, updatePattern p sv)) ps' svs') loc updatePattern (PatternParens pat loc) sv = PatternParens (updatePattern pat sv) loc @@ -1251,9 +1243,14 @@ ++ "but the defunctionaliser expects a monomorphic input program." (tparams', params', body', sv) <- defuncLet (map typeParamName tparams) params body rettype - let rettype' = combineTypeShapes rettype $ anySizes $ toStruct $ typeOf body' globals <- asks fst - let bound_sizes = S.fromList tparams' <> globals + let bound_sizes = foldMap patternNames params' <> S.fromList tparams' <> globals + rettype' = + -- FIXME: dubious that we cannot assume that all sizes in the + -- body are in scope. This is because when we insert + -- applications of lifted functions, we don't properly update + -- the types in the return type annotation. + combineTypeShapes rettype $ first (anyDimIfNotBound bound_sizes) $ toStruct $ typeOf body' (missing_dims, params'') <- sizesForAll bound_sizes params' return ( valbind @@ -1272,18 +1269,17 @@ }, M.singleton name $ Binding - ( Just - ( first - (map typeParamName) - (valBindTypeScheme valbind) - ) - ) + (Just (first (map typeParamName) (valBindTypeScheme valbind))) sv, case sv of DynamicFun {} -> True Dynamic {} -> True _ -> False ) + where + anyDimIfNotBound bound_sizes (NamedDim v) + | qualLeaf v `S.notMember` bound_sizes = AnyDim $ Just $ qualLeaf v + anyDimIfNotBound _ d = d -- | Defunctionalize a list of top-level declarations. defuncVals :: [ValBind] -> DefM () diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -152,16 +152,17 @@ entryPointType (t, ts) | E.Scalar (E.Prim E.Unsigned {}) <- E.entryType t = - [I.TypeUnsigned] + [I.TypeUnsigned u] | E.Array _ _ (E.Prim E.Unsigned {}) _ <- E.entryType t = - [I.TypeUnsigned] + [I.TypeUnsigned u] | E.Scalar E.Prim {} <- E.entryType t = - [I.TypeDirect] + [I.TypeDirect u] | E.Array _ _ E.Prim {} _ <- E.entryType t = - [I.TypeDirect] + [I.TypeDirect u] | otherwise = - [I.TypeOpaque desc $ length ts] + [I.TypeOpaque u desc $ length ts] where + u = foldl max Nonunique $ map I.uniqueness ts desc = maybe (prettyOneLine t') typeExpOpaqueName $ E.entryAscribed t t' = noSizes (E.entryType t) `E.setUniqueness` Nonunique typeExpOpaqueName (TEApply te TypeArgExpDim {} _) = diff --git a/src/Futhark/Internalise/Monad.hs b/src/Futhark/Internalise/Monad.hs --- a/src/Futhark/Internalise/Monad.hs +++ b/src/Futhark/Internalise/Monad.hs @@ -84,7 +84,7 @@ putNameSource src = modify $ \s -> s {stateNameSource = src} instance MonadBinder InternaliseM where - type Lore InternaliseM = SOACS + type Rep InternaliseM = SOACS mkExpDecM pat e = InternaliseM $ mkExpDecM pat e mkBodyM bnds res = InternaliseM $ mkBodyM bnds res mkLetNamesM pat e = InternaliseM $ mkLetNamesM pat e diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs --- a/src/Futhark/Optimise/BlkRegTiling.hs +++ b/src/Futhark/Optimise/BlkRegTiling.hs @@ -24,13 +24,13 @@ import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Sequence as Seq -import Futhark.IR.Kernels +import Futhark.IR.GPU import Futhark.MonadFreshNames import Futhark.Optimise.TileLoops.Shared import Futhark.Tools import Futhark.Transform.Rename -mmBlkRegTiling :: Stm Kernels -> TileM (Maybe (Stms Kernels, Stm Kernels)) +mmBlkRegTiling :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread {} seg_space ts old_kbody)))) | KernelBody () kstms [Returns ResultMaySimplify (Var res_nm)] <- old_kbody, -- check kernel has one result of primitive type @@ -416,7 +416,7 @@ return $ Just (host_stms, new_kernel) mmBlkRegTiling _ = return Nothing -ceilDiv :: MonadBinder m => SubExp -> SubExp -> m (Exp (Lore m)) +ceilDiv :: MonadBinder m => SubExp -> SubExp -> m (Exp (Rep m)) ceilDiv x y = pure $ BasicOp $ BinOp (SDivUp Int64 Unsafe) x y scratch :: MonadBinder m => String -> PrimType -> [SubExp] -> m VName @@ -446,9 +446,9 @@ [VName] -> -- loop inits ( VName -> [VName] -> -- (loop var -> loop inits -> loop body) - Binder Kernels (Body Kernels) + Binder GPU (Body GPU) ) -> - Binder Kernels [VName] + Binder GPU [VName] forLoop' i_bound merge body = do i <- newVName "i" -- could give this as arg to the function let loop_form = ForLoop i Int64 i_bound [] @@ -466,8 +466,8 @@ forLoop :: SubExp -> [VName] -> - (VName -> [VName] -> Binder Kernels (Body Kernels)) -> - Binder Kernels VName + (VName -> [VName] -> Binder GPU (Body GPU)) -> + Binder GPU VName forLoop i_bound merge body = do res_list <- forLoop' i_bound merge body return $ head res_list @@ -479,10 +479,10 @@ -- creates Stms corresponding to binding of new_params, -- lambda body, and binding of lambda result to res_name. rebindLambda :: - Lambda Kernels -> + Lambda GPU -> [VName] -> [VName] -> - Stms Kernels + Stms GPU rebindLambda lam new_params res_names = stmsFromList ( zipWith @@ -513,8 +513,8 @@ -- | Tries to identify the following pattern: -- code followed by some Screma followed by more code. matchCodeStreamCode :: - Stms Kernels -> - (Stms Kernels, Maybe (Stm Kernels), Stms Kernels) + Stms GPU -> + (Stms GPU, Maybe (Stm GPU), Stms GPU) matchCodeStreamCode kstms = let (code1, screma, code2) = foldl @@ -568,9 +568,9 @@ processIndirections :: Names -> -- input arrays to redomap Names -> -- variables on which the result of redomap depends on. - Maybe (Stms Kernels, M.Map VName (Stm Kernels)) -> - Stm Kernels -> - Maybe (Stms Kernels, M.Map VName (Stm Kernels)) + Maybe (Stms GPU, M.Map VName (Stm GPU)) -> + Stm GPU -> + Maybe (Stms GPU, M.Map VName (Stm GPU)) processIndirections arrs _ acc stm@(Let patt _ (BasicOp (Index _ _))) | Just (ss, tab) <- acc, [p] <- patternValueElements patt, @@ -599,7 +599,7 @@ se8 :: SubExp se8 = intConst Int64 8 -getParTiles :: (String, String) -> (Name, Name) -> SubExp -> Binder Kernels (SubExp, SubExp) +getParTiles :: (String, String) -> (Name, Name) -> SubExp -> Binder GPU (SubExp, SubExp) getParTiles (t_str, r_str) (t_name, r_name) len_dim = case len_dim of Constant (IntValue (Int64Value 8)) -> @@ -613,7 +613,7 @@ r <- letSubExp r_str $ Op $ SizeOp $ GetSize r_name SizeRegTile return (t, r) -getSeqTile :: String -> Name -> SubExp -> SubExp -> SubExp -> Binder Kernels SubExp +getSeqTile :: String -> Name -> SubExp -> SubExp -> SubExp -> Binder GPU SubExp getSeqTile tk_str tk_name len_dim ty tx = case (tx, ty) of (Constant (IntValue (Int64Value v_x)), Constant (IntValue (Int64Value v_y))) -> @@ -695,9 +695,9 @@ -- a) each of the statements is a slice that produces one of the -- streamed arrays -- --- mmBlkRegTiling :: Stm Kernels -> TileM (Maybe (Stms Kernels, Stm Kernels)) +-- mmBlkRegTiling :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) -- mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread{} seg_space ts old_kbody)))) -doRegTiling3D :: Stm Kernels -> TileM (Maybe (Stms Kernels, Stm Kernels)) +doRegTiling3D :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) doRegTiling3D (Let pat aux (Op (SegOp old_kernel))) | SegMap SegThread {} space kertp (KernelBody () kstms kres) <- old_kernel, -- build the variance table, that records, for @@ -992,14 +992,14 @@ getResNm (Returns ResultMaySimplify (Var res_nm)) = Just res_nm getResNm _ = Nothing - limitTile :: String -> SubExp -> SubExp -> Binder Kernels SubExp + limitTile :: String -> SubExp -> SubExp -> Binder GPU SubExp limitTile t_str t d_K = letSubExp t_str $ BasicOp $ BinOp (SMin Int64) t d_K insertTranspose :: VarianceTable -> (VName, SubExp) -> - (M.Map VName (Stm Kernels), M.Map VName (PrimType, Stm Kernels)) -> - (VName, Stm Kernels) -> - Binder Kernels (M.Map VName (Stm Kernels), M.Map VName (PrimType, Stm Kernels)) + (M.Map VName (Stm GPU), M.Map VName (PrimType, Stm GPU)) -> + (VName, Stm GPU) -> + Binder GPU (M.Map VName (Stm GPU), M.Map VName (PrimType, Stm GPU)) insertTranspose variance (gidz, _) (tab_inn, tab_out) (p_nm, stm@(Let patt yy (BasicOp (Index arr_nm slc)))) | [p] <- patternValueElements patt, ptp <- elemType $ patElemType p, 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 @@ -46,7 +46,7 @@ removeProgAliases, removeStmAliases, ) -import qualified Futhark.IR.Kernels.Kernel as Kernel +import qualified Futhark.IR.GPU.Kernel as Kernel import qualified Futhark.IR.MC as MC import qualified Futhark.IR.Mem as Memory import Futhark.IR.Prop.Aliases @@ -54,22 +54,22 @@ import Futhark.Pass import Futhark.Transform.Substitute -consumedInStms :: Aliased lore => Stms lore -> Names +consumedInStms :: Aliased rep => Stms rep -> Names consumedInStms = snd . flip mkStmsAliases [] -- | Perform CSE on every function in a program. -- -- If the boolean argument is false, the pass will not perform CSE on --- expressions producing arrays. This should be disabled when the lore has +-- expressions producing arrays. This should be disabled when the rep has -- memory information, since at that point arrays have identity beyond their -- value. performCSE :: - ( ASTLore lore, - CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore)) + ( ASTRep rep, + CanBeAliased (Op rep), + CSEInOp (OpWithAliases (Op rep)) ) => Bool -> - Pass lore lore + Pass rep rep performCSE cse_arrays = Pass "CSE" "Combine common subexpressions." $ fmap removeProgAliases @@ -87,34 +87,34 @@ -- | Perform CSE on a single function. -- -- If the boolean argument is false, the pass will not perform CSE on --- expressions producing arrays. This should be disabled when the lore has +-- expressions producing arrays. This should be disabled when the rep has -- memory information, since at that point arrays have identity beyond their -- value. performCSEOnFunDef :: - ( ASTLore lore, - CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore)) + ( ASTRep rep, + CanBeAliased (Op rep), + CSEInOp (OpWithAliases (Op rep)) ) => Bool -> - FunDef lore -> - FunDef lore + FunDef rep -> + FunDef rep performCSEOnFunDef cse_arrays = removeFunDefAliases . cseInFunDef cse_arrays . analyseFun -- | Perform CSE on some statements. -- -- If the boolean argument is false, the pass will not perform CSE on --- expressions producing arrays. This should be disabled when the lore has +-- expressions producing arrays. This should be disabled when the rep has -- memory information, since at that point arrays have identity beyond their -- value. performCSEOnStms :: - ( ASTLore lore, - CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore)) + ( ASTRep rep, + CanBeAliased (Op rep), + CSEInOp (OpWithAliases (Op rep)) ) => Bool -> - Stms lore -> - Stms lore + Stms rep -> + Stms rep performCSEOnStms cse_arrays = fmap removeStmAliases . f . fst . analyseStms mempty where @@ -129,10 +129,10 @@ (newCSEState cse_arrays) cseInFunDef :: - (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + (ASTRep rep, Aliased rep, CSEInOp (Op rep)) => Bool -> - FunDef lore -> - FunDef lore + FunDef rep -> + FunDef rep cseInFunDef cse_arrays fundec = fundec { funDefBody = @@ -150,13 +150,13 @@ | primType $ declExtTypeOf t = Observe | otherwise = Consume -type CSEM lore = Reader (CSEState lore) +type CSEM rep = Reader (CSEState rep) cseInBody :: - (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + (ASTRep rep, Aliased rep, CSEInOp (Op rep)) => [Diet] -> - Body lore -> - CSEM lore (Body lore) + Body rep -> + CSEM rep (Body rep) cseInBody ds (Body bodydec stms res) = do (stms', res') <- cseInStms (res_cons <> stms_cons) (stmsToList stms) $ do @@ -170,19 +170,19 @@ consumeResult _ _ = mempty cseInLambda :: - (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - Lambda lore -> - CSEM lore (Lambda lore) + (ASTRep rep, Aliased rep, CSEInOp (Op rep)) => + Lambda rep -> + CSEM rep (Lambda rep) cseInLambda lam = do body' <- cseInBody (map (const Observe) $ lambdaReturnType lam) $ lambdaBody lam return lam {lambdaBody = body'} cseInStms :: - (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => + (ASTRep rep, Aliased rep, CSEInOp (Op rep)) => Names -> - [Stm lore] -> - CSEM lore a -> - CSEM lore (Stms lore, a) + [Stm rep] -> + CSEM rep a -> + CSEM rep (Stms rep, a) cseInStms _ [] m = do a <- m return (mempty, a) @@ -208,11 +208,11 @@ | otherwise = Observe cseInStm :: - ASTLore lore => + ASTRep rep => Names -> - Stm lore -> - ([Stm lore] -> CSEM lore a) -> - CSEM lore a + Stm rep -> + ([Stm rep] -> CSEM rep a) -> + CSEM rep a cseInStm consumed (Let pat (StmAux cs attrs edec) e) m = do CSEState (esubsts, nsubsts) cse_arrays <- ask let e' = substituteNames nsubsts e @@ -239,70 +239,70 @@ | patElemName pe `nameIn` consumed = True | otherwise = False -type ExpressionSubstitutions lore = +type ExpressionSubstitutions rep = M.Map - (ExpDec lore, Exp lore) - (Pattern lore) + (ExpDec rep, Exp rep) + (Pattern rep) type NameSubstitutions = M.Map VName VName -data CSEState lore = CSEState - { _cseSubstitutions :: (ExpressionSubstitutions lore, NameSubstitutions), +data CSEState rep = CSEState + { _cseSubstitutions :: (ExpressionSubstitutions rep, NameSubstitutions), _cseArrays :: Bool } -newCSEState :: Bool -> CSEState lore +newCSEState :: Bool -> CSEState rep newCSEState = CSEState (M.empty, M.empty) mkSubsts :: PatternT dec -> PatternT dec -> M.Map VName VName mkSubsts pat vs = M.fromList $ zip (patternNames pat) (patternNames vs) -addNameSubst :: PatternT dec -> PatternT dec -> CSEState lore -> CSEState lore +addNameSubst :: PatternT dec -> PatternT dec -> CSEState rep -> CSEState rep addNameSubst pat subpat (CSEState (esubsts, nsubsts) cse_arrays) = CSEState (esubsts, mkSubsts pat subpat `M.union` nsubsts) cse_arrays addExpSubst :: - ASTLore lore => - Pattern lore -> - ExpDec lore -> - Exp lore -> - CSEState lore -> - CSEState lore + ASTRep rep => + Pattern rep -> + ExpDec rep -> + 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 -- | The operations that permit CSE. class CSEInOp op where -- | Perform CSE within any nested expressions. - cseInOp :: op -> CSEM lore op + cseInOp :: op -> CSEM rep op instance CSEInOp () where cseInOp () = return () -subCSE :: CSEM lore r -> CSEM otherlore r +subCSE :: CSEM rep r -> CSEM otherrep r subCSE m = do CSEState _ cse_arrays <- ask return $ runReader m $ newCSEState cse_arrays instance - ( ASTLore lore, - Aliased lore, - CSEInOp (Op lore), + ( ASTRep rep, + Aliased rep, + CSEInOp (Op rep), CSEInOp op ) => - CSEInOp (Kernel.HostOp lore op) + CSEInOp (Kernel.HostOp rep op) where cseInOp (Kernel.SegOp op) = Kernel.SegOp <$> cseInOp op cseInOp (Kernel.OtherOp op) = Kernel.OtherOp <$> cseInOp op cseInOp x = return x instance - ( ASTLore lore, - Aliased lore, - CSEInOp (Op lore), + ( ASTRep rep, + Aliased rep, + CSEInOp (Op rep), CSEInOp op ) => - CSEInOp (MC.MCOp lore op) + CSEInOp (MC.MCOp rep op) where cseInOp (MC.ParOp par_op op) = MC.ParOp <$> traverse cseInOp par_op <*> cseInOp op @@ -310,8 +310,8 @@ MC.OtherOp <$> cseInOp op instance - (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - CSEInOp (Kernel.SegOp lvl lore) + (ASTRep rep, Aliased rep, CSEInOp (Op rep)) => + CSEInOp (Kernel.SegOp lvl rep) where cseInOp = subCSE @@ -319,9 +319,9 @@ (Kernel.SegOpMapper return cseInLambda cseInKernelBody return return) cseInKernelBody :: - (ASTLore lore, Aliased lore, CSEInOp (Op lore)) => - Kernel.KernelBody lore -> - CSEM lore (Kernel.KernelBody lore) + (ASTRep rep, Aliased rep, CSEInOp (Op rep)) => + Kernel.KernelBody rep -> + CSEM rep (Kernel.KernelBody rep) cseInKernelBody (Kernel.KernelBody bodydec bnds res) = do Body _ bnds' _ <- cseInBody (map (const Observe) res) $ Body bodydec bnds [] return $ Kernel.KernelBody bodydec bnds' res @@ -331,10 +331,10 @@ cseInOp (Memory.Inner k) = Memory.Inner <$> subCSE (cseInOp k) instance - ( ASTLore lore, - CanBeAliased (Op lore), - CSEInOp (OpWithAliases (Op lore)) + ( ASTRep rep, + CanBeAliased (Op rep), + CSEInOp (OpWithAliases (Op rep)) ) => - CSEInOp (SOAC.SOAC (Aliases lore)) + CSEInOp (SOAC.SOAC (Aliases rep)) where cseInOp = subCSE . SOAC.mapSOACM (SOAC.SOACMapper return cseInLambda return) diff --git a/src/Futhark/Optimise/DoubleBuffer.hs b/src/Futhark/Optimise/DoubleBuffer.hs --- a/src/Futhark/Optimise/DoubleBuffer.hs +++ b/src/Futhark/Optimise/DoubleBuffer.hs @@ -22,7 +22,7 @@ -- value. This has the effect of making the memory block returned by -- the array non-existential, which is important for later memory -- expansion to work. -module Futhark.Optimise.DoubleBuffer (doubleBufferKernels, doubleBufferMC) where +module Futhark.Optimise.DoubleBuffer (doubleBufferGPU, doubleBufferMC) where import Control.Monad.Reader import Control.Monad.State @@ -31,24 +31,24 @@ import qualified Data.Map.Strict as M import Data.Maybe import Futhark.Construct -import Futhark.IR.KernelsMem as Kernels +import Futhark.IR.GPUMem as GPU import Futhark.IR.MCMem as MC import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.Pass import Futhark.Pass.ExplicitAllocations (arraySizeInBytesExp) -import Futhark.Pass.ExplicitAllocations.Kernels () +import Futhark.Pass.ExplicitAllocations.GPU () import Futhark.Util (maybeHead) -- | The pass for GPU kernels. -doubleBufferKernels :: Pass KernelsMem KernelsMem -doubleBufferKernels = doubleBuffer optimiseKernelsOp +doubleBufferGPU :: Pass GPUMem GPUMem +doubleBufferGPU = doubleBuffer optimiseGPUOp -- | The pass for multicore doubleBufferMC :: Pass MCMem MCMem doubleBufferMC = doubleBuffer optimiseMCOp -- | The double buffering pass definition. -doubleBuffer :: Mem lore => OptimiseOp lore -> Pass lore lore +doubleBuffer :: Mem rep => OptimiseOp rep -> Pass rep rep doubleBuffer onOp = Pass { passName = "Double buffer", @@ -66,51 +66,51 @@ env = Env mempty doNotTouchLoop onOp doNotTouchLoop ctx val body = return (mempty, ctx, val, body) -type OptimiseLoop lore = - [(FParam lore, SubExp)] -> - [(FParam lore, SubExp)] -> - Body lore -> +type OptimiseLoop rep = + [(FParam rep, SubExp)] -> + [(FParam rep, SubExp)] -> + Body rep -> DoubleBufferM - lore - ( [Stm lore], - [(FParam lore, SubExp)], - [(FParam lore, SubExp)], - Body lore + rep + ( [Stm rep], + [(FParam rep, SubExp)], + [(FParam rep, SubExp)], + Body rep ) -type OptimiseOp lore = - Op lore -> DoubleBufferM lore (Op lore) +type OptimiseOp rep = + Op rep -> DoubleBufferM rep (Op rep) -data Env lore = Env - { envScope :: Scope lore, - envOptimiseLoop :: OptimiseLoop lore, - envOptimiseOp :: OptimiseOp lore +data Env rep = Env + { envScope :: Scope rep, + envOptimiseLoop :: OptimiseLoop rep, + envOptimiseOp :: OptimiseOp rep } -newtype DoubleBufferM lore a = DoubleBufferM - { runDoubleBufferM :: ReaderT (Env lore) (State VNameSource) a +newtype DoubleBufferM rep a = DoubleBufferM + { runDoubleBufferM :: ReaderT (Env rep) (State VNameSource) a } - deriving (Functor, Applicative, Monad, MonadReader (Env lore), MonadFreshNames) + deriving (Functor, Applicative, Monad, MonadReader (Env rep), MonadFreshNames) -instance ASTLore lore => HasScope lore (DoubleBufferM lore) where +instance ASTRep rep => HasScope rep (DoubleBufferM rep) where askScope = asks envScope -instance ASTLore lore => LocalScope lore (DoubleBufferM lore) where +instance ASTRep rep => LocalScope rep (DoubleBufferM rep) where localScope scope = local $ \env -> env {envScope = envScope env <> scope} -optimiseBody :: ASTLore lore => Body lore -> DoubleBufferM lore (Body lore) +optimiseBody :: ASTRep rep => Body rep -> DoubleBufferM rep (Body rep) optimiseBody body = do bnds' <- optimiseStms $ stmsToList $ bodyStms body return $ body {bodyStms = stmsFromList bnds'} -optimiseStms :: ASTLore lore => [Stm lore] -> DoubleBufferM lore [Stm lore] +optimiseStms :: ASTRep rep => [Stm rep] -> DoubleBufferM rep [Stm rep] optimiseStms [] = return [] optimiseStms (e : es) = do e_es <- optimiseStm e es' <- localScope (castScope $ scopeOf e_es) $ optimiseStms es return $ e_es ++ es' -optimiseStm :: forall lore. ASTLore lore => Stm lore -> DoubleBufferM lore [Stm lore] +optimiseStm :: forall rep. ASTRep rep => Stm rep -> DoubleBufferM rep [Stm rep] optimiseStm (Let pat aux (DoLoop ctx val form body)) = do body' <- localScope (scopeOf form <> scopeOfFParams (map fst $ ctx ++ val)) $ @@ -125,12 +125,12 @@ optimise onOp = identityMapper { mapOnBody = \_ x -> - optimiseBody x :: DoubleBufferM lore (Body lore), + optimiseBody x :: DoubleBufferM rep (Body rep), mapOnOp = onOp } -optimiseKernelsOp :: OptimiseOp KernelsMem -optimiseKernelsOp (Inner (SegOp op)) = +optimiseGPUOp :: OptimiseOp GPUMem +optimiseGPUOp (Inner (SegOp op)) = local inSegOp $ Inner . SegOp <$> mapSegOpM mapper op where mapper = @@ -139,7 +139,7 @@ mapOnSegOpBody = optimiseKernelBody } inSegOp env = env {envOptimiseLoop = optimiseLoop} -optimiseKernelsOp op = return op +optimiseGPUOp op = return op optimiseMCOp :: OptimiseOp MCMem optimiseMCOp (Inner (ParOp par_op op)) = @@ -156,34 +156,34 @@ optimiseMCOp op = return op optimiseKernelBody :: - ASTLore lore => - KernelBody lore -> - DoubleBufferM lore (KernelBody lore) + ASTRep rep => + KernelBody rep -> + DoubleBufferM rep (KernelBody rep) optimiseKernelBody kbody = do stms' <- optimiseStms $ stmsToList $ kernelBodyStms kbody return $ kbody {kernelBodyStms = stmsFromList stms'} optimiseLambda :: - ASTLore lore => - Lambda lore -> - DoubleBufferM lore (Lambda lore) + ASTRep rep => + Lambda rep -> + DoubleBufferM rep (Lambda rep) optimiseLambda lam = do body <- localScope (castScope $ scopeOf lam) $ optimiseBody $ lambdaBody lam return lam {lambdaBody = body} -type Constraints lore = - ( ASTLore lore, - FParamInfo lore ~ FParamMem, - LParamInfo lore ~ LParamMem, - RetType lore ~ RetTypeMem, - LetDec lore ~ LetDecMem, - BranchType lore ~ BranchTypeMem, - ExpDec lore ~ (), - BodyDec lore ~ (), - OpReturns lore +type Constraints rep = + ( ASTRep rep, + FParamInfo rep ~ FParamMem, + LParamInfo rep ~ LParamMem, + RetType rep ~ RetTypeMem, + LetDec rep ~ LetDecMem, + BranchType rep ~ BranchTypeMem, + ExpDec rep ~ (), + BodyDec rep ~ (), + OpReturns rep ) -optimiseLoop :: (Constraints lore, Op lore ~ MemOp inner, BinderOps lore) => OptimiseLoop lore +optimiseLoop :: (Constraints rep, Op rep ~ MemOp inner, BinderOps rep) => OptimiseLoop rep optimiseLoop ctx val body = do -- We start out by figuring out which of the merge variables should -- be double-buffered. @@ -273,10 +273,10 @@ _ -> return NoBuffer allocStms :: - (Constraints lore, Op lore ~ MemOp inner, BinderOps lore) => - [(FParam lore, SubExp)] -> + (Constraints rep, Op rep ~ MemOp inner, BinderOps rep) => + [(FParam rep, SubExp)] -> [DoubleBuffer] -> - DoubleBufferM lore ([(FParam lore, SubExp)], [Stm lore]) + DoubleBufferM rep ([(FParam rep, SubExp)], [Stm rep]) allocStms merge = runWriterT . zipWithM allocation merge where allocation m@(Param pname _, _) (BufferAlloc name size space b) = do @@ -310,11 +310,11 @@ return (f, se) doubleBufferResult :: - (Constraints lore) => - [FParam lore] -> + (Constraints rep) => + [FParam rep] -> [DoubleBuffer] -> - Body lore -> - Body lore + Body rep -> + Body rep doubleBufferResult valparams buffered (Body _ bnds res) = let (ctx_res, val_res) = splitAt (length res - length valparams) res (copybnds, val_res') = diff --git a/src/Futhark/Optimise/Fusion/Composing.hs b/src/Futhark/Optimise/Fusion/Composing.hs --- a/src/Futhark/Optimise/Fusion/Composing.hs +++ b/src/Futhark/Optimise/Fusion/Composing.hs @@ -44,11 +44,11 @@ -- The result is the fused function, and a list of the array inputs -- expected by the SOAC containing the fused function. fuseMaps :: - Bindable lore => + Bindable rep => -- | The producer var names that still need to be returned Names -> -- | Function of SOAC to be fused. - Lambda lore -> + Lambda rep -> -- | Input of SOAC to be fused. [SOAC.Input] -> -- | Output of SOAC to be fused. The @@ -58,12 +58,12 @@ -- bind a single element of that output. [(VName, Ident)] -> -- | Function to be fused with. - Lambda lore -> + Lambda rep -> -- | Input of SOAC to be fused with. [SOAC.Input] -> -- | The fused lambda and the inputs of -- the resulting SOAC. - (Lambda lore, [SOAC.Input]) + (Lambda rep, [SOAC.Input]) fuseMaps unfus_nms lam1 inp1 out1 lam2 inp2 = (lam2', M.elems inputmap) where lam2' = @@ -96,19 +96,19 @@ --(unfus_accpat, unfus_arrpat) = splitAt (length unfus_accs) unfus_pat fuseInputs :: - Bindable lore => + Bindable rep => Names -> - Lambda lore -> + Lambda rep -> [SOAC.Input] -> [(VName, Ident)] -> - Lambda lore -> + Lambda rep -> [SOAC.Input] -> ( [Ident], [Ident], [Ident], M.Map Ident SOAC.Input, - Body lore -> Body lore, - Body lore -> Body lore + Body rep -> Body rep, + Body rep -> Body rep ) fuseInputs unfus_nms lam1 inp1 out1 lam2 inp2 = (lam2redparams, unfus_vars, outbnds, inputmap, makeCopies, makeCopiesInner) @@ -172,9 +172,9 @@ _ -> (m, ra) removeDuplicateInputs :: - Bindable lore => + Bindable rep => M.Map Ident SOAC.Input -> - (M.Map Ident SOAC.Input, Body lore -> Body lore) + (M.Map Ident SOAC.Input, Body rep -> Body rep) removeDuplicateInputs = fst . M.foldlWithKey' comb ((M.empty, id), M.empty) where comb ((parmap, inner), arrmap) par arr = @@ -192,19 +192,19 @@ `insertStm` b fuseRedomap :: - Bindable lore => + Bindable rep => Names -> [VName] -> - Lambda lore -> + Lambda rep -> [SubExp] -> [SubExp] -> [SOAC.Input] -> [(VName, Ident)] -> - Lambda lore -> + Lambda rep -> [SubExp] -> [SubExp] -> [SOAC.Input] -> - (Lambda lore, [SOAC.Input]) + (Lambda rep, [SOAC.Input]) fuseRedomap unfus_nms outVars @@ -282,7 +282,7 @@ } in (res_lam', new_inp) -mergeReduceOps :: Lambda lore -> Lambda lore -> Lambda lore +mergeReduceOps :: Lambda rep -> Lambda rep -> Lambda rep mergeReduceOps (Lambda par1 bdy1 rtp1) (Lambda par2 bdy2 rtp2) = let body' = Body diff --git a/src/Futhark/Optimise/Fusion/LoopKernel.hs b/src/Futhark/Optimise/Fusion/LoopKernel.hs --- a/src/Futhark/Optimise/Fusion/LoopKernel.hs +++ b/src/Futhark/Optimise/Fusion/LoopKernel.hs @@ -429,7 +429,7 @@ (body_p, body_c) = (lambdaBody lam_p, lambdaBody lam_c) body' = Body - { bodyDec = bodyDec body_p, -- body_p and body_c have the same lores + { bodyDec = bodyDec body_p, -- body_p and body_c have the same decorations bodyStms = bodyStms body_p <> bodyStms body_c, bodyResult = take c_num_buckets (bodyResult body_c) @@ -461,7 +461,7 @@ let (body_p, body_c) = (lambdaBody lam_p, lambdaBody lam_c) let body' = Body - { bodyDec = bodyDec body_p, -- body_p and body_c have the same lores + { bodyDec = bodyDec body_p, -- body_p and body_c have the same decorations bodyStms = bodyStms body_p <> bodyStms body_c, bodyResult = zipW as_c (bodyResult body_c) as_p (bodyResult body_p) } @@ -547,7 +547,7 @@ --------------------------------- _ -> fail "Cannot fuse" -getStreamOrder :: StreamForm lore -> StreamOrd +getStreamOrder :: StreamForm rep -> StreamOrd getStreamOrder (Parallel o _ _) = o getStreamOrder Sequential = InOrder diff --git a/src/Futhark/Optimise/InPlaceLowering.hs b/src/Futhark/Optimise/InPlaceLowering.hs --- a/src/Futhark/Optimise/InPlaceLowering.hs +++ b/src/Futhark/Optimise/InPlaceLowering.hs @@ -62,7 +62,7 @@ -- FIXME: the implementation is not finished yet. Specifically, not -- all of the above conditions are checked. module Futhark.Optimise.InPlaceLowering - ( inPlaceLoweringKernels, + ( inPlaceLoweringGPU, inPlaceLoweringSeq, inPlaceLoweringMC, ) @@ -73,15 +73,15 @@ import Futhark.Analysis.Alias import Futhark.Binder import Futhark.IR.Aliases -import Futhark.IR.Kernels +import Futhark.IR.GPU import Futhark.IR.MC import Futhark.IR.Seq (Seq) import Futhark.Optimise.InPlaceLowering.LowerIntoStm import Futhark.Pass -- | Apply the in-place lowering optimisation to the given program. -inPlaceLoweringKernels :: Pass Kernels Kernels -inPlaceLoweringKernels = inPlaceLowering onKernelOp lowerUpdateKernels +inPlaceLoweringGPU :: Pass GPU GPU +inPlaceLoweringGPU = inPlaceLowering onKernelOp lowerUpdateGPU -- | Apply the in-place lowering optimisation to the given program. inPlaceLoweringSeq :: Pass Seq Seq @@ -93,10 +93,10 @@ -- | Apply the in-place lowering optimisation to the given program. inPlaceLowering :: - Constraints lore => - OnOp lore -> - LowerUpdate lore (ForwardingM lore) -> - Pass lore lore + Constraints rep => + OnOp rep -> + LowerUpdate rep (ForwardingM rep) -> + Pass rep rep inPlaceLowering onOp lower = Pass "In-place lowering" "Lower in-place updates into loops" $ fmap removeProgAliases @@ -119,12 +119,12 @@ descend [] m = m descend (stm : stms) m = bindingStm stm $ descend stms m -type Constraints lore = (Bindable lore, CanBeAliased (Op lore)) +type Constraints rep = (Bindable rep, CanBeAliased (Op rep)) optimiseBody :: - Constraints lore => - Body (Aliases lore) -> - ForwardingM lore (Body (Aliases lore)) + Constraints rep => + Body (Aliases rep) -> + ForwardingM rep (Body (Aliases rep)) optimiseBody (Body als bnds res) = do bnds' <- deepen $ @@ -136,10 +136,10 @@ seen (Var v) = seenVar v optimiseStms :: - Constraints lore => - [Stm (Aliases lore)] -> - ForwardingM lore () -> - ForwardingM lore [Stm (Aliases lore)] + Constraints rep => + [Stm (Aliases rep)] -> + ForwardingM rep () -> + ForwardingM rep [Stm (Aliases rep)] optimiseStms [] m = m >> return [] optimiseStms (bnd : bnds) m = do (bnds', bup) <- tapBottomUp $ bindingStm bnd $ optimiseStms bnds m @@ -179,11 +179,11 @@ maybeForward ve v dec cs src slice checkIfForwardableUpdate _ = return () -optimiseInStm :: Constraints lore => Stm (Aliases lore) -> ForwardingM lore (Stm (Aliases lore)) +optimiseInStm :: Constraints rep => Stm (Aliases rep) -> ForwardingM rep (Stm (Aliases rep)) optimiseInStm (Let pat dec e) = Let pat dec <$> optimiseExp e -optimiseExp :: Constraints lore => Exp (Aliases lore) -> ForwardingM lore (Exp (Aliases lore)) +optimiseExp :: Constraints rep => Exp (Aliases rep) -> ForwardingM rep (Exp (Aliases rep)) optimiseExp (DoLoop ctx val form body) = bindingScope (scopeOf form) $ bindingFParams (map fst $ ctx ++ val) $ @@ -199,9 +199,9 @@ } onSegOp :: - (Bindable lore, CanBeAliased (Op lore)) => - SegOp lvl (Aliases lore) -> - ForwardingM lore (SegOp lvl (Aliases lore)) + (Bindable rep, CanBeAliased (Op rep)) => + SegOp lvl (Aliases rep) -> + ForwardingM rep (SegOp lvl (Aliases rep)) onSegOp op = bindingScope (scopeOfSegSpace (segSpace op)) $ do let mapper = identitySegOpMapper {mapOnSegOpBody = onKernelBody} @@ -217,63 +217,63 @@ onMCOp (ParOp par_op op) = ParOp <$> traverse onSegOp par_op <*> onSegOp op onMCOp op = return op -onKernelOp :: OnOp Kernels +onKernelOp :: OnOp GPU onKernelOp (SegOp op) = SegOp <$> onSegOp op onKernelOp op = return op -data Entry lore = Entry +data Entry rep = Entry { entryNumber :: Int, entryAliases :: Names, entryDepth :: Int, entryOptimisable :: Bool, - entryType :: NameInfo (Aliases lore) + entryType :: NameInfo (Aliases rep) } -type VTable lore = M.Map VName (Entry lore) +type VTable rep = M.Map VName (Entry rep) -type OnOp lore = Op (Aliases lore) -> ForwardingM lore (Op (Aliases lore)) +type OnOp rep = Op (Aliases rep) -> ForwardingM rep (Op (Aliases rep)) -data TopDown lore = TopDown +data TopDown rep = TopDown { topDownCounter :: Int, - topDownTable :: VTable lore, + topDownTable :: VTable rep, topDownDepth :: Int, - topLowerUpdate :: LowerUpdate lore (ForwardingM lore), - topOnOp :: OnOp lore + topLowerUpdate :: LowerUpdate rep (ForwardingM rep), + topOnOp :: OnOp rep } -data BottomUp lore = BottomUp +data BottomUp rep = BottomUp { bottomUpSeen :: Names, - forwardThese :: [DesiredUpdate (LetDec (Aliases lore))] + forwardThese :: [DesiredUpdate (LetDec (Aliases rep))] } -instance Semigroup (BottomUp lore) where +instance Semigroup (BottomUp rep) where BottomUp seen1 forward1 <> BottomUp seen2 forward2 = BottomUp (seen1 <> seen2) (forward1 <> forward2) -instance Monoid (BottomUp lore) where +instance Monoid (BottomUp rep) where mempty = BottomUp mempty mempty -newtype ForwardingM lore a = ForwardingM (RWS (TopDown lore) (BottomUp lore) VNameSource a) +newtype ForwardingM rep a = ForwardingM (RWS (TopDown rep) (BottomUp rep) VNameSource a) deriving ( Monad, Applicative, Functor, - MonadReader (TopDown lore), - MonadWriter (BottomUp lore), + MonadReader (TopDown rep), + MonadWriter (BottomUp rep), MonadState VNameSource ) -instance MonadFreshNames (ForwardingM lore) where +instance MonadFreshNames (ForwardingM rep) where getNameSource = get putNameSource = put -instance Constraints lore => HasScope (Aliases lore) (ForwardingM lore) where +instance Constraints rep => HasScope (Aliases rep) (ForwardingM rep) where askScope = M.map entryType <$> asks topDownTable runForwardingM :: - LowerUpdate lore (ForwardingM lore) -> - OnOp lore -> - ForwardingM lore a -> + LowerUpdate rep (ForwardingM rep) -> + OnOp rep -> + ForwardingM rep a -> VNameSource -> (a, VNameSource) runForwardingM f g (ForwardingM m) src = @@ -290,10 +290,10 @@ } bindingParams :: - (dec -> NameInfo (Aliases lore)) -> + (dec -> NameInfo (Aliases rep)) -> [Param dec] -> - ForwardingM lore a -> - ForwardingM lore a + ForwardingM rep a -> + ForwardingM rep a bindingParams f params = local $ \(TopDown n vtable d x y) -> let entry fparam = ( paramName fparam, @@ -303,15 +303,15 @@ in TopDown (n + 1) (M.union entries vtable) d x y bindingFParams :: - [FParam (Aliases lore)] -> - ForwardingM lore a -> - ForwardingM lore a + [FParam (Aliases rep)] -> + ForwardingM rep a -> + ForwardingM rep a bindingFParams = bindingParams FParamName bindingScope :: - Scope (Aliases lore) -> - ForwardingM lore a -> - ForwardingM lore a + Scope (Aliases rep) -> + ForwardingM rep a -> + ForwardingM rep a bindingScope scope = local $ \(TopDown n vtable d x y) -> let entries = M.map entry scope infoAliases (LetName (aliases, _)) = unAliases aliases @@ -320,9 +320,9 @@ in TopDown (n + 1) (entries <> vtable) d x y bindingStm :: - Stm (Aliases lore) -> - ForwardingM lore a -> - ForwardingM lore a + Stm (Aliases rep) -> + ForwardingM rep a -> + ForwardingM rep a bindingStm (Let pat _ _) = local $ \(TopDown n vtable d x y) -> let entries = M.fromList $ map entry $ patternElements pat entry patElem = @@ -332,7 +332,7 @@ ) in TopDown (n + 1) (M.union entries vtable) d x y -bindingNumber :: VName -> ForwardingM lore Int +bindingNumber :: VName -> ForwardingM rep Int bindingNumber name = do res <- asks $ fmap entryNumber . M.lookup name . topDownTable case res of @@ -343,16 +343,16 @@ ++ pretty name ++ " not found." -deepen :: ForwardingM lore a -> ForwardingM lore a +deepen :: ForwardingM rep a -> ForwardingM rep a deepen = local $ \env -> env {topDownDepth = topDownDepth env + 1} -areAvailableBefore :: Names -> VName -> ForwardingM lore Bool +areAvailableBefore :: Names -> VName -> ForwardingM rep Bool areAvailableBefore names point = do pointN <- bindingNumber point nameNs <- mapM bindingNumber $ namesToList names return $ all (< pointN) nameNs -isInCurrentBody :: VName -> ForwardingM lore Bool +isInCurrentBody :: VName -> ForwardingM rep Bool isInCurrentBody name = do current <- asks topDownDepth res <- asks $ fmap entryDepth . M.lookup name . topDownTable @@ -364,7 +364,7 @@ ++ pretty name ++ " not found." -isOptimisable :: VName -> ForwardingM lore Bool +isOptimisable :: VName -> ForwardingM rep Bool isOptimisable name = do res <- asks $ fmap entryOptimisable . M.lookup name . topDownTable case res of @@ -375,7 +375,7 @@ ++ pretty name ++ " not found." -seenVar :: VName -> ForwardingM lore () +seenVar :: VName -> ForwardingM rep () seenVar name = do aliases <- asks $ @@ -384,20 +384,20 @@ . topDownTable tell $ mempty {bottomUpSeen = oneName name <> aliases} -tapBottomUp :: ForwardingM lore a -> ForwardingM lore (a, BottomUp lore) +tapBottomUp :: ForwardingM rep a -> ForwardingM rep (a, BottomUp rep) tapBottomUp m = do (x, bup) <- listen m return (x, bup) maybeForward :: - Constraints lore => + Constraints rep => VName -> VName -> - LetDec (Aliases lore) -> + LetDec (Aliases rep) -> Certificates -> VName -> Slice SubExp -> - ForwardingM lore () + ForwardingM rep () maybeForward v dest_nm dest_dec cs src slice = do -- Checks condition (2) available <- diff --git a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs --- a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs +++ b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs @@ -2,7 +2,7 @@ {-# LANGUAGE TypeFamilies #-} module Futhark.Optimise.InPlaceLowering.LowerIntoStm - ( lowerUpdateKernels, + ( lowerUpdateGPU, lowerUpdate, LowerUpdate, DesiredUpdate (..), @@ -17,7 +17,7 @@ import Futhark.Analysis.PrimExp.Convert import Futhark.Construct import Futhark.IR.Aliases -import Futhark.IR.Kernels +import Futhark.IR.GPU import Futhark.Optimise.InPlaceLowering.SubstituteIndices data DesiredUpdate dec = DesiredUpdate @@ -38,19 +38,19 @@ updateHasValue :: VName -> DesiredUpdate dec -> Bool updateHasValue name = (name ==) . updateValue -type LowerUpdate lore m = - Scope (Aliases lore) -> - Stm (Aliases lore) -> - [DesiredUpdate (LetDec (Aliases lore))] -> - Maybe (m [Stm (Aliases lore)]) +type LowerUpdate rep m = + Scope (Aliases rep) -> + Stm (Aliases rep) -> + [DesiredUpdate (LetDec (Aliases rep))] -> + Maybe (m [Stm (Aliases rep)]) lowerUpdate :: ( MonadFreshNames m, - Bindable lore, - LetDec lore ~ Type, - CanBeAliased (Op lore) + Bindable rep, + LetDec rep ~ Type, + CanBeAliased (Op rep) ) => - LowerUpdate lore m + LowerUpdate rep m lowerUpdate scope (Let pat aux (DoLoop ctx val form body)) updates = do canDo <- lowerUpdateIntoLoop scope updates pat ctx val form body Just $ do @@ -76,8 +76,8 @@ lowerUpdate _ _ _ = Nothing -lowerUpdateKernels :: MonadFreshNames m => LowerUpdate Kernels m -lowerUpdateKernels +lowerUpdateGPU :: MonadFreshNames m => LowerUpdate GPU m +lowerUpdateGPU scope (Let pat aux (Op (SegOp (SegMap lvl space ts kbody)))) updates @@ -99,20 +99,20 @@ source_used_in_kbody = mconcat (map (`lookupAliases` scope) (namesToList (freeIn kbody))) `namesIntersect` mconcat (map ((`lookupAliases` scope) . updateSource) updates) -lowerUpdateKernels scope stm updates = lowerUpdate scope stm updates +lowerUpdateGPU scope stm updates = lowerUpdate scope stm updates lowerUpdatesIntoSegMap :: MonadFreshNames m => - Scope (Aliases Kernels) -> - Pattern (Aliases Kernels) -> - [DesiredUpdate (LetDec (Aliases Kernels))] -> + Scope (Aliases GPU) -> + Pattern (Aliases GPU) -> + [DesiredUpdate (LetDec (Aliases GPU))] -> SegSpace -> - KernelBody (Aliases Kernels) -> + KernelBody (Aliases GPU) -> Maybe ( m - ( Pattern (Aliases Kernels), - KernelBody (Aliases Kernels), - Stms (Aliases Kernels) + ( Pattern (Aliases GPU), + KernelBody (Aliases GPU), + Stms (Aliases GPU) ) ) lowerUpdatesIntoSegMap scope pat updates kspace kbody = do @@ -167,28 +167,28 @@ Just $ return (pe, mempty, ret, mempty) lowerUpdateIntoLoop :: - ( Bindable lore, - BinderOps lore, - Aliased lore, - LetDec lore ~ (als, Type), + ( Bindable rep, + BinderOps rep, + Aliased rep, + LetDec rep ~ (als, Type), MonadFreshNames m ) => - Scope lore -> - [DesiredUpdate (LetDec lore)] -> - Pattern lore -> - [(FParam lore, SubExp)] -> - [(FParam lore, SubExp)] -> - LoopForm lore -> - Body lore -> + Scope rep -> + [DesiredUpdate (LetDec rep)] -> + Pattern rep -> + [(FParam rep, SubExp)] -> + [(FParam rep, SubExp)] -> + LoopForm rep -> + Body rep -> Maybe ( m - ( [Stm lore], - [Stm lore], + ( [Stm rep], + [Stm rep], [Ident], [Ident], - [(FParam lore, SubExp)], - [(FParam lore, SubExp)], - Body lore + [(FParam rep, SubExp)], + [(FParam rep, SubExp)], + Body rep ) ) lowerUpdateIntoLoop scope updates pat ctx val form body = do @@ -239,9 +239,9 @@ resmap = zip (bodyResult body) $ patternValueIdents pat mkMerges :: - (MonadFreshNames m, Bindable lore) => + (MonadFreshNames m, Bindable rep) => [LoopResultSummary (als, Type)] -> - m ([(Param DeclType, SubExp)], [Stm lore], [Stm lore]) + m ([(Param DeclType, SubExp)], [Stm rep], [Stm rep]) mkMerges summaries = do ((origmerge, extramerge), (prebnds, postbnds)) <- runWriterT $ partitionEithers <$> mapM mkMerge summaries @@ -292,10 +292,10 @@ Left (inPatternAs summary) summariseLoop :: - ( Aliased lore, + ( Aliased rep, MonadFreshNames m ) => - Scope lore -> + Scope rep -> [DesiredUpdate (als, Type)] -> Names -> [(SubExp, Ident)] -> @@ -354,10 +354,10 @@ return (name, (cs, nm, dec, is)) manipulateResult :: - (Bindable lore, MonadFreshNames m) => - [LoopResultSummary (LetDec lore)] -> - IndexSubstitutions (LetDec lore) -> - m (Result, Stms lore) + (Bindable rep, MonadFreshNames m) => + [LoopResultSummary (LetDec rep)] -> + IndexSubstitutions (LetDec rep) -> + m (Result, Stms rep) manipulateResult summaries substs = do let (orig_ses, updated_ses) = partitionEithers $ map unchangedRes summaries (subst_ses, res_bnds) <- runWriterT $ zipWithM substRes updated_ses substs diff --git a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs --- a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs +++ b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs @@ -25,9 +25,9 @@ type IndexSubstitutions dec = [(VName, IndexSubstitution dec)] typeEnvFromSubstitutions :: - LetDec lore ~ dec => + LetDec rep ~ dec => IndexSubstitutions dec -> - Scope lore + Scope rep typeEnvFromSubstitutions = M.fromList . map (fromSubstitution . snd) where fromSubstitution (_, name, t, _) = @@ -36,42 +36,42 @@ -- | Perform the substitution. substituteIndices :: ( MonadFreshNames m, - BinderOps lore, - Bindable lore, - Aliased lore, - LetDec lore ~ dec + BinderOps rep, + Bindable rep, + Aliased rep, + LetDec rep ~ dec ) => IndexSubstitutions dec -> - Stms lore -> - m (IndexSubstitutions dec, Stms lore) + Stms rep -> + m (IndexSubstitutions dec, Stms rep) substituteIndices substs bnds = runBinderT (substituteIndicesInStms substs bnds) types where types = typeEnvFromSubstitutions substs substituteIndicesInStms :: - (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => - IndexSubstitutions (LetDec (Lore m)) -> - Stms (Lore m) -> - m (IndexSubstitutions (LetDec (Lore m))) + (MonadBinder m, Bindable (Rep m), Aliased (Rep m)) => + IndexSubstitutions (LetDec (Rep m)) -> + Stms (Rep m) -> + m (IndexSubstitutions (LetDec (Rep m))) substituteIndicesInStms = foldM substituteIndicesInStm substituteIndicesInStm :: - (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => - IndexSubstitutions (LetDec (Lore m)) -> - Stm (Lore m) -> - m (IndexSubstitutions (LetDec (Lore m))) -substituteIndicesInStm substs (Let pat lore e) = do + (MonadBinder m, Bindable (Rep m), Aliased (Rep m)) => + IndexSubstitutions (LetDec (Rep m)) -> + Stm (Rep m) -> + m (IndexSubstitutions (LetDec (Rep m))) +substituteIndicesInStm substs (Let pat rep e) = do e' <- substituteIndicesInExp substs e (substs', pat') <- substituteIndicesInPattern substs pat - addStm $ Let pat' lore e' + addStm $ Let pat' rep e' return substs' substituteIndicesInPattern :: - (MonadBinder m, LetDec (Lore m) ~ dec) => - IndexSubstitutions (LetDec (Lore m)) -> + (MonadBinder m, LetDec (Rep m) ~ dec) => + IndexSubstitutions (LetDec (Rep m)) -> PatternT dec -> - m (IndexSubstitutions (LetDec (Lore m)), PatternT dec) + m (IndexSubstitutions (LetDec (Rep m)), PatternT dec) substituteIndicesInPattern substs pat = do (substs', context) <- mapAccumLM sub substs $ patternContextElements pat (substs'', values) <- mapAccumLM sub substs' $ patternValueElements pat @@ -81,13 +81,13 @@ substituteIndicesInExp :: ( MonadBinder m, - Bindable (Lore m), - Aliased (Lore m), - LetDec (Lore m) ~ dec + Bindable (Rep m), + Aliased (Rep m), + LetDec (Rep m) ~ dec ) => - IndexSubstitutions (LetDec (Lore m)) -> - Exp (Lore m) -> - m (Exp (Lore m)) + IndexSubstitutions (LetDec (Rep m)) -> + Exp (Rep m) -> + m (Exp (Rep m)) substituteIndicesInExp substs (Op op) = do let used_in_op = filter ((`nameIn` freeIn op) . fst) substs var_substs <- fmap mconcat $ @@ -137,7 +137,7 @@ substituteIndicesInSubExp :: MonadBinder m => - IndexSubstitutions (LetDec (Lore m)) -> + IndexSubstitutions (LetDec (Rep m)) -> SubExp -> m SubExp substituteIndicesInSubExp substs (Var v) = @@ -147,7 +147,7 @@ substituteIndicesInVar :: MonadBinder m => - IndexSubstitutions (LetDec (Lore m)) -> + IndexSubstitutions (LetDec (Rep m)) -> VName -> m VName substituteIndicesInVar substs v @@ -161,10 +161,10 @@ return v substituteIndicesInBody :: - (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => - IndexSubstitutions (LetDec (Lore m)) -> - Body (Lore m) -> - m (Body (Lore m)) + (MonadBinder m, Bindable (Rep m), Aliased (Rep m)) => + IndexSubstitutions (LetDec (Rep m)) -> + Body (Rep m) -> + m (Body (Rep m)) substituteIndicesInBody substs (Body _ stms res) = do (substs', stms') <- inScopeOf stms $ diff --git a/src/Futhark/Optimise/InliningDeadFun.hs b/src/Futhark/Optimise/InliningDeadFun.hs --- a/src/Futhark/Optimise/InliningDeadFun.hs +++ b/src/Futhark/Optimise/InliningDeadFun.hs @@ -26,7 +26,7 @@ simplifyFun, ) import Futhark.Optimise.CSE -import Futhark.Optimise.Simplify.Lore (addScopeWisdom) +import Futhark.Optimise.Simplify.Rep (addScopeWisdom) import Futhark.Pass import Futhark.Transform.CopyPropagate ( copyPropagateInFun, diff --git a/src/Futhark/Optimise/ReuseAllocations.hs b/src/Futhark/Optimise/ReuseAllocations.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ReuseAllocations.hs @@ -0,0 +1,255 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE TypeFamilies #-} + +-- | This module implements an optimization that tries to statically reuse +-- kernel-level allocations. The goal is to lower the static memory usage, which +-- might allow more programs to run using intra-group parallelism. +module Futhark.Optimise.ReuseAllocations (optimise) where + +import Control.Exception +import Control.Monad.Reader +import Control.Monad.State.Strict +import Data.Function ((&)) +import Data.Map (Map, (!)) +import qualified Data.Map as M +import Data.Set (Set) +import qualified Data.Set as S +import qualified Futhark.Analysis.Interference as Interference +import qualified Futhark.Analysis.LastUse as LastUse +import Futhark.Binder.Class +import Futhark.Construct +import Futhark.IR.GPUMem +import qualified Futhark.Optimise.ReuseAllocations.GreedyColoring as GreedyColoring +import Futhark.Pass (Pass (..), PassM) +import qualified Futhark.Pass as Pass +import Futhark.Util (invertMap) + +-- | A mapping from allocation names to their size and space. +type Allocs = Map VName (SubExp, Space) + +getAllocsStm :: Stm GPUMem -> Allocs +getAllocsStm (Let (Pattern [] [PatElem name _]) _ (Op (Alloc se sp))) = + M.singleton name (se, sp) +getAllocsStm (Let _ _ (Op (Alloc _ _))) = error "impossible" +getAllocsStm (Let _ _ (If _ then_body else_body _)) = + foldMap getAllocsStm (bodyStms then_body) + <> foldMap getAllocsStm (bodyStms else_body) +getAllocsStm (Let _ _ (DoLoop _ _ _ body)) = + foldMap getAllocsStm (bodyStms body) +getAllocsStm _ = mempty + +getAllocsSegOp :: SegOp lvl GPUMem -> Allocs +getAllocsSegOp (SegMap _ _ _ body) = + foldMap getAllocsStm (kernelBodyStms body) +getAllocsSegOp (SegRed _ _ _ _ body) = + foldMap getAllocsStm (kernelBodyStms body) +getAllocsSegOp (SegScan _ _ _ _ body) = + foldMap getAllocsStm (kernelBodyStms body) +getAllocsSegOp (SegHist _ _ _ _ body) = + foldMap getAllocsStm (kernelBodyStms body) + +setAllocsStm :: Map VName SubExp -> Stm GPUMem -> Stm GPUMem +setAllocsStm m stm@(Let (Pattern [] [PatElem name _]) _ (Op (Alloc _ _))) + | Just s <- M.lookup name m = + stm {stmExp = BasicOp $ SubExp s} +setAllocsStm _ stm@(Let _ _ (Op (Alloc _ _))) = stm +setAllocsStm m stm@(Let _ _ (Op (Inner (SegOp segop)))) = + stm {stmExp = Op $ Inner $ SegOp $ setAllocsSegOp m segop} +setAllocsStm m stm@(Let _ _ (If cse then_body else_body dec)) = + stm + { stmExp = + If + cse + (then_body {bodyStms = setAllocsStm m <$> bodyStms then_body}) + (else_body {bodyStms = setAllocsStm m <$> bodyStms else_body}) + dec + } +setAllocsStm m stm@(Let _ _ (DoLoop ctx vals form body)) = + stm + { stmExp = + DoLoop + ctx + vals + form + (body {bodyStms = setAllocsStm m <$> bodyStms body}) + } +setAllocsStm _ stm = stm + +setAllocsSegOp :: + Map VName SubExp -> + SegOp lvl GPUMem -> + SegOp lvl GPUMem +setAllocsSegOp m (SegMap lvl sp tps body) = + SegMap lvl sp tps $ + body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body} +setAllocsSegOp m (SegRed lvl sp segbinops tps body) = + SegRed lvl sp segbinops tps $ + body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body} +setAllocsSegOp m (SegScan lvl sp segbinops tps body) = + SegScan lvl sp segbinops tps $ + body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body} +setAllocsSegOp m (SegHist lvl sp segbinops tps body) = + SegHist lvl sp segbinops tps $ + body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body} + +maxSubExp :: MonadBinder m => Set SubExp -> m SubExp +maxSubExp = helper . S.toList + where + helper (s1 : s2 : sexps) = do + z <- letSubExp "maxSubHelper" $ BasicOp $ BinOp (UMax Int64) s1 s2 + helper (z : sexps) + helper [s] = + return s + helper [] = error "impossible" + +definedInExp :: Exp GPUMem -> Set VName +definedInExp (Op (Inner (SegOp segop))) = + definedInSegOp segop +definedInExp (If _ then_body else_body _) = + foldMap definedInStm (bodyStms then_body) + <> foldMap definedInStm (bodyStms else_body) +definedInExp (DoLoop _ _ _ body) = + foldMap definedInStm $ bodyStms body +definedInExp _ = mempty + +definedInStm :: Stm GPUMem -> Set VName +definedInStm Let {stmPattern = Pattern ctx vals, stmExp} = + let definedInside = + ctx <> vals + & fmap patElemName + & S.fromList + in definedInExp stmExp <> definedInside + +definedInSegOp :: SegOp lvl GPUMem -> Set VName +definedInSegOp (SegMap _ _ _ body) = + foldMap definedInStm $ kernelBodyStms body +definedInSegOp (SegRed _ _ _ _ body) = + foldMap definedInStm $ kernelBodyStms body +definedInSegOp (SegScan _ _ _ _ body) = + foldMap definedInStm $ kernelBodyStms body +definedInSegOp (SegHist _ _ _ _ body) = + foldMap definedInStm $ kernelBodyStms body + +isKernelInvariant :: SegOp lvl GPUMem -> (SubExp, space) -> Bool +isKernelInvariant segop (Var vname, _) = + not $ vname `S.member` definedInSegOp segop +isKernelInvariant _ _ = True + +onKernelBodyStms :: + MonadBinder m => + SegOp lvl GPUMem -> + (Stms GPUMem -> m (Stms GPUMem)) -> + m (SegOp lvl GPUMem) +onKernelBodyStms (SegMap lvl space ts body) f = do + stms <- f $ kernelBodyStms body + return $ SegMap lvl space ts $ body {kernelBodyStms = stms} +onKernelBodyStms (SegRed lvl space binops ts body) f = do + stms <- f $ kernelBodyStms body + return $ SegRed lvl space binops ts $ body {kernelBodyStms = stms} +onKernelBodyStms (SegScan lvl space binops ts body) f = do + stms <- f $ kernelBodyStms body + return $ SegScan lvl space binops ts $ body {kernelBodyStms = stms} +onKernelBodyStms (SegHist lvl space binops ts body) f = do + stms <- f $ kernelBodyStms body + return $ SegHist lvl space binops ts $ body {kernelBodyStms = stms} + +-- | This is the actual optimiser. Given an interference graph and a `SegOp`, +-- replace allocations and references to memory blocks inside with a (hopefully) +-- reduced number of allocations. +optimiseKernel :: + (MonadBinder m, Rep m ~ GPUMem) => + Interference.Graph VName -> + SegOp lvl GPUMem -> + m (SegOp lvl GPUMem) +optimiseKernel graph segop0 = do + segop <- onKernelBodyStms segop0 $ onKernels $ optimiseKernel graph + let allocs = M.filter (isKernelInvariant segop) $ getAllocsSegOp segop + (colorspaces, coloring) = + GreedyColoring.colorGraph + (fmap snd allocs) + graph + (maxes, maxstms) <- + invertMap coloring + & M.elems + & mapM (maxSubExp . S.map (fst . (allocs !))) + & collectStms + (colors, stms) <- + assert (length maxes == M.size colorspaces) maxes + & zip [0 ..] + & mapM (\(i, x) -> letSubExp "color" $ Op $ Alloc x $ colorspaces ! i) + & collectStms + let segop' = setAllocsSegOp (fmap (colors !!) coloring) segop + return $ case segop' of + SegMap lvl sp tps body -> + SegMap lvl sp tps $ + body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body} + SegRed lvl sp binops tps body -> + SegRed lvl sp binops tps $ + body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body} + SegScan lvl sp binops tps body -> + SegScan lvl sp binops tps $ + body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body} + SegHist lvl sp binops tps body -> + SegHist lvl sp binops tps $ + body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body} + +-- | Helper function that modifies kernels found inside some statements. +onKernels :: + LocalScope GPUMem m => + (SegOp SegLevel GPUMem -> m (SegOp SegLevel GPUMem)) -> + Stms GPUMem -> + m (Stms GPUMem) +onKernels f = + mapM helper + where + helper stm@Let {stmExp = Op (Inner (SegOp segop))} = + inScopeOf stm $ do + exp' <- f segop + return $ stm {stmExp = Op $ Inner $ SegOp exp'} + helper stm@Let {stmExp = If c then_body else_body dec} = + inScopeOf stm $ do + then_body_stms <- f `onKernels` bodyStms then_body + else_body_stms <- f `onKernels` bodyStms else_body + return $ + stm + { stmExp = + If + c + (then_body {bodyStms = then_body_stms}) + (else_body {bodyStms = else_body_stms}) + dec + } + helper stm@Let {stmExp = DoLoop ctx vals form body} = + inScopeOf stm $ do + stms <- f `onKernels` bodyStms body + return $ stm {stmExp = DoLoop ctx vals form (body {bodyStms = stms})} + helper stm = + inScopeOf stm $ return stm + +-- | Perform the reuse-allocations optimization. +optimise :: Pass GPUMem GPUMem +optimise = + Pass "reuse allocations" "reuse allocations" $ \prog -> + let (lumap, _) = LastUse.analyseProg prog + graph = + foldMap + ( \f -> + runReader + ( Interference.analyseGPU lumap $ + bodyStms $ funDefBody f + ) + $ scopeOf f + ) + $ progFuns prog + in Pass.intraproceduralTransformation (onStms graph) prog + where + onStms :: + Interference.Graph VName -> + Scope GPUMem -> + Stms GPUMem -> + PassM (Stms GPUMem) + onStms graph scope stms = do + let m = localScope scope $ optimiseKernel graph `onKernels` stms + fmap fst $ modifyNameSource $ runState (runBinderT m mempty) diff --git a/src/Futhark/Optimise/ReuseAllocations/GreedyColoring.hs b/src/Futhark/Optimise/ReuseAllocations/GreedyColoring.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ReuseAllocations/GreedyColoring.hs @@ -0,0 +1,61 @@ +-- | Provides a greedy graph-coloring algorithm. +module Futhark.Optimise.ReuseAllocations.GreedyColoring (colorGraph, Coloring) where + +import Data.Function ((&)) +import Data.Map (Map, (!?)) +import qualified Data.Map as M +import Data.Maybe (fromMaybe) +import Data.Set (Set) +import qualified Data.Set as S +import qualified Futhark.Analysis.Interference as Interference + +-- | A map of values to their color, identified by an integer. +type Coloring a = Map a Int + +-- | A map of values to the set "neighbors" in the graph +type Neighbors a = Map a (Set a) + +-- | Computes the neighbor map of a graph. +neighbors :: Ord a => Interference.Graph a -> Neighbors a +neighbors = + S.foldr + ( \(x, y) acc -> + acc + & M.insertWith S.union x (S.singleton y) + & M.insertWith S.union y (S.singleton x) + ) + M.empty + +firstAvailable :: Eq space => Map Int space -> Set Int -> Int -> space -> (Map Int space, Int) +firstAvailable spaces xs i sp = + case (i `S.member` xs, spaces !? i) of + (False, Just sp') | sp' == sp -> (spaces, i) + (False, Nothing) -> (M.insert i sp spaces, i) + _ -> firstAvailable spaces xs (i + 1) sp + +colorNode :: + (Ord a, Eq space) => + Neighbors a -> + (a, space) -> + (Map Int space, Coloring a) -> + (Map Int space, Coloring a) +colorNode nbs (x, sp) (spaces, coloring) = + let nb_colors = + foldMap (maybe S.empty S.singleton . (coloring !?)) $ + fromMaybe mempty (nbs !? x) + (spaces', color) = firstAvailable spaces nb_colors 0 sp + in (spaces', M.insert x color coloring) + +-- | Graph coloring that takes into account the `space` of values. Two values +-- can only share the same color if they live in the same space. The result is +-- map from each color to a space and a map from each value in the input graph +-- to it's new color. +colorGraph :: + (Ord a, Ord space) => + Map a space -> + Interference.Graph a -> + (Map Int space, Coloring a) +colorGraph spaces graph = + let nodes = S.fromList $ M.toList spaces + nbs = neighbors graph + in S.foldr (colorNode nbs) mempty nodes diff --git a/src/Futhark/Optimise/Simplify.hs b/src/Futhark/Optimise/Simplify.hs --- a/src/Futhark/Optimise/Simplify.hs +++ b/src/Futhark/Optimise/Simplify.hs @@ -15,7 +15,7 @@ Engine.bindableSimpleOps, Engine.noExtraHoistBlockers, Engine.neverHoist, - Engine.SimplifiableLore, + Engine.SimplifiableRep, Engine.HoistBlockers, RuleBook, ) @@ -27,7 +27,7 @@ import Futhark.IR import Futhark.MonadFreshNames import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Optimise.Simplify.Rule import Futhark.Pass @@ -35,12 +35,12 @@ -- output, meaningful simplification may not have taken place - the -- order of bindings may simply have been rearranged. simplifyProg :: - Engine.SimplifiableLore lore => - Engine.SimpleOps lore -> - RuleBook (Engine.Wise lore) -> - Engine.HoistBlockers lore -> - Prog lore -> - PassM (Prog lore) + Engine.SimplifiableRep rep => + Engine.SimpleOps rep -> + RuleBook (Engine.Wise rep) -> + Engine.HoistBlockers rep -> + Prog rep -> + PassM (Prog rep) simplifyProg simpl rules blockers (Prog consts funs) = do (consts_vtable, consts') <- simplifyConsts @@ -84,13 +84,13 @@ -- | Run a simplification operation to convergence. simplifySomething :: - (MonadFreshNames m, Engine.SimplifiableLore lore) => - (a -> Engine.SimpleM lore b) -> + (MonadFreshNames m, Engine.SimplifiableRep rep) => + (a -> Engine.SimpleM rep b) -> (b -> a) -> - Engine.SimpleOps lore -> - RuleBook (Wise lore) -> - Engine.HoistBlockers lore -> - ST.SymbolTable (Wise lore) -> + Engine.SimpleOps rep -> + RuleBook (Wise rep) -> + Engine.HoistBlockers rep -> + ST.SymbolTable (Wise rep) -> a -> m a simplifySomething f g simpl rules blockers vtable x = do @@ -104,26 +104,26 @@ -- order of bindings may simply have been rearranged. Runs in a loop -- until convergence. simplifyFun :: - (MonadFreshNames m, Engine.SimplifiableLore lore) => - Engine.SimpleOps lore -> - RuleBook (Engine.Wise lore) -> - Engine.HoistBlockers lore -> - ST.SymbolTable (Wise lore) -> - FunDef lore -> - m (FunDef lore) + (MonadFreshNames m, Engine.SimplifiableRep rep) => + Engine.SimpleOps rep -> + RuleBook (Engine.Wise rep) -> + Engine.HoistBlockers rep -> + ST.SymbolTable (Wise rep) -> + FunDef rep -> + m (FunDef rep) simplifyFun = simplifySomething Engine.simplifyFun removeFunDefWisdom -- | Simplify just a single t'Lambda'. simplifyLambda :: ( MonadFreshNames m, - HasScope lore m, - Engine.SimplifiableLore lore + HasScope rep m, + Engine.SimplifiableRep rep ) => - Engine.SimpleOps lore -> - RuleBook (Engine.Wise lore) -> - Engine.HoistBlockers lore -> - Lambda lore -> - m (Lambda lore) + Engine.SimpleOps rep -> + RuleBook (Engine.Wise rep) -> + Engine.HoistBlockers rep -> + Lambda rep -> + m (Lambda rep) simplifyLambda simpl rules blockers orig_lam = do vtable <- ST.fromScope . addScopeWisdom <$> askScope simplifySomething @@ -137,13 +137,13 @@ -- | Simplify a sequence of 'Stm's. simplifyStms :: - (MonadFreshNames m, Engine.SimplifiableLore lore) => - Engine.SimpleOps lore -> - RuleBook (Engine.Wise lore) -> - Engine.HoistBlockers lore -> - Scope lore -> - Stms lore -> - m (ST.SymbolTable (Wise lore), Stms lore) + (MonadFreshNames m, Engine.SimplifiableRep rep) => + Engine.SimpleOps rep -> + RuleBook (Engine.Wise rep) -> + Engine.HoistBlockers rep -> + Scope rep -> + Stms rep -> + m (ST.SymbolTable (Wise rep), Stms rep) simplifyStms simpl rules blockers scope = simplifySomething f g simpl rules blockers vtable . (mempty,) where @@ -153,10 +153,10 @@ g = second $ fmap removeStmWisdom loopUntilConvergence :: - (MonadFreshNames m, Engine.SimplifiableLore lore) => - Engine.Env lore -> - Engine.SimpleOps lore -> - (a -> Engine.SimpleM lore b) -> + (MonadFreshNames m, Engine.SimplifiableRep rep) => + Engine.Env rep -> + Engine.SimpleOps rep -> + (a -> Engine.SimpleM rep b) -> (b -> a) -> a -> m a diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -48,7 +48,7 @@ localVtable, -- * Building blocks - SimplifiableLore, + SimplifiableRep, Simplifiable (..), simplifyStms, simplifyFun, @@ -61,7 +61,7 @@ hoistStms, blockIf, enterLoop, - module Futhark.Optimise.Simplify.Lore, + module Futhark.Optimise.Simplify.Rep, ) where @@ -75,35 +75,35 @@ import Futhark.Construct import Futhark.IR import Futhark.IR.Prop.Aliases -import Futhark.Optimise.Simplify.Lore +import Futhark.Optimise.Simplify.Rep import Futhark.Optimise.Simplify.Rule import Futhark.Util (nubOrd, splitFromEnd) -data HoistBlockers lore = HoistBlockers +data HoistBlockers rep = HoistBlockers { -- | Blocker for hoisting out of parallel loops. - blockHoistPar :: BlockPred (Wise lore), + blockHoistPar :: BlockPred (Wise rep), -- | Blocker for hoisting out of sequential loops. - blockHoistSeq :: BlockPred (Wise lore), + blockHoistSeq :: BlockPred (Wise rep), -- | Blocker for hoisting out of branches. - blockHoistBranch :: BlockPred (Wise lore), - isAllocation :: Stm (Wise lore) -> Bool + blockHoistBranch :: BlockPred (Wise rep), + isAllocation :: Stm (Wise rep) -> Bool } -noExtraHoistBlockers :: HoistBlockers lore +noExtraHoistBlockers :: HoistBlockers rep noExtraHoistBlockers = HoistBlockers neverBlocks neverBlocks neverBlocks (const False) -neverHoist :: HoistBlockers lore +neverHoist :: HoistBlockers rep neverHoist = HoistBlockers alwaysBlocks alwaysBlocks alwaysBlocks (const False) -data Env lore = Env - { envRules :: RuleBook (Wise lore), - envHoistBlockers :: HoistBlockers lore, - envVtable :: ST.SymbolTable (Wise lore) +data Env rep = Env + { envRules :: RuleBook (Wise rep), + envHoistBlockers :: HoistBlockers rep, + envVtable :: ST.SymbolTable (Wise rep) } -emptyEnv :: RuleBook (Wise lore) -> HoistBlockers lore -> Env lore +emptyEnv :: RuleBook (Wise rep) -> HoistBlockers rep -> Env rep emptyEnv rules blockers = Env { envRules = rules, @@ -111,33 +111,33 @@ envVtable = mempty } -type Protect m = SubExp -> Pattern (Lore m) -> Op (Lore m) -> Maybe (m ()) +type Protect m = SubExp -> Pattern (Rep m) -> Op (Rep m) -> Maybe (m ()) -data SimpleOps lore = SimpleOps +data SimpleOps rep = SimpleOps { mkExpDecS :: - ST.SymbolTable (Wise lore) -> - Pattern (Wise lore) -> - Exp (Wise lore) -> - SimpleM lore (ExpDec (Wise lore)), + ST.SymbolTable (Wise rep) -> + Pattern (Wise rep) -> + Exp (Wise rep) -> + SimpleM rep (ExpDec (Wise rep)), mkBodyS :: - ST.SymbolTable (Wise lore) -> - Stms (Wise lore) -> + ST.SymbolTable (Wise rep) -> + Stms (Wise rep) -> Result -> - SimpleM lore (Body (Wise lore)), + SimpleM rep (Body (Wise rep)), -- | Make a hoisted Op safe. The SubExp is a boolean -- that is true when the value of the statement will -- actually be used. - protectHoistedOpS :: Protect (Binder (Wise lore)), - opUsageS :: Op (Wise lore) -> UT.UsageTable, - simplifyOpS :: SimplifyOp lore (Op lore) + protectHoistedOpS :: Protect (Binder (Wise rep)), + opUsageS :: Op (Wise rep) -> UT.UsageTable, + simplifyOpS :: SimplifyOp rep (Op rep) } -type SimplifyOp lore op = op -> SimpleM lore (OpWithWisdom op, Stms (Wise lore)) +type SimplifyOp rep op = op -> SimpleM rep (OpWithWisdom op, Stms (Wise rep)) bindableSimpleOps :: - (SimplifiableLore lore, Bindable lore) => - SimplifyOp lore (Op lore) -> - SimpleOps lore + (SimplifiableRep rep, Bindable rep) => + SimplifyOp rep (Op rep) -> + SimpleOps rep bindableSimpleOps = SimpleOps mkExpDecS' mkBodyS' protectHoistedOpS' (const mempty) where @@ -145,10 +145,10 @@ mkBodyS' _ bnds res = return $ mkBody bnds res protectHoistedOpS' _ _ _ = Nothing -newtype SimpleM lore a +newtype SimpleM rep a = SimpleM ( ReaderT - (SimpleOps lore, Env lore) + (SimpleOps rep, Env rep) (State (VNameSource, Bool, Certificates)) a ) @@ -156,15 +156,15 @@ ( Applicative, Functor, Monad, - MonadReader (SimpleOps lore, Env lore), + MonadReader (SimpleOps rep, Env rep), MonadState (VNameSource, Bool, Certificates) ) -instance MonadFreshNames (SimpleM lore) where +instance MonadFreshNames (SimpleM rep) where putNameSource src = modify $ \(_, b, c) -> (src, b, c) getNameSource = gets $ \(a, _, _) -> a -instance SimplifiableLore lore => HasScope (Wise lore) (SimpleM lore) where +instance SimplifiableRep rep => HasScope (Wise rep) (SimpleM rep) where askScope = ST.toScope <$> askVtable lookupType name = do vtable <- askVtable @@ -177,37 +177,37 @@ ++ " in symbol table." instance - SimplifiableLore lore => - LocalScope (Wise lore) (SimpleM lore) + SimplifiableRep rep => + LocalScope (Wise rep) (SimpleM rep) where localScope types = localVtable (<> ST.fromScope types) runSimpleM :: - SimpleM lore a -> - SimpleOps lore -> - Env lore -> + SimpleM rep a -> + SimpleOps rep -> + Env rep -> VNameSource -> ((a, Bool), VNameSource) runSimpleM (SimpleM m) simpl env src = let (x, (src', b, _)) = runState (runReaderT m (simpl, env)) (src, False, mempty) in ((x, b), src') -askEngineEnv :: SimpleM lore (Env lore) +askEngineEnv :: SimpleM rep (Env rep) askEngineEnv = asks snd -asksEngineEnv :: (Env lore -> a) -> SimpleM lore a +asksEngineEnv :: (Env rep -> a) -> SimpleM rep a asksEngineEnv f = f <$> askEngineEnv -askVtable :: SimpleM lore (ST.SymbolTable (Wise lore)) +askVtable :: SimpleM rep (ST.SymbolTable (Wise rep)) askVtable = asksEngineEnv envVtable localVtable :: - (ST.SymbolTable (Wise lore) -> ST.SymbolTable (Wise lore)) -> - SimpleM lore a -> - SimpleM lore a + (ST.SymbolTable (Wise rep) -> ST.SymbolTable (Wise rep)) -> + SimpleM rep a -> + SimpleM rep a localVtable f = local $ \(ops, env) -> (ops, env {envVtable = f $ envVtable env}) -collectCerts :: SimpleM lore a -> SimpleM lore (a, Certificates) +collectCerts :: SimpleM rep a -> SimpleM rep (a, Certificates) collectCerts m = do x <- m (a, b, cs) <- get @@ -216,40 +216,40 @@ -- | Mark that we have changed something and it would be a good idea -- to re-run the simplifier. -changed :: SimpleM lore () +changed :: SimpleM rep () changed = modify $ \(src, _, cs) -> (src, True, cs) -usedCerts :: Certificates -> SimpleM lore () +usedCerts :: Certificates -> SimpleM rep () usedCerts cs = modify $ \(a, b, c) -> (a, b, cs <> c) -- | Indicate in the symbol table that we have descended into a loop. -enterLoop :: SimpleM lore a -> SimpleM lore a +enterLoop :: SimpleM rep a -> SimpleM rep a enterLoop = localVtable ST.deepen -bindFParams :: SimplifiableLore lore => [FParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a +bindFParams :: SimplifiableRep rep => [FParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a bindFParams params = localVtable $ ST.insertFParams params -bindLParams :: SimplifiableLore lore => [LParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a +bindLParams :: SimplifiableRep rep => [LParam (Wise rep)] -> SimpleM rep a -> SimpleM rep a bindLParams params = localVtable $ \vtable -> foldr ST.insertLParam vtable params bindArrayLParams :: - SimplifiableLore lore => - [LParam (Wise lore)] -> - SimpleM lore a -> - SimpleM lore a + SimplifiableRep rep => + [LParam (Wise rep)] -> + SimpleM rep a -> + SimpleM rep a bindArrayLParams params = localVtable $ \vtable -> foldl' (flip ST.insertLParam) vtable params bindMerge :: - SimplifiableLore lore => - [(FParam (Wise lore), SubExp, SubExp)] -> - SimpleM lore a -> - SimpleM lore a + SimplifiableRep rep => + [(FParam (Wise rep), SubExp, SubExp)] -> + SimpleM rep a -> + SimpleM rep a bindMerge = localVtable . ST.insertLoopMerge -bindLoopVar :: SimplifiableLore lore => VName -> IntType -> SubExp -> SimpleM lore a -> SimpleM lore a +bindLoopVar :: SimplifiableRep rep => VName -> IntType -> SubExp -> SimpleM rep a -> SimpleM rep a bindLoopVar var it bound = localVtable $ ST.insertLoopVar var it bound @@ -258,14 +258,14 @@ -- them. (This means such hoisting is not worth it unless they are in -- turn hoisted out of a loop somewhere.) protectIfHoisted :: - SimplifiableLore lore => + SimplifiableRep rep => -- | Branch condition. SubExp -> -- | Which side of the branch are we -- protecting here? Bool -> - SimpleM lore (a, Stms (Wise lore)) -> - SimpleM lore (a, Stms (Wise lore)) + SimpleM rep (a, Stms (Wise rep)) -> + SimpleM rep (a, Stms (Wise rep)) protectIfHoisted cond side m = do (x, stms) <- m ops <- asks $ protectHoistedOpS . fst @@ -286,12 +286,12 @@ -- loops, but they most be protected by adding a branch on top of -- them. protectLoopHoisted :: - SimplifiableLore lore => - [(FParam (Wise lore), SubExp)] -> - [(FParam (Wise lore), SubExp)] -> - LoopForm (Wise lore) -> - SimpleM lore (a, Stms (Wise lore)) -> - SimpleM lore (a, Stms (Wise lore)) + SimplifiableRep rep => + [(FParam (Wise rep), SubExp)] -> + [(FParam (Wise rep), SubExp)] -> + LoopForm (Wise rep) -> + SimpleM rep (a, Stms (Wise rep)) -> + SimpleM rep (a, Stms (Wise rep)) protectLoopHoisted ctx val form m = do (x, stms) <- m ops <- asks $ protectHoistedOpS . fst @@ -317,9 +317,9 @@ protectIf :: MonadBinder m => Protect m -> - (Exp (Lore m) -> Bool) -> + (Exp (Rep m) -> Bool) -> SubExp -> - Stm (Lore m) -> + Stm (Rep m) -> m () protectIf _ @@ -362,7 +362,7 @@ protectIf _ _ _ stm = addStm stm -makeSafe :: Exp lore -> Maybe (Exp lore) +makeSafe :: Exp rep -> Maybe (Exp rep) makeSafe (BasicOp (BinOp (SDiv t _) x y)) = Just $ BasicOp (BinOp (SDiv t Safe) x y) makeSafe (BasicOp (BinOp (SDivUp t _) x y)) = @@ -382,7 +382,7 @@ makeSafe _ = Nothing -emptyOfType :: MonadBinder m => [VName] -> Type -> m (Exp (Lore m)) +emptyOfType :: MonadBinder m => [VName] -> Type -> m (Exp (Rep m)) emptyOfType _ Mem {} = error "emptyOfType: Cannot hoist non-existential memory." emptyOfType _ Acc {} = @@ -399,21 +399,21 @@ -- | Statements that are not worth hoisting out of loops, because they -- are unsafe, and added safety (by 'protectLoopHoisted') may inhibit -- further optimisation.. -notWorthHoisting :: ASTLore lore => BlockPred lore +notWorthHoisting :: ASTRep rep => BlockPred rep notWorthHoisting _ _ (Let pat _ e) = not (safeExp e) && any ((> 0) . arrayRank) (patternTypes pat) hoistStms :: - SimplifiableLore lore => - RuleBook (Wise lore) -> - BlockPred (Wise lore) -> - ST.SymbolTable (Wise lore) -> + SimplifiableRep rep => + RuleBook (Wise rep) -> + BlockPred (Wise rep) -> + ST.SymbolTable (Wise rep) -> UT.UsageTable -> - Stms (Wise lore) -> + Stms (Wise rep) -> SimpleM - lore - ( Stms (Wise lore), - Stms (Wise lore) + rep + ( Stms (Wise rep), + Stms (Wise rep) ) hoistStms rules block vtable uses orig_stms = do (blocked, hoisted) <- simplifyStmsBottomUp vtable uses orig_stms @@ -465,9 +465,9 @@ return (uses'', stms' ++ stms) blockUnhoistedDeps :: - ASTLore lore => - [Either (Stm lore) (Stm lore)] -> - [Either (Stm lore) (Stm lore)] + ASTRep rep => + [Either (Stm rep) (Stm rep)] -> + [Either (Stm rep) (Stm rep)] blockUnhoistedDeps = snd . mapAccumL block mempty where block blocked (Left need) = @@ -478,15 +478,15 @@ | otherwise = (blocked, Right need) -provides :: Stm lore -> [VName] +provides :: Stm rep -> [VName] provides = patternNames . stmPattern expandUsage :: - (ASTLore lore, Aliased lore) => - (Stm lore -> UT.UsageTable) -> - ST.SymbolTable lore -> + (ASTRep rep, Aliased rep) => + (Stm rep -> UT.UsageTable) -> + ST.SymbolTable rep -> UT.UsageTable -> - Stm lore -> + Stm rep -> UT.UsageTable expandUsage usageInStm vtable utable stm@(Let pat _ e) = UT.expand (`ST.lookupAliases` vtable) (usageInStm stm <> usageThroughAliases) @@ -504,47 +504,47 @@ uses <- UT.lookup name utable return $ mconcat $ map (`UT.usage` uses) $ namesToList aliases -type BlockPred lore = ST.SymbolTable lore -> UT.UsageTable -> Stm lore -> Bool +type BlockPred rep = ST.SymbolTable rep -> UT.UsageTable -> Stm rep -> Bool -neverBlocks :: BlockPred lore +neverBlocks :: BlockPred rep neverBlocks _ _ _ = False -alwaysBlocks :: BlockPred lore +alwaysBlocks :: BlockPred rep alwaysBlocks _ _ _ = True -isFalse :: Bool -> BlockPred lore +isFalse :: Bool -> BlockPred rep isFalse b _ _ _ = not b -orIf :: BlockPred lore -> BlockPred lore -> BlockPred lore +orIf :: BlockPred rep -> BlockPred rep -> BlockPred rep orIf p1 p2 body vtable need = p1 body vtable need || p2 body vtable need -andAlso :: BlockPred lore -> BlockPred lore -> BlockPred lore +andAlso :: BlockPred rep -> BlockPred rep -> BlockPred rep andAlso p1 p2 body vtable need = p1 body vtable need && p2 body vtable need -isConsumed :: BlockPred lore +isConsumed :: BlockPred rep isConsumed _ utable = any (`UT.isConsumed` utable) . patternNames . stmPattern -isOp :: BlockPred lore +isOp :: BlockPred rep isOp _ _ (Let _ _ Op {}) = True isOp _ _ _ = False constructBody :: - SimplifiableLore lore => - Stms (Wise lore) -> + SimplifiableRep rep => + Stms (Wise rep) -> Result -> - SimpleM lore (Body (Wise lore)) + SimpleM rep (Body (Wise rep)) constructBody stms res = fmap fst . runBinder . buildBody_ $ do addStms stms pure res -type SimplifiedBody lore a = ((a, UT.UsageTable), Stms (Wise lore)) +type SimplifiedBody rep a = ((a, UT.UsageTable), Stms (Wise rep)) blockIf :: - SimplifiableLore lore => - BlockPred (Wise lore) -> - SimpleM lore (SimplifiedBody lore a) -> - SimpleM lore ((Stms (Wise lore), a), Stms (Wise lore)) + SimplifiableRep rep => + BlockPred (Wise rep) -> + SimpleM rep (SimplifiedBody rep a) -> + SimpleM rep ((Stms (Wise rep), a), Stms (Wise rep)) blockIf block m = do ((x, usages), stms) <- m vtable <- askVtable @@ -552,10 +552,10 @@ (blocked, hoisted) <- hoistStms rules block vtable usages stms return ((blocked, x), hoisted) -hasFree :: ASTLore lore => Names -> BlockPred lore +hasFree :: ASTRep rep => Names -> BlockPred rep hasFree ks _ _ need = ks `namesIntersect` freeIn need -isNotSafe :: ASTLore lore => BlockPred lore +isNotSafe :: ASTRep rep => BlockPred rep isNotSafe _ _ = not . safeExp . stmExp isInPlaceBound :: BlockPred m @@ -564,13 +564,13 @@ isUpdate (BasicOp Update {}) = True isUpdate _ = False -isNotCheap :: ASTLore lore => BlockPred lore +isNotCheap :: ASTRep rep => BlockPred rep isNotCheap _ _ = not . cheapStm -cheapStm :: ASTLore lore => Stm lore -> Bool +cheapStm :: ASTRep rep => Stm rep -> Bool cheapStm = cheapExp . stmExp -cheapExp :: ASTLore lore => Exp lore -> Bool +cheapExp :: ASTRep rep => Exp rep -> Bool cheapExp (BasicOp BinOp {}) = True cheapExp (BasicOp SubExp {}) = True cheapExp (BasicOp UnOp {}) = True @@ -587,24 +587,24 @@ cheapExp _ = True -- Used to be False, but -- let's try it out. -stmIs :: (Stm lore -> Bool) -> BlockPred lore +stmIs :: (Stm rep -> Bool) -> BlockPred rep stmIs f _ _ = f -loopInvariantStm :: ASTLore lore => ST.SymbolTable lore -> Stm lore -> Bool +loopInvariantStm :: ASTRep rep => ST.SymbolTable rep -> Stm rep -> Bool loopInvariantStm vtable = all (`nameIn` ST.availableAtClosestLoop vtable) . namesToList . freeIn hoistCommon :: - SimplifiableLore lore => + SimplifiableRep rep => SubExp -> IfSort -> - SimplifiedBody lore Result -> - SimplifiedBody lore Result -> + SimplifiedBody rep Result -> + SimplifiedBody rep Result -> SimpleM - lore - ( Body (Wise lore), - Body (Wise lore), - Stms (Wise lore) + rep + ( Body (Wise rep), + Body (Wise rep), + Stms (Wise rep) ) hoistCommon cond ifsort ((res1, usages1), stms1) ((res2, usages2), stms2) = do is_alloc_fun <- asksEngineEnv $ isAllocation . envHoistBlockers @@ -663,10 +663,10 @@ -- | Simplify a single body. The @[Diet]@ only covers the value -- elements, because the context cannot be consumed. simplifyBody :: - SimplifiableLore lore => + SimplifiableRep rep => [Diet] -> - Body lore -> - SimpleM lore (SimplifiedBody lore Result) + Body rep -> + SimpleM rep (SimplifiedBody rep Result) simplifyBody ds (Body _ bnds res) = simplifyStms bnds $ do res' <- simplifyResult ds res @@ -675,10 +675,10 @@ -- | Simplify a single 'Result'. The @[Diet]@ only covers the value -- elements, because the context cannot be consumed. simplifyResult :: - SimplifiableLore lore => + SimplifiableRep rep => [Diet] -> Result -> - SimpleM lore (Result, UT.UsageTable) + SimpleM rep (Result, UT.UsageTable) simplifyResult ds res = do let (ctx_res, val_res) = splitFromEnd (length ds) res -- Copy propagation is a little trickier here, because there is no @@ -712,10 +712,10 @@ checkForVar _ = mempty simplifyStms :: - SimplifiableLore lore => - Stms lore -> - SimpleM lore (a, Stms (Wise lore)) -> - SimpleM lore (a, Stms (Wise lore)) + SimplifiableRep rep => + Stms rep -> + SimpleM rep (a, Stms (Wise rep)) -> + SimpleM rep (a, Stms (Wise rep)) simplifyStms stms m = case stmsHead stms of Nothing -> inspectStms mempty m @@ -729,17 +729,17 @@ simplifyStms stms' m inspectStm :: - SimplifiableLore lore => - Stm (Wise lore) -> - SimpleM lore (a, Stms (Wise lore)) -> - SimpleM lore (a, Stms (Wise lore)) + SimplifiableRep rep => + Stm (Wise rep) -> + SimpleM rep (a, Stms (Wise rep)) -> + SimpleM rep (a, Stms (Wise rep)) inspectStm = inspectStms . oneStm inspectStms :: - SimplifiableLore lore => - Stms (Wise lore) -> - SimpleM lore (a, Stms (Wise lore)) -> - SimpleM lore (a, Stms (Wise lore)) + SimplifiableRep rep => + Stms (Wise rep) -> + SimpleM rep (a, Stms (Wise rep)) -> + SimpleM rep (a, Stms (Wise rep)) inspectStms stms m = case stmsHead stms of Nothing -> m @@ -753,15 +753,15 @@ (x, stms'') <- localVtable (ST.insertStm stm) $ inspectStms stms' m return (x, oneStm stm <> stms'') -simplifyOp :: Op lore -> SimpleM lore (Op (Wise lore), Stms (Wise lore)) +simplifyOp :: Op rep -> SimpleM rep (Op (Wise rep), Stms (Wise rep)) simplifyOp op = do f <- asks $ simplifyOpS . fst f op simplifyExp :: - SimplifiableLore lore => - Exp lore -> - SimpleM lore (Exp (Wise lore), Stms (Wise lore)) + SimplifiableRep rep => + Exp rep -> + SimpleM rep (Exp (Wise rep), Stms (Wise rep)) simplifyExp (If cond tbranch fbranch (IfDec ts ifsort)) = do -- Here, we have to check whether 'cond' puts a bound on some free -- variable, and if so, chomp it. We should also try to do CSE @@ -862,9 +862,9 @@ return (e', mempty) simplifyExpBase :: - SimplifiableLore lore => - Exp lore -> - SimpleM lore (Exp (Wise lore)) + SimplifiableRep rep => + Exp rep -> + SimpleM rep (Exp (Wise rep)) simplifyExpBase = mapExpM hoist where hoist = @@ -887,21 +887,21 @@ error "Unhandled Op in simplification engine." } -type SimplifiableLore lore = - ( ASTLore lore, - Simplifiable (LetDec lore), - Simplifiable (FParamInfo lore), - Simplifiable (LParamInfo lore), - Simplifiable (RetType lore), - Simplifiable (BranchType lore), - CanBeWise (Op lore), - ST.IndexOp (OpWithWisdom (Op lore)), - BinderOps (Wise lore), - IsOp (Op lore) +type SimplifiableRep rep = + ( ASTRep rep, + Simplifiable (LetDec rep), + Simplifiable (FParamInfo rep), + Simplifiable (LParamInfo rep), + Simplifiable (RetType rep), + Simplifiable (BranchType rep), + CanBeWise (Op rep), + ST.IndexOp (OpWithWisdom (Op rep)), + BinderOps (Wise rep), + IsOp (Op rep) ) class Simplifiable e where - simplify :: SimplifiableLore lore => e -> SimpleM lore e + simplify :: SimplifiableRep rep => e -> SimpleM rep e instance (Simplifiable a, Simplifiable b) => Simplifiable (a, b) where simplify (x, y) = (,) <$> simplify x <*> simplify y @@ -940,15 +940,15 @@ return $ Constant v simplifyPattern :: - (SimplifiableLore lore, Simplifiable dec) => + (SimplifiableRep rep, Simplifiable dec) => PatternT dec -> - SimpleM lore (PatternT dec) + SimpleM rep (PatternT dec) simplifyPattern pat = Pattern <$> mapM inspect (patternContextElements pat) <*> mapM inspect (patternValueElements pat) where - inspect (PatElem name lore) = PatElem name <$> simplify lore + inspect (PatElem name rep) = PatElem name <$> simplify rep instance Simplifiable () where simplify = pure @@ -992,25 +992,25 @@ simplify (DimSlice i n s) = DimSlice <$> simplify i <*> simplify n <*> simplify s simplifyLambda :: - SimplifiableLore lore => - Lambda lore -> - SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) + SimplifiableRep rep => + Lambda rep -> + SimpleM rep (Lambda (Wise rep), Stms (Wise rep)) simplifyLambda lam = do par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers simplifyLambdaMaybeHoist par_blocker lam simplifyLambdaNoHoisting :: - SimplifiableLore lore => - Lambda lore -> - SimpleM lore (Lambda (Wise lore)) + SimplifiableRep rep => + Lambda rep -> + SimpleM rep (Lambda (Wise rep)) simplifyLambdaNoHoisting lam = fst <$> simplifyLambdaMaybeHoist (isFalse False) lam simplifyLambdaMaybeHoist :: - SimplifiableLore lore => - BlockPred (Wise lore) -> - Lambda lore -> - SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) + SimplifiableRep rep => + BlockPred (Wise rep) -> + Lambda rep -> + SimpleM rep (Lambda (Wise rep), Stms (Wise rep)) simplifyLambdaMaybeHoist blocked lam@(Lambda params body rettype) = do params' <- mapM (traverse simplify) params let paramnames = namesFromList $ boundByLambda lam @@ -1041,15 +1041,15 @@ _ -> return [idd] insertAllStms :: - SimplifiableLore lore => - SimpleM lore (SimplifiedBody lore Result) -> - SimpleM lore (Body (Wise lore)) + SimplifiableRep rep => + SimpleM rep (SimplifiedBody rep Result) -> + SimpleM rep (Body (Wise rep)) insertAllStms = uncurry constructBody . fst <=< blockIf (isFalse False) simplifyFun :: - SimplifiableLore lore => - FunDef lore -> - SimpleM lore (FunDef (Wise lore)) + SimplifiableRep rep => + FunDef rep -> + SimpleM rep (FunDef (Wise rep)) simplifyFun (FunDef entry attrs fname rettype params body) = do rettype' <- simplify rettype params' <- mapM (traverse simplify) params diff --git a/src/Futhark/Optimise/Simplify/Lore.hs b/src/Futhark/Optimise/Simplify/Lore.hs deleted file mode 100644 --- a/src/Futhark/Optimise/Simplify/Lore.hs +++ /dev/null @@ -1,284 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE UndecidableInstances #-} - --- | Definition of the lore used by the simplification engine. -module Futhark.Optimise.Simplify.Lore - ( Wise, - VarWisdom (..), - ExpWisdom, - removeStmWisdom, - removeLambdaWisdom, - removeFunDefWisdom, - removeExpWisdom, - removePatternWisdom, - removeBodyWisdom, - removeScopeWisdom, - addScopeWisdom, - addWisdomToPattern, - mkWiseBody, - mkWiseLetStm, - mkWiseExpDec, - CanBeWise (..), - ) -where - -import Control.Category -import Control.Monad.Identity -import Control.Monad.Reader -import qualified Data.Kind -import qualified Data.Map.Strict as M -import Futhark.Analysis.Rephrase -import Futhark.Binder -import Futhark.IR -import Futhark.IR.Aliases - ( AliasDec (..), - ConsumedInExp, - VarAliases, - unAliases, - ) -import qualified Futhark.IR.Aliases as Aliases -import Futhark.IR.Prop.Aliases -import Futhark.Transform.Rename -import Futhark.Transform.Substitute -import Futhark.Util.Pretty -import Prelude hiding (id, (.)) - -data Wise lore - --- | The wisdom of the let-bound variable. -newtype VarWisdom = VarWisdom {varWisdomAliases :: VarAliases} - deriving (Eq, Ord, Show) - -instance Rename VarWisdom where - rename = substituteRename - -instance Substitute VarWisdom where - substituteNames substs (VarWisdom als) = - VarWisdom (substituteNames substs als) - -instance FreeIn VarWisdom where - freeIn' (VarWisdom als) = freeIn' als - --- | Wisdom about an expression. -data ExpWisdom = ExpWisdom - { _expWisdomConsumed :: ConsumedInExp, - expWisdomFree :: AliasDec - } - deriving (Eq, Ord, Show) - -instance FreeIn ExpWisdom where - freeIn' = mempty - -instance FreeDec ExpWisdom where - precomputed = const . fvNames . unAliases . expWisdomFree - -instance Substitute ExpWisdom where - substituteNames substs (ExpWisdom cons free) = - ExpWisdom - (substituteNames substs cons) - (substituteNames substs free) - -instance Rename ExpWisdom where - rename = substituteRename - --- | Wisdom about a body. -data BodyWisdom = BodyWisdom - { bodyWisdomAliases :: [VarAliases], - bodyWisdomConsumed :: ConsumedInExp, - bodyWisdomFree :: AliasDec - } - deriving (Eq, Ord, Show) - -instance Rename BodyWisdom where - rename = substituteRename - -instance Substitute BodyWisdom where - substituteNames substs (BodyWisdom als cons free) = - BodyWisdom - (substituteNames substs als) - (substituteNames substs cons) - (substituteNames substs free) - -instance FreeIn BodyWisdom where - freeIn' (BodyWisdom als cons free) = - freeIn' als <> freeIn' cons <> freeIn' free - -instance FreeDec BodyWisdom where - precomputed = const . fvNames . unAliases . bodyWisdomFree - -instance - ( Decorations lore, - CanBeWise (Op lore) - ) => - Decorations (Wise lore) - where - type LetDec (Wise lore) = (VarWisdom, LetDec lore) - type ExpDec (Wise lore) = (ExpWisdom, ExpDec lore) - type BodyDec (Wise lore) = (BodyWisdom, BodyDec lore) - type FParamInfo (Wise lore) = FParamInfo lore - type LParamInfo (Wise lore) = LParamInfo lore - type RetType (Wise lore) = RetType lore - type BranchType (Wise lore) = BranchType lore - type Op (Wise lore) = OpWithWisdom (Op lore) - -withoutWisdom :: - (HasScope (Wise lore) m, Monad m) => - ReaderT (Scope lore) m a -> - m a -withoutWisdom m = do - scope <- asksScope removeScopeWisdom - runReaderT m scope - -instance (ASTLore lore, CanBeWise (Op lore)) => ASTLore (Wise lore) where - expTypesFromPattern = - withoutWisdom . expTypesFromPattern . removePatternWisdom - -instance Pretty VarWisdom where - ppr _ = ppr () - -instance (PrettyLore lore, CanBeWise (Op lore)) => PrettyLore (Wise lore) where - ppExpLore (_, dec) = ppExpLore dec . removeExpWisdom - -instance AliasesOf (VarWisdom, dec) where - aliasesOf = unAliases . varWisdomAliases . fst - -instance (ASTLore lore, CanBeWise (Op lore)) => Aliased (Wise lore) where - bodyAliases = map unAliases . bodyWisdomAliases . fst . bodyDec - consumedInBody = unAliases . bodyWisdomConsumed . fst . bodyDec - -removeWisdom :: CanBeWise (Op lore) => Rephraser Identity (Wise lore) lore -removeWisdom = - Rephraser - { rephraseExpLore = return . snd, - rephraseLetBoundLore = return . snd, - rephraseBodyLore = return . snd, - rephraseFParamLore = return, - rephraseLParamLore = return, - rephraseRetType = return, - rephraseBranchType = return, - rephraseOp = return . removeOpWisdom - } - -removeScopeWisdom :: Scope (Wise lore) -> Scope lore -removeScopeWisdom = M.map unAlias - where - unAlias (LetName (_, dec)) = LetName dec - unAlias (FParamName dec) = FParamName dec - unAlias (LParamName dec) = LParamName dec - unAlias (IndexName it) = IndexName it - -addScopeWisdom :: Scope lore -> Scope (Wise lore) -addScopeWisdom = M.map alias - where - alias (LetName dec) = LetName (VarWisdom mempty, dec) - alias (FParamName dec) = FParamName dec - alias (LParamName dec) = LParamName dec - alias (IndexName it) = IndexName it - -removeFunDefWisdom :: CanBeWise (Op lore) => FunDef (Wise lore) -> FunDef lore -removeFunDefWisdom = runIdentity . rephraseFunDef removeWisdom - -removeStmWisdom :: CanBeWise (Op lore) => Stm (Wise lore) -> Stm lore -removeStmWisdom = runIdentity . rephraseStm removeWisdom - -removeLambdaWisdom :: CanBeWise (Op lore) => Lambda (Wise lore) -> Lambda lore -removeLambdaWisdom = runIdentity . rephraseLambda removeWisdom - -removeBodyWisdom :: CanBeWise (Op lore) => Body (Wise lore) -> Body lore -removeBodyWisdom = runIdentity . rephraseBody removeWisdom - -removeExpWisdom :: CanBeWise (Op lore) => Exp (Wise lore) -> Exp lore -removeExpWisdom = runIdentity . rephraseExp removeWisdom - -removePatternWisdom :: PatternT (VarWisdom, a) -> PatternT a -removePatternWisdom = runIdentity . rephrasePattern (return . snd) - -addWisdomToPattern :: - (ASTLore lore, CanBeWise (Op lore)) => - Pattern lore -> - Exp (Wise lore) -> - Pattern (Wise lore) -addWisdomToPattern pat e = - Pattern (map f ctx) (map f val) - where - (ctx, val) = Aliases.mkPatternAliases pat e - f pe = - let (als, dec) = patElemDec pe - in pe `setPatElemLore` (VarWisdom als, dec) - -mkWiseBody :: - (ASTLore lore, CanBeWise (Op lore)) => - BodyDec lore -> - Stms (Wise lore) -> - Result -> - Body (Wise lore) -mkWiseBody innerlore bnds res = - Body - ( BodyWisdom aliases consumed (AliasDec $ freeIn $ freeInStmsAndRes bnds res), - innerlore - ) - bnds - res - where - (aliases, consumed) = Aliases.mkBodyAliases bnds res - -mkWiseLetStm :: - (ASTLore lore, CanBeWise (Op lore)) => - Pattern lore -> - StmAux (ExpDec lore) -> - Exp (Wise lore) -> - Stm (Wise lore) -mkWiseLetStm pat (StmAux cs attrs dec) e = - let pat' = addWisdomToPattern pat e - in Let pat' (StmAux cs attrs $ mkWiseExpDec pat' dec e) e - -mkWiseExpDec :: - (ASTLore lore, CanBeWise (Op lore)) => - Pattern (Wise lore) -> - ExpDec lore -> - Exp (Wise lore) -> - ExpDec (Wise lore) -mkWiseExpDec pat explore e = - ( ExpWisdom - (AliasDec $ consumedInExp e) - (AliasDec $ freeIn pat <> freeIn explore <> freeIn e), - explore - ) - -instance - ( Bindable lore, - CanBeWise (Op lore) - ) => - Bindable (Wise lore) - where - mkExpPat ctx val e = - addWisdomToPattern (mkExpPat ctx val $ removeExpWisdom e) e - - mkExpDec pat e = - mkWiseExpDec pat (mkExpDec (removePatternWisdom pat) $ removeExpWisdom e) e - - mkLetNames names e = do - env <- asksScope removeScopeWisdom - flip runReaderT env $ do - Let pat dec _ <- mkLetNames names $ removeExpWisdom e - return $ mkWiseLetStm pat dec e - - mkBody bnds res = - let Body bodylore _ _ = mkBody (fmap removeStmWisdom bnds) res - in mkWiseBody bodylore bnds res - -class - ( AliasedOp (OpWithWisdom op), - IsOp (OpWithWisdom op) - ) => - CanBeWise op - where - type OpWithWisdom op :: Data.Kind.Type - removeOpWisdom :: OpWithWisdom op -> op - -instance CanBeWise () where - type OpWithWisdom () = () - removeOpWisdom () = () diff --git a/src/Futhark/Optimise/Simplify/Rep.hs b/src/Futhark/Optimise/Simplify/Rep.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rep.hs @@ -0,0 +1,278 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | Representation used by the simplification engine. +module Futhark.Optimise.Simplify.Rep + ( Wise, + VarWisdom (..), + ExpWisdom, + removeStmWisdom, + removeLambdaWisdom, + removeFunDefWisdom, + removeExpWisdom, + removePatternWisdom, + removeBodyWisdom, + removeScopeWisdom, + addScopeWisdom, + addWisdomToPattern, + mkWiseBody, + mkWiseLetStm, + mkWiseExpDec, + CanBeWise (..), + ) +where + +import Control.Category +import Control.Monad.Identity +import Control.Monad.Reader +import qualified Data.Kind +import qualified Data.Map.Strict as M +import Futhark.Analysis.Rephrase +import Futhark.Binder +import Futhark.IR +import Futhark.IR.Aliases + ( AliasDec (..), + ConsumedInExp, + VarAliases, + unAliases, + ) +import qualified Futhark.IR.Aliases as Aliases +import Futhark.IR.Prop.Aliases +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import Futhark.Util.Pretty +import Prelude hiding (id, (.)) + +data Wise rep + +-- | The wisdom of the let-bound variable. +newtype VarWisdom = VarWisdom {varWisdomAliases :: VarAliases} + deriving (Eq, Ord, Show) + +instance Rename VarWisdom where + rename = substituteRename + +instance Substitute VarWisdom where + substituteNames substs (VarWisdom als) = + VarWisdom (substituteNames substs als) + +instance FreeIn VarWisdom where + freeIn' (VarWisdom als) = freeIn' als + +-- | Wisdom about an expression. +data ExpWisdom = ExpWisdom + { _expWisdomConsumed :: ConsumedInExp, + expWisdomFree :: AliasDec + } + deriving (Eq, Ord, Show) + +instance FreeIn ExpWisdom where + freeIn' = mempty + +instance FreeDec ExpWisdom where + precomputed = const . fvNames . unAliases . expWisdomFree + +instance Substitute ExpWisdom where + substituteNames substs (ExpWisdom cons free) = + ExpWisdom + (substituteNames substs cons) + (substituteNames substs free) + +instance Rename ExpWisdom where + rename = substituteRename + +-- | Wisdom about a body. +data BodyWisdom = BodyWisdom + { bodyWisdomAliases :: [VarAliases], + bodyWisdomConsumed :: ConsumedInExp, + bodyWisdomFree :: AliasDec + } + deriving (Eq, Ord, Show) + +instance Rename BodyWisdom where + rename = substituteRename + +instance Substitute BodyWisdom where + substituteNames substs (BodyWisdom als cons free) = + BodyWisdom + (substituteNames substs als) + (substituteNames substs cons) + (substituteNames substs free) + +instance FreeIn BodyWisdom where + freeIn' (BodyWisdom als cons free) = + freeIn' als <> freeIn' cons <> freeIn' free + +instance FreeDec BodyWisdom where + precomputed = const . fvNames . unAliases . bodyWisdomFree + +instance (RepTypes rep, CanBeWise (Op rep)) => RepTypes (Wise rep) where + type LetDec (Wise rep) = (VarWisdom, LetDec rep) + type ExpDec (Wise rep) = (ExpWisdom, ExpDec rep) + type BodyDec (Wise rep) = (BodyWisdom, BodyDec rep) + type FParamInfo (Wise rep) = FParamInfo rep + type LParamInfo (Wise rep) = LParamInfo rep + type RetType (Wise rep) = RetType rep + type BranchType (Wise rep) = BranchType rep + type Op (Wise rep) = OpWithWisdom (Op rep) + +withoutWisdom :: + (HasScope (Wise rep) m, Monad m) => + ReaderT (Scope rep) m a -> + m a +withoutWisdom m = do + scope <- asksScope removeScopeWisdom + runReaderT m scope + +instance (ASTRep rep, CanBeWise (Op rep)) => ASTRep (Wise rep) where + expTypesFromPattern = + withoutWisdom . expTypesFromPattern . removePatternWisdom + +instance Pretty VarWisdom where + ppr _ = ppr () + +instance (PrettyRep rep, CanBeWise (Op rep)) => PrettyRep (Wise rep) where + ppExpDec (_, dec) = ppExpDec dec . removeExpWisdom + +instance AliasesOf (VarWisdom, dec) where + aliasesOf = unAliases . varWisdomAliases . fst + +instance (ASTRep rep, CanBeWise (Op rep)) => Aliased (Wise rep) where + bodyAliases = map unAliases . bodyWisdomAliases . fst . bodyDec + consumedInBody = unAliases . bodyWisdomConsumed . fst . bodyDec + +removeWisdom :: CanBeWise (Op rep) => Rephraser Identity (Wise rep) rep +removeWisdom = + Rephraser + { rephraseExpDec = return . snd, + rephraseLetBoundDec = return . snd, + rephraseBodyDec = return . snd, + rephraseFParamDec = return, + rephraseLParamDec = return, + rephraseRetType = return, + rephraseBranchType = return, + rephraseOp = return . removeOpWisdom + } + +removeScopeWisdom :: Scope (Wise rep) -> Scope rep +removeScopeWisdom = M.map unAlias + where + unAlias (LetName (_, dec)) = LetName dec + unAlias (FParamName dec) = FParamName dec + unAlias (LParamName dec) = LParamName dec + unAlias (IndexName it) = IndexName it + +addScopeWisdom :: Scope rep -> Scope (Wise rep) +addScopeWisdom = M.map alias + where + alias (LetName dec) = LetName (VarWisdom mempty, dec) + alias (FParamName dec) = FParamName dec + alias (LParamName dec) = LParamName dec + alias (IndexName it) = IndexName it + +removeFunDefWisdom :: CanBeWise (Op rep) => FunDef (Wise rep) -> FunDef rep +removeFunDefWisdom = runIdentity . rephraseFunDef removeWisdom + +removeStmWisdom :: CanBeWise (Op rep) => Stm (Wise rep) -> Stm rep +removeStmWisdom = runIdentity . rephraseStm removeWisdom + +removeLambdaWisdom :: CanBeWise (Op rep) => Lambda (Wise rep) -> Lambda rep +removeLambdaWisdom = runIdentity . rephraseLambda removeWisdom + +removeBodyWisdom :: CanBeWise (Op rep) => Body (Wise rep) -> Body rep +removeBodyWisdom = runIdentity . rephraseBody removeWisdom + +removeExpWisdom :: CanBeWise (Op rep) => Exp (Wise rep) -> Exp rep +removeExpWisdom = runIdentity . rephraseExp removeWisdom + +removePatternWisdom :: PatternT (VarWisdom, a) -> PatternT a +removePatternWisdom = runIdentity . rephrasePattern (return . snd) + +addWisdomToPattern :: + (ASTRep rep, CanBeWise (Op rep)) => + Pattern rep -> + Exp (Wise rep) -> + Pattern (Wise rep) +addWisdomToPattern pat e = + Pattern (map f ctx) (map f val) + where + (ctx, val) = Aliases.mkPatternAliases pat e + f pe = + let (als, dec) = patElemDec pe + in pe `setPatElemDec` (VarWisdom als, dec) + +mkWiseBody :: + (ASTRep rep, CanBeWise (Op rep)) => + BodyDec rep -> + Stms (Wise rep) -> + Result -> + Body (Wise rep) +mkWiseBody dec bnds res = + Body + ( BodyWisdom aliases consumed (AliasDec $ freeIn $ freeInStmsAndRes bnds res), + dec + ) + bnds + res + where + (aliases, consumed) = Aliases.mkBodyAliases bnds res + +mkWiseLetStm :: + (ASTRep rep, CanBeWise (Op rep)) => + Pattern rep -> + StmAux (ExpDec rep) -> + Exp (Wise rep) -> + Stm (Wise rep) +mkWiseLetStm pat (StmAux cs attrs dec) e = + let pat' = addWisdomToPattern pat e + in Let pat' (StmAux cs attrs $ mkWiseExpDec pat' dec e) e + +mkWiseExpDec :: + (ASTRep rep, CanBeWise (Op rep)) => + Pattern (Wise rep) -> + ExpDec rep -> + Exp (Wise rep) -> + ExpDec (Wise rep) +mkWiseExpDec pat expdec e = + ( ExpWisdom + (AliasDec $ consumedInExp e) + (AliasDec $ freeIn pat <> freeIn expdec <> freeIn e), + expdec + ) + +instance + ( Bindable rep, + CanBeWise (Op rep) + ) => + Bindable (Wise rep) + where + mkExpPat ctx val e = + addWisdomToPattern (mkExpPat ctx val $ removeExpWisdom e) e + + mkExpDec pat e = + mkWiseExpDec pat (mkExpDec (removePatternWisdom pat) $ removeExpWisdom e) e + + mkLetNames names e = do + env <- asksScope removeScopeWisdom + flip runReaderT env $ do + Let pat dec _ <- mkLetNames names $ removeExpWisdom e + return $ mkWiseLetStm pat dec e + + mkBody bnds res = + let Body bodyrep _ _ = mkBody (fmap removeStmWisdom bnds) res + in mkWiseBody bodyrep bnds res + +class + ( AliasedOp (OpWithWisdom op), + IsOp (OpWithWisdom op) + ) => + CanBeWise op + where + type OpWithWisdom op :: Data.Kind.Type + removeOpWisdom :: OpWithWisdom op -> op + +instance CanBeWise () where + type OpWithWisdom () = () + removeOpWisdom () = () 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 @@ -62,18 +62,18 @@ import Futhark.IR -- | The monad in which simplification rules are evaluated. -newtype RuleM lore a = RuleM (BinderT lore (StateT VNameSource Maybe) a) +newtype RuleM rep a = RuleM (BinderT rep (StateT VNameSource Maybe) a) deriving ( Functor, Applicative, Monad, MonadFreshNames, - HasScope lore, - LocalScope lore + HasScope rep, + LocalScope rep ) -instance (ASTLore lore, BinderOps lore) => MonadBinder (RuleM lore) where - type Lore (RuleM lore) = lore +instance (ASTRep rep, BinderOps rep) => MonadBinder (RuleM rep) where + type Rep (RuleM rep) = rep mkExpDecM pat e = RuleM $ mkExpDecM pat e mkBodyM bnds res = RuleM $ mkBodyM bnds res mkLetNamesM pat e = RuleM $ mkLetNamesM pat e @@ -84,141 +84,141 @@ -- | Execute a 'RuleM' action. If succesful, returns the result and a -- list of new bindings. simplify :: - Scope lore -> + Scope rep -> VNameSource -> - Rule lore -> - Maybe (Stms lore, VNameSource) + Rule rep -> + Maybe (Stms rep, VNameSource) simplify _ _ Skip = Nothing simplify scope src (Simplify (RuleM m)) = runStateT (runBinderT_ m scope) src -cannotSimplify :: RuleM lore a +cannotSimplify :: RuleM rep a cannotSimplify = RuleM $ lift $ lift Nothing -liftMaybe :: Maybe a -> RuleM lore a +liftMaybe :: Maybe a -> RuleM rep a liftMaybe Nothing = cannotSimplify liftMaybe (Just x) = return x -- | An efficient way of encoding whether a simplification rule should even be attempted. -data Rule lore +data Rule rep = -- | Give it a shot. - Simplify (RuleM lore ()) + Simplify (RuleM rep ()) | -- | Don't bother. Skip -type RuleGeneric lore a = a -> Stm lore -> Rule lore +type RuleGeneric rep a = a -> Stm rep -> Rule rep -type RuleBasicOp lore a = +type RuleBasicOp rep a = ( a -> - Pattern lore -> - StmAux (ExpDec lore) -> + Pattern rep -> + StmAux (ExpDec rep) -> BasicOp -> - Rule lore + Rule rep ) -type RuleIf lore a = +type RuleIf rep a = a -> - Pattern lore -> - StmAux (ExpDec lore) -> + Pattern rep -> + StmAux (ExpDec rep) -> ( SubExp, - BodyT lore, - BodyT lore, - IfDec (BranchType lore) + BodyT rep, + BodyT rep, + IfDec (BranchType rep) ) -> - Rule lore + Rule rep -type RuleDoLoop lore a = +type RuleDoLoop rep a = a -> - Pattern lore -> - StmAux (ExpDec lore) -> - ( [(FParam lore, SubExp)], - [(FParam lore, SubExp)], - LoopForm lore, - BodyT lore + Pattern rep -> + StmAux (ExpDec rep) -> + ( [(FParam rep, SubExp)], + [(FParam rep, SubExp)], + LoopForm rep, + BodyT rep ) -> - Rule lore + Rule rep -type RuleOp lore a = +type RuleOp rep a = a -> - Pattern lore -> - StmAux (ExpDec lore) -> - Op lore -> - Rule lore + Pattern rep -> + StmAux (ExpDec rep) -> + Op rep -> + Rule rep -- | A simplification rule takes some argument and a statement, and -- tries to simplify the statement. -data SimplificationRule lore a - = RuleGeneric (RuleGeneric lore a) - | RuleBasicOp (RuleBasicOp lore a) - | RuleIf (RuleIf lore a) - | RuleDoLoop (RuleDoLoop lore a) - | RuleOp (RuleOp lore a) +data SimplificationRule rep a + = RuleGeneric (RuleGeneric rep a) + | RuleBasicOp (RuleBasicOp rep a) + | RuleIf (RuleIf rep a) + | RuleDoLoop (RuleDoLoop rep a) + | RuleOp (RuleOp rep a) -- | A collection of rules grouped by which forms of statements they -- may apply to. -data Rules lore a = Rules - { rulesAny :: [SimplificationRule lore a], - rulesBasicOp :: [SimplificationRule lore a], - rulesIf :: [SimplificationRule lore a], - rulesDoLoop :: [SimplificationRule lore a], - rulesOp :: [SimplificationRule lore a] +data Rules rep a = Rules + { rulesAny :: [SimplificationRule rep a], + rulesBasicOp :: [SimplificationRule rep a], + rulesIf :: [SimplificationRule rep a], + rulesDoLoop :: [SimplificationRule rep a], + rulesOp :: [SimplificationRule rep a] } -instance Semigroup (Rules lore a) where +instance Semigroup (Rules rep a) where Rules as1 bs1 cs1 ds1 es1 <> Rules as2 bs2 cs2 ds2 es2 = Rules (as1 <> as2) (bs1 <> bs2) (cs1 <> cs2) (ds1 <> ds2) (es1 <> es2) -instance Monoid (Rules lore a) where +instance Monoid (Rules rep a) where mempty = Rules mempty mempty mempty mempty mempty -- | Context for a rule applied during top-down traversal of the -- program. Takes a symbol table as argument. -type TopDown lore = ST.SymbolTable lore +type TopDown rep = ST.SymbolTable rep -type TopDownRuleGeneric lore = RuleGeneric lore (TopDown lore) +type TopDownRuleGeneric rep = RuleGeneric rep (TopDown rep) -type TopDownRuleBasicOp lore = RuleBasicOp lore (TopDown lore) +type TopDownRuleBasicOp rep = RuleBasicOp rep (TopDown rep) -type TopDownRuleIf lore = RuleIf lore (TopDown lore) +type TopDownRuleIf rep = RuleIf rep (TopDown rep) -type TopDownRuleDoLoop lore = RuleDoLoop lore (TopDown lore) +type TopDownRuleDoLoop rep = RuleDoLoop rep (TopDown rep) -type TopDownRuleOp lore = RuleOp lore (TopDown lore) +type TopDownRuleOp rep = RuleOp rep (TopDown rep) -type TopDownRule lore = SimplificationRule lore (TopDown lore) +type TopDownRule rep = SimplificationRule rep (TopDown rep) -- | Context for a rule applied during bottom-up traversal of the -- program. Takes a symbol table and usage table as arguments. -type BottomUp lore = (ST.SymbolTable lore, UT.UsageTable) +type BottomUp rep = (ST.SymbolTable rep, UT.UsageTable) -type BottomUpRuleGeneric lore = RuleGeneric lore (BottomUp lore) +type BottomUpRuleGeneric rep = RuleGeneric rep (BottomUp rep) -type BottomUpRuleBasicOp lore = RuleBasicOp lore (BottomUp lore) +type BottomUpRuleBasicOp rep = RuleBasicOp rep (BottomUp rep) -type BottomUpRuleIf lore = RuleIf lore (BottomUp lore) +type BottomUpRuleIf rep = RuleIf rep (BottomUp rep) -type BottomUpRuleDoLoop lore = RuleDoLoop lore (BottomUp lore) +type BottomUpRuleDoLoop rep = RuleDoLoop rep (BottomUp rep) -type BottomUpRuleOp lore = RuleOp lore (BottomUp lore) +type BottomUpRuleOp rep = RuleOp rep (BottomUp rep) -type BottomUpRule lore = SimplificationRule lore (BottomUp lore) +type BottomUpRule rep = SimplificationRule rep (BottomUp rep) -- | A collection of top-down rules. -type TopDownRules lore = Rules lore (TopDown lore) +type TopDownRules rep = Rules rep (TopDown rep) -- | A collection of bottom-up rules. -type BottomUpRules lore = Rules lore (BottomUp lore) +type BottomUpRules rep = Rules rep (BottomUp rep) -- | A collection of both top-down and bottom-up rules. -data RuleBook lore = RuleBook - { bookTopDownRules :: TopDownRules lore, - bookBottomUpRules :: BottomUpRules lore +data RuleBook rep = RuleBook + { bookTopDownRules :: TopDownRules rep, + bookBottomUpRules :: BottomUpRules rep } -instance Semigroup (RuleBook lore) where +instance Semigroup (RuleBook rep) where RuleBook ts1 bs1 <> RuleBook ts2 bs2 = RuleBook (ts1 <> ts2) (bs1 <> bs2) -instance Monoid (RuleBook lore) where +instance Monoid (RuleBook rep) where mempty = RuleBook mempty mempty -- | Construct a rule book from a collection of rules. @@ -259,11 +259,11 @@ -- 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 lore m) => - RuleBook lore -> - ST.SymbolTable lore -> - Stm lore -> - m (Maybe (Stms lore)) + (MonadFreshNames m, HasScope rep m) => + RuleBook rep -> + ST.SymbolTable rep -> + Stm rep -> + m (Maybe (Stms rep)) topDownSimplifyStm = applyRules . bookTopDownRules -- | @simplifyStm uses bnd@ performs simplification of the binding @@ -272,14 +272,14 @@ -- original binding (and possibly more, for intermediate results). -- The first argument is the set of names used after this binding. bottomUpSimplifyStm :: - (MonadFreshNames m, HasScope lore m) => - RuleBook lore -> - (ST.SymbolTable lore, UT.UsageTable) -> - Stm lore -> - m (Maybe (Stms lore)) + (MonadFreshNames m, HasScope rep m) => + RuleBook rep -> + (ST.SymbolTable rep, UT.UsageTable) -> + Stm rep -> + m (Maybe (Stms rep)) bottomUpSimplifyStm = applyRules . bookBottomUpRules -rulesForStm :: Stm lore -> Rules lore a -> [SimplificationRule lore a] +rulesForStm :: Stm rep -> Rules rep a -> [SimplificationRule rep a] rulesForStm stm = case stmExp stm of BasicOp {} -> rulesBasicOp DoLoop {} -> rulesDoLoop @@ -287,7 +287,7 @@ If {} -> rulesIf _ -> rulesAny -applyRule :: SimplificationRule lore a -> a -> Stm lore -> Rule lore +applyRule :: SimplificationRule rep a -> a -> Stm rep -> Rule rep applyRule (RuleGeneric f) a stm = f a stm applyRule (RuleBasicOp f) a (Let pat aux (BasicOp e)) = f a pat aux e applyRule (RuleDoLoop f) a (Let pat aux (DoLoop ctx val form body)) = @@ -300,11 +300,11 @@ Skip applyRules :: - (MonadFreshNames m, HasScope lore m) => - Rules lore a -> + (MonadFreshNames m, HasScope rep m) => + Rules rep a -> a -> - Stm lore -> - m (Maybe (Stms lore)) + Stm rep -> + m (Maybe (Stms rep)) applyRules all_rules context stm = do scope <- askScope 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 @@ -37,7 +37,7 @@ import Futhark.Optimise.Simplify.Rules.Loop import Futhark.Util -topDownRules :: BinderOps lore => [TopDownRule lore] +topDownRules :: BinderOps rep => [TopDownRule rep] topDownRules = [ RuleGeneric constantFoldPrimFun, RuleIf ruleIf, @@ -45,7 +45,7 @@ RuleGeneric withAccTopDown ] -bottomUpRules :: BinderOps lore => [BottomUpRule lore] +bottomUpRules :: BinderOps rep => [BottomUpRule rep] bottomUpRules = [ RuleIf removeDeadBranchResult, RuleGeneric withAccBottomUp, @@ -55,14 +55,14 @@ -- | A set of standard simplification rules. These assume pure -- functional semantics, and so probably should not be applied after -- memory block merging. -standardRules :: (BinderOps lore, Aliased lore) => RuleBook lore +standardRules :: (BinderOps rep, Aliased rep) => RuleBook rep standardRules = ruleBook topDownRules bottomUpRules <> loopRules <> basicOpRules -- | Turn @copy(x)@ into @x@ iff @x@ is not used after this copy -- statement and it can be consumed. -- -- This simplistic rule is only valid before we introduce memory. -removeUnnecessaryCopy :: (BinderOps lore, Aliased lore) => BottomUpRuleBasicOp lore +removeUnnecessaryCopy :: (BinderOps rep, Aliased rep) => BottomUpRuleBasicOp rep removeUnnecessaryCopy (vtable, used) (Pattern [] [d]) _ (Copy v) | not (v `UT.isConsumed` used), (not (v `UT.used` used) && consumable) || not (patElemName d `UT.isConsumed` used) = @@ -84,7 +84,7 @@ pure True removeUnnecessaryCopy _ _ _ _ = Skip -constantFoldPrimFun :: BinderOps lore => TopDownRuleGeneric lore +constantFoldPrimFun :: BinderOps rep => TopDownRuleGeneric rep constantFoldPrimFun _ (Let pat (StmAux cs attrs _) (Apply fname args _ _)) | Just args' <- mapM (isConst . fst) args, Just (_, _, fun) <- M.lookup (nameToString fname) primFuns, @@ -98,7 +98,7 @@ isConst _ = Nothing constantFoldPrimFun _ _ = Skip -simplifyIndex :: BinderOps lore => BottomUpRuleBasicOp lore +simplifyIndex :: BinderOps rep => BottomUpRuleBasicOp rep simplifyIndex (vtable, used) pat@(Pattern [] [pe]) (StmAux cs attrs _) (Index idd inds) | Just m <- simplifyIndexing vtable seType idd inds consumed = Simplify $ do res <- m @@ -115,7 +115,7 @@ seType (Constant v) = Just $ Prim $ primValueType v simplifyIndex _ _ _ _ = Skip -ruleIf :: BinderOps lore => TopDownRuleIf lore +ruleIf :: BinderOps rep => TopDownRuleIf rep ruleIf _ pat _ (e1, tb, fb, IfDec _ ifsort) | Just branch <- checkBranch, ifsort /= IfFallback || isCt1 e1 = Simplify $ do @@ -194,7 +194,7 @@ -- | Move out results of a conditional expression whose computation is -- either invariant to the branches (only done for results in the -- context), or the same in both branches. -hoistBranchInvariant :: BinderOps lore => TopDownRuleIf lore +hoistBranchInvariant :: BinderOps rep => TopDownRuleIf rep hoistBranchInvariant _ pat _ (cond, tb, fb, IfDec ret ifsort) = Simplify $ do let tses = bodyResult tb fses = bodyResult fb @@ -278,7 +278,7 @@ -- after a branch. Standard dead code removal can remove the branch -- if *none* of the return values are used, but this rule is more -- precise. -removeDeadBranchResult :: BinderOps lore => BottomUpRuleIf lore +removeDeadBranchResult :: BinderOps rep => BottomUpRuleIf rep removeDeadBranchResult (_, used) pat _ (e1, tb, fb, IfDec rettype ifsort) | -- Only if there is no existential context... patternSize pat == length rettype, @@ -300,7 +300,7 @@ in Simplify $ letBind (Pattern [] pat') $ If e1 tb' fb' $ IfDec rettype' ifsort | otherwise = Skip -withAccTopDown :: BinderOps lore => TopDownRuleGeneric lore +withAccTopDown :: BinderOps rep => TopDownRuleGeneric rep -- A WithAcc with no accumulators is sent to Valhalla. withAccTopDown _ (Let pat aux (WithAcc [] lam)) = Simplify . auxing aux $ do lam_res <- bodyBind $ lambdaBody lam @@ -359,7 +359,7 @@ pure $ Just x withAccTopDown _ _ = Skip -withAccBottomUp :: BinderOps lore => BottomUpRuleGeneric lore +withAccBottomUp :: BinderOps rep => BottomUpRuleGeneric rep -- Eliminate dead results. withAccBottomUp (_, utable) (Let pat aux (WithAcc inputs lam)) | not $ all (`UT.used` utable) $ patternNames pat = Simplify $ do 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 @@ -35,7 +35,7 @@ | ArgReplicate [SubExp] SubExp | ArgVar VName -toConcatArg :: ST.SymbolTable lore -> VName -> (ConcatArg, Certificates) +toConcatArg :: ST.SymbolTable rep -> VName -> (ConcatArg, Certificates) toConcatArg vtable v = case ST.lookupBasicOp v vtable of Just (ArrayLit ses _, cs) -> @@ -79,7 +79,7 @@ fuseConcatArg xs y = y : xs -simplifyConcat :: BinderOps lore => BottomUpRuleBasicOp lore +simplifyConcat :: BinderOps rep => BottomUpRuleBasicOp rep -- concat@1(transpose(x),transpose(y)) == transpose(concat@0(x,y)) simplifyConcat (vtable, _) pat _ (Concat i x xs new_d) | Just r <- arrayRank <$> ST.lookupType x vtable, @@ -142,7 +142,7 @@ forSingleArray ys = ys simplifyConcat _ _ _ _ = Skip -ruleBasicOp :: BinderOps lore => TopDownRuleBasicOp lore +ruleBasicOp :: BinderOps rep => TopDownRuleBasicOp rep ruleBasicOp vtable pat aux op | Just (op', cs) <- applySimpleRules defOf seType op = Simplify $ certifying (cs <> stmAuxCerts aux) $ letBind pat $ BasicOp op' @@ -375,17 +375,17 @@ ruleBasicOp _ _ _ _ = Skip -topDownRules :: BinderOps lore => [TopDownRule lore] +topDownRules :: BinderOps rep => [TopDownRule rep] topDownRules = [ RuleBasicOp ruleBasicOp ] -bottomUpRules :: BinderOps lore => [BottomUpRule lore] +bottomUpRules :: BinderOps rep => [BottomUpRule rep] bottomUpRules = [ RuleBasicOp simplifyConcat ] -- | A set of simplification rules for 'BasicOp's. Includes rules -- from "Futhark.Optimise.Simplify.Rules.Simple". -basicOpRules :: (BinderOps lore, Aliased lore) => RuleBook lore +basicOpRules :: (BinderOps rep, Aliased rep) => RuleBook rep basicOpRules = ruleBook topDownRules bottomUpRules <> loopRules diff --git a/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs b/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs --- a/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs +++ b/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs @@ -39,13 +39,13 @@ -- | @foldClosedForm look foldfun accargs arrargs@ determines whether -- each of the results of @foldfun@ can be expressed in a closed form. foldClosedForm :: - (ASTLore lore, BinderOps lore) => - VarLookup lore -> - Pattern lore -> - Lambda lore -> + (ASTRep rep, BinderOps rep) => + VarLookup rep -> + Pattern rep -> + Lambda rep -> [SubExp] -> [VName] -> - RuleM lore () + RuleM rep () foldClosedForm look pat lam accs arrs = do inputsize <- arraysSize 0 <$> mapM lookupType arrs @@ -78,14 +78,14 @@ -- | @loopClosedForm pat respat merge bound bodys@ determines whether -- the do-loop can be expressed in a closed form. loopClosedForm :: - (ASTLore lore, BinderOps lore) => - Pattern lore -> - [(FParam lore, SubExp)] -> + (ASTRep rep, BinderOps rep) => + Pattern rep -> + [(FParam rep, SubExp)] -> Names -> IntType -> SubExp -> - Body lore -> - RuleM lore () + Body rep -> + RuleM rep () loopClosedForm pat merge i it bound body = do t <- case patternTypes pat of [Prim t] -> return t @@ -118,7 +118,7 @@ knownBnds = M.fromList $ zip mergenames mergeexp checkResults :: - BinderOps lore => + BinderOps rep => [VName] -> SubExp -> Names -> @@ -126,9 +126,9 @@ M.Map VName SubExp -> -- | Lambda-bound [VName] -> - Body lore -> + Body rep -> [SubExp] -> - RuleM lore (Body lore) + RuleM rep (Body rep) checkResults pat size untouchable it knownBnds params body accs = do ((), bnds) <- collectStms $ @@ -191,8 +191,8 @@ BasicOp $ ConvOp (SIToFP it t) size determineKnownBindings :: - VarLookup lore -> - Lambda lore -> + VarLookup rep -> + Lambda rep -> [SubExp] -> [VName] -> M.Map VName SubExp @@ -215,7 +215,7 @@ Just (p, ve) isReplicate _ = Nothing -makeBindMap :: Body lore -> M.Map VName (Exp lore) +makeBindMap :: Body rep -> M.Map VName (Exp rep) makeBindMap = M.fromList . mapMaybe isSingletonStm . stmsToList . bodyStms where isSingletonStm (Let pat _ e) = case patternNames pat of diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs --- a/src/Futhark/Optimise/Simplify/Rules/Index.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Index.hs @@ -33,7 +33,7 @@ -- | Try to simplify an index operation. simplifyIndexing :: MonadBinder m => - ST.SymbolTable (Lore m) -> + ST.SymbolTable (Rep m) -> TypeLookup -> VName -> Slice SubExp -> diff --git a/src/Futhark/Optimise/Simplify/Rules/Loop.hs b/src/Futhark/Optimise/Simplify/Rules/Loop.hs --- a/src/Futhark/Optimise/Simplify/Rules/Loop.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Loop.hs @@ -25,7 +25,7 @@ -- I do not claim that the current implementation of this rule is -- perfect, but it should suffice for many cases, and should never -- generate wrong code. -removeRedundantMergeVariables :: BinderOps lore => BottomUpRuleDoLoop lore +removeRedundantMergeVariables :: BinderOps rep => BottomUpRuleDoLoop rep removeRedundantMergeVariables (_, used) pat aux (ctx, val, form, body) | not $ all (usedAfterLoop . fst) val, null ctx -- FIXME: things get tricky if we can remove all vals @@ -105,7 +105,7 @@ -- We may change the type of the loop if we hoist out a shape -- annotation, in which case we also need to tweak the bound pattern. -hoistLoopInvariantMergeVariables :: BinderOps lore => TopDownRuleDoLoop lore +hoistLoopInvariantMergeVariables :: BinderOps rep => TopDownRuleDoLoop rep hoistLoopInvariantMergeVariables vtable pat aux (ctx, val, form, loopbody) = -- Figure out which of the elements of loopresult are -- loop-invariant, and hoist them out. @@ -206,12 +206,12 @@ not (name `nameIn` namesOfMergeParams) || name `nameIn` namesOfInvariant -simplifyClosedFormLoop :: BinderOps lore => TopDownRuleDoLoop lore +simplifyClosedFormLoop :: BinderOps rep => TopDownRuleDoLoop rep simplifyClosedFormLoop _ pat _ ([], val, ForLoop i it bound [], body) = Simplify $ loopClosedForm pat val (oneName i) it bound body simplifyClosedFormLoop _ _ _ _ = Skip -simplifyLoopVariables :: (BinderOps lore, Aliased lore) => TopDownRuleDoLoop lore +simplifyLoopVariables :: (BinderOps rep, Aliased rep) => TopDownRuleDoLoop rep simplifyLoopVariables vtable pat aux (ctx, val, form@(ForLoop i it num_iters loop_vars), body) | simplifiable <- map checkIfSimplifiable loop_vars, not $ all isNothing simplifiable = Simplify $ do @@ -290,7 +290,7 @@ -- instead. We then move the sign extension inside the loop instead. -- This addresses loops of the form @for i in x.. TopDownRuleDoLoop lore +narrowLoopType :: (BinderOps rep) => TopDownRuleDoLoop rep narrowLoopType vtable pat aux (ctx, val, ForLoop i Int64 n [], body) | Just (n', it', cs) <- smallerType = Simplify $ do @@ -315,13 +315,13 @@ narrowLoopType _ _ _ _ = Skip unroll :: - BinderOps lore => + BinderOps rep => Integer -> - [(FParam lore, SubExp)] -> + [(FParam rep, SubExp)] -> (VName, IntType, Integer) -> - [(LParam lore, VName)] -> - Body lore -> - RuleM lore [SubExp] + [(LParam rep, VName)] -> + Body rep -> + RuleM rep [SubExp] unroll n merge (iv, it, i) loop_vars body | i >= n = return $ map snd merge @@ -348,7 +348,7 @@ let merge' = zip (map fst merge) $ bodyResult iter_body' unroll n merge' (iv, it, i + 1) loop_vars body -simplifyKnownIterationLoop :: BinderOps lore => TopDownRuleDoLoop lore +simplifyKnownIterationLoop :: BinderOps rep => TopDownRuleDoLoop rep simplifyKnownIterationLoop _ pat aux (ctx, val, ForLoop i it (Constant iters) loop_vars, body) | IntValue n <- iters, zeroIshInt n || oneIshInt n || "unroll" `inAttrs` stmAuxAttrs aux = Simplify $ do @@ -358,7 +358,7 @@ simplifyKnownIterationLoop _ _ _ _ = Skip -topDownRules :: (BinderOps lore, Aliased lore) => [TopDownRule lore] +topDownRules :: (BinderOps rep, Aliased rep) => [TopDownRule rep] topDownRules = [ RuleDoLoop hoistLoopInvariantMergeVariables, RuleDoLoop simplifyClosedFormLoop, @@ -367,11 +367,11 @@ RuleDoLoop narrowLoopType ] -bottomUpRules :: BinderOps lore => [BottomUpRule lore] +bottomUpRules :: BinderOps rep => [BottomUpRule rep] bottomUpRules = [ RuleDoLoop removeRedundantMergeVariables ] -- | Standard loop simplification rules. -loopRules :: (BinderOps lore, Aliased lore) => RuleBook lore +loopRules :: (BinderOps rep, Aliased rep) => RuleBook rep loopRules = ruleBook topDownRules bottomUpRules diff --git a/src/Futhark/Optimise/Simplify/Rules/Simple.hs b/src/Futhark/Optimise/Simplify/Rules/Simple.hs --- a/src/Futhark/Optimise/Simplify/Rules/Simple.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Simple.hs @@ -14,14 +14,14 @@ import Futhark.IR -- | A function that, given a variable name, returns its definition. -type VarLookup lore = VName -> Maybe (Exp lore, Certificates) +type VarLookup rep = VName -> Maybe (Exp rep, Certificates) -- | A function that, given a subexpression, returns its type. type TypeLookup = SubExp -> Maybe Type -- | A simple rule is a top-down rule that can be expressed as a pure -- function. -type SimpleRule lore = VarLookup lore -> TypeLookup -> BasicOp -> Maybe (BasicOp, Certificates) +type SimpleRule rep = VarLookup rep -> TypeLookup -> BasicOp -> Maybe (BasicOp, Certificates) isCt1 :: SubExp -> Bool isCt1 (Constant v) = oneIsh v @@ -31,7 +31,7 @@ isCt0 (Constant v) = zeroIsh v isCt0 _ = False -simplifyCmpOp :: SimpleRule lore +simplifyCmpOp :: SimpleRule rep simplifyCmpOp _ _ (CmpOp cmp e1 e2) | e1 == e2 = constRes $ BoolValue $ @@ -55,7 +55,7 @@ _ -> Just (SubExp (Constant (BoolValue False)), cs) simplifyCmpOp _ _ _ = Nothing -simplifyBinOp :: SimpleRule lore +simplifyBinOp :: SimpleRule rep simplifyBinOp _ _ (BinOp op (Constant v1) (Constant v2)) | Just res <- doBinOp op v1 v2 = constRes res @@ -204,7 +204,7 @@ subExpRes :: SubExp -> Maybe (BasicOp, Certificates) subExpRes = Just . (,mempty) . SubExp -simplifyUnOp :: SimpleRule lore +simplifyUnOp :: SimpleRule rep simplifyUnOp _ _ (UnOp op (Constant v)) = constRes =<< doUnOp op v simplifyUnOp defOf _ (UnOp Not (Var v)) @@ -213,7 +213,7 @@ simplifyUnOp _ _ _ = Nothing -simplifyConvOp :: SimpleRule lore +simplifyConvOp :: SimpleRule rep simplifyConvOp _ _ (ConvOp op (Constant v)) = constRes =<< doConvOp op v simplifyConvOp _ _ (ConvOp op se) @@ -244,13 +244,13 @@ Nothing -- If expression is true then just replace assertion. -simplifyAssert :: SimpleRule lore +simplifyAssert :: SimpleRule rep simplifyAssert _ _ (Assert (Constant (BoolValue True)) _ _) = constRes UnitValue simplifyAssert _ _ _ = Nothing -simplifyIdentityReshape :: SimpleRule lore +simplifyIdentityReshape :: SimpleRule rep simplifyIdentityReshape _ seType (Reshape newshape v) | Just t <- seType $ Var v, newDims newshape == arrayDims t -- No-op reshape. @@ -258,19 +258,19 @@ subExpRes $ Var v simplifyIdentityReshape _ _ _ = Nothing -simplifyReshapeReshape :: SimpleRule lore +simplifyReshapeReshape :: SimpleRule rep simplifyReshapeReshape defOf _ (Reshape newshape v) | Just (BasicOp (Reshape oldshape v2), v_cs) <- defOf v = Just (Reshape (fuseReshape oldshape newshape) v2, v_cs) simplifyReshapeReshape _ _ _ = Nothing -simplifyReshapeScratch :: SimpleRule lore +simplifyReshapeScratch :: SimpleRule rep simplifyReshapeScratch defOf _ (Reshape newshape v) | Just (BasicOp (Scratch bt _), v_cs) <- defOf v = Just (Scratch bt $ newDims newshape, v_cs) simplifyReshapeScratch _ _ _ = Nothing -simplifyReshapeReplicate :: SimpleRule lore +simplifyReshapeReplicate :: SimpleRule rep simplifyReshapeReplicate defOf seType (Reshape newshape v) | Just (BasicOp (Replicate _ se), v_cs) <- defOf v, Just oldshape <- arrayShape <$> seType se, @@ -281,7 +281,7 @@ in Just (Replicate (Shape new) se, v_cs) simplifyReshapeReplicate _ _ _ = Nothing -simplifyReshapeIota :: SimpleRule lore +simplifyReshapeIota :: SimpleRule rep simplifyReshapeIota defOf _ (Reshape newshape v) | Just (BasicOp (Iota _ offset stride it), v_cs) <- defOf v, [n] <- newDims newshape = @@ -297,7 +297,7 @@ -- If we are size-coercing a slice, then we might as well just use a -- different slice instead. -simplifyReshapeIndex :: SimpleRule lore +simplifyReshapeIndex :: SimpleRule rep simplifyReshapeIndex defOf _ (Reshape newshape v) | Just ds <- shapeCoercion newshape, Just (BasicOp (Index v' slice), v_cs) <- defOf v, @@ -308,7 +308,7 @@ -- If we are updating a slice with the result of a size coercion, we -- instead use the original array and update the slice dimensions. -simplifyUpdateReshape :: SimpleRule lore +simplifyUpdateReshape :: SimpleRule rep simplifyUpdateReshape defOf seType (Update dest slice (Var v)) | Just (BasicOp (Reshape newshape v'), v_cs) <- defOf v, Just _ <- shapeCoercion newshape, @@ -318,7 +318,7 @@ Just (Update dest slice' $ Var v', v_cs) simplifyUpdateReshape _ _ _ = Nothing -improveReshape :: SimpleRule lore +improveReshape :: SimpleRule rep improveReshape _ seType (Reshape newshape v) | Just t <- seType $ Var v, newshape' <- informReshape (arrayDims t) newshape, @@ -328,7 +328,7 @@ -- | If we are copying a scratch array (possibly indirectly), just turn it into a scratch by -- itself. -copyScratchToScratch :: SimpleRule lore +copyScratchToScratch :: SimpleRule rep copyScratchToScratch defOf seType (Copy src) = do t <- seType $ Var src if isActuallyScratch src @@ -344,7 +344,7 @@ copyScratchToScratch _ _ _ = Nothing -simpleRules :: [SimpleRule lore] +simpleRules :: [SimpleRule rep] simpleRules = [ simplifyBinOp, simplifyCmpOp, @@ -366,7 +366,7 @@ -- and certificates that it must depend on. {-# NOINLINE applySimpleRules #-} applySimpleRules :: - VarLookup lore -> + VarLookup rep -> TypeLookup -> BasicOp -> Maybe (BasicOp, Certificates) diff --git a/src/Futhark/Optimise/Sink.hs b/src/Futhark/Optimise/Sink.hs --- a/src/Futhark/Optimise/Sink.hs +++ b/src/Futhark/Optimise/Sink.hs @@ -40,10 +40,10 @@ -- you ever see this pass as being a compilation speed bottleneck, -- start by caching that a bit. -- --- This pass is defined on the Kernels representation. This is not --- because we do anything kernel-specific here, but simply because --- more explicit indexing is going on after SOACs are gone. -module Futhark.Optimise.Sink (sinkKernels, sinkMC) where +-- This pass is defined on post-SOACS representations. This is not +-- because we do anything GPU-specific here, but simply because more +-- explicit indexing is going on after SOACs are gone. +module Futhark.Optimise.Sink (sinkGPU, sinkMC) where import Control.Monad.State import Data.Bifunctor @@ -52,28 +52,28 @@ import qualified Futhark.Analysis.Alias as Alias import qualified Futhark.Analysis.SymbolTable as ST import Futhark.IR.Aliases -import Futhark.IR.Kernels +import Futhark.IR.GPU import Futhark.IR.MC import Futhark.Pass -type SymbolTable lore = ST.SymbolTable lore +type SymbolTable rep = ST.SymbolTable rep -type Sinking lore = M.Map VName (Stm lore) +type Sinking rep = M.Map VName (Stm rep) type Sunk = Names -type Sinker lore a = SymbolTable lore -> Sinking lore -> a -> (a, Sunk) +type Sinker rep a = SymbolTable rep -> Sinking rep -> a -> (a, Sunk) -type Constraints lore = - ( ASTLore lore, - Aliased lore, - ST.IndexOp (Op lore) +type Constraints rep = + ( ASTRep rep, + Aliased rep, + ST.IndexOp (Op rep) ) -- | Given a statement, compute how often each of its free variables -- are used. Not accurate: what we care about are only 1, and greater -- than 1. -multiplicity :: Constraints lore => Stm lore -> M.Map VName Int +multiplicity :: Constraints rep => Stm rep -> M.Map VName Int multiplicity stm = case stmExp stm of If cond tbranch fbranch _ -> @@ -86,9 +86,9 @@ comb = M.unionWith (+) optimiseBranch :: - Constraints lore => - Sinker lore (Op lore) -> - Sinker lore (Body lore) + Constraints rep => + Sinker rep (Op rep) -> + Sinker rep (Body rep) optimiseBranch onOp vtable sinking (Body dec stms res) = let (stms', stms_sunk) = optimiseStms onOp vtable sinking' (sunk_stms <> stms) $ freeIn res in ( Body dec stms' res, @@ -104,13 +104,13 @@ sunk = namesFromList $ foldMap (patternNames . stmPattern) sunk_stms optimiseStms :: - Constraints lore => - Sinker lore (Op lore) -> - SymbolTable lore -> - Sinking lore -> - Stms lore -> + Constraints rep => + Sinker rep (Op rep) -> + SymbolTable rep -> + Sinking rep -> + Stms rep -> Names -> - (Stms lore, Sunk) + (Stms rep, Sunk) optimiseStms onOp init_vtable init_sinking all_stms free_in_res = let (all_stms', sunk) = optimiseStms' init_vtable init_sinking $ stmsToList all_stms @@ -168,25 +168,25 @@ } optimiseBody :: - Constraints lore => - Sinker lore (Op lore) -> - Sinker lore (Body lore) + Constraints rep => + Sinker rep (Op rep) -> + Sinker rep (Body rep) optimiseBody onOp vtable sinking (Body attr stms res) = let (stms', sunk) = optimiseStms onOp vtable sinking stms $ freeIn res in (Body attr stms' res, sunk) optimiseKernelBody :: - Constraints lore => - Sinker lore (Op lore) -> - Sinker lore (KernelBody lore) + Constraints rep => + Sinker rep (Op rep) -> + Sinker rep (KernelBody rep) optimiseKernelBody onOp vtable sinking (KernelBody attr stms res) = let (stms', sunk) = optimiseStms onOp vtable sinking stms $ freeIn res in (KernelBody attr stms' res, sunk) optimiseSegOp :: - Constraints lore => - Sinker lore (Op lore) -> - Sinker lore (SegOp lvl lore) + Constraints rep => + Sinker rep (Op rep) -> + Sinker rep (SegOp lvl rep) optimiseSegOp onOp vtable sinking op = let scope = scopeOfSegSpace $ segSpace op in runState (mapSegOpM (opMapper scope) op) mempty @@ -208,15 +208,15 @@ where op_vtable = ST.fromScope scope <> vtable -type SinkLore lore = Aliases lore +type SinkRep rep = Aliases rep sink :: - ( ASTLore lore, - CanBeAliased (Op lore), - ST.IndexOp (OpWithAliases (Op lore)) + ( ASTRep rep, + CanBeAliased (Op rep), + ST.IndexOp (OpWithAliases (Op rep)) ) => - Sinker (SinkLore lore) (Op (SinkLore lore)) -> - Pass lore lore + Sinker (SinkRep rep) (Op (SinkRep rep)) -> + Pass rep rep sink onOp = Pass "sink" "move memory loads closer to their uses" $ fmap removeProgAliases @@ -235,10 +235,10 @@ namesFromList $ M.keys $ scopeOf consts -- | Sinking in GPU kernels. -sinkKernels :: Pass Kernels Kernels -sinkKernels = sink onHostOp +sinkGPU :: Pass GPU GPU +sinkGPU = sink onHostOp where - onHostOp :: Sinker (SinkLore Kernels) (Op (SinkLore Kernels)) + onHostOp :: Sinker (SinkRep GPU) (Op (SinkRep GPU)) onHostOp vtable sinking (SegOp op) = first SegOp $ optimiseSegOp onHostOp vtable sinking op onHostOp _ _ op = (op, mempty) @@ -247,7 +247,7 @@ sinkMC :: Pass MC MC sinkMC = sink onHostOp where - onHostOp :: Sinker (SinkLore MC) (Op (SinkLore MC)) + onHostOp :: Sinker (SinkRep MC) (Op (SinkRep MC)) onHostOp vtable sinking (ParOp par_op op) = let (par_op', par_sunk) = maybe diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs --- a/src/Futhark/Optimise/TileLoops.hs +++ b/src/Futhark/Optimise/TileLoops.hs @@ -12,7 +12,7 @@ import Data.Maybe (mapMaybe) import qualified Data.Sequence as Seq import qualified Futhark.Analysis.Alias as Alias -import Futhark.IR.Kernels +import Futhark.IR.GPU import Futhark.IR.Prop.Aliases (consumedInStm) import Futhark.MonadFreshNames import Futhark.Optimise.BlkRegTiling @@ -23,7 +23,7 @@ import Prelude hiding (quot) -- | The pass definition. -tileLoops :: Pass Kernels Kernels +tileLoops :: Pass GPU GPU tileLoops = Pass "tile loops" "Tile stream loops inside kernels" $ intraproceduralTransformation onStms @@ -33,16 +33,16 @@ runState $ runReaderT (optimiseStms stms) scope -optimiseBody :: Body Kernels -> TileM (Body Kernels) +optimiseBody :: Body GPU -> TileM (Body GPU) optimiseBody (Body () stms res) = Body () <$> optimiseStms stms <*> pure res -optimiseStms :: Stms Kernels -> TileM (Stms Kernels) +optimiseStms :: Stms GPU -> TileM (Stms GPU) optimiseStms stms = localScope (scopeOf stms) $ mconcat <$> mapM optimiseStm (stmsToList stms) -optimiseStm :: Stm Kernels -> TileM (Stms Kernels) +optimiseStm :: Stm GPU -> TileM (Stms GPU) optimiseStm stm@(Let pat aux (Op (SegOp (SegMap lvl@SegThread {} space ts kbody)))) = do res3dtiling <- doRegTiling3D stm case res3dtiling of @@ -51,7 +51,7 @@ blkRegTiling_res <- mmBlkRegTiling stm case blkRegTiling_res of Just (extra_bnds, stmt') -> return (extra_bnds <> oneStm stmt') - Nothing -> do + Nothing -> localScope (scopeOfSegSpace space) $ do (host_stms, (lvl', space', kbody')) <- tileInKernelBody mempty initial_variance lvl space ts kbody return $ host_stms <> oneStm (Let pat aux $ Op $ SegOp $ SegMap lvl' space' ts kbody') where @@ -67,8 +67,8 @@ SegLevel -> SegSpace -> [Type] -> - KernelBody Kernels -> - TileM (Stms Kernels, (SegLevel, SegSpace, KernelBody Kernels)) + KernelBody GPU -> + TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU)) tileInKernelBody branch_variant initial_variance lvl initial_kspace ts kbody | Just kbody_res <- mapM isSimpleResult $ kernelBodyResult kbody = do maybe_tiled <- @@ -99,8 +99,8 @@ SegLevel -> SegSpace -> [Type] -> - Body Kernels -> - TileM (Maybe (Stms Kernels, Tiling, TiledBody)) + Body GPU -> + TileM (Maybe (Stms GPU, Tiling, TiledBody)) tileInBody branch_variant initial_variance initial_lvl initial_space res_ts (Body () initial_kstms stms_res) = descend mempty $ stmsToList initial_kstms where @@ -207,10 +207,10 @@ -- the tiled statement anyway. preludeToPostlude :: VarianceTable -> - Stms Kernels -> - Stm Kernels -> - Stms Kernels -> - (Stms Kernels, Stms Kernels) + Stms GPU -> + Stm GPU -> + Stms GPU -> + (Stms GPU, Stms GPU) preludeToPostlude variance prelude stm_to_tile postlude = (prelude_used, prelude_not_used <> postlude) where @@ -247,10 +247,10 @@ -- in memory). partitionPrelude :: VarianceTable -> - Stms Kernels -> + Stms GPU -> Names -> Names -> - (Stms Kernels, Stms Kernels, Stms Kernels) + (Stms GPU, Stms GPU, Stms GPU) partitionPrelude variance prestms private used_after = (invariant_prestms, precomputed_variant_prestms, recomputed_variant_prestms) where @@ -301,10 +301,10 @@ injectPrelude :: SegSpace -> VarianceTable -> - Stms Kernels -> + Stms GPU -> Names -> - (Stms Kernels, Tiling, TiledBody) -> - (Stms Kernels, Tiling, TiledBody) + (Stms GPU, Tiling, TiledBody) -> + (Stms GPU, Tiling, TiledBody) injectPrelude initial_space variance prestms used (host_stms, tiling, tiledBody) = (host_stms, tiling, tiledBody') where @@ -338,19 +338,19 @@ tileDoLoop :: SegSpace -> VarianceTable -> - Stms Kernels -> + Stms GPU -> Names -> - (Stms Kernels, Tiling, TiledBody) -> + (Stms GPU, Tiling, TiledBody) -> [Type] -> - Pattern Kernels -> - StmAux (ExpDec Kernels) -> - [(FParam Kernels, SubExp)] -> + Pattern GPU -> + StmAux (ExpDec GPU) -> + [(FParam GPU, SubExp)] -> VName -> IntType -> SubExp -> - Stms Kernels -> + Stms GPU -> Result -> - TileM (Stms Kernels, Tiling, TiledBody) + TileM (Stms GPU, Tiling, TiledBody) tileDoLoop initial_space variance prestms used_in_body (host_stms, tiling, tiledBody) res_ts pat aux merge i it bound poststms poststms_res = do let prestms_used = used_in_body <> freeIn poststms <> freeIn poststms_res ( invariant_prestms, @@ -415,7 +415,7 @@ PrivStms mempty indexMergeParams <> privstms <> inloop_privstms loopbody' <- - runBodyBinder $ + localScope (scopeOfFParams mergeparams') . runBodyBinder $ resultBody . map Var <$> tiledBody private' privstms' accs' <- @@ -432,34 +432,29 @@ filter (`notElem` unSegSpace (tilingSpace tiling)) $ unSegSpace initial_space -doPrelude :: Tiling -> PrivStms -> Stms Kernels -> [VName] -> Binder Kernels [VName] +doPrelude :: Tiling -> PrivStms -> Stms GPU -> [VName] -> Binder GPU [VName] doPrelude tiling privstms prestms prestms_live = -- Create a SegMap that takes care of the prelude for every thread. tilingSegMap tiling "prelude" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> do ts <- mapM lookupType prestms_live fmap (map Var) $ - letTupExp "pre" - =<< eIf - (toExp in_bounds) - ( do - addPrivStms slice privstms - addStms prestms - resultBodyM $ map Var prestms_live - ) - (eBody $ map eBlank ts) + protectOutOfBounds "pre" in_bounds ts $ do + addPrivStms slice privstms + addStms prestms + pure $ map Var prestms_live -liveSet :: FreeIn a => Stms Kernels -> a -> Names +liveSet :: FreeIn a => Stms GPU -> a -> Names liveSet stms after = namesFromList (concatMap (patternNames . stmPattern) stms) `namesIntersection` freeIn after tileable :: - Stm Kernels -> + Stm GPU -> Maybe ( SubExp, [VName], - (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) + (Commutativity, Lambda GPU, [SubExp], Lambda GPU) ) tileable stm | Op (OtherOp (Screma w arrs form)) <- stmExp stm, @@ -521,12 +516,12 @@ -- SegMaps. This is for things that cannot efficiently be computed -- once in advance in the prelude SegMap, primarily (exclusively?) -- array slicing operations. -data PrivStms = PrivStms (Stms Kernels) ReadPrelude +data PrivStms = PrivStms (Stms GPU) ReadPrelude -privStms :: Stms Kernels -> PrivStms +privStms :: Stms GPU -> PrivStms privStms stms = PrivStms stms $ const $ return () -addPrivStms :: Slice SubExp -> PrivStms -> Binder Kernels () +addPrivStms :: Slice SubExp -> PrivStms -> Binder GPU () addPrivStms local_slice (PrivStms stms readPrelude) = do readPrelude local_slice addStms stms @@ -541,13 +536,13 @@ instance Monoid PrivStms where mempty = privStms mempty -type ReadPrelude = Slice SubExp -> Binder Kernels () +type ReadPrelude = Slice SubExp -> Binder GPU () data ProcessTileArgs = ProcessTileArgs { processPrivStms :: PrivStms, processComm :: Commutativity, - processRedLam :: Lambda Kernels, - processMapLam :: Lambda Kernels, + processRedLam :: Lambda GPU, + processMapLam :: Lambda GPU, processTiles :: [InputTile], processAcc :: [VName], processTileId :: SubExp @@ -556,8 +551,8 @@ data ResidualTileArgs = ResidualTileArgs { residualPrivStms :: PrivStms, residualComm :: Commutativity, - residualRedLam :: Lambda Kernels, - residualMapLam :: Lambda Kernels, + residualRedLam :: Lambda GPU, + residualMapLam :: Lambda GPU, residualInput :: [InputArray], residualAcc :: [VName], residualInputSize :: SubExp, @@ -572,8 +567,8 @@ String -> SegLevel -> ResultManifest -> - (PrimExp VName -> Slice SubExp -> Binder Kernels [SubExp]) -> - Binder Kernels [VName], + (PrimExp VName -> Slice SubExp -> Binder GPU [SubExp]) -> + Binder GPU [VName], -- The boolean PrimExp indicates whether they are in-bounds. tilingReadTile :: @@ -581,22 +576,22 @@ PrivStms -> SubExp -> [InputArray] -> - Binder Kernels [InputTile], + Binder GPU [InputTile], tilingProcessTile :: ProcessTileArgs -> - Binder Kernels [VName], + Binder GPU [VName], tilingProcessResidualTile :: ResidualTileArgs -> - Binder Kernels [VName], - tilingTileReturns :: VName -> Binder Kernels KernelResult, + Binder GPU [VName], + tilingTileReturns :: VName -> Binder GPU KernelResult, tilingSpace :: SegSpace, tilingTileShape :: Shape, tilingLevel :: SegLevel, - tilingNumWholeTiles :: Binder Kernels SubExp + tilingNumWholeTiles :: Binder GPU SubExp } type DoTiling gtids kdims = - SegLevel -> gtids -> kdims -> SubExp -> Binder Kernels Tiling + SegLevel -> gtids -> kdims -> SubExp -> Binder GPU Tiling scalarLevel :: Tiling -> SegLevel scalarLevel tiling = @@ -608,20 +603,35 @@ String -> PrimExp VName -> [Type] -> - Binder Kernels [SubExp] -> - Binder Kernels [VName] -protectOutOfBounds desc in_bounds ts m = - letTupExp desc =<< eIf (toExp in_bounds) (resultBody <$> m) (eBody $ map eBlank ts) + Binder GPU [SubExp] -> + Binder GPU [VName] +protectOutOfBounds desc in_bounds ts m = do + -- This is more complicated than you might expect, because we need + -- to be able to produce a blank accumulator, which eBlank cannot + -- do. By the linear type rules of accumulators, the body returns + -- an accumulator of type 'acc_t', then a unique variable of type + -- 'acc_t' must also be free in the body. This means we can find it + -- based just on the type. + m_body <- insertStmsM $ resultBody <$> m + let m_body_free = namesToList $ freeIn m_body + t_to_v <- + filter (isAcc . fst) + <$> (zip <$> mapM lookupType m_body_free <*> pure m_body_free) + let blank t = maybe (eBlank t) (pure . BasicOp . SubExp . Var) $ lookup t t_to_v + letTupExp desc =<< eIf (toExp in_bounds) (pure m_body) (eBody $ map blank ts) + where + isAcc Acc {} = True + isAcc _ = False postludeGeneric :: Tiling -> PrivStms -> - Pattern Kernels -> + Pattern GPU -> [VName] -> - Stms Kernels -> + Stms GPU -> Result -> [Type] -> - Binder Kernels [VName] + Binder GPU [VName] postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts = tilingSegMap tiling "thread_res" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> do -- Read our per-thread result from the tiled loop. @@ -640,21 +650,21 @@ addStms poststms return poststms_res -type TiledBody = Names -> PrivStms -> Binder Kernels [VName] +type TiledBody = Names -> PrivStms -> Binder GPU [VName] tileGeneric :: DoTiling gtids kdims -> SegLevel -> [Type] -> - Pattern Kernels -> + Pattern GPU -> gtids -> kdims -> SubExp -> - (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) -> + (Commutativity, Lambda GPU, [SubExp], Lambda GPU) -> [InputArray] -> - Stms Kernels -> + Stms GPU -> Result -> - TileM (Stms Kernels, Tiling, TiledBody) + TileM (Stms GPU, Tiling, TiledBody) tileGeneric doTiling initial_lvl res_ts pat gtids kdims w form inputs poststms poststms_res = do (tiling, tiling_stms) <- runBinder $ doTiling initial_lvl gtids kdims w @@ -662,7 +672,7 @@ where (red_comm, red_lam, red_nes, map_lam) = form - tiledBody :: Tiling -> Names -> PrivStms -> Binder Kernels [VName] + tiledBody :: Tiling -> Names -> PrivStms -> Binder GPU [VName] tiledBody tiling _private privstms = do let tile_shape = tilingTileShape tiling @@ -723,14 +733,14 @@ arr_t <- lookupType arr letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice -tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Binder Kernels KernelResult +tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Binder GPU KernelResult tileReturns dims_on_top dims arr = do let unit_dims = replicate (length dims_on_top) (intConst Int64 1) + arr_t <- lookupType arr arr' <- - if null dims_on_top + if null dims_on_top || null (arrayDims arr_t) -- Second check is for accumulators. then return arr else do - arr_t <- lookupType arr let new_shape = unit_dims ++ arrayDims arr_t letExp (baseString arr) $ BasicOp $ Reshape (map DimNew new_shape) arr let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims @@ -743,34 +753,12 @@ | otherwise = InputDontTile arr -segMap1D :: - String -> - SegLevel -> - ResultManifest -> - (VName -> Binder Kernels [SubExp]) -> - Binder Kernels [VName] -segMap1D desc lvl manifest f = do - ltid <- newVName "ltid" - ltid_flat <- newVName "ltid_flat" - let space = SegSpace ltid_flat [(ltid, unCount $ segGroupSize lvl)] - - ((ts, res), stms) <- runBinder $ do - res <- f ltid - ts <- mapM subExpType res - return (ts, res) - Body _ stms' res' <- renameBody $ mkBody stms res - - letTupExp desc $ - Op $ - SegOp $ - SegMap lvl space ts $ KernelBody () stms' $ map (Returns manifest) res' - reconstructGtids1D :: Count GroupSize SubExp -> VName -> VName -> VName -> - Binder Kernels () + Binder GPU () reconstructGtids1D group_size gtid gid ltid = letBindNames [gtid] =<< toExp (le64 gid * pe64 (unCount group_size) + le64 ltid) @@ -785,7 +773,7 @@ PrivStms -> SubExp -> [InputArray] -> - Binder Kernels [InputTile] + Binder GPU [InputTile] readTile1D tile_size gid gtid num_groups group_size kind privstms tile_id inputs = fmap (inputsToTiles inputs) . segMap1D "full_tile" lvl ResultNoSimplify @@ -826,7 +814,7 @@ Count NumGroups SubExp -> Count GroupSize SubExp -> ProcessTileArgs -> - Binder Kernels [VName] + Binder GPU [VName] processTile1D gid gtid kdim tile_size num_groups group_size tile_args = do let red_comm = processComm tile_args privstms = processPrivStms tile_args @@ -870,7 +858,7 @@ Count NumGroups SubExp -> Count GroupSize SubExp -> ResidualTileArgs -> - Binder Kernels [VName] + Binder GPU [VName] processResidualTile1D gid gtid kdim tile_size num_groups group_size args = do -- The number of residual elements that are not covered by -- the whole tiles. @@ -1001,7 +989,7 @@ (VName, VName) -> (VName, VName) -> (VName, VName) -> - Binder Kernels () + Binder GPU () reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y) = do -- Reconstruct the original gtids from gid_x/gid_y and ltid_x/ltid_y. letBindNames [gtid_x] @@ -1020,7 +1008,7 @@ PrivStms -> SubExp -> [InputArray] -> - Binder Kernels [InputTile] + Binder GPU [InputTile] readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size num_groups group_size kind privstms tile_id inputs = fmap (inputsToTiles inputs) . segMap2D @@ -1075,7 +1063,7 @@ TileFull -> mapM readTileElem arrs_and_perms -findTileSize :: HasScope lore m => [InputTile] -> m SubExp +findTileSize :: HasScope rep m => [InputTile] -> m SubExp findTileSize tiles = case mapMaybe isTiled tiles of v : _ -> arraySize 0 <$> lookupType v @@ -1092,7 +1080,7 @@ Count NumGroups SubExp -> Count GroupSize SubExp -> ProcessTileArgs -> - Binder Kernels [VName] + Binder GPU [VName] processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size num_groups group_size tile_args = do let privstms = processPrivStms tile_args red_comm = processComm tile_args @@ -1148,7 +1136,7 @@ Count NumGroups SubExp -> Count GroupSize SubExp -> ResidualTileArgs -> - Binder Kernels [VName] + Binder GPU [VName] processResidualTile2D gids gtids diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs --- a/src/Futhark/Optimise/TileLoops/Shared.hs +++ b/src/Futhark/Optimise/TileLoops/Shared.hs @@ -1,5 +1,6 @@ module Futhark.Optimise.TileLoops.Shared ( TileM, + segMap1D, segMap2D, segMap3D, segScatter2D, @@ -13,29 +14,51 @@ import Control.Monad.State import Data.List (foldl', zip4) import qualified Data.Map as M -import Futhark.IR.Kernels +import Futhark.IR.GPU import Futhark.MonadFreshNames import Futhark.Tools import Futhark.Transform.Rename -type TileM = ReaderT (Scope Kernels) (State VNameSource) +type TileM = ReaderT (Scope GPU) (State VNameSource) +segMap1D :: + String -> + SegLevel -> + ResultManifest -> + (VName -> Binder GPU [SubExp]) -> + Binder GPU [VName] +segMap1D desc lvl manifest f = do + ltid <- newVName "ltid" + ltid_flat <- newVName "ltid_flat" + let space = SegSpace ltid_flat [(ltid, unCount $ segGroupSize lvl)] + + ((ts, res), stms) <- localScope (scopeOfSegSpace space) . runBinder $ do + res <- f ltid + ts <- mapM subExpType res + return (ts, res) + Body _ stms' res' <- renameBody $ mkBody stms res + + letTupExp desc $ + Op $ + SegOp $ + SegMap lvl space ts $ KernelBody () stms' $ map (Returns manifest) res' + segMap2D :: String -> -- desc SegLevel -> -- lvl ResultManifest -> -- manifest (SubExp, SubExp) -> -- (dim_x, dim_y) ( (VName, VName) -> -- f - Binder Kernels [SubExp] + Binder GPU [SubExp] ) -> - Binder Kernels [VName] + Binder GPU [VName] segMap2D desc lvl manifest (dim_y, dim_x) f = do ltid_xx <- newVName "ltid_x" ltid_flat <- newVName "ltid_flat" ltid_yy <- newVName "ltid_y" let segspace = SegSpace ltid_flat [(ltid_yy, dim_y), (ltid_xx, dim_x)] - ((ts, res), stms) <- runBinder $ do + ((ts, res), stms) <- localScope (scopeOfSegSpace segspace) . runBinder $ do res <- f (ltid_yy, ltid_xx) ts <- mapM subExpType res return (ts, res) @@ -51,9 +74,9 @@ ResultManifest -> -- manifest (SubExp, SubExp, SubExp) -> -- (dim_z, dim_y, dim_x) ( (VName, VName, VName) -> -- f - Binder Kernels [SubExp] + Binder GPU [SubExp] ) -> - Binder Kernels [VName] + Binder GPU [VName] segMap3D desc lvl manifest (dim_z, dim_y, dim_x) f = do ltid_x <- newVName "ltid_x" ltid_flat <- newVName "ltid_flat" @@ -61,7 +84,7 @@ ltid_z <- newVName "ltid_z" let segspace = SegSpace ltid_flat [(ltid_z, dim_z), (ltid_y, dim_y), (ltid_x, dim_x)] - ((ts, res), stms) <- runBinder $ do + ((ts, res), stms) <- localScope (scopeOfSegSpace segspace) . runBinder $ do res <- f (ltid_z, ltid_y, ltid_x) ts <- mapM subExpType res return (ts, res) @@ -77,8 +100,8 @@ VName -> SegLevel -> -- lvl (SubExp, SubExp) -> -- (dim_y, dim_x) - ((VName, VName) -> Binder Kernels (SubExp, SubExp)) -> -- f - Binder Kernels [VName] + ((VName, VName) -> Binder GPU (SubExp, SubExp)) -> -- f + Binder GPU [VName] segScatter2D desc arr_size updt_arr lvl (dim_x, dim_y) f = do ltid_x <- newVName "ltid_x" ltid_y <- newVName "ltid_y" @@ -103,11 +126,11 @@ type VarianceTable = M.Map VName Names isTileableRedomap :: - Stm Kernels -> + Stm GPU -> Maybe ( SubExp, [VName], - (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) + (Commutativity, Lambda GPU, [SubExp], Lambda GPU) ) isTileableRedomap stm | Op (OtherOp (Screma w arrs form)) <- stmExp stm, @@ -123,7 +146,7 @@ | otherwise = Nothing -defVarianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable +defVarianceInStm :: VarianceTable -> Stm GPU -> VarianceTable defVarianceInStm variance bnd = foldl' add variance $ patternNames $ stmPattern bnd where @@ -133,7 +156,7 @@ -- just in case you need the Screma being treated differently than -- by default; previously Cosmin had to enhance it when dealing with stream. -varianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable +varianceInStm :: VarianceTable -> Stm GPU -> VarianceTable varianceInStm v0 bnd@(Let _ _ (Op (OtherOp Screma {}))) | Just (_, arrs, (_, red_lam, red_nes, map_lam)) <- isTileableRedomap bnd = let v = defVarianceInStm v0 bnd @@ -156,5 +179,5 @@ in varianceInStms v' stm_lam varianceInStm v0 bnd = defVarianceInStm v0 bnd -varianceInStms :: VarianceTable -> Stms Kernels -> VarianceTable +varianceInStms :: VarianceTable -> Stms GPU -> VarianceTable varianceInStms = foldl' varianceInStm diff --git a/src/Futhark/Optimise/Unstream.hs b/src/Futhark/Optimise/Unstream.hs --- a/src/Futhark/Optimise/Unstream.hs +++ b/src/Futhark/Optimise/Unstream.hs @@ -18,13 +18,13 @@ -- simplified away, but only *before* they are turned into loops. In -- principle this pass could be split into multiple, but for now it is -- kept together. -module Futhark.Optimise.Unstream (unstreamKernels, unstreamMC) where +module Futhark.Optimise.Unstream (unstreamGPU, unstreamMC) where import Control.Monad.Reader import Control.Monad.State -import Futhark.IR.Kernels -import qualified Futhark.IR.Kernels as Kernels -import Futhark.IR.Kernels.Simplify (simplifyKernels) +import Futhark.IR.GPU +import qualified Futhark.IR.GPU as GPU +import Futhark.IR.GPU.Simplify (simplifyGPU) import Futhark.IR.MC import qualified Futhark.IR.MC as MC import Futhark.MonadFreshNames @@ -33,8 +33,8 @@ import qualified Futhark.Transform.FirstOrderTransform as FOT -- | The pass for GPU kernels. -unstreamKernels :: Pass Kernels Kernels -unstreamKernels = unstream onHostOp simplifyKernels +unstreamGPU :: Pass GPU GPU +unstreamGPU = unstream onHostOp simplifyGPU -- | The pass for multicore. unstreamMC :: Pass MC MC @@ -43,10 +43,10 @@ data Stage = SeqStreams | SeqAll unstream :: - ASTLore lore => - (Stage -> OnOp lore) -> - (Prog lore -> PassM (Prog lore)) -> - Pass lore lore + ASTRep rep => + (Stage -> OnOp rep) -> + (Prog rep -> PassM (Prog rep)) -> + Pass rep rep unstream onOp simplify = Pass "unstream" "sequentialise remaining SOACs" $ intraproceduralTransformation (optimise SeqStreams) @@ -58,33 +58,33 @@ runState $ runReaderT (optimiseStms (onOp stage) stms) scope -type UnstreamM lore = ReaderT (Scope lore) (State VNameSource) +type UnstreamM rep = ReaderT (Scope rep) (State VNameSource) -type OnOp lore = - Pattern lore -> StmAux (ExpDec lore) -> Op lore -> UnstreamM lore [Stm lore] +type OnOp rep = + Pattern rep -> StmAux (ExpDec rep) -> Op rep -> UnstreamM rep [Stm rep] optimiseStms :: - ASTLore lore => - OnOp lore -> - Stms lore -> - UnstreamM lore (Stms lore) + ASTRep rep => + OnOp rep -> + Stms rep -> + UnstreamM rep (Stms rep) optimiseStms onOp stms = localScope (scopeOf stms) $ stmsFromList . concat <$> mapM (optimiseStm onOp) (stmsToList stms) optimiseBody :: - ASTLore lore => - OnOp lore -> - Body lore -> - UnstreamM lore (Body lore) + ASTRep rep => + OnOp rep -> + Body rep -> + UnstreamM rep (Body rep) optimiseBody onOp (Body aux stms res) = Body aux <$> optimiseStms onOp stms <*> pure res optimiseKernelBody :: - ASTLore lore => - OnOp lore -> - KernelBody lore -> - UnstreamM lore (KernelBody lore) + ASTRep rep => + OnOp rep -> + KernelBody rep -> + UnstreamM rep (KernelBody rep) optimiseKernelBody onOp (KernelBody attr stms res) = localScope (scopeOf stms) $ KernelBody attr @@ -92,19 +92,19 @@ <*> pure res optimiseLambda :: - ASTLore lore => - OnOp lore -> - Lambda lore -> - UnstreamM lore (Lambda lore) + ASTRep rep => + OnOp rep -> + Lambda rep -> + UnstreamM rep (Lambda rep) optimiseLambda onOp lam = localScope (scopeOfLParams $ lambdaParams lam) $ do body <- optimiseBody onOp $ lambdaBody lam return lam {lambdaBody = body} optimiseStm :: - ASTLore lore => - OnOp lore -> - Stm lore -> - UnstreamM lore [Stm lore] + ASTRep rep => + OnOp rep -> + Stm rep -> + UnstreamM rep [Stm rep] optimiseStm onOp (Let pat aux (Op op)) = onOp pat aux op optimiseStm onOp (Let pat aux e) = @@ -117,10 +117,10 @@ } optimiseSegOp :: - ASTLore lore => - OnOp lore -> - SegOp lvl lore -> - UnstreamM lore (SegOp lvl lore) + ASTRep rep => + OnOp rep -> + SegOp lvl rep -> + UnstreamM rep (SegOp lvl rep) optimiseSegOp onOp op = localScope (scopeOfSegSpace $ segSpace op) $ mapSegOpM optimise op where @@ -150,13 +150,13 @@ { mapOnSOACLambda = optimiseLambda (onMCOp stage) } -sequentialise :: Stage -> SOAC lore -> Bool +sequentialise :: Stage -> SOAC rep -> Bool sequentialise SeqStreams Stream {} = True sequentialise SeqStreams _ = False sequentialise SeqAll _ = True -onHostOp :: Stage -> OnOp Kernels -onHostOp stage pat aux (Kernels.OtherOp soac) +onHostOp :: Stage -> OnOp GPU +onHostOp stage pat aux (GPU.OtherOp soac) | sequentialise stage soac = do stms <- runBinder_ $ FOT.transformSOAC pat soac fmap concat $ @@ -164,7 +164,7 @@ mapM (optimiseStm (onHostOp stage)) $ stmsToList stms | otherwise = -- Still sequentialise whatever's inside. - pure <$> (Let pat aux . Op . Kernels.OtherOp <$> mapSOACM optimise soac) + pure <$> (Let pat aux . Op . GPU.OtherOp <$> mapSOACM optimise soac) where optimise = identitySOACMapper diff --git a/src/Futhark/Pass.hs b/src/Futhark/Pass.hs --- a/src/Futhark/Pass.hs +++ b/src/Futhark/Pass.hs @@ -3,7 +3,7 @@ {-# LANGUAGE Strict #-} -- | Definition of a polymorphic (generic) pass that can work with --- programs of any lore. +-- programs of any rep. module Futhark.Pass ( PassM, runPassM, @@ -58,9 +58,9 @@ liftEitherM :: Show err => PassM (Either err a) -> PassM a liftEitherM m = liftEither =<< m --- | A compiler pass transforming a 'Prog' of a given lore to a 'Prog' --- of another lore. -data Pass fromlore tolore = Pass +-- | A compiler pass transforming a 'Prog' of a given rep to a 'Prog' +-- of another rep. +data Pass fromrep torep = Pass { -- | Name of the pass. Keep this short and simple. It will -- be used to automatically generate a command-line option -- name via 'passLongOption'. @@ -68,12 +68,12 @@ -- | A slightly longer description, which will show up in the -- command-line help text. passDescription :: String, - passFunction :: Prog fromlore -> PassM (Prog tolore) + passFunction :: Prog fromrep -> PassM (Prog torep) } -- | Take the name of the pass, turn spaces into dashes, and make all -- characters lowercase. -passLongOption :: Pass fromlore tolore -> String +passLongOption :: Pass fromrep torep -> String passLongOption = map (spaceToDash . toLower) . passName where spaceToDash ' ' = '-' @@ -101,10 +101,10 @@ -- The function definition transformations are run in parallel (with -- 'parPass'), since they cannot affect each other. intraproceduralTransformationWithConsts :: - (Stms fromlore -> PassM (Stms tolore)) -> - (Stms tolore -> FunDef fromlore -> PassM (FunDef tolore)) -> - Prog fromlore -> - PassM (Prog tolore) + (Stms fromrep -> PassM (Stms torep)) -> + (Stms torep -> FunDef fromrep -> PassM (FunDef torep)) -> + Prog fromrep -> + PassM (Prog torep) intraproceduralTransformationWithConsts ct ft (Prog consts funs) = do consts' <- ct consts funs' <- parPass (ft consts') funs @@ -113,9 +113,9 @@ -- | Like 'intraproceduralTransformationWithConsts', but do not change -- the top-level constants, and simply pass along their 'Scope'. intraproceduralTransformation :: - (Scope lore -> Stms lore -> PassM (Stms lore)) -> - Prog lore -> - PassM (Prog lore) + (Scope rep -> Stms rep -> PassM (Stms rep)) -> + Prog rep -> + PassM (Prog rep) intraproceduralTransformation f = intraproceduralTransformationWithConsts (f mempty) f' where diff --git a/src/Futhark/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs --- a/src/Futhark/Pass/ExpandAllocations.hs +++ b/src/Futhark/Pass/ExpandAllocations.hs @@ -16,15 +16,15 @@ import qualified Futhark.Analysis.SymbolTable as ST import Futhark.Error import Futhark.IR -import qualified Futhark.IR.Kernels.Simplify as Kernels -import Futhark.IR.KernelsMem +import qualified Futhark.IR.GPU.Simplify as GPU +import Futhark.IR.GPUMem import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.MonadFreshNames -import Futhark.Optimise.Simplify.Lore (addScopeWisdom) +import Futhark.Optimise.Simplify.Rep (addScopeWisdom) import Futhark.Pass -import Futhark.Pass.ExplicitAllocations.Kernels (explicitAllocationsInStms) +import Futhark.Pass.ExplicitAllocations.GPU (explicitAllocationsInStms) import Futhark.Pass.ExtractKernels.BlockedKernel (nonSegRed) -import Futhark.Pass.ExtractKernels.ToKernels (segThread) +import Futhark.Pass.ExtractKernels.ToGPU (segThread) import Futhark.Tools import Futhark.Transform.CopyPropagate (copyPropagateInFun) import Futhark.Transform.Rename (renameStm) @@ -33,7 +33,7 @@ import Prelude hiding (quot) -- | The memory expansion pass definition. -expandAllocations :: Pass KernelsMem KernelsMem +expandAllocations :: Pass GPUMem GPUMem expandAllocations = Pass "expand allocations" "Expand allocations" $ \(Prog consts funs) -> do @@ -45,19 +45,19 @@ -- duplicate size keys (which are not fixed by renamer, and size -- keys must currently be globally unique). -type ExpandM = ReaderT (Scope KernelsMem) (StateT VNameSource (Either String)) +type ExpandM = ReaderT (Scope GPUMem) (StateT VNameSource (Either String)) limitationOnLeft :: Either String a -> a limitationOnLeft = either compilerLimitationS id transformFunDef :: - Scope KernelsMem -> - FunDef KernelsMem -> - PassM (FunDef KernelsMem) + Scope GPUMem -> + FunDef GPUMem -> + PassM (FunDef GPUMem) transformFunDef scope fundec = do body' <- modifyNameSource $ limitationOnLeft . runStateT (runReaderT m mempty) copyPropagateInFun - simpleKernelsMem + simpleGPUMem (ST.fromScope (addScopeWisdom scope)) fundec {funDefBody = body'} where @@ -66,20 +66,20 @@ inScopeOf fundec $ transformBody $ funDefBody fundec -transformBody :: Body KernelsMem -> ExpandM (Body KernelsMem) +transformBody :: Body GPUMem -> ExpandM (Body GPUMem) transformBody (Body () stms res) = Body () <$> transformStms stms <*> pure res -transformLambda :: Lambda KernelsMem -> ExpandM (Lambda KernelsMem) +transformLambda :: Lambda GPUMem -> ExpandM (Lambda GPUMem) transformLambda (Lambda params body ret) = Lambda params <$> localScope (scopeOfLParams params) (transformBody body) <*> pure ret -transformStms :: Stms KernelsMem -> ExpandM (Stms KernelsMem) +transformStms :: Stms GPUMem -> ExpandM (Stms GPUMem) transformStms stms = inScopeOf stms $ mconcat <$> mapM transformStm (stmsToList stms) -transformStm :: Stm KernelsMem -> ExpandM (Stms KernelsMem) +transformStm :: Stm GPUMem -> ExpandM (Stms GPUMem) -- It is possible that we are unable to expand allocations in some -- code versions. If so, we can remove the offending branch. Only if -- both versions fail do we propagate the error. @@ -110,7 +110,7 @@ { mapOnBody = \scope -> localScope scope . transformBody } -transformExp :: Exp KernelsMem -> ExpandM (Stms KernelsMem, Exp KernelsMem) +transformExp :: Exp GPUMem -> ExpandM (Stms GPUMem, Exp GPUMem) transformExp (Op (Inner (SegOp (SegMap lvl space ts kbody)))) = do (alloc_stms, (_, kbody')) <- transformScanRed lvl space [] kbody return @@ -190,9 +190,9 @@ transformScanRed :: SegLevel -> SegSpace -> - [Lambda KernelsMem] -> - KernelBody KernelsMem -> - ExpandM (Stms KernelsMem, ([Lambda KernelsMem], KernelBody KernelsMem)) + [Lambda GPUMem] -> + KernelBody GPUMem -> + ExpandM (Stms GPUMem, ([Lambda GPUMem], KernelBody GPUMem)) transformScanRed lvl space ops kbody = do bound_outside <- asks $ namesFromList . M.keys let user = (lvl, [le64 $ segFlat space]) @@ -230,7 +230,7 @@ namesFromList (M.keys $ scopeOfSegSpace space) <> boundInKernelBody kbody -boundInKernelBody :: KernelBody KernelsMem -> Names +boundInKernelBody :: KernelBody GPUMem -> Names boundInKernelBody = namesFromList . M.keys . scopeOf . kernelBodyStms allocsForBody :: @@ -238,8 +238,8 @@ Extraction -> SegLevel -> SegSpace -> - KernelBody KernelsMem -> - (Stms KernelsMem -> KernelBody KernelsMem -> OffsetM b) -> + KernelBody GPUMem -> + (Stms GPUMem -> KernelBody GPUMem -> OffsetM b) -> ExpandM b allocsForBody variant_allocs invariant_allocs lvl space kbody' m = do (alloc_offsets, alloc_stms) <- @@ -260,10 +260,10 @@ memoryRequirements :: SegLevel -> SegSpace -> - Stms KernelsMem -> + Stms GPUMem -> Extraction -> Extraction -> - ExpandM (RebaseMap, Stms KernelsMem) + ExpandM (RebaseMap, Stms GPUMem) memoryRequirements lvl space kstms variant_allocs invariant_allocs = do (num_threads, num_threads_stms) <- runBinder . letSubExp "num_threads" . BasicOp $ @@ -305,8 +305,8 @@ User -> Names -> Names -> - KernelBody KernelsMem -> - ( KernelBody KernelsMem, + KernelBody GPUMem -> + ( KernelBody GPUMem, Extraction ) extractKernelBodyAllocations lvl bound_outside bound_kernel = @@ -317,8 +317,8 @@ User -> Names -> Names -> - Body KernelsMem -> - (Body KernelsMem, Extraction) + Body GPUMem -> + (Body GPUMem, Extraction) extractBodyAllocations user bound_outside bound_kernel = extractGenericBodyAllocations user bound_outside bound_kernel bodyStms $ \stms body -> body {bodyStms = stms} @@ -327,8 +327,8 @@ User -> Names -> Names -> - Lambda KernelsMem -> - (Lambda KernelsMem, Extraction) + Lambda GPUMem -> + (Lambda GPUMem, Extraction) extractLambdaAllocations user bound_outside bound_kernel lam = (lam {lambdaBody = body'}, allocs) where (body', allocs) = extractBodyAllocations user bound_outside bound_kernel $ lambdaBody lam @@ -337,8 +337,8 @@ User -> Names -> Names -> - (body -> Stms KernelsMem) -> - (Stms KernelsMem -> body -> body) -> + (body -> Stms GPUMem) -> + (Stms GPUMem -> body -> body) -> body -> ( body, Extraction @@ -363,8 +363,8 @@ User -> Names -> Names -> - Stm KernelsMem -> - Writer Extraction (Maybe (Stm KernelsMem)) + Stm GPUMem -> + Writer Extraction (Maybe (Stm GPUMem)) extractStmAllocations user bound_outside bound_kernel (Let (Pattern [] [patElem]) _ (Op (Alloc size space))) | expandable space && expandableSize size -- FIXME: the '&& notScalar space' part is a hack because we @@ -417,7 +417,7 @@ return lam {lambdaBody = body} genericExpandedInvariantAllocations :: - (User -> (Shape, [TPrimExp Int64 VName])) -> Extraction -> ExpandM (Stms KernelsMem, RebaseMap) + (User -> (Shape, [TPrimExp Int64 VName])) -> Extraction -> ExpandM (Stms GPUMem, RebaseMap) genericExpandedInvariantAllocations getNumUsers invariant_allocs = do -- We expand the invariant allocations by adding an inner dimension -- equal to the number of kernel threads. @@ -459,7 +459,7 @@ Count NumGroups SubExp -> Count GroupSize SubExp -> Extraction -> - ExpandM (Stms KernelsMem, RebaseMap) + ExpandM (Stms GPUMem, RebaseMap) expandedInvariantAllocations num_threads (Count num_groups) (Count group_size) = genericExpandedInvariantAllocations getNumUsers where @@ -471,9 +471,9 @@ expandedVariantAllocations :: SubExp -> SegSpace -> - Stms KernelsMem -> + Stms GPUMem -> Extraction -> - ExpandM (Stms KernelsMem, RebaseMap) + ExpandM (Stms GPUMem, RebaseMap) expandedVariantAllocations _ _ _ variant_allocs | null variant_allocs = return (mempty, mempty) expandedVariantAllocations num_threads kspace kstms variant_allocs = do @@ -538,7 +538,7 @@ newtype OffsetM a = OffsetM ( ReaderT - (Scope KernelsMem) + (Scope GPUMem) (ReaderT RebaseMap (Either String)) a ) @@ -546,12 +546,12 @@ ( Applicative, Functor, Monad, - HasScope KernelsMem, - LocalScope KernelsMem, + HasScope GPUMem, + LocalScope GPUMem, MonadError String ) -runOffsetM :: Scope KernelsMem -> RebaseMap -> OffsetM a -> Either String a +runOffsetM :: Scope GPUMem -> RebaseMap -> OffsetM a -> Either String a runOffsetM scope offsets (OffsetM m) = runReaderT (runReaderT m scope) offsets @@ -563,7 +563,7 @@ offsets <- askRebaseMap return $ ($ x) <$> M.lookup name offsets -offsetMemoryInKernelBody :: KernelBody KernelsMem -> OffsetM (KernelBody KernelsMem) +offsetMemoryInKernelBody :: KernelBody GPUMem -> OffsetM (KernelBody GPUMem) offsetMemoryInKernelBody kbody = do scope <- askScope stms' <- @@ -574,7 +574,7 @@ (stmsToList $ kernelBodyStms kbody) return kbody {kernelBodyStms = stms'} -offsetMemoryInBody :: Body KernelsMem -> OffsetM (Body KernelsMem) +offsetMemoryInBody :: Body GPUMem -> OffsetM (Body GPUMem) offsetMemoryInBody (Body dec stms res) = do scope <- askScope stms' <- @@ -585,7 +585,7 @@ (stmsToList stms) return $ Body dec stms' res -offsetMemoryInStm :: Stm KernelsMem -> OffsetM (Scope KernelsMem, Stm KernelsMem) +offsetMemoryInStm :: Stm GPUMem -> OffsetM (Scope GPUMem, Stm GPUMem) offsetMemoryInStm (Let pat dec e) = do pat' <- offsetMemoryInPattern pat e' <- localScope (scopeOfPattern pat') $ offsetMemoryInExp e @@ -618,7 +618,7 @@ inst Ext {} = Nothing inst (Free x) = return x -offsetMemoryInPattern :: Pattern KernelsMem -> OffsetM (Pattern KernelsMem) +offsetMemoryInPattern :: Pattern GPUMem -> OffsetM (Pattern GPUMem) offsetMemoryInPattern (Pattern ctx vals) = do mapM_ inspectCtx ctx Pattern ctx <$> mapM inspectVal vals @@ -664,12 +664,12 @@ IxFun.rebase (fmap (fmap Free) new_base') ixfun offsetMemoryInBodyReturns br = return br -offsetMemoryInLambda :: Lambda KernelsMem -> OffsetM (Lambda KernelsMem) +offsetMemoryInLambda :: Lambda GPUMem -> OffsetM (Lambda GPUMem) offsetMemoryInLambda lam = inScopeOf lam $ do body <- offsetMemoryInBody $ lambdaBody lam return $ lam {lambdaBody = body} -offsetMemoryInExp :: Exp KernelsMem -> OffsetM (Exp KernelsMem) +offsetMemoryInExp :: Exp GPUMem -> OffsetM (Exp GPUMem) offsetMemoryInExp (DoLoop ctx val form body) = do let (ctxparams, ctxinit) = unzip ctx (valparams, valinit) = unzip val @@ -698,8 +698,8 @@ ---- Slicing allocation sizes out of a kernel. -unAllocKernelsStms :: Stms KernelsMem -> Either String (Stms Kernels.Kernels) -unAllocKernelsStms = unAllocStms False +unAllocGPUStms :: Stms GPUMem -> Either String (Stms GPU.GPU) +unAllocGPUStms = unAllocStms False where unAllocBody (Body dec stms res) = Body dec <$> unAllocStms True stms <*> pure res @@ -765,7 +765,7 @@ unMem (MemAcc acc ispace ts u) = Just $ Acc acc ispace ts u unMem MemMem {} = Nothing -unAllocScope :: Scope KernelsMem -> Scope Kernels.Kernels +unAllocScope :: Scope GPUMem -> Scope GPU.GPU unAllocScope = M.mapMaybe unInfo where unInfo (LetName dec) = LetName <$> unMem dec @@ -784,10 +784,10 @@ SubExp -> [SubExp] -> SegSpace -> - Stms KernelsMem -> - ExpandM (Stms Kernels.Kernels, [VName], [VName]) + Stms GPUMem -> + ExpandM (Stms GPU.GPU, [VName], [VName]) sliceKernelSizes num_threads sizes space kstms = do - kstms' <- either throwError return $ unAllocKernelsStms kstms + kstms' <- either throwError return $ unAllocGPUStms kstms let num_sizes = length sizes i64s = replicate num_sizes $ Prim int64 @@ -825,7 +825,7 @@ return sizes localScope (scopeOfSegSpace space) $ - Kernels.simplifyLambda (Lambda [flat_gtid_lparam] (Body () stms zs) i64s) + GPU.simplifyLambda (Lambda [flat_gtid_lparam] (Body () stms zs) i64s) ((maxes_per_thread, size_sums), slice_stms) <- flip runBinderT kernels_scope $ do pat <- diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -54,7 +54,7 @@ import Futhark.MonadFreshNames import Futhark.Optimise.Simplify.Engine (SimpleOps (..)) import qualified Futhark.Optimise.Simplify.Engine as Engine -import Futhark.Optimise.Simplify.Lore (mkWiseBody) +import Futhark.Optimise.Simplify.Rep (mkWiseBody) import Futhark.Pass import Futhark.Tools import Futhark.Util (splitAt3, splitFromEnd, takeLast) @@ -66,7 +66,7 @@ deriving (Eq, Ord, Show) bindAllocStm :: - (MonadBinder m, Op (Lore m) ~ MemOp inner) => + (MonadBinder m, Op (Rep m) ~ MemOp inner) => AllocStm -> m () bindAllocStm (SizeComputation name pe) = @@ -77,15 +77,15 @@ letBindNames [name] $ BasicOp $ Copy src class - (MonadFreshNames m, LocalScope lore m, Mem lore) => - Allocator lore m + (MonadFreshNames m, LocalScope rep m, Mem rep) => + Allocator rep m where addAllocStm :: AllocStm -> m () askDefaultSpace :: m Space default addAllocStm :: - ( Allocable fromlore lore, - m ~ AllocM fromlore lore + ( Allocable fromrep rep, + m ~ AllocM fromrep rep ) => AllocStm -> m () @@ -100,7 +100,7 @@ -- allocate space for. See 'ChunkMap' comment. dimAllocationSize :: SubExp -> m SubExp default dimAllocationSize :: - m ~ AllocM fromlore lore => + m ~ AllocM fromrep rep => SubExp -> m SubExp dimAllocationSize (Var v) = @@ -113,11 +113,11 @@ -- | Get those names that are known to be constants at run-time. askConsts :: m (S.Set VName) - expHints :: Exp lore -> m [ExpHint] + expHints :: Exp rep -> m [ExpHint] expHints = defaultExpHints allocateMemory :: - Allocator lore m => + Allocator rep m => String -> SubExp -> Space -> @@ -128,7 +128,7 @@ return v computeSize :: - Allocator lore m => + Allocator rep m => String -> PrimExp VName -> m SubExp @@ -137,19 +137,19 @@ addAllocStm $ SizeComputation v se return $ Var v -type Allocable fromlore tolore = - ( PrettyLore fromlore, - PrettyLore tolore, - Mem tolore, - FParamInfo fromlore ~ DeclType, - LParamInfo fromlore ~ Type, - BranchType fromlore ~ ExtType, - RetType fromlore ~ DeclExtType, - BodyDec fromlore ~ (), - BodyDec tolore ~ (), - ExpDec tolore ~ (), - SizeSubst (Op tolore), - BinderOps tolore +type Allocable fromrep torep = + ( PrettyRep fromrep, + PrettyRep torep, + Mem torep, + FParamInfo fromrep ~ DeclType, + LParamInfo fromrep ~ Type, + BranchType fromrep ~ ExtType, + RetType fromrep ~ DeclExtType, + BodyDec fromrep ~ (), + BodyDec torep ~ (), + ExpDec torep ~ (), + SizeSubst (Op torep), + BinderOps torep ) -- | A mapping from chunk names to their maximum size. XXX FIXME @@ -158,7 +158,7 @@ -- analysis yet (it should). type ChunkMap = M.Map VName SubExp -data AllocEnv fromlore tolore = AllocEnv +data AllocEnv fromrep torep = AllocEnv { chunkMap :: ChunkMap, -- | Aggressively try to reuse memory in do-loops - -- should be True inside kernels, False outside. @@ -170,28 +170,28 @@ -- | The set of names that are known to be constants at -- kernel compile time. envConsts :: S.Set VName, - allocInOp :: Op fromlore -> AllocM fromlore tolore (Op tolore), - envExpHints :: Exp tolore -> AllocM fromlore tolore [ExpHint] + allocInOp :: Op fromrep -> AllocM fromrep torep (Op torep), + envExpHints :: Exp torep -> AllocM fromrep torep [ExpHint] } -- | Monad for adding allocations to an entire program. -newtype AllocM fromlore tolore a - = AllocM (BinderT tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a) +newtype AllocM fromrep torep a + = AllocM (BinderT torep (ReaderT (AllocEnv fromrep torep) (State VNameSource)) a) deriving ( Applicative, Functor, Monad, MonadFreshNames, - HasScope tolore, - LocalScope tolore, - MonadReader (AllocEnv fromlore tolore) + HasScope torep, + LocalScope torep, + MonadReader (AllocEnv fromrep torep) ) instance - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - MonadBinder (AllocM fromlore tolore) + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => + MonadBinder (AllocM fromrep torep) where - type Lore (AllocM fromlore tolore) = tolore + type Rep (AllocM fromrep torep) = torep mkExpDecM _ _ = return () @@ -205,8 +205,8 @@ collectStms (AllocM m) = AllocM $ collectStms m instance - (Allocable fromlore tolore) => - Allocator tolore (AllocM fromlore tolore) + (Allocable fromrep torep) => + Allocator torep (AllocM fromrep torep) where expHints e = do f <- asks envExpHints @@ -217,9 +217,9 @@ runAllocM :: MonadFreshNames m => - (Op fromlore -> AllocM fromlore tolore (Op tolore)) -> - (Exp tolore -> AllocM fromlore tolore [ExpHint]) -> - AllocM fromlore tolore a -> + (Op fromrep -> AllocM fromrep torep (Op torep)) -> + (Exp torep -> AllocM fromrep torep [ExpHint]) -> + AllocM fromrep torep a -> m a runAllocM handleOp hints (AllocM m) = fmap fst $ modifyNameSource $ runState $ runReaderT (runBinderT m mempty) env @@ -235,10 +235,10 @@ } -- | Monad for adding allocations to a single pattern. -newtype PatAllocM lore a +newtype PatAllocM rep a = PatAllocM ( RWS - (Scope lore) + (Scope rep) [AllocStm] VNameSource a @@ -247,13 +247,13 @@ ( Applicative, Functor, Monad, - HasScope lore, - LocalScope lore, + HasScope rep, + LocalScope rep, MonadWriter [AllocStm], MonadFreshNames ) -instance Mem lore => Allocator lore (PatAllocM lore) where +instance Mem rep => Allocator rep (PatAllocM rep) where addAllocStm = tell . pure dimAllocationSize = return askDefaultSpace = return DefaultSpace @@ -261,8 +261,8 @@ runPatAllocM :: MonadFreshNames m => - PatAllocM lore a -> - Scope lore -> + PatAllocM rep a -> + Scope rep -> m (a, [AllocStm]) runPatAllocM (PatAllocM m) mems = modifyNameSource $ frob . runRWS m mems @@ -276,19 +276,22 @@ arraySizeInBytesExp t = untyped $ foldl' (*) (elemSize t) $ map pe64 (arrayDims t) -arraySizeInBytesExpM :: Allocator lore m => Type -> m (PrimExp VName) +arraySizeInBytesExpM :: Allocator rep m => Type -> m (PrimExp VName) arraySizeInBytesExpM t = do dims <- mapM dimAllocationSize (arrayDims t) - let dim_prod = product $ map pe64 dims - elm_size = elemSize t - return $ untyped $ dim_prod * elm_size + let dim_prod_i64 = product $ map pe64 dims + elm_size_i64 = elemSize t + return $ + BinOpExp (SMax Int64) (ValueExp $ IntValue $ Int64Value 0) $ + untyped $ + dim_prod_i64 * elm_size_i64 -arraySizeInBytes :: Allocator lore m => Type -> m SubExp +arraySizeInBytes :: Allocator rep m => Type -> m SubExp arraySizeInBytes = computeSize "bytes" <=< arraySizeInBytesExpM -- | Allocate memory for a value of the given type. allocForArray :: - Allocator lore m => + Allocator rep m => Type -> Space -> m VName @@ -297,11 +300,11 @@ allocateMemory "mem" size space allocsForStm :: - (Allocator lore m, ExpDec lore ~ ()) => + (Allocator rep m, ExpDec rep ~ ()) => [Ident] -> [Ident] -> - Exp lore -> - m (Stm lore) + Exp rep -> + m (Stm rep) allocsForStm sizeidents validents e = do rts <- expReturns e hints <- expHints e @@ -309,10 +312,10 @@ return $ Let (Pattern ctxElems valElems) (defAux ()) e patternWithAllocations :: - (Allocator lore m, ExpDec lore ~ ()) => + (Allocator rep m, ExpDec rep ~ ()) => [VName] -> - Exp lore -> - m (Pattern lore) + Exp rep -> + m (Pattern rep) patternWithAllocations names e = do (ts', sizes) <- instantiateShapes' =<< expExtType e let identForBindage name t = @@ -321,14 +324,14 @@ stmPattern <$> allocsForStm sizes vals e allocsForPattern :: - Allocator lore m => + Allocator rep m => [Ident] -> [Ident] -> [ExpReturns] -> [ExpHint] -> m - ( [PatElem lore], - [PatElem lore] + ( [PatElem rep], + [PatElem rep] ) allocsForPattern sizeidents validents rts hints = do let sizes' = [PatElem size $ MemPrim int64 | size <- map identName sizeidents] @@ -432,7 +435,7 @@ inst (Free x) = return x summaryForBindage :: - Allocator lore m => + Allocator rep m => Type -> ExpHint -> m (MemBound NoUniqueness) @@ -456,7 +459,7 @@ m <- allocateMemory "mem" bytes space return $ MemArray pt (arrayShape t) NoUniqueness $ ArrayIn m ixfun -lookupMemSpace :: (HasScope lore m, Monad m) => VName -> m Space +lookupMemSpace :: (HasScope rep m, Monad m) => VName -> m Space lookupMemSpace v = do t <- lookupType v case t of @@ -469,10 +472,10 @@ in MemArray bt shape u $ ArrayIn mem ixf allocInFParams :: - (Allocable fromlore tolore) => - [(FParam fromlore, Space)] -> - ([FParam tolore] -> AllocM fromlore tolore a) -> - AllocM fromlore tolore a + (Allocable fromrep torep) => + [(FParam fromrep, Space)] -> + ([FParam torep] -> AllocM fromrep torep a) -> + AllocM fromrep torep a allocInFParams params m = do (valparams, (ctxparams, memparams)) <- runWriterT $ mapM (uncurry allocInFParam) params @@ -481,13 +484,13 @@ localScope summary $ m params' allocInFParam :: - (Allocable fromlore tolore) => - FParam fromlore -> + (Allocable fromrep torep) => + FParam fromrep -> Space -> WriterT - ([FParam tolore], [FParam tolore]) - (AllocM fromlore tolore) - (FParam tolore) + ([FParam torep], [FParam torep]) + (AllocM fromrep torep) + (FParam torep) allocInFParam param pspace = case paramDeclType param of Array pt shape u -> do @@ -504,16 +507,16 @@ return param {paramDec = MemAcc acc ispace ts u} allocInMergeParams :: - ( Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore) + ( Allocable fromrep torep, + Allocator torep (AllocM fromrep torep) ) => - [(FParam fromlore, SubExp)] -> - ( [FParam tolore] -> - [FParam tolore] -> - ([SubExp] -> AllocM fromlore tolore ([SubExp], [SubExp])) -> - AllocM fromlore tolore a + [(FParam fromrep, SubExp)] -> + ( [FParam torep] -> + [FParam torep] -> + ([SubExp] -> AllocM fromrep torep ([SubExp], [SubExp])) -> + AllocM fromrep torep a ) -> - AllocM fromlore tolore a + AllocM fromrep torep a allocInMergeParams merge m = do ((valparams, handle_loop_subexps), (ctx_params, mem_params)) <- runWriterT $ unzip <$> mapM allocInMergeParam merge @@ -528,12 +531,12 @@ localScope summary $ m (ctx_params <> mem_params) valparams mk_loop_res where allocInMergeParam :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => (Param DeclType, SubExp) -> WriterT - ([FParam tolore], [FParam tolore]) - (AllocM fromlore tolore) - (FParam tolore, SubExp -> WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp) + ([FParam torep], [FParam torep]) + (AllocM fromrep torep) + (FParam torep, SubExp -> WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp) allocInMergeParam (mergeparam, Var v) | Array pt shape u <- paramDeclType mergeparam = do (mem', _) <- lift $ lookupArraySummary v @@ -577,10 +580,10 @@ -- Returns the existentialized index function, the list of substituted values and the memory location. existentializeArray :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => Space -> VName -> - AllocM fromlore tolore (SubExp, ExtIxFun, [TPrimExp Int64 VName], VName) + AllocM fromrep torep (SubExp, ExtIxFun, [TPrimExp Int64 VName], VName) existentializeArray ScalarSpace {} v = do (mem', ixfun) <- lookupArraySummary v return (Var v, fmap (fmap Free) ixfun, mempty, mem') @@ -599,12 +602,12 @@ return (subexp, fromJust ext_ixfun, substs, mem) ensureArrayIn :: - ( Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore) + ( Allocable fromrep torep, + Allocator torep (AllocM fromrep torep) ) => Space -> SubExp -> - WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp + WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp ensureArrayIn _ (Constant v) = error $ "ensureArrayIn: " ++ pretty v ++ " cannot be an array." ensureArrayIn space (Var v) = do @@ -623,12 +626,12 @@ return sub_exp ensureDirectArray :: - ( Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore) + ( Allocable fromrep torep, + Allocator torep (AllocM fromrep torep) ) => Maybe Space -> VName -> - AllocM fromlore tolore (VName, SubExp) + AllocM fromrep torep (VName, SubExp) ensureDirectArray space_ok v = do (mem, ixfun) <- lookupArraySummary v mem_space <- lookupMemSpace mem @@ -643,11 +646,11 @@ allocLinearArray space (baseString v) v allocLinearArray :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => Space -> String -> VName -> - AllocM fromlore tolore (VName, SubExp) + AllocM fromrep torep (VName, SubExp) allocLinearArray space s v = do t <- lookupType v case t of @@ -662,11 +665,11 @@ error $ "allocLinearArray: " ++ pretty t funcallArgs :: - ( Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore) + ( Allocable fromrep torep, + Allocator torep (AllocM fromrep torep) ) => [(SubExp, Diet)] -> - AllocM fromlore tolore [(SubExp, Diet)] + AllocM fromrep torep [(SubExp, Diet)] funcallArgs args = do (valargs, (ctx_args, mem_and_size_args)) <- runWriterT $ forM args $ \(arg, d) -> do @@ -677,13 +680,13 @@ return $ map (,Observe) (ctx_args <> mem_and_size_args) <> valargs linearFuncallArg :: - ( Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore) + ( Allocable fromrep torep, + Allocator torep (AllocM fromrep torep) ) => Type -> Space -> SubExp -> - WriterT ([SubExp], [SubExp]) (AllocM fromlore tolore) SubExp + WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp linearFuncallArg Array {} space (Var v) = do (mem, arg') <- lift $ ensureDirectArray (Just space) v tell ([], [Var mem]) @@ -692,12 +695,12 @@ return arg explicitAllocationsGeneric :: - ( Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore) + ( Allocable fromrep torep, + Allocator torep (AllocM fromrep torep) ) => - (Op fromlore -> AllocM fromlore tolore (Op tolore)) -> - (Exp tolore -> AllocM fromlore tolore [ExpHint]) -> - Pass fromlore tolore + (Op fromrep -> AllocM fromrep torep (Op torep)) -> + (Exp torep -> AllocM fromrep torep [ExpHint]) -> + Pass fromrep torep explicitAllocationsGeneric handleOp hints = Pass "explicit allocations" "Transform program to explicit memory representation" $ intraproceduralTransformationWithConsts onStms allocInFun @@ -717,13 +720,13 @@ explicitAllocationsInStmsGeneric :: ( MonadFreshNames m, - HasScope tolore m, - Allocable fromlore tolore + HasScope torep m, + Allocable fromrep torep ) => - (Op fromlore -> AllocM fromlore tolore (Op tolore)) -> - (Exp tolore -> AllocM fromlore tolore [ExpHint]) -> - Stms fromlore -> - m (Stms tolore) + (Op fromrep -> AllocM fromrep torep (Op torep)) -> + (Exp torep -> AllocM fromrep torep [ExpHint]) -> + Stms fromrep -> + m (Stms torep) explicitAllocationsInStmsGeneric handleOp hints stms = do scope <- askScope runAllocM handleOp hints $ @@ -749,9 +752,9 @@ startOfFreeIDRange = S.size . shapeContext bodyReturnMemCtx :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => SubExp -> - AllocM fromlore tolore [SubExp] + AllocM fromrep torep [SubExp] bodyReturnMemCtx Constant {} = return [] bodyReturnMemCtx (Var v) = do @@ -763,10 +766,10 @@ MemArray _ _ _ (ArrayIn mem _) -> return [Var mem] allocInFunBody :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => [Maybe Space] -> - Body fromlore -> - AllocM fromlore tolore (Body tolore) + Body fromrep -> + AllocM fromrep torep (Body torep) allocInFunBody space_oks (Body _ bnds res) = buildBody_ . allocInStms bnds $ do res' <- zipWithM ensureDirect space_oks' res @@ -778,10 +781,10 @@ space_oks' = replicate (length res - num_vals) Nothing ++ space_oks ensureDirect :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => Maybe Space -> SubExp -> - AllocM fromlore tolore SubExp + AllocM fromrep torep SubExp ensureDirect space_ok se = do se_info <- subExpMemInfo se case (se_info, se) of @@ -792,10 +795,10 @@ return se allocInStms :: - (Allocable fromlore tolore) => - Stms fromlore -> - AllocM fromlore tolore a -> - AllocM fromlore tolore a + (Allocable fromrep torep) => + Stms fromrep -> + AllocM fromrep torep a -> + AllocM fromrep torep a allocInStms origstms m = allocInStms' $ stmsToList origstms where allocInStms' [] = m @@ -812,9 +815,9 @@ local f $ allocInStms' stms allocInStm :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Stm fromlore -> - AllocM fromlore tolore () + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => + Stm fromrep -> + AllocM fromrep torep () allocInStm (Let (Pattern sizeElems valElems) _ e) = do e' <- allocInExp e let sizeidents = map patElemIdent sizeElems @@ -823,18 +826,18 @@ addStm bnd allocInLambda :: - Allocable fromlore tolore => - [LParam tolore] -> - Body fromlore -> - AllocM fromlore tolore (Lambda tolore) + Allocable fromrep torep => + [LParam torep] -> + Body fromrep -> + AllocM fromrep torep (Lambda torep) allocInLambda params body = mkLambda params . allocInStms (bodyStms body) $ pure $ bodyResult body allocInExp :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => - Exp fromlore -> - AllocM fromlore tolore (Exp tolore) + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => + Exp fromrep -> + AllocM fromrep torep (Exp torep) allocInExp (DoLoop ctx val form (Body () bodybnds bodyres)) = allocInMergeParams ctx $ \_ ctxparams' _ -> allocInMergeParams val $ @@ -920,10 +923,10 @@ selectSub f (Just (ixfn, m)) = Just (ixfn, map f m) selectSub _ Nothing = Nothing allocInIfBody :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => Int -> - Body fromlore -> - AllocM fromlore tolore (Body tolore, [Maybe IxFun]) + Body fromrep -> + AllocM fromrep torep (Body torep, [Maybe IxFun]) allocInIfBody num_vals (Body _ bnds res) = buildBody . allocInStms bnds $ do let (_, val_res) = splitFromEnd num_vals res @@ -995,9 +998,9 @@ } subExpIxFun :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => SubExp -> - AllocM fromlore tolore (Maybe IxFun) + AllocM fromrep torep (Maybe IxFun) subExpIxFun Constant {} = return Nothing subExpIxFun (Var v) = do info <- lookupMemInfo v @@ -1006,12 +1009,12 @@ _ -> return Nothing addResCtxInIfBody :: - (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + (Allocable fromrep torep, Allocator torep (AllocM fromrep torep)) => [ExtType] -> - Body tolore -> + Body torep -> [Maybe Space] -> [Maybe (ExtIxFun, [TPrimExp Int64 VName])] -> - AllocM fromlore tolore (Body tolore, [BodyReturns]) + AllocM fromrep torep (Body torep, [BodyReturns]) addResCtxInIfBody ifrets (Body _ bnds res) spaces substs = do let num_vals = length ifrets (ctx_res, val_res) = splitFromEnd num_vals res @@ -1090,9 +1093,9 @@ adjustExtPE k = fmap (adjustExtV k) mkSpaceOks :: - (Mem tolore, LocalScope tolore m) => + (Mem torep, LocalScope torep m) => Int -> - Body tolore -> + Body torep -> m [Maybe Space] mkSpaceOks num_vals (Body _ stms res) = inScopeOf stms $ @@ -1110,11 +1113,11 @@ mkSpaceOK _ = return Nothing allocInLoopForm :: - ( Allocable fromlore tolore, - Allocator tolore (AllocM fromlore tolore) + ( Allocable fromrep torep, + Allocator torep (AllocM fromrep torep) ) => - LoopForm fromlore -> - AllocM fromlore tolore (LoopForm tolore) + LoopForm fromrep -> + AllocM fromrep torep (LoopForm torep) allocInLoopForm (WhileLoop v) = return $ WhileLoop v allocInLoopForm (ForLoop i it n loopvars) = ForLoop i it n <$> mapM allocInLoopVar loopvars @@ -1150,41 +1153,41 @@ opIsConst (Inner op) = opIsConst op opIsConst _ = False -sizeSubst :: SizeSubst (Op lore) => Stm lore -> ChunkMap +sizeSubst :: SizeSubst (Op rep) => Stm rep -> ChunkMap sizeSubst (Let pat _ (Op op)) = opSizeSubst pat op sizeSubst _ = mempty -stmConsts :: SizeSubst (Op lore) => Stm lore -> S.Set VName +stmConsts :: SizeSubst (Op rep) => Stm rep -> S.Set VName stmConsts (Let pat _ (Op op)) | opIsConst op = S.fromList $ patternNames pat stmConsts _ = mempty mkLetNamesB' :: - ( Op (Lore m) ~ MemOp inner, + ( Op (Rep m) ~ MemOp inner, MonadBinder m, - ExpDec (Lore m) ~ (), - Allocator (Lore m) (PatAllocM (Lore m)) + ExpDec (Rep m) ~ (), + Allocator (Rep m) (PatAllocM (Rep m)) ) => - ExpDec (Lore m) -> + ExpDec (Rep m) -> [VName] -> - Exp (Lore m) -> - m (Stm (Lore m)) + Exp (Rep m) -> + m (Stm (Rep m)) mkLetNamesB' dec names e = do scope <- askScope pat <- bindPatternWithAllocations scope names e return $ Let pat (defAux dec) e mkLetNamesB'' :: - ( Op (Lore m) ~ MemOp inner, - ExpDec lore ~ (), - HasScope (Engine.Wise lore) m, - Allocator lore (PatAllocM lore), + ( Op (Rep m) ~ MemOp inner, + ExpDec rep ~ (), + HasScope (Engine.Wise rep) m, + Allocator rep (PatAllocM rep), MonadBinder m, - Engine.CanBeWise (Op lore) + Engine.CanBeWise (Op rep) ) => [VName] -> - Exp (Engine.Wise lore) -> - m (Stm (Engine.Wise lore)) + Exp (Engine.Wise rep) -> + m (Stm (Engine.Wise rep)) mkLetNamesB'' names e = do scope <- Engine.removeScopeWisdom <$> askScope (pat, prestms) <- runPatAllocM (patternWithAllocations names $ Engine.removeExpWisdom e) scope @@ -1194,15 +1197,15 @@ return $ Let pat' (defAux dec) e simplifiable :: - ( Engine.SimplifiableLore lore, - ExpDec lore ~ (), - BodyDec lore ~ (), - Op lore ~ MemOp inner, - Allocator lore (PatAllocM lore) + ( Engine.SimplifiableRep rep, + ExpDec rep ~ (), + BodyDec rep ~ (), + Op rep ~ MemOp inner, + Allocator rep (PatAllocM rep) ) => (Engine.OpWithWisdom inner -> UT.UsageTable) -> - (inner -> Engine.SimpleM lore (Engine.OpWithWisdom inner, Stms (Engine.Wise lore))) -> - SimpleOps lore + (inner -> Engine.SimpleM rep (Engine.OpWithWisdom inner, Stms (Engine.Wise rep))) -> + SimpleOps rep simplifiable innerUsage simplifyInnerOp = SimpleOps mkExpDecS' mkBodyS' protectOp opUsage simplifyOp where @@ -1235,14 +1238,14 @@ bindPatternWithAllocations :: ( MonadBinder m, - ExpDec lore ~ (), - Op (Lore m) ~ MemOp inner, - Allocator lore (PatAllocM lore) + ExpDec rep ~ (), + Op (Rep m) ~ MemOp inner, + Allocator rep (PatAllocM rep) ) => - Scope lore -> + Scope rep -> [VName] -> - Exp lore -> - m (Pattern lore) + Exp rep -> + m (Pattern rep) bindPatternWithAllocations types names e = do (pat, prebnds) <- runPatAllocM (patternWithAllocations names e) types mapM_ bindAllocStm prebnds @@ -1252,5 +1255,5 @@ = NoHint | Hint IxFun Space -defaultExpHints :: (Monad m, ASTLore lore) => Exp lore -> m [ExpHint] +defaultExpHints :: (Monad m, ASTRep rep) => Exp rep -> m [ExpHint] defaultExpHints e = return $ replicate (expExtTypeSize e) NoHint diff --git a/src/Futhark/Pass/ExplicitAllocations/GPU.hs b/src/Futhark/Pass/ExplicitAllocations/GPU.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExplicitAllocations/GPU.hs @@ -0,0 +1,192 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +-- | Facilities for converting a 'GPU' program to 'GPUMem'. +module Futhark.Pass.ExplicitAllocations.GPU + ( explicitAllocations, + explicitAllocationsInStms, + ) +where + +import qualified Data.Map as M +import qualified Data.Set as S +import Futhark.IR.GPU +import Futhark.IR.GPUMem +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.Pass.ExplicitAllocations +import Futhark.Pass.ExplicitAllocations.SegOp + +instance SizeSubst (HostOp rep op) where + opSizeSubst (Pattern _ [size]) (SizeOp (SplitSpace _ _ _ elems_per_thread)) = + M.singleton (patElemName size) elems_per_thread + opSizeSubst _ _ = mempty + + opIsConst (SizeOp GetSize {}) = True + opIsConst (SizeOp GetSizeMax {}) = True + opIsConst _ = False + +allocAtLevel :: SegLevel -> AllocM fromrep trep a -> AllocM fromrep trep a +allocAtLevel lvl = local $ \env -> + env + { allocSpace = space, + aggressiveReuse = True + } + where + space = case lvl of + SegThread {} -> DefaultSpace + SegGroup {} -> Space "local" + +handleSegOp :: + SegOp SegLevel GPU -> + AllocM GPU GPUMem (SegOp SegLevel GPUMem) +handleSegOp op = do + num_threads <- + letSubExp "num_threads" $ + BasicOp $ + BinOp + (Mul Int64 OverflowUndef) + (unCount (segNumGroups lvl)) + (unCount (segGroupSize lvl)) + allocAtLevel lvl $ mapSegOpM (mapper num_threads) op + where + scope = scopeOfSegSpace $ segSpace op + lvl = segLevel op + mapper num_threads = + identitySegOpMapper + { mapOnSegOpBody = + localScope scope . local f . allocInKernelBody, + mapOnSegOpLambda = + local inThread + . allocInBinOpLambda num_threads (segSpace op) + } + f = case segLevel op of + SegThread {} -> inThread + SegGroup {} -> inGroup + inThread env = env {envExpHints = inThreadExpHints} + inGroup env = env {envExpHints = inGroupExpHints} + +handleHostOp :: + HostOp GPU (SOAC GPU) -> + AllocM GPU GPUMem (MemOp (HostOp GPUMem ())) +handleHostOp (SizeOp op) = + return $ Inner $ SizeOp op +handleHostOp (OtherOp op) = + error $ "Cannot allocate memory in SOAC: " ++ pretty op +handleHostOp (SegOp op) = + Inner . SegOp <$> handleSegOp op + +kernelExpHints :: Allocator GPUMem m => Exp GPUMem -> m [ExpHint] +kernelExpHints (BasicOp (Manifest perm v)) = do + dims <- arrayDims <$> lookupType v + let perm_inv = rearrangeInverse perm + dims' = rearrangeShape perm dims + ixfun = IxFun.permute (IxFun.iota $ map pe64 dims') perm_inv + return [Hint ixfun DefaultSpace] +kernelExpHints (Op (Inner (SegOp (SegMap lvl@SegThread {} space ts body)))) = + zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body +kernelExpHints (Op (Inner (SegOp (SegRed lvl@SegThread {} space reds ts body)))) = + (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint lvl space) (drop num_reds ts) map_res + where + num_reds = segBinOpResults reds + (red_res, map_res) = splitAt num_reds $ kernelBodyResult body +kernelExpHints e = + return $ replicate (expExtTypeSize e) NoHint + +mapResultHint :: + Allocator rep m => + SegLevel -> + SegSpace -> + Type -> + KernelResult -> + m ExpHint +mapResultHint lvl space = hint + where + num_threads = + pe64 (unCount $ segNumGroups lvl) * pe64 (unCount $ segGroupSize lvl) + + -- Heuristic: do not rearrange for returned arrays that are + -- sufficiently small. + coalesceReturnOfShape _ [] = False + coalesceReturnOfShape bs [Constant (IntValue (Int64Value d))] = bs * d > 4 + coalesceReturnOfShape _ _ = True + + hint t Returns {} + | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t = do + let space_dims = segSpaceDims space + t_dims <- mapM dimAllocationSize $ arrayDims t + return $ Hint (innermost space_dims t_dims) DefaultSpace + hint t (ConcatReturns SplitStrided {} w _ _) = do + t_dims <- mapM dimAllocationSize $ arrayDims t + return $ Hint (innermost [w] t_dims) DefaultSpace + hint Prim {} (ConcatReturns SplitContiguous w elems_per_thread _) = do + let ixfun_base = IxFun.iota [sExt64 num_threads, pe64 elems_per_thread] + ixfun_tr = IxFun.permute ixfun_base [1, 0] + ixfun = IxFun.reshape ixfun_tr $ map (DimNew . pe64) [w] + return $ Hint ixfun DefaultSpace + hint _ _ = return NoHint + +innermost :: [SubExp] -> [SubExp] -> IxFun +innermost space_dims t_dims = + let r = length t_dims + dims = space_dims ++ t_dims + perm = + [length space_dims .. length space_dims + r -1] + ++ [0 .. length space_dims -1] + perm_inv = rearrangeInverse perm + dims_perm = rearrangeShape perm dims + ixfun_base = IxFun.iota $ map pe64 dims_perm + ixfun_rearranged = IxFun.permute ixfun_base perm_inv + in ixfun_rearranged + +semiStatic :: S.Set VName -> SubExp -> Bool +semiStatic _ Constant {} = True +semiStatic consts (Var v) = v `S.member` consts + +inGroupExpHints :: Allocator GPUMem m => Exp GPUMem -> m [ExpHint] +inGroupExpHints (Op (Inner (SegOp (SegMap _ space ts body)))) + | any private $ kernelBodyResult body = do + consts <- askConsts + return $ do + (t, r) <- zip ts $ kernelBodyResult body + return $ + if private r && all (semiStatic consts) (arrayDims t) + then + let seg_dims = map pe64 $ segSpaceDims space + dims = seg_dims ++ map pe64 (arrayDims t) + nilSlice d = DimSlice 0 d 0 + in Hint + ( IxFun.slice (IxFun.iota dims) $ + fullSliceNum dims $ map nilSlice seg_dims + ) + $ ScalarSpace (arrayDims t) $ elemType t + else NoHint + where + private (Returns ResultPrivate _) = True + private _ = False +inGroupExpHints e = return $ replicate (expExtTypeSize e) NoHint + +inThreadExpHints :: Allocator GPUMem m => Exp GPUMem -> m [ExpHint] +inThreadExpHints e = do + consts <- askConsts + mapM (maybePrivate consts) =<< expExtType e + where + maybePrivate consts t + | Just (Array pt shape _) <- hasStaticShape t, + all (semiStatic consts) $ shapeDims shape = do + let ixfun = IxFun.iota $ map pe64 $ shapeDims shape + return $ Hint ixfun $ ScalarSpace (shapeDims shape) pt + | otherwise = + return NoHint + +-- | The pass from 'GPU' to 'GPUMem'. +explicitAllocations :: Pass GPU GPUMem +explicitAllocations = explicitAllocationsGeneric handleHostOp kernelExpHints + +-- | Convert some 'GPU' stms to 'GPUMem'. +explicitAllocationsInStms :: + (MonadFreshNames m, HasScope GPUMem m) => + Stms GPU -> + m (Stms GPUMem) +explicitAllocationsInStms = explicitAllocationsInStmsGeneric handleHostOp kernelExpHints diff --git a/src/Futhark/Pass/ExplicitAllocations/Kernels.hs b/src/Futhark/Pass/ExplicitAllocations/Kernels.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExplicitAllocations/Kernels.hs +++ /dev/null @@ -1,192 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeFamilies #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} - --- | Facilities for converting a 'Kernels' program to 'KernelsMem'. -module Futhark.Pass.ExplicitAllocations.Kernels - ( explicitAllocations, - explicitAllocationsInStms, - ) -where - -import qualified Data.Map as M -import qualified Data.Set as S -import Futhark.IR.Kernels -import Futhark.IR.KernelsMem -import qualified Futhark.IR.Mem.IxFun as IxFun -import Futhark.Pass.ExplicitAllocations -import Futhark.Pass.ExplicitAllocations.SegOp - -instance SizeSubst (HostOp lore op) where - opSizeSubst (Pattern _ [size]) (SizeOp (SplitSpace _ _ _ elems_per_thread)) = - M.singleton (patElemName size) elems_per_thread - opSizeSubst _ _ = mempty - - opIsConst (SizeOp GetSize {}) = True - opIsConst (SizeOp GetSizeMax {}) = True - opIsConst _ = False - -allocAtLevel :: SegLevel -> AllocM fromlore tlore a -> AllocM fromlore tlore a -allocAtLevel lvl = local $ \env -> - env - { allocSpace = space, - aggressiveReuse = True - } - where - space = case lvl of - SegThread {} -> DefaultSpace - SegGroup {} -> Space "local" - -handleSegOp :: - SegOp SegLevel Kernels -> - AllocM Kernels KernelsMem (SegOp SegLevel KernelsMem) -handleSegOp op = do - num_threads <- - letSubExp "num_threads" $ - BasicOp $ - BinOp - (Mul Int64 OverflowUndef) - (unCount (segNumGroups lvl)) - (unCount (segGroupSize lvl)) - allocAtLevel lvl $ mapSegOpM (mapper num_threads) op - where - scope = scopeOfSegSpace $ segSpace op - lvl = segLevel op - mapper num_threads = - identitySegOpMapper - { mapOnSegOpBody = - localScope scope . local f . allocInKernelBody, - mapOnSegOpLambda = - local inThread - . allocInBinOpLambda num_threads (segSpace op) - } - f = case segLevel op of - SegThread {} -> inThread - SegGroup {} -> inGroup - inThread env = env {envExpHints = inThreadExpHints} - inGroup env = env {envExpHints = inGroupExpHints} - -handleHostOp :: - HostOp Kernels (SOAC Kernels) -> - AllocM Kernels KernelsMem (MemOp (HostOp KernelsMem ())) -handleHostOp (SizeOp op) = - return $ Inner $ SizeOp op -handleHostOp (OtherOp op) = - error $ "Cannot allocate memory in SOAC: " ++ pretty op -handleHostOp (SegOp op) = - Inner . SegOp <$> handleSegOp op - -kernelExpHints :: Allocator KernelsMem m => Exp KernelsMem -> m [ExpHint] -kernelExpHints (BasicOp (Manifest perm v)) = do - dims <- arrayDims <$> lookupType v - let perm_inv = rearrangeInverse perm - dims' = rearrangeShape perm dims - ixfun = IxFun.permute (IxFun.iota $ map pe64 dims') perm_inv - return [Hint ixfun DefaultSpace] -kernelExpHints (Op (Inner (SegOp (SegMap lvl@SegThread {} space ts body)))) = - zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body -kernelExpHints (Op (Inner (SegOp (SegRed lvl@SegThread {} space reds ts body)))) = - (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint lvl space) (drop num_reds ts) map_res - where - num_reds = segBinOpResults reds - (red_res, map_res) = splitAt num_reds $ kernelBodyResult body -kernelExpHints e = - return $ replicate (expExtTypeSize e) NoHint - -mapResultHint :: - Allocator lore m => - SegLevel -> - SegSpace -> - Type -> - KernelResult -> - m ExpHint -mapResultHint lvl space = hint - where - num_threads = - pe64 (unCount $ segNumGroups lvl) * pe64 (unCount $ segGroupSize lvl) - - -- Heuristic: do not rearrange for returned arrays that are - -- sufficiently small. - coalesceReturnOfShape _ [] = False - coalesceReturnOfShape bs [Constant (IntValue (Int64Value d))] = bs * d > 4 - coalesceReturnOfShape _ _ = True - - hint t Returns {} - | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t = do - let space_dims = segSpaceDims space - t_dims <- mapM dimAllocationSize $ arrayDims t - return $ Hint (innermost space_dims t_dims) DefaultSpace - hint t (ConcatReturns SplitStrided {} w _ _) = do - t_dims <- mapM dimAllocationSize $ arrayDims t - return $ Hint (innermost [w] t_dims) DefaultSpace - hint Prim {} (ConcatReturns SplitContiguous w elems_per_thread _) = do - let ixfun_base = IxFun.iota [sExt64 num_threads, pe64 elems_per_thread] - ixfun_tr = IxFun.permute ixfun_base [1, 0] - ixfun = IxFun.reshape ixfun_tr $ map (DimNew . pe64) [w] - return $ Hint ixfun DefaultSpace - hint _ _ = return NoHint - -innermost :: [SubExp] -> [SubExp] -> IxFun -innermost space_dims t_dims = - let r = length t_dims - dims = space_dims ++ t_dims - perm = - [length space_dims .. length space_dims + r -1] - ++ [0 .. length space_dims -1] - perm_inv = rearrangeInverse perm - dims_perm = rearrangeShape perm dims - ixfun_base = IxFun.iota $ map pe64 dims_perm - ixfun_rearranged = IxFun.permute ixfun_base perm_inv - in ixfun_rearranged - -semiStatic :: S.Set VName -> SubExp -> Bool -semiStatic _ Constant {} = True -semiStatic consts (Var v) = v `S.member` consts - -inGroupExpHints :: Allocator KernelsMem m => Exp KernelsMem -> m [ExpHint] -inGroupExpHints (Op (Inner (SegOp (SegMap _ space ts body)))) - | any private $ kernelBodyResult body = do - consts <- askConsts - return $ do - (t, r) <- zip ts $ kernelBodyResult body - return $ - if private r && all (semiStatic consts) (arrayDims t) - then - let seg_dims = map pe64 $ segSpaceDims space - dims = seg_dims ++ map pe64 (arrayDims t) - nilSlice d = DimSlice 0 d 0 - in Hint - ( IxFun.slice (IxFun.iota dims) $ - fullSliceNum dims $ map nilSlice seg_dims - ) - $ ScalarSpace (arrayDims t) $ elemType t - else NoHint - where - private (Returns ResultPrivate _) = True - private _ = False -inGroupExpHints e = return $ replicate (expExtTypeSize e) NoHint - -inThreadExpHints :: Allocator KernelsMem m => Exp KernelsMem -> m [ExpHint] -inThreadExpHints e = do - consts <- askConsts - mapM (maybePrivate consts) =<< expExtType e - where - maybePrivate consts t - | Just (Array pt shape _) <- hasStaticShape t, - all (semiStatic consts) $ shapeDims shape = do - let ixfun = IxFun.iota $ map pe64 $ shapeDims shape - return $ Hint ixfun $ ScalarSpace (shapeDims shape) pt - | otherwise = - return NoHint - --- | The pass from 'Kernels' to 'KernelsMem'. -explicitAllocations :: Pass Kernels KernelsMem -explicitAllocations = explicitAllocationsGeneric handleHostOp kernelExpHints - --- | Convert some 'Kernels' stms to 'KernelsMem'. -explicitAllocationsInStms :: - (MonadFreshNames m, HasScope KernelsMem m) => - Stms Kernels -> - m (Stms KernelsMem) -explicitAllocationsInStms = explicitAllocationsInStmsGeneric handleHostOp kernelExpHints diff --git a/src/Futhark/Pass/ExplicitAllocations/MC.hs b/src/Futhark/Pass/ExplicitAllocations/MC.hs --- a/src/Futhark/Pass/ExplicitAllocations/MC.hs +++ b/src/Futhark/Pass/ExplicitAllocations/MC.hs @@ -10,7 +10,7 @@ import Futhark.Pass.ExplicitAllocations import Futhark.Pass.ExplicitAllocations.SegOp -instance SizeSubst (MCOp lore op) where +instance SizeSubst (MCOp rep op) where opSizeSubst _ _ = mempty handleSegOp :: SegOp () MC -> AllocM MC MCMem (SegOp () MCMem) diff --git a/src/Futhark/Pass/ExplicitAllocations/SegOp.hs b/src/Futhark/Pass/ExplicitAllocations/SegOp.hs --- a/src/Futhark/Pass/ExplicitAllocations/SegOp.hs +++ b/src/Futhark/Pass/ExplicitAllocations/SegOp.hs @@ -9,38 +9,38 @@ ) where -import Futhark.IR.KernelsMem +import Futhark.IR.GPUMem import qualified Futhark.IR.Mem.IxFun as IxFun import Futhark.Pass.ExplicitAllocations -instance SizeSubst (SegOp lvl lore) where +instance SizeSubst (SegOp lvl rep) where opSizeSubst _ _ = mempty allocInKernelBody :: - Allocable fromlore tolore => - KernelBody fromlore -> - AllocM fromlore tolore (KernelBody tolore) + Allocable fromrep torep => + KernelBody fromrep -> + AllocM fromrep torep (KernelBody torep) allocInKernelBody (KernelBody () stms res) = uncurry (flip (KernelBody ())) <$> collectStms (allocInStms stms (pure res)) allocInLambda :: - Allocable fromlore tolore => - [LParam tolore] -> - Body fromlore -> - AllocM fromlore tolore (Lambda tolore) + Allocable fromrep torep => + [LParam torep] -> + Body fromrep -> + AllocM fromrep torep (Lambda torep) allocInLambda params body = mkLambda params . allocInStms (bodyStms body) $ pure $ bodyResult body allocInBinOpParams :: - Allocable fromlore tolore => + Allocable fromrep torep => SubExp -> TPrimExp Int64 VName -> TPrimExp Int64 VName -> - [LParam fromlore] -> - [LParam fromlore] -> - AllocM fromlore tolore ([LParam tolore], [LParam tolore]) + [LParam fromrep] -> + [LParam fromrep] -> + AllocM fromrep torep ([LParam torep], [LParam torep]) allocInBinOpParams num_threads my_id other_id xs ys = unzip <$> zipWithM alloc xs ys where alloc x y = @@ -83,11 +83,11 @@ ) allocInBinOpLambda :: - Allocable fromlore tolore => + Allocable fromrep torep => SubExp -> SegSpace -> - Lambda fromlore -> - AllocM fromlore tolore (Lambda tolore) + Lambda fromrep -> + AllocM fromrep torep (Lambda torep) allocInBinOpLambda num_threads (SegSpace flat _) lam = do let (acc_params, arr_params) = splitAt (length (lambdaParams lam) `div` 2) $ lambdaParams lam diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs --- a/src/Futhark/Pass/ExtractKernels.hs +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -165,8 +165,8 @@ import Control.Monad.Reader import Data.Bifunctor (first) import Data.Maybe -import qualified Futhark.IR.Kernels as Out -import Futhark.IR.Kernels.Kernel +import qualified Futhark.IR.GPU as Out +import Futhark.IR.GPU.Kernel import Futhark.IR.SOACS import Futhark.IR.SOACS.Simplify (simplifyStms) import Futhark.MonadFreshNames @@ -177,7 +177,7 @@ import Futhark.Pass.ExtractKernels.ISRWIM import Futhark.Pass.ExtractKernels.Intragroup import Futhark.Pass.ExtractKernels.StreamKernel -import Futhark.Pass.ExtractKernels.ToKernels +import Futhark.Pass.ExtractKernels.ToGPU import Futhark.Tools import qualified Futhark.Transform.FirstOrderTransform as FOT import Futhark.Transform.Rename @@ -186,7 +186,7 @@ -- | Transform a program using SOACs to a program using explicit -- kernels, using the kernel extraction transformation. -extractKernels :: Pass SOACS Out.Kernels +extractKernels :: Pass SOACS Out.GPU extractKernels = Pass { passName = "extract kernels", @@ -194,7 +194,7 @@ passFunction = transformProg } -transformProg :: Prog SOACS -> PassM (Prog Out.Kernels) +transformProg :: Prog SOACS -> PassM (Prog Out.GPU) transformProg (Prog consts funs) = do consts' <- runDistribM $ transformStms mempty $ stmsToList consts funs' <- mapM (transformFunDef $ scopeOf consts') funs @@ -208,13 +208,13 @@ stateThresholdCounter :: Int } -newtype DistribM a = DistribM (RWS (Scope Out.Kernels) Log State a) +newtype DistribM a = DistribM (RWS (Scope Out.GPU) Log State a) deriving ( Functor, Applicative, Monad, - HasScope Out.Kernels, - LocalScope Out.Kernels, + HasScope Out.GPU, + LocalScope Out.GPU, MonadState State, MonadLogger ) @@ -236,23 +236,23 @@ transformFunDef :: (MonadFreshNames m, MonadLogger m) => - Scope Out.Kernels -> + Scope Out.GPU -> FunDef SOACS -> - m (Out.FunDef Out.Kernels) + m (Out.FunDef Out.GPU) transformFunDef scope (FunDef entry attrs name rettype params body) = runDistribM $ do body' <- localScope (scope <> scopeOfFParams params) $ transformBody mempty body return $ FunDef entry attrs name rettype params body' -type KernelsStms = Stms Out.Kernels +type GPUStms = Stms Out.GPU -transformBody :: KernelPath -> Body -> DistribM (Out.Body Out.Kernels) +transformBody :: KernelPath -> Body -> DistribM (Out.Body Out.GPU) transformBody path body = do bnds <- transformStms path $ stmsToList $ bodyStms body return $ mkBody bnds $ bodyResult body -transformStms :: KernelPath -> [Stm] -> DistribM KernelsStms +transformStms :: KernelPath -> [Stm] -> DistribM GPUStms transformStms _ [] = return mempty transformStms path (bnd : bnds) = @@ -308,7 +308,7 @@ String -> Out.SizeClass -> [SubExp] -> - DistribM ((SubExp, Name), Out.Stms Out.Kernels) + DistribM ((SubExp, Name), Out.Stms Out.GPU) cmpSizeLe desc size_class to_what = do x <- gets stateThresholdCounter modify $ \s -> s {stateThresholdCounter = x + 1} @@ -321,11 +321,11 @@ return (cmp_res, size_key) kernelAlternatives :: - (MonadFreshNames m, HasScope Out.Kernels m) => - Out.Pattern Out.Kernels -> - Out.Body Out.Kernels -> - [(SubExp, Out.Body Out.Kernels)] -> - m (Out.Stms Out.Kernels) + (MonadFreshNames m, HasScope Out.GPU m) => + Out.Pattern Out.GPU -> + Out.Body Out.GPU -> + [(SubExp, Out.Body Out.GPU)] -> + m (Out.Stms Out.GPU) kernelAlternatives pat default_body [] = runBinder_ $ do ses <- bodyBind default_body forM_ (zip (patternNames pat) ses) $ \(name, se) -> @@ -343,13 +343,13 @@ If cond alt alt_body $ IfDec (staticShapes (patternTypes pat)) IfEquiv -transformLambda :: KernelPath -> Lambda -> DistribM (Out.Lambda Out.Kernels) +transformLambda :: KernelPath -> Lambda -> DistribM (Out.Lambda Out.GPU) transformLambda path (Lambda params body ret) = Lambda params <$> localScope (scopeOfLParams params) (transformBody path body) <*> pure ret -transformStm :: KernelPath -> Stm -> DistribM KernelsStms +transformStm :: KernelPath -> Stm -> DistribM GPUStms transformStm _ stm | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) = runBinder_ $ FOT.transformStmRecursively stm @@ -366,7 +366,7 @@ <$> (WithAcc (map transformInput inputs) <$> transformLambda path lam) where transformInput (shape, arrs, op) = - (shape, arrs, fmap (first soacsLambdaToKernels) op) + (shape, arrs, fmap (first soacsLambdaToGPU) op) transformStm path (Let pat aux (DoLoop ctx val form body)) = localScope ( castScope (scopeOf form) @@ -392,9 +392,9 @@ | Just (scans, map_lam) <- isScanomapSOAC form = runBinder_ $ do scan_ops <- forM scans $ \(Scan scan_lam nes) -> do (scan_lam', nes', shape) <- determineReduceOp scan_lam nes - let scan_lam'' = soacsLambdaToKernels scan_lam' + let scan_lam'' = soacsLambdaToGPU scan_lam' return $ SegBinOp Noncommutative scan_lam'' nes' shape - let map_lam_sequential = soacsLambdaToKernels map_lam + let map_lam_sequential = soacsLambdaToGPU map_lam lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt addStms . fmap (certify cs) =<< segScan lvl res_pat w scan_ops map_lam_sequential arrs [] [] @@ -415,9 +415,9 @@ let comm' | commutativeLambda red_lam' = Commutative | otherwise = comm - red_lam'' = soacsLambdaToKernels red_lam' + red_lam'' = soacsLambdaToGPU red_lam' return $ SegBinOp comm' red_lam'' nes' shape - let map_lam_sequential = soacsLambdaToKernels map_lam + let map_lam_sequential = soacsLambdaToGPU map_lam lvl <- segThreadCapped [w] "segred" $ NoRecommendation SegNoVirt addStms . fmap (certify cs) =<< nonSegRed lvl pat w red_ops map_lam_sequential arrs @@ -482,7 +482,7 @@ | not $ all primType $ lambdaReturnType red_fun = do -- Split into a chunked map and a reduction, with the latter -- further transformed. - let fold_fun' = soacsLambdaToKernels fold_fun + let fold_fun' = soacsLambdaToGPU fold_fun let (red_pat_elems, concat_pat_elems) = splitAt (length nes) $ patternValueElements pat @@ -509,8 +509,8 @@ Op (Screma num_threads red_results reduce_soac) ) | otherwise = do - let red_fun_sequential = soacsLambdaToKernels red_fun - fold_fun_sequential = soacsLambdaToKernels fold_fun + let red_fun_sequential = soacsLambdaToGPU red_fun + fold_fun_sequential = soacsLambdaToGPU fold_fun fmap (certify cs) <$> streamRed segThreadCapped @@ -551,7 +551,7 @@ transformStms path . stmsToList . snd =<< runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types transformStm _ (Let pat (StmAux cs _ _) (Op (Scatter w lam ivs as))) = runBinder_ $ do - let lam' = soacsLambdaToKernels lam + let lam' = soacsLambdaToGPU lam write_i <- newVName "write_i" let (as_ws, _, _) = unzip3 as kstms = bodyStms $ lambdaBody lam' @@ -574,7 +574,7 @@ addStms stms letBind pat $ Op $ SegOp kernel transformStm _ (Let orig_pat (StmAux cs _ _) (Op (Hist w ops bucket_fun imgs))) = do - let bfun' = soacsLambdaToKernels bucket_fun + let bfun' = soacsLambdaToGPU bucket_fun -- It is important not to launch unnecessarily many threads for -- histograms, because it may mean we unnecessarily need to reduce @@ -583,7 +583,7 @@ lvl <- segThreadCapped [w] "seghist" $ NoRecommendation SegNoVirt addStms =<< histKernel onLambda lvl orig_pat [] [] cs w ops bfun' imgs where - onLambda = pure . soacsLambdaToKernels + onLambda = pure . soacsLambdaToGPU transformStm _ bnd = runBinder_ $ FOT.transformStmRecursively bnd @@ -592,7 +592,7 @@ [SubExp] -> KernelPath -> Maybe Int64 -> - DistribM ((SubExp, Name), Out.Stms Out.Kernels) + DistribM ((SubExp, Name), Out.Stms Out.GPU) sufficientParallelism desc ws path def = cmpSizeLe desc (Out.SizeThreshold path def) ws @@ -654,6 +654,8 @@ 0 -- Basically a map. | DoLoop _ _ ForLoop {} body <- stmExp stm = bodyInterest body * 10 + | WithAcc _ withacc_lam <- stmExp stm = + bodyInterest (lambdaBody withacc_lam) | Op (Screma _ _ form@(ScremaForm _ _ lam')) <- stmExp stm = 1 + bodyInterest (lambdaBody lam') + @@ -672,11 +674,11 @@ onTopLevelStms :: KernelPath -> Stms SOACS -> - DistNestT Out.Kernels DistribM KernelsStms + DistNestT Out.GPU DistribM GPUStms onTopLevelStms path stms = liftInner $ transformStms path $ stmsToList stms -onMap :: KernelPath -> MapLoop -> DistribM KernelsStms +onMap :: KernelPath -> MapLoop -> DistribM GPUStms onMap path (MapLoop pat aux w lam arrs) = do types <- askScope let loopnest = MapNesting pat aux w $ zip (lambdaParams lam) arrs @@ -685,20 +687,20 @@ { distNest = singleNesting (Nesting mempty loopnest), distScope = scopeOfPattern pat - <> scopeForKernels (scopeOf lam) + <> scopeForGPU (scopeOf lam) <> types, distOnInnerMap = onInnerMap path', distOnTopLevelStms = onTopLevelStms path', distSegLevel = segThreadCapped, - distOnSOACSStms = pure . oneStm . soacsStmToKernels, - distOnSOACSLambda = pure . soacsLambdaToKernels + distOnSOACSStms = pure . oneStm . soacsStmToGPU, + distOnSOACSLambda = pure . soacsLambdaToGPU } exploitInnerParallelism path' = runDistNestT (env path') $ distributeMapBodyStms acc (bodyStms $ lambdaBody lam) let exploitOuterParallelism path' = do - let lam' = soacsLambdaToKernels lam + let lam' = soacsLambdaToGPU lam runDistNestT (env path') $ distribute $ addStmsToAcc (bodyStms $ lambdaBody lam') acc @@ -736,11 +738,11 @@ onMap' :: KernelNest -> KernelPath -> - (KernelPath -> DistribM (Out.Stms Out.Kernels)) -> - (KernelPath -> DistribM (Out.Stms Out.Kernels)) -> + (KernelPath -> DistribM (Out.Stms Out.GPU)) -> + (KernelPath -> DistribM (Out.Stms Out.GPU)) -> Pattern -> Lambda -> - DistribM (Out.Stms Out.Kernels) + DistribM (Out.Stms Out.GPU) onMap' loopnest path mk_seq_stms mk_par_stms pat lam = do -- Some of the control flow here looks a bit convoluted because we -- are trying to avoid generating unneeded threshold parameters, @@ -865,8 +867,8 @@ onInnerMap :: KernelPath -> MapLoop -> - DistAcc Out.Kernels -> - DistNestT Out.Kernels DistribM (DistAcc Out.Kernels) + DistAcc Out.GPU -> + DistNestT Out.GPU DistribM (DistAcc Out.GPU) onInnerMap path maploop@(MapLoop pat aux w lam arrs) acc | unbalancedLambda lam, lambdaContainsParallelism lam = @@ -920,7 +922,7 @@ -- versioning does not take place down that branch (it currently -- does not). (sequentialised_kernel, nestw_bnds) <- localScope extra_scope $ do - let sequentialised_lam = soacsLambdaToKernels lam' + let sequentialised_lam = soacsLambdaToGPU lam' constructKernel segThreadCapped nest' $ lambdaBody sequentialised_lam let outer_pat = loopNestingPattern $ fst nest diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs --- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs @@ -3,7 +3,7 @@ {-# LANGUAGE TypeFamilies #-} module Futhark.Pass.ExtractKernels.BlockedKernel - ( DistLore, + ( DistRep, MkSegLevel, ThreadRecommendation (..), segRed, @@ -31,32 +31,32 @@ import Prelude hiding (quot) -- | Constraints pertinent to performing distribution/flattening. -type DistLore lore = - ( Bindable lore, - HasSegOp lore, - BinderOps lore, - LetDec lore ~ Type, - ExpDec lore ~ (), - BodyDec lore ~ (), - CanBeAliased (Op lore) +type DistRep rep = + ( Bindable rep, + HasSegOp rep, + BinderOps rep, + LetDec rep ~ Type, + ExpDec rep ~ (), + BodyDec rep ~ (), + CanBeAliased (Op rep) ) data ThreadRecommendation = ManyThreads | NoRecommendation SegVirt -type MkSegLevel lore m = - [SubExp] -> String -> ThreadRecommendation -> BinderT lore m (SegOpLevel lore) +type MkSegLevel rep m = + [SubExp] -> String -> ThreadRecommendation -> BinderT rep m (SegOpLevel rep) mkSegSpace :: MonadFreshNames m => [(VName, SubExp)] -> m SegSpace mkSegSpace dims = SegSpace <$> newVName "phys_tid" <*> pure dims prepareRedOrScan :: - (MonadBinder m, DistLore (Lore m)) => + (MonadBinder m, DistRep (Rep m)) => SubExp -> - Lambda (Lore m) -> + Lambda (Rep m) -> [VName] -> [(VName, SubExp)] -> [KernelInput] -> - m (SegSpace, KernelBody (Lore m)) + m (SegSpace, KernelBody (Rep m)) prepareRedOrScan w map_lam arrs ispace inps = do gtid <- newVName "gtid" space <- mkSegSpace $ ispace ++ [(gtid, w)] @@ -64,25 +64,24 @@ runBinder $ localScope (scopeOfSegSpace space) $ do mapM_ readKernelInput inps - forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do - arr_t <- lookupType arr - letBindNames [paramName p] $ - BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] + mapM_ readKernelInput $ do + (p, arr) <- zip (lambdaParams map_lam) arrs + pure $ KernelInput (paramName p) (paramType p) arr [Var gtid] map (Returns ResultMaySimplify) <$> bodyBind (lambdaBody map_lam) return (space, kbody) segRed :: - (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore -> - Pattern lore -> + (MonadFreshNames m, DistRep rep, HasScope rep m) => + SegOpLevel rep -> + Pattern rep -> SubExp -> -- segment size - [SegBinOp lore] -> - Lambda lore -> + [SegBinOp rep] -> + Lambda rep -> [VName] -> [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this reduction [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace - m (Stms lore) + m (Stms rep) segRed lvl pat w ops map_lam arrs ispace inps = runBinder_ $ do (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps letBind pat $ @@ -91,16 +90,16 @@ SegRed lvl kspace ops (lambdaReturnType map_lam) kbody segScan :: - (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore -> - Pattern lore -> + (MonadFreshNames m, DistRep rep, HasScope rep m) => + SegOpLevel rep -> + Pattern rep -> SubExp -> -- segment size - [SegBinOp lore] -> - Lambda lore -> + [SegBinOp rep] -> + Lambda rep -> [VName] -> [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this scan [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace - m (Stms lore) + m (Stms rep) segScan lvl pat w ops map_lam arrs ispace inps = runBinder_ $ do (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps letBind pat $ @@ -109,15 +108,15 @@ SegScan lvl kspace ops (lambdaReturnType map_lam) kbody segMap :: - (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore -> - Pattern lore -> + (MonadFreshNames m, DistRep rep, HasScope rep m) => + SegOpLevel rep -> + Pattern rep -> SubExp -> -- segment size - Lambda lore -> + Lambda rep -> [VName] -> [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this map [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace - m (Stms lore) + m (Stms rep) segMap lvl pat w map_lam arrs ispace inps = runBinder_ $ do (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps letBind pat $ @@ -126,9 +125,9 @@ SegMap lvl kspace (lambdaReturnType map_lam) kbody dummyDim :: - (MonadFreshNames m, MonadBinder m, DistLore (Lore m)) => - Pattern (Lore m) -> - m (Pattern (Lore m), [(VName, SubExp)], m ()) + (MonadFreshNames m, MonadBinder m, DistRep (Rep m)) => + Pattern (Rep m) -> + m (Pattern (Rep m), [(VName, SubExp)], m ()) dummyDim pat = do -- We add a unit-size segment on top to ensure that the result -- of the SegRed is an array, which we then immediately index. @@ -153,31 +152,31 @@ ) nonSegRed :: - (MonadFreshNames m, DistLore lore, HasScope lore m) => - SegOpLevel lore -> - Pattern lore -> + (MonadFreshNames m, DistRep rep, HasScope rep m) => + SegOpLevel rep -> + Pattern rep -> SubExp -> - [SegBinOp lore] -> - Lambda lore -> + [SegBinOp rep] -> + Lambda rep -> [VName] -> - m (Stms lore) + m (Stms rep) nonSegRed lvl pat w ops map_lam arrs = runBinder_ $ do (pat', ispace, read_dummy) <- dummyDim pat addStms =<< segRed lvl pat' w ops map_lam arrs ispace [] read_dummy segHist :: - (DistLore lore, MonadFreshNames m, HasScope lore m) => - SegOpLevel lore -> - Pattern lore -> + (DistRep rep, MonadFreshNames m, HasScope rep m) => + SegOpLevel rep -> + Pattern rep -> SubExp -> -- | Segment indexes and sizes. [(VName, SubExp)] -> [KernelInput] -> - [HistOp lore] -> - Lambda lore -> + [HistOp rep] -> + Lambda rep -> [VName] -> - m (Stms lore) + m (Stms rep) segHist lvl pat arr_w ispace inps ops lam arrs = runBinder_ $ do gtid <- newVName "gtid" space <- mkSegSpace $ ispace ++ [(gtid, arr_w)] @@ -195,10 +194,10 @@ letBind pat $ Op $ segOp $ SegHist lvl space ops (lambdaReturnType lam) kbody mapKernelSkeleton :: - (DistLore lore, HasScope lore m, MonadFreshNames m) => + (DistRep rep, HasScope rep m, MonadFreshNames m) => [(VName, SubExp)] -> [KernelInput] -> - m (SegSpace, Stms lore) + m (SegSpace, Stms rep) mapKernelSkeleton ispace inputs = do read_input_bnds <- runBinder_ $ mapM readKernelInput inputs @@ -206,13 +205,13 @@ return (space, read_input_bnds) mapKernel :: - (DistLore lore, HasScope lore m, MonadFreshNames m) => - MkSegLevel lore m -> + (DistRep rep, HasScope rep m, MonadFreshNames m) => + MkSegLevel rep m -> [(VName, SubExp)] -> [KernelInput] -> [Type] -> - KernelBody lore -> - m (SegOp (SegOpLevel lore) lore, Stms lore) + KernelBody rep -> + m (SegOp (SegOpLevel rep) rep, Stms rep) mapKernel mk_lvl ispace inputs rts (KernelBody () kstms krets) = runBinderT' $ do (space, read_input_stms) <- mapKernelSkeleton ispace inputs @@ -238,7 +237,7 @@ deriving (Show) readKernelInput :: - (DistLore (Lore m), MonadBinder m) => + (DistRep (Rep m), MonadBinder m) => KernelInput -> m () readKernelInput inp = do diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs --- a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs +++ b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs @@ -62,7 +62,7 @@ import Futhark.Util import Futhark.Util.Log -scopeForSOACs :: SameScope lore SOACS => Scope lore -> Scope SOACS +scopeForSOACs :: SameScope rep SOACS => Scope rep -> Scope SOACS scopeForSOACs = castScope data MapLoop = MapLoop SOACS.Pattern (StmAux ()) SubExp SOACS.Lambda [VName] @@ -71,80 +71,80 @@ mapLoopStm (MapLoop pat aux w lam arrs) = Let pat aux $ Op $ Screma w arrs $ mapSOAC lam -data DistEnv lore m = DistEnv +data DistEnv rep m = DistEnv { distNest :: Nestings, - distScope :: Scope lore, - distOnTopLevelStms :: Stms SOACS -> DistNestT lore m (Stms lore), + distScope :: Scope rep, + distOnTopLevelStms :: Stms SOACS -> DistNestT rep m (Stms rep), distOnInnerMap :: MapLoop -> - DistAcc lore -> - DistNestT lore m (DistAcc lore), - distOnSOACSStms :: Stm SOACS -> Binder lore (Stms lore), - distOnSOACSLambda :: Lambda SOACS -> Binder lore (Lambda lore), - distSegLevel :: MkSegLevel lore m + DistAcc rep -> + DistNestT rep m (DistAcc rep), + distOnSOACSStms :: Stm SOACS -> Binder rep (Stms rep), + distOnSOACSLambda :: Lambda SOACS -> Binder rep (Lambda rep), + distSegLevel :: MkSegLevel rep m } -data DistAcc lore = DistAcc +data DistAcc rep = DistAcc { distTargets :: Targets, - distStms :: Stms lore + distStms :: Stms rep } -data DistRes lore = DistRes - { accPostStms :: PostStms lore, +data DistRes rep = DistRes + { accPostStms :: PostStms rep, accLog :: Log } -instance Semigroup (DistRes lore) where +instance Semigroup (DistRes rep) where DistRes ks1 log1 <> DistRes ks2 log2 = DistRes (ks1 <> ks2) (log1 <> log2) -instance Monoid (DistRes lore) where +instance Monoid (DistRes rep) where mempty = DistRes mempty mempty -newtype PostStms lore = PostStms {unPostStms :: Stms lore} +newtype PostStms rep = PostStms {unPostStms :: Stms rep} -instance Semigroup (PostStms lore) where +instance Semigroup (PostStms rep) where PostStms xs <> PostStms ys = PostStms $ ys <> xs -instance Monoid (PostStms lore) where +instance Monoid (PostStms rep) where mempty = PostStms mempty -typeEnvFromDistAcc :: DistLore lore => DistAcc lore -> Scope lore +typeEnvFromDistAcc :: DistRep rep => DistAcc rep -> Scope rep typeEnvFromDistAcc = scopeOfPattern . fst . outerTarget . distTargets -addStmsToAcc :: Stms lore -> DistAcc lore -> DistAcc lore +addStmsToAcc :: Stms rep -> DistAcc rep -> DistAcc rep addStmsToAcc stms acc = acc {distStms = stms <> distStms acc} addStmToAcc :: - (MonadFreshNames m, DistLore lore) => + (MonadFreshNames m, DistRep rep) => Stm SOACS -> - DistAcc lore -> - DistNestT lore m (DistAcc lore) + DistAcc rep -> + DistNestT rep m (DistAcc rep) addStmToAcc stm acc = do onSoacs <- asks distOnSOACSStms (stm', _) <- runBinder $ onSoacs stm return acc {distStms = stm' <> distStms acc} soacsLambda :: - (MonadFreshNames m, DistLore lore) => + (MonadFreshNames m, DistRep rep) => Lambda SOACS -> - DistNestT lore m (Lambda lore) + DistNestT rep m (Lambda rep) soacsLambda lam = do onLambda <- asks distOnSOACSLambda fst <$> runBinder (onLambda lam) -newtype DistNestT lore m a - = DistNestT (ReaderT (DistEnv lore m) (WriterT (DistRes lore) m) a) +newtype DistNestT rep m a + = DistNestT (ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a) deriving ( Functor, Applicative, Monad, - MonadReader (DistEnv lore m), - MonadWriter (DistRes lore) + MonadReader (DistEnv rep m), + MonadWriter (DistRes rep) ) -liftInner :: (LocalScope lore m, DistLore lore) => m a -> DistNestT lore m a +liftInner :: (LocalScope rep m, DistRep rep) => m a -> DistNestT rep m a liftInner m = do outer_scope <- askScope DistNestT $ @@ -153,25 +153,25 @@ inner_scope <- askScope localScope (outer_scope `M.difference` inner_scope) m -instance MonadFreshNames m => MonadFreshNames (DistNestT lore m) where +instance MonadFreshNames m => MonadFreshNames (DistNestT rep m) where getNameSource = DistNestT $ lift getNameSource putNameSource = DistNestT . lift . putNameSource -instance (Monad m, ASTLore lore) => HasScope lore (DistNestT lore m) where +instance (Monad m, ASTRep rep) => HasScope rep (DistNestT rep m) where askScope = asks distScope -instance (Monad m, ASTLore lore) => LocalScope lore (DistNestT lore m) where +instance (Monad m, ASTRep rep) => LocalScope rep (DistNestT rep m) where localScope types = local $ \env -> env {distScope = types <> distScope env} -instance Monad m => MonadLogger (DistNestT lore m) where +instance Monad m => MonadLogger (DistNestT rep m) where addLog msgs = tell mempty {accLog = msgs} runDistNestT :: - (MonadLogger m, DistLore lore) => - DistEnv lore m -> - DistNestT lore m (DistAcc lore) -> - m (Stms lore) + (MonadLogger m, DistRep rep) => + DistEnv rep m -> + DistNestT rep m (DistAcc rep) -> + m (Stms rep) runDistNestT env (DistNestT m) = do (acc, res) <- runWriterT $ runReaderT m env addLog $ accLog res @@ -200,17 +200,17 @@ BasicOp $ Replicate (Shape [loopNestingWidth outermost]) se -addPostStms :: Monad m => PostStms lore -> DistNestT lore m () +addPostStms :: Monad m => PostStms rep -> DistNestT rep m () addPostStms ks = tell $ mempty {accPostStms = ks} -postStm :: Monad m => Stms lore -> DistNestT lore m () +postStm :: Monad m => Stms rep -> DistNestT rep m () postStm stms = addPostStms $ PostStms stms withStm :: - (Monad m, DistLore lore) => + (Monad m, DistRep rep) => Stm SOACS -> - DistNestT lore m a -> - DistNestT lore m a + DistNestT rep m a -> + DistNestT rep m a withStm stm = local $ \env -> env { distScope = @@ -223,9 +223,9 @@ provided = namesFromList $ patternNames $ stmPattern stm leavingNesting :: - (MonadFreshNames m, DistLore lore) => - DistAcc lore -> - DistNestT lore m (DistAcc lore) + (MonadFreshNames m, DistRep rep) => + DistAcc rep -> + DistNestT rep m (DistAcc rep) leavingNesting acc = case popInnerTarget $ distTargets acc of Nothing -> @@ -278,14 +278,14 @@ return $ acc {distTargets = newtargets, distStms = stms} mapNesting :: - (MonadFreshNames m, DistLore lore) => + (MonadFreshNames m, DistRep rep) => PatternT Type -> StmAux () -> SubExp -> Lambda SOACS -> [VName] -> - DistNestT lore m (DistAcc lore) -> - DistNestT lore m (DistAcc lore) + DistNestT rep m (DistAcc rep) -> + DistNestT rep m (DistAcc rep) mapNesting pat aux w lam arrs m = local extend $ leavingNesting =<< m where @@ -300,10 +300,10 @@ } inNesting :: - (Monad m, DistLore lore) => + (Monad m, DistRep rep) => KernelNest -> - DistNestT lore m a -> - DistNestT lore m a + DistNestT rep m a -> + DistNestT rep m a inNesting (outer, nests) = local $ \env -> env { distNest = (inner, nests'), @@ -334,10 +334,10 @@ lambdaContainsParallelism = bodyContainsParallelism . lambdaBody distributeMapBodyStms :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> + (MonadFreshNames m, LocalScope rep m, DistRep rep) => + DistAcc rep -> Stms SOACS -> - DistNestT lore m (DistAcc lore) + DistNestT rep m (DistAcc rep) distributeMapBodyStms orig_acc = distribute <=< onStms orig_acc . stmsToList where onStms acc [] = return acc @@ -354,21 +354,21 @@ -- situation that stm is in scope of itself. withStm stm $ maybeDistributeStm stm =<< onStms acc stms -onInnerMap :: Monad m => MapLoop -> DistAcc lore -> DistNestT lore m (DistAcc lore) +onInnerMap :: Monad m => MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep) onInnerMap loop acc = do f <- asks distOnInnerMap f loop acc -onTopLevelStms :: Monad m => Stms SOACS -> DistNestT lore m () +onTopLevelStms :: Monad m => Stms SOACS -> DistNestT rep m () onTopLevelStms stms = do f <- asks distOnTopLevelStms postStm =<< f stms maybeDistributeStm :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => Stm SOACS -> - DistAcc lore -> - DistNestT lore m (DistAcc lore) + DistAcc rep -> + DistNestT rep m (DistAcc rep) maybeDistributeStm stm acc | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) = addStmToAcc stm acc @@ -650,12 +650,12 @@ addStmToAcc bnd acc distributeSingleUnaryStm :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> + (MonadFreshNames m, LocalScope rep m, DistRep rep) => + DistAcc rep -> Stm SOACS -> VName -> - (KernelNest -> PatternT Type -> VName -> DistNestT lore m (Stms lore)) -> - DistNestT lore m (DistAcc lore) + (KernelNest -> PatternT Type -> VName -> DistNestT rep m (Stms rep)) -> + DistNestT rep m (DistAcc rep) distributeSingleUnaryStm acc stm stm_arr f = distributeSingleStm acc stm >>= \case Just (kernels, res, nest, acc') @@ -682,15 +682,15 @@ False distribute :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> - DistNestT lore m (DistAcc lore) + (MonadFreshNames m, LocalScope rep m, DistRep rep) => + DistAcc rep -> + DistNestT rep m (DistAcc rep) distribute acc = fromMaybe acc <$> distributeIfPossible acc mkSegLevel :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistNestT lore m (MkSegLevel lore (DistNestT lore m)) + (MonadFreshNames m, LocalScope rep m, DistRep rep) => + DistNestT rep m (MkSegLevel rep (DistNestT rep m)) mkSegLevel = do mk_lvl <- asks distSegLevel return $ \w desc r -> do @@ -699,9 +699,9 @@ return lvl distributeIfPossible :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> - DistNestT lore m (Maybe (DistAcc lore)) + (MonadFreshNames m, LocalScope rep m, DistRep rep) => + DistAcc rep -> + DistNestT rep m (Maybe (DistAcc rep)) distributeIfPossible acc = do nest <- asks distNest mk_lvl <- mkSegLevel @@ -717,17 +717,17 @@ } distributeSingleStm :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => - DistAcc lore -> + (MonadFreshNames m, LocalScope rep m, DistRep rep) => + DistAcc rep -> Stm SOACS -> DistNestT - lore + rep m ( Maybe - ( PostStms lore, + ( PostStms rep, Result, KernelNest, - DistAcc lore + DistAcc rep ) ) distributeSingleStm acc bnd = do @@ -751,16 +751,16 @@ ) segmentedScatterKernel :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => KernelNest -> [Int] -> PatternT Type -> Certificates -> SubExp -> - Lambda lore -> + Lambda rep -> [VName] -> [(Shape, Int, VName)] -> - DistNestT lore m (Stms lore) + DistNestT rep m (Stms rep) segmentedScatterKernel nest perm scatter_pat cs scatter_w lam ivs dests = do -- We replicate some of the checking done by 'isSegmentedOp', but -- things are different because a scatter is not a reduction or @@ -833,14 +833,14 @@ (gtids, ws) = unzip ispace segmentedUpdateKernel :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => KernelNest -> [Int] -> Certificates -> VName -> Slice SubExp -> VName -> - DistNestT lore m (Stms lore) + DistNestT rep m (Stms rep) segmentedUpdateKernel nest perm cs arr slice v = do (base_ispace, kernel_inps) <- flatKernel nest let slice_dims = sliceDims slice @@ -889,12 +889,12 @@ letBind pat $ Op $ segOp k segmentedGatherKernel :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => KernelNest -> Certificates -> VName -> Slice SubExp -> - DistNestT lore m (Stms lore) + DistNestT rep m (Stms rep) segmentedGatherKernel nest cs arr slice = do let slice_dims = sliceDims slice slice_gtids <- replicateM (length slice_dims) (newVName "gtid_slice") @@ -925,15 +925,15 @@ letBind pat $ Op $ segOp k segmentedHistKernel :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => KernelNest -> [Int] -> Certificates -> SubExp -> [SOACS.HistOp SOACS] -> - Lambda lore -> + Lambda rep -> [VName] -> - DistNestT lore m (Stms lore) + DistNestT rep m (Stms rep) segmentedHistKernel nest perm cs hist_w ops lam arrs = do -- We replicate some of the checking done by 'isSegmentedOp', but -- things are different because a Hist is not a reduction or @@ -970,18 +970,18 @@ bad = error "Ill-typed nested Hist encountered." histKernel :: - (MonadBinder m, DistLore (Lore m)) => - (Lambda SOACS -> m (Lambda (Lore m))) -> - SegOpLevel (Lore m) -> + (MonadBinder m, DistRep (Rep m)) => + (Lambda SOACS -> m (Lambda (Rep m))) -> + SegOpLevel (Rep m) -> PatternT Type -> [(VName, SubExp)] -> [KernelInput] -> Certificates -> SubExp -> [SOACS.HistOp SOACS] -> - Lambda (Lore m) -> + Lambda (Rep m) -> [VName] -> - m (Stms (Lore m)) + m (Stms (Rep m)) histKernel onLambda lvl orig_pat ispace inputs cs hist_w ops lam arrs = runBinderT'_ $ do ops' <- forM ops $ \(SOACS.HistOp num_bins rf dests nes op) -> do (op', nes', shape) <- determineReduceOp op nes @@ -1029,15 +1029,15 @@ | otherwise = (mempty, lam) segmentedScanomapKernel :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => KernelNest -> [Int] -> SubExp -> - Lambda lore -> - Lambda lore -> + Lambda rep -> + Lambda rep -> [SubExp] -> [VName] -> - DistNestT lore m (Maybe (Stms lore)) + DistNestT rep m (Maybe (Stms rep)) segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = do mk_lvl <- asks distSegLevel isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $ @@ -1048,16 +1048,16 @@ =<< segScan lvl pat segment_size [scan_op] map_lam arrs ispace inps regularSegmentedRedomapKernel :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => KernelNest -> [Int] -> SubExp -> Commutativity -> - Lambda lore -> - Lambda lore -> + Lambda rep -> + Lambda rep -> [SubExp] -> [VName] -> - DistNestT lore m (Maybe (Stms lore)) + DistNestT rep m (Maybe (Stms rep)) regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs = do mk_lvl <- asks distSegLevel isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $ @@ -1068,7 +1068,7 @@ =<< segRed lvl pat segment_size [red_op] map_lam arrs ispace inps isSegmentedOp :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => KernelNest -> [Int] -> Names -> @@ -1080,9 +1080,9 @@ [KernelInput] -> [SubExp] -> [VName] -> - BinderT lore m () + BinderT rep m () ) -> - DistNestT lore m (Maybe (Stms lore)) + DistNestT rep m (Maybe (Stms rep)) isSegmentedOp nest perm free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do -- We must verify that array inputs to the operation are inputs to -- the outermost loop nesting or free in the loop nest. Nothing @@ -1181,14 +1181,14 @@ } kernelOrNot :: - (MonadFreshNames m, DistLore lore) => + (MonadFreshNames m, DistRep rep) => Certificates -> Stm SOACS -> - DistAcc lore -> - PostStms lore -> - DistAcc lore -> - Maybe (Stms lore) -> - DistNestT lore m (DistAcc lore) + DistAcc rep -> + PostStms rep -> + DistAcc rep -> + Maybe (Stms rep) -> + DistNestT rep m (DistAcc rep) kernelOrNot cs bnd acc _ _ Nothing = addStmToAcc (certify cs bnd) acc kernelOrNot cs _ _ kernels acc' (Just bnds) = do @@ -1197,10 +1197,10 @@ return acc' distributeMap :: - (MonadFreshNames m, LocalScope lore m, DistLore lore) => + (MonadFreshNames m, LocalScope rep m, DistRep rep) => MapLoop -> - DistAcc lore -> - DistNestT lore m (DistAcc lore) + DistAcc rep -> + DistNestT rep m (DistAcc rep) distributeMap (MapLoop pat aux w lam arrs) acc = distribute =<< mapNesting diff --git a/src/Futhark/Pass/ExtractKernels/Distribution.hs b/src/Futhark/Pass/ExtractKernels/Distribution.hs --- a/src/Futhark/Pass/ExtractKernels/Distribution.hs +++ b/src/Futhark/Pass/ExtractKernels/Distribution.hs @@ -53,7 +53,7 @@ import Futhark.IR.SegOp import Futhark.MonadFreshNames import Futhark.Pass.ExtractKernels.BlockedKernel - ( DistLore, + ( DistRep, KernelInput (..), MkSegLevel, mapKernel, @@ -112,10 +112,10 @@ x : xs -> Just (t, Targets x $ reverse xs) [] -> Nothing -targetScope :: DistLore lore => Target -> Scope lore +targetScope :: DistRep rep => Target -> Scope rep targetScope = scopeOfPattern . fst -targetsScope :: DistLore lore => Targets -> Scope lore +targetsScope :: DistRep rep => Targets -> Scope rep targetsScope (Targets t ts) = mconcat $ map targetScope $ t : ts data LoopNesting = MapNesting @@ -126,7 +126,7 @@ } deriving (Show) -scopeOfLoopNesting :: DistLore lore => LoopNesting -> Scope lore +scopeOfLoopNesting :: DistRep rep => LoopNesting -> Scope rep scopeOfLoopNesting = scopeOfLParams . map fst . loopNestingParamsAndArrs ppLoopNesting :: LoopNesting -> String @@ -249,11 +249,11 @@ kernelNestWidths = map loopNestingWidth . kernelNestLoops constructKernel :: - (DistLore lore, MonadFreshNames m, LocalScope lore m) => - MkSegLevel lore m -> + (DistRep rep, MonadFreshNames m, LocalScope rep m) => + MkSegLevel rep m -> KernelNest -> - Body lore -> - m (Stm lore, Stms lore) + Body rep -> + m (Stm rep, Stms rep) constructKernel mk_lvl kernel_nest inner_body = runBinderT' $ do (ispace, inps) <- flatKernel kernel_nest let aux = loopNestingAux first_nest @@ -331,9 +331,9 @@ distributionInnerPattern = fst . innerTarget . distributionTarget distributionBodyFromStms :: - ASTLore lore => + ASTRep rep => Targets -> - Stms lore -> + Stms rep -> (DistributionBody, Result) distributionBodyFromStms (Targets (inner_pat, inner_res) targets) stms = let bound_by_stms = namesFromList $ M.keys $ scopeOf stms @@ -349,16 +349,16 @@ ) distributionBodyFromStm :: - ASTLore lore => + ASTRep rep => Targets -> - Stm lore -> + Stm rep -> (DistributionBody, Result) distributionBodyFromStm targets bnd = distributionBodyFromStms targets $ oneStm bnd createKernelNest :: - forall lore m. - (MonadFreshNames m, HasScope lore m) => + forall rep m. + (MonadFreshNames m, HasScope rep m) => Nestings -> DistributionBody -> m (Maybe (Targets, KernelNest)) @@ -551,16 +551,16 @@ in (pat', res', identity_map, expand_target) tryDistribute :: - ( DistLore lore, + ( DistRep rep, MonadFreshNames m, - LocalScope lore m, + LocalScope rep m, MonadLogger m ) => - MkSegLevel lore m -> + MkSegLevel rep m -> Nestings -> Targets -> - Stms lore -> - m (Maybe (Targets, Stms lore)) + Stms rep -> + m (Maybe (Targets, Stms rep)) tryDistribute _ _ targets stms | null stms = -- No point in distributing an empty kernel. @@ -590,10 +590,10 @@ (dist_body, inner_body_res) = distributionBodyFromStms targets stms tryDistributeStm :: - (MonadFreshNames m, HasScope t m, ASTLore lore) => + (MonadFreshNames m, HasScope t m, ASTRep rep) => Nestings -> Targets -> - Stm lore -> + Stm rep -> m (Maybe (Result, Targets, KernelNest)) tryDistributeStm nest targets bnd = fmap addRes <$> createKernelNest nest dist_body diff --git a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs --- a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs +++ b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs @@ -18,7 +18,7 @@ -- | Interchange Scan With Inner Map. Tries to turn a @scan(map)@ into a -- @map(scan) iswim :: - (MonadBinder m, Lore m ~ SOACS) => + (MonadBinder m, Rep m ~ SOACS) => Pattern -> SubExp -> Lambda -> @@ -81,7 +81,7 @@ -- | Interchange Reduce With Inner Map. Tries to turn a @reduce(map)@ into a -- @map(reduce) irwim :: - (MonadBinder m, Lore m ~ SOACS) => + (MonadBinder m, Rep m ~ SOACS) => Pattern -> SubExp -> Commutativity -> diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs --- a/src/Futhark/Pass/ExtractKernels/Interchange.hs +++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs @@ -185,7 +185,7 @@ Let pat (defAux ()) $ WithAcc inputs lam interchangeWithAcc1 :: - (MonadBinder m, Lore m ~ SOACS) => + (MonadBinder m, Rep m ~ SOACS) => WithAccStm -> LoopNesting -> m WithAccStm @@ -193,21 +193,30 @@ (WithAccStm perm _withacc_pat inputs acc_lam) (MapNesting map_pat map_aux w params_and_arrs) = do inputs' <- mapM onInput inputs - let lam_params = lambdaParams acc_lam + lam_params' <- newAccLamParams $ lambdaParams acc_lam iota_p <- newParam "iota_p" $ Prim int64 - acc_lam' <- trLam (Var (paramName iota_p)) <=< mkLambda lam_params $ do + acc_lam' <- trLam (Var (paramName iota_p)) <=< mkLambda lam_params' $ do + let acc_params = drop (length inputs) lam_params' + orig_acc_params = drop (length inputs) $ lambdaParams acc_lam iota_w <- letExp "acc_inter_iota" . BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64 let (params, arrs) = unzip params_and_arrs maplam_ret = lambdaReturnType acc_lam - maplam = Lambda (iota_p : params) (lambdaBody acc_lam) maplam_ret + maplam = Lambda (iota_p : orig_acc_params ++ params) (lambdaBody acc_lam) maplam_ret auxing map_aux . letTupExp' "withacc_inter" $ - Op $ Screma w (iota_w : arrs) (mapSOAC maplam) + Op $ Screma w (iota_w : map paramName acc_params ++ arrs) (mapSOAC maplam) let pat = Pattern [] $ rearrangeShape perm $ patternValueElements map_pat perm' = [0 .. patternSize pat -1] pure $ WithAccStm perm' pat inputs' acc_lam' where + newAccLamParams ps = do + let (cert_ps, acc_ps) = splitAt (length ps `div` 2) ps + -- Should not rename the certificates. + acc_ps' <- forM acc_ps $ \(Param v t) -> + newParam (baseString v) t + pure $ cert_ps <> acc_ps' + num_accs = length inputs acc_certs = map paramName $ take num_accs $ lambdaParams acc_lam onArr v = diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs --- a/src/Futhark/Pass/ExtractKernels/Intragroup.hs +++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs @@ -13,14 +13,14 @@ import qualified Data.Map.Strict as M import qualified Data.Set as S import Futhark.Analysis.PrimExp.Convert -import qualified Futhark.IR.Kernels as Out -import Futhark.IR.Kernels.Kernel hiding (HistOp) +import qualified Futhark.IR.GPU as Out +import Futhark.IR.GPU.Kernel hiding (HistOp) import Futhark.IR.SOACS import Futhark.MonadFreshNames import Futhark.Pass.ExtractKernels.BlockedKernel import Futhark.Pass.ExtractKernels.DistributeNests import Futhark.Pass.ExtractKernels.Distribution -import Futhark.Pass.ExtractKernels.ToKernels +import Futhark.Pass.ExtractKernels.ToGPU import Futhark.Tools import qualified Futhark.Transform.FirstOrderTransform as FOT import Futhark.Util.Log @@ -39,7 +39,7 @@ -- We distinguish between "minimum group size" and "maximum -- exploitable parallelism". intraGroupParallelise :: - (MonadFreshNames m, LocalScope Out.Kernels m) => + (MonadFreshNames m, LocalScope Out.GPU m) => KernelNest -> Lambda -> m @@ -47,8 +47,8 @@ ( (SubExp, SubExp), SubExp, Log, - Out.Stms Out.Kernels, - Out.Stms Out.Kernels + Out.Stms Out.GPU, + Out.Stms Out.GPU ) ) intraGroupParallelise knest lam = runMaybeT $ do @@ -136,7 +136,7 @@ aux = loopNestingAux first_nest readGroupKernelInput :: - (DistLore (Lore m), MonadBinder m) => + (DistRep (Rep m), MonadBinder m) => KernelInput -> m () readGroupKernelInput inp @@ -161,15 +161,15 @@ mempty = IntraAcc mempty mempty mempty type IntraGroupM = - BinderT Out.Kernels (RWS () IntraAcc VNameSource) + BinderT Out.GPU (RWS () IntraAcc VNameSource) instance MonadLogger IntraGroupM where addLog log = tell mempty {accLog = log} runIntraGroupM :: - (MonadFreshNames m, HasScope Out.Kernels m) => + (MonadFreshNames m, HasScope Out.GPU m) => IntraGroupM () -> - m (IntraAcc, Out.Stms Out.Kernels) + m (IntraAcc, Out.Stms Out.GPU) runIntraGroupM m = do scope <- castScope <$> askScope modifyNameSource $ \src -> @@ -184,7 +184,7 @@ accAvailPar = S.singleton ws } -intraGroupBody :: SegLevel -> Body -> IntraGroupM (Out.Body Out.Kernels) +intraGroupBody :: SegLevel -> Body -> IntraGroupM (Out.Body Out.GPU) intraGroupBody lvl body = do stms <- collectStms_ $ intraGroupStms lvl $ bodyStms body return $ mkBody stms $ bodyResult body @@ -223,7 +223,7 @@ singleNesting $ Nesting mempty loopnest, distScope = scopeOfPattern pat - <> scopeForKernels (scopeOf lam) + <> scopeForGPU (scopeOf lam) <> scope, distOnInnerMap = distributeMap, @@ -233,9 +233,9 @@ lift $ parallelMin minw return lvl, distOnSOACSStms = - pure . oneStm . soacsStmToKernels, + pure . oneStm . soacsStmToGPU, distOnSOACSLambda = - pure . soacsLambdaToKernels + pure . soacsLambdaToGPU } acc = DistAcc @@ -248,26 +248,26 @@ Op (Screma w arrs form) | Just (scans, mapfun) <- isScanomapSOAC form, Scan scanfun nes <- singleScan scans -> do - let scanfun' = soacsLambdaToKernels scanfun - mapfun' = soacsLambdaToKernels mapfun + let scanfun' = soacsLambdaToGPU scanfun + mapfun' = soacsLambdaToGPU mapfun certifying (stmAuxCerts aux) $ addStms =<< segScan lvl' pat w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] [] parallelMin [w] Op (Screma w arrs form) | Just (reds, map_lam) <- isRedomapSOAC form, Reduce comm red_lam nes <- singleReduce reds -> do - let red_lam' = soacsLambdaToKernels red_lam - map_lam' = soacsLambdaToKernels map_lam + let red_lam' = soacsLambdaToGPU red_lam + map_lam' = soacsLambdaToGPU map_lam certifying (stmAuxCerts aux) $ addStms =<< segRed lvl' pat w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] [] parallelMin [w] Op (Hist w ops bucket_fun arrs) -> do ops' <- forM ops $ \(HistOp num_bins rf dests nes op) -> do (op', nes', shape) <- determineReduceOp op nes - let op'' = soacsLambdaToKernels op' + let op'' = soacsLambdaToGPU op' return $ Out.HistOp num_bins rf dests nes' shape op'' - let bucket_fun' = soacsLambdaToKernels bucket_fun + let bucket_fun' = soacsLambdaToGPU bucket_fun certifying (stmAuxCerts aux) $ addStms =<< segHist lvl' pat w [] [] ops' bucket_fun' arrs parallelMin [w] @@ -285,7 +285,7 @@ write_i <- newVName "write_i" space <- mkSegSpace [(write_i, w)] - let lam' = soacsLambdaToKernels lam + let lam' = soacsLambdaToGPU lam (dests_ws, _, _) = unzip3 dests krets = do (a_w, a, is_vs) <- @@ -307,16 +307,16 @@ parallelMin [w] _ -> - addStm $ soacsStmToKernels stm + addStm $ soacsStmToGPU stm intraGroupStms :: SegLevel -> Stms SOACS -> IntraGroupM () intraGroupStms lvl = mapM_ (intraGroupStm lvl) intraGroupParalleliseBody :: - (MonadFreshNames m, HasScope Out.Kernels m) => + (MonadFreshNames m, HasScope Out.GPU m) => SegLevel -> Body -> - m ([[SubExp]], [[SubExp]], Log, Out.KernelBody Out.Kernels) + m ([[SubExp]], [[SubExp]], Log, Out.KernelBody Out.GPU) intraGroupParalleliseBody lvl body = do (IntraAcc min_ws avail_ws log, kstms) <- runIntraGroupM $ intraGroupStms lvl $ bodyStms body diff --git a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs --- a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs @@ -14,7 +14,7 @@ import Data.List () import Futhark.Analysis.PrimExp import Futhark.IR -import Futhark.IR.Kernels hiding +import Futhark.IR.GPU hiding ( BasicOp, Body, Exp, @@ -30,7 +30,7 @@ ) import Futhark.MonadFreshNames import Futhark.Pass.ExtractKernels.BlockedKernel -import Futhark.Pass.ExtractKernels.ToKernels +import Futhark.Pass.ExtractKernels.ToGPU import Futhark.Tools import Prelude hiding (quot) @@ -43,7 +43,7 @@ deriving (Eq, Ord, Show) numberOfGroups :: - (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => + (MonadBinder m, Op (Rep m) ~ HostOp (Rep m) inner) => String -> SubExp -> SubExp -> @@ -59,7 +59,7 @@ return (num_groups, num_threads) blockedKernelSize :: - (MonadBinder m, Lore m ~ Kernels) => + (MonadBinder m, Rep m ~ GPU) => String -> SubExp -> m KernelSize @@ -75,7 +75,7 @@ return $ KernelSize per_thread_elements num_threads splitArrays :: - (MonadBinder m, Lore m ~ Kernels) => + (MonadBinder m, Rep m ~ GPU) => VName -> [VName] -> SplitOrdering -> @@ -113,12 +113,12 @@ error "partitionChunkedKernelFoldParameters: lambda takes too few parameters" blockedPerThread :: - (MonadBinder m, Lore m ~ Kernels) => + (MonadBinder m, Rep m ~ GPU) => VName -> SubExp -> KernelSize -> StreamOrd -> - Lambda (Lore m) -> + Lambda (Rep m) -> Int -> [VName] -> m ([PatElemT Type], [PatElemT Type]) @@ -172,8 +172,8 @@ kerneliseLambda :: MonadFreshNames m => [SubExp] -> - Lambda Kernels -> - m (Lambda Kernels) + Lambda GPU -> + m (Lambda GPU) kerneliseLambda nes lam = do thread_index <- newVName "thread_index" let thread_index_param = Param thread_index $ Prim int64 @@ -197,15 +197,15 @@ } prepareStream :: - (MonadBinder m, Lore m ~ Kernels) => + (MonadBinder m, Rep m ~ GPU) => KernelSize -> [(VName, SubExp)] -> SubExp -> Commutativity -> - Lambda Kernels -> + Lambda GPU -> [SubExp] -> [VName] -> - m (SubExp, SegSpace, [Type], KernelBody Kernels) + m (SubExp, SegSpace, [Type], KernelBody GPU) prepareStream size ispace w comm fold_lam nes arrs = do let (KernelSize elems_per_thread num_threads) = size let (ordering, split_ordering) = @@ -235,16 +235,16 @@ return (num_threads, space, ts, kbody) streamRed :: - (MonadFreshNames m, HasScope Kernels m) => - MkSegLevel Kernels m -> - Pattern Kernels -> + (MonadFreshNames m, HasScope GPU m) => + MkSegLevel GPU m -> + Pattern GPU -> SubExp -> Commutativity -> - Lambda Kernels -> - Lambda Kernels -> + Lambda GPU -> + Lambda GPU -> [SubExp] -> [VName] -> - m (Stms Kernels) + m (Stms GPU) streamRed mk_lvl pat w comm red_lam fold_lam nes arrs = runBinderT'_ $ do -- The strategy here is to rephrase the stream reduction as a -- non-segmented SegRed that does explicit chunking within its body. @@ -272,16 +272,16 @@ -- Similar to streamRed, but without the last reduction. streamMap :: - (MonadFreshNames m, HasScope Kernels m) => - MkSegLevel Kernels m -> + (MonadFreshNames m, HasScope GPU m) => + MkSegLevel GPU m -> [String] -> - [PatElem Kernels] -> + [PatElem GPU] -> SubExp -> Commutativity -> - Lambda Kernels -> + Lambda GPU -> [SubExp] -> [VName] -> - m ((SubExp, [VName]), Stms Kernels) + m ((SubExp, [VName]), Stms GPU) streamMap mk_lvl out_desc mapout_pes w comm fold_lam nes arrs = runBinderT' $ do size <- blockedKernelSize "stream_map" w @@ -301,7 +301,7 @@ -- | Like 'segThread', but cap the thread count to the input size. -- This is more efficient for small kernels, e.g. summing a small -- array. -segThreadCapped :: MonadFreshNames m => MkSegLevel Kernels m +segThreadCapped :: MonadFreshNames m => MkSegLevel GPU m segThreadCapped ws desc r = do w <- letSubExp "nest_size" diff --git a/src/Futhark/Pass/ExtractKernels/ToGPU.hs b/src/Futhark/Pass/ExtractKernels/ToGPU.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/ToGPU.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +module Futhark.Pass.ExtractKernels.ToGPU + ( getSize, + segThread, + soacsLambdaToGPU, + soacsStmToGPU, + scopeForGPU, + scopeForSOACs, + injectSOACS, + ) +where + +import Control.Monad.Identity +import Data.List () +import Futhark.Analysis.Rephrase +import Futhark.IR +import Futhark.IR.GPU +import Futhark.IR.SOACS (SOACS) +import qualified Futhark.IR.SOACS.SOAC as SOAC +import Futhark.Tools + +getSize :: + (MonadBinder m, Op (Rep m) ~ HostOp (Rep m) inner) => + String -> + SizeClass -> + m SubExp +getSize desc size_class = do + size_key <- nameFromString . pretty <$> newVName desc + letSubExp desc $ Op $ SizeOp $ GetSize size_key size_class + +segThread :: + (MonadBinder m, Op (Rep m) ~ HostOp (Rep m) inner) => + String -> + m SegLevel +segThread desc = + SegThread + <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups) + <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup) + <*> pure SegVirt + +injectSOACS :: + ( Monad m, + SameScope from to, + ExpDec from ~ ExpDec to, + BodyDec from ~ BodyDec to, + RetType from ~ RetType to, + BranchType from ~ BranchType to, + Op from ~ SOAC from + ) => + (SOAC to -> Op to) -> + Rephraser m from to +injectSOACS f = + Rephraser + { rephraseExpDec = return, + rephraseBodyDec = return, + rephraseLetBoundDec = return, + rephraseFParamDec = return, + rephraseLParamDec = return, + rephraseOp = fmap f . onSOAC, + rephraseRetType = return, + rephraseBranchType = return + } + where + onSOAC = SOAC.mapSOACM mapper + mapper = + SOAC.SOACMapper + { SOAC.mapOnSOACSubExp = return, + SOAC.mapOnSOACVName = return, + SOAC.mapOnSOACLambda = rephraseLambda $ injectSOACS f + } + +soacsStmToGPU :: Stm SOACS -> Stm GPU +soacsStmToGPU = runIdentity . rephraseStm (injectSOACS OtherOp) + +soacsLambdaToGPU :: Lambda SOACS -> Lambda GPU +soacsLambdaToGPU = runIdentity . rephraseLambda (injectSOACS OtherOp) + +scopeForSOACs :: Scope GPU -> Scope SOACS +scopeForSOACs = castScope + +scopeForGPU :: Scope SOACS -> Scope GPU +scopeForGPU = castScope diff --git a/src/Futhark/Pass/ExtractKernels/ToKernels.hs b/src/Futhark/Pass/ExtractKernels/ToKernels.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/ToKernels.hs +++ /dev/null @@ -1,85 +0,0 @@ -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE TypeFamilies #-} - -module Futhark.Pass.ExtractKernels.ToKernels - ( getSize, - segThread, - soacsLambdaToKernels, - soacsStmToKernels, - scopeForKernels, - scopeForSOACs, - injectSOACS, - ) -where - -import Control.Monad.Identity -import Data.List () -import Futhark.Analysis.Rephrase -import Futhark.IR -import Futhark.IR.Kernels -import Futhark.IR.SOACS (SOACS) -import qualified Futhark.IR.SOACS.SOAC as SOAC -import Futhark.Tools - -getSize :: - (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => - String -> - SizeClass -> - m SubExp -getSize desc size_class = do - size_key <- nameFromString . pretty <$> newVName desc - letSubExp desc $ Op $ SizeOp $ GetSize size_key size_class - -segThread :: - (MonadBinder m, Op (Lore m) ~ HostOp (Lore m) inner) => - String -> - m SegLevel -segThread desc = - SegThread - <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups) - <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup) - <*> pure SegVirt - -injectSOACS :: - ( Monad m, - SameScope from to, - ExpDec from ~ ExpDec to, - BodyDec from ~ BodyDec to, - RetType from ~ RetType to, - BranchType from ~ BranchType to, - Op from ~ SOAC from - ) => - (SOAC to -> Op to) -> - Rephraser m from to -injectSOACS f = - Rephraser - { rephraseExpLore = return, - rephraseBodyLore = return, - rephraseLetBoundLore = return, - rephraseFParamLore = return, - rephraseLParamLore = return, - rephraseOp = fmap f . onSOAC, - rephraseRetType = return, - rephraseBranchType = return - } - where - onSOAC = SOAC.mapSOACM mapper - mapper = - SOAC.SOACMapper - { SOAC.mapOnSOACSubExp = return, - SOAC.mapOnSOACVName = return, - SOAC.mapOnSOACLambda = rephraseLambda $ injectSOACS f - } - -soacsStmToKernels :: Stm SOACS -> Stm Kernels -soacsStmToKernels = runIdentity . rephraseStm (injectSOACS OtherOp) - -soacsLambdaToKernels :: Lambda SOACS -> Lambda Kernels -soacsLambdaToKernels = runIdentity . rephraseLambda (injectSOACS OtherOp) - -scopeForSOACs :: Scope Kernels -> Scope SOACS -scopeForSOACs = castScope - -scopeForKernels :: Scope SOACS -> Scope Kernels -scopeForKernels = castScope diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs --- a/src/Futhark/Pass/ExtractMulticore.hs +++ b/src/Futhark/Pass/ExtractMulticore.hs @@ -25,7 +25,7 @@ import qualified Futhark.IR.SOACS.Simplify as SOACS import Futhark.Pass import Futhark.Pass.ExtractKernels.DistributeNests -import Futhark.Pass.ExtractKernels.ToKernels (injectSOACS) +import Futhark.Pass.ExtractKernels.ToGPU (injectSOACS) import Futhark.Tools import qualified Futhark.Transform.FirstOrderTransform as FOT import Futhark.Transform.Rename (Rename, renameSomething) diff --git a/src/Futhark/Pass/FirstOrderTransform.hs b/src/Futhark/Pass/FirstOrderTransform.hs --- a/src/Futhark/Pass/FirstOrderTransform.hs +++ b/src/Futhark/Pass/FirstOrderTransform.hs @@ -24,10 +24,10 @@ import Futhark.IR.SOACS (SOACS, scopeOf) import Futhark.Pass -import Futhark.Transform.FirstOrderTransform (FirstOrderLore, transformConsts, transformFunDef) +import Futhark.Transform.FirstOrderTransform (FirstOrderRep, transformConsts, transformFunDef) -- | The first-order transformation pass. -firstOrderTransform :: FirstOrderLore lore => Pass SOACS lore +firstOrderTransform :: FirstOrderRep rep => Pass SOACS rep firstOrderTransform = Pass "first order transform" diff --git a/src/Futhark/Pass/KernelBabysitting.hs b/src/Futhark/Pass/KernelBabysitting.hs --- a/src/Futhark/Pass/KernelBabysitting.hs +++ b/src/Futhark/Pass/KernelBabysitting.hs @@ -11,7 +11,7 @@ import qualified Data.Map.Strict as M import Data.Maybe import Futhark.IR -import Futhark.IR.Kernels hiding +import Futhark.IR.GPU hiding ( BasicOp, Body, Exp, @@ -31,7 +31,7 @@ import Futhark.Util -- | The pass definition. -babysitKernels :: Pass Kernels Kernels +babysitKernels :: Pass GPU GPU babysitKernels = Pass "babysit kernels" @@ -42,12 +42,12 @@ let m = localScope scope $ transformStms mempty stms fmap fst $ modifyNameSource $ runState (runBinderT m M.empty) -type BabysitM = Binder Kernels +type BabysitM = Binder GPU -transformStms :: ExpMap -> Stms Kernels -> BabysitM (Stms Kernels) +transformStms :: ExpMap -> Stms GPU -> BabysitM (Stms GPU) transformStms expmap stms = collectStms_ $ foldM_ transformStm expmap stms -transformBody :: ExpMap -> Body Kernels -> BabysitM (Body Kernels) +transformBody :: ExpMap -> Body GPU -> BabysitM (Body GPU) transformBody expmap (Body () stms res) = do stms' <- transformStms expmap stms return $ Body () stms' res @@ -57,7 +57,7 @@ -- funky in memory (and we'd prefer it not to be). If we cannot find -- it in the map, we just assume it's all good. HACK and FIXME, I -- suppose. We really should do this at the memory level. -type ExpMap = M.Map VName (Stm Kernels) +type ExpMap = M.Map VName (Stm GPU) nonlinearInMemory :: VName -> ExpMap -> Maybe (Maybe [Int]) nonlinearInMemory name m = @@ -80,7 +80,7 @@ return $ Just $ rearrangeInverse $ [inner_r .. inner_r + outer_r -1] ++ [0 .. inner_r -1] | otherwise = Nothing -transformStm :: ExpMap -> Stm Kernels -> BabysitM ExpMap +transformStm :: ExpMap -> Stm GPU -> BabysitM ExpMap transformStm expmap (Let pat aux (Op (SegOp op))) -- FIXME: We only make coalescing optimisations for SegThread -- SegOps, because that's what the analysis assumes. For SegGroup @@ -102,7 +102,7 @@ addStm bnd' return $ M.fromList [(name, bnd') | name <- patternNames pat] <> expmap -transform :: ExpMap -> Mapper Kernels Kernels BabysitM +transform :: ExpMap -> Mapper GPU GPU BabysitM transform expmap = identityMapper {mapOnBody = \scope -> localScope scope . transformBody expmap} @@ -110,8 +110,8 @@ ExpMap -> SegLevel -> SegSpace -> - KernelBody Kernels -> - BabysitM (KernelBody Kernels) + KernelBody GPU -> + BabysitM (KernelBody GPU) transformKernelBody expmap lvl space kbody = do -- Go spelunking for accesses to arrays that are defined outside the -- kernel body and where the indices are kernel thread indices. @@ -143,7 +143,7 @@ (VName -> Bool) -> -- thread local? (VName -> SubExp -> Bool) -> -- variant to a certain gid (given as first param)? (SubExp -> Maybe SubExp) -> -- split substitution? - Scope Kernels -> -- type environment + Scope GPU -> -- type environment VName -> Slice SubExp -> m (Maybe (VName, Slice SubExp)) @@ -152,10 +152,10 @@ (Applicative f, Monad f) => Names -> Names -> - Scope Kernels -> + Scope GPU -> ArrayIndexTransform f -> - KernelBody Kernels -> - f (KernelBody Kernels) + KernelBody GPU -> + f (KernelBody GPU) traverseKernelBodyArrayIndexes free_ker_vars thread_variant outer_scope f (KernelBody () kstms kres) = KernelBody () . stmsFromList <$> mapM @@ -513,10 +513,10 @@ type VarianceTable = M.Map VName Names -varianceInStms :: VarianceTable -> Stms Kernels -> VarianceTable +varianceInStms :: VarianceTable -> Stms GPU -> VarianceTable varianceInStms t = foldl varianceInStm t . stmsToList -varianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable +varianceInStm :: VarianceTable -> Stm GPU -> VarianceTable varianceInStm variance bnd = foldl' add variance $ patternNames $ stmPattern bnd where diff --git a/src/Futhark/Pass/Simplify.hs b/src/Futhark/Pass/Simplify.hs --- a/src/Futhark/Pass/Simplify.hs +++ b/src/Futhark/Pass/Simplify.hs @@ -5,15 +5,15 @@ simplifySOACS, simplifySeq, simplifyMC, - simplifyKernels, - simplifyKernelsMem, + simplifyGPU, + simplifyGPUMem, simplifySeqMem, simplifyMCMem, ) where -import qualified Futhark.IR.Kernels.Simplify as Kernels -import qualified Futhark.IR.KernelsMem as KernelsMem +import qualified Futhark.IR.GPU.Simplify as GPU +import qualified Futhark.IR.GPUMem as GPUMem import qualified Futhark.IR.MC as MC import qualified Futhark.IR.MCMem as MCMem import qualified Futhark.IR.SOACS.Simplify as SOACS @@ -23,15 +23,15 @@ import Futhark.Pass simplify :: - (Prog lore -> PassM (Prog lore)) -> - Pass lore lore + (Prog rep -> PassM (Prog rep)) -> + Pass rep rep simplify = Pass "simplify" "Perform simple enabling optimisations." simplifySOACS :: Pass SOACS.SOACS SOACS.SOACS simplifySOACS = simplify SOACS.simplifySOACS -simplifyKernels :: Pass Kernels.Kernels Kernels.Kernels -simplifyKernels = simplify Kernels.simplifyKernels +simplifyGPU :: Pass GPU.GPU GPU.GPU +simplifyGPU = simplify GPU.simplifyGPU simplifySeq :: Pass Seq.Seq Seq.Seq simplifySeq = simplify Seq.simplifyProg @@ -39,8 +39,8 @@ simplifyMC :: Pass MC.MC MC.MC simplifyMC = simplify MC.simplifyProg -simplifyKernelsMem :: Pass KernelsMem.KernelsMem KernelsMem.KernelsMem -simplifyKernelsMem = simplify KernelsMem.simplifyProg +simplifyGPUMem :: Pass GPUMem.GPUMem GPUMem.GPUMem +simplifyGPUMem = simplify GPUMem.simplifyProg simplifySeqMem :: Pass SeqMem.SeqMem SeqMem.SeqMem simplifySeqMem = simplify SeqMem.simplifyProg diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -13,8 +13,8 @@ where import Control.Category ((>>>)) -import Futhark.IR.Kernels (Kernels) -import Futhark.IR.KernelsMem (KernelsMem) +import Futhark.IR.GPU (GPU) +import Futhark.IR.GPUMem (GPUMem) import Futhark.IR.MC (MC) import Futhark.IR.MCMem (MCMem) import Futhark.IR.SOACS (SOACS) @@ -25,11 +25,12 @@ import Futhark.Optimise.Fusion import Futhark.Optimise.InPlaceLowering import Futhark.Optimise.InliningDeadFun +import qualified Futhark.Optimise.ReuseAllocations as ReuseAllocations import Futhark.Optimise.Sink import Futhark.Optimise.TileLoops import Futhark.Optimise.Unstream import Futhark.Pass.ExpandAllocations -import qualified Futhark.Pass.ExplicitAllocations.Kernels as Kernels +import qualified Futhark.Pass.ExplicitAllocations.GPU as GPU import qualified Futhark.Pass.ExplicitAllocations.MC as MC import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq import Futhark.Pass.ExtractKernels @@ -57,19 +58,19 @@ removeDeadFunctions ] -kernelsPipeline :: Pipeline SOACS Kernels +kernelsPipeline :: Pipeline SOACS GPU kernelsPipeline = standardPipeline >>> onePass extractKernels >>> passes - [ simplifyKernels, + [ simplifyGPU, babysitKernels, tileLoops, - unstreamKernels, + unstreamGPU, performCSE True, - simplifyKernels, - sinkKernels, - inPlaceLoweringKernels + simplifyGPU, + sinkGPU, + inPlaceLoweringGPU ] sequentialPipeline :: Pipeline SOACS Seq @@ -91,18 +92,20 @@ simplifySeqMem ] -gpuPipeline :: Pipeline SOACS KernelsMem +gpuPipeline :: Pipeline SOACS GPUMem gpuPipeline = kernelsPipeline - >>> onePass Kernels.explicitAllocations + >>> onePass GPU.explicitAllocations >>> passes - [ simplifyKernelsMem, + [ simplifyGPUMem, performCSE False, - simplifyKernelsMem, - doubleBufferKernels, - simplifyKernelsMem, + simplifyGPUMem, + doubleBufferGPU, + simplifyGPUMem, + ReuseAllocations.optimise, + simplifyGPUMem, expandAllocations, - simplifyKernelsMem + simplifyGPUMem ] mcPipeline :: Pipeline SOACS MC diff --git a/src/Futhark/Pipeline.hs b/src/Futhark/Pipeline.hs --- a/src/Futhark/Pipeline.hs +++ b/src/Futhark/Pipeline.hs @@ -37,7 +37,7 @@ import Data.Time.Clock import qualified Futhark.Analysis.Alias as Alias import Futhark.Error -import Futhark.IR (PrettyLore, Prog) +import Futhark.IR (PrettyRep, Prog) import Futhark.MonadFreshNames import Futhark.Pass import Futhark.TypeCheck @@ -103,10 +103,10 @@ newEnv = FutharkEnv verbose -- | A compilation always ends with some kind of action. -data Action lore = Action +data Action rep = Action { actionName :: String, actionDescription :: String, - actionProcedure :: Prog lore -> FutharkM () + actionProcedure :: Prog rep -> FutharkM () } -- | Configuration object for running a compiler pipeline. @@ -118,7 +118,7 @@ -- | A compiler pipeline is conceptually a function from programs to -- programs, where the actual representation may change. Pipelines -- can be composed using their 'Category' instance. -newtype Pipeline fromlore tolore = Pipeline {unPipeline :: PipelineConfig -> Prog fromlore -> FutharkM (Prog tolore)} +newtype Pipeline fromrep torep = Pipeline {unPipeline :: PipelineConfig -> Prog fromrep -> FutharkM (Prog torep)} instance Category Pipeline where id = Pipeline $ const return @@ -129,17 +129,17 @@ -- | Run the pipeline on the given program. runPipeline :: - Pipeline fromlore tolore -> + Pipeline fromrep torep -> PipelineConfig -> - Prog fromlore -> - FutharkM (Prog tolore) + Prog fromrep -> + FutharkM (Prog torep) runPipeline = unPipeline -- | Construct a pipeline from a single compiler pass. onePass :: - Checkable tolore => - Pass fromlore tolore -> - Pipeline fromlore tolore + Checkable torep => + Pass fromrep torep -> + Pipeline fromrep torep onePass pass = Pipeline perform where perform cfg prog = do @@ -156,15 +156,15 @@ -- | Create a pipeline from a list of passes. passes :: - Checkable lore => - [Pass lore lore] -> - Pipeline lore lore + Checkable rep => + [Pass rep rep] -> + Pipeline rep rep passes = foldl (>>>) id . map onePass validationError :: - PrettyLore lore => - Pass fromlore tolore -> - Prog lore -> + PrettyRep rep => + Pass fromrep torep -> + Prog rep -> String -> FutharkM a validationError pass prog err = @@ -173,9 +173,9 @@ msg = "Type error after pass '" <> T.pack (passName pass) <> "':\n" <> T.pack err runPass :: - Pass fromlore tolore -> - Prog fromlore -> - FutharkM (Prog tolore) + Pass fromrep torep -> + Prog fromrep -> + FutharkM (Prog torep) runPass pass prog = do (prog', logged) <- runPassM (passFunction pass prog) verb <- asks $ (>= VeryVerbose) . futharkVerbose diff --git a/src/Futhark/Script.hs b/src/Futhark/Script.hs --- a/src/Futhark/Script.hs +++ b/src/Futhark/Script.hs @@ -34,8 +34,6 @@ where import Control.Monad.Except -import qualified Data.Binary as Bin -import qualified Data.ByteString.Lazy.Char8 as LBS import Data.Char import Data.Foldable (toList) import Data.Functor @@ -46,13 +44,12 @@ import qualified Data.Text as T import Data.Traversable import Data.Void +import qualified Futhark.Data.Parser as V import Futhark.Server +import Futhark.Server.Values (getValue, putValue) import qualified Futhark.Test.Values as V -import qualified Futhark.Test.Values.Parser as V import Futhark.Util (nubOrd) import Futhark.Util.Pretty hiding (float, line, sep, string, (), (<|>)) -import System.IO -import System.IO.Temp import Text.Megaparsec import Text.Megaparsec.Char.Lexer (charLiteral) @@ -69,9 +66,9 @@ -- | Start a server, execute an action, then shut down the server. -- Similar to 'withServer'. -withScriptServer :: FilePath -> [FilePath] -> (ScriptServer -> IO a) -> IO a -withScriptServer prog options f = - withServer prog options $ flip withScriptServer' f +withScriptServer :: ServerCfg -> (ScriptServer -> IO a) -> IO a +withScriptServer cfg f = + withServer cfg $ flip withScriptServer' f -- | A function called in a 'Call' expression can be either a Futhark -- function or a builtin function. @@ -101,7 +98,7 @@ instance Pretty Exp where ppr = pprPrec 0 pprPrec _ (ServerVar _ v) = "$" <> ppr v - pprPrec _ (Const v) = ppr v + pprPrec _ (Const v) = strictText $ V.valueText v pprPrec i (Let pat e1 e2) = parensIf (i > 0) $ "let" <+> pat' <+> equals <+> ppr e1 <+> "in" <+> ppr e2 where @@ -137,11 +134,8 @@ <*> pPattern <* lexeme sep "=" <*> parseExp sep <* lexeme sep "in" <*> parseExp sep, - inParens sep (mkTuple <$> (parseExp sep `sepBy` pComma)), - inBraces sep (Record <$> (pField `sepBy` pComma)), - Call <$> parseFunc <*> many (parseExp sep), - Const <$> V.parseValue sep, - StringLit . T.pack <$> lexeme sep ("\"" *> manyTill charLiteral "\"") + try $ Call <$> parseFunc <*> many pAtom, + pAtom ] "expression" where @@ -151,6 +145,16 @@ mkTuple [v] = v mkTuple vs = Tuple vs + pAtom = + choice + [ try $ inParens sep (mkTuple <$> (parseExp sep `sepBy` pComma)), + inParens sep $ parseExp sep, + inBraces sep (Record <$> (pField `sepBy` pComma)), + Const <$> V.parseValue sep, + StringLit . T.pack <$> lexeme sep ("\"" *> manyTill charLiteral "\""), + Call <$> parseFunc <*> pure [] + ] + pPattern = choice [ inParens sep $ pVarName `sepBy` pComma, @@ -172,35 +176,13 @@ where constituent c = isAlphaNum c || c == '_' -prettyFailure :: CmdFailure -> T.Text -prettyFailure (CmdFailure bef aft) = - T.unlines $ bef ++ aft - readVar :: (MonadError T.Text m, MonadIO m) => Server -> VarName -> m V.Value readVar server v = - either throwError pure <=< liftIO $ - withSystemTempFile "futhark-server-read" $ \tmpf tmpf_h -> do - hClose tmpf_h - store_res <- cmdStore server tmpf [v] - case store_res of - Just err -> pure $ Left $ prettyFailure err - Nothing -> do - s <- LBS.readFile tmpf - case V.readValues s of - Just [val] -> pure $ Right val - Just [] -> pure $ Left "Cannot read opaque value from Futhark server." - _ -> pure $ Left "Invalid data file produced by Futhark server." + either throwError pure =<< liftIO (getValue server v) writeVar :: (MonadError T.Text m, MonadIO m) => Server -> VarName -> V.Value -> m () writeVar server v val = - cmdMaybe . liftIO . withSystemTempFile "futhark-server-write" $ \tmpf tmpf_h -> do - LBS.hPutStr tmpf_h $ Bin.encode val - hClose tmpf_h - -- We are not using prettyprinting for the type, because we don't - -- want the sizes of the dimensions. - let V.ValueType dims t = V.valueType val - t' = mconcat (map (const "[]") dims) <> prettyText t - cmdRestore server tmpf [(v, t')] + cmdMaybe $ liftIO (putValue server v val) -- | A ScriptValue is either a base value or a partially applied -- function. We don't have real first-class functions in @@ -327,7 +309,7 @@ valToInterVal :: V.CompoundValue -> ExpValue valToInterVal = fmap $ \v -> - SValue (V.prettyValueTypeNoDims (V.valueType v)) $ VVal v + SValue (V.valueTypeTextNoDims (V.valueType v)) $ VVal v simpleType (V.ValueAtom (STValue _)) = True simpleType _ = False @@ -355,8 +337,8 @@ throwError $ "Locally bound name cannot be invoked as a function: " <> prettyText name pure e evalExp' vtable (Call (FuncFut name) es) = do - in_types <- cmdEither $ cmdInputs server name - out_types <- cmdEither $ cmdOutputs server name + in_types <- fmap (map inputType) $ cmdEither $ cmdInputs server name + out_types <- fmap (map outputType) $ cmdEither $ cmdOutputs server name es' <- mapM (evalExp' vtable) es let es_types = map (fmap scriptValueType) es' @@ -382,17 +364,17 @@ then do outs <- replicateM (length out_types) $ newVar "out" void $ cmdEither $ cmdCall server name outs ins - pure $ V.mkCompound $ zipWith SValue out_types $ map VVar outs + pure $ V.mkCompound $ map V.ValueAtom $ zipWith SValue out_types $ map VVar outs else pure . V.ValueAtom . SFun name in_types out_types $ zipWith SValue in_types $ map VVar ins evalExp' _ (StringLit s) = case V.putValue s of Just s' -> - pure $ V.ValueAtom $ SValue (prettyText (V.valueType s')) $ VVal s' + pure $ V.ValueAtom $ SValue (V.valueTypeText (V.valueType s')) $ VVal s' Nothing -> error $ "Unable to write value " ++ pretty s evalExp' _ (Const val) = - pure $ V.ValueAtom $ SValue (V.prettyValueTypeNoDims (V.valueType val)) $ VVal val + pure $ V.ValueAtom $ SValue (V.valueTypeTextNoDims (V.valueType val)) $ VVal val evalExp' vtable (Tuple es) = V.ValueTuple <$> mapM (evalExp' vtable) es evalExp' vtable e@(Record m) = do diff --git a/src/Futhark/Server.hs b/src/Futhark/Server.hs deleted file mode 100644 --- a/src/Futhark/Server.hs +++ /dev/null @@ -1,211 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} - --- | Haskell code for interacting with a Futhark server program (a --- program compiled with @--server@). -module Futhark.Server - ( Server, - withServer, - CmdFailure (..), - VarName, - TypeName, - EntryName, - cmdRestore, - cmdStore, - cmdCall, - cmdFree, - cmdRename, - cmdInputs, - cmdOutputs, - cmdClear, - cmdReport, - cmdMaybe, - cmdEither, - ) -where - -import Control.Exception -import Control.Monad -import Control.Monad.Except -import Data.Text (Text) -import qualified Data.Text as T -import qualified Data.Text.IO as T -import Futhark.Util (isEnvVarAtLeast) -import System.Directory (removeFile) -import System.Exit -import System.IO hiding (stdin, stdout) -import System.IO.Temp -import qualified System.Process as P - --- | A handle to a running server. -data Server = Server - { serverStdin :: Handle, - serverStdout :: Handle, - serverErrLog :: FilePath, - serverProc :: P.ProcessHandle, - serverDebug :: Bool - } - -startServer :: FilePath -> [FilePath] -> IO Server -startServer prog options = do - tmpdir <- getCanonicalTemporaryDirectory - (err_log_f, err_log_h) <- openTempFile tmpdir "futhark-server-stderr.log" - (Just stdin, Just stdout, Nothing, phandle) <- - P.createProcess - ( (P.proc prog options) - { P.std_err = P.UseHandle err_log_h, - P.std_in = P.CreatePipe, - P.std_out = P.CreatePipe - } - ) - - code <- P.getProcessExitCode phandle - case code of - Just (ExitFailure e) -> - error $ "Cannot start " ++ prog ++ ": error " ++ show e - _ -> do - let server = - Server - { serverStdin = stdin, - serverStdout = stdout, - serverProc = phandle, - serverDebug = isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 1, - serverErrLog = err_log_f - } - void (responseLines server) `catch` onStartupError server - pure server - where - onStartupError :: Server -> IOError -> IO a - onStartupError s _ = do - code <- P.waitForProcess $ serverProc s - stderr_s <- readFile $ serverErrLog s - removeFile $ serverErrLog s - error $ - "Command failed with " ++ show code ++ ":\n" - ++ unwords (prog : options) - ++ "\nStderr:\n" - ++ stderr_s - -stopServer :: Server -> IO () -stopServer s = do - hClose $ serverStdin s - void $ P.waitForProcess $ serverProc s - removeFile $ serverErrLog s - --- | Start a server, execute an action, then shut down the server. -withServer :: FilePath -> [FilePath] -> (Server -> IO a) -> IO a -withServer prog options = bracket (startServer prog options) stopServer - --- Read lines of response until the next %%% OK (which is what --- indicates that the server is ready for new instructions). -responseLines :: Server -> IO [Text] -responseLines s = do - l <- T.hGetLine $ serverStdout s - when (serverDebug s) $ - T.hPutStrLn stderr $ "<<< " <> l - case l of - "%%% OK" -> pure [] - _ -> (l :) <$> responseLines s - --- | The command failed, and this is why. The first 'Text' is any --- output before the failure indincator, and the second Text is the --- output after the indicator. -data CmdFailure = CmdFailure {failureLog :: [Text], failureMsg :: [Text]} - deriving (Eq, Ord, Show) - --- Figure out whether the response is a failure, and if so, return the --- failure message. -checkForFailure :: [Text] -> Either CmdFailure [Text] -checkForFailure [] = Right [] -checkForFailure ("%%% FAILURE" : ls) = Left $ CmdFailure mempty ls -checkForFailure (l : ls) = - case checkForFailure ls of - Left (CmdFailure xs ys) -> Left $ CmdFailure (l : xs) ys - Right ls' -> Right $ l : ls' - --- Words with spaces in them must be quoted. -quoteWord :: Text -> Text -quoteWord t - | Just _ <- T.find (== ' ') t = - "\"" <> t <> "\"" - | otherwise = t - -sendCommand :: Server -> [Text] -> IO (Either CmdFailure [Text]) -sendCommand s command = do - let command' = T.unwords $ map quoteWord command - - when (serverDebug s) $ - T.hPutStrLn stderr $ ">>> " <> command' - - T.hPutStrLn (serverStdin s) command' - hFlush $ serverStdin s - checkForFailure <$> responseLines s `catch` onError - where - onError :: IOError -> IO a - onError e = do - code <- P.getProcessExitCode $ serverProc s - let code_msg = - case code of - Just (ExitFailure x) -> - "\nServer process exited unexpectedly with exit code: " ++ show x - _ -> mempty - stderr_s <- readFile $ serverErrLog s - error $ - "After sending command " ++ show command ++ " to server process:" - ++ show e - ++ code_msg - ++ "\nServer stderr:\n" - ++ stderr_s - --- | The name of a server-side variable. -type VarName = Text - --- | The name of a server-side type. -type TypeName = Text - --- | The name of an entry point. -type EntryName = Text - -helpCmd :: Server -> [Text] -> IO (Maybe CmdFailure) -helpCmd s cmd = - either Just (const Nothing) <$> sendCommand s cmd - -cmdRestore :: Server -> FilePath -> [(VarName, TypeName)] -> IO (Maybe CmdFailure) -cmdRestore s fname vars = helpCmd s $ "restore" : T.pack fname : concatMap f vars - where - f (v, t) = [v, t] - -cmdStore :: Server -> FilePath -> [VarName] -> IO (Maybe CmdFailure) -cmdStore s fname vars = helpCmd s $ "store" : T.pack fname : vars - -cmdCall :: Server -> EntryName -> [VarName] -> [VarName] -> IO (Either CmdFailure [T.Text]) -cmdCall s entry outs ins = - sendCommand s $ "call" : entry : outs ++ ins - -cmdFree :: Server -> [VarName] -> IO (Maybe CmdFailure) -cmdFree s vs = helpCmd s $ "free" : vs - -cmdRename :: Server -> VarName -> VarName -> IO (Maybe CmdFailure) -cmdRename s oldname newname = helpCmd s ["rename", oldname, newname] - -cmdInputs :: Server -> EntryName -> IO (Either CmdFailure [TypeName]) -cmdInputs s entry = - sendCommand s ["inputs", entry] - -cmdOutputs :: Server -> EntryName -> IO (Either CmdFailure [TypeName]) -cmdOutputs s entry = - sendCommand s ["outputs", entry] - -cmdClear :: Server -> IO (Maybe CmdFailure) -cmdClear s = helpCmd s ["clear"] - -cmdReport :: Server -> IO (Either CmdFailure [T.Text]) -cmdReport s = sendCommand s ["report"] - --- | Turn a 'Maybe'-producing command into a monadic action. -cmdMaybe :: (MonadError T.Text m, MonadIO m) => IO (Maybe CmdFailure) -> m () -cmdMaybe = maybe (pure ()) (throwError . T.unlines . failureMsg) <=< liftIO - --- | Turn an 'Either'-producing command into a monadic action. -cmdEither :: (MonadError T.Text m, MonadIO m) => IO (Either CmdFailure a) -> m a -cmdEither = either (throwError . T.unlines . failureMsg) pure <=< liftIO diff --git a/src/Futhark/Test.hs b/src/Futhark/Test.hs --- a/src/Futhark/Test.hs +++ b/src/Futhark/Test.hs @@ -15,7 +15,7 @@ getValues, getValuesBS, valuesAsVars, - compareValues, + V.compareValues, checkResult, testRunReferenceOutput, getExpectedResult, @@ -25,7 +25,8 @@ ensureReferenceOutput, determineTuning, binaryName, - Mismatch, + futharkServerCfg, + V.Mismatch, ProgramTest (..), StructureTest (..), StructurePipeline (..), @@ -37,7 +38,8 @@ ExpectedResult (..), Success (..), Values (..), - Value, + V.Value, + V.valueText, ) where @@ -62,15 +64,14 @@ import qualified Data.Text.IO as T import Data.Void import Futhark.Analysis.Metrics.Type -import Futhark.IR.Primitive (floatByteSize, intByteSize) +import Futhark.Data.Parser +import qualified Futhark.Data.Parser as V import qualified Futhark.Script as Script import Futhark.Server -import Futhark.Test.Values -import Futhark.Test.Values.Parser -import Futhark.Util (directoryContents, pmapIO) -import Futhark.Util.Pretty (pretty, prettyOneLine, prettyText, prettyTextOneLine) -import Language.Futhark.Prop (primByteSize, primValueType) -import Language.Futhark.Syntax (PrimType (..), PrimValue (..)) +import Futhark.Server.Values +import qualified Futhark.Test.Values as V +import Futhark.Util (directoryContents, isEnvVarAtLeast, pmapIO) +import Futhark.Util.Pretty (prettyOneLine, prettyText, prettyTextOneLine) import System.Directory import System.Exit import System.FilePath @@ -151,7 +152,7 @@ -- | Several values - either literally, or by reference to a file, or -- to be generated on demand. All paths are relative to test program. data Values - = Values [Value] + = Values [V.Value] | InFile FilePath | GenValues [GenValue] | ScriptValues Script.Exp @@ -161,18 +162,18 @@ data GenValue = -- | Generate a value of the given rank and primitive -- type. Scalars are considered 0-ary arrays. - GenValue ValueType + GenValue V.ValueType | -- | A fixed non-randomised primitive value. - GenPrim PrimValue + GenPrim V.Value deriving (Show) -- | A prettyprinted representation of type of value produced by a -- 'GenValue'. genValueType :: GenValue -> String -genValueType (GenValue (ValueType ds t)) = - concatMap (\d -> "[" ++ show d ++ "]") ds ++ pretty t +genValueType (GenValue (V.ValueType ds t)) = + concatMap (\d -> "[" ++ show d ++ "]") ds ++ T.unpack (V.primTypeText t) genValueType (GenPrim v) = - pretty v + T.unpack $ V.valueText v -- | How a test case is expected to terminate. data ExpectedResult values @@ -309,7 +310,7 @@ -- Turn linebreaks into space. "#" ++ show i ++ " (\"" ++ unwords (lines vs') ++ "\")" where - vs' = case unwords (map pretty vs) of + vs' = case unwords $ map (T.unpack . V.valueText) vs of s | length s > 50 -> take 50 s ++ "..." | otherwise -> s @@ -352,7 +353,7 @@ parseGenValue = choice [ GenValue <$> lexeme parseType, - GenPrim <$> lexeme parsePrimValue + GenPrim <$> lexeme V.parsePrimValue ] parseValues :: Parser Values @@ -481,9 +482,9 @@ -- | Try to parse a several values from a byte string. The 'String' -- parameter is used for error messages. -valuesFromByteString :: String -> BS.ByteString -> Either String [Value] +valuesFromByteString :: String -> BS.ByteString -> Either String [V.Value] valuesFromByteString srcname = - maybe (Left $ "Cannot parse values from '" ++ srcname ++ "'") Right . readValues + maybe (Left $ "Cannot parse values from '" ++ srcname ++ "'") Right . V.readValues -- | The @futhark@ executable we are using. This is merely a wrapper -- around the underlying file path, because we will be using a lot of @@ -495,7 +496,7 @@ -- specification. The first 'FilePath' is the path of the @futhark@ -- executable, and the second is the directory which file paths are -- read relative to. -getValues :: (MonadFail m, MonadIO m) => FutharkExe -> FilePath -> Values -> m [Value] +getValues :: (MonadFail m, MonadIO m) => FutharkExe -> FilePath -> Values -> m [V.Value] getValues _ _ (Values vs) = return vs getValues futhark dir v = do @@ -521,7 +522,7 @@ -- guarantee that the resulting byte string yields a readable value. getValuesBS :: (MonadFail m, MonadIO m) => FutharkExe -> FilePath -> Values -> m BS.ByteString getValuesBS _ _ (Values vs) = - return $ BS.fromStrict $ T.encodeUtf8 $ T.unlines $ map prettyText vs + return $ BS.fromStrict $ T.encodeUtf8 $ T.unlines $ map V.valueText vs getValuesBS _ dir (InFile file) = case takeExtension file of ".gz" -> liftIO $ do @@ -543,14 +544,10 @@ (MonadError T.Text m, MonadIO m) => Server -> VarName -> - TypeName -> - Value -> + V.Value -> m () -valueAsVar server v t val = do - cmdMaybe . withSystemTempFile "futhark-input" $ \tmpf tmpf_h -> do - BS.hPutStr tmpf_h $ Bin.encode val - hClose tmpf_h - cmdRestore server tmpf [(v, t)] +valueAsVar server v val = + cmdMaybe $ putValue server v val -- Frees the expression on error. scriptValueAsVars :: @@ -560,24 +557,24 @@ Script.ExpValue -> m () scriptValueAsVars server names_and_types val - | vals <- unCompound val, + | vals <- V.unCompound val, length names_and_types == length vals, Just loads <- zipWithM f names_and_types vals = sequence_ loads where - f (v, t0) (ValueAtom (Script.SValue t1 sval)) + f (v, t0) (V.ValueAtom (Script.SValue t1 sval)) | t0 == t1 = Just $ case sval of Script.VVar oldname -> cmdMaybe $ cmdRename server oldname v Script.VVal sval' -> - valueAsVar server v t0 sval' + valueAsVar server v sval' f _ _ = Nothing scriptValueAsVars server names_and_types val = do cmdMaybe $ cmdFree server $ S.toList $ Script.serverVarsInValue val throwError $ "Expected value of type: " - <> prettyTextOneLine (mkCompound (map snd names_and_types)) + <> prettyTextOneLine (V.mkCompound (map (V.ValueAtom . snd) names_and_types)) <> "\nBut got value of type: " <> prettyTextOneLine (fmap Script.scriptValueType val) <> notes @@ -703,16 +700,12 @@ genFileSize = genSize where header_size = 1 + 1 + 1 + 4 -- 'b' - genSize (GenValue (ValueType ds t)) = - header_size + toInteger (length ds) * 8 - + product (map toInteger ds) * primSize t + genSize (GenValue (V.ValueType ds t)) = + toInteger $ + header_size + length ds * 8 + + product ds * V.primTypeBytes t genSize (GenPrim v) = - header_size + primByteSize (primValueType v) - - primSize (Signed it) = intByteSize it - primSize (Unsigned it) = intByteSize it - primSize (FloatType ft) = floatByteSize ft - primSize Bool = 1 + toInteger $ header_size + product (V.valueShape v) * V.primTypeBytes (V.valueElemType v) -- | When/if generating a reference output file for this run, what -- should it be called? Includes the "data/" folder. @@ -737,7 +730,7 @@ FilePath -> T.Text -> TestRun -> - m (ExpectedResult [Value]) + m (ExpectedResult [V.Value]) getExpectedResult futhark prog entry tr = case runExpectedResult tr of (Succeeds (Just (SuccessValues vals))) -> @@ -821,23 +814,9 @@ (MonadIO m, MonadError T.Text m) => Server -> [VarName] -> - FilePath -> - m [Value] -readResults server outs program = - join . liftIO . withSystemTempFile "futhark-output" $ \outputf outputh -> do - hClose outputh - store_r <- cmdStore server outputf outs - case store_r of - Just (CmdFailure _ err) -> - pure $ throwError $ T.unlines err - Nothing -> do - bytes <- BS.readFile outputf - case valuesFromByteString "output" bytes of - Left e -> do - let actualf = program `addExtension` "actual" - liftIO $ BS.writeFile actualf bytes - pure $ throwError $ T.pack e <> "\n(See " <> T.pack actualf <> ")" - Right vs -> pure $ pure vs + m [V.Value] +readResults server = + mapM (either throwError pure <=< liftIO . getValue server) -- | Ensure that any reference output files exist, or create them (by -- compiling the program with the reference compiler and running it on @@ -910,11 +889,11 @@ checkResult :: (MonadError T.Text m, MonadIO m) => FilePath -> - [Value] -> - [Value] -> + [V.Value] -> + [V.Value] -> m () checkResult program expected_vs actual_vs = - case compareValues actual_vs expected_vs of + case V.compareSeveralValues (V.Tolerance 0.002) actual_vs expected_vs of mismatch : mismatches -> do let actualf = program <.> "actual" expectedf = program <.> "expected" @@ -929,3 +908,11 @@ else "\n...and " <> prettyText (length mismatches) <> " other mismatches." [] -> return () + +-- | Create a Futhark server configuration suitable for use when +-- testing/benchmarking Futhark programs. +futharkServerCfg :: FilePath -> [String] -> ServerCfg +futharkServerCfg prog opts = + (newServerCfg prog opts) + { cfgDebug = isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 1 + } diff --git a/src/Futhark/Test/Values.hs b/src/Futhark/Test/Values.hs --- a/src/Futhark/Test/Values.hs +++ b/src/Futhark/Test/Values.hs @@ -1,203 +1,34 @@ {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE Strict #-} {-# LANGUAGE Trustworthy #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} --- | This module defines an efficient value representation as well as --- parsing and comparison functions. This is because the standard --- Futhark parser is not able to cope with large values (like arrays --- that are tens of megabytes in size). The representation defined --- here does not support tuples, so don't use those as input/output --- for your test programs. +-- | This module provides an efficient value representation as well as +-- parsing and comparison functions. module Futhark.Test.Values - ( Value (..), + ( module Futhark.Data, + module Futhark.Data.Compare, + module Futhark.Data.Reader, Compound (..), CompoundValue, - Vector, - - -- * Reading Values - readValues, - - -- * Types of values - ValueType (..), - prettyValueTypeNoDims, - valueType, - valueShape, - - -- * Manipulating values - valueElems, mkCompound, unCompound, - - -- * Comparing Values - compareValues, - Mismatch, - - -- * Converting values - GetValue (..), - PutValue (..), ) where -import Control.Monad -import Control.Monad.ST -import Data.Binary -import Data.Binary.Get -import Data.Binary.Put -import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy.Char8 as LBS -import Data.Char (chr, isSpace, ord) -import Data.Int (Int16, Int32, Int64, Int8) -import Data.List (intercalate) import qualified Data.Map as M import qualified Data.Text as T -import qualified Data.Text.Encoding as T import Data.Traversable -import Data.Vector.Generic (freeze) -import qualified Data.Vector.Storable as SVec -import Data.Vector.Storable.ByteString (byteStringToVector, vectorToByteString) -import qualified Data.Vector.Unboxed as UVec -import qualified Data.Vector.Unboxed.Mutable as UMVec -import Futhark.IR.Primitive (PrimValue) -import Futhark.IR.Prop.Constants (IsValue (..)) -import Futhark.IR.Prop.Reshape (unflattenIndex) -import Futhark.Util.IntegralExp -import Futhark.Util.Loc (Pos (..)) +import Futhark.Data +import Futhark.Data.Compare +import Futhark.Data.Reader import Futhark.Util.Pretty -import qualified Futhark.Util.Pretty as PP -import Language.Futhark.Parser.Lexer -import Language.Futhark.Pretty () -import qualified Language.Futhark.Syntax as F -type STVector s = UMVec.STVector s - --- | The value vector type. -type Vector = SVec.Vector - --- | An efficiently represented Futhark value. Use 'pretty' to get a --- human-readable representation, and v'put' to obtain binary a --- representation. -data Value - = Int8Value (Vector Int) (Vector Int8) - | Int16Value (Vector Int) (Vector Int16) - | Int32Value (Vector Int) (Vector Int32) - | Int64Value (Vector Int) (Vector Int64) - | Word8Value (Vector Int) (Vector Word8) - | Word16Value (Vector Int) (Vector Word16) - | Word32Value (Vector Int) (Vector Word32) - | Word64Value (Vector Int) (Vector Word64) - | Float32Value (Vector Int) (Vector Float) - | Float64Value (Vector Int) (Vector Double) - | BoolValue (Vector Int) (Vector Bool) - deriving (Show) - -binaryFormatVersion :: Word8 -binaryFormatVersion = 2 - -instance Binary Value where - put (Int8Value shape vs) = putBinaryValue " i8" shape vs - put (Int16Value shape vs) = putBinaryValue " i16" shape vs - put (Int32Value shape vs) = putBinaryValue " i32" shape vs - put (Int64Value shape vs) = putBinaryValue " i64" shape vs - put (Word8Value shape vs) = putBinaryValue " u8" shape vs - put (Word16Value shape vs) = putBinaryValue " u16" shape vs - put (Word32Value shape vs) = putBinaryValue " u32" shape vs - put (Word64Value shape vs) = putBinaryValue " u64" shape vs - put (Float32Value shape vs) = putBinaryValue " f32" shape vs - put (Float64Value shape vs) = putBinaryValue " f64" shape vs - -- Bool must be treated specially because the Storable instance - -- uses four bytes. - put (BoolValue shape vs) = putBinaryValue "bool" shape $ SVec.map boolToInt8 vs - where - boolToInt8 True = 1 :: Int8 - boolToInt8 False = 0 - - get = do - first <- getInt8 - version <- getWord8 - rank <- getInt8 - - unless (chr (fromIntegral first) == 'b') $ - fail "Input does not begin with ASCII 'b'." - unless (version == binaryFormatVersion) $ - fail $ "Expecting binary format version 1; found version: " ++ show version - unless (rank >= 0) $ - fail $ "Rank must be non-negative, but is: " ++ show rank - - type_f <- getLazyByteString 4 - - shape <- replicateM (fromIntegral rank) $ fromIntegral <$> getInt64le - let num_elems = product shape - shape' = SVec.fromList shape - - case LBS.unpack type_f of - " i8" -> get' (Int8Value shape') num_elems 1 - " i16" -> get' (Int16Value shape') num_elems 2 - " i32" -> get' (Int32Value shape') num_elems 4 - " i64" -> get' (Int64Value shape') num_elems 8 - " u8" -> get' (Word8Value shape') num_elems 1 - " u16" -> get' (Word16Value shape') num_elems 2 - " u32" -> get' (Word32Value shape') num_elems 4 - " u64" -> get' (Word64Value shape') num_elems 8 - " f32" -> get' (Float32Value shape') num_elems 4 - " f64" -> get' (Float64Value shape') num_elems 8 - -- Bool must be treated specially because the Storable instance - -- uses four bytes. - "bool" -> BoolValue shape' . SVec.map int8ToBool . byteStringToVector . BS.copy <$> getByteString num_elems - s -> fail $ "Cannot parse binary values of type " ++ show s - where - -- The copy is to ensure that the bytestring is properly - -- aligned. - get' mk num_elems elem_size = - mk . byteStringToVector . BS.copy <$> getByteString (num_elems * elem_size) - - int8ToBool :: Int8 -> Bool - int8ToBool = (/= 0) - -putBinaryValue :: - SVec.Storable a => - String -> - Vector Int -> - Vector a -> - Put -putBinaryValue tstr shape vs = do - putInt8 $ fromIntegral $ ord 'b' - putWord8 binaryFormatVersion - putWord8 $ fromIntegral $ SVec.length shape - mapM_ (putInt8 . fromIntegral . ord) tstr - putByteString $ vectorToByteString shape - putByteString $ vectorToByteString vs - -instance PP.Pretty Value where - ppr v - | product (valueShape v) == 0 = - text "empty" - <> parens (dims <> ppr (valueElemType v)) - where - dims = mconcat $ map (brackets . ppr) $ valueShape v - ppr (Int8Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Int16Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Int32Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Int64Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Word8Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Word16Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Word32Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Word64Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Float32Value shape vs) = pprArray (SVec.toList shape) vs - ppr (Float64Value shape vs) = pprArray (SVec.toList shape) vs - ppr (BoolValue shape vs) = pprArray (SVec.toList shape) vs +instance Pretty Value where + ppr = strictText . valueText -pprArray :: (SVec.Storable a, F.IsPrimValue a) => [Int] -> SVec.Vector a -> Doc -pprArray [] vs = - ppr $ F.primValue $ SVec.head vs -pprArray (d : ds) vs = - brackets $ cat $ punctuate separator $ map (pprArray ds . slice) [0 .. d -1] - where - slice_size = product ds - slice i = SVec.slice (i * slice_size) slice_size vs - separator - | null ds = comma <> space - | otherwise = comma <> line +instance Pretty ValueType where + ppr = strictText . valueTypeText -- | The structure of a compound value, parameterised over the actual -- values. For most cases you probably want 'CompoundValue'. @@ -227,9 +58,9 @@ field (k, v) = ppr k <> equals <> ppr v -- | Create a tuple for a non-unit list, and otherwise a 'ValueAtom' -mkCompound :: [v] -> Compound v -mkCompound [v] = ValueAtom v -mkCompound vs = ValueTuple $ map ValueAtom vs +mkCompound :: [Compound v] -> Compound v +mkCompound [v] = v +mkCompound vs = ValueTuple vs -- | If the value is a tuple, extract the components, otherwise return -- a singleton list of the value. @@ -241,633 +72,3 @@ -- supported by raw values. You cannot parse or read these in -- standard ways, and they cannot be elements of arrays. type CompoundValue = Compound Value - --- | A representation of the simple values we represent in this module. -data ValueType = ValueType [Int] F.PrimType - deriving (Eq, Ord, Show) - -instance PP.Pretty ValueType where - ppr (ValueType ds t) = mconcat (map pprDim ds) <> ppr t - where - pprDim d = brackets $ ppr d - --- | Prettyprint a value type with empty dimensions. This is needed --- for Futhark server programs, whose types are un-sized. -prettyValueTypeNoDims :: ValueType -> T.Text -prettyValueTypeNoDims (ValueType dims t) = - mconcat (replicate (length dims) "[]") <> prettyText t - --- | Get the type of a value. -valueType :: Value -> ValueType -valueType v = ValueType (valueShape v) $ valueElemType v - -valueElemType :: Value -> F.PrimType -valueElemType Int8Value {} = F.Signed F.Int8 -valueElemType Int16Value {} = F.Signed F.Int16 -valueElemType Int32Value {} = F.Signed F.Int32 -valueElemType Int64Value {} = F.Signed F.Int64 -valueElemType Word8Value {} = F.Unsigned F.Int8 -valueElemType Word16Value {} = F.Unsigned F.Int16 -valueElemType Word32Value {} = F.Unsigned F.Int32 -valueElemType Word64Value {} = F.Unsigned F.Int64 -valueElemType Float32Value {} = F.FloatType F.Float32 -valueElemType Float64Value {} = F.FloatType F.Float64 -valueElemType BoolValue {} = F.Bool - --- | The shape of a value. Empty list in case of a scalar. -valueShape :: Value -> [Int] -valueShape (Int8Value shape _) = SVec.toList shape -valueShape (Int16Value shape _) = SVec.toList shape -valueShape (Int32Value shape _) = SVec.toList shape -valueShape (Int64Value shape _) = SVec.toList shape -valueShape (Word8Value shape _) = SVec.toList shape -valueShape (Word16Value shape _) = SVec.toList shape -valueShape (Word32Value shape _) = SVec.toList shape -valueShape (Word64Value shape _) = SVec.toList shape -valueShape (Float32Value shape _) = SVec.toList shape -valueShape (Float64Value shape _) = SVec.toList shape -valueShape (BoolValue shape _) = SVec.toList shape - --- | Produce a list of the immediate elements of the value. That is, --- a 2D array will produce a list of 1D values. While lists are of --- course inefficient, the actual values are just slices of the --- original value, which makes them fairly efficient. -valueElems :: Value -> [Value] -valueElems v - | n : ns <- valueShape v = - let k = product ns - slices mk vs = - [ mk (SVec.fromList ns) $ - SVec.slice (k * i) k vs - | i <- [0 .. n -1] - ] - in case v of - Int8Value _ vs -> slices Int8Value vs - Int16Value _ vs -> slices Int16Value vs - Int32Value _ vs -> slices Int32Value vs - Int64Value _ vs -> slices Int64Value vs - Word8Value _ vs -> slices Word8Value vs - Word16Value _ vs -> slices Word16Value vs - Word32Value _ vs -> slices Word32Value vs - Word64Value _ vs -> slices Word64Value vs - Float32Value _ vs -> slices Float32Value vs - Float64Value _ vs -> slices Float64Value vs - BoolValue _ vs -> slices BoolValue vs - | otherwise = - [] - --- The parser - -dropRestOfLine, dropSpaces :: LBS.ByteString -> LBS.ByteString -dropRestOfLine = LBS.drop 1 . LBS.dropWhile (/= '\n') -dropSpaces t = case LBS.dropWhile isSpace t of - t' - | "--" `LBS.isPrefixOf` t' -> dropSpaces $ dropRestOfLine t' - | otherwise -> t' - -type ReadValue v = LBS.ByteString -> Maybe (v, LBS.ByteString) - -symbol :: Char -> LBS.ByteString -> Maybe LBS.ByteString -symbol c t - | Just (c', t') <- LBS.uncons t, c' == c = Just $ dropSpaces t' - | otherwise = Nothing - -lexeme :: LBS.ByteString -> LBS.ByteString -> Maybe LBS.ByteString -lexeme l t - | l `LBS.isPrefixOf` t = Just $ dropSpaces $ LBS.drop (LBS.length l) t - | otherwise = Nothing - --- (Used elements, shape, elements, remaining input) -type State s v = (Int, Vector Int, STVector s v, LBS.ByteString) - -readArrayElemsST :: - UMVec.Unbox v => - Int -> - Int -> - ReadValue v -> - State s v -> - ST s (Maybe (Int, State s v)) -readArrayElemsST j r rv s = do - ms <- readRankedArrayOfST r rv s - case ms of - Just (i, shape, arr, t) - | Just t' <- symbol ',' t -> do - next <- readArrayElemsST (j + 1) r rv (i, shape, arr, t') - -- Not OK to have zero values after a comma. - case next of - Just (0, _) -> return Nothing - _ -> return next - | otherwise -> return $ Just (j, (i, shape, arr, t)) - _ -> - return $ Just (0, s) - -updateShape :: Int -> Int -> Vector Int -> Maybe (Vector Int) -updateShape d n shape - | old_n < 0 = Just $ shape SVec.// [(r - d, n)] - | old_n == n = Just shape - | otherwise = Nothing - where - r = SVec.length shape - old_n = shape SVec.! (r - d) - -growIfFilled :: UVec.Unbox v => Int -> STVector s v -> ST s (STVector s v) -growIfFilled i arr = - if i >= capacity - then UMVec.grow arr capacity - else return arr - where - capacity = UMVec.length arr - -readRankedArrayOfST :: - UMVec.Unbox v => - Int -> - ReadValue v -> - State s v -> - ST s (Maybe (State s v)) -readRankedArrayOfST 0 rv (i, shape, arr, t) - | Just (v, t') <- rv t = do - arr' <- growIfFilled i arr - UMVec.write arr' i v - return $ Just (i + 1, shape, arr', t') -readRankedArrayOfST r rv (i, shape, arr, t) - | Just t' <- symbol '[' t = do - ms <- readArrayElemsST 1 (r -1) rv (i, shape, arr, t') - return $ do - (j, s) <- ms - closeArray r j s -readRankedArrayOfST _ _ _ = - return Nothing - -closeArray :: Int -> Int -> State s v -> Maybe (State s v) -closeArray r j (i, shape, arr, t) = do - t' <- symbol ']' t - shape' <- updateShape r j shape - return (i, shape', arr, t') - -readRankedArrayOf :: - (UMVec.Unbox v, SVec.Storable v) => - Int -> - ReadValue v -> - LBS.ByteString -> - Maybe (Vector Int, Vector v, LBS.ByteString) -readRankedArrayOf r rv t = runST $ do - arr <- UMVec.new 1024 - ms <- readRankedArrayOfST r rv (0, SVec.replicate r (-1), arr, t) - case ms of - Just (i, shape, arr', t') -> do - arr'' <- freeze (UMVec.slice 0 i arr') - return $ Just (shape, UVec.convert arr'', t') - Nothing -> - return Nothing - --- | A character that can be part of a value. This doesn't work for --- string and character literals. -constituent :: Char -> Bool -constituent ',' = False -constituent ']' = False -constituent ')' = False -constituent c = not $ isSpace c - -readIntegral :: Integral int => (Token -> Maybe int) -> ReadValue int -readIntegral f t = do - v <- case fst <$> scanTokens (Pos "" 1 1 0) a of - Right [L _ NEGATE, L _ (INTLIT x)] -> Just $ negate $ fromIntegral x - Right [L _ (INTLIT x)] -> Just $ fromIntegral x - Right [L _ tok] -> f tok - Right [L _ NEGATE, L _ tok] -> negate <$> f tok - _ -> Nothing - return (v, dropSpaces b) - where - (a, b) = LBS.span constituent t - -readInt8 :: ReadValue Int8 -readInt8 = readIntegral f - where - f (I8LIT x) = Just x - f _ = Nothing - -readInt16 :: ReadValue Int16 -readInt16 = readIntegral f - where - f (I16LIT x) = Just x - f _ = Nothing - -readInt32 :: ReadValue Int32 -readInt32 = readIntegral f - where - f (I32LIT x) = Just x - f _ = Nothing - -readInt64 :: ReadValue Int64 -readInt64 = readIntegral f - where - f (I64LIT x) = Just x - f _ = Nothing - -readWord8 :: ReadValue Word8 -readWord8 = readIntegral f - where - f (U8LIT x) = Just x - f _ = Nothing - -readWord16 :: ReadValue Word16 -readWord16 = readIntegral f - where - f (U16LIT x) = Just x - f _ = Nothing - -readWord32 :: ReadValue Word32 -readWord32 = readIntegral f - where - f (U32LIT x) = Just x - f _ = Nothing - -readWord64 :: ReadValue Word64 -readWord64 = readIntegral f - where - f (U64LIT x) = Just x - f _ = Nothing - -readFloat :: RealFloat float => ([Token] -> Maybe float) -> ReadValue float -readFloat f t = do - v <- case map unLoc . fst <$> scanTokens (Pos "" 1 1 0) a of - Right [NEGATE, FLOATLIT x] -> Just $ negate $ fromDouble x - Right [FLOATLIT x] -> Just $ fromDouble x - Right (NEGATE : toks) -> negate <$> f toks - Right toks -> f toks - _ -> Nothing - return (v, dropSpaces b) - where - (a, b) = LBS.span constituent t - fromDouble = uncurry encodeFloat . decodeFloat - unLoc (L _ x) = x - -readFloat32 :: ReadValue Float -readFloat32 = readFloat lexFloat32 - where - lexFloat32 [F32LIT x] = Just x - lexFloat32 [ID "f32", DOT, ID "inf"] = Just $ 1 / 0 - lexFloat32 [ID "f32", DOT, ID "nan"] = Just $ 0 / 0 - lexFloat32 _ = Nothing - -readFloat64 :: ReadValue Double -readFloat64 = readFloat lexFloat64 - where - lexFloat64 [F64LIT x] = Just x - lexFloat64 [ID "f64", DOT, ID "inf"] = Just $ 1 / 0 - lexFloat64 [ID "f64", DOT, ID "nan"] = Just $ 0 / 0 - lexFloat64 _ = Nothing - -readBool :: ReadValue Bool -readBool t = do - v <- case fst <$> scanTokens (Pos "" 1 1 0) a of - Right [L _ TRUE] -> Just True - Right [L _ FALSE] -> Just False - _ -> Nothing - return (v, dropSpaces b) - where - (a, b) = LBS.span constituent t - -readPrimType :: ReadValue String -readPrimType t = do - pt <- case fst <$> scanTokens (Pos "" 1 1 0) a of - Right [L _ (ID s)] -> Just $ F.nameToString s - _ -> Nothing - return (pt, dropSpaces b) - where - (a, b) = LBS.span constituent t - -readEmptyArrayOfShape :: [Int] -> LBS.ByteString -> Maybe (Value, LBS.ByteString) -readEmptyArrayOfShape shape t - | Just t' <- symbol '[' t, - Just (d, t'') <- readIntegral (const Nothing) t', - Just t''' <- symbol ']' t'' = - readEmptyArrayOfShape (shape ++ [d]) t''' - | otherwise = do - (pt, t') <- readPrimType t - guard $ elem 0 shape - v <- case pt of - "i8" -> Just $ Int8Value (SVec.fromList shape) SVec.empty - "i16" -> Just $ Int16Value (SVec.fromList shape) SVec.empty - "i32" -> Just $ Int32Value (SVec.fromList shape) SVec.empty - "i64" -> Just $ Int64Value (SVec.fromList shape) SVec.empty - "u8" -> Just $ Word8Value (SVec.fromList shape) SVec.empty - "u16" -> Just $ Word16Value (SVec.fromList shape) SVec.empty - "u32" -> Just $ Word32Value (SVec.fromList shape) SVec.empty - "u64" -> Just $ Word64Value (SVec.fromList shape) SVec.empty - "f32" -> Just $ Float32Value (SVec.fromList shape) SVec.empty - "f64" -> Just $ Float64Value (SVec.fromList shape) SVec.empty - "bool" -> Just $ BoolValue (SVec.fromList shape) SVec.empty - _ -> Nothing - return (v, t') - -readEmptyArray :: LBS.ByteString -> Maybe (Value, LBS.ByteString) -readEmptyArray t = do - t' <- symbol '(' =<< lexeme "empty" t - (v, t'') <- readEmptyArrayOfShape [] t' - t''' <- symbol ')' t'' - return (v, t''') - -readValue :: LBS.ByteString -> Maybe (Value, LBS.ByteString) -readValue full_t - | Right (t', _, v) <- decodeOrFail full_t = - Just (v, dropSpaces t') - | otherwise = readEmptyArray full_t `mplus` insideBrackets 0 full_t - where - insideBrackets r t = maybe (tryValueAndReadValue r t) (insideBrackets (r + 1)) $ symbol '[' t - tryWith f mk r t - | Just _ <- f t = do - (shape, arr, rest_t) <- readRankedArrayOf r f full_t - return (mk shape arr, rest_t) - | otherwise = Nothing - tryValueAndReadValue r t = - -- 32-bit signed integers come first such that we parse - -- unsuffixed integer constants as of that type. - tryWith readInt32 Int32Value r t - `mplus` tryWith readInt8 Int8Value r t - `mplus` tryWith readInt16 Int16Value r t - `mplus` tryWith readInt64 Int64Value r t - `mplus` tryWith readWord8 Word8Value r t - `mplus` tryWith readWord16 Word16Value r t - `mplus` tryWith readWord32 Word32Value r t - `mplus` tryWith readWord64 Word64Value r t - `mplus` tryWith readFloat64 Float64Value r t - `mplus` tryWith readFloat32 Float32Value r t - `mplus` tryWith readBool BoolValue r t - --- | Parse Futhark values from the given bytestring. -readValues :: LBS.ByteString -> Maybe [Value] -readValues = readValues' . dropSpaces - where - readValues' t - | LBS.null t = Just [] - | otherwise = do - (a, t') <- readValue t - (a :) <$> readValues' t' - --- Comparisons - --- | Two values differ in some way. The 'Show' instance produces a --- human-readable explanation. -data Mismatch - = -- | The position the value number and a flat index - -- into the array. - PrimValueMismatch Int [Int] PrimValue PrimValue - | ArrayShapeMismatch Int [Int] [Int] - | TypeMismatch Int String String - | ValueCountMismatch Int Int - -instance Show Mismatch where - show (PrimValueMismatch vi [] got expected) = - explainMismatch (show vi) "" got expected - show (PrimValueMismatch vi js got expected) = - explainMismatch (show vi ++ " index [" ++ intercalate "," (map show js) ++ "]") "" got expected - show (ArrayShapeMismatch i got expected) = - explainMismatch (show i) "array of shape " got expected - show (TypeMismatch i got expected) = - explainMismatch (show i) "value of type " got expected - show (ValueCountMismatch got expected) = - "Expected " ++ show expected ++ " values, got " ++ show got - --- | A human-readable description of how two values are not the same. -explainMismatch :: (PP.Pretty a) => String -> String -> a -> a -> String -explainMismatch i what got expected = - "Value #" ++ i ++ ": expected " ++ what ++ PP.pretty expected ++ ", got " ++ PP.pretty got - --- | Compare two sets of Futhark values for equality. Shapes and --- types must also match. -compareValues :: [Value] -> [Value] -> [Mismatch] -compareValues got expected - | n /= m = [ValueCountMismatch n m] - | otherwise = concat $ zipWith3 compareValue [0 ..] got expected - where - n = length got - m = length expected - -compareValue :: Int -> Value -> Value -> [Mismatch] -compareValue i got_v expected_v - | valueShape got_v == valueShape expected_v = - case (got_v, expected_v) of - (Int8Value _ got_vs, Int8Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Int16Value _ got_vs, Int16Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Int32Value _ got_vs, Int32Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Int64Value _ got_vs, Int64Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Word8Value _ got_vs, Word8Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Word16Value _ got_vs, Word16Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Word32Value _ got_vs, Word32Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Word64Value _ got_vs, Word64Value _ expected_vs) -> - compareNum 1 got_vs expected_vs - (Float32Value _ got_vs, Float32Value _ expected_vs) -> - compareFloat (tolerance expected_vs) got_vs expected_vs - (Float64Value _ got_vs, Float64Value _ expected_vs) -> - compareFloat (tolerance expected_vs) got_vs expected_vs - (BoolValue _ got_vs, BoolValue _ expected_vs) -> - compareGen compareBool got_vs expected_vs - _ -> - [TypeMismatch i (pretty $ valueElemType got_v) (pretty $ valueElemType expected_v)] - | otherwise = - [ArrayShapeMismatch i (valueShape got_v) (valueShape expected_v)] - where - unflatten = - map wrappedValue . unflattenIndex (map Wrapped (valueShape got_v)) . Wrapped - - {-# INLINE compareGen #-} - {-# INLINE compareNum #-} - {-# INLINE compareFloat #-} - {-# INLINE compareFloatElement #-} - {-# INLINE compareElement #-} - compareNum tol = compareGen $ compareElement tol - compareFloat tol = compareGen $ compareFloatElement tol - - compareGen cmp got expected = - let l = SVec.length got - check acc j - | j < l = - case cmp j (got SVec.! j) (expected SVec.! j) of - Just mismatch -> - check (mismatch : acc) (j + 1) - Nothing -> - check acc (j + 1) - | otherwise = - acc - in reverse $ check [] 0 - - compareElement tol j got expected - | comparePrimValue tol got expected = Nothing - | otherwise = Just $ PrimValueMismatch i (unflatten j) (value got) (value expected) - - compareFloatElement tol j got expected - | isNaN got, - isNaN expected = - Nothing - | isInfinite got, - isInfinite expected, - signum got == signum expected = - Nothing - | otherwise = - compareElement tol j got expected - - compareBool j got expected - | got == expected = Nothing - | otherwise = Just $ PrimValueMismatch i (unflatten j) (value got) (value expected) - -comparePrimValue :: - (Ord num, Num num) => - num -> - num -> - num -> - Bool -comparePrimValue tol x y = - diff < tol - where - diff = abs $ x - y - -minTolerance :: Fractional a => a -minTolerance = 0.002 -- 0.2% - -tolerance :: (RealFloat a, SVec.Storable a) => Vector a -> a -tolerance = SVec.foldl tolerance' minTolerance . SVec.filter (not . nanOrInf) - where - tolerance' t v = max t $ minTolerance * v - nanOrInf x = isInfinite x || isNaN x - --- | A class for Haskell values that can be retrieved from 'Value'. --- This is a convenience facility - don't expect it to be fast. -class GetValue t where - getValue :: Value -> Maybe t - -instance GetValue t => GetValue [t] where - getValue = mapM getValue . valueElems - -instance GetValue Bool where - getValue (BoolValue shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Int8 where - getValue (Int8Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Int16 where - getValue (Int16Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Int32 where - getValue (Int32Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Int64 where - getValue (Int64Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Word8 where - getValue (Word8Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Word16 where - getValue (Word16Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Word32 where - getValue (Word32Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - -instance GetValue Word64 where - getValue (Word64Value shape vs) - | [] <- SVec.toList shape = - Just $ vs SVec.! 0 - getValue _ = Nothing - --- | A class for Haskell values that can be converted to 'Value'. --- This is a convenience facility - don't expect it to be fast. -class PutValue t where - -- | This may fail for cases such as irregular arrays. - putValue :: t -> Maybe Value - -instance PutValue Word8 where - putValue = Just . Word8Value mempty . SVec.singleton - -instance PutValue F.PrimValue where - putValue (F.SignedValue (F.Int8Value x)) = - Just $ Int8Value mempty $ SVec.singleton x - putValue (F.SignedValue (F.Int16Value x)) = - Just $ Int16Value mempty $ SVec.singleton x - putValue (F.SignedValue (F.Int32Value x)) = - Just $ Int32Value mempty $ SVec.singleton x - putValue (F.SignedValue (F.Int64Value x)) = - Just $ Int64Value mempty $ SVec.singleton x - putValue (F.UnsignedValue (F.Int8Value x)) = - Just $ Word8Value mempty $ SVec.singleton $ fromIntegral x - putValue (F.UnsignedValue (F.Int16Value x)) = - Just $ Word16Value mempty $ SVec.singleton $ fromIntegral x - putValue (F.UnsignedValue (F.Int32Value x)) = - Just $ Word32Value mempty $ SVec.singleton $ fromIntegral x - putValue (F.UnsignedValue (F.Int64Value x)) = - Just $ Word64Value mempty $ SVec.singleton $ fromIntegral x - putValue (F.FloatValue (F.Float32Value x)) = - Just $ Float32Value mempty $ SVec.singleton x - putValue (F.FloatValue (F.Float64Value x)) = - Just $ Float64Value mempty $ SVec.singleton x - putValue (F.BoolValue b) = - Just $ BoolValue mempty $ SVec.singleton b - -instance PutValue [Value] where - putValue [] = Nothing - putValue (x : xs) = do - let res_shape = SVec.fromList $ length (x : xs) : valueShape x - guard $ all ((== valueType x) . valueType) xs - Just $ case x of - Int8Value {} -> Int8Value res_shape $ foldMap getVec (x : xs) - Int16Value {} -> Int16Value res_shape $ foldMap getVec (x : xs) - Int32Value {} -> Int32Value res_shape $ foldMap getVec (x : xs) - Int64Value {} -> Int64Value res_shape $ foldMap getVec (x : xs) - Word8Value {} -> Word8Value res_shape $ foldMap getVec (x : xs) - Word16Value {} -> Word16Value res_shape $ foldMap getVec (x : xs) - Word32Value {} -> Word32Value res_shape $ foldMap getVec (x : xs) - Word64Value {} -> Word64Value res_shape $ foldMap getVec (x : xs) - Float32Value {} -> Float32Value res_shape $ foldMap getVec (x : xs) - Float64Value {} -> Float64Value res_shape $ foldMap getVec (x : xs) - BoolValue {} -> BoolValue res_shape $ foldMap getVec (x : xs) - where - getVec (Int8Value _ vec) = SVec.unsafeCast vec - getVec (Int16Value _ vec) = SVec.unsafeCast vec - getVec (Int32Value _ vec) = SVec.unsafeCast vec - getVec (Int64Value _ vec) = SVec.unsafeCast vec - getVec (Word8Value _ vec) = SVec.unsafeCast vec - getVec (Word16Value _ vec) = SVec.unsafeCast vec - getVec (Word32Value _ vec) = SVec.unsafeCast vec - getVec (Word64Value _ vec) = SVec.unsafeCast vec - getVec (Float32Value _ vec) = SVec.unsafeCast vec - getVec (Float64Value _ vec) = SVec.unsafeCast vec - getVec (BoolValue _ vec) = SVec.unsafeCast vec - -instance PutValue T.Text where - putValue = putValue . T.encodeUtf8 - -instance PutValue BS.ByteString where - putValue bs = - Just $ Word8Value size $ byteStringToVector bs - where - size = SVec.fromList [fromIntegral (BS.length bs)] diff --git a/src/Futhark/Test/Values/Parser.hs b/src/Futhark/Test/Values/Parser.hs deleted file mode 100644 --- a/src/Futhark/Test/Values/Parser.hs +++ /dev/null @@ -1,167 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} - --- | Megaparsec-based parser for primitive 'Value's. The difference --- between this and the parser defined in "Futhark.Test.Values" is --- that we don't try to handle both the textual and binary format - --- only the former. On the other hand, this parser has (much) better --- error messages and can be easily used by other parsers (like the --- ones for FutharkScript or test blocks. -module Futhark.Test.Values.Parser - ( parsePrimType, - parseType, - parsePrimValue, - parseValue, - ) -where - -import Control.Monad.Except -import Data.Functor -import qualified Data.Set as S -import qualified Data.Text as T -import qualified Data.Vector.Storable as SVec -import Data.Void -import Futhark.Test.Values -import qualified Language.Futhark.Syntax as F -import Text.Megaparsec -import Text.Megaparsec.Char.Lexer - ( binary, - decimal, - float, - hexadecimal, - signed, - ) - -type Parser = Parsec Void T.Text - --- | Parse the name of a primitive type. Does *not* consume any --- trailing whitespace, nor does it permit any internal whitespace. -parsePrimType :: Parser F.PrimType -parsePrimType = - choice - [ "i8" $> F.Signed F.Int8, - "i16" $> F.Signed F.Int16, - "i32" $> F.Signed F.Int32, - "i64" $> F.Signed F.Int64, - "u8" $> F.Unsigned F.Int8, - "u16" $> F.Unsigned F.Int16, - "u32" $> F.Unsigned F.Int32, - "u64" $> F.Unsigned F.Int64, - "f32" $> F.FloatType F.Float32, - "f64" $> F.FloatType F.Float64, - "bool" $> F.Bool - ] - -parseInteger :: Parser Integer -parseInteger = - signed (pure ()) $ - choice - [ "0b" *> binary, - "0x" *> hexadecimal, - decimal - ] - -parseIntConst :: Parser F.PrimValue -parseIntConst = do - x <- parseInteger - notFollowedBy $ "f32" <|> "f64" <|> "." <|> "e" - choice - [ signedV F.Int8Value x "i8", - signedV F.Int16Value x "i16", - signedV F.Int32Value x "i32", - signedV F.Int64Value x "i64", - unsignedV F.Int8Value x "u8", - unsignedV F.Int16Value x "u16", - unsignedV F.Int32Value x "u32", - unsignedV F.Int64Value x "u64", - signedV F.Int32Value x "" - ] - where - signedV mk x suffix = - suffix $> F.SignedValue (mk (fromInteger x)) - unsignedV mk x suffix = - suffix $> F.UnsignedValue (mk (fromInteger x)) - -parseFloatConst :: Parser F.PrimValue -parseFloatConst = - choice - [ "f32.nan" $> F.FloatValue (F.Float32Value (0 / 0)), - "f64.nan" $> F.FloatValue (F.Float64Value (0 / 0)), - "f32.inf" $> F.FloatValue (F.Float32Value (1 / 0)), - "f64.inf" $> F.FloatValue (F.Float64Value (1 / 0)), - "-f32.inf" $> F.FloatValue (F.Float32Value (-1 / 0)), - "-f64.inf" $> F.FloatValue (F.Float64Value (-1 / 0)), - numeric - ] - where - numeric = do - x <- - signed (pure ()) $ choice [try float, fromInteger <$> decimal] - choice - [ floatV F.Float32Value x "f32", - floatV F.Float64Value x "f64", - floatV F.Float64Value x "" - ] - - floatV mk x suffix = - suffix $> F.FloatValue (mk (realToFrac (x :: Double))) - --- | Parse a primitive value. Does *not* consume any trailing --- whitespace, nor does it permit any internal whitespace. -parsePrimValue :: Parser F.PrimValue -parsePrimValue = - choice - [ try parseIntConst, - parseFloatConst, - "true" $> F.BoolValue True, - "false" $> F.BoolValue False - ] - -lexeme :: Parser () -> Parser a -> Parser a -lexeme sep p = p <* sep - -inBrackets :: Parser () -> Parser a -> Parser a -inBrackets sep = between (lexeme sep "[") (lexeme sep "]") - --- | Parse a type. Does *not* consume any trailing whitespace, nor --- does it permit any internal whitespace. -parseType :: Parser ValueType -parseType = ValueType <$> many parseDim <*> parsePrimType - where - parseDim = fromInteger <$> ("[" *> parseInteger <* "]") - -parseEmpty :: Parser Value -parseEmpty = do - ValueType dims t <- parseType - unless (product dims == 0) $ fail "Expected at least one empty dimension" - pure $ case t of - F.Signed F.Int8 -> Int8Value (SVec.fromList dims) mempty - F.Signed F.Int16 -> Int16Value (SVec.fromList dims) mempty - F.Signed F.Int32 -> Int32Value (SVec.fromList dims) mempty - F.Signed F.Int64 -> Int64Value (SVec.fromList dims) mempty - F.Unsigned F.Int8 -> Word8Value (SVec.fromList dims) mempty - F.Unsigned F.Int16 -> Word16Value (SVec.fromList dims) mempty - F.Unsigned F.Int32 -> Word32Value (SVec.fromList dims) mempty - F.Unsigned F.Int64 -> Word64Value (SVec.fromList dims) mempty - F.FloatType F.Float32 -> Float32Value (SVec.fromList dims) mempty - F.FloatType F.Float64 -> Float64Value (SVec.fromList dims) mempty - F.Bool -> BoolValue (SVec.fromList dims) mempty - --- | Parse a value, given a post-lexeme parser for whitespace. -parseValue :: Parser () -> Parser Value -parseValue sep = - choice - [ putValue' $ lexeme sep parsePrimValue, - putValue' $ inBrackets sep (parseValue sep `sepBy` lexeme sep ","), - lexeme sep $ "empty(" *> parseEmpty <* ")" - ] - where - putValue' :: PutValue v => Parser v -> Parser Value - putValue' p = do - o <- getOffset - x <- p - case putValue x of - Nothing -> - parseError . FancyError o . S.singleton $ - ErrorFail "array is irregular or has elements of multiple types." - Just v -> - pure v diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -31,19 +31,19 @@ -- Only handles a pattern with an empty 'patternContextElements'. redomapToMapAndReduce :: ( MonadFreshNames m, - Bindable lore, - ExpDec lore ~ (), - Op lore ~ SOAC lore + Bindable rep, + ExpDec rep ~ (), + Op rep ~ SOAC rep ) => - Pattern lore -> + Pattern rep -> ( SubExp, Commutativity, - LambdaT lore, - LambdaT lore, + LambdaT rep, + LambdaT rep, [SubExp], [VName] ) -> - m (Stm lore, Stm lore) + m (Stm rep, Stm rep) redomapToMapAndReduce (Pattern [] patelems) (w, comm, redlam, map_lam, accs, arrs) = do @@ -62,7 +62,7 @@ (Typed dec, MonadFreshNames m) => [PatElemT dec] -> SubExp -> - LambdaT lore -> + LambdaT rep -> [SubExp] -> m ([Ident], PatternT dec, [(SubExp, VName)]) splitScanOrRedomap patelems w map_lam accs = do @@ -84,12 +84,12 @@ -- In essense, what happens is the opposite of horisontal fusion. dissectScrema :: ( MonadBinder m, - Op (Lore m) ~ SOAC (Lore m), - Bindable (Lore m) + Op (Rep m) ~ SOAC (Rep m), + Bindable (Rep m) ) => - Pattern (Lore m) -> + Pattern (Rep m) -> SubExp -> - ScremaForm (Lore m) -> + ScremaForm (Rep m) -> [VName] -> m () dissectScrema pat w (ScremaForm scans reds map_lam) arrs = do @@ -110,11 +110,11 @@ -- | Turn a stream SOAC into statements that apply the stream lambda -- to the entire input. sequentialStreamWholeArray :: - (MonadBinder m, Bindable (Lore m)) => - Pattern (Lore m) -> + (MonadBinder m, Bindable (Rep m)) => + Pattern (Rep m) -> SubExp -> [SubExp] -> - LambdaT (Lore m) -> + LambdaT (Rep m) -> [VName] -> m () sequentialStreamWholeArray pat w nes lam arrs = do diff --git a/src/Futhark/Transform/CopyPropagate.hs b/src/Futhark/Transform/CopyPropagate.hs --- a/src/Futhark/Transform/CopyPropagate.hs +++ b/src/Futhark/Transform/CopyPropagate.hs @@ -15,31 +15,31 @@ import Futhark.IR import Futhark.MonadFreshNames import Futhark.Optimise.Simplify -import Futhark.Optimise.Simplify.Lore (Wise) +import Futhark.Optimise.Simplify.Rep (Wise) import Futhark.Pass -- | Run copy propagation on an entire program. copyPropagateInProg :: - SimplifiableLore lore => - SimpleOps lore -> - Prog lore -> - PassM (Prog lore) + SimplifiableRep rep => + SimpleOps rep -> + Prog rep -> + PassM (Prog rep) copyPropagateInProg simpl = simplifyProg simpl mempty neverHoist -- | Run copy propagation on some statements. copyPropagateInStms :: - (MonadFreshNames m, SimplifiableLore lore) => - SimpleOps lore -> - Scope lore -> - Stms lore -> - m (ST.SymbolTable (Wise lore), Stms lore) + (MonadFreshNames m, SimplifiableRep rep) => + SimpleOps rep -> + Scope rep -> + Stms rep -> + m (ST.SymbolTable (Wise rep), Stms rep) copyPropagateInStms simpl = simplifyStms simpl mempty neverHoist -- | Run copy propagation on a function. copyPropagateInFun :: - (MonadFreshNames m, SimplifiableLore lore) => - SimpleOps lore -> - ST.SymbolTable (Wise lore) -> - FunDef lore -> - m (FunDef lore) + (MonadFreshNames m, SimplifiableRep rep) => + SimpleOps rep -> + ST.SymbolTable (Wise rep) -> + FunDef rep -> + m (FunDef rep) copyPropagateInFun simpl = simplifyFun simpl mempty neverHoist diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs --- a/src/Futhark/Transform/FirstOrderTransform.hs +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -10,7 +10,7 @@ module Futhark.Transform.FirstOrderTransform ( transformFunDef, transformConsts, - FirstOrderLore, + FirstOrderRep, Transformer, transformStmRecursively, transformLambda, @@ -30,23 +30,23 @@ import Futhark.Tools import Futhark.Util (chunks, splitAt3) --- | The constraints that must hold for a lore in order to be the +-- | The constraints that must hold for a rep in order to be the -- target of first-order transformation. -type FirstOrderLore lore = - ( Bindable lore, - BinderOps lore, - LetDec SOACS ~ LetDec lore, - LParamInfo SOACS ~ LParamInfo lore, - CanBeAliased (Op lore) +type FirstOrderRep rep = + ( Bindable rep, + BinderOps rep, + LetDec SOACS ~ LetDec rep, + LParamInfo SOACS ~ LParamInfo rep, + CanBeAliased (Op rep) ) -- | First-order-transform a single function, with the given scope -- provided by top-level constants. transformFunDef :: - (MonadFreshNames m, FirstOrderLore tolore) => - Scope tolore -> + (MonadFreshNames m, FirstOrderRep torep) => + Scope torep -> FunDef SOACS -> - m (AST.FunDef tolore) + m (AST.FunDef torep) transformFunDef consts_scope (FunDef entry attrs fname rettype params body) = do (body', _) <- modifyNameSource $ runState $ runBinderT m consts_scope return $ FunDef entry attrs fname rettype params body' @@ -55,9 +55,9 @@ -- | First-order-transform these top-level constants. transformConsts :: - (MonadFreshNames m, FirstOrderLore tolore) => + (MonadFreshNames m, FirstOrderRep torep) => Stms SOACS -> - m (AST.Stms tolore) + m (AST.Stms torep) transformConsts stms = fmap snd $ modifyNameSource $ runState $ runBinderT m mempty where @@ -67,17 +67,17 @@ -- first-order transformation. type Transformer m = ( MonadBinder m, - LocalScope (Lore m) m, - Bindable (Lore m), - BinderOps (Lore m), - LParamInfo SOACS ~ LParamInfo (Lore m), - CanBeAliased (Op (Lore m)) + LocalScope (Rep m) m, + Bindable (Rep m), + BinderOps (Rep m), + LParamInfo SOACS ~ LParamInfo (Rep m), + CanBeAliased (Op (Rep m)) ) transformBody :: - (Transformer m, LetDec (Lore m) ~ LetDec SOACS) => + (Transformer m, LetDec (Rep m) ~ LetDec SOACS) => Body -> - m (AST.Body (Lore m)) + m (AST.Body (Rep m)) transformBody (Body () stms res) = buildBody_ $ do mapM_ transformStmRecursively stms pure res @@ -85,7 +85,7 @@ -- | First transform any nested t'Body' or t'Lambda' elements, then -- apply 'transformSOAC' if the expression is a SOAC. transformStmRecursively :: - (Transformer m, LetDec (Lore m) ~ LetDec SOACS) => + (Transformer m, LetDec (Rep m) ~ LetDec SOACS) => Stm -> m () transformStmRecursively (Let pat aux (Op soac)) = @@ -119,11 +119,11 @@ -- | Transform a single 'SOAC' into a do-loop. The body of the lambda -- is untouched, and may or may not contain further 'SOAC's depending --- on the given lore. +-- on the given rep. transformSOAC :: Transformer m => - AST.Pattern (Lore m) -> - SOAC (Lore m) -> + AST.Pattern (Rep m) -> + SOAC (Rep m) -> m () transformSOAC pat (Screma w arrs form@(ScremaForm scans reds map_lam)) = do -- See Note [Translation of Screma]. @@ -360,15 +360,15 @@ -- | Recursively first-order-transform a lambda. transformLambda :: ( MonadFreshNames m, - Bindable lore, - BinderOps lore, - LocalScope somelore m, - SameScope somelore lore, - LetDec lore ~ LetDec SOACS, - CanBeAliased (Op lore) + Bindable rep, + BinderOps rep, + LocalScope somerep m, + SameScope somerep rep, + LetDec rep ~ LetDec SOACS, + CanBeAliased (Op rep) ) => Lambda -> - m (AST.Lambda lore) + m (AST.Lambda rep) transformLambda (Lambda params body rettype) = do body' <- runBodyBinder $ @@ -389,8 +389,8 @@ bindLambda :: Transformer m => - AST.Lambda (Lore m) -> - [AST.Exp (Lore m)] -> + AST.Lambda (Rep m) -> + [AST.Exp (Rep m)] -> m [SubExp] bindLambda (Lambda params body _) args = do forM_ (zip params args) $ \(param, arg) -> diff --git a/src/Futhark/Transform/Rename.hs b/src/Futhark/Transform/Rename.hs --- a/src/Futhark/Transform/Rename.hs +++ b/src/Futhark/Transform/Rename.hs @@ -53,9 +53,9 @@ -- correct to begin with. In particular, the renaming may make an -- invalid program valid. renameProg :: - (Renameable lore, MonadFreshNames m) => - Prog lore -> - m (Prog lore) + (Renameable rep, MonadFreshNames m) => + Prog rep -> + m (Prog rep) renameProg prog = modifyNameSource $ runRenamer $ renamingStms (progConsts prog) $ \consts -> do @@ -67,9 +67,9 @@ -- expression was correct to begin with. Any free variables are left -- untouched. renameExp :: - (Renameable lore, MonadFreshNames m) => - Exp lore -> - m (Exp lore) + (Renameable rep, MonadFreshNames m) => + Exp rep -> + m (Exp rep) renameExp = modifyNameSource . runRenamer . rename -- | Rename bound variables such that each is unique. The semantics @@ -77,9 +77,9 @@ -- binding was correct to begin with. Any free variables are left -- untouched, as are the names in the pattern of the binding. renameStm :: - (Renameable lore, MonadFreshNames m) => - Stm lore -> - m (Stm lore) + (Renameable rep, MonadFreshNames m) => + Stm rep -> + m (Stm rep) renameStm binding = do e <- renameExp $ stmExp binding return binding {stmExp = e} @@ -88,9 +88,9 @@ -- of the body is unaffected, under the assumption that the body was -- correct to begin with. Any free variables are left untouched. renameBody :: - (Renameable lore, MonadFreshNames m) => - Body lore -> - m (Body lore) + (Renameable rep, MonadFreshNames m) => + Body rep -> + m (Body rep) renameBody = modifyNameSource . runRenamer . rename -- | Rename bound variables such that each is unique. The semantics @@ -98,9 +98,9 @@ -- correct to begin with. Any free variables are left untouched. -- Note in particular that the parameters of the lambda are renamed. renameLambda :: - (Renameable lore, MonadFreshNames m) => - Lambda lore -> - m (Lambda lore) + (Renameable rep, MonadFreshNames m) => + Lambda rep -> + m (Lambda rep) renameLambda = modifyNameSource . runRenamer . rename -- | Produce an equivalent pattern but with each pattern element given @@ -196,7 +196,7 @@ -- | Rename some statements, then execute an action with the name -- substitutions induced by the statements active. -renamingStms :: Renameable lore => Stms lore -> (Stms lore -> RenameM a) -> RenameM a +renamingStms :: Renameable rep => Stms rep -> (Stms rep -> RenameM a) -> RenameM a renamingStms stms m = descend mempty stms where descend stms' rem_stms = case stmsHead rem_stms of @@ -205,7 +205,7 @@ stm' <- rename stm descend (stms' <> oneStm stm') rem_stms' -instance Renameable lore => Rename (FunDef lore) where +instance Renameable rep => Rename (FunDef rep) where rename (FunDef entry attrs fname ret params body) = bind (map paramName params) $ do params' <- mapM rename params @@ -236,16 +236,16 @@ rename (StmAux cs attrs dec) = StmAux <$> rename cs <*> rename attrs <*> rename dec -instance Renameable lore => Rename (Body lore) where +instance Renameable rep => Rename (Body rep) where rename (Body dec stms res) = do dec' <- rename dec renamingStms stms $ \stms' -> Body dec' stms' <$> rename res -instance Renameable lore => Rename (Stm lore) where - rename (Let pat elore e) = Let <$> rename pat <*> rename elore <*> rename e +instance Renameable rep => Rename (Stm rep) where + rename (Let pat dec e) = Let <$> rename pat <*> rename dec <*> rename e -instance Renameable lore => Rename (Exp lore) where +instance Renameable rep => Rename (Exp rep) where rename (WithAcc inputs lam) = WithAcc <$> rename inputs <*> rename lam rename (DoLoop ctx val form loopbody) = do @@ -315,7 +315,7 @@ rename (Acc acc ispace ts u) = Acc <$> rename acc <*> rename ispace <*> rename ts <*> pure u -instance Renameable lore => Rename (Lambda lore) where +instance Renameable rep => Rename (Lambda rep) where rename (Lambda params body ret) = bind (map paramName params) $ do params' <- mapM rename params @@ -343,14 +343,14 @@ rename (DimFix i) = DimFix <$> rename i rename (DimSlice i n s) = DimSlice <$> rename i <*> rename n <*> rename s --- | Lores in which all annotations are renameable. -type Renameable lore = - ( Rename (LetDec lore), - Rename (ExpDec lore), - Rename (BodyDec lore), - Rename (FParamInfo lore), - Rename (LParamInfo lore), - Rename (RetType lore), - Rename (BranchType lore), - Rename (Op lore) +-- | Representations in which all decorations are renameable. +type Renameable rep = + ( Rename (LetDec rep), + Rename (ExpDec rep), + Rename (BodyDec rep), + Rename (FParamInfo rep), + Rename (LParamInfo rep), + Rename (RetType rep), + Rename (BranchType rep), + Rename (Op rep) ) diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs --- a/src/Futhark/Transform/Substitute.hs +++ b/src/Futhark/Transform/Substitute.hs @@ -37,7 +37,7 @@ instance Substitute a => Substitute [a] where substituteNames substs = map $ substituteNames substs -instance Substitute (Stm lore) => Substitute (Stms lore) where +instance Substitute (Stm rep) => Substitute (Stms rep) where substituteNames substs = fmap $ substituteNames substs instance (Substitute a, Substitute b) => Substitute (a, b) where @@ -72,7 +72,7 @@ substituteNames substs (Var v) = Var $ substituteNames substs v substituteNames _ (Constant v) = Constant v -instance Substitutable lore => Substitute (Exp lore) where +instance Substitutable rep => Substitute (Exp rep) where substituteNames substs = mapExp $ replace substs instance Substitute dec => Substitute (PatElemT dec) where @@ -103,21 +103,21 @@ substituteNames substs (Certificates cs) = Certificates $ substituteNames substs cs -instance Substitutable lore => Substitute (Stm lore) where +instance Substitutable rep => Substitute (Stm rep) where substituteNames substs (Let pat annot e) = Let (substituteNames substs pat) (substituteNames substs annot) (substituteNames substs e) -instance Substitutable lore => Substitute (Body lore) where +instance Substitutable rep => Substitute (Body rep) where substituteNames substs (Body dec stms res) = Body (substituteNames substs dec) (substituteNames substs stms) (substituteNames substs res) -replace :: Substitutable lore => M.Map VName VName -> Mapper lore lore Identity +replace :: Substitutable rep => M.Map VName VName -> Mapper rep rep Identity replace substs = Mapper { mapOnVName = return . substituteNames substs, @@ -164,7 +164,7 @@ substituteNames _ (Mem space) = Mem space -instance Substitutable lore => Substitute (Lambda lore) where +instance Substitutable rep => Substitute (Lambda rep) where substituteNames substs (Lambda params body rettype) = Lambda (substituteNames substs params) @@ -191,7 +191,7 @@ substituteNames substs = TPrimExp . fmap (substituteNames substs) . untyped -instance Substitutable lore => Substitute (NameInfo lore) where +instance Substitutable rep => Substitute (NameInfo rep) where substituteNames subst (LetName dec) = LetName $ substituteNames subst dec substituteNames subst (FParamName dec) = @@ -204,16 +204,16 @@ instance Substitute FV where substituteNames subst = fvNames . substituteNames subst . freeIn --- | Lores in which all annotations support name +-- | Representations in which all annotations support name -- substitution. -type Substitutable lore = - ( Decorations lore, - Substitute (ExpDec lore), - Substitute (BodyDec lore), - Substitute (LetDec lore), - Substitute (FParamInfo lore), - Substitute (LParamInfo lore), - Substitute (RetType lore), - Substitute (BranchType lore), - Substitute (Op lore) +type Substitutable rep = + ( RepTypes rep, + Substitute (ExpDec rep), + Substitute (BodyDec rep), + Substitute (LetDec rep), + Substitute (FParamInfo rep), + Substitute (LParamInfo rep), + Substitute (RetType rep), + Substitute (BranchType rep), + Substitute (Op rep) ) diff --git a/src/Futhark/TypeCheck.hs b/src/Futhark/TypeCheck.hs --- a/src/Futhark/TypeCheck.hs +++ b/src/Futhark/TypeCheck.hs @@ -70,14 +70,14 @@ -- | Information about an error during type checking. The 'Show' -- instance for this type produces a human-readable description. -data ErrorCase lore +data ErrorCase rep = TypeError String - | UnexpectedType (Exp lore) Type [Type] + | UnexpectedType (Exp rep) Type [Type] | ReturnTypeError Name [ExtType] [ExtType] | DupDefinitionError Name | DupParamError Name VName | DupPatternError VName - | InvalidPatternError (Pattern (Aliases lore)) [ExtType] (Maybe String) + | InvalidPatternError (Pattern (Aliases rep)) [ExtType] (Maybe String) | UnknownVariableError VName | UnknownFunctionError Name | ParameterMismatch (Maybe Name) [Type] [Type] @@ -88,7 +88,7 @@ | NotAnArray VName Type | PermutationError [Int] Int (Maybe VName) -instance Checkable lore => Show (ErrorCase lore) where +instance Checkable rep => Show (ErrorCase rep) where show (TypeError msg) = "Type error:\n" ++ msg show (UnexpectedType e _ []) = @@ -177,9 +177,9 @@ name' = maybe "" ((++ " ") . pretty) name -- | A type error. -data TypeError lore = Error [String] (ErrorCase lore) +data TypeError rep = Error [String] (ErrorCase rep) -instance Checkable lore => Show (TypeError lore) where +instance Checkable rep => Show (TypeError rep) where show (Error [] err) = show err show (Error msgs err) = @@ -187,9 +187,9 @@ -- | A tuple of a return type and a list of parameters, possibly -- named. -type FunBinding lore = ([RetType (Aliases lore)], [FParam (Aliases lore)]) +type FunBinding rep = ([RetType (Aliases rep)], [FParam (Aliases rep)]) -type VarBinding lore = NameInfo (Aliases lore) +type VarBinding rep = NameInfo (Aliases rep) data Usage = Consumed @@ -271,10 +271,10 @@ -- function table is only initialised at the very beginning, but the -- variable table will be extended during type-checking when -- let-expressions are encountered. -data Env lore = Env - { envVtable :: M.Map VName (VarBinding lore), - envFtable :: M.Map Name (FunBinding lore), - envCheckOp :: OpWithAliases (Op lore) -> TypeM lore (), +data Env rep = Env + { envVtable :: M.Map VName (VarBinding rep), + envFtable :: M.Map Name (FunBinding rep), + envCheckOp :: OpWithAliases (Op rep) -> TypeM rep (), envContext :: [String] } @@ -284,24 +284,24 @@ } -- | The type checker runs in this monad. -newtype TypeM lore a +newtype TypeM rep a = TypeM ( ReaderT - (Env lore) - (StateT TState (Either (TypeError lore))) + (Env rep) + (StateT TState (Either (TypeError rep))) a ) deriving ( Monad, Functor, Applicative, - MonadReader (Env lore), + MonadReader (Env rep), MonadState TState ) instance - Checkable lore => - HasScope (Aliases lore) (TypeM lore) + Checkable rep => + HasScope (Aliases rep) (TypeM rep) where lookupType = fmap typeOf . lookupVar askScope = asks $ M.fromList . mapMaybe varType . M.toList . envVtable @@ -309,18 +309,18 @@ varType (name, dec) = Just (name, dec) runTypeM :: - Env lore -> - TypeM lore a -> - Either (TypeError lore) (a, Consumption) + Env rep -> + TypeM rep a -> + Either (TypeError rep) (a, Consumption) runTypeM env (TypeM m) = second stateCons <$> runStateT (runReaderT m env) (TState mempty mempty) -bad :: ErrorCase lore -> TypeM lore a +bad :: ErrorCase rep -> TypeM rep a bad e = do messages <- asks envContext TypeM $ lift $ lift $ Left $ Error (reverse messages) e -tell :: Consumption -> TypeM lore () +tell :: Consumption -> TypeM rep () tell cons = modify $ \s -> s {stateCons = stateCons s <> cons} -- | Add information about what is being type-checked to the current @@ -329,8 +329,8 @@ -- 'bad'. context :: String -> - TypeM lore a -> - TypeM lore a + TypeM rep a -> + TypeM rep a context s = local $ \env -> env {envContext = s : envContext env} message :: @@ -344,35 +344,35 @@ -- | Mark a name as bound. If the name has been bound previously in -- the program, report a type error. -bound :: VName -> TypeM lore () +bound :: VName -> TypeM rep () bound name = do already_seen <- gets $ nameIn name . stateNames when already_seen $ bad $ TypeError $ "Name " ++ pretty name ++ " bound twice" modify $ \s -> s {stateNames = oneName name <> stateNames s} -occur :: Occurences -> TypeM lore () +occur :: Occurences -> TypeM rep () occur = tell . Consumption . filter (not . nullOccurence) -- | Proclaim that we have made read-only use of the given variable. -- No-op unless the variable is array-typed. observe :: - Checkable lore => + Checkable rep => VName -> - TypeM lore () + TypeM rep () observe name = do dec <- lookupVar name unless (primType $ typeOf dec) $ occur [observation $ oneName name <> aliases dec] -- | Proclaim that we have written to the given variables. -consume :: Checkable lore => Names -> TypeM lore () +consume :: Checkable rep => Names -> TypeM rep () consume als = do scope <- askScope let isArray = maybe False (not . primType . typeOf) . (`M.lookup` scope) occur [consumption $ namesFromList $ filter isArray $ namesToList als] -collectOccurences :: TypeM lore a -> TypeM lore (a, Occurences) +collectOccurences :: TypeM rep a -> TypeM rep (a, Occurences) collectOccurences m = do old <- gets stateCons modify $ \s -> s {stateCons = mempty} @@ -383,16 +383,16 @@ pure (x, o) checkOpWith :: - (OpWithAliases (Op lore) -> TypeM lore ()) -> - TypeM lore a -> - TypeM lore a + (OpWithAliases (Op rep) -> TypeM rep ()) -> + TypeM rep a -> + TypeM rep a checkOpWith checker = local $ \env -> env {envCheckOp = checker} -checkConsumption :: Consumption -> TypeM lore Occurences +checkConsumption :: Consumption -> TypeM rep Occurences checkConsumption (ConsumptionError e) = bad $ TypeError e checkConsumption (Consumption os) = return os -alternative :: TypeM lore a -> TypeM lore b -> TypeM lore (a, b) +alternative :: TypeM rep a -> TypeM rep b -> TypeM rep (a, b) alternative m1 m2 = do (x, os1) <- collectOccurences m1 (y, os2) <- collectOccurences m2 @@ -403,7 +403,7 @@ -- names is consumed, the consumption will be rewritten to be a -- consumption of the corresponding alias set. Consumption of -- anything else will result in a type error. -consumeOnlyParams :: [(VName, Names)] -> TypeM lore a -> TypeM lore a +consumeOnlyParams :: [(VName, Names)] -> TypeM rep a -> TypeM rep a consumeOnlyParams consumable m = do (x, os) <- collectOccurences m tell . Consumption =<< mapM inspect os @@ -427,7 +427,7 @@ -- | Given the immediate aliases, compute the full transitive alias -- set (including the immediate aliases). -expandAliases :: Names -> Env lore -> Names +expandAliases :: Names -> Env rep -> Names expandAliases names env = names <> aliasesOfAliases where aliasesOfAliases = mconcat . map look . namesToList $ names @@ -436,10 +436,10 @@ _ -> mempty binding :: - Checkable lore => - Scope (Aliases lore) -> - TypeM lore a -> - TypeM lore a + Checkable rep => + Scope (Aliases rep) -> + TypeM rep a -> + TypeM rep a binding bnds = check . local (`bindVars` bnds) where bindVars = M.foldlWithKey' bindVar @@ -464,14 +464,14 @@ tell $ Consumption $ unOccur (namesFromList boundnames) os return a -lookupVar :: VName -> TypeM lore (NameInfo (Aliases lore)) +lookupVar :: VName -> TypeM rep (NameInfo (Aliases rep)) lookupVar name = do bnd <- asks $ M.lookup name . envVtable case bnd of Nothing -> bad $ UnknownVariableError name Just dec -> return dec -lookupAliases :: Checkable lore => VName -> TypeM lore Names +lookupAliases :: Checkable rep => VName -> TypeM rep Names lookupAliases name = do info <- lookupVar name return $ @@ -479,19 +479,19 @@ then mempty else oneName name <> aliases info -aliases :: NameInfo (Aliases lore) -> Names +aliases :: NameInfo (Aliases rep) -> Names aliases (LetName (als, _)) = unAliases als aliases _ = mempty -subExpAliasesM :: Checkable lore => SubExp -> TypeM lore Names +subExpAliasesM :: Checkable rep => SubExp -> TypeM rep Names subExpAliasesM Constant {} = return mempty subExpAliasesM (Var v) = lookupAliases v lookupFun :: - Checkable lore => + Checkable rep => Name -> [SubExp] -> - TypeM lore ([RetType lore], [DeclType]) + TypeM rep ([RetType rep], [DeclType]) lookupFun fname args = do bnd <- asks $ M.lookup fname . envFtable case bnd of @@ -510,27 +510,27 @@ String -> Type -> Type -> - TypeM lore () + TypeM rep () checkAnnotation desc t1 t2 | t2 == t1 = return () | otherwise = bad $ BadAnnotation desc t1 t2 -- | @require ts se@ causes a '(TypeError vn)' if the type of @se@ is -- not a subtype of one of the types in @ts@. -require :: Checkable lore => [Type] -> SubExp -> TypeM lore () +require :: Checkable rep => [Type] -> SubExp -> TypeM rep () require ts se = do t <- checkSubExp se unless (t `elem` ts) $ bad $ UnexpectedType (BasicOp $ SubExp se) t ts -- | Variant of 'require' working on variable names. -requireI :: Checkable lore => [Type] -> VName -> TypeM lore () +requireI :: Checkable rep => [Type] -> VName -> TypeM rep () requireI ts ident = require ts $ Var ident checkArrIdent :: - Checkable lore => + Checkable rep => VName -> - TypeM lore Type + TypeM rep Type checkArrIdent v = do t <- lookupType v case t of @@ -538,9 +538,9 @@ _ -> bad $ NotAnArray v t checkAccIdent :: - Checkable lore => + Checkable rep => VName -> - TypeM lore (Shape, [Type]) + TypeM rep (Shape, [Type]) checkAccIdent v = do t <- lookupType v case t of @@ -556,9 +556,9 @@ -- yielding either a type error or a program with complete type -- information. checkProg :: - Checkable lore => - Prog (Aliases lore) -> - Either (TypeError lore) () + Checkable rep => + Prog (Aliases rep) -> + Either (TypeError rep) () checkProg (Prog consts funs) = do let typeenv = Env @@ -588,8 +588,8 @@ return $ M.insert name (ret, params) ftable initialFtable :: - Checkable lore => - TypeM lore (M.Map Name (FunBinding lore)) + Checkable rep => + TypeM rep (M.Map Name (FunBinding rep)) initialFtable = fmap M.fromList $ mapM addBuiltin $ M.toList builtInFunctions where addBuiltin (fname, (t, ts)) = do @@ -598,9 +598,9 @@ name = VName (nameFromString "x") 0 checkFun :: - Checkable lore => - FunDef (Aliases lore) -> - TypeM lore () + Checkable rep => + FunDef (Aliases rep) -> + TypeM rep () checkFun (FunDef _ _ fname rettype params body) = context ("In function " ++ nameToString fname) $ checkFun' @@ -621,40 +621,40 @@ ] funParamsToNameInfos :: - [FParam lore] -> - [(VName, NameInfo (Aliases lore))] -funParamsToNameInfos = map nameTypeAndLore + [FParam rep] -> + [(VName, NameInfo (Aliases rep))] +funParamsToNameInfos = map nameTypeAndDec where - nameTypeAndLore fparam = + nameTypeAndDec fparam = ( paramName fparam, FParamName $ paramDec fparam ) checkFunParams :: - Checkable lore => - [FParam lore] -> - TypeM lore () + Checkable rep => + [FParam rep] -> + TypeM rep () checkFunParams = mapM_ $ \param -> context ("In function parameter " ++ pretty param) $ - checkFParamLore (paramName param) (paramDec param) + checkFParamDec (paramName param) (paramDec param) checkLambdaParams :: - Checkable lore => - [LParam lore] -> - TypeM lore () + Checkable rep => + [LParam rep] -> + TypeM rep () checkLambdaParams = mapM_ $ \param -> context ("In lambda parameter " ++ pretty param) $ - checkLParamLore (paramName param) (paramDec param) + checkLParamDec (paramName param) (paramDec param) checkFun' :: - Checkable lore => + Checkable rep => ( Name, [DeclExtType], - [(VName, NameInfo (Aliases lore))] + [(VName, NameInfo (Aliases rep))] ) -> Maybe [(VName, Names)] -> - TypeM lore [Names] -> - TypeM lore () + TypeM rep [Names] -> + TypeM rep () checkFun' (fname, rettype, params) consumable check = do checkNoDuplicateParams binding (M.fromList params) $ @@ -698,7 +698,7 @@ zip (reverse (map uniqueness expected) ++ repeat Nonunique) $ reverse got -checkSubExp :: Checkable lore => SubExp -> TypeM lore Type +checkSubExp :: Checkable rep => SubExp -> TypeM rep Type checkSubExp (Constant val) = return $ Prim $ primValueType val checkSubExp (Var ident) = context ("In subexp " ++ pretty ident) $ do @@ -706,10 +706,10 @@ lookupType ident checkStms :: - Checkable lore => - Stms (Aliases lore) -> - TypeM lore a -> - TypeM lore a + Checkable rep => + Stms (Aliases rep) -> + TypeM rep a -> + TypeM rep a checkStms origbnds m = delve $ stmsToList origbnds where delve (stm@(Let pat _ e) : bnds) = do @@ -721,18 +721,18 @@ m checkResult :: - Checkable lore => + Checkable rep => Result -> - TypeM lore () + TypeM rep () checkResult = mapM_ checkSubExp checkFunBody :: - Checkable lore => - [RetType lore] -> - Body (Aliases lore) -> - TypeM lore [Names] -checkFunBody rt (Body (_, lore) bnds res) = do - checkBodyLore lore + Checkable rep => + [RetType rep] -> + Body (Aliases rep) -> + TypeM rep [Names] +checkFunBody rt (Body (_, rep) bnds res) = do + checkBodyDec rep checkStms bnds $ do context "When checking body result" $ checkResult res context "When matching declared return type to result of body" $ @@ -742,12 +742,12 @@ bound_here = namesFromList $ M.keys $ scopeOf bnds checkLambdaBody :: - Checkable lore => + Checkable rep => [Type] -> - Body (Aliases lore) -> - TypeM lore [Names] -checkLambdaBody ret (Body (_, lore) bnds res) = do - checkBodyLore lore + Body (Aliases rep) -> + TypeM rep [Names] +checkLambdaBody ret (Body (_, rep) bnds res) = do + checkBodyDec rep checkStms bnds $ do checkLambdaResult ret res map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res @@ -755,10 +755,10 @@ bound_here = namesFromList $ M.keys $ scopeOf bnds checkLambdaResult :: - Checkable lore => + Checkable rep => [Type] -> Result -> - TypeM lore () + TypeM rep () checkLambdaResult ts es | length ts /= length es = bad $ @@ -780,18 +780,18 @@ ++ pretty t checkBody :: - Checkable lore => - Body (Aliases lore) -> - TypeM lore [Names] -checkBody (Body (_, lore) bnds res) = do - checkBodyLore lore + Checkable rep => + Body (Aliases rep) -> + TypeM rep [Names] +checkBody (Body (_, rep) bnds res) = do + checkBodyDec rep checkStms bnds $ do checkResult res map (`namesSubtract` bound_here) <$> mapM subExpAliasesM res where bound_here = namesFromList $ M.keys $ scopeOf bnds -checkBasicOp :: Checkable lore => BasicOp -> TypeM lore () +checkBasicOp :: Checkable rep => BasicOp -> TypeM rep () checkBasicOp (SubExp es) = void $ checkSubExp es checkBasicOp (Opaque es) = @@ -934,11 +934,11 @@ consume =<< lookupAliases acc matchLoopResultExt :: - Checkable lore => + Checkable rep => [Param DeclType] -> [Param DeclType] -> [SubExp] -> - TypeM lore () + TypeM rep () matchLoopResultExt ctx val loopres = do let rettype_ext = existentialiseExtTypes (map paramName ctx) $ @@ -962,9 +962,9 @@ (staticShapes bodyt) checkExp :: - Checkable lore => - Exp (Aliases lore) -> - TypeM lore () + Checkable rep => + Exp (Aliases rep) -> + TypeM rep () checkExp (BasicOp op) = checkBasicOp op checkExp (If e1 e2 e3 info) = do require [Prim Bool] e1 @@ -1016,7 +1016,7 @@ (Just consumable) $ do checkFunParams mergepat - checkBodyLore $ snd $ bodyDec loopbody + checkBodyDec $ snd $ bodyDec loopbody checkStms (bodyStms loopbody) $ do checkResult $ bodyResult loopbody @@ -1037,7 +1037,7 @@ observe a case peelArray 1 a_t of Just a_t_r -> do - checkLParamLore (paramName p) $ paramDec p + checkLParamDec (paramName p) $ paramDec p unless (a_t_r `subtypeOf` typeOf (paramDec p)) $ bad $ TypeError $ @@ -1127,10 +1127,10 @@ checker op checkSOACArrayArgs :: - Checkable lore => + Checkable rep => SubExp -> [VName] -> - TypeM lore [Arg] + TypeM rep [Arg] checkSOACArrayArgs width = mapM checkSOACArrayArg where checkSOACArrayArg v = do @@ -1151,9 +1151,9 @@ "SOAC argument " ++ pretty v ++ " is not an array" checkType :: - Checkable lore => + Checkable rep => TypeBase Shape u -> - TypeM lore () + TypeM rep () checkType (Mem (ScalarSpace d _)) = mapM_ (require [Prim int64]) d checkType (Acc cert shape ts _) = do requireI [Prim Unit] cert @@ -1162,20 +1162,20 @@ checkType t = mapM_ checkSubExp $ arrayDims t checkExtType :: - Checkable lore => + Checkable rep => TypeBase ExtShape u -> - TypeM lore () + TypeM rep () checkExtType = mapM_ checkExtDim . shapeDims . arrayShape where checkExtDim (Free se) = void $ checkSubExp se checkExtDim (Ext _) = return () checkCmpOp :: - Checkable lore => + Checkable rep => CmpOp -> SubExp -> SubExp -> - TypeM lore () + TypeM rep () checkCmpOp (CmpEq t) x y = do require [Prim t] x require [Prim t] y @@ -1189,38 +1189,38 @@ checkCmpOp CmpLle x y = checkBinOpArgs Bool x y checkBinOpArgs :: - Checkable lore => + Checkable rep => PrimType -> SubExp -> SubExp -> - TypeM lore () + TypeM rep () checkBinOpArgs t e1 e2 = do require [Prim t] e1 require [Prim t] e2 checkPatElem :: - Checkable lore => - PatElemT (LetDec lore) -> - TypeM lore () + Checkable rep => + PatElemT (LetDec rep) -> + TypeM rep () checkPatElem (PatElem name dec) = context ("When checking pattern element " ++ pretty name) $ - checkLetBoundLore name dec + checkLetBoundDec name dec checkDimIndex :: - Checkable lore => + Checkable rep => DimIndex SubExp -> - TypeM lore () + TypeM rep () checkDimIndex (DimFix i) = require [Prim int64] i checkDimIndex (DimSlice i n s) = mapM_ (require [Prim int64]) [i, n, s] checkStm :: - Checkable lore => - Stm (Aliases lore) -> - TypeM lore a -> - TypeM lore a + Checkable rep => + Stm (Aliases rep) -> + TypeM rep a -> + TypeM rep a checkStm stm@(Let pat (StmAux (Certificates cs) _ (_, dec)) e) m = do context "When checking certificates" $ mapM_ (requireI [Prim Unit]) cs - context "When checking expression annotation" $ checkExpLore dec + context "When checking expression annotation" $ checkExpDec dec context ("When matching\n" ++ message " " pat ++ "\nwith\n" ++ message " " e) $ matchPattern pat e binding (maybeWithoutAliases $ scopeOf stm) $ do @@ -1242,37 +1242,37 @@ withoutAliases info = info matchExtPattern :: - Checkable lore => - Pattern (Aliases lore) -> + Checkable rep => + Pattern (Aliases rep) -> [ExtType] -> - TypeM lore () + TypeM rep () matchExtPattern pat ts = unless (expExtTypesFromPattern pat == ts) $ bad $ InvalidPatternError pat ts Nothing matchExtReturnType :: - Checkable lore => + Checkable rep => [ExtType] -> Result -> - TypeM lore () + TypeM rep () matchExtReturnType rettype res = do ts <- mapM subExpType res matchExtReturns rettype res ts matchExtBranchType :: - Checkable lore => + Checkable rep => [ExtType] -> - Body (Aliases lore) -> - TypeM lore () + Body (Aliases rep) -> + TypeM rep () matchExtBranchType rettype (Body _ stms res) = do ts <- extendedScope (traverse subExpType res) stmscope matchExtReturns rettype res ts where stmscope = scopeOf stms -matchExtReturns :: [ExtType] -> Result -> [Type] -> TypeM lore () +matchExtReturns :: [ExtType] -> Result -> [Type] -> TypeM rep () matchExtReturns rettype res ts = do - let problem :: TypeM lore a + let problem :: TypeM rep a problem = bad $ TypeError $ @@ -1337,9 +1337,9 @@ noArgAliases (t, _) = (t, mempty) checkArg :: - Checkable lore => + Checkable rep => SubExp -> - TypeM lore Arg + TypeM rep Arg checkArg arg = do argt <- checkSubExp arg als <- subExpAliasesM arg @@ -1349,7 +1349,7 @@ Maybe Name -> [DeclType] -> [Arg] -> - TypeM lore () + TypeM rep () checkFuncall fname paramts args = do let argts = map argType args unless (validApply paramts argts) $ @@ -1359,7 +1359,7 @@ consumeArgs :: [DeclType] -> [Arg] -> - TypeM lore () + TypeM rep () consumeArgs paramts args = forM_ (zip (map diet paramts) args) $ \(d, (_, als)) -> occur [consumption (consumeArg als d)] @@ -1370,7 +1370,7 @@ -- The boolean indicates whether we only allow consumption of -- parameters. checkAnyLambda :: - Checkable lore => Bool -> Lambda (Aliases lore) -> [Arg] -> TypeM lore () + Checkable rep => Bool -> Lambda (Aliases rep) -> [Arg] -> TypeM rep () checkAnyLambda soac (Lambda params body rettype) args = do let fname = nameFromString "" if length params == length args @@ -1407,10 +1407,10 @@ ++ show (length args) ++ " arguments." -checkLambda :: Checkable lore => Lambda (Aliases lore) -> [Arg] -> TypeM lore () +checkLambda :: Checkable rep => Lambda (Aliases rep) -> [Arg] -> TypeM rep () checkLambda = checkAnyLambda True -checkPrimExp :: Checkable lore => PrimExp VName -> TypeM lore () +checkPrimExp :: Checkable rep => PrimExp VName -> TypeM rep () checkPrimExp ValueExp {} = return () checkPrimExp (LeafExp v pt) = requireI [Prim pt] v checkPrimExp (BinOpExp op x y) = do @@ -1442,7 +1442,7 @@ ++ pretty h_ret zipWithM_ requirePrimExp h_ts args -requirePrimExp :: Checkable lore => PrimType -> PrimExp VName -> TypeM lore () +requirePrimExp :: Checkable rep => PrimType -> PrimExp VName -> TypeM rep () requirePrimExp t e = context ("in PrimExp " ++ pretty e) $ do checkPrimExp e unless (primExpType e == t) $ @@ -1450,62 +1450,62 @@ TypeError $ pretty e ++ " must have type " ++ pretty t -class ASTLore lore => CheckableOp lore where - checkOp :: OpWithAliases (Op lore) -> TypeM lore () +class ASTRep rep => CheckableOp rep where + checkOp :: OpWithAliases (Op rep) -> TypeM rep () -- ^ Used at top level; can be locally changed with 'checkOpWith'. --- | The class of lores that can be type-checked. -class (ASTLore lore, CanBeAliased (Op lore), CheckableOp lore) => Checkable lore where - checkExpLore :: ExpDec lore -> TypeM lore () - checkBodyLore :: BodyDec lore -> TypeM lore () - checkFParamLore :: VName -> FParamInfo lore -> TypeM lore () - checkLParamLore :: VName -> LParamInfo lore -> TypeM lore () - checkLetBoundLore :: VName -> LetDec lore -> TypeM lore () - checkRetType :: [RetType lore] -> TypeM lore () - matchPattern :: Pattern (Aliases lore) -> Exp (Aliases lore) -> TypeM lore () - primFParam :: VName -> PrimType -> TypeM lore (FParam (Aliases lore)) - matchReturnType :: [RetType lore] -> Result -> TypeM lore () - matchBranchType :: [BranchType lore] -> Body (Aliases lore) -> TypeM lore () +-- | The class of representations that can be type-checked. +class (ASTRep rep, CanBeAliased (Op rep), CheckableOp rep) => Checkable rep where + checkExpDec :: ExpDec rep -> TypeM rep () + checkBodyDec :: BodyDec rep -> TypeM rep () + checkFParamDec :: VName -> FParamInfo rep -> TypeM rep () + checkLParamDec :: VName -> LParamInfo rep -> TypeM rep () + checkLetBoundDec :: VName -> LetDec rep -> TypeM rep () + checkRetType :: [RetType rep] -> TypeM rep () + matchPattern :: Pattern (Aliases rep) -> Exp (Aliases rep) -> TypeM rep () + primFParam :: VName -> PrimType -> TypeM rep (FParam (Aliases rep)) + matchReturnType :: [RetType rep] -> Result -> TypeM rep () + matchBranchType :: [BranchType rep] -> Body (Aliases rep) -> TypeM rep () matchLoopResult :: - [FParam (Aliases lore)] -> - [FParam (Aliases lore)] -> + [FParam (Aliases rep)] -> + [FParam (Aliases rep)] -> [SubExp] -> - TypeM lore () + TypeM rep () - default checkExpLore :: ExpDec lore ~ () => ExpDec lore -> TypeM lore () - checkExpLore = return + default checkExpDec :: ExpDec rep ~ () => ExpDec rep -> TypeM rep () + checkExpDec = return - default checkBodyLore :: BodyDec lore ~ () => BodyDec lore -> TypeM lore () - checkBodyLore = return + default checkBodyDec :: BodyDec rep ~ () => BodyDec rep -> TypeM rep () + checkBodyDec = return - default checkFParamLore :: FParamInfo lore ~ DeclType => VName -> FParamInfo lore -> TypeM lore () - checkFParamLore _ = checkType + default checkFParamDec :: FParamInfo rep ~ DeclType => VName -> FParamInfo rep -> TypeM rep () + checkFParamDec _ = checkType - default checkLParamLore :: LParamInfo lore ~ Type => VName -> LParamInfo lore -> TypeM lore () - checkLParamLore _ = checkType + default checkLParamDec :: LParamInfo rep ~ Type => VName -> LParamInfo rep -> TypeM rep () + checkLParamDec _ = checkType - default checkLetBoundLore :: LetDec lore ~ Type => VName -> LetDec lore -> TypeM lore () - checkLetBoundLore _ = checkType + default checkLetBoundDec :: LetDec rep ~ Type => VName -> LetDec rep -> TypeM rep () + checkLetBoundDec _ = checkType - default checkRetType :: RetType lore ~ DeclExtType => [RetType lore] -> TypeM lore () + default checkRetType :: RetType rep ~ DeclExtType => [RetType rep] -> TypeM rep () checkRetType = mapM_ $ checkExtType . declExtTypeOf - default matchPattern :: Pattern (Aliases lore) -> Exp (Aliases lore) -> TypeM lore () + default matchPattern :: Pattern (Aliases rep) -> Exp (Aliases rep) -> TypeM rep () matchPattern pat = matchExtPattern pat <=< expExtType - default primFParam :: FParamInfo lore ~ DeclType => VName -> PrimType -> TypeM lore (FParam (Aliases lore)) + default primFParam :: FParamInfo rep ~ DeclType => VName -> PrimType -> TypeM rep (FParam (Aliases rep)) primFParam name t = return $ Param name (Prim t) - default matchReturnType :: RetType lore ~ DeclExtType => [RetType lore] -> Result -> TypeM lore () + default matchReturnType :: RetType rep ~ DeclExtType => [RetType rep] -> Result -> TypeM rep () matchReturnType = matchExtReturnType . map fromDecl - default matchBranchType :: BranchType lore ~ ExtType => [BranchType lore] -> Body (Aliases lore) -> TypeM lore () + default matchBranchType :: BranchType rep ~ ExtType => [BranchType rep] -> Body (Aliases rep) -> TypeM rep () matchBranchType = matchExtBranchType default matchLoopResult :: - FParamInfo lore ~ DeclType => - [FParam (Aliases lore)] -> - [FParam (Aliases lore)] -> + FParamInfo rep ~ DeclType => + [FParam (Aliases rep)] -> + [FParam (Aliases rep)] -> [SubExp] -> - TypeM lore () + TypeM rep () matchLoopResult = matchLoopResultExt diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs --- a/src/Futhark/Util.hs +++ b/src/Futhark/Util.hs @@ -50,22 +50,30 @@ EncodedString, zEncodeString, atMostChars, + invertMap, ) where +import Control.Arrow (first) import Control.Concurrent import Control.Exception import Control.Monad import qualified Data.ByteString as BS import Data.Char import Data.Either +import Data.Function ((&)) import Data.List (foldl', genericDrop, genericSplitAt, sort) import qualified Data.List.NonEmpty as NE +import Data.Map (Map) +import qualified Data.Map as Map import Data.Maybe +import Data.Set (Set) +import qualified Data.Set as Set import qualified Data.Text as T import qualified Data.Text.Encoding as T import qualified Data.Text.Encoding.Error as T import qualified Data.Text.IO as T +import Data.Tuple (swap) import Numeric import qualified System.Directory.Tree as Dir import System.Environment @@ -431,3 +439,9 @@ atMostChars n s | length s > n = take (n -3) s ++ "..." | otherwise = s + +invertMap :: (Ord v, Ord k) => Map k v -> Map v (Set k) +invertMap m = + Map.toList m + & fmap (swap . first Set.singleton) + & foldr (uncurry $ Map.insertWith (<>)) mempty diff --git a/src/Futhark/Util/Loc.hs b/src/Futhark/Util/Loc.hs --- a/src/Futhark/Util/Loc.hs +++ b/src/Futhark/Util/Loc.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE Trustworthy #-} - -- | A Safe Haskell-trusted re-export of the @srcloc@ package. module Futhark.Util.Loc (module Data.Loc) where diff --git a/src/Futhark/Util/Pretty.hs b/src/Futhark/Util/Pretty.hs --- a/src/Futhark/Util/Pretty.hs +++ b/src/Futhark/Util/Pretty.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE Trustworthy #-} {-# OPTIONS_GHC -fno-warn-orphans #-} -- | A re-export of the prettyprinting library, along with some convenience functions. diff --git a/src/Language/Futhark/Core.hs b/src/Language/Futhark/Core.hs --- a/src/Language/Futhark/Core.hs +++ b/src/Language/Futhark/Core.hs @@ -7,7 +7,6 @@ -- representation. module Language.Futhark.Core ( Uniqueness (..), - Commutativity (..), -- * Location utilities SrcLoc, @@ -28,7 +27,6 @@ baseTag, baseName, baseString, - pretty, quote, pquote, @@ -75,19 +73,6 @@ instance Pretty Uniqueness where ppr Unique = star ppr Nonunique = empty - --- | Whether some operator is commutative or not. The 'Monoid' --- instance returns the least commutative of its arguments. -data Commutativity - = Noncommutative - | Commutative - deriving (Eq, Ord, Show) - -instance Semigroup Commutativity where - (<>) = min - -instance Monoid Commutativity where - mempty = Commutative -- | The name of the default program entry point (main). defaultEntryPoint :: Name diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs --- a/src/Language/Futhark/Syntax.hs +++ b/src/Language/Futhark/Syntax.hs @@ -12,6 +12,7 @@ -- module may be a little hard to understand. module Language.Futhark.Syntax ( module Language.Futhark.Core, + pretty, -- * Types Uniqueness (..), diff --git a/unittests/Futhark/IR/Syntax/CoreTests.hs b/unittests/Futhark/IR/Syntax/CoreTests.hs --- a/unittests/Futhark/IR/Syntax/CoreTests.hs +++ b/unittests/Futhark/IR/Syntax/CoreTests.hs @@ -4,7 +4,7 @@ module Futhark.IR.Syntax.CoreTests (tests) where import Control.Applicative -import Futhark.IR.Pretty () +import Futhark.IR.Pretty (pretty) import Futhark.IR.PrimitiveTests () import Futhark.IR.Syntax.Core import Language.Futhark.CoreTests () diff --git a/unittests/Futhark/Optimise/ReuseAllocations/GreedyColoringTests.hs b/unittests/Futhark/Optimise/ReuseAllocations/GreedyColoringTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Optimise/ReuseAllocations/GreedyColoringTests.hs @@ -0,0 +1,70 @@ +module Futhark.Optimise.ReuseAllocations.GreedyColoringTests + ( tests, + ) +where + +import Control.Arrow ((***)) +import Data.Function ((&)) +import qualified Data.Map as Map +import qualified Data.Set as Set +import qualified Futhark.Optimise.ReuseAllocations.GreedyColoring as GreedyColoring +import Test.Tasty +import Test.Tasty.HUnit + +tests :: TestTree +tests = + testGroup + "GreedyColoringTests" + [psumTest, allIntersect, emptyGraph, noIntersections, differentSpaces] + +psumTest :: TestTree +psumTest = + testCase "psumTest" $ + assertEqual + "Color simple 1-2-3 using two colors" + ([(0, "local"), (1, "local")], [(1 :: Int, 0), (2, 1), (3, 0)]) + $ (Map.toList *** Map.toList) $ + GreedyColoring.colorGraph + (Map.fromList [(1, "local"), (2, "local"), (3, "local")]) + $ Set.fromList [(1, 2), (2, 3)] + +allIntersect :: TestTree +allIntersect = + testCase "allIntersect" $ + assertEqual + "Color a graph where all values intersect" + ([(0, "local"), (1, "local"), (2, "local")], [(1 :: Int, 2), (2, 1), (3, 0)]) + $ (Map.toList *** Map.toList) $ + GreedyColoring.colorGraph + (Map.fromList [(1, "local"), (2, "local"), (3, "local")]) + $ Set.fromList [(1, 2), (2, 3), (1, 3)] + +emptyGraph :: TestTree +emptyGraph = + testCase "emptyGraph" $ + assertEqual + "Color an empty graph" + ([] :: [(Int, Char)], [] :: [(Int, Int)]) + $ (Map.toList *** Map.toList) $ GreedyColoring.colorGraph Map.empty $ Set.fromList [] + +noIntersections :: TestTree +noIntersections = + GreedyColoring.colorGraph + (Map.fromList [(1, "local"), (2, "local"), (3, "local")]) + (Set.fromList []) + & Map.toList *** Map.toList + & assertEqual + "Color nodes with no intersections" + ([(0, "local")], [(1, 0), (2, 0), (3, 0)] :: [(Int, Int)]) + & testCase "noIntersections" + +differentSpaces :: TestTree +differentSpaces = + GreedyColoring.colorGraph + (Map.fromList [(1, "a"), (2, "b"), (3, "c")]) + (Set.fromList []) + & Map.toList *** Map.toList + & assertEqual + "Color nodes with no intersections but in different spaces" + ([(0, "c"), (1, "b"), (2, "a")], [(1, 2), (2, 1), (3, 0)] :: [(Int, Int)]) + & testCase "differentSpaces" diff --git a/unittests/futhark_tests.hs b/unittests/futhark_tests.hs --- a/unittests/futhark_tests.hs +++ b/unittests/futhark_tests.hs @@ -5,6 +5,7 @@ import qualified Futhark.IR.PrimitiveTests import qualified Futhark.IR.PropTests import qualified Futhark.IR.Syntax.CoreTests +import qualified Futhark.Optimise.ReuseAllocations.GreedyColoringTests import qualified Futhark.Pkg.SolveTests import qualified Language.Futhark.SyntaxTests import Test.Tasty @@ -19,7 +20,8 @@ Futhark.IR.Syntax.CoreTests.tests, Futhark.Pkg.SolveTests.tests, Futhark.IR.Mem.IxFunTests.tests, - Futhark.IR.PrimitiveTests.tests + Futhark.IR.PrimitiveTests.tests, + Futhark.Optimise.ReuseAllocations.GreedyColoringTests.tests ] main :: IO ()