diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,17 @@ +ISC License + +Copyright (c) 2013-2018. DIKU, University of Copenhagen + +Permission to use, copy, modify, and/or distribute this software for +any purpose with or without fee is hereby granted, provided that the +above copyright notice and this permission notice appear in all +copies. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL +WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE +AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL +DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR +PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER +TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,7 @@ +#!/usr/bin/env runhaskell + +import Distribution.Simple + +main :: IO () +main = defaultMainWithHooks myHooks + where myHooks = simpleUserHooks diff --git a/futhark.cabal b/futhark.cabal new file mode 100644 --- /dev/null +++ b/futhark.cabal @@ -0,0 +1,1133 @@ +-- This file has been generated from package.yaml by hpack version 0.28.2. +-- +-- see: https://github.com/sol/hpack +-- +-- hash: 3b4f60b44c2ac9e9de368d6116299bafb8b535f9a6ae7a26613bca65b22ad1d9 + +name: futhark +version: 0.7.3 +synopsis: An optimising compiler for a functional, array-oriented language. +description: See the website at https://futhark-lang.org +category: Language +homepage: https://futhark-lang.org +bug-reports: https://github.com/diku-dk/futhark/issues +maintainer: Troels Henriksen athas@sigkill.dk +license: ISC +license-file: LICENSE +build-type: Simple +cabal-version: >= 1.10 +extra-source-files: + futlib/array.fut + futlib/functional.fut + futlib/math.fut + futlib/prelude.fut + futlib/soacs.fut + futlib/zip.fut + rts/c/lock.h + rts/c/opencl.h + rts/c/panic.h + rts/c/timing.h + rts/c/values.h + rts/csharp/exceptions.cs + rts/csharp/functions.cs + rts/csharp/memory.cs + rts/csharp/memory_opencl.cs + rts/csharp/opencl.cs + rts/csharp/panic.cs + rts/csharp/reader.cs + rts/csharp/scalar.cs + rts/python/__init__.py + rts/python/memory.py + rts/python/opencl.py + rts/python/panic.py + rts/python/scalar.py + rts/python/values.py + +source-repository head + type: git + location: https://github.com/diku-dk/futhark + +library + hs-source-dirs: + src + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + build-tools: + alex + , happy + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + exposed-modules: + Futhark.Actions + Futhark.Analysis.AlgSimplify + Futhark.Analysis.Alias + Futhark.Analysis.CallGraph + Futhark.Analysis.DataDependencies + Futhark.Analysis.HORepresentation.MapNest + Futhark.Analysis.HORepresentation.SOAC + Futhark.Analysis.Metrics + Futhark.Analysis.PrimExp + Futhark.Analysis.PrimExp.Convert + Futhark.Analysis.PrimExp.Simplify + Futhark.Analysis.Range + Futhark.Analysis.Rephrase + Futhark.Analysis.ScalExp + Futhark.Analysis.SymbolTable + Futhark.Analysis.Usage + Futhark.Analysis.UsageTable + Futhark.Binder + Futhark.Binder.Class + Futhark.CodeGen.Backends.COpenCL + Futhark.CodeGen.Backends.COpenCL.Boilerplate + Futhark.CodeGen.Backends.CSOpenCL + Futhark.CodeGen.Backends.CSOpenCL.Boilerplate + Futhark.CodeGen.Backends.GenericC + Futhark.CodeGen.Backends.GenericC.Options + Futhark.CodeGen.Backends.GenericCSharp + Futhark.CodeGen.Backends.GenericCSharp.AST + Futhark.CodeGen.Backends.GenericCSharp.Definitions + Futhark.CodeGen.Backends.GenericCSharp.Options + Futhark.CodeGen.Backends.GenericPython + Futhark.CodeGen.Backends.GenericPython.AST + Futhark.CodeGen.Backends.GenericPython.Definitions + Futhark.CodeGen.Backends.GenericPython.Options + Futhark.CodeGen.Backends.PyOpenCL + Futhark.CodeGen.Backends.PyOpenCL.Boilerplate + Futhark.CodeGen.Backends.SequentialC + Futhark.CodeGen.Backends.SequentialCSharp + Futhark.CodeGen.Backends.SequentialPython + Futhark.CodeGen.Backends.SimpleRepresentation + Futhark.CodeGen.ImpCode + Futhark.CodeGen.ImpCode.Kernels + Futhark.CodeGen.ImpCode.OpenCL + Futhark.CodeGen.ImpCode.Sequential + Futhark.CodeGen.ImpGen + Futhark.CodeGen.ImpGen.Kernels + Futhark.CodeGen.ImpGen.Kernels.ToOpenCL + Futhark.CodeGen.ImpGen.OpenCL + Futhark.CodeGen.ImpGen.Sequential + Futhark.CodeGen.OpenCL.Kernels + Futhark.CodeGen.SetDefaultSpace + Futhark.Compiler + Futhark.Compiler.CLI + Futhark.Compiler.Program + Futhark.Construct + Futhark.Doc.Generator + Futhark.Doc.Html + Futhark.Error + Futhark.FreshNames + Futhark.Internalise + Futhark.Internalise.AccurateSizes + Futhark.Internalise.Bindings + Futhark.Internalise.Defunctionalise + Futhark.Internalise.Defunctorise + Futhark.Internalise.Lambdas + Futhark.Internalise.Monad + Futhark.Internalise.Monomorphise + Futhark.Internalise.TypesValues + Futhark.MonadFreshNames + Futhark.Optimise.CSE + Futhark.Optimise.DoubleBuffer + Futhark.Optimise.Fusion + Futhark.Optimise.Fusion.Composing + Futhark.Optimise.Fusion.LoopKernel + Futhark.Optimise.Fusion.TryFusion + Futhark.Optimise.InliningDeadFun + Futhark.Optimise.InPlaceLowering + Futhark.Optimise.InPlaceLowering.LowerIntoStm + Futhark.Optimise.InPlaceLowering.SubstituteIndices + Futhark.Optimise.MemoryBlockMerging + Futhark.Optimise.MemoryBlockMerging.ActualVariables + Futhark.Optimise.MemoryBlockMerging.AllExpVars + Futhark.Optimise.MemoryBlockMerging.AuxiliaryInfo + Futhark.Optimise.MemoryBlockMerging.Coalescing + Futhark.Optimise.MemoryBlockMerging.Coalescing.AllocationMovingUp + Futhark.Optimise.MemoryBlockMerging.Coalescing.Core + Futhark.Optimise.MemoryBlockMerging.Coalescing.Exps + Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition2 + Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition3 + Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition5 + Futhark.Optimise.MemoryBlockMerging.CrudeMovingUp + Futhark.Optimise.MemoryBlockMerging.Existentials + Futhark.Optimise.MemoryBlockMerging.Liveness.FirstUse + Futhark.Optimise.MemoryBlockMerging.Liveness.Interference + Futhark.Optimise.MemoryBlockMerging.Liveness.LastUse + Futhark.Optimise.MemoryBlockMerging.MemoryAliases + Futhark.Optimise.MemoryBlockMerging.MemoryUpdater + Futhark.Optimise.MemoryBlockMerging.Miscellaneous + Futhark.Optimise.MemoryBlockMerging.PrimExps + Futhark.Optimise.MemoryBlockMerging.Reuse + Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizeMovingUp + Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizes + Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizeUses + Futhark.Optimise.MemoryBlockMerging.Reuse.Core + Futhark.Optimise.MemoryBlockMerging.Types + Futhark.Optimise.MemoryBlockMerging.VariableAliases + Futhark.Optimise.MemoryBlockMerging.VariableMemory + Futhark.Optimise.Simplify + Futhark.Optimise.Simplify.ClosedForm + Futhark.Optimise.Simplify.Engine + Futhark.Optimise.Simplify.Lore + Futhark.Optimise.Simplify.Rule + Futhark.Optimise.Simplify.Rules + Futhark.Optimise.TileLoops + Futhark.Optimise.Unstream + Futhark.Pass + Futhark.Pass.ExpandAllocations + Futhark.Pass.ExplicitAllocations + Futhark.Pass.ExtractKernels + Futhark.Pass.ExtractKernels.BlockedKernel + Futhark.Pass.ExtractKernels.Distribution + Futhark.Pass.ExtractKernels.Interchange + Futhark.Pass.ExtractKernels.Intragroup + Futhark.Pass.ExtractKernels.ISRWIM + Futhark.Pass.ExtractKernels.Kernelise + Futhark.Pass.ExtractKernels.Segmented + Futhark.Pass.FirstOrderTransform + Futhark.Pass.KernelBabysitting + Futhark.Pass.ResolveAssertions + Futhark.Pass.Simplify + Futhark.Passes + Futhark.Pipeline + Futhark.Pkg.Info + Futhark.Pkg.Solve + Futhark.Pkg.Types + Futhark.Representation.Aliases + Futhark.Representation.AST + Futhark.Representation.AST.Annotations + Futhark.Representation.AST.Attributes + Futhark.Representation.AST.Attributes.Aliases + Futhark.Representation.AST.Attributes.Constants + Futhark.Representation.AST.Attributes.Names + Futhark.Representation.AST.Attributes.Patterns + Futhark.Representation.AST.Attributes.Ranges + Futhark.Representation.AST.Attributes.Rearrange + Futhark.Representation.AST.Attributes.Reshape + Futhark.Representation.AST.Attributes.Scope + Futhark.Representation.AST.Attributes.TypeOf + Futhark.Representation.AST.Attributes.Types + Futhark.Representation.AST.Pretty + Futhark.Representation.AST.RetType + Futhark.Representation.AST.Syntax + Futhark.Representation.AST.Syntax.Core + Futhark.Representation.AST.Traversals + Futhark.Representation.ExplicitMemory + Futhark.Representation.ExplicitMemory.IndexFunction + Futhark.Representation.ExplicitMemory.Lmad + Futhark.Representation.ExplicitMemory.Simplify + Futhark.Representation.Kernels + Futhark.Representation.Kernels.Kernel + Futhark.Representation.Kernels.KernelExp + Futhark.Representation.Kernels.Simplify + Futhark.Representation.Kernels.Sizes + Futhark.Representation.Primitive + Futhark.Representation.Ranges + Futhark.Representation.SOACS + Futhark.Representation.SOACS.Simplify + Futhark.Representation.SOACS.SOAC + Futhark.Test + Futhark.Test.Values + Futhark.Tools + Futhark.Transform.CopyPropagate + Futhark.Transform.FirstOrderTransform + Futhark.Transform.Rename + Futhark.Transform.Substitute + Futhark.TypeCheck + Futhark.Util + Futhark.Util.IntegralExp + Futhark.Util.Log + Futhark.Util.Options + Futhark.Util.Pretty + Futhark.Util.Table + Futhark.Version + Language.Futhark + Language.Futhark.Attributes + Language.Futhark.Core + Language.Futhark.Futlib + Language.Futhark.Interpreter + Language.Futhark.Parser + Language.Futhark.Pretty + Language.Futhark.Semantic + Language.Futhark.Syntax + Language.Futhark.Traversals + Language.Futhark.TypeChecker + Language.Futhark.TypeChecker.Monad + Language.Futhark.TypeChecker.Terms + Language.Futhark.TypeChecker.Types + Language.Futhark.TypeChecker.Unify + Language.Futhark.Warnings + other-modules: + Language.Futhark.Parser.Parser + Language.Futhark.Parser.Lexer + Paths_futhark + default-language: Haskell2010 + +executable futhark + main-is: src/futhark.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-bench + main-is: src/futhark-bench.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-c + main-is: src/futhark-c.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-cs + main-is: src/futhark-cs.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-csopencl + main-is: src/futhark-csopencl.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-dataset + main-is: src/futhark-dataset.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-doc + main-is: src/futhark-doc.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-opencl + main-is: src/futhark-opencl.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-pkg + main-is: src/futhark-pkg.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-py + main-is: src/futhark-py.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-pyopencl + main-is: src/futhark-pyopencl.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futhark-test + main-is: src/futhark-test.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , json + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , random + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , temporary + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +executable futharki + main-is: src/futharki.hs + other-modules: + Paths_futhark + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" + build-depends: + ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , haskeline + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , template-haskell >=2.11.1 + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + default-language: Haskell2010 + +test-suite unit + type: exitcode-stdio-1.0 + main-is: futhark_tests.hs + hs-source-dirs: + unittests + ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists + build-depends: + HUnit + , QuickCheck >=2.8 + , ansi-terminal >=0.6.3.1 + , array >=0.4 + , base >=4 && <5 + , bifunctors >=5.4.2 + , binary >=0.8.3 + , blaze-html >=0.9.0.1 + , bytestring >=0.10.8 + , containers >=0.5 + , data-binary-ieee754 >=0.1 + , directory >=1.3.0.0 + , directory-tree >=0.12.1 + , dlist >=0.6.0.1 + , extra >=1.5.3 + , file-embed >=0.0.9 + , filepath >=1.4.1.1 + , free >=4.12.4 + , futhark + , gitrev >=1.2.0 + , http-client >=0.5.7.0 + , http-client-tls >=0.3.5.1 + , http-conduit >=2.2.4 + , language-c-quote >=0.12 + , mainland-pretty >=0.6.1 + , markdown >=0.1.16 + , megaparsec >=7.0.1 + , mtl >=2.2.1 + , neat-interpolation >=0.3 + , parallel >=3.2.1.0 + , parser-combinators >=1.0.0 + , process >=1.4.3.0 + , process-extras >=0.7.2 + , raw-strings-qq >=1.1 + , regex-tdfa >=1.2 + , srcloc >=0.4 + , tasty + , tasty-hunit + , tasty-quickcheck + , template-haskell >=2.11.1 + , text >=1.2.2.2 + , th-lift-instances >=0.1.11 + , time >=1.6.0.1 + , transformers >=0.3 + , vector >=0.12 + , vector-binary-instances >=0.2.2.0 + , versions >=3.3.1 + , zip-archive >=0.3.1.1 + , zlib >=0.6.1.2 + if !impl(ghc >= 8.0) + build-depends: + semigroups ==0.18.* + other-modules: + Futhark.Analysis.ScalExpTests + Futhark.Optimise.AlgSimplifyTests + Futhark.Pkg.SolveTests + Futhark.Representation.AST.Attributes.RearrangeTests + Futhark.Representation.AST.Attributes.ReshapeTests + Futhark.Representation.AST.AttributesTests + Futhark.Representation.AST.Syntax.CoreTests + Futhark.Representation.AST.SyntaxTests + Futhark.Representation.PrimitiveTests + Language.Futhark.CoreTests + Language.Futhark.SyntaxTests + Paths_futhark + default-language: Haskell2010 diff --git a/futlib/array.fut b/futlib/array.fut new file mode 100644 --- /dev/null +++ b/futlib/array.fut @@ -0,0 +1,144 @@ +-- | Utility functions for arrays. + +import "math" +import "soacs" +import "functional" +open import "zip" -- Rexport. + +-- | The size of the outer dimension of an array. +let length [n] 't (_: [n]t) = n + +-- | Is the array empty? +let null [n] 't (_: [n]t) = n == 0 + +-- | The first element of the array. +let head [n] 't (x: [n]t) = x[0] + +-- | The last element of the array. +let last [n] 't (x: [n]t) = x[n-1] + +-- | Everything but the first element of the array. +let tail [n] 't (x: [n]t) = x[1:] + +-- | Everything but the last element of the array. +let init [n] 't (x: [n]t) = x[0:n-1] + +-- | Take some number of elements from the head of the array. +let take [n] 't (i: i32) (x: [n]t): [i]t = x[0:i] + +-- | Remove some number of elements from the head of the array. +let drop [n] 't (i: i32) (x: [n]t) = x[i:] + +-- | Split an array at a given position. +let split 't (n: i32) (xs: []t): ([n]t, []t) = + (xs[:n], xs[n:]) + +-- | Split an array at two given positions. +let split2 't (i: i32) (j: i32) (xs: []t): ([i]t, []t, []t) = + (xs[:i], xs[i:j], xs[j:]) + +-- | Return the elements of the array in reverse order. +let reverse [n] 't (x: [n]t): [n]t = x[::-1] + +-- | Concatenate two arrays. Warning: never try to perform a reduction +-- with this operator; it will not work. +let (++) 't (xs: []t) (ys: []t): *[]t = intrinsics.concat (xs, ys) + +-- | An old-fashioned way of saying `++`. +let concat 't (xs: []t) (ys: []t): *[]t = xs ++ ys + +-- | Rotate an array some number of elements to the left. A negative +-- rotation amount is also supported. +-- +-- For example, if `b==rotate 1 i a`, then `b[x,y+1] = a[x,y]`. +let rotate 't (r: i32) (xs: []t) = intrinsics.rotate (r, xs) + +-- | Replace an element of the array with a new value. +let update [n] 't (xs: *[n]t) (i: i32) (x: t): *[n]t = xs with [i] = x + +-- | Construct an array of consecutive integers of the given length, +-- starting at 0. +let iota (n: i32): *[n]i32 = + i32.iota n + +-- | Construct an array of the given length containing the given +-- value. +let replicate 't (n: i32) (x: t): *[n]t = + i32.replicate n x + +-- | Copy a value. The result will not alias anything. +let copy 't (a: t): *t = + ([a])[0] + +-- | Combines the outer two dimensions of an array. +let flatten [n][m] 't (xs: [n][m]t): []t = + intrinsics.flatten xs + +-- | Combines the outer three dimensions of an array. +let flatten_3d [n][m][l] 't (xs: [n][m][l]t): []t = + flatten (flatten xs) + +-- | Combines the outer four dimensions of an array. +let flatten_4d [n][m][l][k] 't (xs: [n][m][l][k]t): []t = + flatten (flatten_3d xs) + +-- | Splits the outer dimension of an array in two. +let unflatten 't (n: i32) (m: i32) (xs: []t): [n][m]t = + intrinsics.unflatten (n, m, xs) + +-- | Splits the outer dimension of an array in three. +let unflatten_3d 't (n: i32) (m: i32) (l: i32) (xs: []t): [n][m][l]t = + unflatten n m (unflatten (n*m) l xs) + +-- | Splits the outer dimension of an array in four. +let unflatten_4d 't (n: i32) (m: i32) (l: i32) (k: i32) (xs: []t): [n][m][l][k]t = + unflatten n m (unflatten_3d (n*m) l k xs) + +let intersperse [n] 't (x: t) (xs: [n]t): *[]t = + map (\i -> if i % 2 == 1 && i != 2*n then x + else unsafe xs[i/2]) + (iota (i32.max (2*n-1) 0)) + +let intercalate [n] [m] 't (x: [m]t) (xs: [n][m]t): []t = + unsafe flatten (intersperse x xs) + +let transpose [n] [m] 't (a: [n][m]t): [m][n]t = + intrinsics.transpose a + +let steps (start: i32) (num_steps: i32) (step: i32): [num_steps]i32 = + map (start+) (map (step*) (iota num_steps)) + +let range (start: i32) (end: i32) (step: i32): []i32 = + let w = (end-start)/step + in steps start w step + +-- | True if all of the input elements are true. Produces true on an +-- empty array. +let and: []bool -> bool = all id + +-- | True if any of the input elements are true. Produces false on an +-- empty array. +let or: []bool -> bool = any id + +let pick [n] 't (flags: [n]bool) (xs: [n]t) (ys: [n]t): *[n]t = + map3 (\flag x y -> if flag then x else y) flags xs ys + +-- | Perform a *sequential* left-fold of an array. +let foldl 'a 'b (f: a -> b -> a) (acc: a) (bs: []b): a = + loop acc for b in bs do f acc b + +-- | Perform a *sequential* right-fold of an array. +let foldr 'a 'b (f: b -> a -> a) (acc: a) (bs: []b): a = + foldl (flip f) acc (reverse bs) + +-- | Create a value for each point in a one-dimensional index space. +let tabulate 'a (n: i32) (f: i32 -> a): *[n]a = + map1 f (iota n) + +-- | Create a value for each point in a two-dimensional index space. +let tabulate_2d 'a (n: i32) (m: i32) (f: i32 -> i32 -> a): *[n][m]a = + map1 (f >-> tabulate m) (iota n) + +-- | Create a value for each point in a three-dimensional index space. +let tabulate_3d 'a (n: i32) (m: i32) (o: i32) (f: i32 -> i32 -> i32 -> a): *[n][m][o]a = + map1 (f >-> tabulate_2d m o) (iota n) diff --git a/futlib/functional.fut b/futlib/functional.fut new file mode 100644 --- /dev/null +++ b/futlib/functional.fut @@ -0,0 +1,53 @@ +-- | Simple functional combinators. + +-- | Left-to-right application. Particularly useful for describing +-- computation pipelines: +-- +-- x |> f |> g |> h +let (|>) '^a '^b (x: a) (f: a -> b): b = f x + +-- | Right to left application. +let (<|) '^a '^b (f: a -> b) (x: a) = f x + +-- | Function composition, with values flowing from left to right. +let (>->) '^a '^b '^c (f: a -> b) (g: b -> c) (x: a): c = g (f x) + +-- | Function composition, with values flowing from right to left. +-- This is the same as the `∘` operator known from mathematics. +let (<-<) '^a '^b '^c (g: b -> c) (f: a -> b) (x: a): c = g (f x) + +-- | Flip the arguments passed to a function. +-- +-- f x y == flip f y x +let flip '^a '^b '^c (f: a -> b -> c) (b: b) (a: a): c = + f a b + +-- | Transform a function taking a pair into a function taking two +-- arguments. +let curry '^a '^b '^c (f: (a, b) -> c) (a: a) (b: b): c = + f (a, b) + +-- | Transform a function taking two arguments in a function taking a +-- pair. +let uncurry '^a '^b '^c (f: a -> b -> c) (a: a, b: b): c = + f a b + +-- | The constant function. +let const '^a '^b (x: a) (_: b): a = x + +-- | The identity function. +let id '^a (x: a) = x + +-- | Apply a function some number of times. +let iterate 'a (n: i32) (f: a -> a) (x: a) = + loop x for _i < n do f x + +-- | Keep applying `f` until `p` returns true for the input value. +-- May apply zero times. *Note*: may not terminate. +let iterate_until 'a (p: a -> bool) (f: a -> a) (x: a) = + loop x while ! (p x) do f x + +-- | Keep applying `f` while `p` returns true for the input value. +-- May apply zero times. *Note*: may not terminate. +let iterate_while 'a (p: a -> bool) (f: a -> a) (x: a) = + loop x while p x do f x diff --git a/futlib/math.fut b/futlib/math.fut new file mode 100644 --- /dev/null +++ b/futlib/math.fut @@ -0,0 +1,984 @@ +-- | Basic mathematical modules and functions. + +import "soacs" + +local let const 'a 'b (x: a) (_: b): a = x + +-- | Describes types of values that can be created from the primitive +-- numeric types (and bool). +module type from_prim = { + type t + + val i8: i8 -> t + val i16: i16 -> t + val i32: i32 -> t + val i64: i64 -> t + + val u8: u8 -> t + val u16: u16 -> t + val u32: u32 -> t + val u64: u64 -> t + + val f32: f32 -> t + val f64: f64 -> t + + val bool: bool -> t +} + +-- | A basic numeric module type that can be implemented for both +-- integers and rational numbers. +module type numeric = { + include from_prim + + val +: t -> t -> t + val -: t -> t -> t + val *: t -> t -> t + val /: t -> t -> t + val **: t -> t -> t + + val to_i64: t -> i64 + + val ==: t -> t -> bool + val <: t -> t -> bool + val >: t -> t -> bool + val <=: t -> t -> bool + val >=: t -> t -> bool + val !=: t -> t -> bool + + val negate: t-> t + val max: t -> t -> t + val min: t -> t -> t + + val abs: t -> t + + val sgn: t -> t + + -- | The highest representable number. + val highest: t + + -- | The lowest representable number. + val lowest: t + + -- | Returns zero on empty input. + val sum: []t -> t + + -- | Returns one on empty input. + val product: []t -> t + + -- | Returns `lowest` on empty input. + val maximum: []t -> t + -- | Returns `highest` on empty input. + val minimum: []t -> t +} + +-- | An extension of `numeric`@mtype that provides facilities that are +-- only meaningful for integral types. +module type integral = { + include numeric + + val %: t -> t -> t + val //: t -> t -> t + val %%: t -> t -> t + + val &: t -> t -> t + val |: t -> t -> t + val ^: t -> t -> t + val ~: t -> t + + val <<: t -> t -> t + val >>: t -> t -> t + val >>>: t -> t -> t + + val num_bits: i32 + val get_bit: i32 -> t -> i32 + val set_bit: i32 -> t -> i32 -> t +} + +-- | An extension of `size`@mtype that further includes facilities for +-- constructing arrays where the size is provided as a value of the +-- given integral type. +module type size = { + include integral + + val iota: t -> *[]t + val replicate 'v: t -> v -> *[]v +} + +-- | Numbers that model real numbers to some degree. +module type real = { + include numeric + + val from_fraction: i32 -> i32 -> t + val to_i32: t -> i32 + val to_i64: t -> i64 + val to_f64: t -> f64 + + val sqrt: t -> t + val exp: t -> t + val cos: t -> t + val sin: t -> t + val tan: t -> t + val asin: t -> t + val acos: t -> t + val atan: t -> t + val atan2: t -> t -> t + + -- | Natural logarithm. + val log: t -> t + -- | Base-2 logarithm. + val log2: t -> t + -- | Base-10 logarithm. + val log10: t -> t + + val ceil : t -> t + val floor : t -> t + val trunc : t -> t + + -- | Round to the nearest integer, with alfway cases rounded to the + -- nearest even integer. Note that this differs from `round()` in + -- C, but matches more modern languages. + val round : t -> t + + val isinf: t -> bool + val isnan: t -> bool + + val inf: t + val nan: t + + val pi: t + val e: t +} + +-- | An extension of `real`@mtype that further gives access to the +-- bitwise representation of the underlying number. It is presumed +-- that this will be some form of IEEE float. +module type float = { + include real + + -- | An unsigned integer type containing the same number of bits as + -- 't'. + type int_t + + val from_bits: int_t -> t + val to_bits: t -> int_t + + val num_bits: i32 + val get_bit: i32 -> t -> i32 + val set_bit: i32 -> t -> i32 -> t +} + +-- | Boolean numbers. When converting from a number to `bool`, 0 is +-- considered `false` and any other value is `true`. +module bool: from_prim with t = bool = { + type t = bool + + let i8 = intrinsics.itob_i8_bool + let i16 = intrinsics.itob_i16_bool + let i32 = intrinsics.itob_i32_bool + let i64 = intrinsics.itob_i64_bool + + let u8 (x: u8) = intrinsics.itob_i8_bool (intrinsics.sign_i8 x) + let u16 (x: u16) = intrinsics.itob_i16_bool (intrinsics.sign_i16 x) + let u32 (x: u32) = intrinsics.itob_i32_bool (intrinsics.sign_i32 x) + let u64 (x: u64) = intrinsics.itob_i64_bool (intrinsics.sign_i64 x) + + let f32 (x: f32) = x != 0f32 + let f64 (x: f64) = x != 0f64 + + let bool (x: bool) = x +} + +module i8: (size with t = i8) = { + type t = i8 + + let (x: i8) + (y: i8) = intrinsics.add8 x y + let (x: i8) - (y: i8) = intrinsics.sub8 x y + let (x: i8) * (y: i8) = intrinsics.mul8 x y + let (x: i8) / (y: i8) = intrinsics.sdiv8 x y + let (x: i8) ** (y: i8) = intrinsics.pow8 x y + let (x: i8) % (y: i8) = intrinsics.smod8 x y + let (x: i8) // (y: i8) = intrinsics.squot8 x y + let (x: i8) %% (y: i8) = intrinsics.srem8 x y + + let (x: i8) & (y: i8) = intrinsics.and8 x y + let (x: i8) | (y: i8) = intrinsics.or8 x y + let (x: i8) ^ (y: i8) = intrinsics.xor8 x y + let ~ (x: i8) = intrinsics.complement8 x + + let (x: i8) << (y: i8) = intrinsics.shl8 x y + let (x: i8) >> (y: i8) = intrinsics.ashr8 x y + let (x: i8) >>> (y: i8) = intrinsics.lshr8 x y + + let i8 (x: i8) = intrinsics.sext_i8_i8 x + let i16 (x: i16) = intrinsics.sext_i16_i8 x + let i32 (x: i32) = intrinsics.sext_i32_i8 x + let i64 (x: i64) = intrinsics.sext_i64_i8 x + + let u8 (x: u8) = intrinsics.zext_i8_i8 (intrinsics.sign_i8 x) + let u16 (x: u16) = intrinsics.zext_i16_i8 (intrinsics.sign_i16 x) + let u32 (x: u32) = intrinsics.zext_i32_i8 (intrinsics.sign_i32 x) + let u64 (x: u64) = intrinsics.zext_i64_i8 (intrinsics.sign_i64 x) + + let f32 (x: f32) = intrinsics.fptosi_f32_i8 x + let f64 (x: f64) = intrinsics.fptosi_f64_i8 x + + let bool = intrinsics.btoi_bool_i8 + + let to_i32(x: i8) = intrinsics.sext_i8_i32 x + let to_i64(x: i8) = intrinsics.sext_i8_i64 x + + let (x: i8) == (y: i8) = intrinsics.eq_i8 x y + let (x: i8) < (y: i8) = intrinsics.slt8 x y + let (x: i8) > (y: i8) = intrinsics.slt8 y x + let (x: i8) <= (y: i8) = intrinsics.sle8 x y + let (x: i8) >= (y: i8) = intrinsics.sle8 y x + let (x: i8) != (y: i8) = ! (x == y) + + let sgn (x: i8) = intrinsics.ssignum8 x + let abs (x: i8) = intrinsics.abs8 x + + let negate (x: t) = -x + let max (x: t) (y: t) = intrinsics.smax8 x y + let min (x: t) (y: t) = intrinsics.smin8 x y + + let highest = 127i8 + let lowest = highest + 1i8 + + let num_bits = 8 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + + let iota (n: i8) = 0i8..1i8..> (y: i16) = intrinsics.ashr16 x y + let (x: i16) >>> (y: i16) = intrinsics.lshr16 x y + + let i8 (x: i8) = intrinsics.sext_i8_i16 x + let i16 (x: i16) = intrinsics.sext_i16_i16 x + let i32 (x: i32) = intrinsics.sext_i32_i16 x + let i64 (x: i64) = intrinsics.sext_i64_i16 x + + let u8 (x: u8) = intrinsics.zext_i8_i16 (intrinsics.sign_i8 x) + let u16 (x: u16) = intrinsics.zext_i16_i16 (intrinsics.sign_i16 x) + let u32 (x: u32) = intrinsics.zext_i32_i16 (intrinsics.sign_i32 x) + let u64 (x: u64) = intrinsics.zext_i64_i16 (intrinsics.sign_i64 x) + + let f32 (x: f32) = intrinsics.fptosi_f32_i16 x + let f64 (x: f64) = intrinsics.fptosi_f64_i16 x + + let bool = intrinsics.btoi_bool_i16 + + let to_i32(x: i16) = intrinsics.sext_i16_i32 x + let to_i64(x: i16) = intrinsics.sext_i16_i64 x + + let (x: i16) == (y: i16) = intrinsics.eq_i16 x y + let (x: i16) < (y: i16) = intrinsics.slt16 x y + let (x: i16) > (y: i16) = intrinsics.slt16 y x + let (x: i16) <= (y: i16) = intrinsics.sle16 x y + let (x: i16) >= (y: i16) = intrinsics.sle16 y x + let (x: i16) != (y: i16) = ! (x == y) + + let sgn (x: i16) = intrinsics.ssignum16 x + let abs (x: i16) = intrinsics.abs16 x + + let negate (x: t) = -x + let max (x: t) (y: t) = intrinsics.smax16 x y + let min (x: t) (y: t) = intrinsics.smin16 x y + + let highest = 32767i16 + let lowest = highest + 1i16 + + let num_bits = 16 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + + let iota (n: i16) = 0i16..1i16..> (y: i32) = intrinsics.ashr32 x y + let (x: i32) >>> (y: i32) = intrinsics.lshr32 x y + + let i8 (x: i8) = intrinsics.sext_i8_i32 x + let i16 (x: i16) = intrinsics.sext_i16_i32 x + let i32 (x: i32) = intrinsics.sext_i32_i32 x + let i64 (x: i64) = intrinsics.sext_i64_i32 x + + let u8 (x: u8) = intrinsics.zext_i8_i32 (intrinsics.sign_i8 x) + let u16 (x: u16) = intrinsics.zext_i16_i32 (intrinsics.sign_i16 x) + let u32 (x: u32) = intrinsics.zext_i32_i32 (intrinsics.sign_i32 x) + let u64 (x: u64) = intrinsics.zext_i64_i32 (intrinsics.sign_i64 x) + + let f32 (x: f32) = intrinsics.fptosi_f32_i32 x + let f64 (x: f64) = intrinsics.fptosi_f64_i32 x + + let bool = intrinsics.btoi_bool_i32 + + let to_i32(x: i32) = intrinsics.sext_i32_i32 x + let to_i64(x: i32) = intrinsics.sext_i32_i64 x + + let (x: i32) == (y: i32) = intrinsics.eq_i32 x y + let (x: i32) < (y: i32) = intrinsics.slt32 x y + let (x: i32) > (y: i32) = intrinsics.slt32 y x + let (x: i32) <= (y: i32) = intrinsics.sle32 x y + let (x: i32) >= (y: i32) = intrinsics.sle32 y x + let (x: i32) != (y: i32) = ! (x == y) + + let sgn (x: i32) = intrinsics.ssignum32 x + let abs (x: i32) = intrinsics.abs32 x + + let negate (x: t) = -x + let max (x: t) (y: t) = intrinsics.smax32 x y + let min (x: t) (y: t) = intrinsics.smin32 x y + + let highest = 2147483647 + let lowest = highest + 1 + + let num_bits = 32 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + + let iota (n: i32) = 0..1..> (y: i64) = intrinsics.ashr64 x y + let (x: i64) >>> (y: i64) = intrinsics.lshr64 x y + + let i8 (x: i8) = intrinsics.sext_i8_i64 x + let i16 (x: i16) = intrinsics.sext_i16_i64 x + let i32 (x: i32) = intrinsics.sext_i32_i64 x + let i64 (x: i64) = intrinsics.sext_i64_i64 x + + let u8 (x: u8) = intrinsics.zext_i8_i64 (intrinsics.sign_i8 x) + let u16 (x: u16) = intrinsics.zext_i16_i64 (intrinsics.sign_i16 x) + let u32 (x: u32) = intrinsics.zext_i32_i64 (intrinsics.sign_i32 x) + let u64 (x: u64) = intrinsics.zext_i64_i64 (intrinsics.sign_i64 x) + + let f32 (x: f32) = intrinsics.fptosi_f32_i64 x + let f64 (x: f64) = intrinsics.fptosi_f64_i64 x + + let bool = intrinsics.btoi_bool_i64 + + let to_i32(x: i64) = intrinsics.sext_i64_i32 x + let to_i64(x: i64) = intrinsics.sext_i64_i64 x + + let (x: i64) == (y: i64) = intrinsics.eq_i64 x y + let (x: i64) < (y: i64) = intrinsics.slt64 x y + let (x: i64) > (y: i64) = intrinsics.slt64 y x + let (x: i64) <= (y: i64) = intrinsics.sle64 x y + let (x: i64) >= (y: i64) = intrinsics.sle64 y x + let (x: i64) != (y: i64) = ! (x == y) + + let sgn (x: i64) = intrinsics.ssignum64 x + let abs (x: i64) = intrinsics.abs64 x + + let negate (x: t) = -x + let max (x: t) (y: t) = intrinsics.smax64 x y + let min (x: t) (y: t) = intrinsics.smin64 x y + + let highest = 9223372036854775807i64 + let lowest = highest + 1i64 + + let num_bits = 64 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | intrinsics.zext_i32_i64 (b intrinsics.<< bit)) + + let iota (n: i64) = 0i64..1i64..> (y: u8) = unsign (intrinsics.ashr8 (sign x) (sign y)) + let (x: u8) >>> (y: u8) = unsign (intrinsics.lshr8 (sign x) (sign y)) + + let u8 (x: u8) = unsign (i8.u8 x) + let u16 (x: u16) = unsign (i8.u16 x) + let u32 (x: u32) = unsign (i8.u32 x) + let u64 (x: u64) = unsign (i8.u64 x) + + let i8 (x: i8) = unsign (intrinsics.zext_i8_i8 x) + let i16 (x: i16) = unsign (intrinsics.zext_i16_i8 x) + let i32 (x: i32) = unsign (intrinsics.zext_i32_i8 x) + let i64 (x: i64) = unsign (intrinsics.zext_i64_i8 x) + + let f32 (x: f32) = unsign (intrinsics.fptoui_f32_i8 x) + let f64 (x: f64) = unsign (intrinsics.fptoui_f64_i8 x) + + let bool x = unsign (intrinsics.btoi_bool_i8 x) + + let to_i32(x: u8) = intrinsics.zext_i8_i32 (sign x) + let to_i64(x: u8) = intrinsics.zext_i8_i64 (sign x) + + let (x: u8) == (y: u8) = intrinsics.eq_i8 (sign x) (sign y) + let (x: u8) < (y: u8) = intrinsics.ult8 (sign x) (sign y) + let (x: u8) > (y: u8) = intrinsics.ult8 (sign y) (sign x) + let (x: u8) <= (y: u8) = intrinsics.ule8 (sign x) (sign y) + let (x: u8) >= (y: u8) = intrinsics.ule8 (sign y) (sign x) + let (x: u8) != (y: u8) = ! (x == y) + + let sgn (x: u8) = unsign (intrinsics.usignum8 (sign x)) + let abs (x: u8) = x + + let negate (x: t) = -x + let max (x: t) (y: t) = unsign (intrinsics.umax8 (sign x) (sign y)) + let min (x: t) (y: t) = unsign (intrinsics.umin8 (sign x) (sign y)) + + let highest = 255u8 + let lowest = 0u8 + + let num_bits = 8 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + + let iota (n: u8) = 0u8..1u8..> (y: u16) = unsign (intrinsics.ashr16 (sign x) (sign y)) + let (x: u16) >>> (y: u16) = unsign (intrinsics.lshr16 (sign x) (sign y)) + + let u8 (x: u8) = unsign (i16.u8 x) + let u16 (x: u16) = unsign (i16.u16 x) + let u32 (x: u32) = unsign (i16.u32 x) + let u64 (x: u64) = unsign (i16.u64 x) + + let i8 (x: i8) = unsign (intrinsics.zext_i8_i16 x) + let i16 (x: i16) = unsign (intrinsics.zext_i16_i16 x) + let i32 (x: i32) = unsign (intrinsics.zext_i32_i16 x) + let i64 (x: i64) = unsign (intrinsics.zext_i64_i16 x) + + let f32 (x: f32) = unsign (intrinsics.fptoui_f32_i16 x) + let f64 (x: f64) = unsign (intrinsics.fptoui_f64_i16 x) + + let bool x = unsign (intrinsics.btoi_bool_i16 x) + + let to_i32(x: u16) = intrinsics.zext_i16_i32 (sign x) + let to_i64(x: u16) = intrinsics.zext_i16_i64 (sign x) + + let (x: u16) == (y: u16) = intrinsics.eq_i16 (sign x) (sign y) + let (x: u16) < (y: u16) = intrinsics.ult16 (sign x) (sign y) + let (x: u16) > (y: u16) = intrinsics.ult16 (sign y) (sign x) + let (x: u16) <= (y: u16) = intrinsics.ule16 (sign x) (sign y) + let (x: u16) >= (y: u16) = intrinsics.ule16 (sign y) (sign x) + let (x: u16) != (y: u16) = ! (x == y) + + let sgn (x: u16) = unsign (intrinsics.usignum16 (sign x)) + let abs (x: u16) = x + + let negate (x: t) = -x + let max (x: t) (y: t) = unsign (intrinsics.umax16 (sign x) (sign y)) + let min (x: t) (y: t) = unsign (intrinsics.umin16 (sign x) (sign y)) + + let highest = 65535u16 + let lowest = 0u16 + + let num_bits = 16 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + + let iota (n: u16) = 0u16..1u16..> (y: u32) = unsign (intrinsics.ashr32 (sign x) (sign y)) + let (x: u32) >>> (y: u32) = unsign (intrinsics.lshr32 (sign x) (sign y)) + + let u8 (x: u8) = unsign (i32.u8 x) + let u16 (x: u16) = unsign (i32.u16 x) + let u32 (x: u32) = unsign (i32.u32 x) + let u64 (x: u64) = unsign (i32.u64 x) + + let i8 (x: i8) = unsign (intrinsics.zext_i8_i32 x) + let i16 (x: i16) = unsign (intrinsics.zext_i16_i32 x) + let i32 (x: i32) = unsign (intrinsics.zext_i32_i32 x) + let i64 (x: i64) = unsign (intrinsics.zext_i64_i32 x) + + let f32 (x: f32) = unsign (intrinsics.fptoui_f32_i32 x) + let f64 (x: f64) = unsign (intrinsics.fptoui_f64_i32 x) + + let bool x = unsign (intrinsics.btoi_bool_i32 x) + + let to_i32(x: u32) = intrinsics.zext_i32_i32 (sign x) + let to_i64(x: u32) = intrinsics.zext_i32_i64 (sign x) + + let (x: u32) == (y: u32) = intrinsics.eq_i32 (sign x) (sign y) + let (x: u32) < (y: u32) = intrinsics.ult32 (sign x) (sign y) + let (x: u32) > (y: u32) = intrinsics.ult32 (sign y) (sign x) + let (x: u32) <= (y: u32) = intrinsics.ule32 (sign x) (sign y) + let (x: u32) >= (y: u32) = intrinsics.ule32 (sign y) (sign x) + let (x: u32) != (y: u32) = ! (x == y) + + let sgn (x: u32) = unsign (intrinsics.usignum32 (sign x)) + let abs (x: u32) = x + + let highest = 4294967295u32 + let lowest = highest + 1u32 + + let negate (x: t) = -x + let max (x: t) (y: t) = unsign (intrinsics.umax32 (sign x) (sign y)) + let min (x: t) (y: t) = unsign (intrinsics.umin32 (sign x) (sign y)) + + let num_bits = 32 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + + let iota (n: u32) = 0u32..1u32..> (y: u64) = unsign (intrinsics.ashr64 (sign x) (sign y)) + let (x: u64) >>> (y: u64) = unsign (intrinsics.lshr64 (sign x) (sign y)) + + let u8 (x: u8) = unsign (i64.u8 x) + let u16 (x: u16) = unsign (i64.u16 x) + let u32 (x: u32) = unsign (i64.u32 x) + let u64 (x: u64) = unsign (i64.u64 x) + + let i8 (x: i8) = unsign (intrinsics.zext_i8_i64 x) + let i16 (x: i16) = unsign (intrinsics.zext_i16_i64 x) + let i32 (x: i32) = unsign (intrinsics.zext_i32_i64 x) + let i64 (x: i64) = unsign (intrinsics.zext_i64_i64 x) + + let f32 (x: f32) = unsign (intrinsics.fptoui_f32_i64 x) + let f64 (x: f64) = unsign (intrinsics.fptoui_f64_i64 x) + + let bool x = unsign (intrinsics.btoi_bool_i64 x) + + let to_i32(x: u64) = intrinsics.zext_i64_i32 (sign x) + let to_i64(x: u64) = intrinsics.zext_i64_i64 (sign x) + + let (x: u64) == (y: u64) = intrinsics.eq_i64 (sign x) (sign y) + let (x: u64) < (y: u64) = intrinsics.ult64 (sign x) (sign y) + let (x: u64) > (y: u64) = intrinsics.ult64 (sign y) (sign x) + let (x: u64) <= (y: u64) = intrinsics.ule64 (sign x) (sign y) + let (x: u64) >= (y: u64) = intrinsics.ule64 (sign y) (sign x) + let (x: u64) != (y: u64) = ! (x == y) + + let sgn (x: u64) = unsign (intrinsics.usignum64 (sign x)) + let abs (x: u64) = x + + let negate (x: t) = -x + let max (x: t) (y: t) = unsign (intrinsics.umax64 (sign x) (sign y)) + let min (x: t) (y: t) = unsign (intrinsics.umin64 (sign x) (sign y)) + + let highest = 18446744073709551615u64 + let lowest = highest + 1u64 + + let num_bits = 64 + let get_bit (bit: i32) (x: t) = to_i32 ((x >> i32 bit) & i32 1) + let set_bit (bit: i32) (x: t) (b: i32) = + ((x & i32 (intrinsics.~(1 intrinsics.<< bit))) | i32 (b intrinsics.<< bit)) + + let iota (n: u64) = 0u64..1u64.. (y: f64) = intrinsics.lt64 y x + let (x: f64) <= (y: f64) = intrinsics.le64 x y + let (x: f64) >= (y: f64) = intrinsics.le64 y x + let (x: f64) != (y: f64) = ! (x == y) + + let negate (x: t) = -x + let max (x: t) (y: t) = intrinsics.fmax64 x y + let min (x: t) (y: t) = intrinsics.fmin64 x y + + let sgn (x: f64) = if x < 0f64 then -1f64 + else if x == 0f64 then 0f64 + else 1f64 + let abs (x: f64) = intrinsics.fabs64 x + + let sqrt (x: f64) = intrinsics.sqrt64 x + + let log (x: f64) = intrinsics.log64 x + let log2 (x: f64) = intrinsics.log2_64 x + let log10 (x: f64) = intrinsics.log10_64 x + let exp (x: f64) = intrinsics.exp64 x + let cos (x: f64) = intrinsics.cos64 x + let sin (x: f64) = intrinsics.sin64 x + let tan (x: f64) = intrinsics.tan64 x + let acos (x: f64) = intrinsics.acos64 x + let asin (x: f64) = intrinsics.asin64 x + let atan (x: f64) = intrinsics.atan64 x + let atan2 (x: f64) (y: f64) = intrinsics.atan2_64 x y + + let ceil (x: f64) : f64 = + let i = to_i64 x + let ix = i64 i + in if x >= 0.0 then + if ix < x then i64 (i i64m.+ 1i64) else x + else if ix > x then ix else x + + let floor (x: f64) : f64 = + let i = to_i64 x + let ix = i64 i + in if x >= 0.0 then + if ix < x then ix else x + else if ix > x then i64 (i i64m.- 1i64) else x + + let trunc (x: f64) : f64 = i64 (i64m.f64 x) + + let even (x: f64) = i64m.f64 x % 2i64 i64m.== 0i64 + + let round = intrinsics.round64 + + let to_bits (x: f64): u64 = u64m.i64 (intrinsics.to_bits64 x) + let from_bits (x: u64): f64 = intrinsics.from_bits64 (intrinsics.sign_i64 x) + + let num_bits = 64 + let get_bit (bit: i32) (x: t) = u64m.get_bit bit (to_bits x) + let set_bit (bit: i32) (x: t) (b: i32) = from_bits (u64m.set_bit bit (to_bits x) b) + + let isinf (x: f64) = intrinsics.isinf64 x + let isnan (x: f64) = intrinsics.isnan64 x + + let inf = 1f64 / 0f64 + let nan = 0f64 / 0f64 + + let highest = inf + let lowest = -inf + + let pi = 3.1415926535897932384626433832795028841971693993751058209749445923078164062f64 + let e = 2.718281828459045235360287471352662497757247093699959574966967627724076630353f64 + + let sum = reduce (+) (i32 0) + let product = reduce (*) (i32 1) + let maximum = reduce max lowest + let minimum = reduce min highest +} + +module f32: (float with t = f32 with int_t = u32) = { + type t = f32 + type int_t = u32 + + module i32m = i32 + module u32m = u32 + module f64m = f64 + + let (x: f32) + (y: f32) = intrinsics.fadd32 x y + let (x: f32) - (y: f32) = intrinsics.fsub32 x y + let (x: f32) * (y: f32) = intrinsics.fmul32 x y + let (x: f32) / (y: f32) = intrinsics.fdiv32 x y + let (x: f32) ** (y: f32) = intrinsics.fpow32 x y + + let u8 (x: u8) = intrinsics.uitofp_i8_f32 (i8.u8 x) + let u16 (x: u16) = intrinsics.uitofp_i16_f32 (i16.u16 x) + let u32 (x: u32) = intrinsics.uitofp_i32_f32 (i32.u32 x) + let u64 (x: u64) = intrinsics.uitofp_i64_f32 (i64.u64 x) + + let i8 (x: i8) = intrinsics.sitofp_i8_f32 x + let i16 (x: i16) = intrinsics.sitofp_i16_f32 x + let i32 (x: i32) = intrinsics.sitofp_i32_f32 x + let i64 (x: i64) = intrinsics.sitofp_i64_f32 x + + let f32 (x: f32) = intrinsics.fpconv_f32_f32 x + let f64 (x: f64) = intrinsics.fpconv_f64_f32 x + + let bool (x: bool) = if x then 1f32 else 0f32 + + let from_fraction (x: i32) (y: i32) = i32 x / i32 y + let to_i32 (x: f32) = intrinsics.fptosi_f32_i32 x + let to_i64 (x: f32) = intrinsics.fptosi_f32_i64 x + let to_f64 (x: f32) = intrinsics.fpconv_f32_f64 x + + let (x: f32) == (y: f32) = intrinsics.eq_f32 x y + let (x: f32) < (y: f32) = intrinsics.lt32 x y + let (x: f32) > (y: f32) = intrinsics.lt32 y x + let (x: f32) <= (y: f32) = intrinsics.le32 x y + let (x: f32) >= (y: f32) = intrinsics.le32 y x + let (x: f32) != (y: f32) = ! (x == y) + + let negate (x: t) = -x + let max (x: t) (y: t) = intrinsics.fmax32 x y + let min (x: t) (y: t) = intrinsics.fmin32 x y + + let sgn (x: f32) = if x < 0f32 then -1f32 + else if x == 0f32 then 0f32 + else 1f32 + let abs (x: f32) = intrinsics.fabs32 x + + let sqrt (x: f32) = intrinsics.sqrt32 x + + let log (x: f32) = intrinsics.log32 x + let log2 (x: f32) = intrinsics.log2_32 x + let log10 (x: f32) = intrinsics.log10_32 x + let exp (x: f32) = intrinsics.exp32 x + let cos (x: f32) = intrinsics.cos32 x + let sin (x: f32) = intrinsics.sin32 x + let tan (x: f32) = intrinsics.tan32 x + let acos (x: f32) = intrinsics.acos32 x + let asin (x: f32) = intrinsics.asin32 x + let atan (x: f32) = intrinsics.atan32 x + let atan2 (x: f32) (y: f32) = intrinsics.atan2_32 x y + + let ceil (x: f32) : f32 = + let i = to_i32 x + let ix = i32 i + in if x >= 0f32 then + if ix < x then i32 (i i32m.+ 1i32) else x + else if ix > x then ix else x + + let floor (x: f32) : f32 = + let i = to_i32 x + let ix = i32 i + in if x >= 0f32 then + if ix < x then ix else x + else if ix > x then i32 (i i32m.- 1i32) else x + + let trunc (x: f32) : f32 = i32 (i32m.f32 x) + + let even (x: f32) = i32m.f32 x % 2i32 i32m.== 0i32 + + let round = intrinsics.round32 + + let to_bits (x: f32): u32 = u32m.i32 (intrinsics.to_bits32 x) + let from_bits (x: u32): f32 = intrinsics.from_bits32 (intrinsics.sign_i32 x) + + let num_bits = 32 + let get_bit (bit: i32) (x: t) = u32m.get_bit bit (to_bits x) + let set_bit (bit: i32) (x: t) (b: i32) = from_bits (u32m.set_bit bit (to_bits x) b) + + let isinf (x: f32) = intrinsics.isinf32 x + let isnan (x: f32) = intrinsics.isnan32 x + + let inf = 1f32 / 0f32 + let nan = 0f32 / 0f32 + + let highest = inf + let lowest = -inf + + let pi = f64 f64m.pi + let e = f64 f64m.e + + let sum = reduce (+) (i32 0) + let product = reduce (*) (i32 1) + let maximum = reduce max lowest + let minimum = reduce min highest +} diff --git a/futlib/prelude.fut b/futlib/prelude.fut new file mode 100644 --- /dev/null +++ b/futlib/prelude.fut @@ -0,0 +1,27 @@ +-- | The default prelude that is implicitly available in all Futhark +-- files. + +open import "soacs" +open import "array" +open import "math" +open import "functional" + +-- | Create single-precision float from integer. +let r32 (x: i32): f32 = f32.i32 x +-- | Create integer from single-precision float. +let t32 (x: f32): i32 = i32.f32 x + +-- | Create double-precision float from integer. +let r64 (x: i32): f64 = f64.i32 x +-- | Create integer from double-precision float. +let t64 (x: f64): i32 = i32.f64 x + +-- | Semantically just identity, but in `futharki` the argument value +-- will be printed. +let trace 't (x: t): t = + intrinsics.trace x + +-- | Semantically just identity, but acts as a break point in +-- `futharki`. +let break 't (x: t): t = + intrinsics.break x diff --git a/futlib/soacs.fut b/futlib/soacs.fut new file mode 100644 --- /dev/null +++ b/futlib/soacs.fut @@ -0,0 +1,250 @@ +-- | Various Second-Order Array Combinators that are operationally +-- parallel in a way that can be exploited by the compiler. +-- +-- The functions here are all recognised specially by the compiler (or +-- built on those that are). The asymptotic [work and +-- span](https://en.wikipedia.org/wiki/Analysis_of_parallel_algorithms) +-- is provided for each function, but note that this easily hides very +-- substantial constant factors. For example, `scan`@term is *much* +-- slower than `reduce`@term, although they have the same asymptotic +-- complexity. +-- +-- *Reminder on terminology*: A function `op` is said to be +-- *associative* if +-- +-- (x `op` y) `op` z == x `op` (y `op` z) +-- +-- for all `x`, `y`, `z`. Similarly, it is *commutative* if +-- +-- x `op` y == y `op` x +-- +-- The value `o` is a *neutral element* if +-- +-- x `op` o == o `op` x == x +-- +-- for any `x`. + +import "zip" + +-- | Apply the given function to each element of an array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let map 'a [n] 'x (f: a -> x) (as: [n]a): *[n]x = + intrinsics.map (f, as) + +-- | Apply the given function to each element of a single array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let map1 'a [n] 'x (f: a -> x) (as: [n]a): *[n]x = + map f as + +-- | As `map1`@term, but with one more array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let map2 'a 'b [n] 'x (f: a -> b -> x) (as: [n]a) (bs: [n]b): *[n]x = + map (\(a, b) -> f a b) (zip2 as bs) + +-- | As `map2`@term, but with one more array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let map3 'a 'b 'c [n] 'x (f: a -> b -> c -> x) (as: [n]a) (bs: [n]b) (cs: [n]c): *[n]x = + map (\(a, b, c) -> f a b c) (zip3 as bs cs) + +-- | As `map3`@term, but with one more array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let map4 'a 'b 'c 'd [n] 'x (f: a -> b -> c -> d -> x) (as: [n]a) (bs: [n]b) (cs: [n]c) (ds: [n]d): *[n]x = + map (\(a, b, c, d) -> f a b c d) (zip4 as bs cs ds) + +-- | As `map4`@term, but with one more array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let map5 'a 'b 'c 'd 'e [n] 'x (f: a -> b -> c -> d -> e -> x) (as: [n]a) (bs: [n]b) (cs: [n]c) (ds: [n]d) (es: [n]e): *[n]x = + map (\(a, b, c, d, e) -> f a b c d e) (zip5 as bs cs ds es) + +-- | Reduce the array `as` with `op`, with `ne` as the neutral +-- element for `op`. The function `op` must be associative. If +-- it is not, the return value is unspecified. If the value returned +-- by the operator is an array, it must have the exact same size as +-- the neutral element, and that must again have the same size as the +-- elements of the input array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let reduce 'a (op: a -> a -> a) (ne: a) (as: []a): a = + intrinsics.reduce (op, ne, as) + +-- | As `reduce`, but the operator must also be commutative. This +-- is potentially faster than `reduce`. For simple built-in +-- operators, like addition, the compiler already knows that the +-- operator is associative. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let reduce_comm 'a (op: a -> a -> a) (ne: a) (as: []a): a = + intrinsics.reduce_comm (op, ne, as) + +-- | `reduce_by_index dest f ne is as` returns `dest`, but with each +-- element given by the indices of `is` updated by applying `f` to the +-- current value in `dest` and the corresponding value in `as`. The +-- `ne` value must be a neutral element for `op`. If `is` has +-- duplicates, `f` may be applied multiple times, and hence must be +-- associative and commutative. Out-of-bounds indices in `is` are +-- ignored. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(n)* in the worst case (all updates to same position), +-- but *O(1)* in the best case. +-- +-- In practice, the *O(n)* behaviour only occurs if *m* is also very +-- large. +let reduce_by_index 'a [m] [n] (dest : *[m]a) (f : a -> a -> a) (ne : a) (is : [n]i32) (as : [n]a) : *[m]a = + intrinsics.gen_reduce (dest, f, ne, is, as) + +-- | Inclusive prefix scan. Has the same caveats with respect to +-- associativity as `reduce`. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let scan [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): *[n]a = + intrinsics.scan (op, ne, as) + +-- | Remove all those elements of `as` that do not satisfy the +-- predicate `p`. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let filter 'a (p: a -> bool) (as: []a): *[]a = + let (as', is) = intrinsics.partition (1, \x -> if p x then 0 else 1, as) + in as'[:is[0]] + +-- | Split an array into those elements that satisfy the given +-- predicate, and those that do not. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let partition [n] 'a (p: a -> bool) (as: [n]a): ([]a, []a) = + let p' x = if p x then 0 else 1 + let (as', is) = intrinsics.partition (2, p', as) + in (as'[0:is[0]], as'[is[0]:n]) + +-- | Split an array by two predicates, producing three arrays. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let partition2 [n] 'a (p1: a -> bool) (p2: a -> bool) (as: [n]a): ([]a, []a, []a) = + let p' x = if p1 x then 0 else if p2 x then 1 else 2 + let (as', is) = intrinsics.partition (3, p', as) + in (as'[0:is[0]], as'[is[0]:is[0]+is[1]], as'[is[0]+is[1]:n]) + +-- | `stream_red op f as` splits `as` into chunks, applies `f` to each +-- of these in parallel, and uses `op` (which must be associative) to +-- combine the per-chunk results into a final result. This SOAC is +-- useful when `f` can be given a particularly work-efficient +-- sequential implementation. Operationally, we can imagine that `as` +-- is divided among as many threads as necessary to saturate the +-- machine, with each thread operating sequentially. +-- +-- A chunk may be empty, `f []` must produce the neutral element for +-- `op`. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let stream_red 'a 'b (op: b -> b -> b) (f: []a -> b) (as: []a): b = + intrinsics.stream_red (op, f, as) + +-- | As `stream_red`@term, but the chunks do not necessarily +-- correspond to subsequences of the original array (they may be +-- interleaved). +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let stream_red_per 'a 'b (op: b -> b -> b) (f: []a -> b) (as: []a): b = + intrinsics.stream_red_per (op, f, as) + +-- | Similar to `stream_red`@term, except that each chunk must produce +-- an array *of the same size*. The per-chunk results are +-- concatenated. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let stream_map 'a 'b (f: []a -> []b) (as: []a): *[]b = + intrinsics.stream_map (f, as) + +-- | Similar to `stream_map`@term, but the chunks do not necessarily +-- correspond to subsequences of the original array (they may be +-- interleaved). +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let stream_map_per 'a 'b (f: []a -> []b) (as: []a): *[]b = + intrinsics.stream_map_per (f, as) + +-- | Return `true` if the given function returns `true` for all +-- elements in the array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let all 'a (f: a -> bool) (as: []a): bool = + reduce (&&) true (map f as) + +-- | Return `true` if the given function returns `true` for any +-- elements in the array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(log(n))* +let any 'a (f: a -> bool) (as: []a): bool = + reduce (||) false (map f as) + +-- | The `scatter as is vs` expression calculates the equivalent of +-- this imperative code: +-- +-- ``` +-- for index in 0..length is-1: +-- i = is[index] +-- v = vs[index] +-- as[i] = v +-- ``` +-- +-- The `is` and `vs` arrays must have the same outer size. `scatter` +-- acts in-place and consumes the `as` array, returning a new array +-- that has the same type and elements as `as`, except for the indices +-- in `is`. If `is` contains duplicates (i.e. several writes are +-- performed to the same location), the result is unspecified. It is +-- not guaranteed that one of the duplicate writes will complete +-- atomically - they may be interleaved. See `reduce_by_index`@term +-- for a function that can handle this case deterministically. +-- +-- This is technically not a second-order operation, but it is defined +-- here because it is closely related to the SOACs. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let scatter 't [m] [n] (dest: *[m]t) (is: [n]i32) (vs: [n]t): *[m]t = + intrinsics.scatter (dest, is, vs) diff --git a/futlib/zip.fut b/futlib/zip.fut new file mode 100644 --- /dev/null +++ b/futlib/zip.fut @@ -0,0 +1,58 @@ +-- | Transforming arrays of tuples into tuples of arrays and back +-- again. These are generally very cheap operations, as the internal +-- compiler representation is always tuples of arrays. + +-- The main reason this module exists is that we need it to define +-- SOACs like `map2`@term@"/futlib/soacs". + +-- We need a map to define some of the zip variants, but this file is +-- depended upon by soacs.fut. So we just define a quick-and-dirty +-- internal one here that uses the intrinsic version. +local let internal_map 'a [n] 'x (f: a -> x) (as: [n]a): [n]x = + intrinsics.map (f, as) + +-- | Construct an array of pairs from two arrays. +let zip [n] 'a 'b (as: [n]a) (bs: [n]b): [n](a,b) = + intrinsics.zip (as, bs) + +-- | Construct an array of pairs from two arrays. +let zip2 [n] 'a 'b (as: [n]a) (bs: [n]b): [n](a,b) = + zip as bs + +-- | As `zip2`@term, but with one more array. +let zip3 [n] 'a 'b 'c (as: [n]a) (bs: [n]b) (cs: [n]c): [n](a,b,c) = + internal_map (\(a,(b,c)) -> (a,b,c)) (zip as (zip2 bs cs)) + +-- | As `zip3`@term, but with one more array. +let zip4 [n] 'a 'b 'c 'd (as: [n]a) (bs: [n]b) (cs: [n]c) (ds: [n]d): [n](a,b,c,d) = + internal_map (\(a,(b,c,d)) -> (a,b,c,d)) (zip as (zip3 bs cs ds)) + +-- | As `zip4`@term, but with one more array. +let zip5 [n] 'a 'b 'c 'd 'e (as: [n]a) (bs: [n]b) (cs: [n]c) (ds: [n]d) (es: [n]e): [n](a,b,c,d,e) = + internal_map (\(a,(b,c,d,e)) -> (a,b,c,d,e)) (zip as (zip4 bs cs ds es)) + +-- | Turn an array of pairs into two arrays. +let unzip [n] 'a 'b (xs: [n](a,b)): ([n]a, [n]b) = + intrinsics.unzip xs + +-- | Turn an array of pairs into two arrays. +let unzip2 [n] 'a 'b (xs: [n](a,b)): ([n]a, [n]b) = + unzip xs + +-- | As `unzip2`@term, but with one more array. +let unzip3 [n] 'a 'b 'c (xs: [n](a,b,c)): ([n]a, [n]b, [n]c) = + let (as, bcs) = unzip (internal_map (\(a,b,c) -> (a,(b,c))) xs) + let (bs, cs) = unzip bcs + in (as, bs, cs) + +-- | As `unzip3`@term, but with one more array. +let unzip4 [n] 'a 'b 'c 'd (xs: [n](a,b,c,d)): ([n]a, [n]b, [n]c, [n]d) = + let (as, bs, cds) = unzip3 (internal_map (\(a,b,c,d) -> (a,b,(c,d))) xs) + let (cs, ds) = unzip cds + in (as, bs, cs, ds) + +-- | As `unzip4`@term, but with one more array. +let unzip5 [n] 'a 'b 'c 'd 'e (xs: [n](a,b,c,d,e)): ([n]a, [n]b, [n]c, [n]d, [n]e) = + let (as, bs, cs, des) = unzip4 (internal_map (\(a,b,c,d,e) -> (a,b,c,(d,e))) xs) + let (ds, es) = unzip des + in (as, bs, cs, ds, es) diff --git a/rts/c/lock.h b/rts/c/lock.h new file mode 100644 --- /dev/null +++ b/rts/c/lock.h @@ -0,0 +1,57 @@ +/* A very simple cross-platform implementation of locks. Uses + pthreads on Unix and some Windows thing there. Futhark's + host-level code is not multithreaded, but user code may be, so we + need some mechanism for ensuring atomic access to API functions. + This is that mechanism. It is not exposed to user code at all, so + we do not have to worry about name collisions. */ + +#ifdef _WIN32 + +typedef HANDLE lock_t; + +static lock_t create_lock(lock_t *lock) { + *lock = CreateMutex(NULL, /* Default security attributes. */ + FALSE, /* Initially unlocked. */ + NULL); /* Unnamed. */ +} + +static void lock_lock(lock_t *lock) { + assert(WaitForSingleObject(*lock, INFINITE) == WAIT_OBJECT_0); +} + +static void lock_unlock(lock_t *lock) { + assert(ReleaseMutex(*lock)); +} + +static void free_lock(lock_t *lock) { + CloseHandle(*lock); +} + +#else +/* Assuming POSIX */ + +#include + +typedef pthread_mutex_t lock_t; + +static void create_lock(lock_t *lock) { + int r = pthread_mutex_init(lock, NULL); + assert(r == 0); +} + +static void lock_lock(lock_t *lock) { + int r = pthread_mutex_lock(lock); + assert(r == 0); +} + +static void lock_unlock(lock_t *lock) { + int r = pthread_mutex_unlock(lock); + assert(r == 0); +} + +static void free_lock(lock_t *lock) { + /* Nothing to do for pthreads. */ + lock = lock; +} + +#endif diff --git a/rts/c/opencl.h b/rts/c/opencl.h new file mode 100644 --- /dev/null +++ b/rts/c/opencl.h @@ -0,0 +1,854 @@ +/* The simple OpenCL runtime framework used by Futhark. */ + +#define CL_USE_DEPRECATED_OPENCL_1_2_APIS + +#ifdef __APPLE__ + #include +#else + #include +#endif + +#define OPENCL_SUCCEED_FATAL(e) opencl_succeed_fatal(e, #e, __FILE__, __LINE__) +#define OPENCL_SUCCEED_NONFATAL(e) opencl_succeed_nonfatal(e, #e, __FILE__, __LINE__) +// Take care not to override an existing error. +#define OPENCL_SUCCEED_OR_RETURN(e) { \ + char *error = OPENCL_SUCCEED_NONFATAL(e); \ + if (error) { \ + if (!ctx->error) { \ + ctx->error = error; \ + return bad; \ + } else { \ + free(error); \ + } \ + } \ + } + +// OPENCL_SUCCEED_OR_RETURN returns the value of the variable 'bad' in +// scope. By default, it will be this one. Create a local variable +// of some other type if needed. This is a bit of a hack, but it +// saves effort in the code generator. +static const int bad = 1; + +struct opencl_config { + int debugging; + int logging; + int preferred_device_num; + const char *preferred_platform; + const char *preferred_device; + + const char* dump_program_to; + const char* load_program_from; + + size_t default_group_size; + size_t default_num_groups; + size_t default_tile_size; + size_t default_threshold; + size_t transpose_block_dim; + + int default_group_size_changed; + int default_tile_size_changed; + + int num_sizes; + const char **size_names; + size_t *size_values; + const char **size_classes; + const char **size_entry_points; +}; + +void opencl_config_init(struct opencl_config *cfg, + int num_sizes, + const char *size_names[], + size_t *size_values, + const char *size_classes[], + const char *size_entry_points[]) { + cfg->debugging = 0; + cfg->logging = 0; + cfg->preferred_device_num = 0; + cfg->preferred_platform = ""; + cfg->preferred_device = ""; + cfg->dump_program_to = NULL; + cfg->load_program_from = NULL; + + cfg->default_group_size = 256; + cfg->default_num_groups = 128; + cfg->default_tile_size = 32; + cfg->default_threshold = 32*1024; + cfg->transpose_block_dim = 16; + + cfg->default_group_size_changed = 0; + cfg->default_tile_size_changed = 0; + + cfg->num_sizes = num_sizes; + cfg->size_names = size_names; + cfg->size_values = size_values; + cfg->size_classes = size_classes; + cfg->size_entry_points = size_entry_points; +} + +/* An entry in the free list. May be invalid, to avoid having to + deallocate entries as soon as they are removed. There is also a + tag, to help with memory reuse. */ +struct opencl_free_list_entry { + size_t size; + cl_mem mem; + const char *tag; + unsigned char valid; +}; + +struct opencl_free_list { + struct opencl_free_list_entry *entries; // Pointer to entries. + int capacity; // Number of entries. + int used; // Number of valid entries. +}; + +void free_list_init(struct opencl_free_list *l) { + l->capacity = 30; // Picked arbitrarily. + l->used = 0; + l->entries = malloc(sizeof(struct opencl_free_list_entry) * l->capacity); + for (int i = 0; i < l->capacity; i++) { + l->entries[i].valid = 0; + } +} + +/* Remove invalid entries from the free list. */ +void free_list_pack(struct opencl_free_list *l) { + int p = 0; + for (int i = 0; i < l->capacity; i++) { + if (l->entries[i].valid) { + l->entries[p] = l->entries[i]; + p++; + } + } + // Now p == l->used. + l->entries = realloc(l->entries, l->used * sizeof(struct opencl_free_list_entry)); + l->capacity = l->used; +} + +void free_list_destroy(struct opencl_free_list *l) { + assert(l->used == 0); + free(l->entries); +} + +int free_list_find_invalid(struct opencl_free_list *l) { + int i; + for (i = 0; i < l->capacity; i++) { + if (!l->entries[i].valid) { + break; + } + } + return i; +} + +void free_list_insert(struct opencl_free_list *l, size_t size, cl_mem mem, const char *tag) { + int i = free_list_find_invalid(l); + + if (i == l->capacity) { + // List is full; so we have to grow it. + int new_capacity = l->capacity * 2 * sizeof(struct opencl_free_list_entry); + l->entries = realloc(l->entries, new_capacity); + for (int j = 0; j < l->capacity; j++) { + l->entries[j+l->capacity].valid = 0; + } + l->capacity *= 2; + } + + // Now 'i' points to the first invalid entry. + l->entries[i].valid = 1; + l->entries[i].size = size; + l->entries[i].mem = mem; + l->entries[i].tag = tag; + + l->used++; +} + +/* Find and remove a memory block of at least the desired size and + tag. Returns 0 on success. */ +int free_list_find(struct opencl_free_list *l, const char *tag, size_t *size_out, cl_mem *mem_out) { + int i; + for (i = 0; i < l->capacity; i++) { + if (l->entries[i].valid && l->entries[i].tag == tag) { + l->entries[i].valid = 0; + *size_out = l->entries[i].size; + *mem_out = l->entries[i].mem; + l->used--; + return 0; + } + } + + return 1; +} + +/* Remove the first block in the free list. Returns 0 if a block was + removed, and nonzero if the free list was already empty. */ +int free_list_first(struct opencl_free_list *l, cl_mem *mem_out) { + for (int i = 0; i < l->capacity; i++) { + if (l->entries[i].valid) { + l->entries[i].valid = 0; + *mem_out = l->entries[i].mem; + l->used--; + return 0; + } + } + + return 1; +} + +struct opencl_context { + cl_device_id device; + cl_context ctx; + cl_command_queue queue; + + struct opencl_config cfg; + + struct opencl_free_list free_list; + + size_t max_group_size; + size_t max_num_groups; + size_t max_tile_size; + size_t max_threshold; + + size_t lockstep_width; +}; + +struct opencl_device_option { + cl_platform_id platform; + cl_device_id device; + cl_device_type device_type; + char *platform_name; + char *device_name; +}; + +/* This function must be defined by the user. It is invoked by + setup_opencl() after the platform and device has been found, but + before the program is loaded. Its intended use is to tune + constants based on the selected platform and device. */ +static void post_opencl_setup(struct opencl_context*, struct opencl_device_option*); + +static char *strclone(const char *str) { + size_t size = strlen(str) + 1; + char *copy = malloc(size); + if (copy == NULL) { + return NULL; + } + + memcpy(copy, str, size); + return copy; +} + +static const char* opencl_error_string(unsigned int err) +{ + switch (err) { + case CL_SUCCESS: return "Success!"; + case CL_DEVICE_NOT_FOUND: return "Device not found."; + case CL_DEVICE_NOT_AVAILABLE: return "Device not available"; + case CL_COMPILER_NOT_AVAILABLE: return "Compiler not available"; + case CL_MEM_OBJECT_ALLOCATION_FAILURE: return "Memory object allocation failure"; + case CL_OUT_OF_RESOURCES: return "Out of resources"; + case CL_OUT_OF_HOST_MEMORY: return "Out of host memory"; + case CL_PROFILING_INFO_NOT_AVAILABLE: return "Profiling information not available"; + case CL_MEM_COPY_OVERLAP: return "Memory copy overlap"; + case CL_IMAGE_FORMAT_MISMATCH: return "Image format mismatch"; + case CL_IMAGE_FORMAT_NOT_SUPPORTED: return "Image format not supported"; + case CL_BUILD_PROGRAM_FAILURE: return "Program build failure"; + case CL_MAP_FAILURE: return "Map failure"; + case CL_INVALID_VALUE: return "Invalid value"; + case CL_INVALID_DEVICE_TYPE: return "Invalid device type"; + case CL_INVALID_PLATFORM: return "Invalid platform"; + case CL_INVALID_DEVICE: return "Invalid device"; + case CL_INVALID_CONTEXT: return "Invalid context"; + case CL_INVALID_QUEUE_PROPERTIES: return "Invalid queue properties"; + case CL_INVALID_COMMAND_QUEUE: return "Invalid command queue"; + case CL_INVALID_HOST_PTR: return "Invalid host pointer"; + case CL_INVALID_MEM_OBJECT: return "Invalid memory object"; + case CL_INVALID_IMAGE_FORMAT_DESCRIPTOR: return "Invalid image format descriptor"; + case CL_INVALID_IMAGE_SIZE: return "Invalid image size"; + case CL_INVALID_SAMPLER: return "Invalid sampler"; + case CL_INVALID_BINARY: return "Invalid binary"; + case CL_INVALID_BUILD_OPTIONS: return "Invalid build options"; + case CL_INVALID_PROGRAM: return "Invalid program"; + case CL_INVALID_PROGRAM_EXECUTABLE: return "Invalid program executable"; + case CL_INVALID_KERNEL_NAME: return "Invalid kernel name"; + case CL_INVALID_KERNEL_DEFINITION: return "Invalid kernel definition"; + case CL_INVALID_KERNEL: return "Invalid kernel"; + case CL_INVALID_ARG_INDEX: return "Invalid argument index"; + case CL_INVALID_ARG_VALUE: return "Invalid argument value"; + case CL_INVALID_ARG_SIZE: return "Invalid argument size"; + case CL_INVALID_KERNEL_ARGS: return "Invalid kernel arguments"; + case CL_INVALID_WORK_DIMENSION: return "Invalid work dimension"; + case CL_INVALID_WORK_GROUP_SIZE: return "Invalid work group size"; + case CL_INVALID_WORK_ITEM_SIZE: return "Invalid work item size"; + case CL_INVALID_GLOBAL_OFFSET: return "Invalid global offset"; + case CL_INVALID_EVENT_WAIT_LIST: return "Invalid event wait list"; + case CL_INVALID_EVENT: return "Invalid event"; + case CL_INVALID_OPERATION: return "Invalid operation"; + case CL_INVALID_GL_OBJECT: return "Invalid OpenGL object"; + case CL_INVALID_BUFFER_SIZE: return "Invalid buffer size"; + case CL_INVALID_MIP_LEVEL: return "Invalid mip-map level"; + default: return "Unknown"; + } +} + +static void opencl_succeed_fatal(unsigned int ret, + const char *call, + const char *file, + int line) { + if (ret != CL_SUCCESS) { + panic(-1, "%s:%d: OpenCL call\n %s\nfailed with error code %d (%s)\n", + file, line, call, ret, opencl_error_string(ret)); + } +} + +static char* opencl_succeed_nonfatal(unsigned int ret, + const char *call, + const char *file, + int line) { + if (ret != CL_SUCCESS) { + return msgprintf("%s:%d: OpenCL call\n %s\nfailed with error code %d (%s)\n", + file, line, call, ret, opencl_error_string(ret)); + } else { + return NULL; + } +} + +void set_preferred_platform(struct opencl_config *cfg, const char *s) { + cfg->preferred_platform = s; +} + +void set_preferred_device(struct opencl_config *cfg, const char *s) { + int x = 0; + if (*s == '#') { + s++; + while (isdigit(*s)) { + x = x * 10 + (*s++)-'0'; + } + // Skip trailing spaces. + while (isspace(*s)) { + s++; + } + } + cfg->preferred_device = s; + cfg->preferred_device_num = x; +} + +static char* opencl_platform_info(cl_platform_id platform, + cl_platform_info param) { + size_t req_bytes; + char *info; + + OPENCL_SUCCEED_FATAL(clGetPlatformInfo(platform, param, 0, NULL, &req_bytes)); + + info = malloc(req_bytes); + + OPENCL_SUCCEED_FATAL(clGetPlatformInfo(platform, param, req_bytes, info, NULL)); + + return info; +} + +static char* opencl_device_info(cl_device_id device, + cl_device_info param) { + size_t req_bytes; + char *info; + + OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device, param, 0, NULL, &req_bytes)); + + info = malloc(req_bytes); + + OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device, param, req_bytes, info, NULL)); + + return info; +} + +static void opencl_all_device_options(struct opencl_device_option **devices_out, + size_t *num_devices_out) { + size_t num_devices = 0, num_devices_added = 0; + + cl_platform_id *all_platforms; + cl_uint *platform_num_devices; + + cl_uint num_platforms; + + // Find the number of platforms. + OPENCL_SUCCEED_FATAL(clGetPlatformIDs(0, NULL, &num_platforms)); + + // Make room for them. + all_platforms = calloc(num_platforms, sizeof(cl_platform_id)); + platform_num_devices = calloc(num_platforms, sizeof(cl_uint)); + + // Fetch all the platforms. + OPENCL_SUCCEED_FATAL(clGetPlatformIDs(num_platforms, all_platforms, NULL)); + + // Count the number of devices for each platform, as well as the + // total number of devices. + for (cl_uint i = 0; i < num_platforms; i++) { + if (clGetDeviceIDs(all_platforms[i], CL_DEVICE_TYPE_ALL, + 0, NULL, &platform_num_devices[i]) == CL_SUCCESS) { + num_devices += platform_num_devices[i]; + } else { + platform_num_devices[i] = 0; + } + } + + // Make room for all the device options. + struct opencl_device_option *devices = + calloc(num_devices, sizeof(struct opencl_device_option)); + + // Loop through the platforms, getting information about their devices. + for (cl_uint i = 0; i < num_platforms; i++) { + cl_platform_id platform = all_platforms[i]; + cl_uint num_platform_devices = platform_num_devices[i]; + + if (num_platform_devices == 0) { + continue; + } + + char *platform_name = opencl_platform_info(platform, CL_PLATFORM_NAME); + cl_device_id *platform_devices = + calloc(num_platform_devices, sizeof(cl_device_id)); + + // Fetch all the devices. + OPENCL_SUCCEED_FATAL(clGetDeviceIDs(platform, CL_DEVICE_TYPE_ALL, + num_platform_devices, platform_devices, NULL)); + + // Loop through the devices, adding them to the devices array. + for (cl_uint i = 0; i < num_platform_devices; i++) { + char *device_name = opencl_device_info(platform_devices[i], CL_DEVICE_NAME); + devices[num_devices_added].platform = platform; + devices[num_devices_added].device = platform_devices[i]; + OPENCL_SUCCEED_FATAL(clGetDeviceInfo(platform_devices[i], CL_DEVICE_TYPE, + sizeof(cl_device_type), + &devices[num_devices_added].device_type, + NULL)); + // We don't want the structs to share memory, so copy the platform name. + // Each device name is already unique. + devices[num_devices_added].platform_name = strclone(platform_name); + devices[num_devices_added].device_name = device_name; + num_devices_added++; + } + free(platform_devices); + free(platform_name); + } + free(all_platforms); + free(platform_num_devices); + + *devices_out = devices; + *num_devices_out = num_devices; +} + +static int is_blacklisted(const char *platform_name, const char *device_name, + const struct opencl_config *cfg) { + if (strcmp(cfg->preferred_platform, "") != 0 || + strcmp(cfg->preferred_device, "") != 0) { + return 0; + } else if (strstr(platform_name, "Apple") != NULL && + strstr(device_name, "Intel(R) Core(TM)") != NULL) { + return 1; + } else { + return 0; + } +} + +static struct opencl_device_option get_preferred_device(const struct opencl_config *cfg) { + struct opencl_device_option *devices; + size_t num_devices; + + opencl_all_device_options(&devices, &num_devices); + + int num_device_matches = 0; + + for (size_t i = 0; i < num_devices; i++) { + struct opencl_device_option device = devices[i]; + if (!is_blacklisted(device.platform_name, device.device_name, cfg) && + strstr(device.platform_name, cfg->preferred_platform) != NULL && + strstr(device.device_name, cfg->preferred_device) != NULL && + num_device_matches++ == cfg->preferred_device_num) { + // Free all the platform and device names, except the ones we have chosen. + for (size_t j = 0; j < num_devices; j++) { + if (j != i) { + free(devices[j].platform_name); + free(devices[j].device_name); + } + } + free(devices); + return device; + } + } + + panic(1, "Could not find acceptable OpenCL device.\n"); + exit(1); // Never reached +} + +static void describe_device_option(struct opencl_device_option device) { + fprintf(stderr, "Using platform: %s\n", device.platform_name); + fprintf(stderr, "Using device: %s\n", device.device_name); +} + +static cl_build_status build_opencl_program(cl_program program, cl_device_id device, const char* options) { + cl_int ret_val = clBuildProgram(program, 1, &device, options, NULL, NULL); + + // Avoid termination due to CL_BUILD_PROGRAM_FAILURE + if (ret_val != CL_SUCCESS && ret_val != CL_BUILD_PROGRAM_FAILURE) { + assert(ret_val == 0); + } + + cl_build_status build_status; + ret_val = clGetProgramBuildInfo(program, + device, + CL_PROGRAM_BUILD_STATUS, + sizeof(cl_build_status), + &build_status, + NULL); + assert(ret_val == 0); + + if (build_status != CL_SUCCESS) { + char *build_log; + size_t ret_val_size; + ret_val = clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, NULL, &ret_val_size); + assert(ret_val == 0); + + build_log = malloc(ret_val_size+1); + clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, ret_val_size, build_log, NULL); + assert(ret_val == 0); + + // The spec technically does not say whether the build log is zero-terminated, so let's be careful. + build_log[ret_val_size] = '\0'; + + fprintf(stderr, "Build log:\n%s\n", build_log); + + free(build_log); + } + + return build_status; +} + +/* Fields in a bitmask indicating which types we must be sure are + available. */ +enum opencl_required_type { OPENCL_F64 = 1 }; + +// We take as input several strings representing the program, because +// C does not guarantee that the compiler supports particularly large +// literals. Notably, Visual C has a limit of 2048 characters. The +// array must be NULL-terminated. +static cl_program setup_opencl_with_command_queue(struct opencl_context *ctx, + cl_command_queue queue, + const char *srcs[], + int required_types) { + int error; + + ctx->queue = queue; + + OPENCL_SUCCEED_FATAL(clGetCommandQueueInfo(ctx->queue, CL_QUEUE_CONTEXT, sizeof(cl_context), &ctx->ctx, NULL)); + + // Fill out the device info. This is redundant work if we are + // called from setup_opencl() (which is the common case), but I + // doubt it matters much. + struct opencl_device_option device_option; + OPENCL_SUCCEED_FATAL(clGetCommandQueueInfo(ctx->queue, CL_QUEUE_DEVICE, + sizeof(cl_device_id), + &device_option.device, + NULL)); + OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_PLATFORM, + sizeof(cl_platform_id), + &device_option.platform, + NULL)); + OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_TYPE, + sizeof(cl_device_type), + &device_option.device_type, + NULL)); + device_option.platform_name = opencl_platform_info(device_option.platform, CL_PLATFORM_NAME); + device_option.device_name = opencl_device_info(device_option.device, CL_DEVICE_NAME); + + ctx->device = device_option.device; + + if (required_types & OPENCL_F64) { + cl_uint supported; + OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_PREFERRED_VECTOR_WIDTH_DOUBLE, + sizeof(cl_uint), &supported, NULL)); + if (!supported) { + panic(1, "Program uses double-precision floats, but this is not supported on the chosen device: %s", + device_option.device_name); + } + } + + size_t max_group_size; + OPENCL_SUCCEED_FATAL(clGetDeviceInfo(device_option.device, CL_DEVICE_MAX_WORK_GROUP_SIZE, + sizeof(size_t), &max_group_size, NULL)); + + size_t max_tile_size = sqrt(max_group_size); + + if (max_group_size < ctx->cfg.default_group_size) { + if (ctx->cfg.default_group_size_changed) { + fprintf(stderr, "Note: Device limits default group size to %zu (down from %zu).\n", + max_group_size, ctx->cfg.default_group_size); + } + ctx->cfg.default_group_size = max_group_size; + } + + if (max_tile_size < ctx->cfg.default_tile_size) { + if (ctx->cfg.default_tile_size_changed) { + fprintf(stderr, "Note: Device limits default tile size to %zu (down from %zu).\n", + max_tile_size, ctx->cfg.default_tile_size); + } + ctx->cfg.default_tile_size = max_tile_size; + } + + ctx->max_group_size = max_group_size; + ctx->max_tile_size = max_tile_size; // No limit. + ctx->max_threshold = ctx->max_num_groups = 0; // No limit. + + // Now we go through all the sizes, clamp them to the valid range, + // or set them to the default. + for (int i = 0; i < ctx->cfg.num_sizes; i++) { + const char *size_class = ctx->cfg.size_classes[i]; + size_t *size_value = &ctx->cfg.size_values[i]; + const char* size_name = ctx->cfg.size_names[i]; + size_t max_value, default_value; + if (strstr(size_class, "group_size") == size_class) { + max_value = max_group_size; + default_value = ctx->cfg.default_group_size; + } else if (strstr(size_class, "num_groups") == size_class) { + max_value = max_group_size; // Futhark assumes this constraint. + default_value = ctx->cfg.default_num_groups; + } else if (strstr(size_class, "tile_size") == size_class) { + max_value = sqrt(max_group_size); + default_value = ctx->cfg.default_tile_size; + } else if (strstr(size_class, "threshold") == size_class) { + max_value = 0; // No limit. + default_value = ctx->cfg.default_threshold; + } else { + panic(1, "Unknown size class for size '%s': %s\n", size_name, size_class); + } + if (*size_value == 0) { + *size_value = default_value; + } else if (max_value > 0 && *size_value > max_value) { + fprintf(stderr, "Note: Device limits %s to %d (down from %d)\n", + size_name, (int)max_value, (int)*size_value); + *size_value = max_value; + } + } + + // Make sure this function is defined. + post_opencl_setup(ctx, &device_option); + + if (ctx->cfg.logging) { + fprintf(stderr, "Lockstep width: %d\n", (int)ctx->lockstep_width); + fprintf(stderr, "Default group size: %d\n", (int)ctx->cfg.default_group_size); + fprintf(stderr, "Default number of groups: %d\n", (int)ctx->cfg.default_num_groups); + } + + char *fut_opencl_src = NULL; + size_t src_size = 0; + + // Maybe we have to read OpenCL source from somewhere else (used for debugging). + if (ctx->cfg.load_program_from != NULL) { + FILE *f = fopen(ctx->cfg.load_program_from, "r"); + assert(f != NULL); + fseek(f, 0, SEEK_END); + src_size = ftell(f); + fseek(f, 0, SEEK_SET); + fut_opencl_src = malloc(src_size); + assert(fread(fut_opencl_src, 1, src_size, f) == src_size); + fclose(f); + } else { + // Build the OpenCL program. First we have to concatenate all the fragments. + for (const char **src = srcs; src && *src; src++) { + src_size += strlen(*src); + } + + fut_opencl_src = malloc(src_size + 1); + + size_t n, i; + for (i = 0, n = 0; srcs && srcs[i]; i++) { + strncpy(fut_opencl_src+n, srcs[i], src_size-n); + n += strlen(srcs[i]); + } + fut_opencl_src[src_size] = 0; + + } + + cl_program prog; + error = 0; + const char* src_ptr[] = {fut_opencl_src}; + + if (ctx->cfg.dump_program_to != NULL) { + FILE *f = fopen(ctx->cfg.dump_program_to, "w"); + assert(f != NULL); + fputs(fut_opencl_src, f); + fclose(f); + } + + prog = clCreateProgramWithSource(ctx->ctx, 1, src_ptr, &src_size, &error); + assert(error == 0); + + int compile_opts_size = 1024; + for (int i = 0; i < ctx->cfg.num_sizes; i++) { + compile_opts_size += strlen(ctx->cfg.size_names[i]) + 20; + } + char *compile_opts = malloc(compile_opts_size); + + int w = snprintf(compile_opts, compile_opts_size, + "-DFUT_BLOCK_DIM=%d -DLOCKSTEP_WIDTH=%d ", + (int)ctx->cfg.transpose_block_dim, + (int)ctx->lockstep_width); + + for (int i = 0; i < ctx->cfg.num_sizes; i++) { + w += snprintf(compile_opts+w, compile_opts_size-w, + "-D%s=%d ", ctx->cfg.size_names[i], + (int)ctx->cfg.size_values[i]); + } + + OPENCL_SUCCEED_FATAL(build_opencl_program(prog, device_option.device, compile_opts)); + free(compile_opts); + free(fut_opencl_src); + + return prog; +} + +static cl_program setup_opencl(struct opencl_context *ctx, + const char *srcs[], + int required_types) { + + ctx->lockstep_width = 1; + + free_list_init(&ctx->free_list); + + struct opencl_device_option device_option = get_preferred_device(&ctx->cfg); + + if (ctx->cfg.logging) { + describe_device_option(device_option); + } + + // Note that NVIDIA's OpenCL requires the platform property + cl_context_properties properties[] = { + CL_CONTEXT_PLATFORM, + (cl_context_properties)device_option.platform, + 0 + }; + + cl_int error; + + ctx->ctx = clCreateContext(properties, 1, &device_option.device, NULL, NULL, &error); + assert(error == 0); + + cl_command_queue queue = clCreateCommandQueue(ctx->ctx, device_option.device, 0, &error); + assert(error == 0); + + return setup_opencl_with_command_queue(ctx, queue, srcs, required_types); +} + +// Allocate memory from driver. The problem is that OpenCL may perform +// lazy allocation, so we cannot know whether an allocation succeeded +// until the first time we try to use it. Hence we immediately +// perform a write to see if the allocation succeeded. This is slow, +// but the assumption is that this operation will be rare (most things +// will go through the free list). +int opencl_alloc_actual(struct opencl_context *ctx, size_t size, cl_mem *mem_out) { + int error; + *mem_out = clCreateBuffer(ctx->ctx, CL_MEM_READ_WRITE, size, NULL, &error); + + if (error != CL_SUCCESS) { + return error; + } + + int x = 2; + error = clEnqueueWriteBuffer(ctx->queue, *mem_out, 1, 0, sizeof(x), &x, 0, NULL, NULL); + + // No need to wait for completion here. clWaitForEvents() cannot + // return mem object allocation failures. This implies that the + // buffer is faulted onto the device on enqueue. (Observation by + // Andreas Kloeckner.) + + return error; +} + +int opencl_alloc(struct opencl_context *ctx, size_t min_size, const char *tag, cl_mem *mem_out) { + assert(min_size >= 0); + if (min_size < sizeof(int)) { + min_size = sizeof(int); + } + + size_t size; + + if (free_list_find(&ctx->free_list, tag, &size, mem_out) == 0) { + // Successfully found a free block. Is it big enough? + // + // FIXME: we might also want to check whether the block is *too + // big*, to avoid internal fragmentation. However, this can + // sharply impact performance on programs where arrays change size + // frequently. Fortunately, such allocations are usually fairly + // short-lived, as they are necessarily within a loop, so the risk + // of internal fragmentation resulting in an OOM situation is + // limited. However, it would be preferable if we could go back + // and *shrink* oversize allocations when we encounter an OOM + // condition. That is technically feasible, since we do not + // expose OpenCL pointer values directly to the application, but + // instead rely on a level of indirection. + if (size >= min_size) { + return CL_SUCCESS; + } else { + // Not just right - free it. + int error = clReleaseMemObject(*mem_out); + if (error != CL_SUCCESS) { + return error; + } + } + } + + // We have to allocate a new block from the driver. If the + // allocation does not succeed, then we might be in an out-of-memory + // situation. We now start freeing things from the free list until + // we think we have freed enough that the allocation will succeed. + // Since we don't know how far the allocation is from fitting, we + // have to check after every deallocation. This might be pretty + // expensive. Let's hope that this case is hit rarely. + + int error = opencl_alloc_actual(ctx, min_size, mem_out); + + while (error == CL_MEM_OBJECT_ALLOCATION_FAILURE) { + cl_mem mem; + if (free_list_first(&ctx->free_list, &mem) == 0) { + error = clReleaseMemObject(mem); + if (error != CL_SUCCESS) { + return error; + } + } else { + break; + } + error = opencl_alloc_actual(ctx, min_size, mem_out); + } + + return error; +} + +int opencl_free(struct opencl_context *ctx, cl_mem mem, const char *tag) { + size_t size; + cl_mem existing_mem; + + // If there is already a block with this tag, then remove it. + if (free_list_find(&ctx->free_list, tag, &size, &existing_mem) == 0) { + int error = clReleaseMemObject(existing_mem); + if (error != CL_SUCCESS) { + return error; + } + } + + int error = clGetMemObjectInfo(mem, CL_MEM_SIZE, sizeof(size_t), &size, NULL); + + if (error == CL_SUCCESS) { + free_list_insert(&ctx->free_list, size, mem, tag); + } + + return error; +} + +int opencl_free_all(struct opencl_context *ctx) { + cl_mem mem; + free_list_pack(&ctx->free_list); + while (free_list_first(&ctx->free_list, &mem) == 0) { + int error = clReleaseMemObject(mem); + if (error != CL_SUCCESS) { + return error; + } + } + + return CL_SUCCESS; +} diff --git a/rts/c/panic.h b/rts/c/panic.h new file mode 100644 --- /dev/null +++ b/rts/c/panic.h @@ -0,0 +1,28 @@ +/* Crash and burn. */ + +#include + +static const char *fut_progname; + +static void panic(int eval, const char *fmt, ...) +{ + va_list ap; + + va_start(ap, fmt); + fprintf(stderr, "%s: ", fut_progname); + vfprintf(stderr, fmt, ap); + va_end(ap); + exit(eval); +} + +/* For generating arbitrary-sized error messages. It is the callers + responsibility to free the buffer at some point. */ +static char* msgprintf(const char *s, ...) { + va_list vl; + va_start(vl, s); + size_t needed = 1 + vsnprintf(NULL, 0, s, vl); + char *buffer = malloc(needed); + va_start(vl, s); /* Must re-init. */ + vsnprintf(buffer, needed, s, vl); + return buffer; +} diff --git a/rts/c/timing.h b/rts/c/timing.h new file mode 100644 --- /dev/null +++ b/rts/c/timing.h @@ -0,0 +1,30 @@ +/* Some simple utilities for wall-clock timing. + + The function get_wall_time() returns the wall time in microseconds + (with an unspecified offset). +*/ + +#ifdef _WIN32 + +#include + +static int64_t get_wall_time(void) { + LARGE_INTEGER time,freq; + assert(QueryPerformanceFrequency(&freq)); + assert(QueryPerformanceCounter(&time)); + return ((double)time.QuadPart / freq.QuadPart) * 1000000; +} + +#else +/* Assuming POSIX */ + +#include +#include + +static int64_t get_wall_time(void) { + struct timeval time; + assert(gettimeofday(&time,NULL) == 0); + return time.tv_sec * 1000000 + time.tv_usec; +} + +#endif diff --git a/rts/c/values.h b/rts/c/values.h new file mode 100644 --- /dev/null +++ b/rts/c/values.h @@ -0,0 +1,819 @@ +//// Text I/O + +typedef int (*writer)(FILE*, void*); +typedef int (*bin_reader)(void*); +typedef int (*str_reader)(const char *, void*); + +struct array_reader { + char* elems; + int64_t n_elems_space; + int64_t elem_size; + int64_t n_elems_used; + int64_t *shape; + str_reader elem_reader; +}; + +static void skipspaces() { + int c; + do { + c = getchar(); + } while (isspace(c)); + + if (c != EOF) { + ungetc(c, stdin); + } +} + +static int constituent(char c) { + return isalnum(c) || c == '.' || c == '-' || c == '+' || c == '_'; +} + +// Produces an empty token only on EOF. +static void next_token(char *buf, int bufsize) { + start: + skipspaces(); + + int i = 0; + while (i < bufsize) { + int c = getchar(); + buf[i] = c; + + if (c == EOF) { + buf[i] = 0; + return; + } else if (c == '-' && i == 1 && buf[0] == '-') { + // Line comment, so skip to end of line and start over. + for (; c != '\n' && c != EOF; c = getchar()); + goto start; + } else if (!constituent(c)) { + if (i == 0) { + // We permit single-character tokens that are not + // constituents; this lets things like ']' and ',' be + // tokens. + buf[i+1] = 0; + return; + } else { + ungetc(c, stdin); + buf[i] = 0; + return; + } + } + + i++; + } + + buf[bufsize-1] = 0; +} + +static int next_token_is(char *buf, int bufsize, const char* expected) { + next_token(buf, bufsize); + return strcmp(buf, expected) == 0; +} + +static void remove_underscores(char *buf) { + char *w = buf; + + for (char *r = buf; *r; r++) { + if (*r != '_') { + *w++ = *r; + } + } + + *w++ = 0; +} + +static int read_str_elem(char *buf, struct array_reader *reader) { + int ret; + if (reader->n_elems_used == reader->n_elems_space) { + reader->n_elems_space *= 2; + reader->elems = (char*) realloc(reader->elems, + reader->n_elems_space * reader->elem_size); + } + + ret = reader->elem_reader(buf, reader->elems + reader->n_elems_used * reader->elem_size); + + if (ret == 0) { + reader->n_elems_used++; + } + + return ret; +} + +static int read_str_array_elems(char *buf, int bufsize, + struct array_reader *reader, int dims) { + int ret; + int first = 1; + char *knows_dimsize = (char*) calloc(dims,sizeof(char)); + int cur_dim = dims-1; + int64_t *elems_read_in_dim = (int64_t*) calloc(dims,sizeof(int64_t)); + + while (1) { + next_token(buf, bufsize); + + if (strcmp(buf, "]") == 0) { + if (knows_dimsize[cur_dim]) { + if (reader->shape[cur_dim] != elems_read_in_dim[cur_dim]) { + ret = 1; + break; + } + } else { + knows_dimsize[cur_dim] = 1; + reader->shape[cur_dim] = elems_read_in_dim[cur_dim]; + } + if (cur_dim == 0) { + ret = 0; + break; + } else { + cur_dim--; + elems_read_in_dim[cur_dim]++; + } + } else if (strcmp(buf, ",") == 0) { + next_token(buf, bufsize); + if (strcmp(buf, "[") == 0) { + if (cur_dim == dims - 1) { + ret = 1; + break; + } + first = 1; + cur_dim++; + elems_read_in_dim[cur_dim] = 0; + } else if (cur_dim == dims - 1) { + ret = read_str_elem(buf, reader); + if (ret != 0) { + break; + } + elems_read_in_dim[cur_dim]++; + } else { + ret = 1; + break; + } + } else if (strlen(buf) == 0) { + // EOF + ret = 1; + break; + } else if (first) { + if (strcmp(buf, "[") == 0) { + if (cur_dim == dims - 1) { + ret = 1; + break; + } + cur_dim++; + elems_read_in_dim[cur_dim] = 0; + } else { + ret = read_str_elem(buf, reader); + if (ret != 0) { + break; + } + elems_read_in_dim[cur_dim]++; + first = 0; + } + } else { + ret = 1; + break; + } + } + + free(knows_dimsize); + free(elems_read_in_dim); + return ret; +} + +static int read_str_empty_array(char *buf, int bufsize, + const char *type_name, int64_t *shape, int64_t dims) { + if (strlen(buf) == 0) { + // EOF + return 1; + } + + if (strcmp(buf, "empty") != 0) { + return 1; + } + + if (!next_token_is(buf, bufsize, "(")) { + return 1; + } + + for (int i = 0; i < dims-1; i++) { + if (!next_token_is(buf, bufsize, "[")) { + return 1; + } + + if (!next_token_is(buf, bufsize, "]")) { + return 1; + } + } + + if (!next_token_is(buf, bufsize, type_name)) { + return 1; + } + + + if (!next_token_is(buf, bufsize, ")")) { + return 1; + } + + for (int i = 0; i < dims; i++) { + shape[i] = 0; + } + + return 0; +} + +static int read_str_array(int64_t elem_size, str_reader elem_reader, + const char *type_name, + void **data, int64_t *shape, int64_t dims) { + int ret; + struct array_reader reader; + char buf[100]; + + int dims_seen; + for (dims_seen = 0; dims_seen < dims; dims_seen++) { + if (!next_token_is(buf, sizeof(buf), "[")) { + break; + } + } + + if (dims_seen == 0) { + return read_str_empty_array(buf, sizeof(buf), type_name, shape, dims); + } + + if (dims_seen != dims) { + return 1; + } + + reader.shape = shape; + reader.n_elems_used = 0; + reader.elem_size = elem_size; + reader.n_elems_space = 16; + reader.elems = (char*) realloc(*data, elem_size*reader.n_elems_space); + reader.elem_reader = elem_reader; + + ret = read_str_array_elems(buf, sizeof(buf), &reader, dims); + + *data = reader.elems; + + return ret; +} + +#define READ_STR(MACRO, PTR, SUFFIX) \ + remove_underscores(buf); \ + int j; \ + if (sscanf(buf, "%"MACRO"%n", (PTR*)dest, &j) == 1) { \ + return !(strcmp(buf+j, "") == 0 || strcmp(buf+j, SUFFIX) == 0); \ + } else { \ + return 1; \ + } + +static int read_str_i8(char *buf, void* dest) { + /* Some platforms (WINDOWS) does not support scanf %hhd or its + cousin, %SCNi8. Read into int first to avoid corrupting + memory. + + https://gcc.gnu.org/bugzilla/show_bug.cgi?id=63417 */ + remove_underscores(buf); + int j, x; + if (sscanf(buf, "%i%n", &x, &j) == 1) { + *(int8_t*)dest = x; + return !(strcmp(buf+j, "") == 0 || strcmp(buf+j, "i8") == 0); + } else { + return 1; + } +} + +static int read_str_u8(char *buf, void* dest) { + /* Some platforms (WINDOWS) does not support scanf %hhd or its + cousin, %SCNu8. Read into int first to avoid corrupting + memory. + + https://gcc.gnu.org/bugzilla/show_bug.cgi?id=63417 */ + remove_underscores(buf); + int j, x; + if (sscanf(buf, "%i%n", &x, &j) == 1) { + *(uint8_t*)dest = x; + return !(strcmp(buf+j, "") == 0 || strcmp(buf+j, "u8") == 0); + } else { + return 1; + } +} + +static int read_str_i16(char *buf, void* dest) { + READ_STR(SCNi16, int16_t, "i16"); +} + +static int read_str_u16(char *buf, void* dest) { + READ_STR(SCNi16, int16_t, "u16"); +} + +static int read_str_i32(char *buf, void* dest) { + READ_STR(SCNi32, int32_t, "i32"); +} + +static int read_str_u32(char *buf, void* dest) { + READ_STR(SCNi32, int32_t, "u32"); +} + +static int read_str_i64(char *buf, void* dest) { + READ_STR(SCNi64, int64_t, "i64"); +} + +static int read_str_u64(char *buf, void* dest) { + // FIXME: This is not correct, as SCNu64 only permits decimal + // literals. However, SCNi64 does not handle very large numbers + // correctly (it's really for signed numbers, so that's fair). + READ_STR(SCNu64, uint64_t, "u64"); +} + +static int read_str_f32(char *buf, void* dest) { + remove_underscores(buf); + if (strcmp(buf, "f32.nan") == 0) { + *(float*)dest = NAN; + return 0; + } else if (strcmp(buf, "f32.inf") == 0) { + *(float*)dest = INFINITY; + return 0; + } else if (strcmp(buf, "-f32.inf") == 0) { + *(float*)dest = -INFINITY; + return 0; + } else { + READ_STR("f", float, "f32"); + } +} + +static int read_str_f64(char *buf, void* dest) { + remove_underscores(buf); + if (strcmp(buf, "f64.nan") == 0) { + *(double*)dest = NAN; + return 0; + } else if (strcmp(buf, "f64.inf") == 0) { + *(double*)dest = INFINITY; + return 0; + } else if (strcmp(buf, "-f64.inf") == 0) { + *(double*)dest = -INFINITY; + return 0; + } else { + READ_STR("lf", double, "f64"); + } +} + +static int read_str_bool(char *buf, void* dest) { + if (strcmp(buf, "true") == 0) { + *(char*)dest = 1; + return 0; + } else if (strcmp(buf, "false") == 0) { + *(char*)dest = 0; + return 0; + } else { + return 1; + } +} + +static int write_str_i8(FILE *out, int8_t *src) { + return fprintf(out, "%hhdi8", *src); +} + +static int write_str_u8(FILE *out, uint8_t *src) { + return fprintf(out, "%hhuu8", *src); +} + +static int write_str_i16(FILE *out, int16_t *src) { + return fprintf(out, "%hdi16", *src); +} + +static int write_str_u16(FILE *out, uint16_t *src) { + return fprintf(out, "%huu16", *src); +} + +static int write_str_i32(FILE *out, int32_t *src) { + return fprintf(out, "%di32", *src); +} + +static int write_str_u32(FILE *out, uint32_t *src) { + return fprintf(out, "%uu32", *src); +} + +static int write_str_i64(FILE *out, int64_t *src) { + return fprintf(out, "%"PRIi64"i64", *src); +} + +static int write_str_u64(FILE *out, uint64_t *src) { + return fprintf(out, "%"PRIu64"u64", *src); +} + +static int write_str_f32(FILE *out, float *src) { + float x = *src; + if (isnan(x)) { + return fprintf(out, "f32.nan"); + } else if (isinf(x) && x >= 0) { + return fprintf(out, "f32.inf"); + } else if (isinf(x)) { + return fprintf(out, "-f32.inf"); + } else { + return fprintf(out, "%.6ff32", x); + } +} + +static int write_str_f64(FILE *out, double *src) { + double x = *src; + if (isnan(x)) { + return fprintf(out, "f64.nan"); + } else if (isinf(x) && x >= 0) { + return fprintf(out, "f64.inf"); + } else if (isinf(x)) { + return fprintf(out, "-f64.inf"); + } else { + return fprintf(out, "%.6ff64", *src); + } +} + +static int write_str_bool(FILE *out, void *src) { + return fprintf(out, *(char*)src ? "true" : "false"); +} + +//// Binary I/O + +#define BINARY_FORMAT_VERSION 2 +#define IS_BIG_ENDIAN (!*(unsigned char *)&(uint16_t){1}) + +// Reading little-endian byte sequences. On big-endian hosts, we flip +// the resulting bytes. + +static int read_byte(void* dest) { + int num_elems_read = fread(dest, 1, 1, stdin); + return num_elems_read == 1 ? 0 : 1; +} + +static int read_le_2byte(void* dest) { + uint16_t x; + int num_elems_read = fread(&x, 2, 1, stdin); + if (IS_BIG_ENDIAN) { + x = (x>>8) | (x<<8); + } + *(uint16_t*)dest = x; + return num_elems_read == 1 ? 0 : 1; +} + +static int read_le_4byte(void* dest) { + uint32_t x; + int num_elems_read = fread(&x, 4, 1, stdin); + if (IS_BIG_ENDIAN) { + x = + ((x>>24)&0xFF) | + ((x>>8) &0xFF00) | + ((x<<8) &0xFF0000) | + ((x<<24)&0xFF000000); + } + *(uint32_t*)dest = x; + return num_elems_read == 1 ? 0 : 1; +} + +static int read_le_8byte(void* dest) { + uint64_t x; + int num_elems_read = fread(&x, 8, 1, stdin); + if (IS_BIG_ENDIAN) { + x = + ((x>>56)&0xFFull) | + ((x>>40)&0xFF00ull) | + ((x>>24)&0xFF0000ull) | + ((x>>8) &0xFF000000ull) | + ((x<<8) &0xFF00000000ull) | + ((x<<24)&0xFF0000000000ull) | + ((x<<40)&0xFF000000000000ull) | + ((x<<56)&0xFF00000000000000ull); + } + *(uint64_t*)dest = x; + return num_elems_read == 1 ? 0 : 1; +} + +static int write_byte(void* dest) { + int num_elems_written = fwrite(dest, 1, 1, stdin); + return num_elems_written == 1 ? 0 : 1; +} + +static int write_le_2byte(void* dest) { + uint16_t x = *(uint16_t*)dest; + if (IS_BIG_ENDIAN) { + x = (x>>8) | (x<<8); + } + int num_elems_written = fwrite(&x, 2, 1, stdin); + return num_elems_written == 1 ? 0 : 1; +} + +static int write_le_4byte(void* dest) { + uint32_t x = *(uint32_t*)dest; + if (IS_BIG_ENDIAN) { + x = + ((x>>24)&0xFF) | + ((x>>8) &0xFF00) | + ((x<<8) &0xFF0000) | + ((x<<24)&0xFF000000); + } + int num_elems_written = fwrite(&x, 4, 1, stdin); + return num_elems_written == 1 ? 0 : 1; +} + +static int write_le_8byte(void* dest) { + uint64_t x = *(uint64_t*)dest; + if (IS_BIG_ENDIAN) { + x = + ((x>>56)&0xFFull) | + ((x>>40)&0xFF00ull) | + ((x>>24)&0xFF0000ull) | + ((x>>8) &0xFF000000ull) | + ((x<<8) &0xFF00000000ull) | + ((x<<24)&0xFF0000000000ull) | + ((x<<40)&0xFF000000000000ull) | + ((x<<56)&0xFF00000000000000ull); + } + int num_elems_written = fwrite(&x, 8, 1, stdin); + return num_elems_written == 1 ? 0 : 1; +} + +//// Types + +struct primtype_info_t { + const char binname[4]; // Used for parsing binary data. + const char* type_name; // Same name as in Futhark. + const int size; // in bytes + const writer write_str; // Write in text format. + const str_reader read_str; // Read in text format. + const writer write_bin; // Write in binary format. + const bin_reader read_bin; // Read in binary format. +}; + +static const struct primtype_info_t i8_info = + {.binname = " i8", .type_name = "i8", .size = 1, + .write_str = (writer)write_str_i8, .read_str = (str_reader)read_str_i8, + .write_bin = (writer)write_byte, .read_bin = (bin_reader)read_byte}; +static const struct primtype_info_t i16_info = + {.binname = " i16", .type_name = "i16", .size = 2, + .write_str = (writer)write_str_i16, .read_str = (str_reader)read_str_i16, + .write_bin = (writer)write_le_2byte, .read_bin = (bin_reader)read_le_2byte}; +static const struct primtype_info_t i32_info = + {.binname = " i32", .type_name = "i32", .size = 4, + .write_str = (writer)write_str_i32, .read_str = (str_reader)read_str_i32, + .write_bin = (writer)write_le_4byte, .read_bin = (bin_reader)read_le_4byte}; +static const struct primtype_info_t i64_info = + {.binname = " i64", .type_name = "i64", .size = 8, + .write_str = (writer)write_str_i64, .read_str = (str_reader)read_str_i64, + .write_bin = (writer)write_le_8byte, .read_bin = (bin_reader)read_le_8byte}; +static const struct primtype_info_t u8_info = + {.binname = " u8", .type_name = "u8", .size = 1, + .write_str = (writer)write_str_u8, .read_str = (str_reader)read_str_u8, + .write_bin = (writer)write_byte, .read_bin = (bin_reader)read_byte}; +static const struct primtype_info_t u16_info = + {.binname = " u16", .type_name = "u16", .size = 2, + .write_str = (writer)write_str_u16, .read_str = (str_reader)read_str_u16, + .write_bin = (writer)write_le_2byte, .read_bin = (bin_reader)read_le_2byte}; +static const struct primtype_info_t u32_info = + {.binname = " u32", .type_name = "u32", .size = 4, + .write_str = (writer)write_str_u32, .read_str = (str_reader)read_str_u32, + .write_bin = (writer)write_le_4byte, .read_bin = (bin_reader)read_le_4byte}; +static const struct primtype_info_t u64_info = + {.binname = " u64", .type_name = "u64", .size = 8, + .write_str = (writer)write_str_u64, .read_str = (str_reader)read_str_u64, + .write_bin = (writer)write_le_8byte, .read_bin = (bin_reader)read_le_8byte}; +static const struct primtype_info_t f32_info = + {.binname = " f32", .type_name = "f32", .size = 4, + .write_str = (writer)write_str_f32, .read_str = (str_reader)read_str_f32, + .write_bin = (writer)write_le_4byte, .read_bin = (bin_reader)read_le_4byte}; +static const struct primtype_info_t f64_info = + {.binname = " f64", .type_name = "f64", .size = 8, + .write_str = (writer)write_str_f64, .read_str = (str_reader)read_str_f64, + .write_bin = (writer)write_le_8byte, .read_bin = (bin_reader)read_le_8byte}; +static const struct primtype_info_t bool_info = + {.binname = "bool", .type_name = "bool", .size = 1, + .write_str = (writer)write_str_bool, .read_str = (str_reader)read_str_bool, + .write_bin = (writer)write_byte, .read_bin = (bin_reader)read_byte}; + +static const struct primtype_info_t* primtypes[] = { + &i8_info, &i16_info, &i32_info, &i64_info, + &u8_info, &u16_info, &u32_info, &u64_info, + &f32_info, &f64_info, + &bool_info, + NULL // NULL-terminated +}; + +// General value interface. All endian business taken care of at +// lower layers. + +static int read_is_binary() { + skipspaces(); + int c = getchar(); + if (c == 'b') { + int8_t bin_version; + int ret = read_byte(&bin_version); + + if (ret != 0) { panic(1, "binary-input: could not read version.\n"); } + + if (bin_version != BINARY_FORMAT_VERSION) { + panic(1, "binary-input: File uses version %i, but I only understand version %i.\n", + bin_version, BINARY_FORMAT_VERSION); + } + + return 1; + } + ungetc(c, stdin); + return 0; +} + +static const struct primtype_info_t* read_bin_read_type_enum() { + char read_binname[4]; + + int num_matched = scanf("%4c", read_binname); + if (num_matched != 1) { panic(1, "binary-input: Couldn't read element type.\n"); } + + const struct primtype_info_t **type = primtypes; + + for (; *type != NULL; type++) { + // I compare the 4 characters manually instead of using strncmp because + // this allows any value to be used, also NULL bytes + if (memcmp(read_binname, (*type)->binname, 4) == 0) { + return *type; + } + } + panic(1, "binary-input: Did not recognize the type '%s'.\n", read_binname); + return NULL; +} + +static void read_bin_ensure_scalar(const struct primtype_info_t *expected_type) { + int8_t bin_dims; + int ret = read_byte(&bin_dims); + if (ret != 0) { panic(1, "binary-input: Couldn't get dims.\n"); } + + if (bin_dims != 0) { + panic(1, "binary-input: Expected scalar (0 dimensions), but got array with %i dimensions.\n", + bin_dims); + } + + const struct primtype_info_t *bin_type = read_bin_read_type_enum(); + if (bin_type != expected_type) { + panic(1, "binary-input: Expected scalar of type %s but got scalar of type %s.\n", + expected_type->type_name, + bin_type->type_name); + } +} + +//// High-level interface + +static int read_bin_array(const struct primtype_info_t *expected_type, void **data, int64_t *shape, int64_t dims) { + int ret; + + int8_t bin_dims; + ret = read_byte(&bin_dims); + if (ret != 0) { panic(1, "binary-input: Couldn't get dims.\n"); } + + if (bin_dims != dims) { + panic(1, "binary-input: Expected %i dimensions, but got array with %i dimensions.\n", + dims, bin_dims); + } + + const struct primtype_info_t *bin_primtype = read_bin_read_type_enum(); + if (expected_type != bin_primtype) { + panic(1, "binary-input: Expected %iD-array with element type '%s' but got %iD-array with element type '%s'.\n", + dims, expected_type->type_name, dims, bin_primtype->type_name); + } + + uint64_t elem_count = 1; + for (int i=0; isize; + void* tmp = realloc(*data, elem_count * elem_size); + if (tmp == NULL) { + panic(1, "binary-input: Failed to allocate array of size %i.\n", + elem_count * elem_size); + } + *data = tmp; + + size_t num_elems_read = fread(*data, elem_size, elem_count, stdin); + if (num_elems_read != elem_count) { + panic(1, "binary-input: tried to read %i elements of an array, but only got %i elements.\n", + elem_count, num_elems_read); + } + + // If we're on big endian platform we must change all multibyte elements + // from using little endian to big endian + if (IS_BIG_ENDIAN && elem_size != 1) { + char* elems = (char*) *data; + for (uint64_t i=0; isize, (str_reader)expected_type->read_str, expected_type->type_name, data, shape, dims); + } else { + return read_bin_array(expected_type, data, shape, dims); + } +} + +static int write_str_array(FILE *out, const struct primtype_info_t *elem_type, unsigned char *data, int64_t *shape, int8_t rank) { + if (rank==0) { + elem_type->write_str(out, (void*)data); + } else { + int64_t len = shape[0]; + int64_t slice_size = 1; + + int64_t elem_size = elem_type->size; + for (int64_t i = 1; i < rank; i++) { + slice_size *= shape[i]; + } + + if (len*slice_size == 0) { + printf("empty("); + for (int64_t i = 1; i < rank; i++) { + printf("[]"); + } + printf("%s", elem_type->type_name); + printf(")"); + } else if (rank==1) { + putchar('['); + for (int64_t i = 0; i < len; i++) { + elem_type->write_str(out, (void*) (data + i * elem_size)); + if (i != len-1) { + printf(", "); + } + } + putchar(']'); + } else { + putchar('['); + for (int64_t i = 0; i < len; i++) { + write_str_array(out, elem_type, data + i * slice_size * elem_size, shape+1, rank-1); + if (i != len-1) { + printf(", "); + } + } + putchar(']'); + } + } + return 0; +} + +static int write_bin_array(FILE *out, const struct primtype_info_t *elem_type, unsigned char *data, int64_t *shape, int8_t rank) { + int64_t num_elems = 1; + for (int64_t i = 0; i < rank; i++) { + num_elems *= shape[i]; + } + + fputc('b', out); + fputc((char)BINARY_FORMAT_VERSION, out); + fwrite(&rank, sizeof(int8_t), 1, out); + fputs(elem_type->binname, out); + fwrite(shape, sizeof(int64_t), rank, out); + + if (IS_BIG_ENDIAN) { + for (int64_t i = 0; i < num_elems; i++) { + unsigned char *elem = data+i*elem_type->size; + for (int64_t j = 0; j < elem_type->size; j++) { + fwrite(&elem[elem_type->size-j], 1, 1, out); + } + } + } else { + fwrite(data, elem_type->size, num_elems, out); + } + + return 0; +} + +static int write_array(FILE *out, int write_binary, + const struct primtype_info_t *elem_type, void *data, int64_t *shape, int8_t rank) { + if (write_binary) { + return write_bin_array(out, elem_type, data, shape, rank); + } else { + return write_str_array(out, elem_type, data, shape, rank); + } +} + +static int read_scalar(const struct primtype_info_t *expected_type, void *dest) { + if (!read_is_binary()) { + char buf[100]; + next_token(buf, sizeof(buf)); + return expected_type->read_str(buf, dest); + } else { + read_bin_ensure_scalar(expected_type); + return expected_type->read_bin(dest); + } +} + +static int write_scalar(FILE *out, int write_binary, const struct primtype_info_t *type, void *src) { + if (write_binary) { + return write_bin_array(out, type, src, NULL, 0); + } else { + return type->write_str(out, src); + } +} diff --git a/rts/csharp/exceptions.cs b/rts/csharp/exceptions.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/exceptions.cs @@ -0,0 +1,6 @@ +private class ValueError : Exception +{ + public ValueError(){} + public ValueError(string message):base(message){} + public ValueError(string message, Exception inner):base(message, inner){} +} diff --git a/rts/csharp/functions.cs b/rts/csharp/functions.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/functions.cs diff --git a/rts/csharp/memory.cs b/rts/csharp/memory.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/memory.cs @@ -0,0 +1,457 @@ +public struct FlatArray +{ + public long[] shape; + public T[] array; + + public FlatArray(T[] data_array, long[] shape_array) + { + shape = shape_array; + array = data_array; + } + + public FlatArray(T[] data_array) + { + shape = new long[] {data_array.Length}; + array = data_array; + } + + private long getIdx(int[] idxs) + { + long idx = 0; + for (int i = 0; i", desc); + } +} + +private byte[] allocateMem(sbyte size) +{ + return new byte[size]; +} + +private byte[] allocateMem(short size) +{ + return new byte[size]; +} + +private byte[] allocateMem(int size) +{ + return new byte[size]; +} + +private byte[] allocateMem(long size) +{ + return new byte[size]; +} + +private byte[] allocateMem(byte size) +{ + return new byte[size]; +} + +private byte[] allocateMem(ushort size) +{ + return new byte[size]; +} + +private byte[] allocateMem(uint size) +{ + return new byte[size]; +} + +private byte[] allocateMem(ulong size) +{ + return new byte[size]; +} + +private Tuple createTuple_byte(byte[] bytes, long[] shape) +{ + var byteArray = new byte[bytes.Length / sizeof(byte)]; + Buffer.BlockCopy(bytes, 0, byteArray, 0, bytes.Length); + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_ushort(byte[] bytes, long[] shape) +{ + var ushortArray = new ushort[bytes.Length / sizeof(ushort)]; + Buffer.BlockCopy(bytes, 0, ushortArray, 0, bytes.Length); + return Tuple.Create(ushortArray, shape); +} + +private Tuple createTuple_uint(byte[] bytes, long[] shape) +{ + var uintArray = new uint[bytes.Length / sizeof(uint)]; + Buffer.BlockCopy(bytes, 0, uintArray, 0, bytes.Length); + return Tuple.Create(uintArray, shape); +} + +private Tuple createTuple_ulong(byte[] bytes, long[] shape) +{ + var ulongArray = new ulong[bytes.Length / sizeof(ulong)]; + Buffer.BlockCopy(bytes, 0, ulongArray, 0, bytes.Length); + return Tuple.Create(ulongArray, shape); +} + + +private Tuple createTuple_sbyte(byte[] bytes, long[] shape) +{ + var sbyteArray = new sbyte[1]; + if (bytes.Length > 0) + { + sbyteArray = new sbyte[bytes.Length / sizeof(sbyte)]; + } + Buffer.BlockCopy(bytes, 0, sbyteArray, 0, bytes.Length); + return Tuple.Create(sbyteArray, shape); +} + + +private Tuple createTuple_short(byte[] bytes, long[] shape) +{ + var shortArray = new short[1]; + if (bytes.Length > 0) + { + shortArray = new short[bytes.Length / sizeof(short)]; + } + Buffer.BlockCopy(bytes, 0, shortArray, 0, bytes.Length); + return Tuple.Create(shortArray, shape); +} + +private Tuple createTuple_int(byte[] bytes, long[] shape) +{ + var intArray = new int[1]; + if (bytes.Length > 0) + { + intArray = new int[bytes.Length / sizeof(int)]; + } + Buffer.BlockCopy(bytes, 0, intArray, 0, bytes.Length); + return Tuple.Create(intArray, shape); +} + +private Tuple createTuple_long(byte[] bytes, long[] shape) +{ + var longArray = new long[1]; + if (bytes.Length > 0) + { + longArray = new long[bytes.Length / sizeof(long)]; + } + Buffer.BlockCopy(bytes, 0, longArray, 0, bytes.Length); + return Tuple.Create(longArray, shape); +} + +private Tuple createTuple_float(byte[] bytes, long[] shape) +{ + var floatArray = new float[1]; + if (bytes.Length > 0) + { + floatArray = new float[bytes.Length / sizeof(float)]; + } + Buffer.BlockCopy(bytes, 0, floatArray, 0, bytes.Length); + return Tuple.Create(floatArray, shape); +} + + +private Tuple createTuple_double(byte[] bytes, long[] shape) +{ + var doubleArray = new double[1]; + if (bytes.Length > 0) + { + doubleArray = new double[bytes.Length / sizeof(double)]; + } + Buffer.BlockCopy(bytes, 0, doubleArray, 0, bytes.Length); + return Tuple.Create(doubleArray, shape); +} + +private Tuple createTuple_bool(byte[] bytes, long[] shape) +{ + var boolArray = new bool[1]; + if (bytes.Length > 0) + { + boolArray = new bool[bytes.Length / sizeof(bool)]; + } + Buffer.BlockCopy(bytes, 0, boolArray, 0, bytes.Length); + return Tuple.Create(boolArray, shape); +} + +private byte[] unwrapArray(Array src, int obj_size) +{ + var bytes = new byte[src.Length * obj_size]; + Buffer.BlockCopy(src, 0, bytes, 0, bytes.Length); + return bytes; +} + +private byte indexArray_byte(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(byte*) dest_ptr; + } + } +} + +private ushort indexArray_ushort(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(ushort*) dest_ptr; + } + } +} + +private uint indexArray_uint(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(uint*) dest_ptr; + } + } +} + +private ulong indexArray_ulong(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(ulong*) dest_ptr; + } + } +} + +private sbyte indexArray_sbyte(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(sbyte*) dest_ptr; + } + } +} + +private short indexArray_short(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(short*) dest_ptr; + } + } +} + +private int indexArray_int(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(int*) dest_ptr; + } + } +} + +private long indexArray_long(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(long*) dest_ptr; + } + } +} + +private float indexArray_float(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(float*) dest_ptr; + } + } +} + +private double indexArray_double(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(double*) dest_ptr; + } + } +} + +private bool indexArray_bool(byte[] src, int offset) +{ + unsafe + { + fixed (void* dest_ptr = &src[offset]) + { + return *(bool*) dest_ptr; + } + } +} + +private void writeScalarArray(byte[] dest, int offset, sbyte value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(sbyte*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, byte value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(byte*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, short value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(short*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, ushort value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(ushort*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, int value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(int*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, uint value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(uint*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, long value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(long*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, ulong value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(ulong*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, float value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(float*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, double value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(double*) dest_ptr = value; + } + } +} +private void writeScalarArray(byte[] dest, int offset, bool value) +{ + unsafe + { + fixed (byte* dest_ptr = &dest[offset]) + { + *(bool*) dest_ptr = value; + } + } +} diff --git a/rts/csharp/memory_opencl.cs b/rts/csharp/memory_opencl.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/memory_opencl.cs @@ -0,0 +1,231 @@ +private Tuple createTuple_byte(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new byte[1]; + if (nbytes > 0) + { + byteArray = new byte[nbytes / sizeof(byte)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_ushort(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new ushort[1]; + if (nbytes > 0) + { + byteArray = new ushort[nbytes / sizeof(ushort)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_uint(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new uint[1]; + if (nbytes > 0) + { + byteArray = new uint[nbytes / sizeof(uint)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_ulong(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new ulong[1]; + if (nbytes > 0) + { + byteArray = new ulong[nbytes / sizeof(ulong)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_sbyte(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new sbyte[1]; + if (nbytes > 0) + { + byteArray = new sbyte[nbytes / sizeof(sbyte)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_short(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new short[1]; + if (nbytes > 0) + { + byteArray = new short[nbytes / sizeof(short)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_int(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new int[1]; + if (nbytes > 0) + { + byteArray = new int[nbytes / sizeof(int)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_long(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new long[1]; + if (nbytes > 0) + { + byteArray = new long[nbytes / sizeof(long)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_float(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new float[1]; + if (nbytes > 0) + { + byteArray = new float[nbytes / sizeof(float)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_double(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new double[1]; + if (nbytes > 0) + { + byteArray = new double[nbytes / sizeof(double)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + +private Tuple createTuple_bool(CLMemoryHandle mem, CLCommandQueueHandle queue, + int nbytes, long[] shape) +{ + var byteArray = new bool[1]; + if (nbytes > 0) + { + byteArray = new bool[nbytes / sizeof(bool)]; + } + unsafe + { + fixed(void* ptr = &byteArray[0]) + { + CL10.EnqueueReadBuffer(queue, mem, true, + new IntPtr(0), new IntPtr(nbytes), new IntPtr(ptr), + 0, null, null + ); + } + } + return Tuple.Create(byteArray, shape); +} + diff --git a/rts/csharp/opencl.cs b/rts/csharp/opencl.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/opencl.cs @@ -0,0 +1,926 @@ +// Stub code for OpenCL setup. + +private void OPENCL_SUCCEED(int return_code, + [CallerFilePath] string filePath = "", + [CallerLineNumber] int lineNumber = 0) +{ + OpenCLSucceed(return_code, "", filePath, lineNumber); +} + +private void OPENCL_SUCCEED(ComputeErrorCode return_code, + [CallerFilePath] string filePath = "", + [CallerLineNumber] int lineNumber = 0) +{ + OpenCLSucceed((int) return_code, "", filePath, lineNumber); +} + +private void OPENCL_SUCCEED(object return_code, + [CallerFilePath] string filePath = "", + [CallerLineNumber] int lineNumber = 0) +{ + OpenCLSucceed((int) return_code, "", filePath, lineNumber); +} + +public struct OpenCLConfig +{ + public bool Debugging; + public int PreferredDeviceNum; + public string PreferredPlatform; + public string PreferredDevice; + + public string DumpProgramTo; + public string LoadProgramFrom; + + public int DefaultGroupSize; + public int DefaultNumGroups; + public int DefaultTileSize; + public int DefaultThreshold; + public int TransposeBlockDim; + + public int NumSizes; + public string[] SizeNames; + public int[] SizeValues; + public string[] SizeClasses; +} + +private void MemblockUnrefDevice(ref FutharkContext + context, ref OpenCLMemblock block, string desc) +{ + if (!block.IsNull) + { + block.DecreaseRefs(); + if (context.DetailMemory) + { + Console.Error.WriteLine(String.Format( + "Unreferencing block {0} (allocated as {1}) in {2}: {3} references remaining.", + desc, block.Tag, "space 'device'", block.References)); + } + + if (block.References == 0) + { + context.CurrentMemUsageDevice -= block.Size; + OPENCL_SUCCEED(OpenCLFree(ref context, block.Mem, block.Tag)); + block.IsNull = true; + } + + if (context.DetailMemory) + { + Console.Error.WriteLine(String.Format( + "{0} bytes freed (now allocated: {1} bytes)", + block.Size, context.CurrentMemUsageDevice)); + } + } +} + +private void MemblockSetDevice(ref FutharkContext context, + ref OpenCLMemblock lhs, ref OpenCLMemblock rhs, string lhs_desc) +{ + MemblockUnrefDevice(ref context, ref lhs, lhs_desc); + rhs.IncreaseRefs(); + lhs = rhs; +} + +private OpenCLMemblock MemblockAllocDevice(ref FutharkContext context, OpenCLMemblock block, long size, string desc) +{ + if (size < 0) + { + panic(1, String.Format("Negative allocation of {0} bytes attempted for {1} in {2}", + size, desc)); + } + + MemblockUnrefDevice(ref context, ref block, desc); + OPENCL_SUCCEED(OpenCLAlloc(ref context, size, desc, ref block.Mem)); + + block.References = 1; + block.IsNull = false; + block.Size = size; + block.Tag = desc; + context.CurrentMemUsageDevice += size; + + if (context.DetailMemory) + { + Console.Error.Write(String.Format("Allocated {0} bytes for {1} in {2} (now allocated: {3} bytes)", + size, desc, "space 'device'", Ctx.CurrentMemUsageDevice)); + } + + if (context.CurrentMemUsageDevice > context.PeakMemUsageDevice) + { + context.PeakMemUsageDevice = context.CurrentMemUsageDevice; + if (context.DetailMemory) + { + Console.Error.Write(" (new peak).\n"); + } + } + else if (context.DetailMemory) + { + Console.Error.Write(".\n"); + } + + return block; +} + + +private bool FreeListFind(ref OpenCLFreeList free_list, string tag, ref long size_out, ref CLMemoryHandle mem_out) +{ + for (int i = 0; i < free_list.Capacity; i++) + { + if (free_list.Entries[i].Valid && free_list.Entries[i].Tag == tag) + { + free_list.Entries[i].Valid = false; + size_out = free_list.Entries[i].Size; + mem_out = free_list.Entries[i].Mem; + free_list.Used--; + return true; + } + } + + return false; +} + +private bool FreeListFirst(ref OpenCLFreeList free_list, ref CLMemoryHandle mem_out) +{ + for (int i = 0; i < free_list.Capacity; i++) + { + if (free_list.Entries[i].Valid) + { + free_list.Entries[i].Valid = false; + mem_out = free_list.Entries[i].Mem; + free_list.Used--; + return true; + } + } + return false; +} + +private ComputeErrorCode OpenCLAllocActual(ref FutharkContext context, long min_size, ref CLMemoryHandle mem) +{ + ComputeErrorCode error; + mem = CL10.CreateBuffer(context.OpenCL.Context, ComputeMemoryFlags.ReadWrite + , new IntPtr(min_size), IntPtr.Zero, out error); + + if (error != ComputeErrorCode.Success) + { + return error; + } + + int x = 2; + unsafe + { + error = CL10.EnqueueWriteBuffer(Ctx.OpenCL.Queue, mem, true, IntPtr.Zero, new IntPtr(sizeof(int)), new IntPtr(&x), 0, null, null); + } + + return error; +} + +private ComputeErrorCode OpenCLAlloc(ref FutharkContext context, long min_size, string tag, ref CLMemoryHandle mem_out) +{ + if (min_size < 0) + { + panic(1, "Tried to allocate a negative amount of bytes."); + } + + min_size = (min_size < sizeof(int)) ? sizeof(int) : min_size; + + long size = 0; + + if (FreeListFind(ref context.FreeList, tag, ref size, ref mem_out)) + { + if (size >= min_size && size <= min_size * 2) + { + return ComputeErrorCode.Success; + } + else + { + ComputeErrorCode code1 = CL10.ReleaseMemObject(mem_out); + if (code1 != ComputeErrorCode.Success) + { + return code1; + } + } + } + + ComputeErrorCode error = OpenCLAllocActual(ref context, min_size, ref mem_out); + while (error == ComputeErrorCode.MemoryObjectAllocationFailure) + { + CLMemoryHandle mem = Ctx.EMPTY_MEM_HANDLE; + if (FreeListFirst(ref context.FreeList, ref mem)) + { + error = CL10.ReleaseMemObject(mem); + if (error != ComputeErrorCode.Success) + { + return error; + } + } + else + { + break; + } + + error = OpenCLAllocActual(ref context, min_size, ref mem_out); + } + return error; +} + + +private ComputeErrorCode OpenCLFree(ref FutharkContext context, CLMemoryHandle mem, string tag) +{ + long size = 0; + CLMemoryHandle existing_mem = Ctx.EMPTY_MEM_HANDLE; + ComputeErrorCode error = ComputeErrorCode.Success; + if (FreeListFind(ref context.FreeList, tag, ref size, ref existing_mem)) + { + error = CL10.ReleaseMemObject(existing_mem); + if (error != ComputeErrorCode.Success) + { + return error; + } + } + + if (existing_mem.Value == mem.Value) + { + return error; + } + + var trash_null = new IntPtr(0); + unsafe + { + error = CL10.GetMemObjectInfo(mem, ComputeMemoryInfo.Size, + new IntPtr(sizeof(long)), new IntPtr(&size), out trash_null); + } + + if (error == ComputeErrorCode.Success) + { + FreeListInsert(ref context, size, mem, tag); + } + return error; +} + +private void FreeListInsert(ref FutharkContext context, long size, CLMemoryHandle mem, string tag) +{ + int i = FreeListFindInvalid(ref context); + if (i == context.FreeList.Capacity) + { + var cap = context.FreeList.Capacity; + int new_capacity = cap * 2; + Array.Resize(ref context.FreeList.Entries, new_capacity); + for (int j = 0; j < cap; j++) + { + var entry = new OpenCLFreeListEntry(); + entry.Valid = false; + context.FreeList.Entries[cap + j] = entry; + } + + context.FreeList.Capacity *= 2; + } + + context.FreeList.Entries[i].Valid = true; + context.FreeList.Entries[i].Size = size; + context.FreeList.Entries[i].Tag = tag; + context.FreeList.Entries[i].Mem = mem; + context.FreeList.Used++; +} + +private int FreeListFindInvalid(ref FutharkContext context) +{ + int i; + for (i = 0; i < context.FreeList.Capacity; i++) + { + if (!context.FreeList.Entries[i].Valid) + { + break; + } + } + + return i; +} + +private class OpenCLMemblock +{ + public int References; + public CLMemoryHandle Mem; + public long Size; + public string Tag; + public bool IsNull; + + public void IncreaseRefs() + { + this.References += 1; + } + + public void DecreaseRefs() + { + this.References -= 1; + } +} + +private OpenCLMemblock EmptyMemblock(CLMemoryHandle mem) +{ + var block = new OpenCLMemblock(); + block.Mem = mem; + block.References = 0; + block.Tag = ""; + block.Size = 0; + block.IsNull = true; + + return block; +} + +public struct OpenCLFreeListEntry +{ + public bool Valid; + public CLMemoryHandle Mem; + public long Size; + public string Tag; +} + +public struct OpenCLFreeList +{ + public OpenCLFreeListEntry[] Entries; + public int Capacity; + public int Used; +} + + +private OpenCLFreeList OpenCLFreeListInit() +{ + int CAPACITY = 30; // arbitrarily chosen + var free_list = new OpenCLFreeList(); + free_list.Entries = Enumerable.Range(0, CAPACITY) + .Select(_ => + { + var entry = new OpenCLFreeListEntry(); + entry.Valid = false; + return entry; + }).ToArray(); + + free_list.Capacity = CAPACITY; + free_list.Used = 0; + + return free_list; +} + + +private void OpenCLConfigInit(out OpenCLConfig cfg, + int num_sizes, + string[] size_names, + int[] size_values, + string[] size_classes) +{ + cfg.Debugging = false; + cfg.PreferredDeviceNum = 0; + cfg.PreferredPlatform = ""; + cfg.PreferredDevice = ""; + cfg.DumpProgramTo = null; + cfg.LoadProgramFrom = null; + + cfg.DefaultGroupSize = 256; + cfg.DefaultNumGroups = 128; + cfg.DefaultTileSize = 32; + cfg.DefaultThreshold = 32*1024; + cfg.TransposeBlockDim = 16; + + cfg.NumSizes = num_sizes; + cfg.SizeNames = size_names; + cfg.SizeValues = size_values; + cfg.SizeClasses = size_classes; +} + +public struct OpenCLContext { + public CLPlatformHandle Platform; + public CLDeviceHandle Device; + public CLContextHandle Context; + public CLCommandQueueHandle Queue; + + public OpenCLConfig Cfg; + + public int MaxGroupSize; + public int MaxNumGroups; + public int MaxTileSize; + public int MaxThreshold; + + public int LockstepWidth; +} + +public struct OpenCLDeviceOption { + public CLPlatformHandle Platform; + public CLDeviceHandle Device; + public ComputeDeviceTypes DeviceType; + public string PlatformName; + public string DeviceName; +}; + +/* This function must be defined by the user. It is invoked by + setup_opencl() after the platform and device has been found, but + before the program is loaded. Its intended use is to tune + constants based on the selected platform and device. */ + +private string OpenCLErrorString(int err) +{ + switch ((ComputeErrorCode) err) { + case ComputeErrorCode.Success: return "Success!"; + case ComputeErrorCode.DeviceNotFound: return "Device not found."; + case ComputeErrorCode.DeviceNotAvailable: return "Device not available"; + case ComputeErrorCode.CompilerNotAvailable: return "Compiler not available"; + case ComputeErrorCode.MemoryObjectAllocationFailure: return "Memory object allocation failure"; + case ComputeErrorCode.OutOfResources: return "Out of resources"; + case ComputeErrorCode.OutOfHostMemory: return "Out of host memory"; + case ComputeErrorCode.ProfilingInfoNotAvailable: return "Profiling information not available"; + case ComputeErrorCode.MemoryCopyOverlap: return "Memory copy overlap"; + case ComputeErrorCode.ImageFormatMismatch: return "Image format mismatch"; + case ComputeErrorCode.ImageFormatNotSupported: return "Image format not supported"; + case ComputeErrorCode.BuildProgramFailure: return "Program build failure"; + case ComputeErrorCode.MapFailure: return "Map failure"; + case ComputeErrorCode.InvalidValue: return "Invalid value"; + case ComputeErrorCode.InvalidDeviceType: return "Invalid device type"; + case ComputeErrorCode.InvalidPlatform: return "Invalid platform"; + case ComputeErrorCode.InvalidDevice: return "Invalid device"; + case ComputeErrorCode.InvalidContext: return "Invalid context"; + case ComputeErrorCode.InvalidCommandQueueFlags: return "Invalid queue properties"; + case ComputeErrorCode.InvalidCommandQueue: return "Invalid command queue"; + case ComputeErrorCode.InvalidHostPointer: return "Invalid host pointer"; + case ComputeErrorCode.InvalidMemoryObject: return "Invalid memory object"; + case ComputeErrorCode.InvalidImageFormatDescriptor: return "Invalid image format descriptor"; + case ComputeErrorCode.InvalidImageSize: return "Invalid image size"; + case ComputeErrorCode.InvalidSampler: return "Invalid sampler"; + case ComputeErrorCode.InvalidBinary: return "Invalid binary"; + case ComputeErrorCode.InvalidBuildOptions: return "Invalid build options"; + case ComputeErrorCode.InvalidProgram: return "Invalid program"; + case ComputeErrorCode.InvalidProgramExecutable: return "Invalid program executable"; + case ComputeErrorCode.InvalidKernelName: return "Invalid kernel name"; + case ComputeErrorCode.InvalidKernelDefinition: return "Invalid kernel definition"; + case ComputeErrorCode.InvalidKernel: return "Invalid kernel"; + case ComputeErrorCode.InvalidArgumentIndex: return "Invalid argument index"; + case ComputeErrorCode.InvalidArgumentValue: return "Invalid argument value"; + case ComputeErrorCode.InvalidArgumentSize: return "Invalid argument size"; + case ComputeErrorCode.InvalidKernelArguments: return "Invalid kernel arguments"; + case ComputeErrorCode.InvalidWorkDimension: return "Invalid work dimension"; + case ComputeErrorCode.InvalidWorkGroupSize: return "Invalid work group size"; + case ComputeErrorCode.InvalidWorkItemSize: return "Invalid work item size"; + case ComputeErrorCode.InvalidGlobalOffset: return "Invalid global offset"; + case ComputeErrorCode.InvalidEventWaitList: return "Invalid event wait list"; + case ComputeErrorCode.InvalidEvent: return "Invalid event"; + case ComputeErrorCode.InvalidOperation: return "Invalid operation"; + case ComputeErrorCode.InvalidGLObject: return "Invalid OpenGL object"; + case ComputeErrorCode.InvalidBufferSize: return "Invalid buffer size"; + case ComputeErrorCode.InvalidMipLevel: return "Invalid mip-map level"; + default: return "Unknown"; + } +} + +private void OpenCLSucceed(int ret, + string call, + string file, + int line) +{ + if (ret != (int) ComputeErrorCode.Success) + { + panic(-1, "{0}:{1}: OpenCL call\n {2}\nfailed with error code {3} ({4})\n", + file, line, call, ret, OpenCLErrorString(ret)); + } +} + +private void SetPreferredPlatform(ref OpenCLConfig cfg, string s) { + cfg.PreferredPlatform = s; +} + +private void SetPreferredDevice(ref OpenCLConfig cfg, string s) +{ + int x = 0; + int i = 0; + if (s[0] == '#') { + i = 1; + while (i < s.Length && char.IsDigit(s[i])) { + x = x * 10 + (int) (s[i])-'0'; + i++; + } + // Skip trailing spaces. + while (i < s.Length && char.IsWhiteSpace(s[i])) { + i++; + } + } + cfg.PreferredDevice = s.Substring(i); + cfg.PreferredDeviceNum = x; +} + +private string OpenCLPlatformInfo(CLPlatformHandle platform, + ComputePlatformInfo param) { + IntPtr req_bytes; + IntPtr _null = new IntPtr(); + OPENCL_SUCCEED(CL10.GetPlatformInfo(platform, param, _null, _null, out req_bytes)); + + byte[] info = new byte[(int) req_bytes]; + unsafe + { + fixed (byte* ptr = &info[0]) + { + OPENCL_SUCCEED(CL10.GetPlatformInfo(platform, param, req_bytes, new IntPtr(ptr), out _null)); + } + } + + return System.Text.Encoding.Default.GetString(info); +} + +private string OpenCLDeviceInfo(CLDeviceHandle device, + ComputeDeviceInfo param) { + IntPtr req_bytes; + IntPtr _null = new IntPtr(); + OPENCL_SUCCEED(CL10.GetDeviceInfo(device, param, _null, _null, out req_bytes)); + + byte[] info = new byte[(int) req_bytes]; + unsafe + { + fixed (byte* ptr = &info[0]) + { + OPENCL_SUCCEED(CL10.GetDeviceInfo(device, param, req_bytes, new IntPtr(ptr), out _null)); + } + } + return System.Text.Encoding.Default.GetString(info); + +} + +private void OpenCLAllDeviceOptions(out OpenCLDeviceOption[] devices_out, + out int num_devices_out) +{ + int num_devices = 0, num_devices_added = 0; + + CLPlatformHandle[] all_platforms; + int[] platform_num_devices; + + int num_platforms; + + // Find the number of platforms. + OPENCL_SUCCEED(CL10.GetPlatformIDs(0, null, out num_platforms)); + + // Make room for them. + all_platforms = new CLPlatformHandle[num_platforms]; + platform_num_devices = new int[num_platforms]; + + int tmp; + // Fetch all the platforms. + OPENCL_SUCCEED(CL10.GetPlatformIDs(num_platforms, all_platforms, out tmp)); + + // Count the number of devices for each platform, as well as the + // total number of devices. + for (int i = 0; i < num_platforms; i++) + { + if (CL10.GetDeviceIDs(all_platforms[i], ComputeDeviceTypes.All, + 0, null, out platform_num_devices[i]) == ComputeErrorCode.Success) + { + num_devices += platform_num_devices[i]; + } + else + { + platform_num_devices[i] = 0; + } + } + + // Make room for all the device options. + OpenCLDeviceOption[] devices = new OpenCLDeviceOption[num_devices]; + + // Loop through the platforms, getting information about their devices. + for (int i = 0; i < num_platforms; i++) { + CLPlatformHandle platform = all_platforms[i]; + int num_platform_devices = platform_num_devices[i]; + + if (num_platform_devices == 0) { + continue; + } + + string platform_name = OpenCLPlatformInfo(platform, ComputePlatformInfo.Name); + CLDeviceHandle[] platform_devices = new CLDeviceHandle[num_platform_devices]; + + // Fetch all the devices. + OPENCL_SUCCEED(CL10.GetDeviceIDs(platform, ComputeDeviceTypes.All, + num_platform_devices, platform_devices, out tmp)); + + IntPtr tmpptr; + // Loop through the devices, adding them to the devices array. + unsafe + { + for (int j = 0; j < num_platform_devices; j++) { + string device_name = OpenCLDeviceInfo(platform_devices[j], ComputeDeviceInfo.Name); + devices[num_devices_added].Platform = platform; + devices[num_devices_added].Device = platform_devices[j]; + fixed (void* ptr = &devices[num_devices_added].DeviceType) + { + OPENCL_SUCCEED(CL10.GetDeviceInfo(platform_devices[j], + ComputeDeviceInfo.Type, + new IntPtr(sizeof(ComputeDeviceTypes)), + new IntPtr(ptr), + out tmpptr)); + } + // We don't want the structs to share memory, so copy the platform name. + // Each device name is already unique. + devices[num_devices_added].PlatformName = platform_name; + devices[num_devices_added].DeviceName = device_name; + num_devices_added++; + } + } + } + + devices_out = devices; + num_devices_out = num_devices; +} + +private bool IsBlacklisted(string platform_name, string device_name) +{ + return (platform_name.Contains("Apple") && + device_name.Contains("Intel(R) Core(TM)")); +} + +private OpenCLDeviceOption GetPreferredDevice(OpenCLConfig cfg) { + OpenCLDeviceOption[] devices; + int num_devices; + + OpenCLAllDeviceOptions(out devices, out num_devices); + + int num_device_matches = 0; + + for (int i = 0; i < num_devices; i++) + { + OpenCLDeviceOption device = devices[i]; + if (!IsBlacklisted(device.PlatformName, device.DeviceName) && + device.PlatformName.Contains(cfg.PreferredPlatform) && + device.DeviceName.Contains(cfg.PreferredDevice) && + num_device_matches++ == cfg.PreferredDeviceNum) + { + return device; + } + } + + panic(1, "Could not find acceptable OpenCL device.\n"); + // this is never reached + throw new Exception(); + +} + +private void DescribeDeviceOption(OpenCLDeviceOption device) { + Console.Error.WriteLine("Using platform: {0}", device.PlatformName); + Console.Error.WriteLine("Using device: {0}", device.DeviceName); +} + +private ComputeProgramBuildStatus BuildOpenCLProgram(ref CLProgramHandle program, CLDeviceHandle device, string options) { + ComputeErrorCode ret_val = CL10.BuildProgram(program, 1, new []{device}, options, null, IntPtr.Zero); + + // Avoid termination due to CL_BUILD_PROGRAM_FAILURE + if (ret_val != ComputeErrorCode.Success && ret_val != ComputeErrorCode.BuildProgramFailure) { + Debug.Assert((int) ret_val == 0); + } + + ComputeProgramBuildStatus build_status; + unsafe + { + IntPtr _null = new IntPtr(); + ret_val = CL10.GetProgramBuildInfo(program, + device, + ComputeProgramBuildInfo.Status, + new IntPtr(sizeof(int)), + new IntPtr(&build_status), + out _null); + } + Debug.Assert(ret_val == 0); + + if (build_status != ComputeProgramBuildStatus.Success) { + char[] build_log; + IntPtr ret_val_size; + unsafe + { + ret_val = CL10.GetProgramBuildInfo(program, + device, + ComputeProgramBuildInfo.BuildLog, + IntPtr.Zero, + IntPtr.Zero, + out ret_val_size); + } + Debug.Assert(ret_val == 0); + + build_log = new char[((int)ret_val_size)+1]; + unsafe + { + IntPtr _null = new IntPtr(); + fixed (char* ptr = &build_log[0]) + { + CL10.GetProgramBuildInfo(program, + device, + ComputeProgramBuildInfo.BuildLog, + ret_val_size, + new IntPtr(ptr), + out _null); + } + } + Debug.Assert(ret_val == 0); + + // The spec technically does not say whether the build log is zero-terminated, so let's be careful. + build_log[(int)ret_val_size] = '\0'; + Console.Error.Write("Build log:\n{0}\n", new string(build_log)); + } + + return build_status; +} + + +// We take as input several strings representing the program, because +// C does not guarantee that the compiler supports particularly large +// literals. Notably, Visual C has a limit of 2048 characters. The +// array must be NULL-terminated. +private CLProgramHandle SetupOpenCL(ref FutharkContext ctx, + string[] srcs, + bool required_types) { + + ComputeErrorCode error; + CLPlatformHandle platform; + CLDeviceHandle device; + int MaxGroupSize; + + ctx.OpenCL.LockstepWidth = 0; + + OpenCLDeviceOption device_option = GetPreferredDevice(ctx.OpenCL.Cfg); + + if (ctx.Debugging) { + DescribeDeviceOption(device_option); + } + + device = device = device_option.Device; + platform = platform = device_option.Platform; + + if (required_types){ + int supported; + unsafe + { + IntPtr throwaway0 = new IntPtr(); + OPENCL_SUCCEED(CL10.GetDeviceInfo(device, + ComputeDeviceInfo.PreferredVectorWidthDouble, + new IntPtr(sizeof(IntPtr)), + new IntPtr(&supported), + out throwaway0)); + } + if (supported == 0) { + panic(1, + "Program uses double-precision floats, but this is not supported on chosen device: {0}\n", + device_option.DeviceName); + } + } + + unsafe + { + IntPtr throwaway1 = new IntPtr(); + OPENCL_SUCCEED(CL10.GetDeviceInfo(device, + ComputeDeviceInfo.MaxWorkGroupSize, + new IntPtr(sizeof(IntPtr)), + new IntPtr(&MaxGroupSize), + out throwaway1)); + } + + int MaxTileSize = (int) Math.Sqrt(MaxGroupSize); + + if (MaxGroupSize < ctx.OpenCL.Cfg.DefaultGroupSize) { + Console.Error.WriteLine("Note: Device limits default group size to {0} (down from {1}).\n", + MaxGroupSize, ctx.OpenCL.Cfg.DefaultGroupSize); + ctx.OpenCL.Cfg.DefaultGroupSize = MaxGroupSize; + } + + if (MaxTileSize < ctx.OpenCL.Cfg.DefaultTileSize) { + Console.Error.WriteLine("Note: Device limits default tile size to {0} (down from {1}).\n", + MaxTileSize, ctx.OpenCL.Cfg.DefaultTileSize); + ctx.OpenCL.Cfg.DefaultTileSize = MaxTileSize; + } + + ctx.OpenCL.MaxGroupSize = MaxGroupSize; + ctx.OpenCL.MaxTileSize = MaxTileSize; // No limit. + ctx.OpenCL.MaxThreshold = ctx.OpenCL.MaxNumGroups; // No limit. + + // Now we go through all the sizes, clamp them to the valid range, + // or set them to the default. + for (int i = 0; i < ctx.OpenCL.Cfg.NumSizes; i++) { + string size_class = ctx.OpenCL.Cfg.SizeClasses[i]; + int size_value = ctx.OpenCL.Cfg.SizeValues[i]; + string size_name = ctx.OpenCL.Cfg.SizeNames[i]; + int max_value, default_value; + max_value = default_value = 0; + if (size_class == "group_size") { + max_value = MaxGroupSize; + default_value = ctx.OpenCL.Cfg.DefaultGroupSize; + } else if (size_class == "num_groups") { + max_value = MaxGroupSize; // Futhark assumes this constraint. + default_value = ctx.OpenCL.Cfg.DefaultNumGroups; + } else if (size_class == "tile_size"){ + max_value = (int) Math.Sqrt(MaxGroupSize); + default_value = ctx.OpenCL.Cfg.DefaultTileSize; + } else if (size_class == "threshold") { + max_value = 0; // No limit. + default_value = ctx.OpenCL.Cfg.DefaultThreshold; + } else { + panic(1, "Unknown size class for size '{0}': {1}\n", size_name, size_class); + } + if (size_value == 0) { + ctx.OpenCL.Cfg.SizeValues[i] = default_value; + } else if (max_value > 0 && size_value > max_value) { + Console.Error.WriteLine("Note: Device limits {0} to {1} (down from {2})", + size_name, max_value, size_value); + ctx.OpenCL.Cfg.SizeValues[i] = default_value; + } + } + + IntPtr[] properties = new []{ + new IntPtr((int) ComputeContextInfo.Platform), + platform.Value, + IntPtr.Zero + }; + // Note that nVidia's OpenCL requires the platform property + IntPtr _null; + ctx.OpenCL.Context = CL10.CreateContext(properties, 1, new []{device}, null, ctx.NULL, out error); + Debug.Assert(error == 0); + + ctx.OpenCL.Queue = CL10.CreateCommandQueue(ctx.OpenCL.Context, device, 0, out error); + Debug.Assert(error == 0); + + // Make sure this function is defined. + PostOpenCLSetup(ref ctx, ref device_option); + + if (ctx.Debugging) { + Console.Error.WriteLine("Lockstep width: {0}\n", (int)ctx.OpenCL.LockstepWidth); + Console.Error.WriteLine("Default group size: {0}\n", (int)ctx.OpenCL.Cfg.DefaultGroupSize); + Console.Error.WriteLine("Default number of groups: {0}\n", (int)ctx.OpenCL.Cfg.DefaultNumGroups); + } + + string fut_opencl_src; + + // Maybe we have to read OpenCL source from somewhere else (used for debugging). + if (ctx.OpenCL.Cfg.LoadProgramFrom != null) { + fut_opencl_src = File.ReadAllText(ctx.OpenCL.Cfg.LoadProgramFrom); + } else { + // Build the OpenCL program. First we have to concatenate all the fragments. + fut_opencl_src = string.Join("\n", srcs); + } + + CLProgramHandle prog; + error = 0; + string[] src_ptr = new[]{fut_opencl_src}; + IntPtr[] src_size = new []{IntPtr.Zero}; + + if (ctx.OpenCL.Cfg.DumpProgramTo != null) { + File.WriteAllText(ctx.OpenCL.Cfg.DumpProgramTo, fut_opencl_src); + } + + unsafe + { + prog = CL10.CreateProgramWithSource(ctx.OpenCL.Context, 1, src_ptr, src_size, out error); + } + Debug.Assert(error == 0); + + int compile_opts_size = 1024; + + string compile_opts = String.Format("-DFUT_BLOCK_DIM={0} -DLOCKSTEP_WIDTH={1} ", + ctx.OpenCL.Cfg.TransposeBlockDim, + ctx.OpenCL.LockstepWidth); + + for (int i = 0; i < ctx.OpenCL.Cfg.NumSizes; i++) { + compile_opts += String.Format("-D{0}={1} ", + ctx.OpenCL.Cfg.SizeNames[i], + ctx.OpenCL.Cfg.SizeValues[i]); + } + + OPENCL_SUCCEED(BuildOpenCLProgram(ref prog, device, compile_opts)); + + return prog; +} + +private CLMemoryHandle EmptyMemHandle(CLContextHandle context) +{ + ComputeErrorCode tmp; + var cl_mem = CL10.CreateBuffer(context, ComputeMemoryFlags.ReadWrite, + IntPtr.Zero, IntPtr.Zero, + out tmp); + return cl_mem; + +} + +private void FutharkConfigPrintSizes() +{ + int n = FutharkGetNumSizes(); + for (int i = 0; i < n; i++) + { + if (FutharkGetSizeEntry(i) == EntryPoint) + { + Console.WriteLine("{0} ({1})", FutharkGetSizeName(i), + FutharkGetSizeClass(i)); + } + } + Environment.Exit(0); +} + +private void FutharkConfigSetSize(ref FutharkContextConfig config, string optarg) +{ + var name_and_value = optarg.Split('='); + if (name_and_value.Length != 2) + { + panic(1, "Invalid argument for size option: {0}", optarg); + } + + var name = name_and_value[0]; + var value = Convert.ToInt32(name_and_value[1]); + if (!FutharkContextConfigSetSize(ref config, name, value)) + { + panic(1, "Unknown size: {0}", name); + } +} diff --git a/rts/csharp/panic.cs b/rts/csharp/panic.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/panic.cs @@ -0,0 +1,24 @@ +private void panic(int exitcode, string str, params Object[] args) +{ + var prog_name = Environment.GetCommandLineArgs()[0]; + Console.Error.WriteLine(String.Format("{0}:", prog_name)); + Console.Error.WriteLine(String.Format(str, args)); + Environment.Exit(exitcode); +} + +private void FutharkAssert(bool assertion) +{ + if (!assertion) + { + Environment.Exit(1); + } +} + +private void FutharkAssert(bool assertion, string errorMsg) +{ + if (!assertion) + { + Console.Error.WriteLine(errorMsg); + Environment.Exit(1); + } +} diff --git a/rts/csharp/reader.cs b/rts/csharp/reader.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/reader.cs @@ -0,0 +1,857 @@ +private Stream s; +private BinaryReader b; + +// Note that the lookahead buffer does not interact well with +// binary reading. We are careful to not let this become a +// problem. +private Stack LookaheadBuffer = new Stack(); + +private void ResetLookahead(){ + LookaheadBuffer.Clear(); +} + +private void ValueReader(Stream s) +{ + this.s = s; +} + +private void ValueReader() +{ + this.s = Console.OpenStandardInput(); + this.b = new BinaryReader(s); +} + +private char? GetChar() +{ + char c; + if (LookaheadBuffer.Count == 0) + { + c = (char) this.b.ReadByte(); + } + else + { + c = LookaheadBuffer.Pop(); + } + + return c; +} + +private char[] GetChars(int n) +{ + return Enumerable.Range(0, n).Select(_ => GetChar().Value).ToArray(); +} + +private void UngetChar(char c) +{ + LookaheadBuffer.Push(c); +} + +private char PeekChar() +{ + var c = GetChar(); + UngetChar(c.Value); + return c.Value; +} + +private void SkipSpaces() +{ + var c = GetChar(); + while (c.HasValue){ + if (char.IsWhiteSpace(c.Value)) + { + c = GetChar(); + } + else if (c == '-') + { + if (PeekChar() == '-') + { + while (c.Value != '\n') + { + c = GetChar(); + } + } + else + { + break; + } + } + else + { + break; + } + } + + if (c.HasValue) + { + UngetChar(c.Value); + } +} + +private bool ParseSpecificChar(char c) +{ + var got = GetChar(); + if (got.Value != c) + { + UngetChar(got.Value); + throw new ValueError(); + } + return true; +} + +private bool ParseSpecificString(string str) +{ + var read = new List(); + foreach (var c in str.ToCharArray()) + { + try + { + ParseSpecificChar(c); + read.Add(c); + } + catch(ValueError) + { + read.Reverse(); + foreach (var cc in read) + { + UngetChar(cc); + } + throw; + } + } + + return true; +} + +private string Optional(Func p) +{ + string res = null; + try + { + res = p(); + } + catch (Exception) + { + } + + return res; +} + +private bool Optional(Func p, char c) +{ + try + { + return p(c); + } + catch (Exception) + { + } + + return false; +} + +private bool OptionalSpecificString(string s) +{ + var c = PeekChar(); + if (c == s[0]) + { + return ParseSpecificString(s); + } + return false; +} + + +private List sepBy(Func p, Func sep) +{ + var elems = new List(); + var x = Optional(p); + if (!string.IsNullOrWhiteSpace(x)) + { + elems.Add(x); + while (!string.IsNullOrWhiteSpace(Optional(sep))) + { + var y = Optional(p); + elems.Add(y); + } + } + return elems; +} + +private string ParseHexInt() +{ + var s = ""; + var c = GetChar(); + while (c.HasValue) + { + if (Uri.IsHexDigit(c.Value)) + { + s += c.Value; + c = GetChar(); + } + else if (c == '_') + { + c = GetChar(); + } + else + { + UngetChar(c.Value); + break; + } + } + + return Convert.ToString(Convert.ToUInt32(s, 16)); +} + +private string ParseInt() +{ + var s = ""; + var c = GetChar(); + if (c.Value == '0' && "xX".Contains(PeekChar())) + { + GetChar(); + s += ParseHexInt(); + } + else + { + while (c.HasValue) + { + if (char.IsDigit(c.Value)) + { + s += c.Value; + c = GetChar(); + }else if (c == '_') + { + c = GetChar(); + } + else + { + UngetChar(c.Value); + break; + } + } + + } + + if (s.Length == 0) + { + throw new Exception("ValueError"); + } + + return s; +} + +private string ParseIntSigned() +{ + var c = GetChar(); + if (c.Value == '-' && char.IsDigit(PeekChar())) + { + return c + ParseInt(); + } + else + { + if (c.Value != '+') + { + UngetChar(c.Value); + } + + return ParseInt(); + } +} + +private string ReadStrComma() +{ + SkipSpaces(); + ParseSpecificChar(','); + return ","; +} + +private int ReadStrInt(string s) +{ + SkipSpaces(); + var x = Convert.ToInt32(ParseIntSigned()); + OptionalSpecificString(s); + return x; +} + +private ulong ReadStrUInt64(string s) +{ + SkipSpaces(); + var x = Convert.ToUInt64(ParseInt()); + OptionalSpecificString(s); + return x; +} + +private long ReadStrInt64(string s) +{ + SkipSpaces(); + var x = Convert.ToInt64(ParseIntSigned()); + OptionalSpecificString(s); + return x; +} + +private uint ReadStrUInt(string s) +{ + SkipSpaces(); + var x = Convert.ToUInt32(ParseInt()); + OptionalSpecificString(s); + return x; +} + +private int ReadStrI8(){return ReadStrInt("i8");} +private int ReadStrI16(){return ReadStrInt("i16");} +private int ReadStrI32(){return ReadStrInt("i32");} +private long ReadStrI64(){return ReadStrInt64("i64");} +private uint ReadStrU8(){return ReadStrUInt("u8");} +private uint ReadStrU16(){return ReadStrUInt("u16");} +private uint ReadStrU32(){return ReadStrUInt("u32");} +private ulong ReadStrU64(){return ReadStrUInt64("u64");} +private sbyte ReadBinI8(){return (sbyte) b.ReadByte();} +private short ReadBinI16(){return b.ReadInt16();} +private int ReadBinI32(){return b.ReadInt32();} +private long ReadBinI64(){return b.ReadInt64();} +private byte ReadBinU8(){return (byte) b.ReadByte();} +private ushort ReadBinU16(){return b.ReadUInt16();} +private uint ReadBinU32(){return b.ReadUInt32();} +private ulong ReadBinU64(){return b.ReadUInt64();} +private float ReadBinF32(){return b.ReadSingle();} +private double ReadBinF64(){return b.ReadDouble();} +private bool ReadBinBool(){return b.ReadBoolean();} + +private char ReadChar() +{ + SkipSpaces(); + ParseSpecificChar('\''); + var c = GetChar(); + ParseSpecificChar('\''); + return c.Value; +} + +private double ReadStrHexFloat(char sign) +{ + var int_part = ParseHexInt(); + ParseSpecificChar('.'); + var frac_part = ParseHexInt(); + ParseSpecificChar('p'); + var exponent = ParseHexInt(); + + var int_val = Convert.ToInt32(int_part, 16); + var frac_val = Convert.ToSingle(Convert.ToInt32(frac_part, 16)) / Math.Pow(16, frac_part.Length); + var exp_val = Convert.ToInt32(exponent); + + var total_val = (int_val + frac_val) * Math.Pow(2, exp_val); + if (sign == '-') + { + total_val = -1 * total_val; + } + + return Convert.ToDouble(total_val); +} + +private double ReadStrDecimal() +{ + SkipSpaces(); + var c = GetChar(); + char sign; + if (c.Value == '-') + { + sign = '-'; + } + else + { + UngetChar(c.Value); + sign = '+'; + } + + // Check for hexadecimal float + c = GetChar(); + if (c.Value == '0' && "xX".Contains(PeekChar())) + { + GetChar(); + return ReadStrHexFloat(sign); + } + else + { + UngetChar(c.Value); + } + + var bef = Optional(this.ParseInt); + var aft = ""; + if (string.IsNullOrEmpty(bef)) + { + bef = "0"; + ParseSpecificChar('.'); + aft = ParseInt(); + }else if (Optional(ParseSpecificChar, '.')) + { + aft = ParseInt(); + } + else + { + aft = "0"; + } + + var expt = ""; + if (Optional(ParseSpecificChar, 'E') || + Optional(ParseSpecificChar, 'e')) + { + expt = ParseIntSigned(); + } + else + { + expt = "0"; + } + + return Convert.ToDouble(sign + bef + "." + aft + "E" + expt); +} + +private float ReadStrF32() +{ + try + { + ParseSpecificString("f32.nan"); + return Single.NaN; + } + catch (ValueError) + { + try + { + ParseSpecificString("-f32.inf"); + return Single.NegativeInfinity; + } + catch (ValueError) + { + try + { + ParseSpecificString("f32.inf"); + return Single.PositiveInfinity; + } + catch (ValueError) + { + var x = ReadStrDecimal(); + OptionalSpecificString("f32"); + return Convert.ToSingle(x); + } + } + } +} + +private double ReadStrF64() +{ + try + { + ParseSpecificString("f64.nan"); + return Double.NaN; + } + catch (ValueError) + { + try + { + ParseSpecificString("-f64.inf"); + return Double.NegativeInfinity; + } + catch (ValueError) + { + try + { + ParseSpecificString("f64.inf"); + return Double.PositiveInfinity; + } + catch (ValueError) + { + var x = ReadStrDecimal(); + OptionalSpecificString("f64"); + return x; + } + } + } +} +private bool ReadStrBool() +{ + SkipSpaces(); + if (PeekChar() == 't') + { + ParseSpecificString("true"); + return true; + } + + if (PeekChar() == 'f') + { + ParseSpecificString("false"); + return false; + } + + throw new ValueError(); +} + +private (T[], int[]) ReadStrArrayElems(int rank, Func ReadStrScalar) +{ + bool first = true; + bool[] knows_dimsize = new bool[rank]; + int cur_dim = rank-1; + int[] elems_read_in_dim = new int[rank]; + int[] shape = new int[rank]; + + int capacity = 100; + T[] data = new T[capacity]; + int write_ptr = 0; + + while (true) { + SkipSpaces(); + + char c = (char) GetChar(); + if (c == ']') { + if (knows_dimsize[cur_dim]) { + if (shape[cur_dim] != elems_read_in_dim[cur_dim]) { + throw new Exception("Irregular array"); + } + } else { + knows_dimsize[cur_dim] = true; + shape[cur_dim] = elems_read_in_dim[cur_dim]; + } + if (cur_dim == 0) { + break; + } else { + cur_dim--; + elems_read_in_dim[cur_dim]++; + } + } else if (c == ',') { + SkipSpaces(); + c = (char) GetChar(); + if (c == '[') { + if (cur_dim == rank - 1) { + throw new Exception("Array has too many dimensions"); + } + first = true; + cur_dim++; + elems_read_in_dim[cur_dim] = 0; + } else if (cur_dim == rank - 1) { + UngetChar(c); + + data[write_ptr++] = ReadStrScalar(); + if (write_ptr == capacity) { + capacity *= 2; + Array.Resize(ref data, capacity); + } + elems_read_in_dim[cur_dim]++; + } else { + throw new Exception("Unexpected comma when reading array"); + } + } else if (first) { + if (c == '[') { + if (cur_dim == rank - 1) { + throw new Exception("Array has too many dimensions"); + } + cur_dim++; + elems_read_in_dim[cur_dim] = 0; + } else { + UngetChar(c); + data[write_ptr++] = ReadStrScalar(); + if (write_ptr == capacity) { + capacity *= 2; + Array.Resize(ref data, capacity); + } + elems_read_in_dim[cur_dim]++; + first = false; + } + } else { + throw new Exception("Unexpected character in array"); + } + } + Array.Resize(ref data, write_ptr); + return (data, shape); +} + +private (T[], int[]) ReadStrArrayEmpty(int rank, string typeName, Func ReadStrScalar) +{ + ParseSpecificString("empty"); + ParseSpecificChar('('); + for (int i = 0; i < rank-1; i++) { + ParseSpecificString("[]"); + } + ParseSpecificString(typeName); + ParseSpecificChar(')'); + + return (new T[1], new int[rank]); +} + +private (T[], int[]) ReadStrArray(int rank, string typeName, Func ReadStrScalar) +{ + long read_dims = 0; + + while (true) { + SkipSpaces(); + var c = GetChar(); + if (c=='[') { + read_dims++; + } else { + if (c != null) { + UngetChar((char)c); + } + break; + } + } + + if (read_dims == 0) { + return ReadStrArrayEmpty(rank, typeName, ReadStrScalar); + } + + if (read_dims != rank) { + throw new Exception("Wrong number of dimensions"); + } + + return ReadStrArrayElems(rank, ReadStrScalar); +} + +private Dictionary primtypes = new Dictionary +{ + {" i8", "i8"}, + {" i16", "i16"}, + {" i32", "i32"}, + {" i64", "i64"}, + {" u8", "u8"}, + {" u16", "u16"}, + {" u32", "u32"}, + {" u64", "u64"}, + {" f32", "f32"}, + {" f64", "f64"}, + {"bool", "bool"} +}; + +private int BINARY_FORMAT_VERSION = 2; + + +private void read_le_2byte(ref short dest) +{ + dest = b.ReadInt16(); +} + +private void read_le_4byte(ref int dest) +{ + dest = b.ReadInt32(); +} + +private void read_le_8byte(ref long dest) +{ + dest = b.ReadInt64(); +} + +private bool ReadIsBinary() + { + SkipSpaces(); + var c = GetChar(); + if (c == 'b') + { + byte bin_version = new byte(); + try + { + bin_version = (byte) b.ReadByte(); + } + catch + { + Console.WriteLine("binary-input: could not read version"); + Environment.Exit(1); + } + + if (bin_version != BINARY_FORMAT_VERSION) + { + Console.WriteLine(( + "binary-input: File uses version {0}, but I only understand version {1}.", bin_version, + BINARY_FORMAT_VERSION)); + Environment.Exit(1); + } + + return true; + } + UngetChar((char) c); + return false; + } + +private (T[], int[]) ReadArray(int rank, string typeName, Func ReadStrScalar) +{ + if (!ReadIsBinary()) + { + return ReadStrArray(rank, typeName, ReadStrScalar); + } + else + { + return ReadBinArray(rank, typeName, ReadStrScalar); + } +} +private T ReadScalar(string typeName, Func ReadStrScalar, Func ReadBinScalar) +{ + if (!ReadIsBinary()) + { + return ReadStrScalar(); + } + else + { + ReadBinEnsureScalar(typeName); + return ReadBinScalar(); + } +} + +private void ReadBinEnsureScalar(string typeName) +{ + var bin_dims = b.ReadByte(); + if (bin_dims != 0) + { + Console.WriteLine("binary-input: Expected scalar (0 dimensions), but got array with {0} dimensions.", bin_dims); + Environment.Exit(1); + } + + var bin_type = ReadBinReadTypeString(); + if (bin_type != typeName) + { + Console.WriteLine("binary-input: Expected scalar of type {0} but got scalar of type {1}.", typeName, + bin_type); + Environment.Exit(1); + } +} + +private string ReadBinReadTypeString() +{ + var str_bytes = b.ReadBytes(4); + var str = System.Text.Encoding.UTF8.GetString(str_bytes, 0, 4); + return primtypes[str]; +} + +private (T[], int[]) ReadBinArray(int rank, string typeName, Func ReadStrScalar) +{ + var bin_dims = new int(); + var shape = new int[rank]; + try + { + bin_dims = b.ReadByte(); + } + catch + { + Console.WriteLine("binary-input: Couldn't get dims."); + Environment.Exit(1); + } + + if (bin_dims != rank) + { + Console.WriteLine("binary-input: Expected {0} dimensions, but got array with {1} dimensions", rank, + bin_dims); + Environment.Exit(1); + + } + + var bin_primtype = ReadBinReadTypeString(); + if (typeName != bin_primtype) + { + Console.WriteLine("binary-input: Expected {0}D-array with element type '{1}', but got {2}D-array with element type '{3}'.", + rank, typeName, bin_dims, bin_primtype); + Environment.Exit(1); + } + + int elem_count = 1; + for (var i = 0; i < rank; i++) + { + long bin_shape = new long(); + try + { + read_le_8byte(ref bin_shape); + } + catch + { + Console.WriteLine("binary-input: Couldn't read size for dimension {0} of array.", i); + Environment.Exit(1); + } + + elem_count *= (int) bin_shape; + shape[i] = (int) bin_shape; + } + + var elem_size = Marshal.SizeOf(typeof(T)); + var num_bytes = elem_count * elem_size; + var tmp = new byte[num_bytes]; + var data = new T[elem_count]; + + var to_read = num_bytes; + var have_read = 0; + while (to_read > 0) + { + var bytes_read = b.Read(tmp, have_read, to_read); + to_read -= bytes_read; + have_read += bytes_read; + } + + if (!BitConverter.IsLittleEndian && elem_size != 1) + { + for (int i = 0; i < elem_count; i ++) + { + Array.Reverse(tmp, i * elem_size, elem_size); + } + } + Buffer.BlockCopy(tmp,0,data,0,num_bytes); + + /* we should have a proper error message here */ + return (data, shape); +} + + +private sbyte ReadI8() +{ + return (sbyte) ReadStrI8(); +} +private short ReadI16() +{ + return (short) ReadStrI16(); +} +private int ReadI32() +{ + return ReadStrI32(); +} +private long ReadI64() +{ + return ReadStrI64(); +} + +private byte ReadU8() +{ + return (byte) ReadStrU8(); +} +private ushort ReadU16() +{ + return (ushort) ReadStrU16(); +} +private uint ReadU32() +{ + return (uint) ReadStrU32(); +} +private ulong ReadU64() +{ + return (ulong) ReadStrU64(); +} +private bool ReadBool() +{ + return ReadStrBool(); +} +private float ReadF32() +{ + return ReadStrF32(); +} +private double ReadF64() +{ + return ReadStrF64(); +} + +private void WriteValue(bool x){Console.Write(x ? "true" : "false", x);} +private void WriteValue(sbyte x){Console.Write("{0}i8", x);} +private void WriteValue(short x){Console.Write("{0}i16", x);} +private void WriteValue(int x){Console.Write("{0}i32", x);} +private void WriteValue(long x){Console.Write("{0}i64", x);} +private void WriteValue(byte x){Console.Write("{0}u8", x);} +private void WriteValue(ushort x){Console.Write("{0}u16", x);} +private void WriteValue(uint x){Console.Write("{0}u32", x);} +private void WriteValue(ulong x){Console.Write("{0}u64", x);} +private void WriteValue(float x){if (Single.IsNaN(x)) + {Console.Write("f32.nan");} else if (Single.IsNegativeInfinity(x)) + {Console.Write("-f32.inf");} else if (Single.IsPositiveInfinity(x)) + {Console.Write("f32.inf");} else + {Console.Write("{0:0.000000}f32", x);}} +private void WriteValue(double x){if (Double.IsNaN(x)) + {Console.Write("f64.nan");} else if (Double.IsNegativeInfinity(x)) + {Console.Write("-f64.inf");} else if (Double.IsPositiveInfinity(x)) + {Console.Write("f64.inf");} else + {Console.Write("{0:0.000000}f64", x);}} diff --git a/rts/csharp/scalar.cs b/rts/csharp/scalar.cs new file mode 100644 --- /dev/null +++ b/rts/csharp/scalar.cs @@ -0,0 +1,312 @@ +// Scalar functions. +private static sbyte signed(byte x){ return (sbyte) x;} +private static short signed(ushort x){ return (short) x;} +private static int signed(uint x){ return (int) x;} +private static long signed(ulong x){ return (long) x;} + +private static byte unsigned(sbyte x){ return (byte) x;} +private static ushort unsigned(short x){ return (ushort) x;} +private static uint unsigned(int x){ return (uint) x;} +private static ulong unsigned(long x){ return (ulong) x;} + +private static sbyte add8(sbyte x, sbyte y){ return (sbyte) ((byte) x + (byte) y);} +private static short add16(short x, short y){ return (short) ((ushort) x + (ushort) y);} +private static int add32(int x, int y){ return (int) ((uint) x + (uint) y);} +private static long add64(long x, long y){ return (long) ((ulong) x + (ulong) y);} + +private static sbyte sub8(sbyte x, sbyte y){ return (sbyte) ((byte) x - (byte) y);} +private static short sub16(short x, short y){ return (short) ((ushort) x - (ushort) y);} +private static int sub32(int x, int y){ return (int) ((uint) x - (uint) y);} +private static long sub64(long x, long y){ return (long) ((ulong) x - (ulong) y);} + +private static sbyte mul8(sbyte x, sbyte y){ return (sbyte) ((byte) x * (byte) y);} +private static short mul16(short x, short y){ return (short) ((ushort) x * (ushort) y);} +private static int mul32(int x, int y){ return (int) ((uint) x * (uint) y);} +private static long mul64(long x, long y){ return (long) ((ulong) x * (ulong) y);} + +private static sbyte or8(sbyte x, sbyte y){ return (sbyte) (x | y); } +private static short or16(short x, short y){ return (short) (x | y); } +private static int or32(int x, int y){ return x | y; } +private static long or64(long x, long y){ return x | y;} + +private static sbyte xor8(sbyte x, sbyte y){ return (sbyte) (x ^ y); } +private static short xor16(short x, short y){ return (short) (x ^ y); } +private static int xor32(int x, int y){ return x ^ y; } +private static long xor64(long x, long y){ return x ^ y;} + +private static sbyte and8(sbyte x, sbyte y){ return (sbyte) (x & y); } +private static short and16(short x, short y){ return (short) (x & y); } +private static int and32(int x, int y){ return x & y; } +private static long and64(long x, long y){ return x & y;} + +private static sbyte shl8(sbyte x, sbyte y){ return (sbyte) (x << y); } +private static short shl16(short x, short y){ return (short) (x << y); } +private static int shl32(int x, int y){ return x << y; } +private static long shl64(long x, long y){ return x << Convert.ToInt32(y); } + +private static sbyte ashr8(sbyte x, sbyte y){ return (sbyte) (x >> y); } +private static short ashr16(short x, short y){ return (short) (x >> y); } +private static int ashr32(int x, int y){ return x >> y; } +private static long ashr64(long x, long y){ return x >> Convert.ToInt32(y); } + +private static sbyte sdiv8(sbyte x, sbyte y){ + var q = squot8(x,y); + var r = srem8(x,y); + return (sbyte) (q - (((r != (sbyte) 0) && ((r < (sbyte) 0) != (y < (sbyte) 0))) ? (sbyte) 1 : (sbyte) 0)); +} +private static short sdiv16(short x, short y){ + var q = squot16(x,y); + var r = srem16(x,y); + return (short) (q - (((r != (short) 0) && ((r < (short) 0) != (y < (short) 0))) ? (short) 1 : (short) 0)); +} +private static int sdiv32(int x, int y){ + var q = squot32(x,y); + var r = srem32(x,y); + return q - (((r != (int) 0) && ((r < (int) 0) != (y < (int) 0))) ? (int) 1 : (int) 0); +} +private static long sdiv64(long x, long y){ + var q = squot64(x,y); + var r = srem64(x,y); + return q - (((r != (long) 0) && ((r < (long) 0) != (y < (long) 0))) ? (long) 1 : (long) 0); +} + +private static sbyte smod8(sbyte x, sbyte y){ + var r = srem8(x,y); + return (sbyte) (r + ((r == (sbyte) 0 || (x > (sbyte) 0 && y > (sbyte) 0) || (x < (sbyte) 0 && y < (sbyte) 0)) ? (sbyte) 0 : y)); +} +private static short smod16(short x, short y){ + var r = srem16(x,y); + return (short) (r + ((r == (short) 0 || (x > (short) 0 && y > (short) 0) || (x < (short) 0 && y < (short) 0)) ? (short) 0 : y)); +} +private static int smod32(int x, int y){ + var r = srem32(x,y); + return (int) r + ((r == (int) 0 || (x > (int) 0 && y > (int) 0) || (x < (int) 0 && y < (int) 0)) ? (int) 0 : y); +} +private static long smod64(long x, long y){ + var r = srem64(x,y); + return (long) r + ((r == (long) 0 || (x > (long) 0 && y > (long) 0) || (x < (long) 0 && y < (long) 0)) ? (long) 0 : y); +} + +private static sbyte udiv8(sbyte x, sbyte y){ return signed((byte) (unsigned(x) / unsigned(y))); } +private static short udiv16(short x, short y){ return signed((ushort) (unsigned(x) / unsigned(y))); } +private static int udiv32(int x, int y){ return signed(unsigned(x) / unsigned(y)); } +private static long udiv64(long x, long y){ return signed(unsigned(x) / unsigned(y)); } + +private static sbyte umod8(sbyte x, sbyte y){ return signed((byte) (unsigned(x) % unsigned(y))); } +private static short umod16(short x, short y){ return signed((ushort) (unsigned(x) % unsigned(y))); } +private static int umod32(int x, int y){ return signed(unsigned(x) % unsigned(y)); } +private static long umod64(long x, long y){ return signed(unsigned(x) % unsigned(y)); } + +private static sbyte squot8(sbyte x, sbyte y){ return (sbyte) Math.Truncate(ToSingle(x) / ToSingle(y)); } +private static short squot16(short x, short y){ return (short) Math.Truncate(ToSingle(x) / ToSingle(y)); } +private static int squot32(int x, int y){ return (int) Math.Truncate(ToSingle(x) / ToSingle(y)); } +private static long squot64(long x, long y){ return (long) Math.Truncate(ToSingle(x) / ToSingle(y)); } + +// private static Maybe change srem, it calls np.fmod originally so i dont know +private static sbyte srem8(sbyte x, sbyte y){ return (sbyte) ((sbyte) x % (sbyte) y);} +private static short srem16(short x, short y){ return (short) ((short) x % (short) y);} +private static int srem32(int x, int y){ return (int) ((int) x % (int) y);} +private static long srem64(long x, long y){ return (long) ((long) x % (long) y);} + +private static sbyte smin8(sbyte x, sbyte y){ return Math.Min(x,y);} +private static short smin16(short x, short y){ return Math.Min(x,y);} +private static int smin32(int x, int y){ return Math.Min(x,y);} +private static long smin64(long x, long y){ return Math.Min(x,y);} + +private static sbyte smax8(sbyte x, sbyte y){ return Math.Max(x,y);} +private static short smax16(short x, short y){ return Math.Max(x,y);} +private static int smax32(int x, int y){ return Math.Max(x,y);} +private static long smax64(long x, long y){ return Math.Max(x,y);} + +private static sbyte umin8(sbyte x, sbyte y){ return signed(Math.Min(unsigned(x),unsigned(y)));} +private static short umin16(short x, short y){ return signed(Math.Min(unsigned(x),unsigned(y)));} +private static int umin32(int x, int y){ return signed(Math.Min(unsigned(x),unsigned(y)));} +private static long umin64(long x, long y){ return signed(Math.Min(unsigned(x),unsigned(y)));} + +private static sbyte umax8(sbyte x, sbyte y){ return signed(Math.Max(unsigned(x),unsigned(y)));} +private static short umax16(short x, short y){ return signed(Math.Max(unsigned(x),unsigned(y)));} +private static int umax32(int x, int y){ return signed(Math.Max(unsigned(x),unsigned(y)));} +private static long umax64(long x, long y){ return signed(Math.Max(unsigned(x),unsigned(y)));} + +private static float fmin32(float x, float y){ return Math.Min(x,y);} +private static double fmin64(double x, double y){ return Math.Min(x,y);} +private static float fmax32(float x, float y){ return Math.Max(x,y);} +private static double fmax64(double x, double y){ return Math.Max(x,y);} + +private static sbyte pow8(sbyte x, sbyte y){sbyte res = 1;for (var i = 0; i < y; i++){res *= x;}return res;} +private static short pow16(short x, short y){short res = 1;for (var i = 0; i < y; i++){res *= x;}return res;} +private static int pow32(int x, int y){int res = 1;for (var i = 0; i < y; i++){res *= x;}return res;} +private static long pow64(long x, long y){long res = 1;for (var i = 0; i < y; i++){res *= x;}return res;} + +private static float fpow32(float x, float y){ return Convert.ToSingle(Math.Pow(x,y));} +private static double fpow64(double x, double y){ return Convert.ToDouble(Math.Pow(x,y));} + +private static bool sle8(sbyte x, sbyte y){ return x <= y ;} +private static bool sle16(short x, short y){ return x <= y ;} +private static bool sle32(int x, int y){ return x <= y ;} +private static bool sle64(long x, long y){ return x <= y ;} + +private static bool slt8(sbyte x, sbyte y){ return x < y ;} +private static bool slt16(short x, short y){ return x < y ;} +private static bool slt32(int x, int y){ return x < y ;} +private static bool slt64(long x, long y){ return x < y ;} + +private static bool ule8(sbyte x, sbyte y){ return unsigned(x) <= unsigned(y) ;} +private static bool ule16(short x, short y){ return unsigned(x) <= unsigned(y) ;} +private static bool ule32(int x, int y){ return unsigned(x) <= unsigned(y) ;} +private static bool ule64(long x, long y){ return unsigned(x) <= unsigned(y) ;} + +private static bool ult8(sbyte x, sbyte y){ return unsigned(x) < unsigned(y) ;} +private static bool ult16(short x, short y){ return unsigned(x) < unsigned(y) ;} +private static bool ult32(int x, int y){ return unsigned(x) < unsigned(y) ;} +private static bool ult64(long x, long y){ return unsigned(x) < unsigned(y) ;} + +private static sbyte lshr8(sbyte x, sbyte y){ return (sbyte) ((uint) x >> (int) y);} +private static short lshr16(short x, short y){ return (short) ((ushort) x >> (int) y);} +private static int lshr32(int x, int y){ return (int) ((uint) (x) >> (int) y);} +private static long lshr64(long x, long y){ return (long) ((ulong) x >> (int) y);} + +private static sbyte sext_i8_i8(sbyte x){return (sbyte) (x);} +private static short sext_i8_i16(sbyte x){return (short) (x);} +private static int sext_i8_i32(sbyte x){return (int) (x);} +private static long sext_i8_i64(sbyte x){return (long) (x);} + +private static sbyte sext_i16_i8(short x){return (sbyte) (x);} +private static short sext_i16_i16(short x){return (short) (x);} +private static int sext_i16_i32(short x){return (int) (x);} +private static long sext_i16_i64(short x){return (long) (x);} + +private static sbyte sext_i32_i8(int x){return (sbyte) (x);} +private static short sext_i32_i16(int x){return (short) (x);} +private static int sext_i32_i32(int x){return (int) (x);} +private static long sext_i32_i64(int x){return (long) (x);} + +private static sbyte sext_i64_i8(long x){return (sbyte) (x);} +private static short sext_i64_i16(long x){return (short) (x);} +private static int sext_i64_i32(long x){return (int) (x);} +private static long sext_i64_i64(long x){return (long) (x);} + +private static sbyte btoi_bool_i8 (bool x){return (sbyte) (Convert.ToInt32(x));} +private static short btoi_bool_i16(bool x){return (short) (Convert.ToInt32(x));} +private static int btoi_bool_i32(bool x){return (int) (Convert.ToInt32(x));} +private static long btoi_bool_i64(bool x){return (long) (Convert.ToInt32(x));} + +private static bool itob_i8_bool (sbyte x){return x != 0;} +private static bool itob_i16_bool(short x){return x != 0;} +private static bool itob_i32_bool(int x) {return x != 0;} +private static bool itob_i64_bool(long x) {return x != 0;} + +private static sbyte zext_i8_i8(sbyte x) {return (sbyte) ((byte)(x));} +private static short zext_i8_i16(sbyte x) {return (short)((byte)(x));} +private static int zext_i8_i32(sbyte x) {return (int)((byte)(x));} +private static long zext_i8_i64(sbyte x) {return (long)((byte)(x));} + +private static sbyte zext_i16_i8(short x) {return (sbyte) ((ushort)(x));} +private static short zext_i16_i16(short x) {return (short)((ushort)(x));} +private static int zext_i16_i32(short x) {return (int)((ushort)(x));} +private static long zext_i16_i64(short x) {return (long)((ushort)(x));} + +private static sbyte zext_i32_i8(int x){return (sbyte) ((uint)(x));} +private static short zext_i32_i16(int x){return (short)((uint)(x));} +private static int zext_i32_i32(int x){return (int)((uint)(x));} +private static long zext_i32_i64(int x){return (long)((uint)(x));} + +private static sbyte zext_i64_i8(long x){return (sbyte) ((ulong)(x));} +private static short zext_i64_i16(long x){return (short)((ulong)(x));} +private static int zext_i64_i32(long x){return (int)((ulong)(x));} +private static long zext_i64_i64(long x){return (long)((ulong)(x));} + +private static sbyte ssignum(sbyte x){return (sbyte) Math.Sign(x);} +private static short ssignum(short x){return (short) Math.Sign(x);} +private static int ssignum(int x){return Math.Sign(x);} +private static long ssignum(long x){return (long) Math.Sign(x);} + +private static sbyte usignum(sbyte x){return ((byte) x > 0) ? (sbyte) 1 : (sbyte) 0;} +private static short usignum(short x){return ((ushort) x > 0) ? (short) 1 : (short) 0;} +private static int usignum(int x){return ((uint) x > 0) ? (int) 1 : (int) 0;} +private static long usignum(long x){return ((ulong) x > 0) ? (long) 1 : (long) 0;} + +private static float sitofp_i8_f32(sbyte x){return Convert.ToSingle(x);} +private static float sitofp_i16_f32(short x){return Convert.ToSingle(x);} +private static float sitofp_i32_f32(int x){return Convert.ToSingle(x);} +private static float sitofp_i64_f32(long x){return Convert.ToSingle(x);} + +private static double sitofp_i8_f64(sbyte x){return Convert.ToDouble(x);} +private static double sitofp_i16_f64(short x){return Convert.ToDouble(x);} +private static double sitofp_i32_f64(int x){return Convert.ToDouble(x);} +private static double sitofp_i64_f64(long x){return Convert.ToDouble(x);} + + +private static float uitofp_i8_f32(sbyte x){return Convert.ToSingle(unsigned(x));} +private static float uitofp_i16_f32(short x){return Convert.ToSingle(unsigned(x));} +private static float uitofp_i32_f32(int x){return Convert.ToSingle(unsigned(x));} +private static float uitofp_i64_f32(long x){return Convert.ToSingle(unsigned(x));} + +private static double uitofp_i8_f64(sbyte x){return Convert.ToDouble(unsigned(x));} +private static double uitofp_i16_f64(short x){return Convert.ToDouble(unsigned(x));} +private static double uitofp_i32_f64(int x){return Convert.ToDouble(unsigned(x));} +private static double uitofp_i64_f64(long x){return Convert.ToDouble(unsigned(x));} + +private static byte fptoui_f32_i8(float x){return (byte) (Math.Truncate(x));} +private static byte fptoui_f64_i8(double x){return (byte) (Math.Truncate(x));} +private static sbyte fptosi_f32_i8(float x){return (sbyte) (Math.Truncate(x));} +private static sbyte fptosi_f64_i8(double x){return (sbyte) (Math.Truncate(x));} + +private static ushort fptoui_f32_i16(float x){return (ushort) (Math.Truncate(x));} +private static ushort fptoui_f64_i16(double x){return (ushort) (Math.Truncate(x));} +private static short fptosi_f32_i16(float x){return (short) (Math.Truncate(x));} +private static short fptosi_f64_i16(double x){return (short) (Math.Truncate(x));} + +private static uint fptoui_f32_i32(float x){return (uint) (Math.Truncate(x));} +private static uint fptoui_f64_i32(double x){return (uint) (Math.Truncate(x));} +private static int fptosi_f32_i32(float x){return (int) (Math.Truncate(x));} +private static int fptosi_f64_i32(double x){return (int) (Math.Truncate(x));} + +private static ulong fptoui_f32_i64(float x){return (ulong) (Math.Truncate(x));} +private static ulong fptoui_f64_i64(double x){return (ulong) (Math.Truncate(x));} +private static long fptosi_f32_i64(float x){return (long) (Math.Truncate(x));} +private static long fptosi_f64_i64(double x){return (long) (Math.Truncate(x));} + +private static double fpconv_f32_f64(float x){return Convert.ToDouble(x);} +private static float fpconv_f64_f32(double x){return Convert.ToSingle(x);} + +private static double futhark_log64(double x){return Math.Log(x);} +private static double futhark_log2_64(double x){return Math.Log(x,2.0);} +private static double futhark_log10_64(double x){return Math.Log10(x);} +private static double futhark_sqrt64(double x){return Math.Sqrt(x);} +private static double futhark_exp64(double x){return Math.Exp(x);} +private static double futhark_cos64(double x){return Math.Cos(x);} +private static double futhark_sin64(double x){return Math.Sin(x);} +private static double futhark_tan64(double x){return Math.Tan(x);} +private static double futhark_acos64(double x){return Math.Acos(x);} +private static double futhark_asin64(double x){return Math.Asin(x);} +private static double futhark_atan64(double x){return Math.Atan(x);} +private static double futhark_atan2_64(double x, double y){return Math.Atan2(x, y);} +private static bool futhark_isnan64(double x){return double.IsNaN(x);} +private static bool futhark_isinf64(double x){return double.IsInfinity(x);} +private static long futhark_to_bits64(double x){return BitConverter.ToInt64(BitConverter.GetBytes(x),0);} +private static double futhark_from_bits64(long x){return BitConverter.ToDouble(BitConverter.GetBytes(x),0);} + +private static float futhark_log32(float x){return (float) Math.Log(x);} +private static float futhark_log2_32(float x){return (float) Math.Log(x,2.0);} +private static float futhark_log10_32(float x){return (float) Math.Log10(x);} +private static float futhark_sqrt32(float x){return (float) Math.Sqrt(x);} +private static float futhark_exp32(float x){return (float) Math.Exp(x);} +private static float futhark_cos32(float x){return (float) Math.Cos(x);} +private static float futhark_sin32(float x){return (float) Math.Sin(x);} +private static float futhark_tan32(float x){return (float) Math.Tan(x);} +private static float futhark_acos32(float x){return (float) Math.Acos(x);} +private static float futhark_asin32(float x){return (float) Math.Asin(x);} +private static float futhark_atan32(float x){return (float) Math.Atan(x);} +private static float futhark_atan2_32(float x, float y){return (float) Math.Atan2(x, y);} +private static bool futhark_isnan32(float x){return float.IsNaN(x);} +private static bool futhark_isinf32(float x){return float.IsInfinity(x);} +private static int futhark_to_bits32(float x){return BitConverter.ToInt32(BitConverter.GetBytes(x), 0);} +private static float futhark_from_bits32(int x){return BitConverter.ToSingle(BitConverter.GetBytes(x), 0);} + +private static float futhark_round32(float x){return (float) Math.Round(x);} +private static double futhark_round64(double x){return Math.Round(x);} + +private static bool llt (bool x, bool y){return (!x && y);} +private static bool lle (bool x, bool y){return (!x || y);} + diff --git a/rts/python/__init__.py b/rts/python/__init__.py new file mode 100644 --- /dev/null +++ b/rts/python/__init__.py diff --git a/rts/python/memory.py b/rts/python/memory.py new file mode 100644 --- /dev/null +++ b/rts/python/memory.py @@ -0,0 +1,38 @@ +# Helper functions dealing with memory blocks. + +import ctypes as ct + +def addressOffset(x, offset, bt): + return ct.cast(ct.addressof(x.contents)+int(offset), ct.POINTER(bt)) + +def allocateMem(size): + return ct.cast((ct.c_byte * max(0,size))(), ct.POINTER(ct.c_byte)) + +# Copy an array if its is not-None. This is important for treating +# Numpy arrays as flat memory, but has some overhead. +def normaliseArray(x): + if (x.base is x) or (x.base is None): + return x + else: + return x.copy() + +def unwrapArray(x): + return normaliseArray(x).ctypes.data_as(ct.POINTER(ct.c_byte)) + +def createArray(x, dim): + return np.ctypeslib.as_array(x, shape=dim) + +def indexArray(x, offset, bt, nptype): + return nptype(addressOffset(x, offset, bt)[0]) + +def writeScalarArray(x, offset, v): + ct.memmove(ct.addressof(x.contents)+int(offset), ct.addressof(v), ct.sizeof(v)) + +# An opaque Futhark value. +class opaque(object): + def __init__(self, desc, *payload): + self.data = payload + self.desc = desc + + def __repr__(self): + return "".format(self.desc) diff --git a/rts/python/opencl.py b/rts/python/opencl.py new file mode 100644 --- /dev/null +++ b/rts/python/opencl.py @@ -0,0 +1,180 @@ +# Stub code for OpenCL setup. + +import pyopencl as cl +import numpy as np +import sys + +if cl.version.VERSION < (2015,2): + raise Exception('Futhark requires at least PyOpenCL version 2015.2. Installed version is %s.' % + cl.version.VERSION_TEXT) + +def parse_preferred_device(s): + pref_num = 0 + if len(s) > 1 and s[0] == '#': + i = 1 + while i < len(s): + if not s[i].isdigit(): + break + else: + pref_num = pref_num * 10 + int(s[i]) + i += 1 + while i < len(s) and s[i].isspace(): + i += 1 + return (s[i:], pref_num) + else: + return (s, 0) + +def get_prefered_context(interactive=False, platform_pref=None, device_pref=None): + if device_pref != None: + (device_pref, device_num) = parse_preferred_device(device_pref) + else: + device_num = 0 + + if interactive: + return cl.create_some_context(interactive=True) + + def blacklisted(p, d): + return platform_pref == None and device_pref == None and \ + p.name == "Apple" and d.name.find("Intel(R) Core(TM)") >= 0 + def platform_ok(p): + return not platform_pref or p.name.find(platform_pref) >= 0 + def device_ok(d): + return not device_pref or d.name.find(device_pref) >= 0 + + device_matches = 0 + + for p in cl.get_platforms(): + if not platform_ok(p): + continue + for d in p.get_devices(): + if blacklisted(p,d) or not device_ok(d): + continue + if device_matches == device_num: + return cl.Context(devices=[d]) + else: + device_matches += 1 + raise Exception('No OpenCL platform and device matching constraints found.') + +def check_types(self, required_types): + if 'f64' in required_types: + if self.device.get_info(cl.device_info.PREFERRED_VECTOR_WIDTH_DOUBLE) == 0: + raise Exception('Program uses double-precision floats, but this is not supported on chosen device: %s' % self.device.name) + +def apply_size_heuristics(self, size_heuristics, sizes): + for (platform_name, device_type, size, value) in size_heuristics: + if sizes[size] == None \ + and self.platform.name.find(platform_name) >= 0 \ + and self.device.type == device_type: + if type(value) == str: + sizes[size] = self.device.get_info(getattr(cl.device_info,value)) + else: + sizes[size] = value + return sizes + +def initialise_opencl_object(self, + program_src='', + command_queue=None, + interactive=False, + platform_pref=None, + device_pref=None, + default_group_size=None, + default_num_groups=None, + default_tile_size=None, + default_threshold=None, + transpose_block_dim=16, + size_heuristics=[], + required_types=[], + all_sizes={}, + user_sizes={}): + if command_queue is None: + self.ctx = get_prefered_context(interactive, platform_pref, device_pref) + self.queue = cl.CommandQueue(self.ctx) + else: + self.ctx = command_queue.context + self.queue = command_queue + self.device = self.queue.device + self.platform = self.device.platform + self.pool = cl.tools.MemoryPool(cl.tools.ImmediateAllocator(self.queue)) + device_type = self.device.type + + check_types(self, required_types) + + max_group_size = int(self.device.max_work_group_size) + max_tile_size = int(np.sqrt(self.device.max_work_group_size)) + + self.max_group_size = max_group_size + self.max_tile_size = max_tile_size + self.max_threshold = 0 + self.max_num_groups = 0 + self.free_list = {} + + default_group_size_set = default_group_size != None + default_tile_size_set = default_tile_size != None + default_sizes = apply_size_heuristics(self, size_heuristics, + {'group_size': default_group_size, + 'tile_size': default_tile_size, + 'num_groups': default_num_groups, + 'lockstep_width': None, + 'threshold': default_threshold}) + default_group_size = default_sizes['group_size'] + default_num_groups = default_sizes['num_groups'] + default_threshold = default_sizes['threshold'] + default_tile_size = default_sizes['tile_size'] + lockstep_width = default_sizes['lockstep_width'] + + if default_group_size > max_group_size: + if default_group_size_set: + sys.stderr.write('Note: Device limits group size to {} (down from {})\n'. + format(max_tile_size, default_group_size)) + default_group_size = max_group_size + + if default_tile_size > max_tile_size: + if default_tile_size_set: + sys.stderr.write('Note: Device limits tile size to {} (down from {})\n'. + format(max_tile_size, default_tile_size)) + default_tile_size = max_tile_size + + for (k,v) in user_sizes.items(): + if k in all_sizes: + all_sizes[k]['value'] = v + else: + raise Exception('Unknown size: {}'.format(k)) + + self.sizes = {} + for (k,v) in all_sizes.items(): + if v['class'] == 'group_size': + max_value = max_group_size + default_value = default_group_size + elif v['class'] == 'num_groups': + max_value = max_group_size # Intentional! + default_value = default_num_groups + elif v['class'] == 'tile_size': + max_value = max_tile_size + default_value = default_tile_size + elif v['class'].startswith('threshold'): + max_value = None + default_value = default_threshold + else: + raise Exception('Unknown size class for size \'{}\': {}'.format(k, v['class'])) + if v['value'] == None: + self.sizes[k] = default_value + elif max_value != None and v['value'] > max_value: + sys.stderr.write('Note: Device limits {} to {} (down from {}\n'. + format(k, max_value, v['value'])) + self.sizes[k] = max_value + else: + self.sizes[k] = v['value'] + + if (len(program_src) >= 0): + return cl.Program(self.ctx, program_src).build( + ["-DFUT_BLOCK_DIM={}".format(transpose_block_dim), + "-DLOCKSTEP_WIDTH={}".format(lockstep_width)] + + ["-D{}={}".format(s,v) for (s,v) in self.sizes.items()]) + +def opencl_alloc(self, min_size, tag): + min_size = 1 if min_size == 0 else min_size + assert min_size > 0 + return self.pool.allocate(min_size) + +def opencl_free_all(self): + self.pool.free_held() diff --git a/rts/python/panic.py b/rts/python/panic.py new file mode 100644 --- /dev/null +++ b/rts/python/panic.py @@ -0,0 +1,4 @@ +def panic(exitcode, fmt, *args): + sys.stderr.write('%s: ' % sys.argv[0]) + sys.stderr.write(fmt % args) + sys.exit(exitcode) diff --git a/rts/python/scalar.py b/rts/python/scalar.py new file mode 100644 --- /dev/null +++ b/rts/python/scalar.py @@ -0,0 +1,366 @@ +# Scalar functions. + +import numpy as np +import struct + +def signed(x): + if type(x) == np.uint8: + return np.int8(x) + elif type(x) == np.uint16: + return np.int16(x) + elif type(x) == np.uint32: + return np.int32(x) + else: + return np.int64(x) + +def unsigned(x): + if type(x) == np.int8: + return np.uint8(x) + elif type(x) == np.int16: + return np.uint16(x) + elif type(x) == np.int32: + return np.uint32(x) + else: + return np.uint64(x) + +def shlN(x,y): + return x << y + +def ashrN(x,y): + return x >> y + +def sdivN(x,y): + return x // y + +def smodN(x,y): + return x % y + +def udivN(x,y): + return signed(unsigned(x) // unsigned(y)) + +def umodN(x,y): + return signed(unsigned(x) % unsigned(y)) + +def squotN(x,y): + return np.floor_divide(np.abs(x), np.abs(y)) * np.sign(x) * np.sign(y) + +def sremN(x,y): + return np.remainder(np.abs(x), np.abs(y)) * np.sign(x) + +def sminN(x,y): + return min(x,y) + +def smaxN(x,y): + return max(x,y) + +def uminN(x,y): + return signed(min(unsigned(x),unsigned(y))) + +def umaxN(x,y): + return signed(max(unsigned(x),unsigned(y))) + +def fminN(x,y): + return min(x,y) + +def fmaxN(x,y): + return max(x,y) + +def powN(x,y): + return x ** y + +def fpowN(x,y): + return x ** y + +def sleN(x,y): + return x <= y + +def sltN(x,y): + return x < y + +def uleN(x,y): + return unsigned(x) <= unsigned(y) + +def ultN(x,y): + return unsigned(x) < unsigned(y) + +def lshr8(x,y): + return np.int8(np.uint8(x) >> np.uint8(y)) + +def lshr16(x,y): + return np.int16(np.uint16(x) >> np.uint16(y)) + +def lshr32(x,y): + return np.int32(np.uint32(x) >> np.uint32(y)) + +def lshr64(x,y): + return np.int64(np.uint64(x) >> np.uint64(y)) + +def sext_T_i8(x): + return np.int8(x) + +def sext_T_i16(x): + return np.int16(x) + +def sext_T_i32(x): + return np.int32(x) + +def sext_T_i64(x): + return np.int64(x) + +def itob_T_bool(x): + return np.bool(x) + +def btoi_bool_i8(x): + return np.int8(x) + +def btoi_bool_i16(x): + return np.int8(x) + +def btoi_bool_i32(x): + return np.int8(x) + +def btoi_bool_i64(x): + return np.int8(x) + +def zext_i8_i8(x): + return np.int8(np.uint8(x)) + +def zext_i8_i16(x): + return np.int16(np.uint8(x)) + +def zext_i8_i32(x): + return np.int32(np.uint8(x)) + +def zext_i8_i64(x): + return np.int64(np.uint8(x)) + +def zext_i16_i8(x): + return np.int8(np.uint16(x)) + +def zext_i16_i16(x): + return np.int16(np.uint16(x)) + +def zext_i16_i32(x): + return np.int32(np.uint16(x)) + +def zext_i16_i64(x): + return np.int64(np.uint16(x)) + +def zext_i32_i8(x): + return np.int8(np.uint32(x)) + +def zext_i32_i16(x): + return np.int16(np.uint32(x)) + +def zext_i32_i32(x): + return np.int32(np.uint32(x)) + +def zext_i32_i64(x): + return np.int64(np.uint32(x)) + +def zext_i64_i8(x): + return np.int8(np.uint64(x)) + +def zext_i64_i16(x): + return np.int16(np.uint64(x)) + +def zext_i64_i32(x): + return np.int32(np.uint64(x)) + +def zext_i64_i64(x): + return np.int64(np.uint64(x)) + +shl8 = shl16 = shl32 = shl64 = shlN +ashr8 = ashr16 = ashr32 = ashr64 = ashrN +sdiv8 = sdiv16 = sdiv32 = sdiv64 = sdivN +smod8 = smod16 = smod32 = smod64 = smodN +udiv8 = udiv16 = udiv32 = udiv64 = udivN +umod8 = umod16 = umod32 = umod64 = umodN +squot8 = squot16 = squot32 = squot64 = squotN +srem8 = srem16 = srem32 = srem64 = sremN +smax8 = smax16 = smax32 = smax64 = smaxN +smin8 = smin16 = smin32 = smin64 = sminN +umax8 = umax16 = umax32 = umax64 = umaxN +umin8 = umin16 = umin32 = umin64 = uminN +pow8 = pow16 = pow32 = pow64 = powN +fpow32 = fpow64 = fpowN +fmax32 = fmax64 = fmaxN +fmin32 = fmin64 = fminN +sle8 = sle16 = sle32 = sle64 = sleN +slt8 = slt16 = slt32 = slt64 = sltN +ule8 = ule16 = ule32 = ule64 = uleN +ult8 = ult16 = ult32 = ult64 = ultN +sext_i8_i8 = sext_i16_i8 = sext_i32_i8 = sext_i64_i8 = sext_T_i8 +sext_i8_i16 = sext_i16_i16 = sext_i32_i16 = sext_i64_i16 = sext_T_i16 +sext_i8_i32 = sext_i16_i32 = sext_i32_i32 = sext_i64_i32 = sext_T_i32 +sext_i8_i64 = sext_i16_i64 = sext_i32_i64 = sext_i64_i64 = sext_T_i64 +itob_i8_bool = itob_i16_bool = itob_i32_bool = itob_i64_bool = itob_T_bool + +def ssignum(x): + return np.sign(x) + +def usignum(x): + if x < 0: + return ssignum(-x) + else: + return ssignum(x) + +def sitofp_T_f32(x): + return np.float32(x) +sitofp_i8_f32 = sitofp_i16_f32 = sitofp_i32_f32 = sitofp_i64_f32 = sitofp_T_f32 + +def sitofp_T_f64(x): + return np.float64(x) +sitofp_i8_f64 = sitofp_i16_f64 = sitofp_i32_f64 = sitofp_i64_f64 = sitofp_T_f64 + +def uitofp_T_f32(x): + return np.float32(unsigned(x)) +uitofp_i8_f32 = uitofp_i16_f32 = uitofp_i32_f32 = uitofp_i64_f32 = uitofp_T_f32 + +def uitofp_T_f64(x): + return np.float64(unsigned(x)) +uitofp_i8_f64 = uitofp_i16_f64 = uitofp_i32_f64 = uitofp_i64_f64 = uitofp_T_f64 + +def fptosi_T_i8(x): + return np.int8(np.trunc(x)) +fptosi_f32_i8 = fptosi_f64_i8 = fptosi_T_i8 + +def fptosi_T_i16(x): + return np.int16(np.trunc(x)) +fptosi_f32_i16 = fptosi_f64_i16 = fptosi_T_i16 + +def fptosi_T_i32(x): + return np.int32(np.trunc(x)) +fptosi_f32_i32 = fptosi_f64_i32 = fptosi_T_i32 + +def fptosi_T_i64(x): + return np.int64(np.trunc(x)) +fptosi_f32_i64 = fptosi_f64_i64 = fptosi_T_i64 + +def fptoui_T_i8(x): + return np.uint8(np.trunc(x)) +fptoui_f32_i8 = fptoui_f64_i8 = fptoui_T_i8 + +def fptoui_T_i16(x): + return np.uint16(np.trunc(x)) +fptoui_f32_i16 = fptoui_f64_i16 = fptoui_T_i16 + +def fptoui_T_i32(x): + return np.uint32(np.trunc(x)) +fptoui_f32_i32 = fptoui_f64_i32 = fptoui_T_i32 + +def fptoui_T_i64(x): + return np.uint64(np.trunc(x)) +fptoui_f32_i64 = fptoui_f64_i64 = fptoui_T_i64 + +def fpconv_f32_f64(x): + return np.float64(x) + +def fpconv_f64_f32(x): + return np.float32(x) + +def futhark_log64(x): + return np.float64(np.log(x)) + +def futhark_log2_64(x): + return np.float64(np.log2(x)) + +def futhark_log10_64(x): + return np.float64(np.log10(x)) + +def futhark_sqrt64(x): + return np.sqrt(x) + +def futhark_exp64(x): + return np.exp(x) + +def futhark_cos64(x): + return np.cos(x) + +def futhark_sin64(x): + return np.sin(x) + +def futhark_tan64(x): + return np.tan(x) + +def futhark_acos64(x): + return np.arccos(x) + +def futhark_asin64(x): + return np.arcsin(x) + +def futhark_atan64(x): + return np.arctan(x) + +def futhark_atan2_64(x, y): + return np.arctan2(x, y) + +def futhark_round64(x): + return np.round(x) + +def futhark_isnan64(x): + return np.isnan(x) + +def futhark_isinf64(x): + return np.isinf(x) + +def futhark_to_bits64(x): + s = struct.pack('>d', x) + return np.int64(struct.unpack('>q', s)[0]) + +def futhark_from_bits64(x): + s = struct.pack('>q', x) + return np.float64(struct.unpack('>d', s)[0]) + +def futhark_log32(x): + return np.float32(np.log(x)) + +def futhark_log2_32(x): + return np.float32(np.log2(x)) + +def futhark_log10_32(x): + return np.float32(np.log10(x)) + +def futhark_sqrt32(x): + return np.float32(np.sqrt(x)) + +def futhark_exp32(x): + return np.exp(x) + +def futhark_cos32(x): + return np.cos(x) + +def futhark_sin32(x): + return np.sin(x) + +def futhark_tan32(x): + return np.tan(x) + +def futhark_acos32(x): + return np.arccos(x) + +def futhark_asin32(x): + return np.arcsin(x) + +def futhark_atan32(x): + return np.arctan(x) + +def futhark_atan2_32(x, y): + return np.arctan2(x, y) + +def futhark_round32(x): + return np.round(x) + +def futhark_isnan32(x): + return np.isnan(x) + +def futhark_isinf32(x): + return np.isinf(x) + +def futhark_to_bits32(x): + s = struct.pack('>f', x) + return np.int32(struct.unpack('>l', s)[0]) + +def futhark_from_bits32(x): + s = struct.pack('>l', x) + return np.float32(struct.unpack('>f', s)[0]) diff --git a/rts/python/values.py b/rts/python/values.py new file mode 100644 --- /dev/null +++ b/rts/python/values.py @@ -0,0 +1,627 @@ +# Hacky parser/reader/writer for values written in Futhark syntax. +# Used for reading stdin when compiling standalone programs with the +# Python code generator. + +import numpy as np +import string +import struct +import sys + +class ReaderInput: + def __init__(self, f): + self.f = f + self.lookahead_buffer = [] + + def get_char(self): + if len(self.lookahead_buffer) == 0: + return self.f.read(1) + else: + c = self.lookahead_buffer[0] + self.lookahead_buffer = self.lookahead_buffer[1:] + return c + + def unget_char(self, c): + self.lookahead_buffer = [c] + self.lookahead_buffer + + def get_chars(self, n): + s = b'' + for _ in range(n): + s += self.get_char() + return s + + def peek_char(self): + c = self.get_char() + if c: + self.unget_char(c) + return c + +def skip_spaces(f): + c = f.get_char() + while c != None: + if c.isspace(): + c = f.get_char() + elif c == b'-': + # May be line comment. + if f.peek_char() == b'-': + # Yes, line comment. Skip to end of line. + while (c != b'\n' and c != None): + c = f.get_char() + else: + break + else: + break + if c: + f.unget_char(c) + +def parse_specific_char(f, expected): + got = f.get_char() + if got != expected: + f.unget_char(got) + raise ValueError + return True + +def parse_specific_string(f, s): + # This funky mess is intended, and is caused by the fact that if `type(b) == + # bytes` then `type(b[0]) == int`, but we need to match each element with a + # `bytes`, so therefore we make each character an array element + b = s.encode('utf8') + bs = [b[i:i+1] for i in range(len(b))] + read = [] + try: + for c in bs: + parse_specific_char(f, c) + read.append(c) + return True + except ValueError: + map(f.unget_char, read[::-1]) + raise + +def optional(p, *args): + try: + return p(*args) + except ValueError: + return None + +def optional_specific_string(f, s): + c = f.peek_char() + # This funky mess is intended, and is caused by the fact that if `type(b) == + # bytes` then `type(b[0]) == int`, but we need to match each element with a + # `bytes`, so therefore we make each character an array element + b = s.encode('utf8') + bs = [b[i:i+1] for i in range(len(b))] + if c == bs[0]: + return parse_specific_string(f, s) + else: + return False + +def sepBy(p, sep, *args): + elems = [] + x = optional(p, *args) + if x != None: + elems += [x] + while optional(sep, *args) != None: + x = p(*args) + elems += [x] + return elems + +# Assumes '0x' has already been read +def parse_hex_int(f): + s = b'' + c = f.get_char() + while c != None: + if c in string.hexdigits: + s += c + c = f.get_char() + elif c == '_': + c = f.get_char() # skip _ + else: + f.unget_char(c) + break + return str(int(s, 16)) + + +def parse_int(f): + s = b'' + c = f.get_char() + if c == b'0' and f.peek_char() in [b'x', b'X']: + c = f.get_char() # skip X + s += parse_hex_int(f) + else: + while c != None: + if c.isdigit(): + s += c + c = f.get_char() + elif c == '_': + c = f.get_char() # skip _ + else: + f.unget_char(c) + break + if len(s) == 0: + raise ValueError + return s + +def parse_int_signed(f): + s = b'' + c = f.get_char() + + if c == b'-' and f.peek_char().isdigit(): + s = c + parse_int(f) + else: + if c != b'+': + f.unget_char(c) + s = parse_int(f) + + return s + +def read_str_comma(f): + skip_spaces(f) + parse_specific_char(f, b',') + return b',' + +def read_str_int(f, s): + skip_spaces(f) + x = int(parse_int_signed(f)) + optional_specific_string(f, s) + return x + +def read_str_uint(f, s): + skip_spaces(f) + x = int(parse_int(f)) + optional_specific_string(f, s) + return x + +def read_str_i8(f): + return np.int8(read_str_int(f, 'i8')) +def read_str_i16(f): + return np.int16(read_str_int(f, 'i16')) +def read_str_i32(f): + return np.int32(read_str_int(f, 'i32')) +def read_str_i64(f): + return np.int64(read_str_int(f, 'i64')) + +def read_str_u8(f): + return np.uint8(read_str_int(f, 'u8')) +def read_str_u16(f): + return np.uint16(read_str_int(f, 'u16')) +def read_str_u32(f): + return np.uint32(read_str_int(f, 'u32')) +def read_str_u64(f): + return np.uint64(read_str_int(f, 'u64')) + +def read_char(f): + skip_spaces(f) + parse_specific_char(f, b'\'') + c = f.get_char() + parse_specific_char(f, b'\'') + return c + +def read_str_hex_float(f, sign): + int_part = parse_hex_int(f) + parse_specific_char(f, b'.') + frac_part = parse_hex_int(f) + parse_specific_char(f, b'p') + exponent = parse_int(f) + + int_val = int(int_part, 16) + frac_val = float(int(frac_part, 16)) / (16 ** len(frac_part)) + exp_val = int(exponent) + + total_val = (int_val + frac_val) * (2.0 ** exp_val) + if sign == b'-': + total_val = -1 * total_val + + return float(total_val) + + +def read_str_decimal(f): + skip_spaces(f) + c = f.get_char() + if (c == b'-'): + sign = b'-' + else: + f.unget_char(c) + sign = b'' + + # Check for hexadecimal float + c = f.get_char() + if (c == '0' and (f.peek_char() in ['x', 'X'])): + f.get_char() + return read_str_hex_float(f, sign) + else: + f.unget_char(c) + + bef = optional(parse_int, f) + if bef == None: + bef = b'0' + parse_specific_char(f, b'.') + aft = parse_int(f) + elif optional(parse_specific_char, f, b'.'): + aft = parse_int(f) + else: + aft = b'0' + if (optional(parse_specific_char, f, b'E') or + optional(parse_specific_char, f, b'e')): + expt = parse_int_signed(f) + else: + expt = b'0' + return float(sign + bef + b'.' + aft + b'E' + expt) + +def read_str_f32(f): + skip_spaces(f) + try: + parse_specific_string(f, 'f32.nan') + return np.float32(np.nan) + except ValueError: + try: + parse_specific_string(f, 'f32.inf') + return np.float32(np.inf) + except ValueError: + try: + parse_specific_string(f, '-f32.inf') + return np.float32(-np.inf) + except ValueError: + x = read_str_decimal(f) + optional_specific_string(f, 'f32') + return x + +def read_str_f64(f): + skip_spaces(f) + try: + parse_specific_string(f, 'f64.nan') + return np.float64(np.nan) + except ValueError: + try: + parse_specific_string(f, 'f64.inf') + return np.float64(np.inf) + except ValueError: + try: + parse_specific_string(f, '-f64.inf') + return np.float64(-np.inf) + except ValueError: + x = read_str_decimal(f) + optional_specific_string(f, 'f64') + return x + +def read_str_bool(f): + skip_spaces(f) + if f.peek_char() == b't': + parse_specific_string(f, 'true') + return True + elif f.peek_char() == b'f': + parse_specific_string(f, 'false') + return False + else: + raise ValueError + +def read_str_empty_array(f, type_name, rank): + parse_specific_string(f, 'empty') + parse_specific_char(f, b'(') + for i in range(rank): + parse_specific_string(f, '[]') + parse_specific_string(f, type_name) + parse_specific_char(f, b')') + + return None + +def read_str_array_elems(f, elem_reader, type_name, rank): + skip_spaces(f) + try: + parse_specific_char(f, b'[') + except ValueError: + return read_str_empty_array(f, type_name, rank) + else: + xs = sepBy(elem_reader, read_str_comma, f) + skip_spaces(f) + parse_specific_char(f, b']') + return xs + +def read_str_array_helper(f, elem_reader, type_name, rank): + def nested_row_reader(_): + return read_str_array_helper(f, elem_reader, type_name, rank-1) + if rank == 1: + row_reader = elem_reader + else: + row_reader = nested_row_reader + return read_str_array_elems(f, row_reader, type_name, rank-1) + +def expected_array_dims(l, rank): + if rank > 1: + n = len(l) + if n == 0: + elem = [] + else: + elem = l[0] + return [n] + expected_array_dims(elem, rank-1) + else: + return [len(l)] + +def verify_array_dims(l, dims): + if dims[0] != len(l): + raise ValueError + if len(dims) > 1: + for x in l: + verify_array_dims(x, dims[1:]) + +def read_str_array(f, elem_reader, type_name, rank, bt): + elems = read_str_array_helper(f, elem_reader, type_name, rank) + if elems == None: + # Empty array + return np.empty([0]*rank, dtype=bt) + else: + dims = expected_array_dims(elems, rank) + verify_array_dims(elems, dims) + return np.array(elems, dtype=bt) + +################################################################################ + +READ_BINARY_VERSION = 2 + +# struct format specified at +# https://docs.python.org/2/library/struct.html#format-characters + +def mk_bin_scalar_reader(t): + def bin_reader(f): + fmt = FUTHARK_PRIMTYPES[t]['bin_format'] + size = FUTHARK_PRIMTYPES[t]['size'] + return struct.unpack('<' + fmt, f.get_chars(size))[0] + return bin_reader + +read_bin_i8 = mk_bin_scalar_reader('i8') +read_bin_i16 = mk_bin_scalar_reader('i16') +read_bin_i32 = mk_bin_scalar_reader('i32') +read_bin_i64 = mk_bin_scalar_reader('i64') + +read_bin_u8 = mk_bin_scalar_reader('u8') +read_bin_u16 = mk_bin_scalar_reader('u16') +read_bin_u32 = mk_bin_scalar_reader('u32') +read_bin_u64 = mk_bin_scalar_reader('u64') + +read_bin_f32 = mk_bin_scalar_reader('f32') +read_bin_f64 = mk_bin_scalar_reader('f64') + +read_bin_bool = mk_bin_scalar_reader('bool') + +def read_is_binary(f): + skip_spaces(f) + c = f.get_char() + if c == b'b': + bin_version = read_bin_u8(f) + if bin_version != READ_BINARY_VERSION: + panic(1, "binary-input: File uses version %i, but I only understand version %i.\n", + bin_version, READ_BINARY_VERSION) + return True + else: + f.unget_char(c) + return False + +FUTHARK_PRIMTYPES = { + 'i8': {'binname' : b" i8", + 'size' : 1, + 'bin_reader': read_bin_i8, + 'str_reader': read_str_i8, + 'bin_format': 'b', + 'numpy_type': np.int8 }, + + 'i16': {'binname' : b" i16", + 'size' : 2, + 'bin_reader': read_bin_i16, + 'str_reader': read_str_i16, + 'bin_format': 'h', + 'numpy_type': np.int16 }, + + 'i32': {'binname' : b" i32", + 'size' : 4, + 'bin_reader': read_bin_i32, + 'str_reader': read_str_i32, + 'bin_format': 'i', + 'numpy_type': np.int32 }, + + 'i64': {'binname' : b" i64", + 'size' : 8, + 'bin_reader': read_bin_i64, + 'str_reader': read_str_i64, + 'bin_format': 'q', + 'numpy_type': np.int64}, + + 'u8': {'binname' : b" u8", + 'size' : 1, + 'bin_reader': read_bin_u8, + 'str_reader': read_str_u8, + 'bin_format': 'B', + 'numpy_type': np.uint8 }, + + 'u16': {'binname' : b" u16", + 'size' : 2, + 'bin_reader': read_bin_u16, + 'str_reader': read_str_u16, + 'bin_format': 'H', + 'numpy_type': np.uint16 }, + + 'u32': {'binname' : b" u32", + 'size' : 4, + 'bin_reader': read_bin_u32, + 'str_reader': read_str_u32, + 'bin_format': 'I', + 'numpy_type': np.uint32 }, + + 'u64': {'binname' : b" u64", + 'size' : 8, + 'bin_reader': read_bin_u64, + 'str_reader': read_str_u64, + 'bin_format': 'Q', + 'numpy_type': np.uint64 }, + + 'f32': {'binname' : b" f32", + 'size' : 4, + 'bin_reader': read_bin_f32, + 'str_reader': read_str_f32, + 'bin_format': 'f', + 'numpy_type': np.float32 }, + + 'f64': {'binname' : b" f64", + 'size' : 8, + 'bin_reader': read_bin_f64, + 'str_reader': read_str_f64, + 'bin_format': 'd', + 'numpy_type': np.float64 }, + + 'bool': {'binname' : b"bool", + 'size' : 1, + 'bin_reader': read_bin_bool, + 'str_reader': read_str_bool, + 'bin_format': 'b', + 'numpy_type': np.bool } +} + +def read_bin_read_type(f): + read_binname = f.get_chars(4) + + for (k,v) in FUTHARK_PRIMTYPES.items(): + if v['binname'] == read_binname: + return k + panic(1, "binary-input: Did not recognize the type '%s'.\n", read_binname) + +def numpy_type_to_type_name(t): + for (k,v) in FUTHARK_PRIMTYPES.items(): + if v['numpy_type'] == t: + return k + raise Exception('Unknown Numpy type: {}'.format(t)) + +def read_bin_ensure_scalar(f, expected_type): + dims = read_bin_i8(f) + + if dims != 0: + panic(1, "binary-input: Expected scalar (0 dimensions), but got array with %i dimensions.\n", dims) + + bin_type = read_bin_read_type(f) + if bin_type != expected_type: + panic(1, "binary-input: Expected scalar of type %s but got scalar of type %s.\n", + expected_type, bin_type) + +# ------------------------------------------------------------------------------ +# General interface for reading Primitive Futhark Values +# ------------------------------------------------------------------------------ + +def read_scalar(f, ty): + if read_is_binary(f): + read_bin_ensure_scalar(f, ty) + return FUTHARK_PRIMTYPES[ty]['bin_reader'](f) + return FUTHARK_PRIMTYPES[ty]['str_reader'](f) + +def read_array(f, expected_type, rank): + if not read_is_binary(f): + str_reader = FUTHARK_PRIMTYPES[expected_type]['str_reader'] + return read_str_array(f, str_reader, expected_type, rank, + FUTHARK_PRIMTYPES[expected_type]['numpy_type']) + + bin_rank = read_bin_u8(f) + + if bin_rank != rank: + panic(1, "binary-input: Expected %i dimensions, but got array with %i dimensions.\n", + rank, bin_rank) + + bin_type_enum = read_bin_read_type(f) + if expected_type != bin_type_enum: + panic(1, "binary-input: Expected %iD-array with element type '%s' but got %iD-array with element type '%s'.\n", + rank, expected_type, bin_rank, bin_type_enum) + + shape = [] + elem_count = 1 + for i in range(rank): + bin_size = read_bin_u64(f) + elem_count *= bin_size + shape.append(bin_size) + + bin_fmt = FUTHARK_PRIMTYPES[bin_type_enum]['bin_format'] + + # We first read the expected number of types into a bytestring, + # then use np.fromstring. This is because np.fromfile does not + # work on things that are insufficiently file-like, like a network + # stream. + bytes = f.get_chars(elem_count * FUTHARK_PRIMTYPES[expected_type]['size']) + arr = np.fromstring(bytes, dtype='<'+bin_fmt) + arr.shape = shape + + return arr + +if sys.version_info >= (3,0): + input_reader = ReaderInput(sys.stdin.buffer) +else: + input_reader = ReaderInput(sys.stdin) + +import re + +def read_value(type_desc, reader=input_reader): + """Read a value of the given type. The type is a string +representation of the Futhark type.""" + m = re.match(r'((?:\[\])*)([a-z0-9]+)$', type_desc) + if m: + dims = int(len(m.group(1))/2) + basetype = m.group(2) + assert basetype in FUTHARK_PRIMTYPES, "Unknown type: {}".format(type_desc) + if dims > 0: + return read_array(reader, basetype, dims) + else: + return read_scalar(reader, basetype) + return (dims, basetype) + +def write_value(v, out=sys.stdout): + if type(v) == np.uint8: + out.write("%uu8" % v) + elif type(v) == np.uint16: + out.write("%uu16" % v) + elif type(v) == np.uint32: + out.write("%uu32" % v) + elif type(v) == np.uint64: + out.write("%uu64" % v) + elif type(v) == np.int8: + out.write("%di8" % v) + elif type(v) == np.int16: + out.write("%di16" % v) + elif type(v) == np.int32: + out.write("%di32" % v) + elif type(v) == np.int64: + out.write("%di64" % v) + elif type(v) in [np.bool, np.bool_]: + if v: + out.write("true") + else: + out.write("false") + elif type(v) == np.float32: + if np.isnan(v): + out.write('f32.nan') + elif np.isinf(v): + if v >= 0: + out.write('f32.inf') + else: + out.write('-f32.inf') + else: + out.write("%.6ff32" % v) + elif type(v) == np.float64: + if np.isnan(v): + out.write('f64.nan') + elif np.isinf(v): + if v >= 0: + out.write('f64.inf') + else: + out.write('-f64.inf') + else: + out.write("%.6ff64" % v) + elif type(v) == np.ndarray: + if np.product(v.shape) == 0: + tname = numpy_type_to_type_name(v.dtype) + out.write('empty({}{})'.format(''.join(['[]' for _ in v.shape[1:]]), tname)) + else: + first = True + out.write('[') + for x in v: + if not first: out.write(', ') + first = False + write_value(x, out=out) + out.write(']') + else: + raise Exception("Cannot print value of type {}: {}".format(type(v), v)) + +################################################################################ +### end of values.py +################################################################################ diff --git a/src/Futhark/Actions.hs b/src/Futhark/Actions.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Actions.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Actions + ( printAction + , impCodeGenAction + , kernelImpCodeGenAction + , rangeAction + , metricsAction + ) +where + +import Control.Monad.IO.Class + +import Futhark.Pipeline +import Futhark.Analysis.Alias +import Futhark.Analysis.Range +import Futhark.Representation.AST +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.ExplicitMemory (ExplicitMemory) +import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGenSequential +import qualified Futhark.CodeGen.ImpGen.Kernels as ImpGenKernels +import Futhark.Representation.AST.Attributes.Ranges (CanBeRanged) +import Futhark.Analysis.Metrics +import Futhark.Util.Pretty (prettyText) + +printAction :: (Attributes lore, CanBeAliased (Op lore)) => Action lore +printAction = + Action { actionName = "Prettyprint" + , actionDescription = "Prettyprint the resulting internal representation on standard output." + , actionProcedure = liftIO . putStrLn . pretty . aliasAnalysis + } + +rangeAction :: (Attributes lore, CanBeRanged (Op lore)) => Action lore +rangeAction = + Action { actionName = "Range analysis" + , actionDescription = "Print the program with range annotations added." + , actionProcedure = liftIO . putStrLn . pretty . rangeAnalysis + } + +metricsAction :: OpMetrics (Op lore) => Action lore +metricsAction = + Action { actionName = "Compute metrics" + , actionDescription = "Print metrics on the final AST." + , actionProcedure = liftIO . putStr . show . progMetrics + } + +impCodeGenAction :: Action ExplicitMemory +impCodeGenAction = + Action { actionName = "Compile imperative" + , actionDescription = "Translate program into imperative IL and write it on standard output." + , actionProcedure = \prog -> + either (`internalError` prettyText prog) (liftIO . putStrLn . pretty) =<< + ImpGenSequential.compileProg prog + } + +kernelImpCodeGenAction :: Action ExplicitMemory +kernelImpCodeGenAction = + Action { actionName = "Compile imperative kernels" + , actionDescription = "Translate program into imperative IL with kernels and write it on standard output." + , actionProcedure = \prog -> + either (`internalError` prettyText prog) (liftIO . putStrLn . pretty) =<< + ImpGenKernels.compileProg prog + } diff --git a/src/Futhark/Analysis/AlgSimplify.hs b/src/Futhark/Analysis/AlgSimplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/AlgSimplify.hs @@ -0,0 +1,1441 @@ +{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, LambdaCase #-} +module Futhark.Analysis.AlgSimplify + ( ScalExp + , Error + , simplify + , mkSuffConds + , RangesRep + , ppRangesRep + , linFormScalE + , pickSymToElim + ) + where + +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import Data.List +import Control.Monad +import Control.Monad.Reader +import Control.Monad.State + +import Futhark.Representation.AST hiding (SDiv, SMod, SQuot, SRem, SSignum) +import Futhark.Analysis.ScalExp +import qualified Futhark.Representation.Primitive as P + +-- | Ranges are inclusive. +type RangesRep = M.Map VName (Int, Maybe ScalExp, Maybe ScalExp) + +-- | Prettyprint a 'RangesRep'. Do not rely on the format of this +-- string. Does not include the loop nesting depth information. +ppRangesRep :: RangesRep -> String +ppRangesRep = unlines . sort . map ppRange . M.toList + where ppRange (name, (_, lower, upper)) = + pretty name ++ ": " ++ + if lower == upper + then "== " ++ ppBound lower + else "[" ++ ppBound lower ++ ", " ++ + ppBound upper ++ "]" + ppBound Nothing = "?" + ppBound (Just se) = pretty se + +-- | environment recording the position and +-- a list of variable-to-range bindings. +data AlgSimplifyEnv = AlgSimplifyEnv { inSolveLTH0 :: Bool + , ranges :: RangesRep + , maxSteps :: Int + -- ^ The number of + -- simplifications to do before + -- bailing out, to avoid spending + -- too much time. + } + +data Error = StepsExceeded | Error String + +type AlgSimplifyM = StateT Int (ReaderT AlgSimplifyEnv (Either Error)) + +runAlgSimplifier :: Bool -> AlgSimplifyM a -> RangesRep -> Either Error a +runAlgSimplifier s x r = runReaderT (evalStateT x 0) env + where env = AlgSimplifyEnv { inSolveLTH0 = s + , ranges = r + , maxSteps = 100 -- heuristically chosen + } + +step :: AlgSimplifyM () +step = do modify (1+) + exceeded <- pure (>) <*> get <*> asks maxSteps + when exceeded stepsExceeded + +stepsExceeded :: AlgSimplifyM a +stepsExceeded = lift $ lift $ Left StepsExceeded + +badAlgSimplifyM :: String -> AlgSimplifyM a +badAlgSimplifyM = lift . lift . Left . Error + +-- | Binds an array name to the set of used-array vars +markInSolve :: AlgSimplifyEnv -> AlgSimplifyEnv +markInSolve env = + env { inSolveLTH0 = True } + +markGaussLTH0 :: AlgSimplifyM a -> AlgSimplifyM a +markGaussLTH0 = local markInSolve + +----------------------------------------------------------- +-- A Scalar Expression, i.e., ScalExp, is simplified to: -- +-- 1. if numeric: to a normalized sum-of-products form,-- +-- in which on the outermost level there are N-ary -- +-- Min/Max nodes, and the next two levels are a sum -- +-- of products. -- +-- 2. if boolean: to disjunctive normal form -- +-- -- +-- Corresponding Helper Representations are: -- +-- 1. NNumExp, i.e., NSum of NProd of ScalExp -- +-- 2. DNF -- +----------------------------------------------------------- + +data NNumExp = NSum [NNumExp] PrimType + | NProd [ScalExp] PrimType + deriving (Eq, Ord, Show) + +data BTerm = NRelExp RelOp0 NNumExp + | LogCt !Bool + | PosId VName + | NegId VName + deriving (Eq, Ord, Show) +type NAnd = [BTerm] +type DNF = [NAnd ] +--type NOr = [BTerm] +--type CNF = [NOr ] + +-- | Applies Simplification at Expression level: +simplify :: ScalExp -> RangesRep -> ScalExp +simplify e rangesrep = case runAlgSimplifier False (simplifyScal e) rangesrep of + Left (Error err) -> + error $ "Error during algebraic simplification of: " ++ pretty e ++ + "\n" ++ err + Left StepsExceeded -> e + Right e' -> e' + +-- | Given a symbol i and a scalar expression e, it decomposes +-- e = a*i + b and returns (a,b) if possible, otherwise Nothing. +linFormScalE :: VName -> ScalExp -> RangesRep -> Either Error (Maybe (ScalExp,ScalExp)) +linFormScalE i e = runAlgSimplifier False (linearFormScalExp i e) + +-- | Extracts sufficient conditions for a LTH0 relation to hold +mkSuffConds :: ScalExp -> RangesRep -> Either Error [[ScalExp]] +mkSuffConds e = runAlgSimplifier True (gaussElimRel e) + +{- +-- | Test if Simplification engine can handle this kind of expression +canSimplify :: Int -> Either Error ScalExp --[[ScalExp]] +canSimplify i = do + let (h,_,e2) = mkRelExp i + case e2 of + (RelExp LTH0 _) -> do + -- let e1' = trace (pretty e1) e1 + simplify e2 noLoc h +-- runAlgSimplifier False (gaussAllLTH0 False S.empty =<< toNumSofP =<< simplifyScal e) noLoc h + _ -> simplify e2 noLoc h-- badAlgSimplifyM "canSimplify: unimplemented!" +-} +------------------------------------------------------- +--- Assumes the relational expression is simplified -- +--- All uses of gaussiam elimination from simplify -- +--- must use simplifyNRel, which calls markGaussLTH0-- +--- to set the inSolveLTH0 environment var, so that -- +--- we do not enter an infinite recurssion! -- +--- Returns True or False or the input replation,i.e.-- +--- static-only simplification! -- +------------------------------------------------------- +simplifyNRel :: BTerm -> AlgSimplifyM BTerm +simplifyNRel inp_term@(NRelExp LTH0 inp_sofp) = do + term <- cheapSimplifyNRel inp_term + in_gauss <- asks inSolveLTH0 + let tp = typeOfNAlg inp_sofp + + if in_gauss || isTrivialNRel term || tp `notElem` map IntType allIntTypes + then return term + else do ednf <- markGaussLTH0 $ gaussAllLTH0 True S.empty inp_sofp + return $ case ednf of + Val (BoolValue c) -> LogCt c + _ -> term + where + isTrivialNRel (NRelExp _ (NProd [Val _] _)) = True + isTrivialNRel NRelExp{} = False + isTrivialNRel _ = False + + cheapSimplifyNRel :: BTerm -> AlgSimplifyM BTerm + cheapSimplifyNRel (NRelExp rel (NProd [Val v] _)) = + LogCt <$> valLTHEQ0 rel v + cheapSimplifyNRel e = return e +simplifyNRel inp_term = + return inp_term --- TODO: handle more cases. + +--gaussEliminateNRel :: BTerm -> AlgSimplifyM DNF +--gaussEliminateNRel _ = +-- badAlgSimplifyM "gaussElimNRel: unimplemented!" + + +gaussElimRel :: ScalExp -> AlgSimplifyM [[ScalExp]] -- ScalExp +gaussElimRel (RelExp LTH0 e) = do + e_sofp <- toNumSofP =<< simplifyScal e + e_scal<- simplifyScal =<< gaussAllLTH0 False S.empty e_sofp + e_dnf <- toDNF e_scal + mapM (mapM (\case + LogCt c -> return $ Val (BoolValue c) + PosId i -> return $ Id i $ scalExpType e + NegId i -> return $ Id i $ scalExpType e + NRelExp rel ee -> RelExp rel <$> fromNumSofP ee + )) e_dnf + +gaussElimRel _ = + badAlgSimplifyM "gaussElimRel: only LTH0 Int relations please!" + +--ppSyms :: S.Set VName -> String +--ppSyms ss = foldl (\s x -> s ++ " " ++ (baseString x)) "ElimSyms: " (S.toList ss) + + +primScalExpLTH0 :: ScalExp -> Bool +primScalExpLTH0 (Val (IntValue v)) = P.intToInt64 v < 0 +primScalExpLTH0 _ = False +----------------------------------------------------------- +----------------------------------------------------------- +----------------------------------------------------------- +---`gaussAllLTH' --- +--- `static_only':whether only a True/False answer is --- +--- required or actual a sufficient cond--- +--- `el_syms': the list of already eliminated --- +--- symbols, initialy empty --- +--- `sofp': the expression e in sum-of-product --- +--- form that is compared to 0, --- +--- i.e., e < 0. sofp assumed simplified--- +--- Result: is a ScalExp expression, which is --- +--- actually a predicate in DNF form, --- +--- that is a sufficient condition --- +--- for e < 0! --- +--- --- +--- gaussAllLTH0 is implementing the tracking of Min/Max--- +--- terms, and uses `gaussOneDefaultLTH0' --- +--- to implement gaussian-like elimination --- +--- to solve the a*i + b < 0 problem. --- +--- --- +--- IMPORTANT: before calling gaussAllLTH0 from outside --- +--- make sure to set insideSolveLTH0 env --- +--- member to True, via markGaussLTH0; --- +--- otherwise infinite recursion might happen--- +--- w.r.t. `simplifyScal' --- +----------------------------------------------------------- +----------------------------------------------------------- +----------------------------------------------------------- +type Prod = [ScalExp] +gaussAllLTH0 :: Bool -> S.Set VName -> NNumExp -> AlgSimplifyM ScalExp +gaussAllLTH0 static_only el_syms sofp = do + step + let tp = typeOfNAlg sofp + rangesrep <- asks ranges + e_scal <- fromNumSofP sofp + let mi = pickSymToElim rangesrep el_syms e_scal + + case mi of + Nothing -> return $ if primScalExpLTH0 e_scal + then Val (BoolValue True) + else RelExp LTH0 e_scal + Just i -> do + (jmm, fs0, terms) <- findMinMaxTerm i sofp + -- i.e., sofp == fs0 * jmm + terms, where + -- i appears in jmm and jmm = MinMax ... + + fs <- if not (null fs0) then return fs0 + else do one <- getPos1 tp; return [Val one] + + case jmm of + ------------------------------------------------------------------------ + -- A MinMax expression which uses to-be-eliminated symbol i was found -- + ------------------------------------------------------------------------ + Just (MaxMin _ [] ) -> + badAlgSimplifyM "gaussAllLTH0: Empty MinMax Node!" + Just (MaxMin ismin mmts) -> do + mone <- getNeg1 tp + + -- fs_lth0 => fs < 0 +-- fs_lth0 <- if null fs then return $ Val (BoolValue False) +-- else gaussAllLTH0 static_only el_syms (NProd fs tp) + fsm1 <- toNumSofP =<< simplifyScal =<< fromNumSofP + ( NSum [NProd fs tp, NProd [Val mone] tp] tp ) + fs_leq0 <- gaussAllLTH0 static_only el_syms fsm1 -- fs <= 0 + -- mfsm1 = - fs - 1, fs_geq0 => (fs >= 0), + -- i.e., fs_geq0 => (-fs - 1 < 0) + mfsm1 <- toNumSofP =<< simplifyScal =<< fromNumSofP + ( NSum [NProd (Val mone:fs) tp,NProd [Val mone] tp] tp ) + fs_geq0 <- gaussAllLTH0 static_only el_syms mfsm1 + + -- mm_terms are the simplified terms of the MinMax obtained + -- after inlining everything inside the MinMax, i.e., intially + -- terms + fs * MinMax ismin [t1,..,tn] -> [fs*t1+terms, ..., fs*tn+terms] + mm_terms<- mapM (\t -> toNumSofP =<< simplifyScal =<< fromNumSofP + (NSum ( NProd (t:fs) tp:terms ) tp) ) mmts + + -- for every (simplified) `term_i' of the inline MinMax exp, + -- get the sufficient conditions for `term_i < 0' + mms <- mapM (gaussAllLTH0 static_only el_syms) mm_terms + + if static_only + -------------------------------------------------------------------- + -- returns either Val (BoolValue True) or the original ScalExp relat -- + -------------------------------------------------------------------- + then if ( fs_geq0 == Val (BoolValue True) && ismin) || + ( fs_leq0 == Val (BoolValue True) && not ismin) + -- at least one term should be < 0! + then do let is_one_true = Val (BoolValue True ) `elem` mms + let are_all_false = all (== Val (BoolValue False)) mms + return $ if is_one_true then Val (BoolValue True) + else if are_all_false then Val (BoolValue False) + else RelExp LTH0 e_scal + -- otherwise all terms should be all true! + else do let are_all_true = all (== Val (BoolValue True )) mms + let is_one_false = Val (BoolValue False) `elem` mms + return $ if are_all_true then Val (BoolValue True ) + else if is_one_false then Val (BoolValue False) + else RelExp LTH0 e_scal + -------------------------------------------------------------------- + -- returns sufficient conditions for the ScalExp relation to hold -- + -------------------------------------------------------------------- + else do + let mm_fsgeq0 = foldl (if ismin then SLogOr else SLogAnd) + (Val (BoolValue (not ismin))) mms + let mm_fslth0 = foldl (if ismin then SLogAnd else SLogOr) + (Val (BoolValue ismin )) mms + -- the sufficient condition for the original expression, e.g., + -- terms + fs * Min [t1,..,tn] < 0 is + -- (fs >= 0 && (fs*t_1+terms < 0 || ... || fs*t_n+terms < 0) ) || + -- (fs < 0 && (fs*t_1+terms < 0 && ... && fs*t_n+terms < 0) ) + return $ SLogOr (SLogAnd fs_geq0 mm_fsgeq0) (SLogAnd fs_leq0 mm_fslth0) + + Just _ -> badAlgSimplifyM "gaussOneLTH0: (Just MinMax) invariant violated!" + ------------------------------------------------------------------------ + -- A MinMax expression which uses (to-be-elim) symbol i was NOT found -- + ------------------------------------------------------------------------ + Nothing-> do + m_sofp <- gaussOneDefaultLTH0 static_only i el_syms sofp + case m_sofp of + Nothing -> gaussAllLTH0 static_only (S.insert i el_syms) sofp + Just res_eofp -> return res_eofp + where + findMinMaxTerm :: VName -> NNumExp -> AlgSimplifyM (Maybe ScalExp, Prod, [NNumExp]) + findMinMaxTerm _ (NSum [] _) = return (Nothing, [], []) + findMinMaxTerm _ (NSum [NProd [MaxMin ismin e] _] _) = + return (Just (MaxMin ismin e), [], []) + findMinMaxTerm _ (NProd [MaxMin ismin e] _) = + return (Just (MaxMin ismin e), [], []) + + findMinMaxTerm ii t@NProd{} = do (mm, fs) <- findMinMaxFact ii t + return (mm, fs, []) + findMinMaxTerm ii (NSum (t:ts) tp)= do + rangesrep <- asks ranges + case M.lookup ii rangesrep of + Just (_, Just _, Just _) -> do + f <- findMinMaxFact ii t + case f of + (Just mm, fs) -> return (Just mm, fs, ts) + (Nothing, _ ) -> do (mm, fs', ts') <- findMinMaxTerm ii (NSum ts tp) + return (mm, fs', t:ts') + _ -> return (Nothing, [], t:ts) + + findMinMaxFact :: VName -> NNumExp -> AlgSimplifyM (Maybe ScalExp, Prod) + findMinMaxFact _ (NProd [] _ ) = return (Nothing, []) + findMinMaxFact ii (NProd (f:fs) tp) = + case f of + MaxMin ismin ts -> do + let id_set = mconcat $ map freeIn ts + if S.member ii id_set + then return (Just (MaxMin ismin ts), fs) + else do (mm, fs') <- findMinMaxFact ii (NProd fs tp) + return (mm, f:fs') + + _ -> do (mm, fs') <- findMinMaxFact ii (NProd fs tp) + return (mm, f:fs') + findMinMaxFact ii (NSum [f] _) = findMinMaxFact ii f + findMinMaxFact _ (NSum _ _) = + badAlgSimplifyM "findMinMaxFact: NSum argument illegal!" + + + +gaussOneDefaultLTH0 :: Bool -> VName -> S.Set VName -> NNumExp -> AlgSimplifyM (Maybe ScalExp) +gaussOneDefaultLTH0 static_only i elsyms e = do + aipb <- linearForm i e + case aipb of + Nothing -> return Nothing + Just (a, b) -> do + rangesrep <- asks ranges + one <- getPos1 (typeOfNAlg e) + ascal <- fromNumSofP a + mam1 <- toNumSofP =<< simplifyScal (SNeg (SPlus ascal (Val one))) + am1 <- toNumSofP =<< simplifyScal (SMinus ascal (Val one)) + ma <- toNumSofP =<< simplifyScal (SNeg ascal) + + b_scal<- fromNumSofP b + mbm1 <- toNumSofP =<< simplifyScal (SNeg (SPlus b_scal (Val one))) + + aleq0 <- simplifyScal =<< gaussAllLTH0 static_only elsyms am1 + ageq0 <- simplifyScal =<< gaussAllLTH0 static_only elsyms mam1 + + case M.lookup i rangesrep of + Nothing -> + badAlgSimplifyM "gaussOneDefaultLTH0: sym not in ranges!" + Just (_, Nothing, Nothing) -> + badAlgSimplifyM "gaussOneDefaultLTH0: both bounds are undefined!" + + -- only the lower-bound of i is known! + Just (_, Just lb, Nothing) -> do + alpblth0 <- gaussElimHalf static_only elsyms lb a b + and_half <- simplifyScal alpblth0 + case (and_half, aleq0) of + (Val (BoolValue True), Val (BoolValue True)) -> + return $ Just and_half + _ -> do malmbm1lth0 <- gaussElimHalf static_only elsyms lb ma mbm1 + other_half <- simplifyScal malmbm1lth0 + case (other_half, ageq0) of + (Val (BoolValue True), Val (BoolValue True)) -> + return $ Just (Val (BoolValue False)) + _ -> return Nothing + + Just (_, Nothing, Just ub) -> do + aupblth0 <- gaussElimHalf static_only elsyms ub a b + and_half <- simplifyScal aupblth0 + case (and_half, ageq0) of + (Val (BoolValue True), Val (BoolValue True)) -> + return $ Just and_half + _ -> do + maumbm1 <- gaussElimHalf static_only elsyms ub ma mbm1 + other_half <- simplifyScal maumbm1 + case (other_half, aleq0) of + (Val (BoolValue True), Val (BoolValue True)) -> + return $ Just (Val (BoolValue False)) + _ -> return Nothing + + Just (_, Just lb, Just ub) -> + if static_only + then if aleq0 == Val (BoolValue True) + then do alpblth0 <- simplifyScal =<< gaussElimHalf static_only elsyms lb a b + if alpblth0 == Val (BoolValue True) + then return $ Just (Val (BoolValue True)) + else do maubmbm1 <- simplifyScal =<< gaussElimHalf static_only elsyms ub ma mbm1 + return $ if maubmbm1 == Val (BoolValue True) + then Just (Val (BoolValue False)) + else Nothing + else if ageq0 == Val (BoolValue True) + then do aupblth0 <- simplifyScal =<< gaussElimHalf static_only elsyms ub a b + if aupblth0 == Val (BoolValue True) + then return $ Just (Val (BoolValue True)) + else do malbmbm1 <- simplifyScal =<< gaussElimHalf static_only elsyms lb ma mbm1 + return $ if malbmbm1 == Val (BoolValue True) + then Just (Val (BoolValue False)) + else Nothing + else return Nothing + else do + alpblth0 <- gaussElimHalf static_only elsyms lb a b + aupblth0 <- gaussElimHalf static_only elsyms ub a b + res <- simplifyScal $ SLogOr (SLogAnd aleq0 alpblth0) (SLogAnd ageq0 aupblth0) + return $ Just res + + where + gaussElimHalf :: Bool -> S.Set VName -> ScalExp -> NNumExp -> NNumExp -> AlgSimplifyM ScalExp + gaussElimHalf only_static elsyms0 q a b = do + a_scal <- fromNumSofP a + b_scal <- fromNumSofP b + e_num_scal <- simplifyScal (SPlus (STimes a_scal q) b_scal) + e_num <- toNumSofP e_num_scal + gaussAllLTH0 only_static elsyms0 e_num + +-- pos <- asks pos +-- badAlgSimplifyM "gaussOneDefaultLTH0: unimplemented!" + +---------------------------------------------------------- +--- Pick a Symbol to Eliminate & Bring To Linear Form --- +---------------------------------------------------------- + +pickSymToElim :: RangesRep -> S.Set VName -> ScalExp -> Maybe VName +pickSymToElim rangesrep elsyms0 e_scal = +-- ranges <- asks ranges +-- e_scal <- fromNumSofP e0 + let ids0= S.toList $ freeIn e_scal + ids1= filter (\s -> not (S.member s elsyms0)) ids0 + ids2= filter (\s -> case M.lookup s rangesrep of + Nothing -> False + Just _ -> True + ) ids1 + ids = sortBy (\n1 n2 -> let n1p = M.lookup n1 rangesrep + n2p = M.lookup n2 rangesrep + in case (n1p, n2p) of + (Just (p1,_,_), Just (p2,_,_)) -> compare (-p1) (-p2) + (_ , _ ) -> compare (1::Int) (1::Int) + ) ids2 + in case ids of + [] -> Nothing + v:_ -> Just v + + +linearFormScalExp :: VName -> ScalExp -> AlgSimplifyM (Maybe (ScalExp, ScalExp)) +linearFormScalExp sym scl_exp = do + sofp <- toNumSofP =<< simplifyScal scl_exp + ab <- linearForm sym sofp + case ab of + Just (a_sofp, b_sofp) -> do + a <- fromNumSofP a_sofp + b <- fromNumSofP b_sofp + a'<- simplifyScal a + b'<- simplifyScal b + return $ Just (a', b') + Nothing -> + return Nothing + +linearForm :: VName -> NNumExp -> AlgSimplifyM (Maybe (NNumExp, NNumExp)) +linearForm _ (NProd [] _) = + badAlgSimplifyM "linearForm: empty Prod!" +linearForm idd ee@NProd{} = linearForm idd (NSum [ee] (typeOfNAlg ee)) +linearForm _ (NSum [] _) = + badAlgSimplifyM "linearForm: empty Sum!" +linearForm idd (NSum terms tp) = do + terms_d_idd <- mapM (\t -> do t0 <- case t of + NProd (_:_) _ -> return t + _ -> badAlgSimplifyM "linearForm: ILLEGAL111!!!!" + t_scal <- fromNumSofP t0 + simplifyScal $ SDiv t_scal (Id idd (scalExpType t_scal)) + ) terms + let myiota = [1..(length terms)] + let ia_terms= filter (\(_,t)-> case t of + SDiv _ _ -> False + _ -> True + ) (zip myiota terms_d_idd) + let (a_inds, a_terms) = unzip ia_terms + + let (_, b_terms) = unzip $ filter (\(iii,_) -> iii `notElem` a_inds) + (zip myiota terms) + -- check that b_terms do not contain idd + b_succ <- foldM (\acc x -> + case x of + NProd fs _ -> do let fs_scal = case fs of + [] -> Val $ IntValue $ Int32Value 1 + f:fs' -> foldl STimes f fs' + let b_ids = freeIn fs_scal + return $ acc && not (idd `S.member` b_ids) + _ -> badAlgSimplifyM "linearForm: ILLEGAL222!!!!" + ) True b_terms + + case a_terms of + t:ts | b_succ -> do + let a_scal = foldl SPlus t ts + a_terms_sofp <- toNumSofP =<< simplifyScal a_scal + b_terms_sofp <- if null b_terms + then do zero <- getZero tp; return $ NProd [Val zero] tp + else return $ NSum b_terms tp + return $ Just (a_terms_sofp, b_terms_sofp) + _ -> return Nothing + +------------------------------------------------ +------------------------------------------------ +-- Main Helper Function: takes a scalar exp, -- +-- normalizes and simplifies it -- +------------------------------------------------ +------------------------------------------------ + +simplifyScal :: ScalExp -> AlgSimplifyM ScalExp + +simplifyScal (Val v) = return $ Val v +simplifyScal (Id x t) = return $ Id x t + +simplifyScal e@SNot{} = fromDNF =<< simplifyDNF =<< toDNF e +simplifyScal e@SLogAnd{} = fromDNF =<< simplifyDNF =<< toDNF e +simplifyScal e@SLogOr{} = fromDNF =<< simplifyDNF =<< toDNF e +simplifyScal e@RelExp{} = fromDNF =<< simplifyDNF =<< toDNF e + +-------------------------------------- +--- MaxMin related simplifications --- +-------------------------------------- +simplifyScal (MaxMin _ []) = + badAlgSimplifyM "Scalar MaxMin expression with empty arglist." +simplifyScal (MaxMin _ [e]) = simplifyScal e +simplifyScal (MaxMin ismin es) = do -- helperMinMax ismin es pos + -- pos <- asks pos + es0 <- mapM simplifyScal es + let evals = filter isValue es0 + es' = filter (not . isValue) es0 + mvv = case evals of + [] -> Nothing + v:vs -> let myop = if ismin then min else max + myval= getValue v + oneval = (foldl myop myval . map getValue) vs + in Just $ Val oneval + -- flatten the result and remove duplicates, + -- e.g., Max(Max(e1,e2), e1) -> Max(e1,e2,e3) + case (es', mvv) of + ([], Just vv) -> return vv + (_, Just vv) -> return $ MaxMin ismin $ remDups $ foldl flatop [] $ vv:es' + (_, Nothing) -> return $ MaxMin ismin $ remDups $ foldl flatop [] es' + -- ToDo: This can prove very expensive as compile time + -- but, IF e2-e1 <= 0 simplifies to True THEN + -- Min(e1,e2) = e2. Code example: + -- e1me2 <- if isMin + -- then simplifyScal $ AlgOp MINUS e1 e2 pos + -- else simplifyScal $ AlgOp MINUS e2 e1 pos + -- e1me2leq0 <- simplifyNRel $ NRelExp LEQ0 e1me2 pos + -- case e1me2leq0 of + -- NAnd [LogCt True _] _ -> simplifyAlgN e1 + -- NAnd [LogCt False _] _ -> simplifyAlgN e2 + where + isValue :: ScalExp -> Bool + isValue e = case e of + Val _ -> True + _ -> False + getValue :: ScalExp -> PrimValue + getValue se = case se of + Val v -> v + _ -> value (0::Int32) + flatop :: [ScalExp] -> ScalExp -> [ScalExp] + flatop a e@(MaxMin ismin' ses) = + a ++ if ismin == ismin' then ses else [e] + flatop a e = a++[e] + remDups :: [ScalExp] -> [ScalExp] + remDups l = S.toList (S.fromList l) + +--------------------------------------------------- +--- Plus/Minus related simplifications --- +--- BUG: the MinMax pattern matching should --- +--- be performed on the simplified subexps --- +--------------------------------------------------- +simplifyScal (SPlus e1o e2o) = do + e1' <- simplifyScal e1o + e2' <- simplifyScal e2o + if isMaxMin e1' || isMaxMin e2' + then helperPlusMinMax $ SPlus e1' e2' + else normalPlus e1' e2' + + where + normalPlus :: ScalExp -> ScalExp -> AlgSimplifyM ScalExp + normalPlus e1 e2 = do + e1' <- toNumSofP e1 + e2' <- toNumSofP e2 + let tp = scalExpType e1 + let terms = getTerms e1' ++ getTerms e2' + splittedTerms <- mapM splitTerm terms + let sortedTerms = sortBy (\(n1,_) (n2,_) -> compare n1 n2) splittedTerms + -- foldM discriminate: adds together identical terms, and + -- we reverse the list, to keep it in a ascending order. + merged <- reverse <$> foldM discriminate [] sortedTerms + let filtered = filter (\(_,v) -> not $ zeroIsh v ) merged + if null filtered + then do + zero <- getZero tp + fromNumSofP $ NProd [Val zero] tp + else do + terms' <- mapM joinTerm filtered + fromNumSofP $ NSum terms' tp + +simplifyScal (SMinus e1 e2) = do + let tp = scalExpType e1 + if e1 == e2 + then Val <$> getZero tp + else do min_1 <- getNeg1 $ scalExpType e1 + simplifyScal $ SPlus e1 $ STimes (Val min_1) e2 + +simplifyScal (SNeg e) = do + negOne <- getNeg1 $ scalExpType e + simplifyScal $ STimes (Val negOne) e + +simplifyScal (SAbs e) = return $ SAbs e + +simplifyScal (SSignum e) = return $ SSignum e + +--------------------------------------------------- +--- Times related simplifications --- +--- BUG: the MinMax pattern matching should --- +--- be performed on the simplified subexps --- +--------------------------------------------------- +simplifyScal (STimes e1o e2o) = do + e1'' <- simplifyScal e1o + e2'' <- simplifyScal e2o + if isMaxMin e1'' || isMaxMin e2'' + then helperMultMinMax $ STimes e1'' e2'' + else normalTimes e1'' e2'' + + where + normalTimes :: ScalExp -> ScalExp -> AlgSimplifyM ScalExp + normalTimes e1 e2 = do + let tp = scalExpType e1 + e1' <- toNumSofP e1 + e2' <- toNumSofP e2 + case (e1', e2') of + (NProd xs _, y@NProd{}) -> fromNumSofP =<< makeProds xs y + (NProd xs _, y) -> do + prods <- mapM (makeProds xs) $ getTerms y + fromNumSofP $ NSum (sort prods) tp + (x, NProd ys _) -> do + prods <- mapM (makeProds ys) $ getTerms x + fromNumSofP $ NSum (sort prods) tp + (NSum xs _, NSum ys _) -> do + xsMultChildren <- mapM getMultChildren xs + prods <- mapM (\x -> mapM (makeProds x) ys) xsMultChildren + fromNumSofP $ NSum (sort $ concat prods) tp + + makeProds :: [ScalExp] -> NNumExp -> AlgSimplifyM NNumExp + makeProds [] _ = + badAlgSimplifyM " In simplifyAlgN, makeProds: 1st arg is the empty list! " + makeProds _ (NProd [] _) = + badAlgSimplifyM " In simplifyAlgN, makeProds: 2nd arg is the empty list! " + makeProds _ NSum{} = + badAlgSimplifyM " In simplifyAlgN, makeProds: e1 * e2: e2 is a sum of sums! " + makeProds (Val v1:exs) (NProd (Val v2:ys) tp1) = do + v <- mulVals v1 v2 + return $ NProd (Val v : sort (ys++exs) ) tp1 + makeProds (Val v:exs) (NProd ys tp1) = + return $ NProd (Val v : sort (ys++exs) ) tp1 + makeProds exs (NProd (Val v : ys) tp1) = + return $ NProd (Val v : sort (ys++exs) ) tp1 + makeProds exs (NProd ys tp1) = + return $ NProd (sort (ys++exs)) tp1 + +--------------------------------------------------- +--------------------------------------------------- +--- DIvide related simplifications --- +--------------------------------------------------- +--------------------------------------------------- + +simplifyScal (SDiv e1o e2o) = do + e1' <- simplifyScal e1o + e2' <- simplifyScal e2o + + if isMaxMin e1' || isMaxMin e2' + then helperMultMinMax $ SDiv e1' e2' + else normalFloatDiv e1' e2' + + where + normalFloatDiv :: ScalExp -> ScalExp -> AlgSimplifyM ScalExp + normalFloatDiv e1 e2 + | e1 == e2 = do one <- getPos1 $ scalExpType e1 + return $ Val one +-- | e1 == (negateSimplified e2) = do mone<- getNeg1 $ scalExpType e1 +-- return $ Val mone + | otherwise = do + e1' <- toNumSofP e1 + e2' <- toNumSofP e2 + case e2' of + NProd fs tp -> do + e1Split <- mapM splitTerm (getTerms e1') + case e1Split of + [] -> Val <$> getZero tp + _ -> do (fs', e1Split') <- trySimplifyDivRec fs [] e1Split + if length fs' == length fs + then turnBackAndDiv e1' e2' -- insuccess + else do terms_e1' <- mapM joinTerm e1Split' + e1'' <- fromNumSofP $ NSum terms_e1' tp + case fs' of + [] -> return e1'' + _ -> do e2'' <- fromNumSofP $ NProd fs' tp + return $ SDiv e1'' e2'' + + _ -> turnBackAndDiv e1' e2' + + turnBackAndDiv :: NNumExp -> NNumExp -> AlgSimplifyM ScalExp + turnBackAndDiv ee1 ee2 = do + ee1' <- fromNumSofP ee1 + ee2' <- fromNumSofP ee2 + return $ SDiv ee1' ee2' + +simplifyScal (SMod e1o e2o) = + SMod <$> simplifyScal e1o <*> simplifyScal e2o + +simplifyScal (SQuot e1o e2o) = + SQuot <$> simplifyScal e1o <*> simplifyScal e2o + +simplifyScal (SRem e1o e2o) = + SRem <$> simplifyScal e1o <*> simplifyScal e2o + +--------------------------------------------------- +--------------------------------------------------- +--- Power related simplifications --- +--------------------------------------------------- +--------------------------------------------------- + +-- cannot handle 0^a, because if a < 0 then it's an error. +-- Could be extented to handle negative exponents better, if needed +simplifyScal (SPow e1 e2) = do + let tp = scalExpType e1 + e1' <- simplifyScal e1 + e2' <- simplifyScal e2 + + if isCt1 e1' || isCt0 e2' + then Val <$> getPos1 tp + else if isCt1 e2' + then return e1' + else case (e1', e2') of + (Val v1, Val v2) + | Just v <- powVals v1 v2 -> return $ Val v + (_, Val (IntValue n)) -> + if P.intToInt64 n >= 1 + then -- simplifyScal =<< fromNumSofP $ NProd (replicate n e1') tp + do new_e <- fromNumSofP $ NProd (genericReplicate (P.intToInt64 n) e1') tp + simplifyScal new_e + else return $ SPow e1' e2' + (_, _) -> return $ SPow e1' e2' + + where + powVals :: PrimValue -> PrimValue -> Maybe PrimValue + powVals (IntValue v1) (IntValue v2) = IntValue <$> P.doPow v1 v2 + powVals _ _ = Nothing + +----------------------------------------------------- +--- Helpers for simplifyScal: MinMax related, etc --- +----------------------------------------------------- + +isMaxMin :: ScalExp -> Bool +isMaxMin MaxMin{} = True +isMaxMin _ = False + +helperPlusMinMax :: ScalExp -> AlgSimplifyM ScalExp +helperPlusMinMax (SPlus (MaxMin ismin es) e) = + simplifyScal $ MaxMin ismin $ map (`SPlus` e) es +helperPlusMinMax (SPlus e (MaxMin ismin es)) = + simplifyScal $ MaxMin ismin $ map (SPlus e) es +helperPlusMinMax _ = badAlgSimplifyM "helperPlusMinMax: Reached unreachable case!" + +{- +helperMinusMinMax :: ScalExp -> AlgSimplifyM ScalExp +helperMinusMinMax (SMinus (MaxMin ismin es) e) = + simplifyScal $ MaxMin ismin $ map (\x -> SMinus x e) es +helperMinusMinMax (SMinus e (MaxMin ismin es)) = + simplifyScal $ MaxMin ismin $ map (\x -> SMinus e x) es +helperMinusMinMax _ = do + pos <- asks pos + badAlgSimplifyM "helperMinusMinMax: Reached unreachable case!" +-} +helperMultMinMax :: ScalExp -> AlgSimplifyM ScalExp +helperMultMinMax (STimes e em@MaxMin{}) = helperTimesDivMinMax True True em e +helperMultMinMax (STimes em@MaxMin{} e) = helperTimesDivMinMax True False em e +helperMultMinMax (SDiv e em@MaxMin{}) = helperTimesDivMinMax False True em e +helperMultMinMax (SDiv em@MaxMin{} e) = helperTimesDivMinMax False False em e +helperMultMinMax _ = badAlgSimplifyM "helperMultMinMax: Reached unreachable case!" + +helperTimesDivMinMax :: Bool -> Bool -> ScalExp -> ScalExp -> AlgSimplifyM ScalExp +helperTimesDivMinMax isTimes isRev emo@MaxMin{} e = do + em <- simplifyScal emo + case em of + MaxMin ismin es -> do + e' <- simplifyScal e + e'_sop <- toNumSofP e' + p' <- simplifyNRel $ NRelExp LTH0 e'_sop + case p' of + LogCt ctbool -> do +-- let cond = not isTimes && isRev +-- let cond'= if ctbool then cond else not cond +-- let ismin'= if cond' then ismin else not ismin + + let cond = ( isTimes && not ctbool ) || + ( not isTimes && not isRev && not ctbool ) || + ( not isTimes && isRev && ctbool ) + let ismin' = if cond then ismin else not ismin + simplifyScal $ MaxMin ismin' $ map (`mkTimesDiv` e') es + + _ -> if not isTimes then return $ mkTimesDiv em e' + else -- e' * MaxMin{...} + case e'_sop of + NProd _ _ -> return $ mkTimesDiv em e' -- simplifyScal =<< fromNumSofP (NProd (em:fs) tp) + NSum ts tp -> do + new_ts <- + mapM (\case + NProd fs _ -> return $ NProd (em:fs) tp + _ -> badAlgSimplifyM + "helperTimesDivMinMax: SofP invariant violated!" + ) ts + simplifyScal =<< fromNumSofP ( NSum new_ts tp ) + _ -> simplifyScal $ mkTimesDiv em e + where + mkTimesDiv :: ScalExp -> ScalExp -> ScalExp + mkTimesDiv e1 e2 + | not isTimes = if isRev then SDiv e2 e1 else SDiv e1 e2 + | isRev = STimes e2 e1 + | otherwise = STimes e1 e2 + +helperTimesDivMinMax _ _ _ _ = + badAlgSimplifyM "helperTimesDivMinMax: Reached unreachable case!" + + +--------------------------------------------------- +--------------------------------------------------- +--- translating to and simplifying the --- +--- disjunctive normal form: toDNF, simplifyDNF --- +--------------------------------------------------- +--------------------------------------------------- +--isTrueDNF :: DNF -> Bool +--isTrueDNF [[LogCt True]] = True +--isTrueDNF _ = False +-- +--getValueDNF :: DNF -> Maybe Bool +--getValueDNF [[LogCt True]] = Just True +--getValueDNF [[LogCt False]] = Just True +--getValueDNF _ = Nothing + + +negateBTerm :: BTerm -> AlgSimplifyM BTerm +negateBTerm (LogCt v) = return $ LogCt (not v) +negateBTerm (PosId i) = return $ NegId i +negateBTerm (NegId i) = return $ PosId i +negateBTerm (NRelExp rel e) = do + let tp = typeOfNAlg e + case (tp, rel) of + (IntType it, LTH0) -> do + se <- fromNumSofP e + ne <- toNumSofP =<< simplifyScal (SNeg $ SPlus se (Val (value (P.intValue it (1::Int))))) + return $ NRelExp LTH0 ne + _ -> NRelExp (if rel == LEQ0 then LTH0 else LEQ0) <$> + (toNumSofP =<< negateSimplified =<< fromNumSofP e) + +bterm2ScalExp :: BTerm -> AlgSimplifyM ScalExp +bterm2ScalExp (LogCt v) = return $ Val $ BoolValue v +bterm2ScalExp (PosId i) = return $ Id i int32 +bterm2ScalExp (NegId i) = return $ SNot $ Id i int32 +bterm2ScalExp (NRelExp rel e) = RelExp rel <$> fromNumSofP e + +-- translates from DNF to ScalExp +fromDNF :: DNF -> AlgSimplifyM ScalExp +fromDNF [] = badAlgSimplifyM "fromDNF: empty DNF!" +fromDNF (t:ts) = do + t' <- translFact t + foldM (\acc x -> do x' <- translFact x; return $ SLogOr x' acc) t' ts + where + translFact [] = badAlgSimplifyM "fromDNF, translFact empty DNF factor!" + translFact (f:fs) = do + f' <- bterm2ScalExp f + foldM (\acc x -> do x' <- bterm2ScalExp x; return $ SLogAnd x' acc) f' fs + +-- translates (and simplifies numeric expressions?) to DNF form. +toDNF :: ScalExp -> AlgSimplifyM DNF +toDNF (Val (BoolValue v)) = return [[LogCt v]] +toDNF (Id idd _ ) = return [[PosId idd]] +toDNF (RelExp rel e ) = do + let t = scalExpType e + case t of + IntType it -> do + e' <- if rel == LEQ0 + then do m1 <- getNeg1 $ IntType it + return $ SPlus e $ Val m1 + else return e + ne <- toNumSofP =<< simplifyScal e' + nrel <- simplifyNRel $ NRelExp LTH0 ne -- False + return [[nrel]] + + _ -> do ne <- toNumSofP =<< simplifyScal e + nrel <- markGaussLTH0 $ simplifyNRel $ NRelExp rel ne + return [[nrel]] +-- +toDNF (SNot (SNot e)) = toDNF e +toDNF (SNot (Val (BoolValue v))) = return [[LogCt $ not v]] +toDNF (SNot (Id idd _)) = return [[NegId idd]] +toDNF (SNot (RelExp rel e)) = do + let not_rel = if rel == LEQ0 then LTH0 else LEQ0 + neg_e <- simplifyScal (SNeg e) + toDNF $ RelExp not_rel neg_e +-- +toDNF (SLogOr e1 e2 ) = do + e1s <- toDNF e1 + e2s <- toDNF e2 + return $ sort $ e1s ++ e2s +toDNF (SLogAnd e1 e2 ) = do + -- [t1 ++ t2 | t1 <- toDNF e1, t2 <- toDNF e2] + e1s <- toDNF e1 + e2s <- toDNF e2 + let lll = map (\t2-> map (++t2) e1s) e2s + return $ sort $ concat lll +toDNF (SNot (SLogAnd e1 e2)) = do + e1s <- toDNF (SNot e1) + e2s <- toDNF (SNot e2) + return $ sort $ e1s ++ e2s +toDNF (SNot (SLogOr e1 e2)) = do + -- [t1 ++ t2 | t1 <- dnf $ SNot e1, t2 <- dnf $ SNot e2] + e1s <- toDNF $ SNot e1 + e2s <- toDNF $ SNot e2 + let lll = map (\t2-> map (++t2) e1s) e2s + return $ sort $ concat lll +toDNF _ = badAlgSimplifyM "toDNF: not a boolean expression!" + +------------------------------------------------------ +--- Simplifying Boolean Expressions: --- +--- 0. p AND p == p; p OR p == p --- +--- 1. False AND p == False; True OR p == True --- +--- 2. True AND p == p; False OR p == p --- +--- 3.(not p)AND p == FALSE; (not p)OR p == True --- +--- 4. ToDo: p1 AND p2 == p1 if p1 => p2 --- +--- p1 AND p2 == False if p1 => not p2 or--- +--- p2 => not p1 --- +--- Also: p1 OR p2 == p2 if p1 => p2 --- +--- p1 OR p2 == True if not p1 => p2 or --- +--- not p2 => p1 --- +--- This boils down to relations: --- +--- e1 < 0 => e2 < 0 if e2 <= e1 --- +------------------------------------------------------ +simplifyDNF :: DNF -> AlgSimplifyM DNF +simplifyDNF terms0 = do + terms1 <- mapM (simplifyAndOr True) terms0 + let terms' = if [LogCt True] `elem` terms1 then [[LogCt True]] + else S.toList $ S.fromList $ + filter (/= [LogCt False]) terms1 + if null terms' then return [[LogCt False]] + else do + let len1terms = all ((1==) . length) terms' + if not len1terms then return terms' + else do let terms_flat = concat terms' + terms'' <- simplifyAndOr False terms_flat + return $ map (:[]) terms'' + +-- big helper function for simplifyDNF +simplifyAndOr :: Bool -> [BTerm] -> AlgSimplifyM [BTerm] +simplifyAndOr _ [] = badAlgSimplifyM "simplifyAndOr: not a boolean expression!" +simplifyAndOr is_and fs = + if LogCt (not is_and) `elem` fs + -- False AND p == False + then return [LogCt $ not is_and] + -- (i) p AND p == p, (ii) True AND p == p + else do let fs' = S.toList . S.fromList . filter (/=LogCt is_and) $ fs + if null fs' + then return [LogCt is_and] + else do -- IF p1 => p2 THEN p1 AND p2 --> p1 + fs''<- foldM (\l x-> do (addx, l') <- trimImplies is_and x l + return $ if addx then x:l' else l' + ) [] fs' + -- IF p1 => not p2 THEN p1 AND p2 == False + isF <- foldM (\b x -> if b then return b + else do notx <- negateBTerm x + impliesAny is_and x notx fs'' + ) False fs'' + return $ if not isF then fs'' + else if is_and then [LogCt False] + else [LogCt True ] + where + -- e1 => e2 ? + impliesRel :: BTerm -> BTerm -> AlgSimplifyM Bool + impliesRel (LogCt False) _ = return True + impliesRel _ (LogCt True) = return True + impliesRel (LogCt True) e = do + let e' = e -- simplifyNRel e + return $ e' == LogCt True + impliesRel e (LogCt False) = do + e' <- negateBTerm e -- simplifyNRel =<< negateBTerm e + return $ e' == LogCt True + impliesRel (NRelExp rel1 e1) (NRelExp rel2 e2) = do + -- ToDo: implement implies relation! + --not_aggr <- asks cheap + let btp = typeOfNAlg e1 + if btp /= typeOfNAlg e2 + then return False + else do + one <- getPos1 btp + e1' <- fromNumSofP e1 + e2' <- fromNumSofP e2 + case (rel1, rel2, btp) of + (LTH0, LTH0, IntType _) -> do + e2me1m1 <- toNumSofP =<< simplifyScal (SMinus e2' $ SPlus e1' $ Val one) + diffrel <- simplifyNRel $ NRelExp LTH0 e2me1m1 + return $ diffrel == LogCt True + (_, _, IntType _) -> badAlgSimplifyM "impliesRel: LEQ0 for Int!" + (_, _, _) -> badAlgSimplifyM "impliesRel: exp of illegal type!" + impliesRel p1 p2 + | p1 == p2 = return True + | otherwise = return False + + -- trimImplies(true, x, l) performs: p1 AND p2 == p1 if p1 => p2, + -- i.e., removes any p in l such that: + -- (i) x => p orelse if + -- (ii) p => x then indicates that p should not be added to the reuslt + -- trimImplies(false, x, l) performs: p1 OR p2 == p2 if p1 => p2, + -- i.e., removes any p from l such that: + -- (i) p => x orelse if + -- (ii) x => p then indicates that p should not be added to the result + trimImplies :: Bool -> BTerm -> [BTerm] -> AlgSimplifyM (Bool, [BTerm]) + trimImplies _ _ [] = return (True, []) + trimImplies and_case x (p:ps) = do + succc <- impliesRel x p + if succc + then if and_case then trimImplies and_case x ps else return (False, p:ps) + else do suc <- impliesRel p x + if suc then if and_case then return (False, p:ps) else trimImplies and_case x ps + else do (addx, newps) <- trimImplies and_case x ps + return (addx, p:newps) + + -- impliesAny(true, x, notx, l) performs: + -- x AND p == False if p => not x, where p in l, + -- impliesAny(true, x, notx, l) performs: + -- x OR p == True if not x => p + -- BUT only when p != x, i.e., avoids comparing x with notx + impliesAny :: Bool -> BTerm -> BTerm -> [BTerm] -> AlgSimplifyM Bool + impliesAny _ _ _ [] = return False + impliesAny and_case x notx (p:ps) + | x == p = impliesAny and_case x notx ps + | otherwise = do + succ' <- if and_case then impliesRel p notx else impliesRel notx p + if succ' then return True + else impliesAny and_case x notx ps + +------------------------------------------------ +--- Syntax-Directed (Brainless) Translators --- +--- scalExp <-> NNumExp --- +--- and negating a scalar expression --- +------------------------------------------------ + +-- negates an already simplified scalar expression, +-- presumably more efficient than negating and +-- then simplifying it. +negateSimplified :: ScalExp -> AlgSimplifyM ScalExp +negateSimplified (SNeg e) = return e +negateSimplified (SNot e) = return e +negateSimplified (SAbs e) = return $ SAbs e +negateSimplified (SSignum e) = + SSignum <$> negateSimplified e +negateSimplified e@(Val v) = do + m1 <- getNeg1 $ scalExpType e + v' <- mulVals m1 v; return $ Val v' +negateSimplified e@Id{} = do + m1 <- getNeg1 $ scalExpType e + return $ STimes (Val m1) e +negateSimplified (SMinus e1 e2) = do -- return $ SMinus e2 e1 + e1' <- negateSimplified e1 + return $ SPlus e1' e2 +negateSimplified (SPlus e1 e2) = do + e1' <- negateSimplified e1 + e2' <- negateSimplified e2 + return $ SPlus e1' e2' +negateSimplified e@(SPow _ _) = do + m1 <- getNeg1 $ scalExpType e + return $ STimes (Val m1) e +negateSimplified (STimes e1 e2) = do + (e1', e2') <- helperNegateMult e1 e2; return $ STimes e1' e2' +negateSimplified (SDiv e1 e2) = do + (e1', e2') <- helperNegateMult e1 e2; return $ SDiv e1' e2' +negateSimplified (SMod e1 e2) = + return $ SMod e1 e2 +negateSimplified (SQuot e1 e2) = do + (e1', e2') <- helperNegateMult e1 e2; return $ SQuot e1' e2' +negateSimplified (SRem e1 e2) = + return $ SRem e1 e2 +negateSimplified (MaxMin ismin ts) = + MaxMin (not ismin) <$> mapM negateSimplified ts +negateSimplified (RelExp LEQ0 e) = + RelExp LTH0 <$> negateSimplified e +negateSimplified (RelExp LTH0 e) = + RelExp LEQ0 <$> negateSimplified e +negateSimplified SLogAnd{} = badAlgSimplifyM "negateSimplified: SLogAnd unimplemented!" +negateSimplified SLogOr{} = badAlgSimplifyM "negateSimplified: SLogOr unimplemented!" + +helperNegateMult :: ScalExp -> ScalExp -> AlgSimplifyM (ScalExp, ScalExp) +helperNegateMult e1 e2 = + case (e1, e2) of + (Val _, _) -> do e1'<- negateSimplified e1; return (e1', e2) + (STimes (Val v) e1r, _) -> do ev <- negateSimplified (Val v); return (STimes ev e1r, e2) + (_, Val _) -> do e2'<- negateSimplified e2; return (e1, e2') + (_, STimes (Val v) e2r) -> do ev <- negateSimplified (Val v); return (e1, STimes ev e2r) + (_, _) -> do e1'<- negateSimplified e1; return (e1', e2) + + +toNumSofP :: ScalExp -> AlgSimplifyM NNumExp +toNumSofP e@(Val _) = return $ NProd [e] $ scalExpType e +toNumSofP e@(Id _ _) = return $ NProd [e] $ scalExpType e +toNumSofP e@SDiv{} = return $ NProd [e] $ scalExpType e +toNumSofP e@SPow{} = return $ NProd [e] $ scalExpType e +toNumSofP (SMinus _ _) = badAlgSimplifyM "toNumSofP: SMinus is not in SofP form!" +toNumSofP (SNeg _) = badAlgSimplifyM "toNumSofP: SNeg is not in SofP form!" +toNumSofP (STimes e1 e2) = do + e2' <- toNumSofP e2 + case e2' of + NProd es2 t -> return $ NProd (e1:es2) t + _ -> badAlgSimplifyM "toNumSofP: STimes nor in SofP form!" +toNumSofP (SPlus e1 e2) = do + let t = scalExpType e1 + e1' <- toNumSofP e1 + e2' <- toNumSofP e2 + case (e1', e2') of + (NSum es1 _, NSum es2 _) -> return $ NSum (es1++es2) t + (NSum es1 _, NProd{}) -> return $ NSum (es1++[e2']) t + (NProd{}, NSum es2 _) -> return $ NSum (e1':es2) t + (NProd{}, NProd{} ) -> return $ NSum [e1', e2'] t +toNumSofP me@MaxMin{} = + return $ NProd [me] $ scalExpType me +toNumSofP s_e = return $ NProd [s_e] $ scalExpType s_e + + +fromNumSofP :: NNumExp -> AlgSimplifyM ScalExp +fromNumSofP (NSum [ ] t) = + Val <$> getZero t +fromNumSofP (NSum [f] _) = fromNumSofP f +fromNumSofP (NSum (f:fs) t) = do + fs_e <- fromNumSofP $ NSum fs t + f_e <- fromNumSofP f + return $ SPlus f_e fs_e +fromNumSofP (NProd [] _) = + badAlgSimplifyM " In fromNumSofP, empty NProd expression! " +fromNumSofP (NProd [f] _) = return f +fromNumSofP (NProd (f:fs) t) = do + fs_e <- fromNumSofP $ NProd fs t + return $ STimes f fs_e +--fromNumSofP _ = do +-- pos <- asks pos +-- badAlgSimplifyM "fromNumSofP: unimplemented!" +------------------------------------------------------------ +--- Helpers for simplifyScal: getTerms, getMultChildren, --- +--- splitTerm, joinTerm, discriminate +------------------------------------------------------------ + + +-- get the list of terms of an expression +-- BUG for NMinMax -> should convert it back to a ScalExp +getTerms :: NNumExp -> [NNumExp] +getTerms (NSum es _) = es +getTerms e@NProd{} = [e] + +-- get the factors of a term +getMultChildren :: NNumExp -> AlgSimplifyM [ScalExp] +getMultChildren (NSum _ _) = badAlgSimplifyM "getMultChildren, NaryPlus should not be nested 2 levels deep " +getMultChildren (NProd xs _) = return xs + +-- split a term into a (multiplicative) value and the rest of the factors. +splitTerm :: NNumExp -> AlgSimplifyM (NNumExp, PrimValue) +splitTerm (NProd [ ] _) = badAlgSimplifyM "splitTerm: Empty n-ary list of factors." +splitTerm (NProd [f] tp) = do + one <- getPos1 tp + case f of + (Val v) -> return (NProd [Val one] tp, v ) + e -> return (NProd [e] tp, one) +splitTerm ne@(NProd (f:fs) tp) = + case f of + (Val v) -> return (NProd fs tp, v) + _ -> do one <- getPos1 tp + return (ne, one) +splitTerm e = do + one <- getPos1 (typeOfNAlg e) + return (e, one) + +-- join a value with a list of factors into a term. +joinTerm :: (NNumExp, PrimValue) -> AlgSimplifyM NNumExp +joinTerm ( NSum _ _, _) = badAlgSimplifyM "joinTerm: NaryPlus two levels deep." +joinTerm ( NProd [] _, _) = badAlgSimplifyM "joinTerm: Empty NaryProd." +joinTerm ( NProd (Val l:fs) tp, v) = do + v' <- mulVals v l + let v'Lit = Val v' + return $ NProd (v'Lit:sort fs) tp +joinTerm ( e@(NProd fs tp), v) + | P.oneIsh v = return e + | otherwise = let vExp = Val v + in return $ NProd (vExp:sort fs) tp + +-- adds up the values corresponding to identical factors! +discriminate :: [(NNumExp, PrimValue)] -> (NNumExp, PrimValue) -> AlgSimplifyM [(NNumExp, PrimValue)] +discriminate [] e = return [e] +discriminate e@((k,v):t) (k', v') = + if k == k' + then do v'' <- addVals v v' + return ( (k, v'') : t ) + else return ( (k', v') : e ) + +------------------------------------------------------ +--- Trivial Utility Functions --- +------------------------------------------------------ + +getZero :: PrimType -> AlgSimplifyM PrimValue +getZero (IntType t) = return $ value $ intValue t (0::Int) +getZero tp = badAlgSimplifyM ("getZero for type: "++pretty tp) + +getPos1 :: PrimType -> AlgSimplifyM PrimValue +getPos1 (IntType t) = return $ value $ intValue t (1::Int) +getPos1 tp = badAlgSimplifyM ("getOne for type: "++pretty tp) + +getNeg1 :: PrimType -> AlgSimplifyM PrimValue +getNeg1 (IntType t) = return $ value $ intValue t (-1::Int) +getNeg1 tp = badAlgSimplifyM ("getOne for type: "++pretty tp) + +valLTHEQ0 :: RelOp0 -> PrimValue -> AlgSimplifyM Bool +valLTHEQ0 LEQ0 (IntValue iv) = return $ P.intToInt64 iv <= 0 +valLTHEQ0 LTH0 (IntValue iv) = return $ P.intToInt64 iv < 0 +valLTHEQ0 _ _ = badAlgSimplifyM "valLTHEQ0 for non-numeric type!" + +isCt1 :: ScalExp -> Bool +isCt1 (Val v) = P.oneIsh v +isCt1 _ = False + +isCt0 :: ScalExp -> Bool +isCt0 (Val v) = P.zeroIsh v +isCt0 _ = False + + +addVals :: PrimValue -> PrimValue -> AlgSimplifyM PrimValue +addVals (IntValue v1) (IntValue v2) = + return $ IntValue $ P.doAdd v1 v2 +addVals _ _ = + badAlgSimplifyM "addVals: operands not of (the same) numeral type! " + +mulVals :: PrimValue -> PrimValue -> AlgSimplifyM PrimValue +mulVals (IntValue v1) (IntValue v2) = + return $ IntValue $ P.doMul v1 v2 +mulVals v1 v2 = + badAlgSimplifyM $ "mulVals: operands not of (the same) numeral type! "++ + pretty v1++" "++pretty v2 + +divVals :: PrimValue -> PrimValue -> AlgSimplifyM PrimValue +divVals (IntValue v1) (IntValue v2) = + case P.doSDiv v1 v2 of + Just v -> return $ IntValue v + Nothing -> badAlgSimplifyM "Division by zero" +divVals _ _ = + badAlgSimplifyM "divVals: operands not of (the same) numeral type! " + +canDivValsEvenly :: PrimValue -> PrimValue -> AlgSimplifyM Bool +canDivValsEvenly (IntValue v1) (IntValue v2) = + case P.doSMod v1 v2 of + Just v -> return $ P.zeroIsh $ IntValue v + Nothing -> return False +canDivValsEvenly _ _ = + badAlgSimplifyM "canDivValsEvenly: operands not of (the same) numeral type!" + +------------------------------------------------------------- +------------------------------------------------------------- +---- Helpers for the ScalExp and NNumRelLogExp Datatypes ---- +------------------------------------------------------------- +------------------------------------------------------------- + +typeOfNAlg :: NNumExp -> PrimType +typeOfNAlg (NSum _ t) = t +typeOfNAlg (NProd _ t) = t + +---------------------------------------- +---- Helpers for Division and Power ---- +---------------------------------------- +trySimplifyDivRec :: [ScalExp] -> [ScalExp] -> [(NNumExp, PrimValue)] -> + AlgSimplifyM ([ScalExp], [(NNumExp, PrimValue)]) +trySimplifyDivRec [] fs' spl_terms = + return (fs', spl_terms) +trySimplifyDivRec (f:fs) fs' spl_terms = do + res_tmp <- mapM (tryDivProdByOneFact f) spl_terms + let (succs, spl_terms') = unzip res_tmp + if all (==True) succs + then trySimplifyDivRec fs fs' spl_terms' + else trySimplifyDivRec fs (fs'++[f]) spl_terms + + +tryDivProdByOneFact :: ScalExp -> (NNumExp, PrimValue) -> AlgSimplifyM (Bool, (NNumExp, PrimValue)) +tryDivProdByOneFact (Val f) (e, v) = do + succc <- canDivValsEvenly v f + if succc then do vres <- divVals v f + return (True, (e, vres)) + else return (False,(e, v) ) + +tryDivProdByOneFact _ pev@(NProd [] _, _) = return (False, pev) +tryDivProdByOneFact f pev@(NProd (t:tfs) tp, v) = do + (succc, newt) <- tryDivTriv t f + one <- getPos1 tp + if not succc + then do (succ', (tfs', v')) <- tryDivProdByOneFact f (NProd tfs tp, v) + case (succ', tfs') of + (True, NProd (Val vv:tfs'') _) -> do + vres <- mulVals v' vv + return (True, (NProd (t:tfs'') tp, vres)) + (True, NProd tfs'' _) -> return (True, (NProd (t:tfs'') tp, v')) + (_, _) -> return (False, pev) + else case (newt, tfs) of + (Val vv, _) -> do vres <- mulVals vv v + return $ if null tfs + then (True, (NProd [Val one] tp, vres)) + else (True, (NProd tfs tp, vres)) + (_, _) -> return (True, (NProd (newt:tfs) tp, v)) + +tryDivProdByOneFact _ (NSum _ _, _) = + badAlgSimplifyM "tryDivProdByOneFact: unreachable case NSum reached!" + + +tryDivTriv :: ScalExp -> ScalExp -> AlgSimplifyM (Bool, ScalExp) +tryDivTriv (SPow a e1) (SPow d e2) + | a == d && e1 == e2 = do one <- getPos1 $ scalExpType a + return (True, Val one) + | a == d = do + let tp = scalExpType a + one <- getPos1 tp + e1me2 <- simplifyScal $ SMinus e1 e2 + case (tp, e1me2) of + (IntType _, Val v) | P.zeroIsh v -> + return (True, Val one) + (IntType _, Val v) | P.oneIsh v -> + return (True, a) + (IntType _, _) -> do + e2me1 <- negateSimplified e1me2 + e2me1_sop <- toNumSofP e2me1 + p' <- simplifyNRel $ NRelExp LTH0 e2me1_sop + return $ if p' == LogCt True + then (True, SPow a e1me2) + else (False, SDiv (SPow a e1) (SPow d e2)) + + (_, _) -> return (False, SDiv (SPow a e1) (SPow d e2)) + + | otherwise = return (False, SDiv (SPow a e1) (SPow d e2)) + +tryDivTriv (SPow a e1) b + | a == b = do one <- getPos1 $ scalExpType a + tryDivTriv (SPow a e1) (SPow a (Val one)) + | otherwise = return (False, SDiv (SPow a e1) b) + +tryDivTriv b (SPow a e1) + | a == b = do one <- getPos1 $ scalExpType a + tryDivTriv (SPow a (Val one)) (SPow a e1) + | otherwise = return (False, SDiv b (SPow a e1)) + +tryDivTriv t f + | t == f = do one <- getPos1 $ scalExpType t + return (True, Val one) + | otherwise = return (False, SDiv t f) + + +{- +mkRelExp :: Int -> (RangesRep, ScalExp, ScalExp) +mkRelExp 1 = + let (i',j',n',p',m') = (tident "int i", tident "int j", tident "int n", tident "int p", tident "int m") + (i,j,n,p,m) = (Id i', Id j', Id n', Id p', Id m') + one = Val (IntVal 1) + min_p_nm1 = MaxMin True [p, SMinus n one] + hash = M.fromList $ [ (identName n', ( 1::Int, Just (Val (IntVal 1)), Nothing ) ), + (identName p', ( 1::Int, Just (Val (IntVal 0)), Nothing ) ), + (identName i', ( 5::Int, Just (Val (IntVal 0)), Just min_p_nm1 ) ) + , (identName j', ( 9::Int, Just (Val (IntVal 0)), Just i ) ) + ] -- M.Map VName (Int, Maybe ScalExp, Maybe ScalExp) + ij_p_j_p_1_m_m = SMinus (SPlus (STimes i j) (SPlus j one)) m + rel1 = RelExp LTH0 ij_p_j_p_1_m_m + m_ij_m_j_m_2 = SNeg ( SPlus (STimes i j) (SPlus j (Val (IntVal 2))) ) + rel2 = RelExp LTH0 m_ij_m_j_m_2 + + in (hash, rel1, rel2) +mkRelExp 2 = + let (i',a',b',l',u') = (tident "int i", tident "int a", tident "int b", tident "int l", tident "int u") + (i,a,b,l,u) = (Id i', Id a', Id b', Id l', Id u') + hash = M.fromList $ [ (identName i', ( 5::Int, Just l, Just u ) ) ] + ai_p_b = SPlus (STimes a i) b + rel1 = RelExp LTH0 ai_p_b + + in (hash, rel1, rel1) +mkRelExp 3 = + let (i',j',n',m') = (tident "int i", tident "int j", tident "int n", tident "int m") + (i,j,n,m) = (Id i', Id j', Id n', Id m') + one = Val (IntVal 1) + two = Val (IntVal 2) + min_j_nm1 = MaxMin True [MaxMin False [Val (IntVal 0), SMinus i (STimes two n)], SMinus n one] + hash = M.fromList $ [ (identName n', ( 1::Int, Just (Val (IntVal 1)), Nothing ) ), + (identName m', ( 2::Int, Just (Val (IntVal 1)), Nothing ) ), + (identName i', ( 5::Int, Just (Val (IntVal 0)), Just (SMinus m one) ) ) + , (identName j', ( 9::Int, Just (Val (IntVal 0)), Just min_j_nm1 ) ) + ] -- M.Map VName (Int, Maybe ScalExp, Maybe ScalExp) + ij_m_m = SMinus (STimes i j) m + rel1 = RelExp LTH0 ij_m_m +-- rel3 = RelExp LTH0 (SMinus i (SPlus (STimes two n) j)) + m_ij_m_1 = SMinus (Val (IntVal (-1))) (STimes i j) + rel2 = RelExp LTH0 m_ij_m_1 + +-- simpl_exp = SDiv (MaxMin True [SMinus (Val (IntVal 0)) (STimes i j), SNeg (STimes i n) ]) +-- (STimes i j) +-- rel4 = RelExp LTH0 simpl_exp + + in (hash, rel1, rel2) + +mkRelExp _ = let hash = M.empty + rel = RelExp LTH0 (Val (IntVal (-1))) + in (hash, rel, rel) +-} diff --git a/src/Futhark/Analysis/Alias.hs b/src/Futhark/Analysis/Alias.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/Alias.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE FlexibleContexts #-} +-- | Alias analysis of a full Futhark program. Takes as input a +-- program with an arbitrary lore and produces one with aliases. This +-- module does not implement the aliasing logic itself, and derives +-- its information from definitions in +-- "Futhark.Representation.AST.Attributes.Aliases" and +-- "Futhark.Representation.Aliases". +module Futhark.Analysis.Alias + ( aliasAnalysis + -- * Ad-hoc utilities + , analyseFun + , analyseStm + , analyseExp + , analyseBody + , analyseLambda + ) + where + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.Aliases + +-- | Perform alias analysis on a Futhark program. +aliasAnalysis :: (Attributes lore, CanBeAliased (Op lore)) => + Prog lore -> Prog (Aliases lore) +aliasAnalysis = Prog . map analyseFun . progFunctions + +analyseFun :: (Attributes lore, CanBeAliased (Op lore)) => + FunDef lore -> FunDef (Aliases lore) +analyseFun (FunDef entry fname restype params body) = + FunDef entry fname restype params body' + where body' = analyseBody body + +analyseBody :: (Attributes lore, + CanBeAliased (Op lore)) => + Body lore -> Body (Aliases lore) +analyseBody (Body lore origbnds result) = + let bnds' = fmap analyseStm origbnds + in mkAliasedBody lore bnds' result + +analyseStm :: (Attributes lore, CanBeAliased (Op lore)) => + Stm lore -> Stm (Aliases lore) +analyseStm (Let pat (StmAux cs attr) e) = + let e' = analyseExp e + pat' = addAliasesToPattern pat e' + lore' = (Names' $ consumedInExp e', attr) + in Let pat' (StmAux cs lore') e' + +analyseExp :: (Attributes lore, CanBeAliased (Op lore)) => + Exp lore -> Exp (Aliases lore) +analyseExp = mapExp analyse + where analyse = + Mapper { mapOnSubExp = return + , mapOnCertificates = return + , mapOnVName = return + , mapOnBody = const $ return . analyseBody + , mapOnRetType = return + , mapOnBranchType = return + , mapOnFParam = return + , mapOnLParam = return + , mapOnOp = return . addOpAliases + } + +analyseLambda :: (Attributes lore, CanBeAliased (Op lore)) => + Lambda lore -> Lambda (Aliases lore) +analyseLambda lam = + let body = analyseBody $ lambdaBody lam + in lam { lambdaBody = body + , lambdaParams = lambdaParams lam + } diff --git a/src/Futhark/Analysis/CallGraph.hs b/src/Futhark/Analysis/CallGraph.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/CallGraph.hs @@ -0,0 +1,64 @@ +-- | This module exports functionality for generating a call graph of +-- an Futhark program. +module Futhark.Analysis.CallGraph + ( CallGraph + , buildCallGraph + ) + where + +import Control.Monad.Writer.Strict +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe (isJust) +import Data.List + +import Futhark.Representation.SOACS + +type FunctionTable = M.Map Name FunDef + +buildFunctionTable :: Prog -> FunctionTable +buildFunctionTable = foldl expand M.empty . progFunctions + where expand ftab f = M.insert (funDefName f) f ftab + +-- | The call graph is just a mapping from a function name, i.e., the +-- caller, to a list of the names of functions called by the function. +-- The order of this list is not significant. +type CallGraph = M.Map Name (S.Set Name) + +-- | @buildCallGraph prog@ build the program's Call Graph. The representation +-- is a hashtable that maps function names to a list of callee names. +buildCallGraph :: Prog -> CallGraph +buildCallGraph prog = foldl' (buildCGfun ftable) M.empty entry_points + where entry_points = map funDefName $ filter (isJust . funDefEntryPoint) $ progFunctions prog + ftable = buildFunctionTable prog + +-- | @buildCallGraph ftable cg fname@ updates Call Graph @cg@ with the +-- contributions of function @fname@, and recursively, with the +-- contributions of the callees of @fname@. +buildCGfun :: FunctionTable -> CallGraph -> Name -> CallGraph +buildCGfun ftable cg fname = + -- Check if function is a non-builtin that we have not already + -- processed. + case M.lookup fname ftable of + Just f | Nothing <- M.lookup fname cg -> do + let callees = buildCGbody $ funDefBody f + cg' = M.insert fname callees cg + -- recursively build the callees + foldl' (buildCGfun ftable) cg' callees + _ -> cg + +buildCGbody :: Body -> S.Set Name +buildCGbody = mconcat . map (buildCGexp . stmExp) . stmsToList . bodyStms + +buildCGexp :: Exp -> S.Set Name +buildCGexp (Apply fname _ _ _) = S.singleton fname +buildCGexp (Op op) = execWriter $ mapSOACM folder op + where folder = identitySOACMapper { + mapOnSOACLambda = \lam -> do tell $ buildCGbody $ lambdaBody lam + return lam + } +buildCGexp e = execWriter $ mapExpM folder e + where folder = identityMapper { + mapOnBody = \_ body -> do tell $ buildCGbody body + return body + } diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/DataDependencies.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE FlexibleContexts #-} +-- | Facilities for inspecting the data dependencies of a program. +module Futhark.Analysis.DataDependencies + ( Dependencies + , dataDependencies + , findNecessaryForReturned + ) + where + +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Futhark.Representation.AST + +-- | A mapping from a variable name @v@, to those variables on which +-- the value of @v@ is dependent. The intuition is that we could +-- remove all other variables, and @v@ would still be computable. +-- This also includes names bound in loops or by lambdas. +type Dependencies = M.Map VName Names + +-- | Compute the data dependencies for an entire body. +dataDependencies :: Attributes lore => Body lore -> Dependencies +dataDependencies = dataDependencies' M.empty + +dataDependencies' :: Attributes lore => + Dependencies -> Body lore -> Dependencies +dataDependencies' startdeps = foldl grow startdeps . bodyStms + where grow deps (Let pat _ (If c tb fb _)) = + let tdeps = dataDependencies' deps tb + fdeps = dataDependencies' deps fb + cdeps = depsOf deps c + comb (pe, tres, fres) = + (patElemName pe, + S.unions $ [freeIn pe, cdeps, depsOf tdeps tres, depsOf fdeps fres] ++ + map (depsOfVar deps) (S.toList $ freeIn pe)) + branchdeps = + M.fromList $ map comb $ zip3 (patternElements pat) + (bodyResult tb) + (bodyResult fb) + in M.unions [branchdeps, deps, tdeps, fdeps] + + grow deps (Let pat _ e) = + let free = freeIn pat <> freeInExp e + freeDeps = S.unions $ map (depsOfVar deps) $ S.toList free + in M.fromList [ (name, freeDeps) | name <- patternNames pat ] `M.union` deps + +depsOf :: Dependencies -> SubExp -> Names +depsOf _ (Constant _) = S.empty +depsOf deps (Var v) = depsOfVar deps v + +depsOfVar :: Dependencies -> VName -> Names +depsOfVar deps name = S.insert name $ M.findWithDefault S.empty name deps + +findNecessaryForReturned :: (Param attr -> Bool) -> [(Param attr, SubExp)] + -> M.Map VName Names + -> Names +findNecessaryForReturned usedAfterLoop merge_and_res allDependencies = + iterateNecessary mempty + where iterateNecessary prev_necessary + | necessary == prev_necessary = necessary + | otherwise = iterateNecessary necessary + where necessary = mconcat $ map dependencies returnedResultSubExps + usedAfterLoopOrNecessary param = + usedAfterLoop param || paramName param `S.member` prev_necessary + returnedResultSubExps = + map snd $ filter (usedAfterLoopOrNecessary . fst) merge_and_res + dependencies (Constant _) = + S.empty + dependencies (Var v) = + M.findWithDefault (S.singleton v) v allDependencies diff --git a/src/Futhark/Analysis/HORepresentation/MapNest.hs b/src/Futhark/Analysis/HORepresentation/MapNest.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/HORepresentation/MapNest.hs @@ -0,0 +1,179 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TupleSections #-} +module Futhark.Analysis.HORepresentation.MapNest + ( Nesting (..) + , MapNest (..) + , typeOf + , params + , inputs + , setInputs + , fromSOAC + , toSOAC + ) +where + +import Control.Monad +import Data.List +import Data.Maybe +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC +import Futhark.Analysis.HORepresentation.SOAC (SOAC) + +import qualified Futhark.Representation.SOACS.SOAC as Futhark +import Futhark.Transform.Substitute +import Futhark.Representation.AST hiding (typeOf) +import Futhark.MonadFreshNames +import Futhark.Construct + +data Nesting lore = Nesting { + nestingParamNames :: [VName] + , nestingResult :: [VName] + , nestingReturnType :: [Type] + , nestingWidth :: SubExp + } deriving (Eq, Ord, Show) + +data MapNest lore = MapNest SubExp (Lambda lore) [Nesting lore] [SOAC.Input] + deriving (Show) + +typeOf :: MapNest lore -> [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 _ lam [] _) = + map paramName $ lambdaParams lam +params (MapNest _ _ (nest:_) _) = + nestingParamNames nest + +inputs :: MapNest lore -> [SOAC.Input] +inputs (MapNest _ _ _ inps) = inps + +setInputs :: [SOAC.Input] -> MapNest lore -> MapNest lore +setInputs [] (MapNest w body ns _) = MapNest w body ns [] +setInputs (inp:inps) (MapNest _ body ns _) = MapNest w body ns' (inp:inps) + where w = arraySize 0 $ SOAC.inputType inp + ws = drop 1 $ arrayDims $ SOAC.inputType inp + ns' = zipWith setDepth ns ws + setDepth n nw = n { nestingWidth = nw } + +fromSOAC :: (Bindable lore, MonadFreshNames m, + LocalScope lore m, + Op lore ~ Futhark.SOAC lore) => + SOAC lore -> m (Maybe (MapNest lore)) +fromSOAC = fromSOAC' mempty + +fromSOAC' :: (Bindable lore, MonadFreshNames m, + LocalScope lore m, + Op lore ~ Futhark.SOAC lore) => + [Ident] + -> SOAC lore + -> m (Maybe (MapNest lore)) + +fromSOAC' bound (SOAC.Screma w (SOAC.ScremaForm (_, []) (_, _, []) lam) inps) = do + maybenest <- case (stmsToList $ bodyStms $ lambdaBody lam, + bodyResult $ lambdaBody lam) of + ([Let pat _ e], res) | res == map Var (patternNames pat) -> + localScope (scopeOfLParams $ lambdaParams lam) $ + SOAC.fromExp e >>= + either (return . Left) (fmap (Right . fmap (pat,)) . fromSOAC' bound') + _ -> + return $ Right Nothing + + case maybenest of + -- Do we have a nested MapNest? + Right (Just (pat, mn@(MapNest inner_w body' ns' inps'))) -> do + (ps, inps'') <- + unzip <$> + fixInputs w (zip (map paramName $ lambdaParams lam) inps) + (zip (params mn) inps') + let n' = Nesting { + nestingParamNames = ps + , nestingResult = patternNames pat + , nestingReturnType = typeOf mn + , nestingWidth = inner_w + } + return $ Just $ MapNest w body' (n':ns') inps'' + -- No nested MapNest it seems. + _ -> do + let isBound name + | Just param <- find ((name==) . identName) bound = + Just param + | otherwise = + Nothing + boundUsedInBody = + mapMaybe isBound $ S.toList $ freeInLambda lam + newParams <- mapM (newIdent' (++"_wasfree")) boundUsedInBody + let subst = M.fromList $ + zip (map identName boundUsedInBody) (map identName newParams) + inps' = map (substituteNames subst) inps ++ + map (SOAC.addTransform (SOAC.Replicate mempty $ Shape [w]) . SOAC.identInput) + boundUsedInBody + lam' = + lam { lambdaBody = + substituteNames subst $ lambdaBody lam + , lambdaParams = + lambdaParams lam ++ [ Param name t + | Ident name t <- newParams ] + } + return $ Just $ MapNest w lam' [] inps' + where bound' = bound <> map paramIdent (lambdaParams lam) + +fromSOAC' _ _ = return Nothing + +toSOAC :: (MonadFreshNames m, HasScope lore m, + Bindable lore, BinderOps lore, Op lore ~ Futhark.SOAC lore) => + MapNest lore -> m (SOAC lore) +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 + let nparams = zipWith Param npnames $ map SOAC.inputRowType inps + (e,bnds) <- runBinder $ localScope (scopeOfLParams nparams) $ SOAC.toExp =<< + toSOAC (MapNest nw lam ns $ map (SOAC.identInput . paramIdent) nparams) + bnd <- mkLetNames nres e + let outerlam = Lambda { lambdaParams = nparams + , lambdaBody = mkBody (bnds<>oneStm bnd) $ map Var nres + , lambdaReturnType = nrettype + } + return $ SOAC.Screma w (Futhark.mapSOAC outerlam) inps + +fixInputs :: MonadFreshNames m => + SubExp -> [(VName, SOAC.Input)] -> [(VName, SOAC.Input)] + -> m [(VName, SOAC.Input)] +fixInputs w ourInps childInps = + reverse . snd <$> foldM inspect (ourInps, []) childInps + where + isParam x (y, _) = x == y + + findParam :: [(VName, SOAC.Input)] + -> VName + -> Maybe ((VName, SOAC.Input), [(VName, SOAC.Input)]) + findParam remPs v + | ([ourP], remPs') <- partition (isParam v) remPs = Just (ourP, remPs') + | otherwise = Nothing + + inspect :: MonadFreshNames m => + ([(VName, SOAC.Input)], [(VName, SOAC.Input)]) + -> (VName, SOAC.Input) + -> m ([(VName, SOAC.Input)], [(VName, SOAC.Input)]) + inspect (remPs, newInps) (_, SOAC.Input ts v _) + | Just ((p,pInp), remPs') <- findParam remPs v = + let pInp' = SOAC.transformRows ts pInp + in return (remPs', + (p, pInp') : newInps) + + | Just ((p,pInp), _) <- findParam newInps v = do + -- The input corresponds to a variable that has already + -- been used. + p' <- newNameFromString $ baseString p + return (remPs, (p', pInp) : newInps) + + inspect (remPs, newInps) (param, SOAC.Input ts a t) = do + param' <- newNameFromString (baseString param ++ "_rep") + return (remPs, (param', + SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t) : newInps) diff --git a/src/Futhark/Analysis/HORepresentation/SOAC.hs b/src/Futhark/Analysis/HORepresentation/SOAC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/HORepresentation/SOAC.hs @@ -0,0 +1,658 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +-- | High-level representation of SOACs. When performing +-- SOAC-transformations, operating on normal 'Exp' values is somewhat +-- of a nuisance, as they can represent terms that are not proper +-- SOACs. In contrast, this module exposes a SOAC representation that +-- does not enable invalid representations (except for type errors). +-- +-- Furthermore, while standard normalised Futhark requires that the inputs +-- to a SOAC are variables or constants, the representation in this +-- module also supports various index-space transformations, like +-- @replicate@ or @rearrange@. This is also very convenient when +-- implementing transformations. +-- +-- The names exported by this module conflict with the standard Futhark +-- syntax tree constructors, so you are advised to use a qualified +-- import: +-- +-- @ +-- import Futhark.Analysis.HORepresentation.SOAC (SOAC) +-- import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC +-- @ +module Futhark.Analysis.HORepresentation.SOAC + ( + -- * SOACs + SOAC (..) + , Futhark.ScremaForm(..) + , inputs + , setInputs + , lambda + , setLambda + , typeOf + , width + -- ** Converting to and from expressions + , NotSOAC (..) + , fromExp + , toExp + , toSOAC + -- * SOAC inputs + , Input (..) + , varInput + , identInput + , isVarInput + , isVarishInput + , addTransform + , addTransforms + , addInitialTransforms + , inputArray + , inputRank + , inputType + , inputRowType + , transformRows + , transposeInput + -- ** Input transformations + , ArrayTransforms + , noTransforms + , singleTransform + , nullTransforms + , (|>) + , (<|) + , viewf + , ViewF(..) + , viewl + , ViewL(..) + , ArrayTransform(..) + , transformFromExp + , soacToStream + ) + where + +import Data.Foldable as Foldable +import Data.Maybe +import Data.Monoid ((<>)) +import qualified Data.Sequence as Seq +import qualified Data.Semigroup as Sem + +import qualified Futhark.Representation.AST as Futhark +import Futhark.Representation.SOACS.SOAC + (StreamForm(..), ScremaForm(..), scremaType, getStreamAccums, GenReduceOp(..)) +import qualified Futhark.Representation.SOACS.SOAC as Futhark +import Futhark.Representation.AST + hiding (Var, Iota, Rearrange, Reshape, Replicate, typeOf) +import Futhark.Transform.Substitute +import Futhark.Construct hiding (toExp) +import Futhark.Transform.Rename (renameLambda) +import qualified Futhark.Util.Pretty as PP +import Futhark.Util.Pretty (ppr, text) + +-- | A single, simple transformation. If you want several, don't just +-- create a list, use 'ArrayTransforms' instead. +data ArrayTransform = Rearrange Certificates [Int] + -- ^ A permutation of an otherwise valid input. + | Reshape Certificates (ShapeChange SubExp) + -- ^ A reshaping of an otherwise valid input. + | ReshapeOuter Certificates (ShapeChange SubExp) + -- ^ A reshaping of the outer dimension. + | ReshapeInner Certificates (ShapeChange SubExp) + -- ^ A reshaping of everything but the outer dimension. + | Replicate Certificates Shape + -- ^ Replicate the rows of the array a number of times. + deriving (Show, Eq, Ord) + +instance Substitute ArrayTransform where + substituteNames substs (Rearrange cs xs) = + Rearrange (substituteNames substs cs) xs + substituteNames substs (Reshape cs ses) = + Reshape (substituteNames substs cs) (substituteNames substs ses) + substituteNames substs (ReshapeOuter cs ses) = + ReshapeOuter (substituteNames substs cs) (substituteNames substs ses) + substituteNames substs (ReshapeInner cs ses) = + ReshapeInner (substituteNames substs cs) (substituteNames substs ses) + substituteNames substs (Replicate cs se) = + Replicate (substituteNames substs cs) (substituteNames substs se) + +-- | A sequence of array transformations, heavily inspired by +-- "Data.Seq". You can decompose it using 'viewF' and 'viewL', and +-- grow it by using '|>' and '<|'. These correspond closely to the +-- similar operations for sequences, except that appending will try to +-- normalise and simplify the transformation sequence. +-- +-- The data type is opaque in order to enforce normalisation +-- invariants. Basically, when you grow the sequence, the +-- implementation will try to coalesce neighboring permutations, for +-- example by composing permutations and removing identity +-- transformations. +newtype ArrayTransforms = ArrayTransforms (Seq.Seq ArrayTransform) + deriving (Eq, Ord, Show) + +instance Sem.Semigroup ArrayTransforms where + ts1 <> ts2 = case viewf ts2 of + t :< ts2' -> (ts1 |> t) <> ts2' + EmptyF -> ts1 + +instance Monoid ArrayTransforms where + mempty = noTransforms + mappend = (Sem.<>) + +instance Substitute ArrayTransforms where + substituteNames substs (ArrayTransforms ts) = + ArrayTransforms $ substituteNames substs <$> ts + +-- | The empty transformation list. +noTransforms :: ArrayTransforms +noTransforms = ArrayTransforms Seq.empty + +-- | Is it an empty transformation list? +nullTransforms :: ArrayTransforms -> Bool +nullTransforms (ArrayTransforms s) = Seq.null s + +-- | A transformation list containing just a single transformation. +singleTransform :: ArrayTransform -> ArrayTransforms +singleTransform = ArrayTransforms . Seq.singleton + +-- | Decompose the input-end of the transformation sequence. +viewf :: ArrayTransforms -> ViewF +viewf (ArrayTransforms s) = case Seq.viewl s of + t Seq.:< s' -> t :< ArrayTransforms s' + Seq.EmptyL -> EmptyF + +-- | A view of the first transformation to be applied. +data ViewF = EmptyF + | ArrayTransform :< ArrayTransforms + +-- | Decompose the output-end of the transformation sequence. +viewl :: ArrayTransforms -> ViewL +viewl (ArrayTransforms s) = case Seq.viewr s of + s' Seq.:> t -> ArrayTransforms s' :> t + Seq.EmptyR -> EmptyL + +-- | A view of the last transformation to be applied. +data ViewL = EmptyL + | ArrayTransforms :> ArrayTransform + +-- | Add a transform to the end of the transformation list. +(|>) :: ArrayTransforms -> ArrayTransform -> ArrayTransforms +(|>) = flip $ addTransform' extract add $ uncurry (flip (,)) + where extract ts' = case viewl ts' of + EmptyL -> Nothing + ts'' :> t' -> Just (t', ts'') + add t' (ArrayTransforms ts') = ArrayTransforms $ ts' Seq.|> t' + +-- | Add a transform at the beginning of the transformation list. +(<|) :: ArrayTransform -> ArrayTransforms -> ArrayTransforms +(<|) = addTransform' extract add id + where extract ts' = case viewf ts' of + EmptyF -> Nothing + t' :< ts'' -> Just (t', ts'') + add t' (ArrayTransforms ts') = ArrayTransforms $ t' Seq.<| ts' + +addTransform' :: (ArrayTransforms -> Maybe (ArrayTransform, ArrayTransforms)) + -> (ArrayTransform -> ArrayTransforms -> ArrayTransforms) + -> ((ArrayTransform,ArrayTransform) -> (ArrayTransform,ArrayTransform)) + -> ArrayTransform -> ArrayTransforms + -> ArrayTransforms +addTransform' extract add swap t ts = + fromMaybe (t `add` ts) $ do + (t', ts') <- extract ts + combined <- uncurry combineTransforms $ swap (t', t) + Just $ if identityTransform combined then ts' + else addTransform' extract add swap combined ts' + +identityTransform :: ArrayTransform -> Bool +identityTransform (Rearrange _ perm) = + Foldable.and $ zipWith (==) perm [0..] +identityTransform _ = False + +combineTransforms :: ArrayTransform -> ArrayTransform -> Maybe ArrayTransform +combineTransforms (Rearrange cs2 perm2) (Rearrange cs1 perm1) = + Just $ Rearrange (cs1<>cs2) $ perm2 `rearrangeCompose` perm1 +combineTransforms _ _ = Nothing + +-- | Given an expression, determine whether the expression represents +-- 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 cs (BasicOp (Futhark.Rearrange perm v)) = + Just (v, Rearrange cs perm) +transformFromExp cs (BasicOp (Futhark.Reshape shape v)) = + Just (v, Reshape cs shape) +transformFromExp cs (BasicOp (Futhark.Replicate shape (Futhark.Var v))) = + Just (v, Replicate cs shape) +transformFromExp _ _ = Nothing + +-- | One array input to a SOAC - a SOAC may have multiple inputs, but +-- all are of this form. Only the array inputs are expressed with +-- this type; other arguments, such as initial accumulator values, are +-- plain expressions. The transforms are done left-to-right, that is, +-- the first element of the 'ArrayTransform' list is applied first. +data Input = Input ArrayTransforms VName Type + deriving (Show, Eq, Ord) + +instance Substitute Input where + substituteNames substs (Input ts v t) = + Input (substituteNames substs ts) + (substituteNames substs v) (substituteNames substs t) + +-- | Create a plain array variable input with no transformations. +varInput :: HasScope t f => VName -> f Input +varInput v = withType <$> lookupType v + where withType = Input (ArrayTransforms Seq.empty) v + +-- | Create a plain array variable input with no transformations, from an 'Ident'. +identInput :: Ident -> Input +identInput v = Input (ArrayTransforms Seq.empty) (identName v) (identType v) + +-- | If the given input is a plain variable input, with no transforms, +-- return the variable. +isVarInput :: Input -> Maybe VName +isVarInput (Input ts v _) | nullTransforms ts = Just v +isVarInput _ = Nothing + +-- | If the given input is a plain variable input, with no non-vacuous transforms, +-- return the variable. +isVarishInput :: Input -> Maybe VName +isVarishInput (Input ts v t) + | nullTransforms ts = Just v + | Reshape cs [DimCoercion _] :< ts' <- viewf ts, cs == mempty = + isVarishInput $ Input ts' v t +isVarishInput _ = Nothing + +-- | Add a transformation to the end of the transformation list. +addTransform :: ArrayTransform -> Input -> Input +addTransform tr (Input trs a t) = + Input (trs |> tr) a t + +-- | Add several transformations to the end of the transformation +-- list. +addTransforms :: ArrayTransforms -> Input -> Input +addTransforms ts (Input ots a t) = Input (ots <> ts) a t + +-- | Add several transformations to the start of the transformation +-- list. +addInitialTransforms :: ArrayTransforms -> Input -> Input +addInitialTransforms ts (Input ots a t) = Input (ts <> ots) a t + +-- | Convert SOAC inputs to the corresponding expressions. +inputsToSubExps :: (MonadBinder m) => + [Input] -> m [VName] +inputsToSubExps = mapM inputToExp' + where inputToExp' (Input (ArrayTransforms ts) a _) = + foldlM transform a ts + + transform ia (Replicate cs n) = + certifying cs $ + letExp "repeat" $ BasicOp $ Futhark.Replicate n (Futhark.Var ia) + + transform ia (Rearrange cs perm) = + certifying cs $ + letExp "rearrange" $ BasicOp $ Futhark.Rearrange perm ia + + transform ia (Reshape cs shape) = + certifying cs $ + letExp "reshape" $ BasicOp $ Futhark.Reshape shape ia + + transform ia (ReshapeOuter cs shape) = do + shape' <- reshapeOuter shape 1 . arrayShape <$> lookupType ia + certifying cs $ + letExp "reshape_outer" $ BasicOp $ Futhark.Reshape shape' ia + + transform ia (ReshapeInner cs shape) = do + shape' <- reshapeInner shape 1 . arrayShape <$> lookupType ia + certifying cs $ + letExp "reshape_inner" $ BasicOp $ Futhark.Reshape shape' ia + +-- | Return the array name of the input. +inputArray :: Input -> VName +inputArray (Input _ v _) = v + +-- | Return the type of an input. +inputType :: Input -> Type +inputType (Input (ArrayTransforms ts) _ at) = + Foldable.foldl transformType at ts + where transformType t (Replicate _ shape) = + arrayOfShape t shape + transformType t (Rearrange _ perm) = + rearrangeType perm t + transformType t (Reshape _ shape) = + t `setArrayShape` newShape shape + transformType t (ReshapeOuter _ shape) = + let Shape oldshape = arrayShape t + in t `setArrayShape` Shape (newDims shape ++ drop 1 oldshape) + transformType t (ReshapeInner _ shape) = + let Shape oldshape = arrayShape t + in t `setArrayShape` Shape (take 1 oldshape ++ newDims shape) + +-- | Return the row type of an input. Just a convenient alias. +inputRowType :: Input -> Type +inputRowType = rowType . inputType + +-- | Return the array rank (dimensionality) of an input. Just a +-- convenient alias. +inputRank :: Input -> Int +inputRank = arrayRank . inputType + +-- | Apply the transformations to every row of the input. +transformRows :: ArrayTransforms -> Input -> Input +transformRows (ArrayTransforms ts) = + flip (Foldable.foldl transformRows') ts + where transformRows' inp (Rearrange cs perm) = + addTransform (Rearrange cs (0:map (+1) perm)) inp + transformRows' inp (Reshape cs shape) = + addTransform (ReshapeInner cs shape) inp + transformRows' inp (Replicate cs n) + | inputRank inp == 1 = + Rearrange mempty [1,0] `addTransform` + (Replicate cs n `addTransform` inp) + | otherwise = + Rearrange mempty (2:0:1:[3..inputRank inp]) `addTransform` + (Replicate cs n `addTransform` + (Rearrange mempty (1:0:[2..inputRank inp-1]) `addTransform` inp)) + transformRows' inp nts = + error $ "transformRows: Cannot transform this yet:\n" ++ show nts ++ "\n" ++ show inp + +-- | Add to the input a 'Rearrange' transform that performs an @(k,n)@ +-- transposition. The new transform will be at the end of the current +-- transformation list. +transposeInput :: Int -> Int -> Input -> Input +transposeInput k n inp = + 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) [Input] + | Scatter SubExp (Lambda lore) [Input] [(SubExp, Int, VName)] + | Screma SubExp (ScremaForm lore) [Input] + | GenReduce SubExp [GenReduceOp lore] (Lambda lore) [Input] + deriving (Eq, Show) + +instance PP.Pretty Input where + ppr (Input (ArrayTransforms ts) arr _) = foldl f (ppr arr) ts + where f e (Rearrange cs perm) = + text "rearrange" <> ppr cs <> PP.apply [PP.apply (map ppr perm), e] + f e (Reshape cs shape) = + text "reshape" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] + f e (ReshapeOuter cs shape) = + text "reshape_outer" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] + f e (ReshapeInner cs shape) = + text "reshape_inner" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] + f e (Replicate cs ne) = + text "replicate" <> ppr cs <> PP.apply [ppr ne, e] + +instance PrettyLore lore => PP.Pretty (SOAC lore) where + ppr (Screma w form arrs) = Futhark.ppScrema w form arrs + ppr (GenReduce len ops bucket_fun imgs) = + Futhark.ppGenReduce len ops bucket_fun imgs + ppr soac = text $ show soac + +-- | Returns the inputs used in a SOAC. +inputs :: SOAC lore -> [Input] +inputs (Stream _ _ _ arrs) = arrs +inputs (Scatter _len _lam ivs _as) = ivs +inputs (Screma _ _ arrs) = arrs +inputs (GenReduce _ _ _ inps) = inps + +-- | Set the inputs to a SOAC. +setInputs :: [Input] -> SOAC lore -> SOAC lore +setInputs arrs (Stream w form lam _) = + Stream (newWidth arrs w) form lam arrs +setInputs arrs (Scatter w lam _ivs as) = + Scatter (newWidth arrs w) lam arrs as +setInputs arrs (Screma w form _) = + Screma w form arrs +setInputs inps (GenReduce w ops lam _) = + GenReduce w ops lam inps + +newWidth :: [Input] -> SubExp -> SubExp +newWidth [] w = w +newWidth (inp:_) _ = arraySize 0 $ inputType inp + +-- | The lambda used in a given SOAC. +lambda :: SOAC lore -> Lambda lore +lambda (Stream _ _ lam _) = lam +lambda (Scatter _len lam _ivs _as) = lam +lambda (Screma _ (ScremaForm _ _ lam) _) = lam +lambda (GenReduce _ _ lam _) = lam + +-- | Set the lambda used in the SOAC. +setLambda :: Lambda lore -> SOAC lore -> SOAC lore +setLambda lam (Stream w form _ arrs) = + Stream w form lam arrs +setLambda lam (Scatter len _lam ivs as) = + Scatter len lam ivs as +setLambda lam (Screma w (ScremaForm scan red _) arrs) = + Screma w (ScremaForm scan red lam) arrs +setLambda lam (GenReduce w ops _ inps) = + GenReduce w ops lam inps + +-- | The return type of a SOAC. +typeOf :: SOAC lore -> [Type] +typeOf (Stream w form lam _) = + let nes = getStreamAccums form + accrtps = take (length nes) $ lambdaReturnType lam + arrtps = [ arrayOf (stripArray 1 t) (Shape [w]) NoUniqueness + | t <- drop (length nes) (lambdaReturnType lam) ] + in accrtps ++ arrtps +typeOf (Scatter _w lam _ivs dests) = + zipWith arrayOfRow (snd $ splitAt (n `div` 2) lam_ts) aws + where lam_ts = lambdaReturnType lam + n = length lam_ts + (aws, _, _) = unzip3 dests +typeOf (Screma w form _) = + scremaType w form +typeOf (GenReduce _ ops _ _) = do + op <- ops + map (`arrayOfRow` genReduceWidth op) (lambdaReturnType $ genReduceOp op) + +-- | 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 (Stream w _ _ _) = w +width (Scatter len _lam _ivs _as) = len +width (Screma w _ _) = w +width (GenReduce w _ _ _) = w + +-- | Convert a SOAC to the corresponding expression. +toExp :: (MonadBinder m, Op (Lore m) ~ Futhark.SOAC (Lore m)) => + SOAC (Lore m) -> m (Exp (Lore 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)) +toSOAC (Stream w form lam inps) = + Futhark.Stream w form lam <$> inputsToSubExps inps +toSOAC (Scatter len lam ivs dests) = do + ivs' <- inputsToSubExps ivs + return $ Futhark.Scatter len lam ivs' dests +toSOAC (Screma w form arrs) = + Futhark.Screma w form <$> inputsToSubExps arrs +toSOAC (GenReduce w ops lam inps) = + Futhark.GenReduce w ops lam <$> inputsToSubExps inps + +-- | The reason why some expression cannot be converted to a 'SOAC' +-- value. +data NotSOAC = NotSOAC -- ^ The expression is not a (tuple-)SOAC at all. + deriving (Show) + +-- | Either convert an expression to the normalised SOAC +-- representation, or a reason why the expression does not have the +-- valid form. +fromExp :: (Op lore ~ Futhark.SOAC lore, Bindable lore, + HasScope lore m, MonadFreshNames m) => + Exp lore -> m (Either NotSOAC (SOAC lore)) + +fromExp (BasicOp (Copy arr)) = do + arr_t <- lookupType arr + p <- Param <$> newVName "copy_p" <*> pure (rowType arr_t) + let lam = Lambda [p] (mkBody mempty [Futhark.Var $ paramName p]) [rowType arr_t] + Right . Screma (arraySize 0 arr_t) (Futhark.mapSOAC lam) . pure <$> varInput arr +fromExp (Op (Futhark.Stream w form lam as)) = + Right . Stream w form lam <$> traverse varInput as +fromExp (Op (Futhark.Scatter len lam ivs as)) = do + ivs' <- traverse varInput ivs + return $ Right $ Scatter len lam ivs' as +fromExp (Op (Futhark.Screma w form arrs)) = + Right . Screma w form <$> traverse varInput arrs +fromExp (Op (Futhark.GenReduce w ops lam arrs)) = + Right . GenReduce w ops lam <$> traverse varInput arrs +fromExp _ = pure $ Left NotSOAC + +-- | To-Stream translation of SOACs. +-- 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]) +soacToStream soac = do + chunk_param <- newParam "chunk" $ Prim int32 + let chvar= Futhark.Var $ paramName chunk_param + (lam, inps) = (lambda soac, inputs soac) + w = width soac + lam' <- renameLambda lam + let arrrtps= mapType w lam + -- the chunked-outersize of the array result and input types + loutps = [ arrayOfRow t chvar | t <- map rowType arrrtps ] + lintps = [ arrayOfRow t chvar | t <- map inputRowType inps ] + + strm_inpids <- mapM (newParam "inp") lintps + -- Treat each SOAC case individually: + case soac of + Screma _ form _ + | Just _ <- Futhark.isMapSOAC form -> do + -- Map(f,a) => is translated in strem's body to: + -- let strm_resids = map(f,a_ch) in strm_resids + -- + -- array result and input IDs of the stream's lambda + strm_resids <- mapM (newIdent "res") loutps + let insoac = Futhark.Screma chvar (Futhark.mapSOAC lam') $ map paramName strm_inpids + insbnd = mkLet [] strm_resids $ Op insoac + strmbdy= mkBody (oneStm insbnd) $ map (Futhark.Var . identName) strm_resids + strmpar= chunk_param:strm_inpids + strmlam= Lambda strmpar strmbdy loutps + empty_lam = Lambda [] (mkBody mempty []) [] + -- map(f,a) creates a stream with NO accumulators + return (Stream w (Parallel Disorder Commutative empty_lam []) strmlam inps, []) + + | Just (scan_lam, nes, _) <- Futhark.isScanomapSOAC form -> do + -- scanomap(scan_lam,nes,map_lam,a) => is translated in strem's body to: + -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam, nes, map_lam, a_ch) + -- 2. let strm_resids = map (acc `+`,nes, scan0_ids) + -- 3. let outerszm1id = sizeof(0,strm_resids) - 1 + -- 4. let lasteel_ids = if outerszm1id < 0 + -- then nes + -- else strm_resids[outerszm1id] + -- 5. let acc' = acc + lasteel_ids + -- {acc', strm_resids, map_resids} + -- the array and accumulator result types + let scan_arr_ts = map (`arrayOfRow` chvar) $ lambdaReturnType scan_lam + map_arr_ts = drop (length nes) loutps + accrtps = lambdaReturnType scan_lam + + -- array result and input IDs of the stream's lambda + strm_resids <- mapM (newIdent "res") scan_arr_ts + scan0_ids <- mapM (newIdent "resarr0") scan_arr_ts + map_resids <- mapM (newIdent "map_res") map_arr_ts + + lastel_ids <- mapM (newIdent "lstel") accrtps + lastel_tmp_ids <- mapM (newIdent "lstel_tmp") accrtps + empty_arr <- newIdent "empty_arr" $ Prim Bool + inpacc_ids <- mapM (newParam "inpacc") accrtps + outszm1id <- newIdent "szm1" $ Prim int32 + -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam,nes,map_lam,a_ch) + let insbnd = mkLet [] (scan0_ids++map_resids) $ Op $ + Futhark.Screma chvar (Futhark.scanomapSOAC scan_lam nes lam') $ + map paramName strm_inpids + -- 2. let outerszm1id = chunksize - 1 + outszm1bnd = mkLet [] [outszm1id] $ BasicOp $ + BinOp (Sub Int32) + (Futhark.Var $ paramName chunk_param) + (constant (1::Int32)) + -- 3. let lasteel_ids = ... + empty_arr_bnd = mkLet [] [empty_arr] $ BasicOp $ CmpOp (CmpSlt Int32) + (Futhark.Var $ identName outszm1id) + (constant (0::Int32)) + leltmpbnds= zipWith (\ lid arrid -> mkLet [] [lid] $ BasicOp $ + Index (identName arrid) $ + fullSlice (identType arrid) + [DimFix $ Futhark.Var $ identName outszm1id] + ) lastel_tmp_ids scan0_ids + lelbnd = mkLet [] lastel_ids $ + If (Futhark.Var $ identName empty_arr) + (mkBody mempty nes) + (mkBody (stmsFromList leltmpbnds) $ + map (Futhark.Var . identName) lastel_tmp_ids) $ + ifCommon $ map identType lastel_tmp_ids + -- 4. let strm_resids = map (acc `+`,nes, scan0_ids) + maplam <- mkMapPlusAccLam (map (Futhark.Var . paramName) inpacc_ids) scan_lam + let mapbnd = mkLet [] strm_resids $ Op $ + Futhark.Screma chvar (Futhark.mapSOAC maplam) $ + map identName scan0_ids + -- 5. let acc' = acc + lasteel_ids + addlelbdy <- mkPlusBnds scan_lam $ map Futhark.Var $ + map paramName inpacc_ids++map identName lastel_ids + -- Finally, construct the stream + let (addlelbnd,addlelres) = (bodyStms addlelbdy, bodyResult addlelbdy) + strmbdy= mkBody (stmsFromList [insbnd,outszm1bnd,empty_arr_bnd,lelbnd,mapbnd]<>addlelbnd) $ + addlelres ++ map (Futhark.Var . identName) (strm_resids ++ map_resids) + strmpar= chunk_param:inpacc_ids++strm_inpids + strmlam= Lambda strmpar strmbdy (accrtps++loutps) + return (Stream w (Sequential nes) strmlam inps, + map paramIdent inpacc_ids) + + | Just (comm, lamin, nes, _) <- Futhark.isRedomapSOAC form -> do + -- Redomap(+,lam,nes,a) => is translated in strem's body to: + -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in + -- 2. let acc' = acc + acc0_ids in + -- {acc', strm_resids} + + let accrtps= take (length nes) $ lambdaReturnType lam + -- the chunked-outersize of the array result and input types + loutps' = drop (length nes) loutps + -- the lambda with proper index + foldlam = lam' + -- array result and input IDs of the stream's lambda + strm_resids <- mapM (newIdent "res") loutps' + inpacc_ids <- mapM (newParam "inpacc") accrtps + acc0_ids <- mapM (newIdent "acc0" ) accrtps + -- 1. let (acc0_ids,strm_resids) = redomap(+,lam,nes,a_ch) in + let insoac = Futhark.Screma chvar (Futhark.redomapSOAC comm lamin nes foldlam) $ + map paramName strm_inpids + insbnd = mkLet [] (acc0_ids++strm_resids) $ Op insoac + -- 2. let acc' = acc + acc0_ids in + addaccbdy <- mkPlusBnds lamin $ map Futhark.Var $ + map paramName inpacc_ids++map identName acc0_ids + -- Construct the stream + let (addaccbnd,addaccres) = (bodyStms addaccbdy, bodyResult addaccbdy) + strmbdy= mkBody (oneStm insbnd <> addaccbnd) $ + addaccres ++ map (Futhark.Var . identName) strm_resids + strmpar= chunk_param:inpacc_ids++strm_inpids + strmlam= Lambda strmpar strmbdy (accrtps++loutps') + lam0 <- renameLambda lamin + return (Stream w (Parallel InOrder comm lam0 nes) strmlam inps, []) + + -- Otherwise it cannot become a stream. + _ -> return (soac,[]) + where mkMapPlusAccLam :: (MonadFreshNames m, Bindable lore) + => [SubExp] -> Lambda lore -> m (Lambda lore) + mkMapPlusAccLam accs plus = do + let lampars = lambdaParams plus + (accpars, rempars) = ( take (length accs) lampars, + drop (length accs) lampars ) + parbnds = zipWith (\ par se -> mkLet [] [paramIdent par] + (BasicOp $ SubExp se) + ) accpars accs + plus_bdy = lambdaBody plus + newlambdy = Body (bodyAttr plus_bdy) + (stmsFromList parbnds <> bodyStms plus_bdy) + (bodyResult plus_bdy) + renameLambda $ Lambda rempars newlambdy $ lambdaReturnType plus + + mkPlusBnds :: (MonadFreshNames m, Bindable lore) + => Lambda lore -> [SubExp] -> m (Body lore) + mkPlusBnds plus accels = do + plus' <- renameLambda plus + let parbnds = zipWith (\ par se -> mkLet [] [paramIdent par] + (BasicOp $ SubExp se) + ) (lambdaParams plus') accels + body = lambdaBody plus' + return $ body { bodyStms = stmsFromList parbnds <> bodyStms body } diff --git a/src/Futhark/Analysis/Metrics.hs b/src/Futhark/Analysis/Metrics.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/Metrics.hs @@ -0,0 +1,129 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +-- | Abstract Syntax Tree metrics. This is used in the @futhark-test@ program. +module Futhark.Analysis.Metrics + ( AstMetrics(..) + , progMetrics + + -- * Extensibility + , OpMetrics(..) + , seen + , inside + , MetricsM + , bodyMetrics + , bindingMetrics + , lambdaMetrics + ) where + +import Control.Monad.Writer +import Data.Text (Text) +import qualified Data.Text as T +import Data.String +import Data.List +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem + +import Futhark.Representation.AST + +newtype AstMetrics = AstMetrics (M.Map Text Int) + +instance Show AstMetrics where + show (AstMetrics m) = unlines $ map metric $ M.toList m + where metric (k, v) = pretty k ++ " " ++ pretty v + +instance Read AstMetrics where + readsPrec _ s = + maybe [] success $ mapM onLine $ lines s + where onLine l = case words l of + [k, x] | [(n, "")] <- reads x -> Just (T.pack k, n) + _ -> Nothing + success m = [(AstMetrics $ M.fromList m, "")] + +class OpMetrics op where + opMetrics :: op -> MetricsM () + +instance OpMetrics () where + opMetrics () = return () + +newtype CountMetrics = CountMetrics [([Text], Text)] + +instance Sem.Semigroup CountMetrics where + CountMetrics x <> CountMetrics y = CountMetrics $ x <> y + +instance Monoid CountMetrics where + mempty = CountMetrics mempty + mappend = (Sem.<>) + +actualMetrics :: CountMetrics -> AstMetrics +actualMetrics (CountMetrics metrics) = + AstMetrics $ M.fromListWith (+) $ concatMap expand metrics + where expand (ctx, k) = + [ (T.intercalate "/" (ctx' ++ [k]), 1) + | ctx' <- tails $ "" : ctx ] + +newtype MetricsM a = MetricsM { runMetricsM :: Writer CountMetrics a } + deriving (Monad, Applicative, Functor, MonadWriter CountMetrics) + +seen :: Text -> MetricsM () +seen k = tell $ CountMetrics [([], k)] + +inside :: Text -> MetricsM () -> MetricsM () +inside what m = seen what >> censor addWhat m + where addWhat (CountMetrics metrics) = + CountMetrics (map addWhat' metrics) + addWhat' (ctx, k) = (what : ctx, k) + +progMetrics :: OpMetrics (Op lore) => Prog lore -> AstMetrics +progMetrics = actualMetrics . execWriter . runMetricsM . mapM_ funDefMetrics . progFunctions + +funDefMetrics :: OpMetrics (Op lore) => FunDef lore -> MetricsM () +funDefMetrics = bodyMetrics . funDefBody + +bodyMetrics :: OpMetrics (Op lore) => Body lore -> MetricsM () +bodyMetrics = mapM_ bindingMetrics . bodyStms + +bindingMetrics :: OpMetrics (Op lore) => Stm lore -> MetricsM () +bindingMetrics = expMetrics . stmExp + +expMetrics :: OpMetrics (Op lore) => Exp lore -> MetricsM () +expMetrics (BasicOp op) = + seen "BasicOp" >> primOpMetrics op +expMetrics (DoLoop _ _ ForLoop{} body) = + inside "DoLoop" $ seen "ForLoop" >> bodyMetrics body +expMetrics (DoLoop _ _ WhileLoop{} body) = + inside "DoLoop" $ seen "WhileLoop" >> bodyMetrics body +expMetrics (If _ tb fb _) = + inside "If" $ do + inside "True" $ bodyMetrics tb + inside "False" $ bodyMetrics fb +expMetrics (Apply fname _ _ _) = + seen $ "Apply" <> fromString (nameToString fname) +expMetrics (Op op) = + opMetrics op + +primOpMetrics :: BasicOp lore -> MetricsM () +primOpMetrics (SubExp _) = seen "SubExp" +primOpMetrics (Opaque _) = seen "Opaque" +primOpMetrics ArrayLit{} = seen "ArrayLit" +primOpMetrics BinOp{} = seen "BinOp" +primOpMetrics UnOp{} = seen "UnOp" +primOpMetrics ConvOp{} = seen "ConvOp" +primOpMetrics CmpOp{} = seen "ConvOp" +primOpMetrics Assert{} = seen "Assert" +primOpMetrics Index{} = seen "Index" +primOpMetrics Update{} = seen "Update" +primOpMetrics Concat{} = seen "Concat" +primOpMetrics Copy{} = seen "Copy" +primOpMetrics Manifest{} = seen "Manifest" +primOpMetrics Iota{} = seen "Iota" +primOpMetrics Replicate{} = seen "Replicate" +primOpMetrics Repeat{} = seen "Repeat" +primOpMetrics Scratch{} = seen "Scratch" +primOpMetrics Reshape{} = seen "Reshape" +primOpMetrics Rearrange{} = seen "Rearrange" +primOpMetrics Rotate{} = seen "Rotate" +primOpMetrics Partition{} = seen "Partition" + +lambdaMetrics :: OpMetrics (Op lore) => Lambda lore -> MetricsM () +lambdaMetrics = bodyMetrics . lambdaBody diff --git a/src/Futhark/Analysis/PrimExp.hs b/src/Futhark/Analysis/PrimExp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/PrimExp.hs @@ -0,0 +1,268 @@ +-- | A primitive expression is an expression where the non-leaves are +-- primitive operators. Our representation does not guarantee that +-- the expression is type-correct. +module Futhark.Analysis.PrimExp + ( PrimExp (..) + , evalPrimExp + , primExpType + , coerceIntPrimExp + + , module Futhark.Representation.Primitive + ) where + +import Data.Foldable +import Data.Traversable +import qualified Data.Map as M + +import Futhark.Representation.AST.Attributes.Names +import Futhark.Representation.Primitive +import Futhark.Util.IntegralExp +import Futhark.Util.Pretty + +-- | A primitive expression parametrised over the representation of free variables. +data PrimExp v = LeafExp v PrimType + | ValueExp PrimValue + | BinOpExp BinOp (PrimExp v) (PrimExp v) + | CmpOpExp CmpOp (PrimExp v) (PrimExp v) + | UnOpExp UnOp (PrimExp v) + | ConvOpExp ConvOp (PrimExp v) + | FunExp String [PrimExp v] PrimType + deriving (Ord, Show) + +-- The Eq instance upcoerces all integer constants to their largest +-- type before comparing for equality. This is technically not a good +-- idea, but solves annoying problems related to the Num instance +-- always producing Int64s. +instance Eq v => Eq (PrimExp v) where + LeafExp x xt == LeafExp y yt = x == y && xt == yt + ValueExp (IntValue x) == ValueExp (IntValue y) = + intToInt64 x == intToInt64 y + ValueExp x == ValueExp y = + x == y + BinOpExp xop x1 x2 == BinOpExp yop y1 y2 = + xop == yop && x1 == y1 && x2 == y2 + CmpOpExp xop x1 x2 == CmpOpExp yop y1 y2 = + xop == yop && x1 == y1 && x2 == y2 + UnOpExp xop x == UnOpExp yop y = + xop == yop && x == y + ConvOpExp xop x == ConvOpExp yop y = + xop == yop && x == y + FunExp xf xargs _ == FunExp yf yargs _ = + xf == yf && xargs == yargs + _ == _ = False + +instance Functor PrimExp where + fmap = fmapDefault + +instance Foldable PrimExp where + foldMap = foldMapDefault + +instance Traversable PrimExp where + traverse f (LeafExp v t) = + LeafExp <$> f v <*> pure t + traverse _ (ValueExp v) = + pure $ ValueExp v + traverse f (BinOpExp op x y) = + BinOpExp op <$> traverse f x <*> traverse f y + traverse f (CmpOpExp op x y) = + CmpOpExp op <$> traverse f x <*> traverse f y + traverse f (ConvOpExp op x) = + ConvOpExp op <$> traverse f x + traverse f (UnOpExp op x) = + UnOpExp op <$> traverse f x + traverse f (FunExp h args t) = + FunExp h <$> traverse (traverse f) args <*> pure t + +instance FreeIn v => FreeIn (PrimExp v) where + freeIn = foldMap freeIn + +-- The Num instance performs a little bit of magic: whenever an +-- expression and a constant is combined with a binary operator, the +-- type of the constant may be changed to be the type of the +-- expression, if they are not already the same. This permits us to +-- write e.g. @x * 4@, where @x@ is an arbitrary PrimExp, and have the +-- @4@ converted to the proper primitive type. We also support +-- converting integers to floating point values, but not the other way +-- around. All numeric instances assume unsigned integers for such +-- conversions. +-- +-- We also perform simple constant folding, in particular to reduce +-- expressions to constants so that the above works. However, it is +-- still a bit of a hack. +instance Pretty v => Num (PrimExp v) where + x + y | zeroIshExp x = y + | zeroIshExp y = x + | IntType t <- primExpType x, + Just z <- constFold (doBinOp $ Add t) x y = z + | FloatType t <- primExpType x, + Just z <- constFold (doBinOp $ FAdd t) x y = z + | Just z <- msum [asIntOp Add x y, asFloatOp FAdd x y] = z + | otherwise = numBad "+" (x,y) + + x - y | zeroIshExp y = x + | IntType t <- primExpType x, + Just z <- constFold (doBinOp $ Sub t) x y = z + | FloatType t <- primExpType x, + Just z <- constFold (doBinOp $ FSub t) x y = z + | Just z <- msum [asIntOp Sub x y, asFloatOp FSub x y] = z + | otherwise = numBad "-" (x,y) + + x * y | zeroIshExp x = x + | zeroIshExp y = y + | oneIshExp x = y + | oneIshExp y = x + | IntType t <- primExpType x, + Just z <- constFold (doBinOp $ Mul t) x y = z + | FloatType t <- primExpType x, + Just z <- constFold (doBinOp $ FMul t) x y = z + | Just z <- msum [asIntOp Mul x y, asFloatOp FMul x y] = z + | otherwise = numBad "*" (x,y) + + abs x | IntType t <- primExpType x = UnOpExp (Abs t) x + | FloatType t <- primExpType x = UnOpExp (FAbs t) x + | otherwise = numBad "abs" x + + signum x | IntType t <- primExpType x = UnOpExp (SSignum t) x + | otherwise = numBad "signum" x + + fromInteger = fromInt32 . fromInteger + +instance Pretty v => IntegralExp (PrimExp v) where + x `div` y | oneIshExp y = x + | Just z <- msum [asIntOp SDiv x y, asFloatOp FDiv x y] = z + | otherwise = numBad "div" (x,y) + + x `mod` y | Just z <- msum [asIntOp SMod x y] = z + | otherwise = numBad "mod" (x,y) + + x `quot` y | oneIshExp y = x + | Just z <- msum [asIntOp SQuot x y] = z + | otherwise = numBad "quot" (x,y) + + x `rem` y | Just z <- msum [asIntOp SRem x y] = z + | otherwise = numBad "rem" (x,y) + + sgn (ValueExp (IntValue i)) = Just $ signum $ valueIntegral i + sgn _ = Nothing + + fromInt8 = ValueExp . IntValue . Int8Value + fromInt16 = ValueExp . IntValue . Int16Value + fromInt32 = ValueExp . IntValue . Int32Value + fromInt64 = ValueExp . IntValue . Int64Value + +asIntOp :: (IntType -> BinOp) -> PrimExp v -> PrimExp v -> Maybe (PrimExp v) +asIntOp f x y + | IntType t <- primExpType x, + Just y' <- asIntExp t y = Just $ BinOpExp (f t) x y' + | IntType t <- primExpType y, + Just x' <- asIntExp t x = Just $ BinOpExp (f t) x' y + | otherwise = Nothing + +asIntExp :: IntType -> PrimExp v -> Maybe (PrimExp v) +asIntExp t e + | primExpType e == IntType t = Just e +asIntExp t (ValueExp (IntValue v)) = + Just $ ValueExp $ IntValue $ doSExt v t +asIntExp _ _ = + Nothing + +asFloatOp :: (FloatType -> BinOp) -> PrimExp v -> PrimExp v -> Maybe (PrimExp v) +asFloatOp f x y + | FloatType t <- primExpType x, + Just y' <- asFloatExp t y = Just $ BinOpExp (f t) x y' + | FloatType t <- primExpType y, + Just x' <- asFloatExp t x = Just $ BinOpExp (f t) x' y + | otherwise = Nothing + +asFloatExp :: FloatType -> PrimExp v -> Maybe (PrimExp v) +asFloatExp t e + | primExpType e == FloatType t = Just e +asFloatExp t (ValueExp (FloatValue v)) = + Just $ ValueExp $ FloatValue $ doFPConv v t +asFloatExp t (ValueExp (IntValue v)) = + Just $ ValueExp $ FloatValue $ doSIToFP v t +asFloatExp _ _ = + Nothing + +constFold :: (PrimValue -> PrimValue -> Maybe PrimValue) + -> PrimExp v -> PrimExp v + -> Maybe (PrimExp v) +constFold f x y = do x' <- valueExp x + y' <- valueExp y + ValueExp <$> f x' y' + +numBad :: Pretty a => String -> a -> b +numBad s x = + error $ "Invalid argument to PrimExp method " ++ s ++ ": " ++ pretty x + +-- | Evaluate a 'PrimExp' in the given monad. Invokes 'fail' on type +-- errors. +evalPrimExp :: (Pretty v, Monad m) => (v -> m PrimValue) -> PrimExp v -> m PrimValue +evalPrimExp f (LeafExp v _) = f v +evalPrimExp _ (ValueExp v) = return v +evalPrimExp f (BinOpExp op x y) = do + x' <- evalPrimExp f x + y' <- evalPrimExp f y + maybe (evalBad op (x,y)) return $ doBinOp op x' y' +evalPrimExp f (CmpOpExp op x y) = do + x' <- evalPrimExp f x + y' <- evalPrimExp f y + maybe (evalBad op (x,y)) (return . BoolValue) $ doCmpOp op x' y' +evalPrimExp f (UnOpExp op x) = do + x' <- evalPrimExp f x + maybe (evalBad op x) return $ doUnOp op x' +evalPrimExp f (ConvOpExp op x) = do + x' <- evalPrimExp f x + maybe (evalBad op x) return $ doConvOp op x' +evalPrimExp f (FunExp h args _) = do + args' <- mapM (evalPrimExp f) args + maybe (evalBad h args) return $ do (_, _, fun) <- M.lookup h primFuns + fun args' + +evalBad :: (Pretty a, Pretty b, Monad m) => a -> b -> m c +evalBad op arg = fail $ "evalPrimExp: Type error when applying " ++ + pretty op ++ " to " ++ pretty arg + +-- | The type of values returned by a 'PrimExp'. This function +-- returning does not imply that the 'PrimExp' is type-correct. +primExpType :: PrimExp v -> PrimType +primExpType (LeafExp _ t) = t +primExpType (ValueExp v) = primValueType v +primExpType (BinOpExp op _ _) = binOpType op +primExpType CmpOpExp{} = Bool +primExpType (UnOpExp op _) = unOpType op +primExpType (ConvOpExp op _) = snd $ convOpType op +primExpType (FunExp _ _ t) = t + +-- | Is the expression a constant zero of some sort? +zeroIshExp :: PrimExp v -> Bool +zeroIshExp (ValueExp v) = zeroIsh v +zeroIshExp _ = False + +-- | Is the expression a constant one of some sort? +oneIshExp :: PrimExp v -> Bool +oneIshExp (ValueExp v) = oneIsh v +oneIshExp _ = False + +-- | Is the expression a constant value? +valueExp :: PrimExp v -> Maybe PrimValue +valueExp (ValueExp v) = Just v +valueExp _ = Nothing + +-- | If the given 'PrimExp' is a constant of the wrong integer type, +-- coerce it to the given integer type. This is a workaround for an +-- issue in the 'Num' instance. +coerceIntPrimExp :: IntType -> PrimExp v -> PrimExp v +coerceIntPrimExp t (ValueExp (IntValue v)) = ValueExp $ IntValue $ doSExt v t +coerceIntPrimExp _ e = e + +-- Prettyprinting instances + +instance Pretty v => Pretty (PrimExp v) where + ppr (LeafExp v _) = ppr v + ppr (ValueExp v) = ppr v + ppr (BinOpExp op x y) = ppr op <+> parens (ppr x) <+> parens (ppr y) + ppr (CmpOpExp op x y) = ppr op <+> parens (ppr x) <+> parens (ppr y) + ppr (ConvOpExp op x) = ppr op <+> parens (ppr x) + ppr (UnOpExp op x) = ppr op <+> parens (ppr x) + ppr (FunExp h args _) = text h <+> parens (commasep $ map ppr args) diff --git a/src/Futhark/Analysis/PrimExp/Convert.hs b/src/Futhark/Analysis/PrimExp/Convert.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/PrimExp/Convert.hs @@ -0,0 +1,109 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +-- | Converting back and forth between 'PrimExp's. +module Futhark.Analysis.PrimExp.Convert + ( + primExpToExp + , primExpFromExp + , primExpFromSubExp + , primExpFromSubExpM + , replaceInPrimExp + , substituteInPrimExp + + -- * Module reexport + , module Futhark.Analysis.PrimExp + ) where + +import qualified Control.Monad.Fail as Fail +import Data.Loc +import qualified Data.Map.Strict as M +import Data.Maybe + +import Futhark.Analysis.PrimExp +import Futhark.Construct +import Futhark.Representation.AST + +-- | Convert a 'PrimExp' to a Futhark expression. The provided +-- function converts the leaves. +primExpToExp :: MonadBinder m => + (v -> m (Exp (Lore m))) -> PrimExp v -> m (Exp (Lore m)) +primExpToExp f (BinOpExp op x y) = + BasicOp <$> (BinOp op + <$> primExpToSubExp "binop_x" f x + <*> primExpToSubExp "binop_y" f y) +primExpToExp f (CmpOpExp op x y) = + BasicOp <$> (CmpOp op + <$> primExpToSubExp "cmpop_x" f x + <*> primExpToSubExp "cmpop_y" f y) +primExpToExp f (UnOpExp op x) = + BasicOp <$> (UnOp op <$> primExpToSubExp "unop_x" f x) +primExpToExp f (ConvOpExp op x) = + BasicOp <$> (ConvOp op <$> primExpToSubExp "convop_x" f x) +primExpToExp _ (ValueExp v) = + return $ BasicOp $ SubExp $ Constant v +primExpToExp f (FunExp h args t) = + Apply (nameFromString h) <$> args' <*> pure [primRetType t] <*> pure (Safe, noLoc, []) + where args' = zip <$> mapM (primExpToSubExp "apply_arg" f) args <*> pure (repeat Observe) +primExpToExp f (LeafExp v _) = + f v + +instance ToExp v => ToExp (PrimExp v) where + toExp = primExpToExp toExp + +primExpToSubExp :: MonadBinder m => + String -> (v -> m (Exp (Lore m))) -> PrimExp v -> m SubExp +primExpToSubExp s f e = letSubExp s =<< primExpToExp f e + +-- | Convert an expression to a 'PrimExp'. The provided function is +-- used to convert expressions that are not trivially 'PrimExp's. +-- This includes constants and variable names, which are passed as +-- 'SubExp's. +primExpFromExp :: (Fail.MonadFail m, Annotations lore) => + (VName -> m (PrimExp v)) -> Exp lore -> m (PrimExp v) +primExpFromExp f (BasicOp (BinOp op x y)) = + BinOpExp op <$> primExpFromSubExpM f x <*> primExpFromSubExpM f y +primExpFromExp f (BasicOp (CmpOp op x y)) = + CmpOpExp op <$> primExpFromSubExpM f x <*> primExpFromSubExpM f y +primExpFromExp f (BasicOp (UnOp op x)) = + UnOpExp op <$> primExpFromSubExpM f x +primExpFromExp f (BasicOp (ConvOp op x)) = + ConvOpExp op <$> primExpFromSubExpM f x +primExpFromExp _ (BasicOp (SubExp (Constant v))) = + return $ ValueExp v +primExpFromExp f (Apply fname args ts _) + | isBuiltInFunction fname, [Prim t] <- retTypeValues ts = + FunExp (nameToString fname) <$> mapM (primExpFromSubExpM f . fst) args <*> pure t +primExpFromExp _ _ = fail "Not a PrimExp" + +primExpFromSubExpM :: Fail.MonadFail m => + (VName -> m (PrimExp v)) -> SubExp -> m (PrimExp v) +primExpFromSubExpM f (Var v) = f v +primExpFromSubExpM _ (Constant v) = return $ ValueExp v + +-- | Convert 'SubExp's of a given type. +primExpFromSubExp :: PrimType -> SubExp -> PrimExp VName +primExpFromSubExp t (Var v) = LeafExp v t +primExpFromSubExp _ (Constant v) = ValueExp v + +-- | Applying a transformation to the leaves in a 'PrimExp'. +replaceInPrimExp :: (v -> PrimType -> PrimExp v) -> + PrimExp v -> PrimExp v +replaceInPrimExp f (LeafExp v pt) = + f v pt +replaceInPrimExp _ (ValueExp v) = + ValueExp v +replaceInPrimExp f (BinOpExp bop pe1 pe2) = + BinOpExp bop (replaceInPrimExp f pe1) (replaceInPrimExp f pe2) +replaceInPrimExp f (CmpOpExp cop pe1 pe2) = + CmpOpExp cop (replaceInPrimExp f pe1) (replaceInPrimExp f pe2) +replaceInPrimExp f (UnOpExp uop pe) = + UnOpExp uop $ replaceInPrimExp f pe +replaceInPrimExp f (ConvOpExp cop pe) = + ConvOpExp cop $ replaceInPrimExp f pe +replaceInPrimExp f (FunExp h args t) = + FunExp h (map (replaceInPrimExp f) args) t + +-- | Substituting names in a PrimExp with other PrimExps +substituteInPrimExp :: Ord v => M.Map v (PrimExp v) + -> PrimExp v -> PrimExp v +substituteInPrimExp tab = replaceInPrimExp $ \v t -> + fromMaybe (LeafExp v t) $ M.lookup v tab diff --git a/src/Futhark/Analysis/PrimExp/Simplify.hs b/src/Futhark/Analysis/PrimExp/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/PrimExp/Simplify.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE FlexibleContexts #-} +-- | Defines simplification functions for 'PrimExp's. +module Futhark.Analysis.PrimExp.Simplify + (simplifyPrimExp, simplifyExtPrimExp) +where + +import Futhark.Analysis.PrimExp +import Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Representation.AST + +-- | Simplify a 'PrimExp', including copy propagation. If a 'LeafExp' +-- refers to a name that is a 'Constant', the node turns into a +-- 'ValueExp'. +simplifyPrimExp :: SimplifiableLore lore => + PrimExp VName -> SimpleM lore (PrimExp VName) +simplifyPrimExp = simplifyAnyPrimExp onLeaf + where onLeaf v pt = do + se <- simplify $ Var v + case se of + Var v' -> return $ LeafExp v' pt + Constant pv -> return $ ValueExp pv + +-- | Like 'simplifyPrimExp', but where leaves may be 'Ext's. +simplifyExtPrimExp :: SimplifiableLore lore => + PrimExp (Ext VName) -> SimpleM lore (PrimExp (Ext VName)) +simplifyExtPrimExp = simplifyAnyPrimExp onLeaf + where onLeaf (Free v) pt = do + se <- simplify $ Var v + case se of + Var v' -> return $ LeafExp (Free v') pt + Constant pv -> return $ ValueExp pv + onLeaf (Ext i) pt = return $ LeafExp (Ext i) pt + +simplifyAnyPrimExp :: SimplifiableLore lore => + (a -> PrimType -> SimpleM lore (PrimExp a)) + -> PrimExp a -> SimpleM lore (PrimExp a) +simplifyAnyPrimExp f (LeafExp v pt) = f v pt +simplifyAnyPrimExp _ (ValueExp pv) = + return $ ValueExp pv +simplifyAnyPrimExp f (BinOpExp bop e1 e2) = + BinOpExp bop <$> simplifyAnyPrimExp f e1 <*> simplifyAnyPrimExp f e2 +simplifyAnyPrimExp f (CmpOpExp cmp e1 e2) = + CmpOpExp cmp <$> simplifyAnyPrimExp f e1 <*> simplifyAnyPrimExp f e2 +simplifyAnyPrimExp f (UnOpExp op e) = + UnOpExp op <$> simplifyAnyPrimExp f e +simplifyAnyPrimExp f (ConvOpExp conv e) = + ConvOpExp conv <$> simplifyAnyPrimExp f e +simplifyAnyPrimExp f (FunExp h args t) = + FunExp h <$> mapM (simplifyAnyPrimExp f) args <*> pure t diff --git a/src/Futhark/Analysis/Range.hs b/src/Futhark/Analysis/Range.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/Range.hs @@ -0,0 +1,214 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Analysis.Range + ( rangeAnalysis + , runRangeM + , RangeM + , analyseExp + , analyseLambda + , analyseBody + , analyseStms + ) + where + +import qualified Data.Map.Strict as M +import Control.Monad.Reader +import Data.Semigroup ((<>)) +import Data.List + +import qualified Futhark.Analysis.ScalExp as SE +import Futhark.Representation.Ranges +import Futhark.Analysis.AlgSimplify as AS + +-- Entry point + +-- | Perform variable range analysis on the given program, returning a +-- program with embedded range annotations. +rangeAnalysis :: (Attributes lore, CanBeRanged (Op lore)) => + Prog lore -> Prog (Ranges lore) +rangeAnalysis = Prog . map analyseFun . progFunctions + +-- Implementation + +analyseFun :: (Attributes lore, CanBeRanged (Op lore)) => + FunDef lore -> FunDef (Ranges lore) +analyseFun (FunDef entry fname restype params body) = + runRangeM $ bindFunParams params $ + FunDef entry fname restype params <$> analyseBody body + +analyseBody :: (Attributes lore, CanBeRanged (Op lore)) => + Body lore + -> RangeM (Body (Ranges lore)) +analyseBody (Body lore origbnds result) = + analyseStms origbnds $ \bnds' -> + return $ mkRangedBody lore bnds' result + +analyseStms :: (Attributes lore, CanBeRanged (Op lore)) => + Stms lore + -> (Stms (Ranges lore) -> RangeM a) + -> RangeM a +analyseStms = analyseStms' mempty . stmsToList + where analyseStms' acc [] m = + m acc + analyseStms' acc (bnd:bnds) m = do + bnd' <- analyseStm bnd + bindPattern (stmPattern bnd') $ + analyseStms' (acc <> oneStm bnd') bnds m + +analyseStm :: (Attributes lore, CanBeRanged (Op lore)) => + Stm lore -> RangeM (Stm (Ranges lore)) +analyseStm (Let pat lore e) = do + e' <- analyseExp e + pat' <- simplifyPatRanges $ addRangesToPattern pat e' + return $ Let pat' lore e' + +analyseExp :: (Attributes lore, CanBeRanged (Op lore)) => + Exp lore + -> RangeM (Exp (Ranges lore)) +analyseExp = mapExpM analyse + where analyse = + Mapper { mapOnSubExp = return + , mapOnCertificates = return + , mapOnVName = return + , mapOnBody = const analyseBody + , mapOnRetType = return + , mapOnBranchType = return + , mapOnFParam = return + , mapOnLParam = return + , mapOnOp = return . addOpRanges + } + +analyseLambda :: (Attributes lore, CanBeRanged (Op lore)) => + Lambda lore + -> RangeM (Lambda (Ranges lore)) +analyseLambda lam = do + body <- analyseBody $ lambdaBody lam + return $ lam { lambdaBody = body + , lambdaParams = lambdaParams lam + } + +-- Monad and utility definitions + +type RangeEnv = M.Map VName Range + +emptyRangeEnv :: RangeEnv +emptyRangeEnv = M.empty + +type RangeM = Reader RangeEnv + +runRangeM :: RangeM a -> a +runRangeM = flip runReader emptyRangeEnv + +bindFunParams :: Typed attr => + [ParamT attr] -> RangeM a -> RangeM a +bindFunParams [] m = + m +bindFunParams (param:params) m = do + ranges <- rangesRep + local bindFunParam $ + local (refineDimensionRanges ranges dims) $ + bindFunParams params m + where bindFunParam = M.insert (paramName param) unknownRange + dims = arrayDims $ paramType param + +bindPattern :: Typed attr => + PatternT (Range, attr) -> RangeM a -> RangeM a +bindPattern pat m = do + ranges <- rangesRep + local bindPatElems $ + local (refineDimensionRanges ranges dims) + m + where bindPatElems env = + foldl bindPatElem env $ patternElements pat + bindPatElem env patElem = + M.insert (patElemName patElem) (fst $ patElemAttr patElem) env + dims = nub $ concatMap arrayDims $ patternTypes pat + +refineDimensionRanges :: AS.RangesRep -> [SubExp] + -> RangeEnv -> RangeEnv +refineDimensionRanges ranges = flip $ foldl refineShape + where refineShape env (Var dim) = + refineRange ranges dim dimBound env + refineShape env _ = + env + -- A dimension is never negative. + dimBound :: Range + dimBound = (Just $ ScalarBound 0, + Nothing) + +refineRange :: AS.RangesRep -> VName -> Range -> RangeEnv + -> RangeEnv +refineRange = + M.insertWith . refinedRange + +-- New range, old range, result range. +refinedRange :: AS.RangesRep -> Range -> Range -> Range +refinedRange ranges (new_lower, new_upper) (old_lower, old_upper) = + (simplifyBound ranges $ refineLowerBound new_lower old_lower, + simplifyBound ranges $ refineUpperBound new_upper old_upper) + +-- New bound, old bound, result bound. +refineLowerBound :: Bound -> Bound -> Bound +refineLowerBound = flip maximumBound + +-- New bound, old bound, result bound. +refineUpperBound :: Bound -> Bound -> Bound +refineUpperBound = flip minimumBound + +lookupRange :: VName -> RangeM Range +lookupRange = asks . M.findWithDefault unknownRange + +simplifyPatRanges :: PatternT (Range, attr) + -> RangeM (PatternT (Range, attr)) +simplifyPatRanges (Pattern context values) = + Pattern <$> mapM simplifyPatElemRange context <*> mapM simplifyPatElemRange values + where simplifyPatElemRange patElem = do + let (range, innerattr) = patElemAttr patElem + range' <- simplifyRange range + return $ setPatElemLore patElem (range', innerattr) + +simplifyRange :: Range -> RangeM Range +simplifyRange (lower, upper) = do + ranges <- rangesRep + lower' <- simplifyBound ranges <$> betterLowerBound lower + upper' <- simplifyBound ranges <$> betterUpperBound upper + return (lower', upper') + +simplifyBound :: AS.RangesRep -> Bound -> Bound +simplifyBound ranges = fmap $ simplifyKnownBound ranges + +simplifyKnownBound :: AS.RangesRep -> KnownBound -> KnownBound +simplifyKnownBound ranges bound + | Just se <- boundToScalExp bound = + ScalarBound $ AS.simplify se ranges +simplifyKnownBound ranges (MinimumBound b1 b2) = + MinimumBound (simplifyKnownBound ranges b1) (simplifyKnownBound ranges b2) +simplifyKnownBound ranges (MaximumBound b1 b2) = + MaximumBound (simplifyKnownBound ranges b1) (simplifyKnownBound ranges b2) +simplifyKnownBound _ bound = + bound + +betterLowerBound :: Bound -> RangeM Bound +betterLowerBound (Just (ScalarBound (SE.Id v t))) = do + range <- lookupRange v + return $ Just $ case range of + (Just lower, _) -> lower + _ -> ScalarBound $ SE.Id v t +betterLowerBound bound = + return bound + +betterUpperBound :: Bound -> RangeM Bound +betterUpperBound (Just (ScalarBound (SE.Id v t))) = do + range <- lookupRange v + return $ Just $ case range of + (_, Just upper) -> upper + _ -> ScalarBound $ SE.Id v t +betterUpperBound bound = + return bound + +-- The algebraic simplifier requires a loop nesting level for each +-- range. We just put a zero because I don't think it's used for +-- anything in this case. +rangesRep :: RangeM AS.RangesRep +rangesRep = asks $ M.map addLeadingZero + where addLeadingZero (x,y) = + (0, boundToScalExp =<< x, boundToScalExp =<< y) diff --git a/src/Futhark/Analysis/Rephrase.hs b/src/Futhark/Analysis/Rephrase.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/Rephrase.hs @@ -0,0 +1,107 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Facilities for changing the lore of some fragment, with no context. +module Futhark.Analysis.Rephrase + ( rephraseProg + , rephraseFunDef + , rephraseExp + , rephraseBody + , rephraseStm + , rephraseLambda + , rephrasePattern + , rephrasePatElem + , Rephraser (..) + + , castStm + ) +where + +import Futhark.Representation.AST + +data Rephraser m from to + = Rephraser { rephraseExpLore :: ExpAttr from -> m (ExpAttr to) + , rephraseLetBoundLore :: LetAttr from -> m (LetAttr to) + , rephraseFParamLore :: FParamAttr from -> m (FParamAttr to) + , rephraseLParamLore :: LParamAttr from -> m (LParamAttr to) + , rephraseBodyLore :: BodyAttr from -> m (BodyAttr to) + , rephraseRetType :: RetType from -> m (RetType to) + , rephraseBranchType :: BranchType from -> m (BranchType to) + , rephraseOp :: Op from -> m (Op to) + } + +rephraseProg :: Monad m => Rephraser m from to -> Prog from -> m (Prog to) +rephraseProg rephraser = fmap Prog . mapM (rephraseFunDef rephraser) . progFunctions + +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 + rettype' <- mapM (rephraseRetType rephraser) $ funDefRetType fundec + return fundec { funDefBody = body', funDefParams = params', funDefRetType = rettype' } + +rephraseExp :: Monad m => Rephraser m from to -> Exp from -> m (Exp to) +rephraseExp = mapExpM . mapper + +rephraseStm :: Monad m => Rephraser m from to -> Stm from -> m (Stm to) +rephraseStm rephraser (Let pat (StmAux cs attr) e) = + Let <$> + rephrasePattern (rephraseLetBoundLore rephraser) pat <*> + (StmAux cs <$> rephraseExpLore rephraser attr) <*> + rephraseExp rephraser e + +rephrasePattern :: Monad m => + (from -> m to) + -> PatternT from + -> m (PatternT to) +rephrasePattern f (Pattern context values) = + Pattern <$> rephrase context <*> rephrase values + where rephrase = mapM $ rephrasePatElem f + +rephrasePatElem :: Monad m => (from -> m to) -> PatElemT from -> m (PatElemT to) +rephrasePatElem rephraser (PatElem ident from) = + PatElem ident <$> rephraser from + +rephraseParam :: Monad m => (from -> m to) -> ParamT from -> m (ParamT to) +rephraseParam rephraser (Param name from) = + Param name <$> rephraser from + +rephraseBody :: Monad m => Rephraser m from to -> Body from -> m (Body to) +rephraseBody rephraser (Body lore bnds res) = + Body <$> + rephraseBodyLore rephraser lore <*> + (stmsFromList <$> mapM (rephraseStm rephraser) (stmsToList bnds)) <*> + pure res + +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 + return lam { lambdaBody = body', lambdaParams = params' } + +mapper :: Monad m => Rephraser m from to -> Mapper from to m +mapper rephraser = identityMapper { + mapOnBody = const $ rephraseBody rephraser + , mapOnRetType = rephraseRetType rephraser + , mapOnBranchType = rephraseBranchType rephraser + , mapOnFParam = rephraseParam (rephraseFParamLore rephraser) + , mapOnLParam = rephraseParam (rephraseLParamLore rephraser) + , mapOnOp = rephraseOp rephraser + } + +-- | Convert a binding from one lore to another, if possible. +castStm :: (SameScope from to, + ExpAttr from ~ ExpAttr to, + BodyAttr from ~ BodyAttr to, + RetType from ~ RetType to, + BranchType from ~ BranchType to) => + Stm from -> Maybe (Stm to) +castStm = rephraseStm caster + where caster = Rephraser { rephraseExpLore = Just + , rephraseBodyLore = Just + , rephraseLetBoundLore = Just + , rephraseFParamLore = Just + , rephraseLParamLore = Just + , rephraseOp = const Nothing + , rephraseRetType = Just + , rephraseBranchType = Just + } diff --git a/src/Futhark/Analysis/ScalExp.hs b/src/Futhark/Analysis/ScalExp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/ScalExp.hs @@ -0,0 +1,308 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Analysis.ScalExp + ( RelOp0(..) + , ScalExp(..) + , scalExpType + , scalExpSize + , subExpToScalExp + , toScalExp + , expandScalExp + , LookupVar + , module Futhark.Representation.Primitive + ) +where + +import Data.List +import qualified Data.Set as S +import Data.Maybe +import Data.Monoid ((<>)) + +import Futhark.Representation.Primitive hiding (SQuot, SRem, SDiv, SMod, SSignum) +import Futhark.Representation.AST hiding (SQuot, SRem, SDiv, SMod, SSignum) +import qualified Futhark.Representation.AST as AST +import Futhark.Transform.Substitute +import Futhark.Transform.Rename +import Futhark.Util.Pretty hiding (pretty) + +----------------------------------------------------------------- +-- BINARY OPERATORS for Numbers -- +-- Note that MOD, BAND, XOR, BOR, SHIFTR, SHIFTL not supported -- +-- `a SHIFTL/SHIFTR p' can be translated if desired as as -- +-- `a * 2^p' or `a / 2^p -- +----------------------------------------------------------------- + +-- | Relational operators. +data RelOp0 = LTH0 + | LEQ0 + deriving (Eq, Ord, Enum, Bounded, Show) + +-- | Representation of a scalar expression, which is: +-- +-- (i) an algebraic expression, e.g., min(a+b, a*b), +-- +-- (ii) a relational expression: a+b < 5, +-- +-- (iii) a logical expression: e1 and (not (a+b>5) +data ScalExp= Val PrimValue + | Id VName PrimType + | SNeg ScalExp + | SNot ScalExp + | SAbs ScalExp + | SSignum ScalExp + | SPlus ScalExp ScalExp + | SMinus ScalExp ScalExp + | STimes ScalExp ScalExp + | SPow ScalExp ScalExp + | SDiv ScalExp ScalExp + | SMod ScalExp ScalExp + | SQuot ScalExp ScalExp + | SRem ScalExp ScalExp + | MaxMin Bool [ScalExp] + | RelExp RelOp0 ScalExp + | SLogAnd ScalExp ScalExp + | SLogOr ScalExp ScalExp + deriving (Eq, Ord, Show) + +instance Num ScalExp where + 0 + y = y + x + 0 = x + x + y = SPlus x y + + x - 0 = x + x - y = SMinus x y + + 0 * _ = 0 + _ * 0 = 0 + 1 * y = y + y * 1 = y + x * y = STimes x y + + abs = SAbs + signum = SSignum + fromInteger = Val . IntValue . Int32Value . fromInteger -- probably not OK + negate = SNeg + +instance Pretty ScalExp where + pprPrec _ (Val val) = ppr val + pprPrec _ (Id v _) = ppr v + pprPrec _ (SNeg e) = text "-" <> pprPrec 9 e + pprPrec _ (SNot e) = text "not" <+> pprPrec 9 e + pprPrec _ (SAbs e) = text "abs" <+> pprPrec 9 e + pprPrec _ (SSignum e) = text "signum" <+> pprPrec 9 e + pprPrec prec (SPlus x y) = ppBinOp prec "+" 4 4 x y + pprPrec prec (SMinus x y) = ppBinOp prec "-" 4 10 x y + pprPrec prec (SPow x y) = ppBinOp prec "^" 6 6 x y + pprPrec prec (STimes x y) = ppBinOp prec "*" 5 5 x y + pprPrec prec (SDiv x y) = ppBinOp prec "/" 5 10 x y + pprPrec prec (SMod x y) = ppBinOp prec "%" 5 10 x y + pprPrec prec (SQuot x y) = ppBinOp prec "//" 5 10 x y + pprPrec prec (SRem x y) = ppBinOp prec "%%" 5 10 x y + pprPrec prec (SLogOr x y) = ppBinOp prec "||" 0 0 x y + pprPrec prec (SLogAnd x y) = ppBinOp prec "&&" 1 1 x y + pprPrec prec (RelExp LTH0 e) = ppBinOp prec "<" 2 2 e (0::Int) + pprPrec prec (RelExp LEQ0 e) = ppBinOp prec "<=" 2 2 e (0::Int) + pprPrec _ (MaxMin True es) = text "min" <> parens (commasep $ map ppr es) + pprPrec _ (MaxMin False es) = text "max" <> parens (commasep $ map ppr es) + +ppBinOp :: (Pretty a, Pretty b) => Int -> String -> Int -> Int -> a -> b -> Doc +ppBinOp p bop precedence rprecedence x y = + parensIf (p > precedence) $ + pprPrec precedence x <+/> + text bop <+> + pprPrec rprecedence y + +instance Substitute ScalExp where + substituteNames subst e = + case e of Id v t -> Id (substituteNames subst v) t + Val v -> Val v + SNeg x -> SNeg $ substituteNames subst x + SNot x -> SNot $ substituteNames subst x + SAbs x -> SAbs $ substituteNames subst x + SSignum x -> SSignum $ substituteNames subst x + SPlus x y -> substituteNames subst x `SPlus` substituteNames subst y + SMinus x y -> substituteNames subst x `SMinus` substituteNames subst y + SPow x y -> substituteNames subst x `SPow` substituteNames subst y + STimes x y -> substituteNames subst x `STimes` substituteNames subst y + SDiv x y -> substituteNames subst x `SDiv` substituteNames subst y + SMod x y -> substituteNames subst x `SMod` substituteNames subst y + SQuot x y -> substituteNames subst x `SDiv` substituteNames subst y + SRem x y -> substituteNames subst x `SRem` substituteNames subst y + MaxMin m es -> MaxMin m $ map (substituteNames subst) es + RelExp r x -> RelExp r $ substituteNames subst x + SLogAnd x y -> substituteNames subst x `SLogAnd` substituteNames subst y + SLogOr x y -> substituteNames subst x `SLogOr` substituteNames subst y + +instance Rename ScalExp where + rename = substituteRename + +scalExpType :: ScalExp -> PrimType +scalExpType (Val v) = primValueType v +scalExpType (Id _ t) = t +scalExpType (SNeg e) = scalExpType e +scalExpType (SNot _) = Bool +scalExpType (SAbs e) = scalExpType e +scalExpType (SSignum e) = scalExpType e +scalExpType (SPlus e _) = scalExpType e +scalExpType (SMinus e _) = scalExpType e +scalExpType (STimes e _) = scalExpType e +scalExpType (SDiv e _) = scalExpType e +scalExpType (SMod e _) = scalExpType e +scalExpType (SPow e _) = scalExpType e +scalExpType (SQuot e _) = scalExpType e +scalExpType (SRem e _) = scalExpType e +scalExpType (SLogAnd _ _) = Bool +scalExpType (SLogOr _ _) = Bool +scalExpType (RelExp _ _) = Bool +scalExpType (MaxMin _ []) = IntType Int32 -- arbitrary and probably wrong. +scalExpType (MaxMin _ (e:_)) = scalExpType e + +-- | Number of nodes in the scalar expression. +scalExpSize :: ScalExp -> Int +scalExpSize Val{} = 1 +scalExpSize Id{} = 1 +scalExpSize (SNeg e) = scalExpSize e +scalExpSize (SNot e) = scalExpSize e +scalExpSize (SAbs e) = scalExpSize e +scalExpSize (SSignum e) = scalExpSize e +scalExpSize (SPlus x y) = scalExpSize x + scalExpSize y +scalExpSize (SMinus x y) = scalExpSize x + scalExpSize y +scalExpSize (STimes x y) = scalExpSize x + scalExpSize y +scalExpSize (SDiv x y) = scalExpSize x + scalExpSize y +scalExpSize (SMod x y) = scalExpSize x + scalExpSize y +scalExpSize (SPow x y) = scalExpSize x + scalExpSize y +scalExpSize (SQuot x y) = scalExpSize x + scalExpSize y +scalExpSize (SRem x y) = scalExpSize x + scalExpSize y +scalExpSize (SLogAnd x y) = scalExpSize x + scalExpSize y +scalExpSize (SLogOr x y) = scalExpSize x + scalExpSize y +scalExpSize (RelExp _ x) = scalExpSize x +scalExpSize (MaxMin _ []) = 0 +scalExpSize (MaxMin _ es) = sum $ map scalExpSize es + +-- | A function that checks whether a variable name corresponds to a +-- scalar expression. +type LookupVar = VName -> Maybe ScalExp + +-- | Non-recursively convert a subexpression to a 'ScalExp'. The +-- (scalar) type of the subexpression must be given in advance. +subExpToScalExp :: SubExp -> PrimType -> ScalExp +subExpToScalExp (Var v) t = Id v t +subExpToScalExp (Constant val) _ = Val val + +toScalExp :: (HasScope t f, Monad f) => + LookupVar -> Exp lore -> f (Maybe ScalExp) +toScalExp look (BasicOp (SubExp (Var v))) + | Just se <- look v = + return $ Just se + | otherwise = do + t <- lookupType v + case t of + Prim bt | typeIsOK bt -> + return $ Just $ Id v bt + _ -> + return Nothing +toScalExp _ (BasicOp (SubExp (Constant val))) + | typeIsOK $ primValueType val = + return $ Just $ Val val +toScalExp look (BasicOp (CmpOp (CmpSlt _) x y)) = + Just . RelExp LTH0 <$> (sminus <$> subExpToScalExp' look x <*> subExpToScalExp' look y) +toScalExp look (BasicOp (CmpOp (CmpSle _) x y)) = + Just . RelExp LEQ0 <$> (sminus <$> subExpToScalExp' look x <*> subExpToScalExp' look y) +toScalExp look (BasicOp (CmpOp (CmpEq t) x y)) + | typeIsOK t = do + x' <- subExpToScalExp' look x + y' <- subExpToScalExp' look y + return $ Just $ case t of + Bool -> + SLogAnd x' y' `SLogOr` SLogAnd (SNot x') (SNot y') + _ -> + RelExp LEQ0 (x' `sminus` y') `SLogAnd` RelExp LEQ0 (y' `sminus` x') +toScalExp look (BasicOp (BinOp (Sub t) (Constant x) y)) + | typeIsOK $ IntType t, zeroIsh x = + Just . SNeg <$> subExpToScalExp' look y +toScalExp look (BasicOp (UnOp AST.Not e)) = + Just . SNot <$> subExpToScalExp' look e +toScalExp look (BasicOp (BinOp bop x y)) + | Just f <- binOpScalExp bop = + Just <$> (f <$> subExpToScalExp' look x <*> subExpToScalExp' look y) + +toScalExp _ _ = return Nothing + +typeIsOK :: PrimType -> Bool +typeIsOK = (`elem` Bool : map IntType allIntTypes) + +subExpToScalExp' :: HasScope t f => + LookupVar -> SubExp -> f ScalExp +subExpToScalExp' look (Var v) + | Just se <- look v = + pure se + | otherwise = + withType <$> lookupType v + where withType (Prim t) = + subExpToScalExp (Var v) t + withType t = + error $ "Cannot create ScalExp from variable " ++ pretty v ++ + " of type " ++ pretty t +subExpToScalExp' _ (Constant val) = + pure $ Val val + +-- | If you have a scalar expression that has been created with +-- incomplete symbol table information, you can use this function to +-- grow its 'Id' leaves. +expandScalExp :: LookupVar -> ScalExp -> ScalExp +expandScalExp _ (Val v) = Val v +expandScalExp look (Id v t) = fromMaybe (Id v t) $ look v +expandScalExp look (SNeg se) = SNeg $ expandScalExp look se +expandScalExp look (SNot se) = SNot $ expandScalExp look se +expandScalExp look (SAbs se) = SAbs $ expandScalExp look se +expandScalExp look (SSignum se) = SSignum $ expandScalExp look se +expandScalExp look (MaxMin b ses) = MaxMin b $ map (expandScalExp look) ses +expandScalExp look (SPlus x y) = SPlus (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SMinus x y) = SMinus (expandScalExp look x) (expandScalExp look y) +expandScalExp look (STimes x y) = STimes (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SDiv x y) = SDiv (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SMod x y) = SMod (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SQuot x y) = SQuot (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SRem x y) = SRem (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SPow x y) = SPow (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SLogAnd x y) = SLogAnd (expandScalExp look x) (expandScalExp look y) +expandScalExp look (SLogOr x y) = SLogOr (expandScalExp look x) (expandScalExp look y) +expandScalExp look (RelExp relop x) = RelExp relop $ expandScalExp look x + +-- | "Smart constructor" that checks whether we are subtracting zero, +-- and if so just returns the first argument. +sminus :: ScalExp -> ScalExp -> ScalExp +sminus x (Val v) | zeroIsh v = x +sminus x y = x `SMinus` y + + -- XXX: Only integers and booleans, OK? +binOpScalExp :: BinOp -> Maybe (ScalExp -> ScalExp -> ScalExp) +binOpScalExp bop = fmap snd . find ((==bop) . fst) $ + concatMap intOps allIntTypes ++ + [ (LogAnd, SLogAnd), (LogOr, SLogOr) ] + where intOps t = [ (Add t, SPlus) + , (Sub t, SMinus) + , (Mul t, STimes) + , (AST.SDiv t, SDiv) + , (AST.Pow t, SPow) + ] + +instance FreeIn ScalExp where + freeIn (Val _) = mempty + freeIn (Id i _) = S.singleton i + freeIn (SNeg e) = freeIn e + freeIn (SNot e) = freeIn e + freeIn (SAbs e) = freeIn e + freeIn (SSignum e) = freeIn e + freeIn (SPlus x y) = freeIn x <> freeIn y + freeIn (SMinus x y) = freeIn x <> freeIn y + freeIn (SPow x y) = freeIn x <> freeIn y + freeIn (STimes x y) = freeIn x <> freeIn y + freeIn (SDiv x y) = freeIn x <> freeIn y + freeIn (SMod x y) = freeIn x <> freeIn y + freeIn (SQuot x y) = freeIn x <> freeIn y + freeIn (SRem x y) = freeIn x <> freeIn y + freeIn (SLogOr x y) = freeIn x <> freeIn y + freeIn (SLogAnd x y) = freeIn x <> freeIn y + freeIn (RelExp LTH0 e) = freeIn e + freeIn (RelExp LEQ0 e) = freeIn e + freeIn (MaxMin _ es) = mconcat $ map freeIn es diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -0,0 +1,763 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE TypeFamilies #-} +module Futhark.Analysis.SymbolTable + ( SymbolTable (bindings, loopDepth, availableAtClosestLoop) + , empty + , fromScope + , toScope + , castSymbolTable + -- * Entries + , Entry + , deepen + , bindingDepth + , valueRange + , loopVariable + , entryStm + , entryLetBoundAttr + , entryFParamLore + , entryType + , asScalExp + -- * Lookup + , elem + , lookup + , lookupStm + , lookupExp + , lookupBasicOp + , lookupType + , lookupSubExp + , lookupScalExp + , lookupValue + , lookupVar + , lookupAliases + , index + , index' + , IndexOp(..) + -- * Insertion + , insertStm + , insertFParams + , insertLParam + , insertArrayLParam + , insertChunkLParam + , insertLoopVar + -- * Bounds + , updateBounds + , setUpperBound + , setLowerBound + , isAtLeast + -- * Misc + , enclosingLoopVars + , rangesRep + ) + where + +import Control.Arrow (second, (&&&)) +import Control.Monad +import Control.Monad.Reader +import Data.Ord +import Data.Maybe +import Data.Semigroup ((<>)) +import Data.List hiding (elem, lookup) +import qualified Data.List as L +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem + +import Prelude hiding (elem, lookup) + +import Futhark.Analysis.PrimExp.Convert +import Futhark.Representation.AST hiding (FParam, ParamT (..), lookupType) +import qualified Futhark.Representation.AST as AST +import Futhark.Analysis.ScalExp + +import Futhark.Analysis.Rephrase +import qualified Futhark.Analysis.AlgSimplify as AS +import Futhark.Representation.AST.Attributes.Ranges + (Range, ScalExpRange, Ranged) +import qualified Futhark.Representation.AST.Attributes.Ranges as Ranges +import qualified Futhark.Representation.AST.Attributes.Aliases as Aliases + +data SymbolTable lore = SymbolTable { + loopDepth :: Int + , bindings :: M.Map VName (Entry lore) + , availableAtClosestLoop :: Names + -- ^ Which names are available just before the most enclosing + -- loop? + } + +instance Sem.Semigroup (SymbolTable lore) where + table1 <> table2 = + SymbolTable { loopDepth = max (loopDepth table1) (loopDepth table2) + , bindings = bindings table1 <> bindings table2 + , availableAtClosestLoop = availableAtClosestLoop table1 <> + availableAtClosestLoop table2 + } + +instance Monoid (SymbolTable lore) where + mempty = empty + mappend = (Sem.<>) + +empty :: SymbolTable lore +empty = SymbolTable 0 M.empty mempty + +fromScope :: Attributes lore => Scope lore -> SymbolTable lore +fromScope = M.foldlWithKey' insertFreeVar' empty + where insertFreeVar' m k attr = insertFreeVar k attr m + +toScope :: SymbolTable lore -> Scope lore +toScope = M.map entryInfo . bindings + +-- | Try to convert a symbol table for one representation into a +-- symbol table for another. The two symbol tables will have the same +-- keys, but some entries may be diferent (i.e. some expression +-- entries will have been turned into free variable entries). +castSymbolTable :: (SameScope from to, + ExpAttr from ~ ExpAttr to, + BodyAttr from ~ BodyAttr to, + RetType from ~ RetType to, + BranchType from ~ BranchType to) => + SymbolTable from -> SymbolTable to +castSymbolTable table = genCastSymbolTable loopVar letBound fParam lParam freeVar table + where loopVar (LoopVarEntry r d it) = LoopVar $ LoopVarEntry r d it + letBound e + | Just e' <- castStm $ letBoundStm e = + LetBound e { letBoundStm = e' + , letBoundAttr = letBoundAttr e + } + | otherwise = + FreeVar FreeVarEntry + { freeVarAttr = LetInfo $ letBoundAttr e + , freeVarStmDepth = letBoundStmDepth e + , freeVarRange = letBoundRange e + , freeVarIndex = \name is -> index' name is table + } + + fParam e = FParam e { fparamAttr = fparamAttr e } + lParam e = LParam e { lparamAttr = lparamAttr e } + freeVar e = FreeVar e { freeVarAttr = castNameInfo $ freeVarAttr e } + +genCastSymbolTable :: (LoopVarEntry fromlore -> Entry tolore) + -> (LetBoundEntry fromlore -> Entry tolore) + -> (FParamEntry fromlore -> Entry tolore) + -> (LParamEntry fromlore -> Entry tolore) + -> (FreeVarEntry fromlore -> Entry tolore) + -> SymbolTable fromlore + -> SymbolTable tolore +genCastSymbolTable loopVar letBound fParam lParam freeVar (SymbolTable depth entries loopfree) = + SymbolTable depth (M.map onEntry entries) loopfree + where onEntry (LoopVar entry) = loopVar entry + onEntry (LetBound entry) = letBound entry + onEntry (FParam entry) = fParam entry + onEntry (LParam entry) = lParam entry + onEntry (FreeVar entry) = freeVar entry + +deepen :: SymbolTable lore -> SymbolTable lore +deepen vtable = vtable { loopDepth = loopDepth vtable + 1, + availableAtClosestLoop = S.fromList $ M.keys $ bindings vtable + } + +-- | Indexing a delayed array if possible. +type IndexArray = [PrimExp VName] -> Maybe (PrimExp VName, Certificates) + +data Entry lore = LoopVar (LoopVarEntry lore) + | LetBound (LetBoundEntry lore) + | FParam (FParamEntry lore) + | LParam (LParamEntry lore) + | FreeVar (FreeVarEntry lore) + +data LoopVarEntry lore = + LoopVarEntry { loopVarRange :: ScalExpRange + , loopVarStmDepth :: Int + , loopVarType :: IntType + } + +data LetBoundEntry lore = + LetBoundEntry { letBoundRange :: ScalExpRange + , letBoundAttr :: LetAttr lore + , letBoundAliases :: Names + , letBoundStm :: Stm lore + , letBoundStmDepth :: Int + , letBoundScalExp :: Maybe ScalExp + , letBoundIndex :: Int -> IndexArray + -- ^ Index a delayed array, if possible. + } + +data FParamEntry lore = + FParamEntry { fparamRange :: ScalExpRange + , fparamAttr :: FParamAttr lore + , fparamAliases :: Names + , fparamStmDepth :: Int + } + +data LParamEntry lore = + LParamEntry { lparamRange :: ScalExpRange + , lparamAttr :: LParamAttr lore + , lparamStmDepth :: Int + , lparamIndex :: IndexArray + } + +data FreeVarEntry lore = + FreeVarEntry { freeVarAttr :: NameInfo lore + , freeVarStmDepth :: Int + , freeVarRange :: ScalExpRange + , freeVarIndex :: VName -> IndexArray + -- ^ Index a delayed array, if possible. + } + +entryInfo :: Entry lore -> NameInfo lore +entryInfo (LetBound entry) = LetInfo $ letBoundAttr entry +entryInfo (LoopVar entry) = IndexInfo $ loopVarType entry +entryInfo (FParam entry) = FParamInfo $ fparamAttr entry +entryInfo (LParam entry) = LParamInfo $ lparamAttr entry +entryInfo (FreeVar entry) = freeVarAttr entry + +entryType :: Attributes lore => Entry lore -> Type +entryType = typeOf . entryInfo + +isVarBound :: Entry lore -> Maybe (LetBoundEntry lore) +isVarBound (LetBound entry) = Just entry +isVarBound _ = Nothing + +asScalExp :: Entry lore -> Maybe ScalExp +asScalExp = letBoundScalExp <=< isVarBound + +bindingDepth :: Entry lore -> Int +bindingDepth (LetBound entry) = letBoundStmDepth entry +bindingDepth (FParam entry) = fparamStmDepth entry +bindingDepth (LParam entry) = lparamStmDepth entry +bindingDepth (LoopVar entry) = loopVarStmDepth entry +bindingDepth (FreeVar _) = 0 + +setStmDepth :: Int -> Entry lore -> Entry lore +setStmDepth d (LetBound entry) = + LetBound $ entry { letBoundStmDepth = d } +setStmDepth d (FParam entry) = + FParam $ entry { fparamStmDepth = d } +setStmDepth d (LParam entry) = + LParam $ entry { lparamStmDepth = d } +setStmDepth d (LoopVar entry) = + LoopVar $ entry { loopVarStmDepth = d } +setStmDepth d (FreeVar entry) = + FreeVar $ entry { freeVarStmDepth = d } + +valueRange :: Entry lore -> ScalExpRange +valueRange (LetBound entry) = letBoundRange entry +valueRange (FParam entry) = fparamRange entry +valueRange (LParam entry) = lparamRange entry +valueRange (LoopVar entry) = loopVarRange entry +valueRange (FreeVar entry) = freeVarRange entry + +setValueRange :: ScalExpRange -> Entry lore -> Entry lore +setValueRange range (LetBound entry) = + LetBound $ entry { letBoundRange = range } +setValueRange range (FParam entry) = + FParam $ entry { fparamRange = range } +setValueRange range (LParam entry) = + LParam $ entry { lparamRange = range } +setValueRange range (LoopVar entry) = + LoopVar $ entry { loopVarRange = range } +setValueRange range (FreeVar entry) = + FreeVar $ entry { freeVarRange = range } + +entryStm :: Entry lore -> Maybe (Stm lore) +entryStm (LetBound entry) = Just $ letBoundStm entry +entryStm _ = Nothing + +entryLetBoundAttr :: Entry lore -> Maybe (LetAttr lore) +entryLetBoundAttr (LetBound entry) = Just $ letBoundAttr entry +entryLetBoundAttr _ = Nothing + +entryFParamLore :: Entry lore -> Maybe (FParamAttr lore) +entryFParamLore (FParam entry) = Just $ fparamAttr entry +entryFParamLore _ = Nothing + +loopVariable :: Entry lore -> Bool +loopVariable (LoopVar _) = True +loopVariable _ = False + +asStm :: Entry lore -> Maybe (Stm lore) +asStm = fmap letBoundStm . isVarBound + +elem :: VName -> SymbolTable lore -> Bool +elem name = isJust . lookup name + +lookup :: VName -> SymbolTable lore -> Maybe (Entry lore) +lookup name = M.lookup name . bindings + +lookupStm :: VName -> SymbolTable lore -> Maybe (Stm lore) +lookupStm name vtable = asStm =<< lookup name vtable + +lookupExp :: VName -> SymbolTable lore -> Maybe (Exp lore, Certificates) +lookupExp name vtable = (stmExp &&& stmCerts) <$> lookupStm name vtable + +lookupBasicOp :: VName -> SymbolTable lore -> Maybe (BasicOp lore, Certificates) +lookupBasicOp name vtable = case lookupExp name vtable of + Just (BasicOp e, cs) -> Just (e, cs) + _ -> Nothing + +lookupType :: Attributes lore => VName -> SymbolTable lore -> Maybe Type +lookupType name vtable = entryType <$> lookup name vtable + +lookupSubExpType :: Attributes lore => SubExp -> SymbolTable lore -> Maybe Type +lookupSubExpType (Var v) = lookupType v +lookupSubExpType (Constant v) = const $ Just $ Prim $ primValueType v + +lookupSubExp :: VName -> SymbolTable lore -> Maybe (SubExp, Certificates) +lookupSubExp name vtable = do + (e,cs) <- lookupExp name vtable + case e of + BasicOp (SubExp se) -> Just (se,cs) + _ -> Nothing + +lookupScalExp :: Attributes lore => VName -> SymbolTable lore -> Maybe ScalExp +lookupScalExp name vtable = + case (lookup name vtable, lookupRange name vtable) of + -- If we know the lower and upper bound, and these are the same, + -- then we morally know the ScalExp, but only if the variable has + -- the right type. + (Just entry, (Just lower, Just upper)) + | entryType entry == Prim int32, + lower == upper, scalExpType lower == int32 -> + Just $ expandScalExp (`lookupScalExp` vtable) lower + (Just entry, _) -> asScalExp entry + _ -> Nothing + +lookupValue :: VName -> SymbolTable lore -> Maybe (PrimValue, Certificates) +lookupValue name vtable = case lookupSubExp name vtable of + Just (Constant val, cs) -> Just (val, cs) + _ -> Nothing + +lookupVar :: VName -> SymbolTable lore -> Maybe (VName, Certificates) +lookupVar name vtable = case lookupSubExp name vtable of + Just (Var v, cs) -> Just (v, cs) + _ -> Nothing + +lookupAliases :: VName -> SymbolTable lore -> Names +lookupAliases name vtable = case M.lookup name $ bindings vtable of + Just (LetBound e) -> letBoundAliases e + Just (FParam e) -> fparamAliases e + _ -> mempty + +index :: Attributes lore => VName -> [SubExp] -> SymbolTable lore + -> Maybe (PrimExp VName, Certificates) +index name is table = do + is' <- mapM asPrimExp is + index' name is' table + where asPrimExp i = do + Prim t <- lookupSubExpType i table + return $ primExpFromSubExp t i + +index' :: VName -> [PrimExp VName] -> SymbolTable lore + -> Maybe (PrimExp VName, Certificates) +index' name is vtable = do + entry <- lookup name vtable + case entry of + LetBound entry' | + Just k <- elemIndex name $ patternValueNames $ + stmPattern $ letBoundStm entry' -> + letBoundIndex entry' k is + FreeVar entry' -> + freeVarIndex entry' name is + LParam entry' -> lparamIndex entry' is + _ -> Nothing + +lookupRange :: VName -> SymbolTable lore -> ScalExpRange +lookupRange name vtable = + maybe (Nothing, Nothing) valueRange $ lookup name vtable + +enclosingLoopVars :: [VName] -> SymbolTable lore -> [VName] +enclosingLoopVars free vtable = + map fst $ + sortBy (flip (comparing (bindingDepth . snd))) $ + filter (loopVariable . snd) $ mapMaybe fetch free + where fetch name = do e <- lookup name vtable + return (name, e) + +rangesRep :: SymbolTable lore -> AS.RangesRep +rangesRep = M.filter knownRange . M.map toRep . bindings + where toRep entry = (bindingDepth entry, lower, upper) + where (lower, upper) = valueRange entry + knownRange (_, lower, upper) = isJust lower || isJust upper + +class IndexOp op where + indexOp :: (Attributes lore, IndexOp (Op lore)) => + SymbolTable lore -> Int -> op + -> [PrimExp VName] -> Maybe (PrimExp VName, Certificates) + indexOp _ _ _ _ = Nothing + +instance IndexOp () where + +indexExp :: (IndexOp (Op lore), Attributes lore) => + SymbolTable lore -> Exp lore -> Int -> IndexArray + +indexExp vtable (Op op) k is = + indexOp vtable k op is + +indexExp _ (BasicOp (Iota _ x s to_it)) _ [i] + | IntType from_it <- primExpType i = + Just ( ConvOpExp (SExt from_it to_it) i + * primExpFromSubExp (IntType to_it) s + + primExpFromSubExp (IntType to_it) x + , mempty) + +indexExp table (BasicOp (Replicate (Shape ds) v)) _ is + | length ds == length is, + Just (Prim t) <- lookupSubExpType v table = + Just (primExpFromSubExp t v, mempty) + +indexExp table (BasicOp (Replicate (Shape [_]) (Var v))) _ (_:is) = + index' v is table + +indexExp table (BasicOp (Reshape newshape v)) _ is + | Just oldshape <- arrayDims <$> lookupType v table = + let is' = + reshapeIndex (map (primExpFromSubExp int32) oldshape) + (map (primExpFromSubExp int32) $ newDims newshape) + is + in index' v is' table + +indexExp table (BasicOp (Index v slice)) _ is = + index' v (adjust slice is) table + where adjust (DimFix j:js') is' = + pe j : adjust js' is' + adjust (DimSlice j _ s:js') (i:is') = + let i_t_s = i * pe s + j_p_i_t_s = pe j + i_t_s + in j_p_i_t_s : adjust js' is' + adjust _ _ = [] + + pe = primExpFromSubExp (IntType Int32) + +indexExp _ _ _ _ = Nothing + +indexChunk :: SymbolTable lore -> VName -> VName -> IndexArray +indexChunk table offset array (i:is) = + index' array (offset'+i:is) table + where offset' = primExpFromSubExp (IntType Int32) (Var offset) +indexChunk _ _ _ _ = Nothing + +defBndEntry :: (Attributes lore, IndexOp (Op lore)) => + SymbolTable lore + -> PatElem lore + -> Range + -> Names + -> Stm lore + -> LetBoundEntry lore +defBndEntry vtable patElem range als bnd = + LetBoundEntry { + letBoundRange = simplifyRange $ scalExpRange range + , letBoundAttr = patElemAttr patElem + , letBoundAliases = als + , letBoundStm = bnd + , letBoundScalExp = + runReader (toScalExp (`lookupScalExp` vtable) (stmExp bnd)) types + , letBoundStmDepth = 0 + , letBoundIndex = \k -> fmap (second (<>(stmAuxCerts $ stmAux bnd))) . + indexExp vtable (stmExp bnd) k + } + where ranges :: AS.RangesRep + ranges = rangesRep vtable + + types = toScope vtable + + scalExpRange :: Range -> ScalExpRange + scalExpRange (lower, upper) = + (scalExpBound fst =<< lower, + scalExpBound snd =<< upper) + + scalExpBound :: (ScalExpRange -> Maybe ScalExp) + -> Ranges.KnownBound + -> Maybe ScalExp + scalExpBound pick (Ranges.VarBound v) = + pick $ lookupRange v vtable + scalExpBound _ (Ranges.ScalarBound se) = + Just se + scalExpBound _ (Ranges.MinimumBound b1 b2) = do + b1' <- scalExpBound fst b1 + b2' <- scalExpBound fst b2 + return $ MaxMin True [b1', b2'] + scalExpBound _ (Ranges.MaximumBound b1 b2) = do + b1' <- scalExpBound snd b1 + b2' <- scalExpBound snd b2 + return $ MaxMin False [b1', b2'] + + simplifyRange :: ScalExpRange -> ScalExpRange + simplifyRange (lower, upper) = + (simplifyBound lower, + simplifyBound upper) + + simplifyBound (Just se) | scalExpType se == int32 = + Just $ AS.simplify se ranges + simplifyBound _ = + Nothing + +bindingEntries :: (Ranged lore, Aliases.Aliased lore, IndexOp (Op lore)) => + Stm lore -> SymbolTable lore + -> [LetBoundEntry lore] +bindingEntries bnd@(Let pat _ _) vtable = do + pat_elem <- patternElements pat + return $ defBndEntry vtable pat_elem + (Ranges.rangeOf pat_elem) (Aliases.aliasesOf pat_elem) bnd + +insertEntry :: Attributes lore => + VName -> Entry lore -> SymbolTable lore + -> SymbolTable lore +insertEntry name entry = + insertEntries [(name,entry)] + +insertEntries :: Attributes lore => + [(VName, Entry lore)] -> SymbolTable lore + -> SymbolTable lore +insertEntries entries vtable = + let vtable' = vtable { bindings = foldl insertWithDepth (bindings vtable) entries } + in foldr (`isAtLeast` 0) vtable' dim_vars + where insertWithDepth bnds (name, entry) = + let entry' = setStmDepth (loopDepth vtable) entry + in M.insert name entry' bnds + dim_vars = subExpVars $ concatMap (arrayDims . entryType . snd) entries + +insertStm :: (IndexOp (Op lore), Ranged lore, Aliases.Aliased lore) => + Stm lore + -> SymbolTable lore + -> SymbolTable lore +insertStm stm vtable = + foldl' addRevAliases + (insertEntries (zip names $ map LetBound $ bindingEntries stm vtable) vtable) $ + patternElements $ stmPattern stm + where names = patternNames $ stmPattern stm + adjustSeveral f = flip $ foldl' $ flip $ M.adjust f + addRevAliases vtable' pe = + vtable' { bindings = adjustSeveral update inedges $ bindings vtable' } + where inedges = expandAliases (Aliases.aliasesOf pe) vtable' + update (LetBound entry) = + LetBound entry + { letBoundAliases = patElemName pe `S.insert` letBoundAliases entry } + update (FParam entry) = + FParam entry + { fparamAliases = patElemName pe `S.insert` fparamAliases entry } + update e = e + +expandAliases :: Names -> SymbolTable lore -> Names +expandAliases names vtable = names `S.union` aliasesOfAliases + where aliasesOfAliases = + mconcat . map (`lookupAliases` vtable) . S.toList $ names + +insertFParam :: Attributes lore => + AST.FParam lore + -> SymbolTable lore + -> SymbolTable lore +insertFParam fparam = insertEntry name entry + where name = AST.paramName fparam + entry = FParam FParamEntry { fparamRange = (Nothing, Nothing) + , fparamAttr = AST.paramAttr fparam + , fparamAliases = mempty + , fparamStmDepth = 0 + } + +insertFParams :: Attributes lore => + [AST.FParam lore] + -> SymbolTable lore + -> SymbolTable lore +insertFParams fparams symtable = foldr insertFParam symtable fparams + +insertLParamWithRange :: Attributes lore => + LParam lore -> ScalExpRange -> IndexArray -> SymbolTable lore + -> SymbolTable lore +insertLParamWithRange param range indexf vtable = + -- We know that the sizes in the type of param are at least zero, + -- since they are array sizes. + let vtable' = insertEntry name bind vtable + in foldr (`isAtLeast` 0) vtable' sizevars + where bind = LParam LParamEntry { lparamRange = range + , lparamAttr = AST.paramAttr param + , lparamStmDepth = 0 + , lparamIndex = indexf + } + name = AST.paramName param + sizevars = subExpVars $ arrayDims $ AST.paramType param + +insertLParam :: Attributes lore => + LParam lore -> SymbolTable lore -> SymbolTable lore +insertLParam param = + insertLParamWithRange param (Nothing, Nothing) (const Nothing) + +insertArrayLParam :: Attributes lore => + LParam lore -> Maybe VName -> SymbolTable lore + -> SymbolTable lore +insertArrayLParam param (Just array) vtable = + -- We now know that the outer size of 'array' is at least one, and + -- that the inner sizes are at least zero, since they are array + -- sizes. + let vtable' = insertLParamWithRange param (lookupRange array vtable) (const Nothing) vtable + in case arrayDims <$> lookupType array vtable of + Just (Var v:_) -> (v `isAtLeast` 1) vtable' + _ -> vtable' +insertArrayLParam param Nothing vtable = + -- Well, we still know that it's a param... + insertLParam param vtable + +insertChunkLParam :: Attributes lore => + VName -> LParam lore -> VName -> SymbolTable lore + -> SymbolTable lore +insertChunkLParam offset param array vtable = + -- We now know that the outer size of 'array' is at least one, and + -- that the inner sizes are at least zero, since they are array + -- sizes. + let vtable' = insertLParamWithRange param + (lookupRange array vtable) (indexChunk vtable offset array) vtable + in case arrayDims <$> lookupType array vtable of + Just (Var v:_) -> (v `isAtLeast` 1) vtable' + _ -> vtable' + +insertLoopVar :: Attributes lore => VName -> IntType -> SubExp -> SymbolTable lore -> SymbolTable lore +insertLoopVar name it bound = insertEntry name bind + where bind = LoopVar LoopVarEntry { + loopVarRange = (Just 0, + Just $ subExpToScalExp bound (IntType it) - 1) + , loopVarStmDepth = 0 + , loopVarType = it + } + +insertFreeVar :: Attributes lore => VName -> NameInfo lore -> SymbolTable lore -> SymbolTable lore +insertFreeVar name attr = insertEntry name entry + where entry = FreeVar FreeVarEntry { + freeVarAttr = attr + , freeVarRange = (Nothing, Nothing) + , freeVarStmDepth = 0 + , freeVarIndex = \_ _ -> Nothing + } + +updateBounds :: Attributes lore => Bool -> SubExp -> SymbolTable lore -> SymbolTable lore +updateBounds isTrue cond vtable = + case runReader (toScalExp (`lookupScalExp` vtable) $ BasicOp $ SubExp cond) types of + Nothing -> vtable + Just cond' -> + let cond'' | isTrue = cond' + | otherwise = SNot cond' + in updateBounds' cond'' vtable + where types = toScope vtable + +-- | Updating the ranges of all symbols whenever we enter a branch is +-- presently too expensive, and disabled here. +noUpdateBounds :: Bool +noUpdateBounds = True + +-- | Refines the ranges in the symbol table with +-- ranges extracted from branch conditions. +-- `cond' is the condition of the if-branch. +updateBounds' :: ScalExp -> SymbolTable lore -> SymbolTable lore +updateBounds' _ sym_tab | noUpdateBounds = sym_tab +updateBounds' cond sym_tab = + foldr updateBound sym_tab $ mapMaybe solve_leq0 $ + getNotFactorsLEQ0 $ AS.simplify (SNot cond) ranges + where + updateBound (sym,True ,bound) = setUpperBound sym bound + updateBound (sym,False,bound) = setLowerBound sym bound + + ranges = M.filter nonEmptyRange $ M.map toRep $ bindings sym_tab + toRep entry = (bindingDepth entry, lower, upper) + where (lower, upper) = valueRange entry + nonEmptyRange (_, lower, upper) = isJust lower || isJust upper + + -- | Input: a bool exp in DNF form, named `cond' + -- It gets the terms of the argument, + -- i.e., cond = c1 || ... || cn + -- and negates them. + -- Returns [not c1, ..., not cn], i.e., the factors + -- of `not cond' in CNF form: not cond = (not c1) && ... && (not cn) + getNotFactorsLEQ0 :: ScalExp -> [ScalExp] + getNotFactorsLEQ0 (RelExp rel e_scal) = + if scalExpType e_scal /= int32 then [] + else let leq0_escal = if rel == LTH0 + then SMinus 0 e_scal + else SMinus 1 e_scal + + in [AS.simplify leq0_escal ranges] + getNotFactorsLEQ0 (SLogOr e1 e2) = getNotFactorsLEQ0 e1 ++ getNotFactorsLEQ0 e2 + getNotFactorsLEQ0 _ = [] + + -- | Argument is scalar expression `e'. + -- Implementation finds the symbol defined at + -- the highest depth in expression `e', call it `i', + -- and decomposes e = a*i + b. If `a' and `b' are + -- free of `i', AND `a == 1 or -1' THEN the upper/lower + -- bound can be improved. Otherwise Nothing. + -- + -- Returns: Nothing or + -- Just (i, a == 1, -a*b), i.e., (symbol, isUpperBound, bound) + solve_leq0 :: ScalExp -> Maybe (VName, Bool, ScalExp) + solve_leq0 e_scal = do + sym <- pickRefinedSym S.empty e_scal + (a,b) <- either (const Nothing) id $ AS.linFormScalE sym e_scal ranges + case a of + -1 -> + Just (sym, False, b) + 1 -> + let mb = AS.simplify (negate b) ranges + in Just (sym, True, mb) + _ -> Nothing + + -- When picking a symbols, @sym@ whose bound it is to be refined: + -- make sure that @sym@ does not belong to the transitive closure + -- of the symbols apearing in the ranges of all the other symbols + -- in the sclar expression (themselves included). + -- If this does not hold, pick another symbol, rinse and repeat. + pickRefinedSym :: S.Set VName -> ScalExp -> Maybe VName + pickRefinedSym elsyms0 e_scal = do + let candidates = freeIn e_scal + sym0 = AS.pickSymToElim ranges elsyms0 e_scal + case sym0 of + Just sy -> let trclsyms = foldl trClSymsInRange S.empty $ S.toList $ + candidates `S.difference` S.singleton sy + in if S.member sy trclsyms + then pickRefinedSym (S.insert sy elsyms0) e_scal + else sym0 + Nothing -> sym0 + -- computes the transitive closure of the symbols appearing + -- in the ranges of a symbol + trClSymsInRange :: S.Set VName -> VName -> S.Set VName + trClSymsInRange cur_syms sym = + if S.member sym cur_syms then cur_syms + else case M.lookup sym ranges of + Just (_,lb,ub) -> let sym_bds = concatMap (S.toList . freeIn) (catMaybes [lb, ub]) + in foldl trClSymsInRange + (S.insert sym cur_syms) + (S.toList $ S.fromList sym_bds) + Nothing -> S.insert sym cur_syms + +setUpperBound :: VName -> ScalExp -> SymbolTable lore + -> SymbolTable lore +setUpperBound name bound vtable = + vtable { bindings = M.adjust setUpperBound' name $ bindings vtable } + where setUpperBound' entry = + let (oldLowerBound, oldUpperBound) = valueRange entry + in if alreadyTheBound bound True oldUpperBound + then entry + else setValueRange + (oldLowerBound, + Just $ maybe bound (MaxMin True . (:[bound])) oldUpperBound) + entry + +setLowerBound :: VName -> ScalExp -> SymbolTable lore -> SymbolTable lore +setLowerBound name bound vtable = + vtable { bindings = M.adjust setLowerBound' name $ bindings vtable } + where setLowerBound' entry = + let (oldLowerBound, oldUpperBound) = valueRange entry + in if alreadyTheBound bound False oldLowerBound + then entry + else setValueRange + (Just $ maybe bound (MaxMin False . (:[bound])) oldLowerBound, + oldUpperBound) + entry + +alreadyTheBound :: ScalExp -> Bool -> Maybe ScalExp -> Bool +alreadyTheBound _ _ Nothing = False +alreadyTheBound new_bound b1 (Just cur_bound) + | cur_bound == new_bound = True + | MaxMin b2 ses <- cur_bound = b1 == b2 && (new_bound `L.elem` ses) + | otherwise = False + +isAtLeast :: VName -> Int -> SymbolTable lore -> SymbolTable lore +isAtLeast name x = + setLowerBound name $ fromIntegral x diff --git a/src/Futhark/Analysis/Usage.hs b/src/Futhark/Analysis/Usage.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/Usage.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Analysis.Usage + ( usageInStm + , usageInExp + , usageInLambda + + , UsageInOp(..) + ) + where + +import Data.Semigroup ((<>)) +import Data.Foldable +import qualified Data.Set as S + +import Futhark.Representation.AST +import Futhark.Representation.AST.Attributes.Aliases +import qualified Futhark.Analysis.UsageTable as UT + +usageInStm :: (Attributes lore, Aliased lore, UsageInOp (Op lore)) => + Stm lore -> UT.UsageTable +usageInStm (Let pat lore e) = + mconcat [usageInPat, + usageInExpLore, + usageInExp e, + UT.usages (freeInExp e)] + where usageInPat = + UT.usages (mconcat (map freeIn $ patternElements pat) + `S.difference` + S.fromList (patternNames pat)) + usageInExpLore = + UT.usages $ freeIn lore + +usageInExp :: (Aliased lore, UsageInOp (Op lore)) => Exp lore -> UT.UsageTable +usageInExp (Apply _ args _ _) = + mconcat [ mconcat $ map UT.consumedUsage $ + S.toList $ subExpAliases arg + | (arg,d) <- args, d == Consume ] +usageInExp (DoLoop _ merge _ _) = + mconcat [ mconcat $ map UT.consumedUsage $ + S.toList $ subExpAliases se + | (v,se) <- merge, unique $ paramDeclType v ] +usageInExp (If _ tbranch fbranch _) = + fold $ map UT.consumedUsage $ S.toList $ + consumedInBody tbranch <> consumedInBody fbranch +usageInExp (BasicOp (Update src _ _)) = + UT.consumedUsage src +usageInExp (Op op) = + mconcat $ usageInOp op : map UT.consumedUsage (S.toList $ consumedInOp op) +usageInExp _ = UT.empty + +class UsageInOp op where + usageInOp :: op -> UT.UsageTable + +instance UsageInOp () where + usageInOp () = mempty + +usageInLambda :: Aliased lore => + Lambda lore -> [VName] -> UT.UsageTable +usageInLambda lam arrs = + mconcat $ + map (UT.consumedUsage . snd) $ + filter ((`S.member` consumed_in_body) . fst) $ + zip (map paramName arr_params) arrs + where arr_params = snd $ splitAt n $ lambdaParams lam + consumed_in_body = consumedInBody $ lambdaBody lam + n = length arrs diff --git a/src/Futhark/Analysis/UsageTable.hs b/src/Futhark/Analysis/UsageTable.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Analysis/UsageTable.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE Strict #-} +-- | A usage-table is sort of a bottom-up symbol table, describing how +-- (and if) a variable is used. +module Futhark.Analysis.UsageTable + ( UsageTable + , empty + , contains + , without + , lookup + , keys + , used + , expand + , isConsumed + , isInResult + , isUsedDirectly + , allConsumed + , usages + , usage + , consumedUsage + , inResultUsage + , Usages + , leftScope + ) + where + +import Control.Arrow (first) +import Data.Bits +import qualified Data.Foldable as Foldable +import Data.List (foldl') +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Data.Semigroup as Sem + +import Prelude hiding (lookup) + +import Futhark.Transform.Substitute +import Futhark.Representation.AST + +newtype UsageTable = UsageTable (M.Map VName Usages) + deriving (Eq, Show) + +instance Sem.Semigroup UsageTable where + UsageTable table1 <> UsageTable table2 = + UsageTable $ M.unionWith (<>) table1 table2 + +instance Monoid UsageTable where + mempty = empty + mappend = (Sem.<>) + +instance Substitute UsageTable where + substituteNames subst (UsageTable table) + | not $ M.null $ subst `M.intersection` table = + UsageTable $ M.fromList $ + map (first $ substituteNames subst) $ M.toList table + | otherwise = UsageTable table + +empty :: UsageTable +empty = UsageTable M.empty + +contains :: UsageTable -> [VName] -> Bool +contains (UsageTable table) = Foldable.any (`M.member` table) + +without :: UsageTable -> [VName] -> UsageTable +without (UsageTable table) = UsageTable . Foldable.foldl (flip M.delete) table + +lookup :: VName -> UsageTable -> Maybe Usages +lookup name (UsageTable table) = M.lookup name table + +lookupPred :: (Usages -> Bool) -> VName -> UsageTable -> Bool +lookupPred f name = maybe False f . lookup name + +used :: VName -> UsageTable -> Bool +used = lookupPred $ const True + +-- | Expand the usage table based on aliasing information. +expand :: (VName -> Names) -> UsageTable -> UsageTable +expand look (UsageTable m) = UsageTable $ foldl' grow m $ M.toList m + where grow m' (k, v) = foldl' (grow'' $ v `withoutU` presentU) m' $ look k + grow'' v m'' k = M.insertWith (<>) k v m'' + +keys :: UsageTable -> [VName] +keys (UsageTable table) = M.keys table + +is :: Usages -> VName -> UsageTable -> Bool +is = lookupPred . matches + +isConsumed :: VName -> UsageTable -> Bool +isConsumed = is consumedU + +isInResult :: VName -> UsageTable -> Bool +isInResult = is inResultU + +-- | Has the given name been used directly (i.e. could we rename it or +-- remove it without anyone noticing?) +isUsedDirectly :: VName -> UsageTable -> Bool +isUsedDirectly = is presentU + +allConsumed :: UsageTable -> Names +allConsumed (UsageTable m) = + S.fromList . map fst . filter (matches consumedU . snd) $ M.toList m + +usages :: Names -> UsageTable +usages names = UsageTable $ M.fromList [ (name, presentU) | name <- S.toList names ] + +usage :: VName -> Usages -> UsageTable +usage name uses = UsageTable $ M.singleton name uses + +consumedUsage :: VName -> UsageTable +consumedUsage name = UsageTable $ M.singleton name consumedU + +inResultUsage :: VName -> UsageTable +inResultUsage name = UsageTable $ M.singleton name inResultU + +newtype Usages = Usages Int + deriving (Eq, Ord, Show) + +instance Sem.Semigroup Usages where + Usages x <> Usages y = Usages $ x .|. y + +instance Monoid Usages where + mempty = Usages 0 + mappend = (Sem.<>) + +consumedU, inResultU, presentU :: Usages +consumedU = Usages 1 +inResultU = Usages 2 +presentU = Usages 4 + +-- | Check whether the bits that are set in the first argument are +-- also set in the second. +matches :: Usages -> Usages -> Bool +matches (Usages x) (Usages y) = x == (x .&. y) + +-- | x - y, but for Usages. +withoutU :: Usages -> Usages -> Usages +withoutU (Usages x) (Usages y) = Usages $ x .&. complement y + +leftScope :: UsageTable -> UsageTable +leftScope (UsageTable table) = UsageTable $ M.map (`withoutU` inResultU) table diff --git a/src/Futhark/Binder.hs b/src/Futhark/Binder.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Binder.hs @@ -0,0 +1,197 @@ +{-# LANGUAGE FlexibleContexts, GeneralizedNewtypeDeriving, TypeFamilies, FlexibleInstances, MultiParamTypeClasses, UndecidableInstances #-} +{-# LANGUAGE ConstraintKinds #-} +-- | This module defines a convenience monad/typeclass for creating +-- normalised programs. +module Futhark.Binder + ( -- * A concrete @MonadBinder@ monad. + BinderT + , runBinderT + , BinderOps (..) + , bindableMkExpAttrB + , bindableMkBodyB + , bindableMkLetNamesB + , Binder + , runBinder + , runBinder_ + , joinBinder + , runBodyBinder + -- * Non-class interface + , addBinderStms + , collectBinderStms + , certifyingBinder + -- * The 'MonadBinder' typeclass + , module Futhark.Binder.Class + ) +where + +import Control.Arrow (second) +import Control.Monad.Writer +import Control.Monad.State.Strict +import Control.Monad.Reader +import Control.Monad.Error.Class +import qualified Data.Map.Strict as M +import qualified Control.Monad.Fail as Fail + +import Futhark.Binder.Class +import Futhark.Representation.AST + +class Attributes lore => BinderOps lore where + mkExpAttrB :: (MonadBinder m, Lore m ~ lore) => + Pattern lore -> Exp lore -> m (ExpAttr lore) + mkBodyB :: (MonadBinder m, Lore m ~ lore) => + Stms lore -> Result -> m (Body lore) + mkLetNamesB :: (MonadBinder m, Lore m ~ lore) => + [VName] -> Exp lore -> m (Stm lore) + +bindableMkExpAttrB :: (MonadBinder m, Bindable (Lore m)) => + Pattern (Lore m) -> Exp (Lore m) -> m (ExpAttr (Lore m)) +bindableMkExpAttrB pat e = return $ mkExpAttr pat e + +bindableMkBodyB :: (MonadBinder m, Bindable (Lore m)) => + Stms (Lore m) -> Result -> m (Body (Lore m)) +bindableMkBodyB stms res = return $ mkBody stms res + +bindableMkLetNamesB :: (MonadBinder m, Bindable (Lore m)) => + [VName] -> Exp (Lore m) -> m (Stm (Lore m)) +bindableMkLetNamesB = mkLetNames + +newtype BinderT lore m a = BinderT (StateT (Stms lore, Scope lore) m a) + deriving (Functor, Monad, Applicative) + +instance MonadTrans (BinderT lore) where + lift = BinderT . lift + +instance Monad m => Fail.MonadFail (BinderT lore m) where + fail = error . ("BinderT.fail: "++) + +type Binder lore = BinderT lore (State VNameSource) + +instance MonadFreshNames m => MonadFreshNames (BinderT lore m) where + getNameSource = lift getNameSource + putNameSource = lift . putNameSource + +instance (Attributes lore, Monad m) => + HasScope lore (BinderT lore m) where + lookupType name = do + t <- BinderT $ gets $ M.lookup name . snd + case t of + Nothing -> fail $ "BinderT.lookupType: unknown variable " ++ pretty name + Just t' -> return $ typeOf t' + askScope = BinderT $ gets snd + +instance (Attributes lore, Monad m) => + LocalScope lore (BinderT lore m) where + localScope types (BinderT m) = BinderT $ do + modify $ second (M.union types) + x <- m + modify $ second (`M.difference` types) + return x + +instance (Attributes lore, MonadFreshNames m, BinderOps lore) => + MonadBinder (BinderT lore m) where + type Lore (BinderT lore m) = lore + mkExpAttrM = mkExpAttrB + mkBodyM = mkBodyB + mkLetNamesM = mkLetNamesB + + addStms = addBinderStms + collectStms = collectBinderStms + + certifying = certifyingBinder + +runBinderT :: MonadFreshNames m => + BinderT lore m a + -> Scope lore + -> m (a, Stms lore) +runBinderT (BinderT m) scope = do + (x, (stms, _)) <- runStateT m (mempty, scope) + return (x, stms) + +runBinder :: (MonadFreshNames m, + HasScope somelore m, SameScope somelore lore) => + Binder lore a + -> m (a, Stms lore) +runBinder m = do + types <- askScope + modifyNameSource $ runState $ runBinderT m $ castScope types + +-- | Like 'runBinder', but throw away the result and just return the +-- added bindings. +runBinder_ :: (MonadFreshNames m, + HasScope somelore m, SameScope somelore lore) => + Binder lore a + -> m (Stms lore) +runBinder_ = fmap snd . runBinder + +-- | As 'runBinder', but uses 'addStm' to add the returned +-- bindings to the surrounding monad. +joinBinder :: MonadBinder m => Binder (Lore m) a -> m a +joinBinder m = do (x, bnds) <- runBinder m + addStms bnds + return x + +runBodyBinder :: (Bindable lore, MonadFreshNames m, + HasScope somelore m, SameScope somelore lore) => + Binder lore (Body lore) -> m (Body lore) +runBodyBinder = fmap (uncurry $ flip insertStms) . runBinder + +addBinderStms :: Monad m => + Stms lore -> BinderT lore m () +addBinderStms stms = BinderT $ + modify $ \(cur_stms,scope) -> (cur_stms<>stms, + scope `M.union` scopeOf stms) + +collectBinderStms :: Monad m => + BinderT lore m a + -> BinderT lore m (a, Stms lore) +collectBinderStms m = do + (old_stms, old_scope) <- BinderT get + BinderT $ put (mempty, old_scope) + x <- m + (new_stms, _) <- BinderT get + BinderT $ put (old_stms, old_scope) + return (x, new_stms) + +certifyingBinder :: (MonadFreshNames m, BinderOps lore) => + Certificates -> BinderT lore m a + -> BinderT lore m a +certifyingBinder cs m = do + (x, stms) <- collectStms m + addStms $ certify cs <$> stms + return x + +-- Utility instance defintions for MTL classes. These require +-- UndecidableInstances, but save on typing elsewhere. + +mapInner :: Monad m => + (m (a, (Stms lore, Scope lore)) + -> m (b, (Stms lore, Scope lore))) + -> BinderT lore m a -> BinderT lore 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 + ask = BinderT $ lift ask + local f = mapInner $ local f + +instance MonadState s m => MonadState s (BinderT lore m) where + get = BinderT $ lift get + put = BinderT . lift . put + +instance MonadWriter w m => MonadWriter w (BinderT lore m) where + tell = BinderT . lift . tell + pass = mapInner $ \m -> pass $ do + ((x, f), s) <- m + return ((x, s), f) + listen = mapInner $ \m -> do + ((x, s), y) <- listen m + return ((x, y), s) + +instance MonadError e m => MonadError e (BinderT lore m) where + throwError = lift . throwError + catchError (BinderT m) f = + BinderT $ catchError m $ unBinder . f + where unBinder (BinderT m') = m' diff --git a/src/Futhark/Binder/Class.hs b/src/Futhark/Binder/Class.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Binder/Class.hs @@ -0,0 +1,117 @@ +{-# LANGUAGE FlexibleContexts, TypeFamilies #-} +-- | This module defines a convenience typeclass for creating +-- normalised programs. +module Futhark.Binder.Class + ( Bindable (..) + , mkLet + , MonadBinder (..) + , mkLetM + , bodyStms + , insertStms + , insertStm + , letBind + , letBind_ + , letBindNames + , letBindNames_ + , collectStms_ + , bodyBind + + , module Futhark.MonadFreshNames + ) +where + +import Control.Monad.Writer +import qualified Control.Monad.Fail as Fail + +import Futhark.Representation.AST +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 +-- exclusively within a 'MonadBinder' instance, it is acceptable for +-- them to create new bindings, however. +class (Attributes lore, + FParamAttr lore ~ DeclType, + LParamAttr lore ~ Type, + RetType lore ~ DeclExtType, + BranchType lore ~ ExtType, + SetType (LetAttr lore)) => + Bindable lore where + mkExpPat :: [Ident] -> [Ident] -> Exp lore -> Pattern lore + mkExpAttr :: Pattern lore -> Exp lore -> ExpAttr lore + mkBody :: Stms lore -> Result -> Body lore + mkLetNames :: (MonadFreshNames m, HasScope lore m) => + [VName] -> Exp lore -> m (Stm lore) + +-- | A monad that supports the creation of bindings from expressions +-- and bodies from bindings, with a specific lore. This is the main +-- typeclass that a monad must implement in order for it to be useful +-- for generating or modifying Futhark code. +-- +-- 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. It is acceptable for them to create new +-- bindings, however. +class (Attributes (Lore m), + MonadFreshNames m, Applicative m, Monad m, + LocalScope (Lore m) m, + Fail.MonadFail m) => + MonadBinder m where + type Lore m :: * + mkExpAttrM :: Pattern (Lore m) -> Exp (Lore m) -> m (ExpAttr (Lore m)) + mkBodyM :: Stms (Lore m) -> Result -> m (Body (Lore m)) + mkLetNamesM :: [VName] -> Exp (Lore m) -> m (Stm (Lore m)) + addStm :: Stm (Lore m) -> m () + addStm = addStms . oneStm + addStms :: Stms (Lore m) -> m () + collectStms :: m a -> m (a, Stms (Lore m)) + certifying :: Certificates -> m a -> m a + +mkLetM :: MonadBinder m => Pattern (Lore m) -> Exp (Lore m) -> m (Stm (Lore m)) +mkLetM pat e = Let pat <$> (StmAux mempty <$> mkExpAttrM pat e) <*> pure e + +letBind :: MonadBinder m => + Pattern (Lore m) -> Exp (Lore m) -> m [Ident] +letBind pat e = do + bnd <- mkLetM pat e + addStm bnd + return $ patternValueIdents $ stmPattern bnd + +letBind_ :: MonadBinder m => + Pattern (Lore m) -> Exp (Lore m) -> m () +letBind_ pat e = void $ letBind pat e + +mkLet :: Bindable lore => [Ident] -> [Ident] -> Exp lore -> Stm lore +mkLet ctx val e = + let pat = mkExpPat ctx val e + attr = mkExpAttr pat e + in Let pat (StmAux mempty attr) e + +letBindNames :: MonadBinder m => + [VName] -> Exp (Lore m) -> m [Ident] +letBindNames names e = do + bnd <- mkLetNamesM names e + addStm bnd + return $ patternValueIdents $ stmPattern bnd + +letBindNames_ :: MonadBinder m => + [VName] -> Exp (Lore m) -> m () +letBindNames_ names e = void $ letBindNames names e + +collectStms_ :: MonadBinder m => m a -> m (Stms (Lore m)) +collectStms_ = fmap snd . collectStms + +bodyBind :: MonadBinder m => Body (Lore m) -> m [SubExp] +bodyBind (Body _ bnds es) = do + addStms bnds + return es + +-- | Add several bindings at the outermost level of a 'Body'. +insertStms :: Bindable lore => Stms lore -> Body lore -> Body lore +insertStms bnds1 (Body _ bnds2 res) = mkBody (bnds1<>bnds2) res + +-- | Add a single binding at the outermost level of a 'Body'. +insertStm :: Bindable lore => Stm lore -> Body lore -> Body lore +insertStm = insertStms . oneStm diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/COpenCL.hs @@ -0,0 +1,341 @@ +{-# LANGUAGE QuasiQuotes, FlexibleContexts #-} +module Futhark.CodeGen.Backends.COpenCL + ( compileProg + , GC.CParts(..) + , GC.asLibrary + , GC.asExecutable + ) where + +import Control.Monad hiding (mapM) +import Data.List + +import qualified Language.C.Syntax as C +import qualified Language.C.Quote.OpenCL as C + +import Futhark.Error +import Futhark.Representation.ExplicitMemory hiding (GetSize, CmpSizeLe, GetSizeMax) +import Futhark.CodeGen.Backends.COpenCL.Boilerplate +import qualified Futhark.CodeGen.Backends.GenericC as GC +import Futhark.CodeGen.Backends.GenericC.Options +import Futhark.CodeGen.ImpCode.OpenCL +import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen +import Futhark.MonadFreshNames + +compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError GC.CParts) +compileProg prog = do + res <- ImpGen.compileProg prog + case res of + Left err -> return $ Left err + Right (Program opencl_code opencl_prelude kernel_names types sizes prog') -> + Right <$> GC.compileProg operations + (generateBoilerplate opencl_code opencl_prelude kernel_names types sizes) + include_opencl_h [Space "device", Space "local", DefaultSpace] + cliOptions prog' + where operations :: GC.Operations OpenCL () + operations = GC.Operations + { GC.opsCompiler = callKernel + , GC.opsWriteScalar = writeOpenCLScalar + , GC.opsReadScalar = readOpenCLScalar + , GC.opsAllocate = allocateOpenCLBuffer + , GC.opsDeallocate = deallocateOpenCLBuffer + , GC.opsCopy = copyOpenCLMemory + , GC.opsStaticArray = staticOpenCLArray + , GC.opsMemoryType = openclMemoryType + , GC.opsFatMemory = True + } + include_opencl_h = unlines ["#define CL_USE_DEPRECATED_OPENCL_1_2_APIS", + "#ifdef __APPLE__", + "#include ", + "#else", + "#include ", + "#endif"] + +cliOptions :: [Option] +cliOptions = [ Option { optionLongName = "platform" + , optionShortName = Just 'p' + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_set_platform(cfg, optarg);|] + } + , Option { optionLongName = "device" + , optionShortName = Just 'd' + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_set_device(cfg, optarg);|] + } + , Option { optionLongName = "default-group-size" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_set_default_group_size(cfg, atoi(optarg));|] + } + , Option { optionLongName = "default-num-groups" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_set_default_num_groups(cfg, atoi(optarg));|] + } + , Option { optionLongName = "default-tile-size" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_set_default_tile_size(cfg, atoi(optarg));|] + } + , Option { optionLongName = "default-threshold" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_set_default_threshold(cfg, atoi(optarg));|] + } + , Option { optionLongName = "dump-opencl" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_dump_program_to(cfg, optarg);|] + } + , Option { optionLongName = "load-opencl" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [C.cstm|futhark_context_config_load_program_from(cfg, optarg);|] + } + , Option { optionLongName = "print-sizes" + , optionShortName = Nothing + , optionArgument = NoArgument + , optionAction = [C.cstm|{ + int n = futhark_get_num_sizes(); + for (int i = 0; i < n; i++) { + if (strcmp(futhark_get_size_entry(i), entry_point) == 0) { + printf("%s (%s)\n", futhark_get_size_name(i), + futhark_get_size_class(i)); + } + } + exit(0); + }|] + } + , Option { optionLongName = "size" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [C.cstm|{ + char *name = optarg; + char *equals = strstr(optarg, "="); + char *value_str = equals != NULL ? equals+1 : optarg; + int value = atoi(value_str); + if (equals != NULL) { + *equals = 0; + if (futhark_context_config_set_size(cfg, name, value) != 0) { + panic(1, "Unknown size: %s\n", name); + } + } else { + panic(1, "Invalid argument for size option: %s\n", optarg); + }}|] + } + ] + +writeOpenCLScalar :: GC.WriteScalar OpenCL () +writeOpenCLScalar mem i t "device" _ val = do + val' <- newVName "write_tmp" + GC.stm [C.cstm|{$ty:t $id:val' = $exp:val; + OPENCL_SUCCEED_OR_RETURN( + clEnqueueWriteBuffer(ctx->opencl.queue, $exp:mem, CL_TRUE, + $exp:i, sizeof($ty:t), + &$id:val', + 0, NULL, NULL)); + }|] +writeOpenCLScalar _ _ _ space _ _ = + fail $ "Cannot write to '" ++ space ++ "' memory space." + +readOpenCLScalar :: GC.ReadScalar OpenCL () +readOpenCLScalar mem i t "device" _ = do + val <- newVName "read_res" + GC.decl [C.cdecl|$ty:t $id:val;|] + GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN( + clEnqueueReadBuffer(ctx->opencl.queue, $exp:mem, CL_TRUE, + $exp:i, sizeof($ty:t), + &$id:val, + 0, NULL, NULL)); + |] + return [C.cexp|$id:val|] +readOpenCLScalar _ _ _ space _ = + fail $ "Cannot read from '" ++ space ++ "' memory space." + +allocateOpenCLBuffer :: GC.Allocate OpenCL () +allocateOpenCLBuffer mem size tag "device" = + GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(opencl_alloc(&ctx->opencl, $exp:size, $exp:tag, &$exp:mem));|] +allocateOpenCLBuffer _ _ _ "local" = + return () -- Hack - these memory blocks do not actually exist. +allocateOpenCLBuffer _ _ _ space = + fail $ "Cannot allocate in '" ++ space ++ "' space" + +deallocateOpenCLBuffer :: GC.Deallocate OpenCL () +deallocateOpenCLBuffer mem tag "device" = + GC.stm [C.cstm|OPENCL_SUCCEED_OR_RETURN(opencl_free(&ctx->opencl, $exp:mem, $exp:tag));|] +deallocateOpenCLBuffer _ _ "local" = + return () -- Hack - these memory blocks do not actually exist. +deallocateOpenCLBuffer _ _ space = + fail $ "Cannot deallocate in '" ++ space ++ "' space" + + +copyOpenCLMemory :: GC.Copy OpenCL () +-- The read/write/copy-buffer functions fail if the given offset is +-- out of bounds, even if asked to read zero bytes. We protect with a +-- branch to avoid this. +copyOpenCLMemory destmem destidx DefaultSpace srcmem srcidx (Space "device") nbytes = + GC.stm [C.cstm| + if ($exp:nbytes > 0) { + OPENCL_SUCCEED_OR_RETURN( + clEnqueueReadBuffer(ctx->opencl.queue, $exp:srcmem, CL_TRUE, + $exp:srcidx, $exp:nbytes, + $exp:destmem + $exp:destidx, + 0, NULL, NULL)); + } + |] +copyOpenCLMemory destmem destidx (Space "device") srcmem srcidx DefaultSpace nbytes = + GC.stm [C.cstm| + if ($exp:nbytes > 0) { + OPENCL_SUCCEED_OR_RETURN( + clEnqueueWriteBuffer(ctx->opencl.queue, $exp:destmem, CL_TRUE, + $exp:destidx, $exp:nbytes, + $exp:srcmem + $exp:srcidx, + 0, NULL, NULL)); + } + |] +copyOpenCLMemory destmem destidx (Space "device") srcmem srcidx (Space "device") nbytes = + -- Be aware that OpenCL swaps the usual order of operands for + -- memcpy()-like functions. The order below is not a typo. + GC.stm [C.cstm|{ + if ($exp:nbytes > 0) { + OPENCL_SUCCEED_OR_RETURN( + clEnqueueCopyBuffer(ctx->opencl.queue, + $exp:srcmem, $exp:destmem, + $exp:srcidx, $exp:destidx, + $exp:nbytes, + 0, NULL, NULL)); + if (ctx->debugging) { + OPENCL_SUCCEED_FATAL(clFinish(ctx->opencl.queue)); + } + } + }|] +copyOpenCLMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes = + GC.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes +copyOpenCLMemory _ _ destspace _ _ srcspace _ = + error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace + +openclMemoryType :: GC.MemoryType OpenCL () +openclMemoryType "device" = pure [C.cty|typename cl_mem|] +openclMemoryType "local" = pure [C.cty|unsigned char|] -- dummy type +openclMemoryType space = + fail $ "OpenCL backend does not support '" ++ space ++ "' memory space." + +staticOpenCLArray :: GC.StaticArray OpenCL () +staticOpenCLArray name "device" t vs = do + let ct = GC.primTypeToCType t + vs' = [[C.cinit|$exp:v|] | v <- map GC.compilePrimValue vs] + num_elems = length vs + name_realtype <- newVName $ baseString name ++ "_realtype" + GC.libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:num_elems] = {$inits:vs'};|] + -- Fake a memory block. + GC.contextField (pretty name) [C.cty|struct memblock_device|] Nothing + -- During startup, copy the data to where we need it. + GC.atInit [C.cstm|{ + typename cl_int success; + ctx->$id:name.references = NULL; + ctx->$id:name.size = 0; + ctx->$id:name.mem = + clCreateBuffer(ctx->opencl.ctx, CL_MEM_READ_WRITE, + ($int:num_elems > 0 ? $int:num_elems : 1)*sizeof($ty:ct), NULL, + &success); + OPENCL_SUCCEED_OR_RETURN(success); + if ($int:num_elems > 0) { + OPENCL_SUCCEED_OR_RETURN( + clEnqueueWriteBuffer(ctx->opencl.queue, ctx->$id:name.mem, CL_TRUE, + 0, $int:num_elems*sizeof($ty:ct), + $id:name_realtype, + 0, NULL, NULL)); + } + }|] + GC.item [C.citem|struct memblock_device $id:name = ctx->$id:name;|] + +staticOpenCLArray _ space _ _ = + fail $ "OpenCL backend cannot create static array in memory space '" ++ space ++ "'" + +callKernel :: GC.OpCompiler OpenCL () +callKernel (GetSize v key) = + GC.stm [C.cstm|$id:v = ctx->sizes.$id:key;|] +callKernel (CmpSizeLe v key x) = do + x' <- GC.compileExp x + GC.stm [C.cstm|$id:v = ctx->sizes.$id:key <= $exp:x';|] + GC.stm [C.cstm|if (ctx->logging) { + fprintf(stderr, "Compared %s <= %d.\n", $string:(pretty key), $exp:x'); + }|] +callKernel (GetSizeMax v size_class) = + let field = "max_" ++ pretty size_class + in GC.stm [C.cstm|$id:v = ctx->opencl.$id:field;|] +callKernel (HostCode c) = + GC.compileCode c + +callKernel (LaunchKernel name args kernel_size workgroup_size) = do + zipWithM_ setKernelArg [(0::Int)..] args + kernel_size' <- mapM GC.compileExp kernel_size + workgroup_size' <- mapM GC.compileExp workgroup_size + launchKernel name kernel_size' workgroup_size' + where setKernelArg i (ValueKArg e bt) = do + v <- GC.compileExpToName "kernel_arg" bt e + GC.stm [C.cstm| + OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, sizeof($id:v), &$id:v)); + |] + + setKernelArg i (MemKArg v) = do + v' <- GC.rawMem v + GC.stm [C.cstm| + OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, sizeof($exp:v'), &$exp:v')); + |] + + setKernelArg i (SharedMemoryKArg num_bytes) = do + num_bytes' <- GC.compileExp $ innerExp num_bytes + GC.stm [C.cstm| + OPENCL_SUCCEED_OR_RETURN(clSetKernelArg(ctx->$id:name, $int:i, $exp:num_bytes', NULL)); + |] + +launchKernel :: C.ToExp a => + String -> [a] -> [a] -> GC.CompilerM op s () +launchKernel kernel_name kernel_dims workgroup_dims = do + global_work_size <- newVName "global_work_size" + time_start <- newVName "time_start" + time_end <- newVName "time_end" + time_diff <- newVName "time_diff" + local_work_size <- newVName "local_work_size" + + GC.stm [C.cstm| + if ($exp:total_elements != 0) { + const size_t $id:global_work_size[$int:kernel_rank] = {$inits:kernel_dims'}; + const size_t $id:local_work_size[$int:kernel_rank] = {$inits:workgroup_dims'}; + typename int64_t $id:time_start = 0, $id:time_end = 0; + if (ctx->debugging) { + fprintf(stderr, "Launching %s with global work size [", $string:kernel_name); + $stms:(printKernelSize global_work_size) + fprintf(stderr, "] and local work size ["); + $stms:(printKernelSize local_work_size) + fprintf(stderr, "].\n"); + $id:time_start = get_wall_time(); + } + OPENCL_SUCCEED_OR_RETURN( + clEnqueueNDRangeKernel(ctx->opencl.queue, ctx->$id:kernel_name, $int:kernel_rank, NULL, + $id:global_work_size, $id:local_work_size, + 0, NULL, NULL)); + if (ctx->debugging) { + OPENCL_SUCCEED_FATAL(clFinish(ctx->opencl.queue)); + $id:time_end = get_wall_time(); + long int $id:time_diff = $id:time_end - $id:time_start; + ctx->$id:(kernelRuntime kernel_name) += $id:time_diff; + ctx->$id:(kernelRuns kernel_name)++; + fprintf(stderr, "kernel %s runtime: %ldus\n", + $string:kernel_name, $id:time_diff); + } + }|] + where kernel_rank = length kernel_dims + kernel_dims' = map toInit kernel_dims + workgroup_dims' = map toInit workgroup_dims + total_elements = foldl multExp [C.cexp|1|] kernel_dims + + toInit e = [C.cinit|$exp:e|] + multExp x y = [C.cexp|$exp:x * $exp:y|] + + printKernelSize :: VName -> [C.Stm] + printKernelSize work_size = + intercalate [[C.cstm|fprintf(stderr, ", ");|]] $ + map (printKernelDim work_size) [0..kernel_rank-1] + printKernelDim global_work_size i = + [[C.cstm|fprintf(stderr, "%zu", $id:global_work_size[$int:i]);|]] diff --git a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs @@ -0,0 +1,401 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} +module Futhark.CodeGen.Backends.COpenCL.Boilerplate + ( generateBoilerplate + + , kernelRuntime + , kernelRuns + ) where + +import Data.FileEmbed +import qualified Data.Map as M +import qualified Language.C.Syntax as C +import qualified Language.C.Quote.OpenCL as C + +import Futhark.CodeGen.ImpCode.OpenCL +import qualified Futhark.CodeGen.Backends.GenericC as GC +import Futhark.CodeGen.OpenCL.Kernels +import Futhark.Util (chunk) + +generateBoilerplate :: String -> String -> [String] -> [PrimType] + -> M.Map VName (SizeClass, Name) + -> GC.CompilerM OpenCL () () +generateBoilerplate opencl_code opencl_prelude kernel_names types sizes = do + final_inits <- GC.contextFinalInits + + let (ctx_opencl_fields, ctx_opencl_inits, top_decls, later_top_decls) = + openClDecls kernel_names opencl_code opencl_prelude + + GC.earlyDecls top_decls + + let size_name_inits = map (\k -> [C.cinit|$string:(pretty k)|]) $ M.keys sizes + size_class_inits = map (\(c,_) -> [C.cinit|$string:(pretty c)|]) $ M.elems sizes + size_entry_points_inits = map (\(_,e) -> [C.cinit|$string:(pretty e)|]) $ M.elems sizes + num_sizes = M.size sizes + + GC.libDecl [C.cedecl|static const char *size_names[] = { $inits:size_name_inits };|] + GC.libDecl [C.cedecl|static const char *size_classes[] = { $inits:size_class_inits };|] + GC.libDecl [C.cedecl|static const char *size_entry_points[] = { $inits:size_entry_points_inits };|] + + GC.publicDef_ "get_num_sizes" GC.InitDecl $ \s -> + ([C.cedecl|int $id:s(void);|], + [C.cedecl|int $id:s(void) { + return $int:num_sizes; + }|]) + + GC.publicDef_ "get_size_name" GC.InitDecl $ \s -> + ([C.cedecl|const char* $id:s(int);|], + [C.cedecl|const char* $id:s(int i) { + return size_names[i]; + }|]) + + GC.publicDef_ "get_size_class" GC.InitDecl $ \s -> + ([C.cedecl|const char* $id:s(int);|], + [C.cedecl|const char* $id:s(int i) { + return size_classes[i]; + }|]) + + GC.publicDef_ "get_size_entry" GC.InitDecl $ \s -> + ([C.cedecl|const char* $id:s(int);|], + [C.cedecl|const char* $id:s(int i) { + return size_entry_points[i]; + }|]) + + let size_decls = map (\k -> [C.csdecl|size_t $id:k;|]) $ M.keys sizes + GC.libDecl [C.cedecl|struct sizes { $sdecls:size_decls };|] + cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { struct opencl_config opencl; + size_t sizes[$int:num_sizes]; + };|]) + + let size_value_inits = map (\i -> [C.cstm|cfg->sizes[$int:i] = 0;|]) [0..M.size sizes-1] + transposeBlockDim' = transposeBlockDim :: Int + GC.publicDef_ "context_config_new" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:cfg* $id:s(void);|], + [C.cedecl|struct $id:cfg* $id:s(void) { + struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg)); + if (cfg == NULL) { + return NULL; + } + + $stms:size_value_inits + opencl_config_init(&cfg->opencl, $int:num_sizes, + size_names, cfg->sizes, size_classes, size_entry_points); + + cfg->opencl.transpose_block_dim = $int:transposeBlockDim'; + return cfg; + }|]) + + GC.publicDef_ "context_config_free" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { + free(cfg); + }|]) + + GC.publicDef_ "context_config_set_debugging" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { + cfg->opencl.logging = cfg->opencl.debugging = flag; + }|]) + + GC.publicDef_ "context_config_set_logging" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag) { + cfg->opencl.logging = flag; + }|]) + + GC.publicDef_ "context_config_set_device" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { + set_preferred_device(&cfg->opencl, s); + }|]) + + GC.publicDef_ "context_config_set_platform" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *s) { + set_preferred_platform(&cfg->opencl, s); + }|]) + + GC.publicDef_ "context_config_dump_program_to" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + cfg->opencl.dump_program_to = path; + }|]) + + GC.publicDef_ "context_config_load_program_from" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *path) { + cfg->opencl.load_program_from = path; + }|]) + + GC.publicDef_ "context_config_set_default_group_size" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int size);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { + cfg->opencl.default_group_size = size; + cfg->opencl.default_group_size_changed = 1; + }|]) + + GC.publicDef_ "context_config_set_default_num_groups" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int num) { + cfg->opencl.default_num_groups = num; + }|]) + + GC.publicDef_ "context_config_set_default_tile_size" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { + cfg->opencl.default_tile_size = size; + cfg->opencl.default_tile_size_changed = 1; + }|]) + + GC.publicDef_ "context_config_set_default_threshold" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { + cfg->opencl.default_threshold = size; + }|]) + + GC.publicDef_ "context_config_set_size" GC.InitDecl $ \s -> + ([C.cedecl|int $id:s(struct $id:cfg* cfg, const char *size_name, size_t size_value);|], + [C.cedecl|int $id:s(struct $id:cfg* cfg, const char *size_name, size_t size_value) { + + for (int i = 0; i < $int:num_sizes; i++) { + if (strcmp(size_name, size_names[i]) == 0) { + cfg->sizes[i] = size_value; + return 0; + } + } + return 1; + }|]) + + (fields, init_fields) <- GC.contextContents + ctx <- GC.publicDef "context" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { + int detail_memory; + int debugging; + int logging; + typename lock_t lock; + char *error; + $sdecls:fields + $sdecls:ctx_opencl_fields + struct opencl_context opencl; + struct sizes sizes; + };|]) + + mapM_ GC.libDecl later_top_decls + let set_required_types = [ [C.cstm|required_types |= OPENCL_F64; |] + | FloatType Float64 `elem` types ] + set_sizes = zipWith (\i k -> [C.cstm|ctx->sizes.$id:k = cfg->sizes[$int:i];|]) + [(0::Int)..] $ M.keys sizes + + GC.libDecl [C.cedecl|static void init_context_early(struct $id:cfg *cfg, struct $id:ctx* ctx) { + typename cl_int error; + ctx->opencl.cfg = cfg->opencl; + ctx->detail_memory = cfg->opencl.debugging; + ctx->debugging = cfg->opencl.debugging; + ctx->logging = cfg->opencl.logging; + ctx->error = NULL; + create_lock(&ctx->lock); + + $stms:init_fields + $stms:ctx_opencl_inits + }|] + + GC.libDecl [C.cedecl|static int init_context_late(struct $id:cfg *cfg, struct $id:ctx* ctx, typename cl_program prog) { + typename cl_int error; + // Load all the kernels. + $stms:(map (loadKernelByName) kernel_names) + + $stms:final_inits + + $stms:set_sizes + + return 0; + }|] + + GC.publicDef_ "context_new" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|], + [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) { + struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx)); + if (ctx == NULL) { + return NULL; + } + + int required_types = 0; + $stms:set_required_types + + init_context_early(cfg, ctx); + typename cl_program prog = setup_opencl(&ctx->opencl, opencl_program, required_types); + init_context_late(cfg, ctx, prog); + return ctx; + }|]) + + GC.publicDef_ "context_new_with_command_queue" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue);|], + [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg, typename cl_command_queue queue) { + struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx)); + if (ctx == NULL) { + return NULL; + } + + int required_types = 0; + $stms:set_required_types + + init_context_early(cfg, ctx); + typename cl_program prog = setup_opencl_with_command_queue(&ctx->opencl, queue, opencl_program, required_types); + init_context_late(cfg, ctx, prog); + return ctx; + }|]) + + GC.publicDef_ "context_free" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:ctx* ctx);|], + [C.cedecl|void $id:s(struct $id:ctx* ctx) { + free_lock(&ctx->lock); + free(ctx); + }|]) + + GC.publicDef_ "context_sync" GC.InitDecl $ \s -> + ([C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { + ctx->error = OPENCL_SUCCEED_NONFATAL(clFinish(ctx->opencl.queue)); + return ctx->error != NULL; + }|]) + + GC.publicDef_ "context_get_error" GC.InitDecl $ \s -> + ([C.cedecl|char* $id:s(struct $id:ctx* ctx);|], + [C.cedecl|char* $id:s(struct $id:ctx* ctx) { + char* error = ctx->error; + ctx->error = NULL; + return error; + }|]) + + GC.publicDef_ "context_clear_caches" GC.InitDecl $ \s -> + ([C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { + ctx->error = OPENCL_SUCCEED_NONFATAL(opencl_free_all(&ctx->opencl)); + return ctx->error != NULL; + }|]) + + GC.publicDef_ "context_get_command_queue" GC.InitDecl $ \s -> + ([C.cedecl|typename cl_command_queue $id:s(struct $id:ctx* ctx);|], + [C.cedecl|typename cl_command_queue $id:s(struct $id:ctx* ctx) { + return ctx->opencl.queue; + }|]) + + mapM_ GC.debugReport $ openClReport kernel_names + +openClDecls :: [String] -> String -> String + -> ([C.FieldGroup], [C.Stm], [C.Definition], [C.Definition]) +openClDecls kernel_names opencl_program opencl_prelude = + (ctx_fields, ctx_inits, openCL_boilerplate, openCL_load) + where opencl_program_fragments = + -- Some C compilers limit the size of literal strings, so + -- chunk the entire program into small bits here, and + -- concatenate it again at runtime. + [ [C.cinit|$string:s|] | s <- chunk 2000 (opencl_prelude++opencl_program) ] + + ctx_fields = + [ [C.csdecl|int total_runs;|], + [C.csdecl|long int total_runtime;|] ] ++ + concat + [ [ [C.csdecl|typename cl_kernel $id:name;|] + , [C.csdecl|int $id:(kernelRuntime name);|] + , [C.csdecl|int $id:(kernelRuns name);|] + ] + | name <- kernel_names ] + + ctx_inits = + [ [C.cstm|ctx->total_runs = 0;|], + [C.cstm|ctx->total_runtime = 0;|] ] ++ + concat + [ [ [C.cstm|ctx->$id:(kernelRuntime name) = 0;|] + , [C.cstm|ctx->$id:(kernelRuns name) = 0;|] + ] + | name <- kernel_names ] + + openCL_load = [ + [C.cedecl| +void post_opencl_setup(struct opencl_context *ctx, struct opencl_device_option *option) { + $stms:(map sizeHeuristicsCode sizeHeuristicsTable) +}|]] + + openCL_h = $(embedStringFile "rts/c/opencl.h") + + program_fragments = opencl_program_fragments ++ [[C.cinit|NULL|]] + openCL_boilerplate = [C.cunit| + $esc:openCL_h + const char *opencl_program[] = {$inits:program_fragments};|] + +loadKernelByName :: String -> C.Stm +loadKernelByName name = [C.cstm|{ + ctx->$id:name = clCreateKernel(prog, $string:name, &error); + assert(error == 0); + if (ctx->debugging) { + fprintf(stderr, "Created kernel %s.\n", $string:name); + } + }|] + +kernelRuntime :: String -> String +kernelRuntime = (++"_total_runtime") + +kernelRuns :: String -> String +kernelRuns = (++"_runs") + +openClReport :: [String] -> [C.BlockItem] +openClReport names = report_kernels ++ [report_total] + where longest_name = foldl max 0 $ map length names + report_kernels = concatMap reportKernel names + format_string name = + let padding = replicate (longest_name - length name) ' ' + in unwords ["Kernel", + name ++ padding, + "executed %6d times, with average runtime: %6ldus\tand total runtime: %6ldus\n"] + reportKernel name = + let runs = kernelRuns name + total_runtime = kernelRuntime name + in [[C.citem| + fprintf(stderr, + $string:(format_string name), + ctx->$id:runs, + (long int) ctx->$id:total_runtime / (ctx->$id:runs != 0 ? ctx->$id:runs : 1), + (long int) ctx->$id:total_runtime); + |], + [C.citem|ctx->total_runtime += ctx->$id:total_runtime;|], + [C.citem|ctx->total_runs += ctx->$id:runs;|]] + + report_total = [C.citem| + if (ctx->debugging) { + fprintf(stderr, "Ran %d kernels with cumulative runtime: %6ldus\n", + ctx->total_runs, ctx->total_runtime); + } + |] + +sizeHeuristicsCode :: SizeHeuristic -> C.Stm +sizeHeuristicsCode (SizeHeuristic platform_name device_type which what) = + [C.cstm| + if ($exp:which' == 0 && + strstr(option->platform_name, $string:platform_name) != NULL && + option->device_type == $exp:(clDeviceType device_type)) { + $stm:get_size + }|] + where clDeviceType DeviceGPU = [C.cexp|CL_DEVICE_TYPE_GPU|] + clDeviceType DeviceCPU = [C.cexp|CL_DEVICE_TYPE_CPU|] + + which' = case which of + LockstepWidth -> [C.cexp|ctx->lockstep_width|] + NumGroups -> [C.cexp|ctx->cfg.default_num_groups|] + GroupSize -> [C.cexp|ctx->cfg.default_group_size|] + TileSize -> [C.cexp|ctx->cfg.default_tile_size|] + + get_size = case what of + HeuristicConst x -> + [C.cstm|$exp:which' = $int:x;|] + HeuristicDeviceInfo s -> + -- This only works for device info that fits in the variable. + let s' = "CL_DEVICE_" ++ s + in [C.cstm|clGetDeviceInfo(ctx->device, + $id:s', + sizeof($exp:which'), + &$exp:which', + NULL);|] diff --git a/src/Futhark/CodeGen/Backends/CSOpenCL.hs b/src/Futhark/CodeGen/Backends/CSOpenCL.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/CSOpenCL.hs @@ -0,0 +1,416 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.CodeGen.Backends.CSOpenCL + ( compileProg + ) where + +import Control.Monad +import Data.List + + +import Futhark.Error +import Futhark.Representation.ExplicitMemory (Prog, ExplicitMemory) +import Futhark.CodeGen.Backends.CSOpenCL.Boilerplate +import qualified Futhark.CodeGen.Backends.GenericCSharp as CS +import qualified Futhark.CodeGen.ImpCode.OpenCL as Imp +import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen +import Futhark.CodeGen.Backends.GenericCSharp.AST +import Futhark.CodeGen.Backends.GenericCSharp.Options +import Futhark.CodeGen.Backends.GenericCSharp.Definitions +import Futhark.Util.Pretty(pretty) +import Futhark.MonadFreshNames hiding (newVName') + + +compileProg :: MonadFreshNames m => Maybe String + -> Prog ExplicitMemory -> m (Either InternalError String) +compileProg module_name prog = do + res <- ImpGen.compileProg prog + case res of + Left err -> return $ Left err + Right (Imp.Program opencl_code opencl_prelude kernel_names types sizes prog') -> + Right <$> CS.compileProg + module_name + CS.emptyConstructor + imports + defines + operations + () + (generateBoilerplate opencl_code opencl_prelude kernel_names types sizes) + [] + [Imp.Space "device", Imp.Space "local", Imp.DefaultSpace] + cliOptions + prog' + + where operations :: CS.Operations Imp.OpenCL () + operations = CS.defaultOperations + { CS.opsCompiler = callKernel + , CS.opsWriteScalar = writeOpenCLScalar + , CS.opsReadScalar = readOpenCLScalar + , CS.opsAllocate = allocateOpenCLBuffer + , CS.opsCopy = copyOpenCLMemory + , CS.opsStaticArray = staticOpenCLArray + , CS.opsEntryInput = unpackArrayInput + , CS.opsEntryOutput = packArrayOutput + , CS.opsSyncRun = futharkSyncContext + } + imports = [ Using Nothing "System.Runtime.CompilerServices" + , Using Nothing "Cloo" + , Using Nothing "Cloo.Bindings" ] + defines = [ Escape csOpenCL + , Escape csMemoryOpenCL ] +cliOptions :: [Option] +cliOptions = [ Option { optionLongName = "platform" + , optionShortName = Just 'p' + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigSetPlatform(ref Cfg, optarg);"] + } + , Option { optionLongName = "device" + , optionShortName = Just 'd' + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigSetDevice(ref Cfg, optarg);"] + } + , Option { optionLongName = "dump-opencl" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigDumpProgramTo(ref Cfg, optarg);"] + } + , Option { optionLongName = "load-opencl" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigLoadProgramFrom(ref Cfg, optarg);"] + } + , Option { optionLongName = "debugging" + , optionShortName = Just 'D' + , optionArgument = NoArgument + , optionAction = [Escape "FutharkContextConfigSetDebugging(ref Cfg, true);"] + } + , Option { optionLongName = "default-group-size" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigSetDefaultGroupSize(ref Cfg, Convert.ToInt32(optarg));"] + } + , Option { optionLongName = "default-num-groups" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigSetDefaultNumGroups(ref Cfg, Convert.ToInt32(optarg));"] + } + , Option { optionLongName = "default-tile-size" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigSetDefaultTileSize(ref Cfg, Convert.ToInt32(optarg));"] + } + , Option { optionLongName = "default-threshold" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkContextConfigSetDefaultThreshold(ref Cfg, Convert.ToInt32(optarg));"] + } + , Option { optionLongName = "print-sizes" + , optionShortName = Nothing + , optionArgument = NoArgument + , optionAction = [Escape "FutharkConfigPrintSizes();"] + } + , Option { optionLongName = "size" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkConfigSetSize(ref Cfg, optarg);"] + } + ] + + +callKernel :: CS.OpCompiler Imp.OpenCL () +callKernel (Imp.GetSize v key) = + CS.stm $ Reassign (Var (CS.compileName v)) $ + Field (Var "Ctx.Sizes") $ pretty key + +callKernel (Imp.GetSizeMax v size_class) = + CS.stm $ Reassign (Var (CS.compileName v)) $ + Var $ "max_" ++ pretty size_class + +callKernel (Imp.HostCode c) = CS.compileCode c + +callKernel (Imp.LaunchKernel name args kernel_size workgroup_size) = do + kernel_size' <- mapM CS.compileExp kernel_size + let total_elements = foldl mult_exp (Integer 1) kernel_size' + let cond = BinOp "!=" total_elements (Integer 0) + workgroup_size' <- mapM CS.compileExp workgroup_size + body <- CS.collect $ launchKernel name kernel_size' workgroup_size' args + CS.stm $ If cond body [] + where mult_exp = BinOp "*" + +callKernel _ = undefined + +launchKernel :: String -> [CSExp] -> [CSExp] -> [Imp.KernelArg] -> CS.CompilerM op s () +launchKernel kernel_name kernel_dims workgroup_dims args = do + let kernel_name' = "Ctx."++kernel_name + args_stms <- zipWithM (processKernelArg kernel_name') [0..] args + + CS.stm $ Unsafe $ concat args_stms + + global_work_size <- newVName' "GlobalWorkSize" + local_work_size <- newVName' "LocalWorkSize" + stop_watch <- newVName' "StopWatch" + time_diff <- newVName' "TimeDiff" + + let debugStartStmts = + map Exp $ [CS.consoleErrorWrite "Launching {0} with global work size [" [String kernel_name]] ++ + printKernelSize global_work_size ++ + [ CS.consoleErrorWrite "] and local work size [" []] ++ + printKernelSize local_work_size ++ + [ CS.consoleErrorWrite "].\n" [] + , CallMethod (Var stop_watch) (Var "Start") []] + + let ctx = (++) "Ctx." + let debugEndStmts = + [ Exp $ CS.simpleCall "OPENCL_SUCCEED" [ + CS.simpleCall "CL10.Finish" + [Var "Ctx.OpenCL.Queue"]] + , Exp $ CallMethod (Var stop_watch) (Var "Stop") [] + , Assign (Var time_diff) $ asMicroseconds (Var stop_watch) + , AssignOp "+" (Var $ ctx $ kernelRuntime kernel_name) (Var time_diff) + , AssignOp "+" (Var $ ctx $ kernelRuns kernel_name) (Integer 1) + , Exp $ CS.consoleErrorWriteLine "kernel {0} runtime: {1}" [String kernel_name, Var time_diff] + ] + + + CS.stm $ If (BinOp "!=" total_elements (Integer 0)) + ([ Assign (Var global_work_size) (Collection "IntPtr[]" $ map CS.toIntPtr kernel_dims) + , Assign (Var local_work_size) (Collection "IntPtr[]" $ map CS.toIntPtr workgroup_dims) + , Assign (Var stop_watch) $ CS.simpleInitClass "Stopwatch" [] + , If (Var "Ctx.Debugging") debugStartStmts [] + ] + ++ + [ Exp $ CS.simpleCall "OPENCL_SUCCEED" [ + CS.simpleCall "CL10.EnqueueNDRangeKernel" + [ Var "Ctx.OpenCL.Queue", Var kernel_name', Integer kernel_rank, Null + , Var global_work_size, Var local_work_size, Integer 0, Null, Null]]] + ++ + [ If (Var "Ctx.Debugging") debugEndStmts [] ]) [] + finishIfSynchronous + + where processKernelArg :: String + -> Integer + -> Imp.KernelArg + -> CS.CompilerM op s [CSStmt] + processKernelArg kernel argnum (Imp.ValueKArg e bt) = do + let t = CS.compilePrimTypeToAST bt + tmp <- newVName' "kernelArg" + e' <- CS.compileExp e + err <- newVName' "setargErr" + let err_var = Var err + return [ AssignTyped t (Var tmp) (Just e') + , Assign err_var $ getKernelCall kernel argnum (CS.sizeOf t) (Addr $ Var tmp)] + + processKernelArg kernel argnum (Imp.MemKArg v) = do + err <- newVName' "setargErr" + dest <- newVName "kArgDest" + let err_var = Var err + return [ Fixed (Var $ CS.compileName dest) (Addr $ memblockFromMem v) + [ Assign err_var $ getKernelCall kernel argnum (CS.sizeOf $ Primitive IntPtrT) (Var $ CS.compileName dest)] + ] + + processKernelArg kernel argnum (Imp.SharedMemoryKArg (Imp.Count num_bytes)) = do + err <- newVName' "setargErr" + let err_var = Var err + num_bytes' <- CS.compileExp num_bytes + return [ Assign err_var $ getKernelCall kernel argnum num_bytes' Null ] + + kernel_rank = toInteger $ length kernel_dims + total_elements = foldl (BinOp "*") (Integer 1) kernel_dims + + printKernelSize :: String -> [CSExp] + printKernelSize work_size = + intersperse (CS.consoleErrorWrite ", " []) $ map (printKernelDim work_size) [0..kernel_rank-1] + + printKernelDim global_work_size i = + CS.consoleErrorWrite "{0}" [Index (Var global_work_size) (IdxExp (Integer $ toInteger i))] + + asMicroseconds watch = + BinOp "/" (Field watch "ElapsedTicks") + (BinOp "/" (Field (Var "TimeSpan") "TicksPerMillisecond") (Integer 1000)) + + + +getKernelCall :: String -> Integer -> CSExp -> CSExp -> CSExp +getKernelCall kernel arg_num size Null = + CS.simpleCall "CL10.SetKernelArg" [ Var kernel, Integer arg_num, CS.toIntPtr size, Var "Ctx.NULL"] +getKernelCall kernel arg_num size e = + CS.simpleCall "CL10.SetKernelArg" [ Var kernel, Integer arg_num, CS.toIntPtr size, CS.toIntPtr e] + +writeOpenCLScalar :: CS.WriteScalar Imp.OpenCL () +writeOpenCLScalar mem i bt "device" val = do + let bt' = CS.compilePrimTypeToAST bt + scalar <- newVName' "scalar" + ptr <- newVName' "ptr" + CS.stm $ Unsafe + [ AssignTyped bt' (Var scalar) (Just val) + , AssignTyped (PointerT VoidT) (Var ptr) (Just $ Addr $ Var scalar) + , Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer" + [ Var "Ctx.OpenCL.Queue", memblockFromMem mem, Bool True + ,CS.toIntPtr i,CS.toIntPtr $ CS.sizeOf bt',CS.toIntPtr $ Var ptr + , Integer 0, Null, Null] + ] + +writeOpenCLScalar _ _ _ space _ = + fail $ "Cannot write to '" ++ space ++ "' memory space." + +readOpenCLScalar :: CS.ReadScalar Imp.OpenCL () +readOpenCLScalar mem i bt "device" = do + val <- newVName' "read_res" + ptr <- newVName' "ptr" + let bt' = CS.compilePrimTypeToAST bt + CS.stm $ AssignTyped bt' (Var val) (Just $ CS.simpleInitClass (pretty bt') []) + CS.stm $ Unsafe + [ CS.assignScalarPointer (Var val) (Var ptr) + , Exp $ CS.simpleCall "CL10.EnqueueReadBuffer" + [ Var "Ctx.OpenCL.Queue", memblockFromMem mem , Bool True + , CS.toIntPtr i, CS.toIntPtr $ CS.sizeOf bt', CS.toIntPtr $ Var ptr + , Integer 0, Null, Null] + ] + return $ Var val + +readOpenCLScalar _ _ _ space = + fail $ "Cannot read from '" ++ space ++ "' memory space." + +computeErrCodeT :: CSType +computeErrCodeT = CustomT "ComputeErrorCode" + +allocateOpenCLBuffer :: CS.Allocate Imp.OpenCL () +allocateOpenCLBuffer mem size "device" = do + let mem' = CS.compileName mem + errcode <- CS.compileName <$> newVName "errCode" + CS.stm $ AssignTyped computeErrCodeT (Var errcode) Nothing + CS.stm $ Reassign (Var mem') (CS.simpleCall "MemblockAllocDevice" [Ref $ Var "Ctx", Var mem', size, String mem']) + +allocateOpenCLBuffer _ _ space = + fail $ "Cannot allocate in '" ++ space ++ "' space" + +copyOpenCLMemory :: CS.Copy Imp.OpenCL () +copyOpenCLMemory destmem destidx Imp.DefaultSpace srcmem srcidx (Imp.Space "device") nbytes _ = do + let destmem' = Var $ CS.compileName destmem + ptr <- newVName' "ptr" + CS.stm $ Fixed (Var ptr) (Addr $ Index destmem' $ IdxExp $ Integer 0) + [ ifNotZeroSize nbytes $ + Exp $ CS.simpleCall "CL10.EnqueueReadBuffer" + [ Var "Ctx.Opencl.Queue", memblockFromMem srcmem, Bool True + , CS.toIntPtr srcidx, nbytes,CS.toIntPtr $ Var ptr + , CS.toIntPtr destidx, Null, Null] + ] + +copyOpenCLMemory destmem destidx (Imp.Space "device") srcmem srcidx Imp.DefaultSpace nbytes _ = do + let srcmem' = CS.compileName srcmem + ptr <- newVName' "ptr" + CS.stm $ Fixed (Var ptr) (Addr $ Index (Var srcmem') $ IdxExp $ Integer 0) + [ ifNotZeroSize nbytes $ + Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer" + [ Var "Ctx.OpenCL.Queue", memblockFromMem destmem, Bool True + , CS.toIntPtr destidx, CS.toIntPtr nbytes, CS.toIntPtr $ Var ptr + , srcidx, Null, Null] + ] + +copyOpenCLMemory destmem destidx (Imp.Space "device") srcmem srcidx (Imp.Space "device") nbytes _ = do + CS.stm $ ifNotZeroSize nbytes $ + Exp $ CS.simpleCall "CL10.EnqueueCopyBuffer" + [ Var "Ctx.OpenCL.Queue", memblockFromMem srcmem, memblockFromMem destmem + , CS.toIntPtr srcidx, CS.toIntPtr destidx, CS.toIntPtr nbytes + , Integer 0, Null, Null] + finishIfSynchronous + +copyOpenCLMemory destmem destidx Imp.DefaultSpace srcmem srcidx Imp.DefaultSpace nbytes _ = + CS.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes + +copyOpenCLMemory _ _ destspace _ _ srcspace _ _= + error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace + +staticOpenCLArray :: CS.StaticArray Imp.OpenCL () +staticOpenCLArray name "device" t vs = do + let name' = CS.compileName name + CS.staticMemDecl $ AssignTyped (CustomT "OpenCLMemblock") (Var name') Nothing + + -- Create host-side C# array with intended values. + tmp_arr <- newVName' "tmpArr" + let t' = CS.compilePrimTypeToAST t + CS.staticMemDecl $ AssignTyped (Composite $ ArrayT t') (Var tmp_arr) (Just $ CreateArray t' $ map CS.compilePrimValue vs) + + -- Create memory block on the device. + ptr <- newVName' "ptr" + let size = Integer $ genericLength vs * Imp.primByteSize t + + CS.staticMemAlloc $ Reassign (Var name') (CS.simpleCall "EmptyMemblock" [Var "Ctx.EMPTY_MEM_HANDLE"]) + errcode <- CS.compileName <$> newVName "errCode" + CS.staticMemAlloc $ AssignTyped computeErrCodeT (Var errcode) Nothing + CS.staticMemAlloc $ Reassign (Var name') (CS.simpleCall "MemblockAllocDevice" [Ref $ Var "Ctx", Var name', size, String name']) + + -- Copy Numpy array to the device memory block. + CS.staticMemAlloc $ Unsafe [ + Fixed (Var ptr) (Addr $ Index (Var tmp_arr) $ IdxExp $ Integer 0) + [ ifNotZeroSize size $ + Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer" + [ Var "Ctx.OpenCL.Queue", memblockFromMem name, Bool True + , CS.toIntPtr (Integer 0),CS.toIntPtr size + , CS.toIntPtr $ Var ptr, Integer 0, Null, Null ] + ] + ] + +staticOpenCLArray _ space _ _ = + fail $ "CSOpenCL backend cannot create static array in memory space '" ++ space ++ "'" + +memblockFromMem :: VName -> CSExp +memblockFromMem mem = + let mem' = Var $ CS.compileName mem + in Field mem' "Mem" + +packArrayOutput :: CS.EntryOutput Imp.OpenCL () +packArrayOutput mem "device" bt ept dims = do + let size = foldr (BinOp "*") (Integer 1) dims' + let bt' = CS.compilePrimTypeToASText bt ept + let nbytes = BinOp "*" (CS.sizeOf bt') size + let createTuple = "createTuple_"++ pretty bt' + + return $ CS.simpleCall createTuple [ memblockFromMem mem, Var "Ctx.OpenCL.Queue", nbytes + , CreateArray (Primitive $ CSInt Int64T) dims'] + where dims' = map CS.compileDim dims + +packArrayOutput _ sid _ _ _ = + fail $ "Cannot return array from " ++ sid ++ " space." + +unpackArrayInput :: CS.EntryInput Imp.OpenCL () +unpackArrayInput mem memsize "device" t _ dims e = do + let size = foldr (BinOp "*") (Integer 1) dims' + let t' = CS.compilePrimTypeToAST t + let nbytes = BinOp "*" (CS.sizeOf t') size + zipWithM_ (CS.unpackDim e) dims [0..] + ptr <- pretty <$> newVName "ptr" + + CS.stm $ compileMemsize memsize nbytes + + let memsize' = CS.compileDim memsize + + CS.stm $ CS.getDefaultDecl (Imp.MemParam mem (Imp.Space "device")) + allocateOpenCLBuffer mem memsize' "device" + CS.stm $ Unsafe [Fixed (Var ptr) (Addr $ Index (Field e "Item1") $ IdxExp $ Integer 0) + [ ifNotZeroSize memsize' $ + Exp $ CS.simpleCall "CL10.EnqueueWriteBuffer" + [ Var "Ctx.OpenCL.Queue", memblockFromMem mem, Bool True + , CS.toIntPtr (Integer 0), CS.toIntPtr memsize', CS.toIntPtr (Var ptr) + , Integer 0, Null, Null] + ]] + + where dims' = map CS.compileDim dims + compileMemsize (Imp.VarSize v) nbytes = Assign (Var $ CS.compileName v) nbytes + compileMemsize _ _ = Pass + +unpackArrayInput _ _ sid _ _ _ _ = + fail $ "Cannot accept array from " ++ sid ++ " space." + +futharkSyncContext :: CSStmt +futharkSyncContext = Exp $ CS.simpleCall "FutharkContextSync" [] + +ifNotZeroSize :: CSExp -> CSStmt -> CSStmt +ifNotZeroSize e s = + If (BinOp "!=" e (Integer 0)) [s] [] + +finishIfSynchronous :: CS.CompilerM op s () +finishIfSynchronous = + CS.stm $ If (Var "Synchronous") [Exp $ CS.simpleCall "CL10.Finish" [Var "Ctx.OpenCL.Queue"]] [] + +newVName' :: MonadFreshNames f => String -> f String +newVName' s = CS.compileName <$> newVName s diff --git a/src/Futhark/CodeGen/Backends/CSOpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/CSOpenCL/Boilerplate.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/CSOpenCL/Boilerplate.hs @@ -0,0 +1,300 @@ +module Futhark.CodeGen.Backends.CSOpenCL.Boilerplate + ( generateBoilerplate + + , kernelRuntime + , kernelRuns + ) where + +import qualified Data.Map as M + +import Futhark.CodeGen.ImpCode.OpenCL hiding (Index, If) +import Futhark.CodeGen.Backends.GenericCSharp as CS +import Futhark.CodeGen.Backends.GenericCSharp.AST as AST +import Futhark.CodeGen.OpenCL.Kernels + + +intT, longT, stringT, intArrayT, stringArrayT :: CSType +intT = Primitive $ CSInt Int32T +longT = Primitive $ CSInt Int64T +stringT = Primitive StringT +intArrayT = Composite $ ArrayT intT +stringArrayT = Composite $ ArrayT stringT + +generateBoilerplate :: String -> String -> [String] -> [PrimType] + -> M.Map VName (SizeClass, Name) + -> CS.CompilerM OpenCL () () +generateBoilerplate opencl_code opencl_prelude kernel_names types sizes = do + final_inits <- CS.contextFinalInits + + let (opencl_fields, opencl_inits, top_decls, later_top_decls) = + openClDecls kernel_names opencl_code opencl_prelude + + CS.stm top_decls + + CS.stm $ AssignTyped stringArrayT (Var "SizeNames") + (Just $ Collection "string[]" (map (String . pretty) $ M.keys sizes)) + + CS.stm $ AssignTyped stringArrayT (Var "SizeClasses") + (Just $ Collection "string[]" (map (String . pretty . fst) $ M.elems sizes)) + + CS.stm $ AssignTyped stringArrayT (Var "SizeEntryPoints") + (Just $ Collection "string[]" (map (String . pretty . snd) $ M.elems sizes)) + + + let get_num_sizes = CS.publicName "GetNumSizes" + let get_size_name = CS.publicName "GetSizeName" + let get_size_class = CS.publicName "GetSizeClass" + let get_size_entry = CS.publicName "GetSizeEntry" + + + CS.stm $ CS.privateFunDef get_num_sizes intT [] + [ Return $ (Integer . toInteger) $ M.size sizes ] + CS.stm $ CS.privateFunDef get_size_name (Primitive StringT) [(intT, "i")] + [ Return $ Index (Var "SizeNames") (IdxExp $ Var "i") ] + CS.stm $ CS.privateFunDef get_size_class (Primitive StringT) [(intT, "i")] + [ Return $ Index (Var "SizeClasses") (IdxExp $ Var "i") ] + CS.stm $ CS.privateFunDef get_size_entry (Primitive StringT) [(intT, "i")] + [ Return $ Index (Var "SizeEntryPoints") (IdxExp $ Var "i") ] + + let cfg = CS.publicName "ContextConfig" + let new_cfg = CS.publicName "ContextConfigNew" + let cfg_set_debugging = CS.publicName "ContextConfigSetDebugging" + let cfg_set_device = CS.publicName "ContextConfigSetDevice" + let cfg_set_platform = CS.publicName "ContextConfigSetPlatform" + let cfg_dump_program_to = CS.publicName "ContextConfigDumpProgramTo" + let cfg_load_program_from = CS.publicName "ContextConfigLoadProgramFrom" + let cfg_set_default_group_size = CS.publicName "ContextConfigSetDefaultGroupSize" + let cfg_set_default_num_groups = CS.publicName "ContextConfigSetDefaultNumGroups" + let cfg_set_default_tile_size = CS.publicName "ContextConfigSetDefaultTileSize" + let cfg_set_default_threshold = CS.publicName "ContextConfigSetDefaultThreshold" + let cfg_set_size = CS.publicName "ContextConfigSetSize" + + CS.stm $ StructDef "Sizes" (map (\k -> (intT, pretty k)) $ M.keys sizes) + CS.stm $ StructDef cfg [ (CustomT "OpenCLConfig", "OpenCL") + , (intArrayT, "Sizes")] + + let tmp_cfg = Var "tmp_cfg" + CS.stm $ CS.privateFunDef new_cfg (CustomT cfg) [] + [ Assign tmp_cfg $ CS.simpleInitClass cfg [] + , Reassign (Field tmp_cfg "Sizes") (Collection "int[]" (replicate (M.size sizes) (Integer 0))) + , Exp $ CS.simpleCall "OpenCLConfigInit" [ Out $ Field tmp_cfg "OpenCL", (Integer . toInteger) $ M.size sizes + , Var "SizeNames", Field tmp_cfg "Sizes", Var "SizeClasses" ] + , Reassign (Field tmp_cfg "OpenCL.TransposeBlockDim") (Integer transposeBlockDim) + , Return tmp_cfg + ] + + CS.stm $ CS.privateFunDef cfg_set_debugging VoidT [(RefT $ CustomT cfg, "_cfg"),(Primitive BoolT, "flag")] + [Reassign (Var "_cfg.OpenCL.Debugging") (Var "flag")] + + CS.stm $ CS.privateFunDef cfg_set_device VoidT [(RefT $ CustomT cfg, "_cfg"),(stringT, "s")] + [Exp $ CS.simpleCall "SetPreferredDevice" [Ref $ Var "_cfg.OpenCL", Var "s"]] + + CS.stm $ CS.privateFunDef cfg_set_platform VoidT [(RefT $ CustomT cfg, "_cfg"),(stringT, "s")] + [Exp $ CS.simpleCall "SetPreferredPlatform" [Ref $ Var "_cfg.OpenCL", Var "s"]] + + CS.stm $ CS.privateFunDef cfg_dump_program_to VoidT [(RefT $ CustomT cfg, "_cfg"),(stringT, "path")] + [Reassign (Var "_cfg.OpenCL.DumpProgramTo") (Var "path")] + + CS.stm $ CS.privateFunDef cfg_load_program_from VoidT [(RefT $ CustomT cfg, "_cfg"),(stringT, "path")] + [Reassign (Var "_cfg.OpenCL.LoadProgramFrom") (Var "path")] + + CS.stm $ CS.privateFunDef cfg_set_default_group_size VoidT [(RefT $ CustomT cfg, "_cfg"),(intT, "size")] + [Reassign (Var "_cfg.OpenCL.DefaultGroupSize") (Var "size")] + + CS.stm $ CS.privateFunDef cfg_set_default_num_groups VoidT [(RefT $ CustomT cfg, "_cfg"),(intT, "num")] + [Reassign (Var "_cfg.OpenCL.DefaultNumGroups") (Var "num")] + + + CS.stm $ CS.privateFunDef cfg_set_default_tile_size VoidT [(RefT $ CustomT cfg, "_cfg"),(intT, "size")] + [Reassign (Var "_cfg.OpenCL.DefaultTileSize") (Var "size")] + + CS.stm $ CS.privateFunDef cfg_set_default_threshold VoidT [(RefT $ CustomT cfg, "_cfg"),(intT, "size")] + [Reassign (Var "_cfg.OpenCL.DefaultThreshold") (Var "size")] + + CS.stm $ CS.privateFunDef cfg_set_size (Primitive BoolT) [(RefT $ CustomT cfg, "_cfg") + , (stringT, "SizeName") + , (intT, "SizeValue")] + [ AST.For "i" ((Integer . toInteger) $ M.size sizes) + [ If (BinOp "==" (Var "SizeName") (Index (Var "SizeNames") (IdxExp (Var "i")))) + [ Reassign (Index (Var "_cfg.Sizes") (IdxExp (Var "i"))) (Var "SizeValue") + , Return (AST.Bool True)] [] + ] + , Return $ AST.Bool False ] + + + let ctx_ = CS.publicName "Context" + let new_ctx = CS.publicName "ContextNew" + let sync_ctx = CS.publicName "ContextSync" + + CS.stm $ StructDef ctx_ $ + [ (Primitive IntPtrT, "NULL") + , (CustomT "CLMemoryHandle", "EMPTY_MEM_HANDLE") + , (CustomT "OpenCLFreeList", "FreeList") + , (Primitive $ CSInt Int64T, "CurrentMemUsageDevice") + , (Primitive $ CSInt Int64T, "PeakMemUsageDevice") + , (Primitive BoolT, "DetailMemory") + , (Primitive BoolT, "Debugging") + , (CustomT "OpenCLContext", "OpenCL") + , (CustomT "Sizes", "Sizes") ] + ++ opencl_fields + + mapM_ CS.stm later_top_decls + + CS.addMemberDecl $ AssignTyped (CustomT cfg) (Var "Cfg") Nothing + CS.addMemberDecl $ AssignTyped (CustomT ctx_) (Var "Ctx") Nothing + + CS.beforeParse $ Reassign (Var "Cfg") $ CS.simpleCall new_cfg [] + CS.atInit $ Exp $ CS.simpleCall new_ctx [Var "Cfg"] + CS.atInit $ Reassign (Var "Ctx.EMPTY_MEM_HANDLE") $ CS.simpleCall "EmptyMemHandle" [Var "Ctx.OpenCL.Context"] + CS.atInit $ Reassign (Var "Ctx.FreeList") $ CS.simpleCall "OpenCLFreeListInit" [] + + CS.addMemberDecl $ AssignTyped (Primitive BoolT) (Var "Synchronous") (Just $ AST.Bool False) + + let set_required_types = [Reassign (Var "RequiredTypes") (AST.Bool True) + | FloatType Float64 `elem` types] + + set_sizes = zipWith (\i k -> Reassign (Field (Var "Ctx.Sizes") (pretty k)) + (Index (Var "Cfg.Sizes") (IdxExp $ (Integer . toInteger) i))) + [(0::Int)..] $ M.keys sizes + + + CS.stm $ CS.privateFunDef new_ctx VoidT [(CustomT cfg, "Cfg")] $ + [ AssignTyped (CustomT "ComputeErrorCode") (Var "error") Nothing + , Reassign (Var "Ctx.DetailMemory") (Var "Cfg.OpenCL.Debugging") + , Reassign (Var "Ctx.Debugging") (Var "Cfg.OpenCL.Debugging") + , Reassign (Var "Ctx.OpenCL.Cfg") (Var "Cfg.OpenCL")] + ++ opencl_inits ++ + [ Assign (Var "RequiredTypes") (AST.Bool False) ] + ++ set_required_types ++ + [ AssignTyped (CustomT "CLProgramHandle") (Var "prog") + (Just $ CS.simpleCall "SetupOpenCL" [ Ref $ Var "Ctx" + , Var "OpenCLProgram" + , Var "RequiredTypes"])] + ++ concatMap loadKernelByName kernel_names + ++ final_inits + ++ set_sizes + + CS.stm $ CS.privateFunDef sync_ctx intT [] + [ Exp $ CS.simpleCall "OPENCL_SUCCEED" [CS.simpleCall "CL10.Finish" [Var "Ctx.OpenCL.Queue"]] + , Return $ Integer 0 ] + + CS.debugReport $ openClReport kernel_names + + +openClDecls :: [String] -> String -> String + -> ([(CSType, String)], [CSStmt], CSStmt, [CSStmt]) +openClDecls kernel_names opencl_program opencl_prelude = + (ctx_fields, ctx_inits, openCL_boilerplate, openCL_load) + where ctx_fields = + [ (intT, "TotalRuns") + , (Primitive $ CSInt Int64T, "TotalRuntime")] + ++ concatMap (\name -> [(CustomT "CLKernelHandle", name) + ,(longT, kernelRuntime name) + ,(intT, kernelRuns name)]) kernel_names + + ctx_inits = + [ Reassign (Var $ ctx "TotalRuns") (Integer 0) + , Reassign (Var $ ctx "TotalRuntime") (Integer 0) ] + ++ concatMap (\name -> [ Reassign (Var $ (ctx . kernelRuntime) name) (Integer 0) + , Reassign (Var $ (ctx . kernelRuns) name) (Integer 0)] + ) kernel_names + + + futhark_context = CS.publicName "Context" + + openCL_load = [CS.privateFunDef "PostOpenCLSetup" VoidT + [ (RefT $ CustomT futhark_context, "Ctx") + , (RefT $ CustomT "OpenCLDeviceOption", "Option")] $ map sizeHeuristicsCode sizeHeuristicsTable] + + openCL_boilerplate = + AssignTyped stringArrayT (Var "OpenCLProgram") + (Just $ Collection "string[]" [String $ opencl_prelude ++ opencl_program]) + +loadKernelByName :: String -> [CSStmt] +loadKernelByName name = + [ Reassign (Var $ ctx name) + (CS.simpleCall "CL10.CreateKernel" [Var "prog", String name, Out $ Var "error"]) + , AST.Assert (BinOp "==" (Var "error") (Integer 0)) [] + , If (Var "Ctx.Debugging") + [Exp $ consoleErrorWriteLine "Created kernel {0}" [Var $ ctx name]] + [] + ] + +kernelRuntime :: String -> String +kernelRuntime = (++"_TotalRuntime") + +kernelRuns :: String -> String +kernelRuns = (++"_Runs") + +openClReport :: [String] -> CSStmt +openClReport names = + If (Var "Ctx.Debugging") (report_kernels ++ [report_total]) [] + where longest_name = foldl max 0 $ map length names + report_kernels = map reportKernel names + format_string name = + let padding = replicate (longest_name - length name) ' ' + in unwords ["Kernel", + name ++ padding, + "executed {0} times, with average runtime: {1}\tand total runtime: {2}"] + reportKernel name = + let runs = ctx $ kernelRuns name + total_runtime = ctx $ kernelRuntime name + in If (BinOp "!=" (Var runs) (Integer 0)) + [Exp $ consoleErrorWriteLine (format_string name) + [ Var runs + , Ternary (BinOp "!=" + (BinOp "/" + (Cast (Primitive $ CSInt Int64T) (Var total_runtime)) + (Var runs)) + (Integer 0)) + (Var runs) (Integer 1) + , Cast (Primitive $ CSInt Int64T) $ Var total_runtime] + , AssignOp "+" (Var $ ctx "TotalRuntime") (Var total_runtime) + , AssignOp "+" (Var $ ctx "TotalRuns") (Var runs) + ] [] + + ran_text = "Ran {0} kernels with cumulative runtime: {1}" + report_total = Exp $ consoleErrorWriteLine ran_text [ Var $ ctx "TotalRuns" + , Var $ ctx "TotalRuntime"] + +sizeHeuristicsCode :: SizeHeuristic -> CSStmt +sizeHeuristicsCode (SizeHeuristic platform_name device_type which what) = + let which'' = BinOp "==" which' (Integer 0) + option_contains_platform_name = CS.simpleCall "Option.PlatformName.Contains" [String platform_name] + option_contains_device_type = BinOp "==" (Var "Option.DeviceType") (Var $ clDeviceType device_type) + in If (BinOp "&&" which'' + (BinOp "&&" option_contains_platform_name + option_contains_device_type)) + [ get_size ] [] + + where clDeviceType DeviceGPU = "ComputeDeviceTypes.Gpu" + clDeviceType DeviceCPU = "ComputeDeviceTypes.Cpu" + + which' = case which of + LockstepWidth -> Var "Ctx.OpenCL.LockstepWidth" + NumGroups -> Var "Ctx.OpenCL.Cfg.DefaultNumGroups" + GroupSize -> Var "Ctx.OpenCL.Cfg.DefaultGroupSize" + TileSize -> Var "Ctx.OpenCL.Cfg.DefaultTileSize" + + get_size = case what of + HeuristicConst x -> + Reassign which' (Integer $ toInteger x) + + HeuristicDeviceInfo _ -> + -- This only works for device info that fits in the variable. + Unsafe + [ + Fixed (Var "ptr") (Addr which') + [ + Exp $ CS.simpleCall "CL10.GetDeviceInfo" + [ Var "Ctx.OpenCL.Device", Var "ComputeDeviceInfo.MaxComputeUnits" + , CS.simpleCall "new IntPtr" [CS.simpleCall "Marshal.SizeOf" [which']] + , CS.toIntPtr $ Var "ptr", Out ctxNULL ] + ] + ] + +ctx :: String -> String +ctx = (++) "Ctx." + +ctxNULL :: CSExp +ctxNULL = Var "Ctx.NULL" diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -0,0 +1,1901 @@ +{-# LANGUAGE QuasiQuotes, GeneralizedNewtypeDeriving, TypeSynonymInstances, FlexibleInstances #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +-- | C code generator framework. +module Futhark.CodeGen.Backends.GenericC + ( compileProg + , CParts(..) + , asLibrary + , asExecutable + + -- * Pluggable compiler + , Operations (..) + , defaultOperations + , OpCompiler + + , PointerQuals + , MemoryType + , WriteScalar + , writeScalarPointerWithQuals + , ReadScalar + , readScalarPointerWithQuals + , Allocate + , Deallocate + , Copy + , StaticArray + + -- * Monadic compiler interface + , CompilerM + , CompilerState (compUserState) + , getUserState + , putUserState + , modifyUserState + , contextContents + , contextFinalInits + , runCompilerM + , blockScope + , compileFun + , compileCode + , compileExp + , compilePrimExp + , compilePrimValue + , compileExpToName + , dimSizeToExp + , rawMem + , item + , stm + , stms + , decl + , atInit + , headerDecl + , publicDef + , publicDef_ + , debugReport + , HeaderSection(..) + , libDecl + , earlyDecls + , publicName + , contextType + , contextField + + -- * Building Blocks + , primTypeToCType + , copyMemoryDefaultSpace + ) where + +import Control.Monad.Identity +import Control.Monad.State +import Control.Monad.Reader +import Control.Monad.Writer +import Control.Monad.RWS +import Data.Bits (xor, shiftR) +import Data.Char (ord, isDigit, isAlphaNum) +import qualified Data.Map.Strict as M +import qualified Data.DList as DL +import Data.List +import Data.Loc +import Data.Maybe +import Data.FileEmbed +import qualified Data.Semigroup as Sem +import Text.Printf + +import qualified Language.C.Syntax as C +import qualified Language.C.Quote.OpenCL as C + +import Futhark.CodeGen.ImpCode hiding (dimSizeToExp) +import Futhark.MonadFreshNames +import Futhark.CodeGen.Backends.SimpleRepresentation +import Futhark.CodeGen.Backends.GenericC.Options +import Futhark.Util (zEncodeString) +import Futhark.Representation.AST.Attributes (isBuiltInFunction, builtInFunctions) + + +data CompilerState s = CompilerState { + compTypeStructs :: [([Type], (C.Type, C.Definition))] + , compArrayStructs :: [((C.Type, Int), (C.Type, [C.Definition]))] + , compOpaqueStructs :: [(String, (C.Type, [C.Definition]))] + , compEarlyDecls :: DL.DList C.Definition + , compInit :: [C.Stm] + , compNameSrc :: VNameSource + , compUserState :: s + , compHeaderDecls :: M.Map HeaderSection (DL.DList C.Definition) + , compLibDecls :: DL.DList C.Definition + , compCtxFields :: DL.DList (String, C.Type, Maybe C.Exp) + , compDebugItems :: DL.DList C.BlockItem + , compDeclaredMem :: [(VName,Space)] + } + +newCompilerState :: VNameSource -> s -> CompilerState s +newCompilerState src s = CompilerState { compTypeStructs = [] + , compArrayStructs = [] + , compOpaqueStructs = [] + , compEarlyDecls = mempty + , compInit = [] + , compNameSrc = src + , compUserState = s + , compHeaderDecls = mempty + , compLibDecls = mempty + , compCtxFields = mempty + , compDebugItems = mempty + , compDeclaredMem = mempty + } + +-- | In which part of the header file we put the declaration. This is +-- to ensure that the header file remains structured and readable. +data HeaderSection = ArrayDecl String + | OpaqueDecl String + | EntryDecl + | MiscDecl + | InitDecl + deriving (Eq, Ord) + +-- | A substitute expression compiler, tried before the main +-- compilation function. +type OpCompiler op s = op -> CompilerM op s () + +-- | The address space qualifiers for a pointer of the given type with +-- the given annotation. +type PointerQuals op s = String -> CompilerM op s [C.TypeQual] + +-- | The type of a memory block in the given memory space. +type MemoryType op s = SpaceId -> CompilerM op s C.Type + +-- | Write a scalar to the given memory block with the given index and +-- in the given memory space. +type WriteScalar op s = + C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s () + +-- | Read a scalar from the given memory block with the given index and +-- in the given memory space. +type ReadScalar op s = + C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp + +-- | Allocate a memory block of the given size and with the given tag +-- in the given memory space, saving a reference in the given variable +-- name. +type Allocate op s = C.Exp -> C.Exp -> C.Exp -> SpaceId + -> CompilerM op s () + +-- | De-allocate the given memory block with the given tag, which is +-- in the given memory space. +type Deallocate op s = C.Exp -> C.Exp -> SpaceId -> CompilerM op s () + +-- | Create a static array of values - initialised at load time. +type StaticArray op s = VName -> SpaceId -> PrimType -> [PrimValue] -> CompilerM op s () + +-- | Copy from one memory block to another. +type Copy op s = C.Exp -> C.Exp -> Space -> + C.Exp -> C.Exp -> Space -> + C.Exp -> + CompilerM op s () + +data Operations op s = + Operations { opsWriteScalar :: WriteScalar op s + , opsReadScalar :: ReadScalar op s + , opsAllocate :: Allocate op s + , opsDeallocate :: Deallocate op s + , opsCopy :: Copy op s + , opsStaticArray :: StaticArray op s + + , opsMemoryType :: MemoryType op s + , opsCompiler :: OpCompiler op s + + , opsFatMemory :: Bool + -- ^ If true, use reference counting. Otherwise, bare + -- pointers. + } + +-- | A set of operations that fail for every operation involving +-- non-default memory spaces. Uses plain pointers and @malloc@ for +-- memory management. +defaultOperations :: Operations op s +defaultOperations = Operations { opsWriteScalar = defWriteScalar + , opsReadScalar = defReadScalar + , opsAllocate = defAllocate + , opsDeallocate = defDeallocate + , opsCopy = defCopy + , opsStaticArray = defStaticArray + , opsMemoryType = defMemoryType + , opsCompiler = defCompiler + , opsFatMemory = True + } + where defWriteScalar _ _ _ _ _ = + fail "Cannot write to non-default memory space because I am dumb" + defReadScalar _ _ _ _ = + fail "Cannot read from non-default memory space" + defAllocate _ _ _ = + fail "Cannot allocate in non-default memory space" + defDeallocate _ _ = + fail "Cannot deallocate in non-default memory space" + defCopy destmem destoffset DefaultSpace srcmem srcoffset DefaultSpace size = + copyMemoryDefaultSpace destmem destoffset srcmem srcoffset size + defCopy _ _ _ _ _ _ _ = + fail "Cannot copy to or from non-default memory space" + defStaticArray _ _ _ _ = + fail "Cannot create static array in non-default memory space" + defMemoryType _ = + fail "Has no type for non-default memory space" + defCompiler _ = + fail "The default compiler cannot compile extended operations" + +data CompilerEnv op s = CompilerEnv { + envOperations :: Operations op s + , envFtable :: M.Map Name [Type] + } + +newtype CompilerAcc op s = CompilerAcc { + accItems :: DL.DList C.BlockItem + } + +instance Sem.Semigroup (CompilerAcc op s) where + CompilerAcc items1 <> CompilerAcc items2 = + CompilerAcc (items1<>items2) + +instance Monoid (CompilerAcc op s) where + mempty = CompilerAcc mempty + mappend = (Sem.<>) + +envOpCompiler :: CompilerEnv op s -> OpCompiler op s +envOpCompiler = opsCompiler . envOperations + +envMemoryType :: CompilerEnv op s -> MemoryType op s +envMemoryType = opsMemoryType . envOperations + +envReadScalar :: CompilerEnv op s -> ReadScalar op s +envReadScalar = opsReadScalar . envOperations + +envWriteScalar :: CompilerEnv op s -> WriteScalar op s +envWriteScalar = opsWriteScalar . envOperations + +envAllocate :: CompilerEnv op s -> Allocate op s +envAllocate = opsAllocate . envOperations + +envDeallocate :: CompilerEnv op s -> Deallocate op s +envDeallocate = opsDeallocate . envOperations + +envCopy :: CompilerEnv op s -> Copy op s +envCopy = opsCopy . envOperations + +envStaticArray :: CompilerEnv op s -> StaticArray op s +envStaticArray = opsStaticArray . envOperations + +envFatMemory :: CompilerEnv op s -> Bool +envFatMemory = opsFatMemory . envOperations + +newCompilerEnv :: Functions op -> Operations op s + -> CompilerEnv op s +newCompilerEnv (Functions funs) ops = + CompilerEnv { envOperations = ops + , envFtable = ftable <> builtinFtable + } + where ftable = M.fromList $ map funReturn funs + funReturn (name, fun) = + (name, paramsTypes $ functionOutput fun) + builtinFtable = + M.map (map Scalar . snd) builtInFunctions + +tupleDefinitions, arrayDefinitions, opaqueDefinitions :: CompilerState s -> [C.Definition] +tupleDefinitions = map (snd . snd) . compTypeStructs +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 + where isArrayDecl ArrayDecl{} = True + isArrayDecl _ = False +opaqueDecls = concatMap (DL.toList . snd) . filter (isOpaqueDecl . fst) . M.toList . compHeaderDecls + where isOpaqueDecl OpaqueDecl{} = True + isOpaqueDecl _ = False +entryDecls = concatMap (DL.toList . snd) . filter ((==EntryDecl) . fst) . M.toList . compHeaderDecls +miscDecls = concatMap (DL.toList . snd) . filter ((==MiscDecl) . fst) . M.toList . compHeaderDecls + +contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm]) +contextContents = do + (field_names, field_types, field_values) <- gets $ unzip3 . DL.toList . compCtxFields + let fields = [ [C.csdecl|$ty:ty $id:name;|] + | (name, ty) <- zip field_names field_types ] + init_fields = [ [C.cstm|ctx->$id:name = $exp:e;|] + | (name, Just e) <- zip field_names field_values ] + return (fields, init_fields) + +contextFinalInits :: CompilerM op s [C.Stm] +contextFinalInits = gets compInit + +newtype CompilerM op s a = CompilerM (RWS + (CompilerEnv op s) + (CompilerAcc op s) + (CompilerState s) a) + deriving (Functor, Applicative, Monad, + MonadState (CompilerState s), + MonadReader (CompilerEnv op s), + MonadWriter (CompilerAcc op s)) + +instance MonadFreshNames (CompilerM op s) where + getNameSource = gets compNameSrc + putNameSource src = modify $ \s -> s { compNameSrc = src } + +runCompilerM :: Functions op -> Operations op s -> VNameSource -> s + -> CompilerM op s a + -> (a, CompilerState s) +runCompilerM prog ops src userstate (CompilerM m) = + let (x, s, _) = runRWS m (newCompilerEnv prog ops) (newCompilerState src userstate) + in (x, s) + +getUserState :: CompilerM op s s +getUserState = gets compUserState + +putUserState :: s -> CompilerM op s () +putUserState s = modify $ \compstate -> compstate { compUserState = s } + +modifyUserState :: (s -> s) -> CompilerM op s () +modifyUserState f = modify $ \compstate -> + compstate { compUserState = f $ compUserState compstate } + +atInit :: C.Stm -> CompilerM op s () +atInit x = modify $ \s -> + s { compInit = compInit s ++ [x] } + +collect :: CompilerM op s () -> CompilerM op s [C.BlockItem] +collect m = snd <$> collect' m + +collect' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) +collect' m = pass $ do + (x, w) <- listen m + return ((x, DL.toList $ accItems w), + const w { accItems = mempty}) + +item :: C.BlockItem -> CompilerM op s () +item x = tell $ mempty { accItems = DL.singleton x } + +instance C.ToIdent VName where + toIdent = C.toIdent . zEncodeString . pretty + +instance C.ToExp VName where + toExp v _ = [C.cexp|$id:v|] + +instance C.ToExp IntValue where + toExp (Int8Value v) = C.toExp v + toExp (Int16Value v) = C.toExp v + toExp (Int32Value v) = C.toExp v + toExp (Int64Value v) = C.toExp v + +instance C.ToExp FloatValue where + toExp (Float32Value v) = C.toExp v + toExp (Float64Value v) = C.toExp v + +instance C.ToExp PrimValue where + toExp (IntValue v) = C.toExp v + toExp (FloatValue v) = C.toExp v + toExp (BoolValue True) = C.toExp (1::Int8) + toExp (BoolValue False) = C.toExp (0::Int8) + toExp Checked = C.toExp (1::Int8) + +-- | Construct a publicly visible definition using the specified name +-- as the template. The first returned definition is put in the +-- header file, and the second is the implementation. Returns the public +-- name. +publicDef :: String -> HeaderSection -> (String -> (C.Definition, C.Definition)) + -> CompilerM op s String +publicDef s h f = do + s' <- publicName s + let (pub, priv) = f s' + headerDecl h pub + libDecl priv + return s' + +-- | As 'publicDef', but ignores the public name. +publicDef_ :: String -> HeaderSection -> (String -> (C.Definition, C.Definition)) + -> CompilerM op s () +publicDef_ s h f = void $ publicDef s h f + +headerDecl :: HeaderSection -> C.Definition -> CompilerM op s () +headerDecl sec def = modify $ \s -> + s { compHeaderDecls = M.unionWith (<>) (compHeaderDecls s) + (M.singleton sec (DL.singleton def)) } + +libDecl :: C.Definition -> CompilerM op s () +libDecl def = modify $ \s -> + s { compLibDecls = compLibDecls s <> DL.singleton def } + +earlyDecls :: [C.Definition] -> CompilerM op s () +earlyDecls def = modify $ \s -> + s { compEarlyDecls = compEarlyDecls s <> DL.fromList def } + +contextField :: String -> C.Type -> Maybe C.Exp -> CompilerM op s () +contextField name ty initial = modify $ \s -> + s { compCtxFields = compCtxFields s <> DL.singleton (name,ty,initial) } + +debugReport :: C.BlockItem -> CompilerM op s () +debugReport x = modify $ \s -> + s { compDebugItems = compDebugItems s <> DL.singleton x } + +stm :: C.Stm -> CompilerM op s () +stm (C.Block items _) = mapM_ item items +stm (C.Default s _) = stm s +stm s = item [C.citem|$stm:s|] + +stms :: [C.Stm] -> CompilerM op s () +stms = mapM_ stm + +decl :: C.InitGroup -> CompilerM op s () +decl x = item [C.citem|$decl:x;|] + +addrOf :: C.Exp -> C.Exp +addrOf e = [C.cexp|&$exp:e|] + +-- | Public names must have a consitent prefix. +publicName :: String -> CompilerM op s String +publicName s = return $ "futhark_" ++ s + +-- | The generated code must define a struct with this name. +contextType :: CompilerM op s C.Type +contextType = do + name <- publicName "context" + return [C.cty|struct $id:name|] + +memToCType :: Space -> CompilerM op s C.Type +memToCType space = do + refcount <- asks envFatMemory + if refcount + then return $ fatMemType space + else rawMemCType space + +rawMemCType :: Space -> CompilerM op s C.Type +rawMemCType DefaultSpace = return defaultMemBlockType +rawMemCType (Space sid) = join $ asks envMemoryType <*> pure sid + +fatMemType :: Space -> C.Type +fatMemType space = + [C.cty|struct $id:name|] + where name = case space of + DefaultSpace -> "memblock" + Space sid -> "memblock_" ++ sid + +fatMemSet :: Space -> String +fatMemSet DefaultSpace = "memblock_set" +fatMemSet (Space sid) = "memblock_set_" ++ sid + +fatMemAlloc :: Space -> String +fatMemAlloc DefaultSpace = "memblock_alloc" +fatMemAlloc (Space sid) = "memblock_alloc_" ++ sid + +fatMemUnRef :: Space -> String +fatMemUnRef DefaultSpace = "memblock_unref" +fatMemUnRef (Space sid) = "memblock_unref_" ++ sid + +rawMem :: C.ToExp a => a -> CompilerM op s C.Exp +rawMem v = rawMem' <$> asks envFatMemory <*> pure v + +rawMem' :: C.ToExp a => Bool -> a -> C.Exp +rawMem' True e = [C.cexp|$exp:e.mem|] +rawMem' False e = [C.cexp|$exp:e|] + +defineMemorySpace :: Space -> CompilerM op s (C.Definition, [C.Definition], C.BlockItem) +defineMemorySpace space = do + rm <- rawMemCType space + let structdef = + [C.cedecl|struct $id:sname { int *references; + $ty:rm mem; + typename int64_t size; + const char *desc; };|] + + contextField peakname [C.cty|typename int64_t|] $ Just [C.cexp|0|] + contextField usagename [C.cty|typename int64_t|] $ Just [C.cexp|0|] + + -- Unreferencing a memory block consists of decreasing its reference + -- count and freeing the corresponding memory if the count reaches + -- zero. + free <- case space of + Space sid -> do free_mem <- asks envDeallocate + collect $ free_mem [C.cexp|block->mem|] [C.cexp|block->desc|] sid + DefaultSpace -> return [[C.citem|free(block->mem);|]] + ctx_ty <- contextType + let unrefdef = [C.cedecl|static int $id:(fatMemUnRef space) ($ty:ctx_ty *ctx, $ty:mty *block, const char *desc) { + if (block->references != NULL) { + *(block->references) -= 1; + if (ctx->detail_memory) { + fprintf(stderr, $string:("Unreferencing block %s (allocated as %s) in %s: %d references remaining.\n"), + desc, block->desc, $string:spacedesc, *(block->references)); + } + if (*(block->references) == 0) { + ctx->$id:usagename -= block->size; + $items:free + free(block->references); + if (ctx->detail_memory) { + fprintf(stderr, "%lld bytes freed (now allocated: %lld bytes)\n", + (long long) block->size, (long long) ctx->$id:usagename); + } + } + block->references = NULL; + } + return 0; +}|] + + -- When allocating a memory block we initialise the reference count to 1. + alloc <- collect $ + case space of + DefaultSpace -> + stm [C.cstm|block->mem = (char*) malloc(size);|] + Space sid -> + join $ asks envAllocate <*> pure [C.cexp|block->mem|] <*> + pure [C.cexp|size|] <*> pure [C.cexp|desc|] <*> pure sid + let allocdef = [C.cedecl|static int $id:(fatMemAlloc space) ($ty:ctx_ty *ctx, $ty:mty *block, typename int64_t size, const char *desc) { + if (size < 0) { + panic(1, "Negative allocation of %lld bytes attempted for %s in %s.\n", + (long long)size, desc, $string:spacedesc, ctx->$id:usagename); + } + int ret = $id:(fatMemUnRef space)(ctx, block, desc); + $items:alloc + block->references = (int*) malloc(sizeof(int)); + *(block->references) = 1; + block->size = size; + block->desc = desc; + ctx->$id:usagename += size; + if (ctx->detail_memory) { + fprintf(stderr, "Allocated %lld bytes for %s in %s (now allocated: %lld bytes)", + (long long) size, + desc, $string:spacedesc, + (long long) ctx->$id:usagename); + } + if (ctx->$id:usagename > ctx->$id:peakname) { + ctx->$id:peakname = ctx->$id:usagename; + if (ctx->detail_memory) { + fprintf(stderr, " (new peak).\n"); + } + } else if (ctx->detail_memory) { + fprintf(stderr, ".\n"); + } + return ret; + }|] + + -- Memory setting - unreference the destination and increase the + -- count of the source by one. + let setdef = [C.cedecl|static int $id:(fatMemSet space) ($ty:ctx_ty *ctx, $ty:mty *lhs, $ty:mty *rhs, const char *lhs_desc) { + int ret = $id:(fatMemUnRef space)(ctx, lhs, lhs_desc); + (*(rhs->references))++; + *lhs = *rhs; + return ret; +} +|] + + let peakmsg = "Peak memory usage for " ++ spacedesc ++ ": %lld bytes.\n" + return (structdef, + [unrefdef, allocdef, setdef], + [C.citem|fprintf(stderr, $string:peakmsg, + (long long) ctx->$id:peakname);|]) + where mty = fatMemType space + (peakname, usagename, sname, spacedesc) = case space of + DefaultSpace -> ("peak_mem_usage_default", + "cur_mem_usage_default", + "memblock", + "default space") + Space sid -> ("peak_mem_usage_" ++ sid, + "cur_mem_usage_" ++ sid, + "memblock_" ++ sid, + "space '" ++ sid ++ "'") + +declMem :: VName -> Space -> CompilerM op s () +declMem name space = do + ty <- memToCType space + decl [C.cdecl|$ty:ty $id:name;|] + resetMem name + modify $ \s -> s { compDeclaredMem = (name, space) : compDeclaredMem s } + +resetMem :: C.ToExp a => a -> CompilerM op s () +resetMem mem = do + refcount <- asks envFatMemory + when refcount $ + stm [C.cstm|$exp:mem.references = NULL;|] + +setMem :: (C.ToExp a, C.ToExp b) => a -> b -> Space -> CompilerM op s () +setMem dest src space = do + refcount <- asks envFatMemory + let src_s = pretty $ C.toExp src noLoc + if refcount + then stm [C.cstm|if ($id:(fatMemSet space)(ctx, &$exp:dest, &$exp:src, + $string:src_s) != 0) { + return 1; + }|] + else stm [C.cstm|$exp:dest = $exp:src;|] + +unRefMem :: C.ToExp a => a -> Space -> CompilerM op s () +unRefMem mem space = do + refcount <- asks envFatMemory + let mem_s = pretty $ C.toExp mem noLoc + when refcount $ + stm [C.cstm|if ($id:(fatMemUnRef space)(ctx, &$exp:mem, $string:mem_s) != 0) { + return 1; + }|] + +allocMem :: (C.ToExp a, C.ToExp b) => + a -> b -> Space -> C.Stm -> CompilerM op s () +allocMem name size space on_failure = do + refcount <- asks envFatMemory + let name_s = pretty $ C.toExp name noLoc + if refcount + then stm [C.cstm|if ($id:(fatMemAlloc space)(ctx, &$exp:name, $exp:size, + $string:name_s)) { + $stm:on_failure + }|] + else alloc name + where alloc dest = case space of + DefaultSpace -> + stm [C.cstm|$exp:dest = (char*) malloc($exp:size);|] + Space sid -> + join $ asks envAllocate <*> rawMem name <*> + pure [C.cexp|$exp:size|] <*> pure [C.cexp|desc|] <*> pure sid + +primTypeInfo :: PrimType -> Signedness -> C.Exp +primTypeInfo (IntType it) t = case (it, t) of + (Int8, TypeUnsigned) -> [C.cexp|u8_info|] + (Int16, TypeUnsigned) -> [C.cexp|u16_info|] + (Int32, TypeUnsigned) -> [C.cexp|u32_info|] + (Int64, TypeUnsigned) -> [C.cexp|u64_info|] + (Int8, _) -> [C.cexp|i8_info|] + (Int16, _) -> [C.cexp|i16_info|] + (Int32, _) -> [C.cexp|i32_info|] + (Int64, _) -> [C.cexp|i64_info|] +primTypeInfo (FloatType Float32) _ = [C.cexp|f32_info|] +primTypeInfo (FloatType Float64) _ = [C.cexp|f64_info|] +primTypeInfo Bool _ = [C.cexp|bool_info|] +primTypeInfo Cert _ = [C.cexp|bool_info|] + +copyMemoryDefaultSpace :: C.Exp -> C.Exp -> C.Exp -> C.Exp -> C.Exp -> + CompilerM op s () +copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = + stm [C.cstm|memmove($exp:destmem + $exp:destidx, + $exp:srcmem + $exp:srcidx, + $exp:nbytes);|] + +paramsTypes :: [Param] -> [Type] +paramsTypes = map paramType + -- Let's hope we don't need the size for anything, because we are + -- just making something up. + where paramType (MemParam _ space) = Mem (ConstSize 0) space + paramType (ScalarParam _ t) = Scalar t + +--- Entry points. + +arrayName :: PrimType -> Signedness -> Int -> String +arrayName pt signed rank = + prettySigned (signed==TypeUnsigned) pt ++ "_" ++ show rank ++ "d" + +opaqueName :: String -> [ValueDesc] -> String +opaqueName s _ + | valid = "opaque_" ++ s + where valid = head s /= '_' && + not (isDigit $ head s) && + all ok s + ok c = isAlphaNum c || c == '_' +opaqueName s vds = "opaque_" ++ hash (zipWith xor [0..] $ map ord (s ++ concatMap p vds)) + where p (ScalarValue pt signed _) = + show (pt, signed) + p (ArrayValue _ _ space pt signed dims) = + show (space, pt, signed, length dims) + + -- FIXME: a stupid hash algorithm; may have collisions. + hash = printf "%x" . foldl xor 0 . map (iter . (*0x45d9f3b) . + iter . (*0x45d9f3b) . + iter . fromIntegral) + iter x = ((x::Word32) `shiftR` 16) `xor` x + +criticalSection :: [C.BlockItem] -> [C.BlockItem] +criticalSection items = [[C.citem|lock_lock(&ctx->lock);|]] <> + items <> + [[C.citem|lock_unlock(&ctx->lock);|]] + +arrayLibraryFunctions :: Space -> PrimType -> Signedness -> [DimSize] + -> CompilerM op s [C.Definition] +arrayLibraryFunctions space pt signed shape = do + let rank = length shape + pt' = signedPrimTypeToCType signed pt + name = arrayName pt signed rank + arr_name = "futhark_" ++ name + array_type = [C.cty|struct $id:arr_name|] + + new_array <- publicName $ "new_" ++ name + new_raw_array <- publicName $ "new_raw_" ++ name + free_array <- publicName $ "free_" ++ name + values_array <- publicName $ "values_" ++ name + values_raw_array <- publicName $ "values_raw_" ++ name + shape_array <- publicName $ "shape_" ++ name + + let shape_names = [ "dim"++show i | i <- [0..rank-1] ] + shape_params = [ [C.cparam|int $id:k|] | k <- shape_names ] + arr_size = cproduct [ [C.cexp|$id:k|] | k <- shape_names ] + arr_size_array = cproduct [ [C.cexp|arr->shape[$int:i]|] | i <- [0..rank-1] ] + copy <- asks envCopy + + arr_raw_mem <- rawMem [C.cexp|arr->mem|] + memty <- rawMemCType space + + let prepare_new = do + resetMem [C.cexp|arr->mem|] + allocMem [C.cexp|arr->mem|] [C.cexp|$exp:arr_size * sizeof($ty:pt')|] space + [C.cstm|return NULL;|] + forM_ [0..rank-1] $ \i -> + let dim_s = "dim"++show i + in stm [C.cstm|arr->shape[$int:i] = $id:dim_s;|] + + new_body <- collect $ do + prepare_new + copy arr_raw_mem [C.cexp|0|] space + [C.cexp|data|] [C.cexp|0|] DefaultSpace + [C.cexp|$exp:arr_size * sizeof($ty:pt')|] + + new_raw_body <- collect $ do + prepare_new + copy arr_raw_mem [C.cexp|0|] space + [C.cexp|data|] [C.cexp|offset|] space + [C.cexp|$exp:arr_size * sizeof($ty:pt')|] + + free_body <- collect $ unRefMem [C.cexp|arr->mem|] space + + values_body <- collect $ + copy [C.cexp|data|] [C.cexp|0|] DefaultSpace + arr_raw_mem [C.cexp|0|] space + [C.cexp|$exp:arr_size_array * sizeof($ty:pt')|] + + ctx_ty <- contextType + + headerDecl (ArrayDecl name) + [C.cedecl|struct $id:name;|] + headerDecl (ArrayDecl name) + [C.cedecl|$ty:array_type* $id:new_array($ty:ctx_ty *ctx, $ty:pt' *data, $params:shape_params);|] + headerDecl (ArrayDecl name) + [C.cedecl|$ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, $ty:memty data, int offset, $params:shape_params);|] + headerDecl (ArrayDecl name) + [C.cedecl|int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + headerDecl (ArrayDecl name) + [C.cedecl|int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data);|] + headerDecl (ArrayDecl name) + [C.cedecl|$ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + headerDecl (ArrayDecl name) + [C.cedecl|typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + + return [C.cunit| + $ty:array_type* $id:new_array($ty:ctx_ty *ctx, $ty:pt' *data, $params:shape_params) { + $ty:array_type* bad = NULL; + $ty:array_type *arr = malloc(sizeof($ty:array_type)); + if (arr == NULL) { + return bad; + } + $items:(criticalSection new_body) + return arr; + } + + $ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, $ty:memty data, int offset, + $params:shape_params) { + $ty:array_type* bad = NULL; + $ty:array_type *arr = malloc(sizeof($ty:array_type)); + if (arr == NULL) { + return bad; + } + $items:(criticalSection new_raw_body) + return arr; + } + + int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + $items:(criticalSection free_body) + free(arr); + return 0; + } + + int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data) { + $items:(criticalSection values_body) + return 0; + } + + $ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + return $exp:arr_raw_mem; + } + + typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + return arr->shape; + } + |] + +opaqueLibraryFunctions :: String -> [ValueDesc] + -> CompilerM op s [C.Definition] +opaqueLibraryFunctions desc vds = do + name <- publicName $ opaqueName desc vds + free_opaque <- publicName $ "free_" ++ opaqueName desc vds + + let opaque_type = [C.cty|struct $id:name|] + + freeComponent _ ScalarValue{} = + return () + freeComponent i (ArrayValue _ _ _ pt signed shape) = do + let rank = length shape + free_array <- publicName $ "free_" ++ arrayName pt signed rank + stm [C.cstm|if ((tmp = $id:free_array(ctx, obj->$id:(tupleField i))) != 0) { + ret = tmp; + }|] + + ctx_ty <- contextType + + free_body <- collect $ zipWithM_ freeComponent [0..] vds + + headerDecl (OpaqueDecl desc) + [C.cedecl|int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj);|] + + return [C.cunit| + int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj) { + int ret = 0, tmp; + $items:free_body + free(obj); + return ret; + } + |] + +valueDescToCType :: ValueDesc -> CompilerM op s C.Type +valueDescToCType (ScalarValue pt signed _) = + return $ signedPrimTypeToCType signed pt +valueDescToCType (ArrayValue _ _ 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 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 + +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 + where field vd@ScalarValue{} i = do + ct <- valueDescToCType vd + return [C.csdecl|$ty:ct $id:(tupleField i);|] + field vd i = do + ct <- valueDescToCType vd + return [C.csdecl|$ty:ct *$id:(tupleField i);|] + +externalValueToCType :: ExternalValue -> CompilerM op s C.Type +externalValueToCType (TransparentValue vd) = valueDescToCType vd +externalValueToCType (OpaqueValue desc vds) = opaqueToCType desc vds + +prepareEntryInputs :: [ExternalValue] -> CompilerM op s [C.Param] +prepareEntryInputs = zipWithM prepare [(0::Int)..] + where prepare pno (TransparentValue vd) = do + let pname = "in" ++ show pno + ty <- prepareValue [C.cexp|$id:pname|] vd + return [C.cparam|const $ty:ty $id:pname|] + + 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)|] + zipWithM_ prepareValue (zipWith field [0..] vds) vds + return [C.cparam|const $ty:ty *$id:pname|] + + 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 mem_size _ _ _ shape) = do + ty <- valueDescToCType vd + + stm [C.cstm|$exp:mem = $exp:src->mem;|] + case mem_size of + VarSize v -> stm [C.cstm|$id:v = $exp:src->mem.size;|] + ConstSize _ -> return () + + + let rank = length shape + maybeCopyDim (VarSize d) i = + Just [C.cstm|$id:d = $exp:src->shape[$int:i];|] + maybeCopyDim _ _ = Nothing + + stms $ catMaybes $ zipWith maybeCopyDim shape [0..rank-1] + + return [C.cty|$ty:ty*|] + +prepareEntryOutputs :: [ExternalValue] -> CompilerM op s [C.Param] +prepareEntryOutputs = zipWithM prepare [(0::Int)..] + where prepare pno (TransparentValue vd) = do + let pname = "out" ++ show pno + ty <- valueDescToCType vd + + case vd of + ArrayValue{} -> do + stm [C.cstm|assert((*$id:pname = malloc(sizeof($ty:ty))) != NULL);|] + prepareValue [C.cexp|*$id:pname|] vd + return [C.cparam|$ty:ty **$id:pname|] + ScalarValue{} -> do + prepareValue [C.cexp|*$id:pname|] vd + return [C.cparam|$ty:ty *$id:pname|] + + prepare pno (OpaqueValue desc vds) = do + let pname = "out" ++ show pno + ty <- opaqueToCType desc vds + vd_ts <- mapM valueDescToCType vds + + stm [C.cstm|assert((*$id:pname = malloc(sizeof($ty:ty))) != NULL);|] + + + forM_ (zip3 [0..] vd_ts vds) $ \(i,ct,vd) -> do + let field = [C.cexp|(*$id:pname)->$id:(tupleField i)|] + case vd of + ScalarValue{} -> return () + _ -> stm [C.cstm|assert(($exp:field = malloc(sizeof($ty:ct))) != NULL);|] + prepareValue field vd + + return [C.cparam|$ty:ty **$id:pname|] + + prepareValue dest (ScalarValue _ _ name) = + stm [C.cstm|$exp:dest = $id:name;|] + + prepareValue dest (ArrayValue mem _ _ _ _ shape) = do + stm [C.cstm|$exp:dest->mem = $id:mem;|] + + let rank = length shape + maybeCopyDim (ConstSize x) i = + [C.cstm|$exp:dest->shape[$int:i] = $int:x;|] + maybeCopyDim (VarSize d) i = + [C.cstm|$exp:dest->shape[$int:i] = $id:d;|] + stms $ zipWith maybeCopyDim shape [0..rank-1] + +onEntryPoint :: Name -> Function op + -> CompilerM op s (C.Definition, C.Definition, C.Initializer) +onEntryPoint fname function@(Function _ outputs inputs _ results args) = do + let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs + in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs + + inputdecls <- collect $ mapM_ stubParam inputs + outputdecls <- collect $ mapM_ stubParam outputs + + let entry_point_name = nameToString fname + entry_point_function_name <- publicName $ "entry_" ++ entry_point_name + + (entry_point_input_params, unpack_entry_inputs) <- + collect' $ prepareEntryInputs args + (entry_point_output_params, pack_entry_outputs) <- + collect' $ prepareEntryOutputs results + + (cli_entry_point, cli_init) <- cliEntryPoint fname function + + ctx_ty <- contextType + + headerDecl EntryDecl [C.cedecl|int $id:entry_point_function_name + ($ty:ctx_ty *ctx, + $params:entry_point_output_params, + $params:entry_point_input_params);|] + + return ([C.cedecl|int $id:entry_point_function_name + ($ty:ctx_ty *ctx, + $params:entry_point_output_params, + $params:entry_point_input_params) { + $items:inputdecls + $items:outputdecls + + lock_lock(&ctx->lock); + + $items:unpack_entry_inputs + + int ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args); + + if (ret == 0) { + $items:pack_entry_outputs + } + + lock_unlock(&ctx->lock); + + return ret; +} + |], + cli_entry_point, + cli_init) + where stubParam (MemParam name space) = + declMem name space + stubParam (ScalarParam name ty) = do + let ty' = primTypeToCType ty + decl [C.cdecl|$ty:ty' $id:name;|] + +--- CLI interface +-- +-- Our strategy for CLI entry points is to parse everything into +-- host memory ('DefaultSpace') and copy the result into host memory +-- after the entry point has returned. We have some ad-hoc frobbery +-- to copy the host-level memory blocks to another memory space if +-- necessary. This will break if the Futhark entry point uses +-- non-trivial index functions for its input or output. +-- +-- The idea here is to keep the nastyness in the wrapper, whilst not +-- messing up anything else. + +printPrimStm :: (C.ToExp a, C.ToExp b) => a -> b -> PrimType -> Signedness -> C.Stm +printPrimStm dest val bt ept = + [C.cstm|write_scalar($exp:dest, binary_output, &$exp:(primTypeInfo bt ept), &$exp:val);|] + +-- | Return a statement printing the given external value. +printStm :: ExternalValue -> C.Exp -> CompilerM op s C.Stm +printStm (OpaqueValue desc _) _ = + return [C.cstm|printf("#", $string:desc);|] +printStm (TransparentValue (ScalarValue bt ept _)) e = + return $ printPrimStm [C.cexp|stdout|] e bt ept +printStm (TransparentValue (ArrayValue _ _ _ bt ept shape)) e = do + values_array <- publicName $ "values_" ++ name + shape_array <- publicName $ "shape_" ++ name + let num_elems = cproduct [ [C.cexp|$id:shape_array(ctx, $exp:e)[$int:i]|] | i <- [0..rank-1] ] + return [C.cstm|{ + $ty:bt' *arr = calloc(sizeof($ty:bt'), $exp:num_elems); + assert(arr != NULL); + assert($id:values_array(ctx, $exp:e, arr) == 0); + write_array(stdout, binary_output, &$exp:(primTypeInfo bt ept), arr, + $id:shape_array(ctx, $exp:e), $int:rank); + free(arr); + }|] + where rank = length shape + bt' = primTypeToCType bt + name = arrayName bt ept rank + +readPrimStm :: C.ToExp a => a -> Int -> PrimType -> Signedness -> C.Stm +readPrimStm place i t ept = + [C.cstm|if (read_scalar(&$exp:(primTypeInfo t ept),&$exp:place) != 0) { + panic(1, "Error when reading input #%d of type %s (errno: %s).\n", + $int:i, + $exp:(primTypeInfo t ept).type_name, + strerror(errno)); + }|] + +readInputs :: [ExternalValue] -> CompilerM op s [(C.Stm, C.Stm, C.Stm, C.Exp)] +readInputs = zipWithM readInput [0..] + +readInput :: Int -> ExternalValue -> CompilerM op s (C.Stm, C.Stm, C.Stm, C.Exp) +readInput i (OpaqueValue desc _) = do + stm [C.cstm|panic(1, "Cannot read input #%d of type %s\n", $int:i, $string:desc);|] + return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|NULL|]) +readInput i (TransparentValue (ScalarValue t ept _)) = do + dest <- newVName "read_value" + item [C.citem|$ty:(primTypeToCType t) $id:dest;|] + stm $ readPrimStm dest i t ept + return ([C.cstm|;|], [C.cstm|;|], [C.cstm|;|], [C.cexp|$id:dest|]) +readInput i (TransparentValue vd@(ArrayValue _ _ _ t ept dims)) = do + dest <- newVName "read_value" + shape <- newVName "read_shape" + arr <- newVName "read_arr" + ty <- valueDescToCType vd + item [C.citem|$ty:ty *$id:dest;|] + + let t' = primTypeToCType t + rank = length dims + name = arrayName t ept rank + dims_exps = [ [C.cexp|$id:shape[$int:j]|] | j <- [0..rank-1] ] + dims_s = concat $ replicate rank "[]" + + new_array <- publicName $ "new_" ++ name + free_array <- publicName $ "free_" ++ name + + stm [C.cstm|{ + typename int64_t $id:shape[$int:rank]; + $ty:t' *$id:arr = NULL; + errno = 0; + if (read_array(&$exp:(primTypeInfo t ept), + (void**) &$id:arr, + $id:shape, + $int:(length dims)) + != 0) { + panic(1, "Cannot read input #%d of type %s%s (errno: %s).\n", + $int:i, + $string:dims_s, + $exp:(primTypeInfo t ept).type_name, + strerror(errno)); + } + }|] + + return ([C.cstm|assert(($exp:dest = $id:new_array(ctx, $id:arr, $args:dims_exps)) != 0);|], + [C.cstm|assert($id:free_array(ctx, $exp:dest) == 0);|], + [C.cstm|free($id:arr);|], + [C.cexp|$id:dest|]) + +prepareOutputs :: [ExternalValue] -> CompilerM op s [(C.Exp, C.Stm)] +prepareOutputs = mapM prepareResult + where prepareResult ev = do + ty <- externalValueToCType ev + result <- newVName "result" + + case ev of + TransparentValue ScalarValue{} -> do + item [C.citem|$ty:ty $id:result;|] + return ([C.cexp|$id:result|], [C.cstm|;|]) + TransparentValue (ArrayValue _ _ _ t ept dims) -> do + let name = arrayName t ept $ length dims + free_array <- publicName $ "free_" ++ name + item [C.citem|$ty:ty *$id:result;|] + return ([C.cexp|$id:result|], + [C.cstm|assert($id:free_array(ctx, $exp:result) == 0);|]) + OpaqueValue desc vds -> do + free_opaque <- publicName $ "free_" ++ opaqueName desc vds + item [C.citem|$ty:ty *$id:result;|] + return ([C.cexp|$id:result|], + [C.cstm|assert($id:free_opaque(ctx, $exp:result) == 0);|]) + +printResult :: [(ExternalValue,C.Exp)] -> CompilerM op s [C.Stm] +printResult vs = fmap concat $ forM vs $ \(v,e) -> do + p <- printStm v e + return [p, [C.cstm|printf("\n");|]] + +cliEntryPoint :: Name + -> FunctionT a + -> CompilerM op s (C.Definition, C.Initializer) +cliEntryPoint fname (Function _ _ _ _ results args) = do + ((pack_input, free_input, free_parsed, input_args), input_items) <- + collect' $ unzip4 <$> readInputs args + + ((output_vals, free_outputs), output_decls) <- + collect' $ unzip <$> prepareOutputs results + printstms <- printResult $ zip results output_vals + + ctx_ty <- contextType + sync_ctx <- publicName "context_sync" + error_ctx <- publicName "context_get_error" + + let entry_point_name = nameToString fname + cli_entry_point_function_name = "futrts_cli_entry_" ++ entry_point_name + entry_point_function_name <- publicName $ "entry_" ++ entry_point_name + + let run_it = [C.citems| + int r; + /* Run the program once. */ + $stms:pack_input + assert($id:sync_ctx(ctx) == 0); + t_start = get_wall_time(); + r = $id:entry_point_function_name(ctx, + $args:(map addrOf output_vals), + $args:input_args); + if (r != 0) { + panic(1, "%s", $id:error_ctx(ctx)); + } + assert($id:sync_ctx(ctx) == 0); + t_end = get_wall_time(); + long int elapsed_usec = t_end - t_start; + if (time_runs && runtime_file != NULL) { + fprintf(runtime_file, "%lld\n", (long long) elapsed_usec); + } + $stms:free_input + |] + + return ([C.cedecl|static void $id:cli_entry_point_function_name($ty:ctx_ty *ctx) { + typename int64_t t_start, t_end; + int time_runs; + + /* Declare and read input. */ + $items:input_items + $items:output_decls + + /* Warmup run */ + if (perform_warmup) { + time_runs = 0; + $items:run_it + $stms:free_outputs + } + time_runs = 1; + /* Proper run. */ + for (int run = 0; run < num_runs; run++) { + $items:run_it + if (run < num_runs-1) { + $stms:free_outputs + } + } + + /* Free the parsed input. */ + $stms:free_parsed + + /* Print the final result. */ + $stms:printstms + + $stms:free_outputs + } + |], + [C.cinit|{ .name = $string:entry_point_name, + .fun = $id:cli_entry_point_function_name }|] + ) + +benchmarkOptions :: [Option] +benchmarkOptions = + [ Option { optionLongName = "write-runtime-to" + , optionShortName = Just 't' + , optionArgument = RequiredArgument + , optionAction = set_runtime_file + } + , Option { optionLongName = "runs" + , optionShortName = Just 'r' + , optionArgument = RequiredArgument + , optionAction = set_num_runs + } + , Option { optionLongName = "debugging" + , optionShortName = Just 'D' + , optionArgument = NoArgument + , optionAction = [C.cstm|futhark_context_config_set_debugging(cfg, 1);|] + } + , Option { optionLongName = "log" + , optionShortName = Just 'L' + , optionArgument = NoArgument + , optionAction = [C.cstm|futhark_context_config_set_logging(cfg, 1);|] + } + , Option { optionLongName = "entry-point" + , optionShortName = Just 'e' + , optionArgument = RequiredArgument + , optionAction = [C.cstm|entry_point = optarg;|] + } + , Option { optionLongName = "binary-output" + , optionShortName = Just 'b' + , optionArgument = NoArgument + , optionAction = [C.cstm|binary_output = 1;|] + } + ] + where set_runtime_file = [C.cstm|{ + runtime_file = fopen(optarg, "w"); + if (runtime_file == NULL) { + panic(1, "Cannot open %s: %s\n", optarg, strerror(errno)); + } + }|] + set_num_runs = [C.cstm|{ + num_runs = atoi(optarg); + perform_warmup = 1; + if (num_runs <= 0) { + panic(1, "Need a positive number of runs, not %s\n", optarg); + } + }|] + +-- | The result of compilation to C is four parts, which can be put +-- together in various ways. The obvious way is to concatenate all of +-- them, which yields a CLI program. Another is to compile the +-- library part by itself, and use the header file to call into it. +data CParts = CParts { cHeader :: String + , cUtils :: String + -- ^ Utility definitions that must be visible + -- to both CLI and library parts. + , cCLI :: String + , cLib :: String + } + +-- | Produce header and implementation files. +asLibrary :: CParts -> (String, String) +asLibrary parts = (cHeader parts, cUtils parts <> cLib parts) + +-- | As executable with command-line interface. +asExecutable :: CParts -> String +asExecutable (CParts a b c d) = a <> b <> c <> d + +-- | Compile imperative program to a C program. Always uses the +-- function named "main" as entry point, so make sure it is defined. +compileProg :: MonadFreshNames m => + Operations op () + -> CompilerM op () () + -> String + -> [Space] + -> [Option] + -> Functions op + -> m CParts +compileProg ops extra header_extra spaces options prog@(Functions funs) = do + src <- getNameSource + let ((prototypes, definitions, entry_points), endstate) = + runCompilerM prog ops src () compileProg' + (entry_point_decls, cli_entry_point_decls, entry_point_inits) = + unzip3 entry_points + option_parser = generateOptionParser "parse_options" $ benchmarkOptions++options + + let headerdefs = [C.cunit| +$esc:("/*\n * Headers\n*/\n") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:(header_extra) + +$esc:("\n/*\n * Initialisation\n*/\n") +$edecls:(initDecls endstate) + +$esc:("\n/*\n * Arrays\n*/\n") +$edecls:(arrayDecls endstate) + +$esc:("\n/*\n * Opaque values\n*/\n") +$edecls:(opaqueDecls endstate) + +$esc:("\n/*\n * Entry points\n*/\n") +$edecls:(entryDecls endstate) + +$esc:("\n/*\n * Miscellaneous\n*/\n") +$edecls:(miscDecls endstate) + |] + + let utildefs = [C.cunit| +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +/* If NDEBUG is set, the assert() macro will do nothing. Since Futhark + (unfortunately) makes use of assert() for error detection (and even some + side effects), we want to avoid that. */ +$esc:("#undef NDEBUG") +$esc:("#include ") + +$esc:panic_h + +$esc:timing_h +|] + + let clidefs = [C.cunit| +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") + +$esc:values_h + +static int binary_output = 0; +static typename FILE *runtime_file; +static int perform_warmup = 0; +static int num_runs = 1; +static const char *entry_point = "main"; + +$func:option_parser + +$edecls:cli_entry_point_decls + +typedef void entry_point_fun(struct futhark_context*); + +struct entry_point_entry { + const char *name; + entry_point_fun *fun; +}; + +int main(int argc, char** argv) { + fut_progname = argv[0]; + + struct entry_point_entry entry_points[] = { + $inits:entry_point_inits + }; + + struct futhark_context_config *cfg = futhark_context_config_new(); + assert(cfg != NULL); + + int parsed_options = parse_options(cfg, argc, argv); + argc -= parsed_options; + argv += parsed_options; + + if (argc != 0) { + panic(1, "Excess non-option: %s\n", argv[0]); + } + + struct futhark_context *ctx = futhark_context_new(cfg); + assert (ctx != NULL); + + int num_entry_points = sizeof(entry_points) / sizeof(entry_points[0]); + entry_point_fun *entry_point_fun = NULL; + for (int i = 0; i < num_entry_points; i++) { + if (strcmp(entry_points[i].name, entry_point) == 0) { + entry_point_fun = entry_points[i].fun; + break; + } + } + + if (entry_point_fun == NULL) { + fprintf(stderr, "No entry point '%s'. Select another with --entry-point. Options are:\n", + entry_point); + for (int i = 0; i < num_entry_points; i++) { + fprintf(stderr, "%s\n", entry_points[i].name); + } + return 1; + } + + entry_point_fun(ctx); + + if (runtime_file != NULL) { + fclose(runtime_file); + } + + futhark_debugging_report(ctx); + + futhark_context_free(ctx); + futhark_context_config_free(cfg); + return 0; +} + |] + + let early_decls = DL.toList $ compEarlyDecls endstate + let lib_decls = DL.toList $ compLibDecls endstate + let libdefs = [C.cunit| +$esc:("#ifdef _MSC_VER\n#define inline __inline\n#endif") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") +$esc:("#include ") + +$esc:lock_h + +$edecls:early_decls + +$edecls:lib_decls + +$edecls:(tupleDefinitions endstate) + +$edecls:prototypes + +$edecls:builtin + +$edecls:(map funcToDef definitions) + +$edecls:(arrayDefinitions endstate) + +$edecls:(opaqueDefinitions endstate) + +$edecls:entry_point_decls + |] + + return $ CParts (pretty headerdefs) (pretty utildefs) (pretty clidefs) (pretty libdefs) + where compileProg' = do + (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces + + (prototypes, definitions) <- unzip <$> mapM compileFun funs + + mapM_ libDecl memstructs + entry_points <- mapM (uncurry onEntryPoint) $ filter (functionEntry . snd) funs + extra + mapM_ libDecl $ concat memfuns + debugreport <- gets $ DL.toList . compDebugItems + + ctx_ty <- contextType + headerDecl MiscDecl [C.cedecl|void futhark_debugging_report($ty:ctx_ty *ctx);|] + libDecl [C.cedecl|void futhark_debugging_report($ty:ctx_ty *ctx) { + if (ctx->detail_memory) { + $items:memreport + } + if (ctx->debugging) { + $items:debugreport + } +}|] + + return (prototypes, definitions, entry_points) + funcToDef func = C.FuncDef func loc + where loc = case func of + C.OldFunc _ _ _ _ _ _ l -> l + C.Func _ _ _ _ _ l -> l + + builtin = cIntOps ++ cFloat32Ops ++ cFloat64Ops ++ cFloatConvOps ++ + cFloat32Funs ++ cFloat64Funs + + panic_h = $(embedStringFile "rts/c/panic.h") + values_h = $(embedStringFile "rts/c/values.h") + timing_h = $(embedStringFile "rts/c/timing.h") + lock_h = $(embedStringFile "rts/c/lock.h") + +compileFun :: (Name, Function op) -> CompilerM op s (C.Definition, C.Func) +compileFun (fname, Function _ outputs inputs body _ _) = do + (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs + inparams <- mapM compileInput inputs + body' <- blockScope $ compileFunBody out_ptrs outputs body + ctx_ty <- contextType + return ([C.cedecl|static int $id:(funName fname)($ty:ctx_ty *ctx, + $params:outparams, $params:inparams);|], + [C.cfun|static int $id:(funName fname)($ty:ctx_ty *ctx, + $params:outparams, $params:inparams) { + $items:body' + return 0; +}|]) + where compileInput (ScalarParam name bt) = do + let ctp = primTypeToCType bt + return [C.cparam|$ty:ctp $id:name|] + compileInput (MemParam name space) = do + ty <- memToCType space + return [C.cparam|$ty:ty $id:name|] + + compileOutput (ScalarParam name bt) = do + let ctp = primTypeToCType bt + p_name <- newVName $ "out_" ++ baseString name + return ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|]) + compileOutput (MemParam name space) = do + ty <- memToCType space + p_name <- newVName $ baseString name ++ "_p" + return ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|]) + +compilePrimValue :: PrimValue -> C.Exp + +compilePrimValue (IntValue (Int8Value k)) = [C.cexp|$int:k|] +compilePrimValue (IntValue (Int16Value k)) = [C.cexp|$int:k|] +compilePrimValue (IntValue (Int32Value k)) = [C.cexp|$int:k|] +compilePrimValue (IntValue (Int64Value k)) = [C.cexp|$int:k|] + +compilePrimValue (FloatValue (Float64Value x)) + | isInfinite x = + if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|] + | isNaN x = + [C.cexp|NAN|] + | otherwise = + [C.cexp|$double:x|] +compilePrimValue (FloatValue (Float32Value x)) + | isInfinite x = + if x > 0 then [C.cexp|INFINITY|] else [C.cexp|-INFINITY|] + | isNaN x = + [C.cexp|NAN|] + | otherwise = + [C.cexp|$float:x|] + +compilePrimValue (BoolValue b) = + [C.cexp|$int:b'|] + where b' :: Int + b' = if b then 1 else 0 + +compilePrimValue Checked = + [C.cexp|0|] + +dimSizeToExp :: DimSize -> C.Exp +dimSizeToExp (ConstSize x) = [C.cexp|$int:x|] +dimSizeToExp (VarSize v) = [C.cexp|$exp:v|] + +derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp +derefPointer ptr i res_t = + [C.cexp|*(($ty:res_t)&($exp:ptr[$exp:i]))|] + +writeScalarPointerWithQuals :: PointerQuals op s -> WriteScalar op s +writeScalarPointerWithQuals quals_f dest i elemtype space vol v = do + quals <- quals_f space + let quals' = case vol of Volatile -> [C.ctyquals|volatile|] ++ quals + Nonvolatile -> quals + deref = derefPointer dest i + [C.cty|$tyquals:quals' $ty:elemtype*|] + stm [C.cstm|$exp:deref = $exp:v;|] + +readScalarPointerWithQuals :: PointerQuals op s -> ReadScalar op s +readScalarPointerWithQuals quals_f dest i elemtype space vol = do + quals <- quals_f space + let quals' = case vol of Volatile -> [C.ctyquals|volatile|] ++ quals + Nonvolatile -> quals + return $ derefPointer dest i [C.cty|$tyquals:quals' $ty:elemtype*|] + +compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName +compileExpToName _ _ (LeafExp (ScalarVar v) _) = + return v +compileExpToName desc t e = do + desc' <- newVName desc + e' <- compileExp e + decl [C.cdecl|$ty:(primTypeToCType t) $id:desc' = $e';|] + return desc' + +compileExp :: Exp -> CompilerM op s C.Exp + +compileExp = compilePrimExp compileLeaf + where compileLeaf (ScalarVar src) = + return [C.cexp|$id:src|] + + compileLeaf (Index src (Count iexp) restype DefaultSpace vol) = do + src' <- rawMem src + derefPointer src' + <$> compileExp iexp + <*> pure [C.cty|$tyquals:vol' $ty:(primTypeToCType restype)*|] + where vol' = case vol of Volatile -> [C.ctyquals|volatile|] + Nonvolatile -> [] + + compileLeaf (Index src (Count iexp) restype (Space space) vol) = + join $ asks envReadScalar + <*> rawMem src <*> compileExp iexp + <*> pure (primTypeToCType restype) <*> pure space <*> pure vol + + compileLeaf (SizeOf t) = + return [C.cexp|(sizeof($ty:t'))|] + where t' = primTypeToCType t + +-- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. +compilePrimExp :: Monad m => (v -> m C.Exp) -> PrimExp v -> m C.Exp + +compilePrimExp _ (ValueExp val) = + return $ compilePrimValue val + +compilePrimExp f (LeafExp v _) = + f v + +compilePrimExp f (UnOpExp Complement{} x) = do + x' <- compilePrimExp f x + return [C.cexp|~$exp:x'|] + +compilePrimExp f (UnOpExp Not{} x) = do + x' <- compilePrimExp f x + return [C.cexp|!$exp:x'|] + +compilePrimExp f (UnOpExp Abs{} x) = do + x' <- compilePrimExp f x + return [C.cexp|abs($exp:x')|] + +compilePrimExp f (UnOpExp (FAbs Float32) x) = do + x' <- compilePrimExp f x + return [C.cexp|(float)fabs($exp:x')|] + +compilePrimExp f (UnOpExp (FAbs Float64) x) = do + x' <- compilePrimExp f x + return [C.cexp|fabs($exp:x')|] + +compilePrimExp f (UnOpExp SSignum{} x) = do + x' <- compilePrimExp f x + return [C.cexp|($exp:x' > 0) - ($exp:x' < 0)|] + +compilePrimExp f (UnOpExp USignum{} x) = do + x' <- compilePrimExp f x + return [C.cexp|($exp:x' > 0) - ($exp:x' < 0) != 0|] + +compilePrimExp f (CmpOpExp cmp x y) = do + x' <- compilePrimExp f x + y' <- compilePrimExp f y + return $ case cmp of + CmpEq{} -> [C.cexp|$exp:x' == $exp:y'|] + + FCmpLt{} -> [C.cexp|$exp:x' < $exp:y'|] + FCmpLe{} -> [C.cexp|$exp:x' <= $exp:y'|] + + CmpLlt{} -> [C.cexp|$exp:x' < $exp:y'|] + CmpLle{} -> [C.cexp|$exp:x' <= $exp:y'|] + + _ -> [C.cexp|$id:(pretty cmp)($exp:x', $exp:y')|] + +compilePrimExp f (ConvOpExp conv x) = do + x' <- compilePrimExp f x + return [C.cexp|$id:(pretty conv)($exp:x')|] + +compilePrimExp f (BinOpExp bop x y) = do + x' <- compilePrimExp f x + y' <- compilePrimExp f y + return $ case bop of + Add{} -> [C.cexp|$exp:x' + $exp:y'|] + FAdd{} -> [C.cexp|$exp:x' + $exp:y'|] + Sub{} -> [C.cexp|$exp:x' - $exp:y'|] + FSub{} -> [C.cexp|$exp:x' - $exp:y'|] + Mul{} -> [C.cexp|$exp:x' * $exp:y'|] + FMul{} -> [C.cexp|$exp:x' * $exp:y'|] + FDiv{} -> [C.cexp|$exp:x' / $exp:y'|] + Xor{} -> [C.cexp|$exp:x' ^ $exp:y'|] + And{} -> [C.cexp|$exp:x' & $exp:y'|] + Or{} -> [C.cexp|$exp:x' | $exp:y'|] + Shl{} -> [C.cexp|$exp:x' << $exp:y'|] + LogAnd{} -> [C.cexp|$exp:x' && $exp:y'|] + LogOr{} -> [C.cexp|$exp:x' || $exp:y'|] + _ -> [C.cexp|$id:(pretty bop)($exp:x', $exp:y')|] + +compilePrimExp f (FunExp h args _) = do + args' <- mapM (compilePrimExp f) args + return [C.cexp|$id:(funName (nameFromString h))($args:args')|] + +compileCode :: Code op -> CompilerM op s () + +compileCode (Op op) = + join $ asks envOpCompiler <*> pure op + +compileCode Skip = return () + +compileCode (Comment s code) = do + items <- blockScope $ compileCode code + let comment = "// " ++ s + stm [C.cstm|$comment:comment + { $items:items } + |] + +compileCode (DebugPrint s _ e) = do + e' <- compileExp e + stm [C.cstm|if (ctx->debugging) { + fprintf(stderr, "%s: %d\n", $exp:s, (int)$exp:e'); + }|] + +compileCode c + | Just (name, t, e, c') <- declareAndSet c = do + let ct = primTypeToCType t + e' <- compileExp e + item [C.citem|$ty:ct $id:name = $exp:e';|] + compileCode c' + +compileCode (c1 :>>: c2) = compileCode c1 >> compileCode c2 + +compileCode (Assert e (ErrorMsg parts) (loc, locs)) = do + e' <- compileExp e + free_all_mem <- collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem + let onPart (ErrorString s) = return ("%s", [C.cexp|$string:s|]) + onPart (ErrorInt32 x) = ("%d",) <$> compileExp x + (formatstrs, formatargs) <- unzip <$> mapM onPart parts + let formatstr = "Error at %s:\n" <> concat formatstrs <> "\n" + stm [C.cstm|if (!$exp:e') { + ctx->error = msgprintf($string:formatstr, $string:stacktrace, $args:formatargs); + $items:free_all_mem + return 1; + }|] + where stacktrace = intercalate " -> " (reverse $ map locStr $ loc:locs) + +compileCode (Allocate name (Count e) space) = do + size <- compileExp e + allocMem name size space [C.cstm|return 1;|] + +compileCode (Free name space) = + unRefMem name space + +compileCode (For i it bound body) = do + let i' = C.toIdent i + it' = intTypeToCType it + bound' <- compileExp bound + body' <- blockScope $ compileCode body + stm [C.cstm|for ($ty:it' $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) { + $items:body' + }|] + +compileCode (While cond body) = do + cond' <- compileExp cond + body' <- blockScope $ compileCode body + stm [C.cstm|while ($exp:cond') { + $items:body' + }|] + +compileCode (If cond tbranch fbranch) = do + cond' <- compileExp cond + tbranch' <- blockScope $ compileCode tbranch + fbranch' <- blockScope $ compileCode fbranch + stm $ case (tbranch', fbranch') of + (_, []) -> + [C.cstm|if ($exp:cond') { $items:tbranch' }|] + ([], _) -> + [C.cstm|if (!($exp:cond')) { $items:fbranch' }|] + _ -> + [C.cstm|if ($exp:cond') { $items:tbranch' } else { $items:fbranch' }|] + +compileCode (Copy dest (Count destoffset) DefaultSpace src (Count srcoffset) DefaultSpace (Count size)) = do + destoffset' <- compileExp destoffset + srcoffset' <- compileExp srcoffset + size' <- compileExp size + dest' <- rawMem dest + src' <- rawMem src + stm [C.cstm|memmove($exp:dest' + $exp:destoffset', + $exp:src' + $exp:srcoffset', + $exp:size');|] + +compileCode (Copy dest (Count destoffset) destspace src (Count srcoffset) srcspace (Count size)) = do + copy <- asks envCopy + join $ copy + <$> rawMem dest <*> compileExp destoffset <*> pure destspace + <*> rawMem src <*> compileExp srcoffset <*> pure srcspace + <*> compileExp size + +compileCode (Write dest (Count idx) elemtype DefaultSpace vol elemexp) = do + dest' <- rawMem dest + deref <- derefPointer dest' + <$> compileExp idx + <*> pure [C.cty|$tyquals:vol' $ty:(primTypeToCType elemtype)*|] + elemexp' <- compileExp elemexp + stm [C.cstm|$exp:deref = $exp:elemexp';|] + where vol' = case vol of Volatile -> [C.ctyquals|volatile|] + Nonvolatile -> [] + +compileCode (Write dest (Count idx) elemtype (Space space) vol elemexp) = + join $ asks envWriteScalar + <*> rawMem dest + <*> compileExp idx + <*> pure (primTypeToCType elemtype) + <*> pure space + <*> pure vol + <*> compileExp elemexp + +compileCode (DeclareMem name space) = + declMem name space + +compileCode (DeclareScalar name t) = do + let ct = primTypeToCType t + decl [C.cdecl|$ty:ct $id:name;|] + +compileCode (DeclareArray name DefaultSpace t vs) = do + let ct = primTypeToCType t + vs' = [[C.cinit|$exp:(compilePrimValue v)|] | v <- vs] + name_realtype <- newVName $ baseString name ++ "_realtype" + libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs)] = {$inits:vs'};|] + -- Fake a memory block. + contextField (pretty name) + [C.cty|struct memblock|] $ + Just [C.cexp|(struct memblock){NULL, (char*)$id:name_realtype, 0}|] + item [C.citem|struct memblock $id:name = ctx->$id:name;|] + +compileCode (DeclareArray name (Space space) t vs) = + join $ asks envStaticArray <*> + pure name <*> pure space <*> pure t <*> pure vs + +-- For assignments of the form 'x = x OP e', we generate C assignment +-- operators to make the resulting code slightly nicer. This has no +-- effect on performance. +compileCode (SetScalar dest (BinOpExp op (LeafExp (ScalarVar x) _) y)) + | dest == x, Just f <- assignmentOperator op = do + y' <- compileExp y + stm [C.cstm|$exp:(f dest y');|] + +compileCode (SetScalar dest src) = do + src' <- compileExp src + stm [C.cstm|$id:dest = $exp:src';|] + +compileCode (SetMem dest src space) = + setMem dest src space + +compileCode (Call dests fname args) = do + args' <- mapM compileArg args + let out_args = [ [C.cexp|&$id:d|] | d <- dests ] + args'' | isBuiltInFunction fname = args' + | otherwise = [C.cexp|ctx|] : out_args ++ args' + case dests of + [dest] | isBuiltInFunction fname -> + stm [C.cstm|$id:dest = $id:(funName fname)($args:args'');|] + _ -> do + ret <- newVName "call_ret" + item [C.citem|int $id:ret = $id:(funName fname)($args:args'');|] + stm [C.cstm|assert($id:ret == 0);|] + where compileArg (MemArg m) = return [C.cexp|$exp:m|] + compileArg (ExpArg e) = compileExp e + +blockScope :: CompilerM op s () -> CompilerM op s [C.BlockItem] +blockScope = fmap snd . blockScope' + +blockScope' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) +blockScope' m = do + old_allocs <- gets compDeclaredMem + (x, items) <- pass $ do + (x, w) <- listen m + let items = DL.toList $ accItems w + return ((x, items), const mempty) + new_allocs <- gets $ filter (`notElem` old_allocs) . compDeclaredMem + modify $ \s -> s { compDeclaredMem = old_allocs } + releases <- collect $ mapM_ (uncurry unRefMem) new_allocs + return (x, items <> releases) + +compileFunBody :: [C.Exp] -> [Param] -> Code op -> CompilerM op s () +compileFunBody output_ptrs outputs code = do + mapM_ declareOutput outputs + compileCode code + zipWithM_ setRetVal' output_ptrs outputs + where declareOutput (MemParam name space) = + declMem name space + declareOutput (ScalarParam name pt) = do + let ctp = primTypeToCType pt + decl [C.cdecl|$ty:ctp $id:name;|] + + setRetVal' p (MemParam name space) = do + resetMem [C.cexp|*$exp:p|] + setMem [C.cexp|*$exp:p|] name space + setRetVal' p (ScalarParam name _) = + stm [C.cstm|*$exp:p = $id:name;|] + +declareAndSet :: Code op -> Maybe (VName, PrimType, Exp, Code op) +declareAndSet (DeclareScalar name t :>>: (SetScalar dest e :>>: c)) + | name == dest = Just (name, t, e, c) +declareAndSet ((DeclareScalar name t :>>: SetScalar dest e) :>>: c) + | name == dest = Just (name, t, e, c) +declareAndSet _ = Nothing + +assignmentOperator :: BinOp -> Maybe (VName -> C.Exp -> C.Exp) +assignmentOperator Add{} = Just $ \d e -> [C.cexp|$id:d += $exp:e|] +assignmentOperator Sub{} = Just $ \d e -> [C.cexp|$id:d -= $exp:e|] +assignmentOperator Mul{} = Just $ \d e -> [C.cexp|$id:d *= $exp:e|] +assignmentOperator _ = Nothing + +-- | Return an expression multiplying together the given expressions. +-- If an empty list is given, the expression @1@ is returned. +cproduct :: [C.Exp] -> C.Exp +cproduct [] = [C.cexp|1|] +cproduct (e:es) = foldl mult e es + where mult x y = [C.cexp|$exp:x * $exp:y|] diff --git a/src/Futhark/CodeGen/Backends/GenericC/Options.hs b/src/Futhark/CodeGen/Backends/GenericC/Options.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericC/Options.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE QuasiQuotes #-} +-- | This module defines a generator for @getopt_long@ based command +-- line argument parsing. Each option is associated with arbitrary C +-- code that will perform side effects, usually by setting some global +-- variables. +module Futhark.CodeGen.Backends.GenericC.Options + ( Option (..) + , OptionArgument (..) + , generateOptionParser + ) + where + +import Data.Maybe + +import qualified Language.C.Syntax as C +import qualified Language.C.Quote.C as C + +-- | Specification if a single command line option. The option must +-- have a long name, and may also have a short name. +-- +-- In the action, the option argument (if any) is stored as in the +-- @char*@-typed variable @optarg@. +data Option = Option { optionLongName :: String + , optionShortName :: Maybe Char + , optionArgument :: OptionArgument + , optionAction :: C.Stm + } + +-- | Whether an option accepts an argument. +data OptionArgument = NoArgument + | RequiredArgument + | OptionalArgument + +-- | Generate an option parser as a function of the given name, that +-- accepts the given command line options. The result is a function +-- that should be called with @argc@ and @argv@. The function returns +-- the number of @argv@ elements that have been processed. +-- +-- If option parsing fails for any reason, the entire process will +-- terminate with error code 1. +generateOptionParser :: String -> [Option] -> C.Func +generateOptionParser fname options = + [C.cfun|int $id:fname(struct futhark_context_config *cfg, int argc, char* const argv[]) { + int $id:chosen_option; + + static struct option long_options[] = { $inits:option_fields, {0, 0, 0, 0} }; + + while (($id:chosen_option = + getopt_long(argc, argv, $string:option_string, long_options, NULL)) != -1) { + $stms:option_applications + if ($id:chosen_option == ':') { + panic(-1, "Missing argument for option %s\n", argv[optind-1]); + } + if ($id:chosen_option == '?') { + panic(-1, "Unknown option %s\n", argv[optind-1]); + } + } + return optind; + } + |] + where chosen_option = "ch" + option_string = ':' : optionString options + option_applications = optionApplications chosen_option options + option_fields = optionFields options + +optionFields :: [Option] -> [C.Initializer] +optionFields = zipWith field [(1::Int)..] + where field i option = + [C.cinit| { $string:(optionLongName option), $id:arg, NULL, $int:i } |] + where arg = case optionArgument option of + NoArgument -> "no_argument" + RequiredArgument -> "required_argument" + OptionalArgument -> "optional_argument" + +optionApplications :: String -> [Option] -> [C.Stm] +optionApplications chosen_option = zipWith check [(1::Int)..] + where check i option = + [C.cstm|if ($exp:cond) $stm:(optionAction option)|] + where cond = case optionShortName option of + Nothing -> [C.cexp|$id:chosen_option == $int:i|] + Just c -> [C.cexp|($id:chosen_option == $int:i) || + ($id:chosen_option == $char:c)|] +optionString :: [Option] -> String +optionString = concat . mapMaybe optionStringChunk + where optionStringChunk option = do + short <- optionShortName option + return $ short : + case optionArgument option of + NoArgument -> "" + RequiredArgument -> ":" + OptionalArgument -> "::" diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp.hs b/src/Futhark/CodeGen/Backends/GenericCSharp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericCSharp.hs @@ -0,0 +1,1404 @@ +{-# LANGUAGE OverloadedStrings, GeneralizedNewtypeDeriving, LambdaCase #-} +{-# LANGUAGE TupleSections #-} +-- | A generic C# code generator which is polymorphic in the type +-- of the operations. Concretely, we use this to handle both +-- sequential and OpenCL C# code. +module Futhark.CodeGen.Backends.GenericCSharp + ( compileProg + , Constructor (..) + , emptyConstructor + + , assignScalarPointer + , toIntPtr + , compileName + , compileDim + , compileExp + , compileCode + , compilePrimValue + , compilePrimType + , compilePrimTypeExt + , compilePrimTypeToAST + , compilePrimTypeToASText + , contextFinalInits + , debugReport + + , Operations (..) + , defaultOperations + + , unpackDim + + , CompilerM (..) + , OpCompiler + , WriteScalar + , ReadScalar + , Allocate + , Copy + , StaticArray + , EntryOutput + , EntryInput + + , CompilerEnv(..) + , CompilerState(..) + , stm + , stms + , atInit + , staticMemDecl + , staticMemAlloc + , addMemberDecl + , beforeParse + , collect' + , collect + , simpleCall + , callMethod + , simpleInitClass + , parametrizedCall + + , copyMemoryDefaultSpace + , consoleErrorWrite + , consoleErrorWriteLine + , consoleWrite + , consoleWriteLine + + , publicName + , sizeOf + , privateFunDef + , publicFunDef + , getDefaultDecl + ) where + +import Control.Monad.Identity +import Control.Monad.State +import Control.Monad.Reader +import Control.Monad.Writer +import Control.Monad.RWS +import Control.Arrow((&&&)) +import Data.Maybe +import Data.List +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem + +import Futhark.Representation.Primitive hiding (Bool) +import Futhark.MonadFreshNames +import Futhark.Representation.AST.Syntax (Space(..)) +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.CodeGen.Backends.GenericCSharp.AST +import Futhark.CodeGen.Backends.GenericCSharp.Options +import Futhark.CodeGen.Backends.GenericCSharp.Definitions +import Futhark.Util.Pretty(pretty) +import Futhark.Util (zEncodeString) +import Futhark.Representation.AST.Attributes (builtInFunctions) +import Text.Printf (printf) + +-- | A substitute expression compiler, tried before the main +-- compilation function. +type OpCompiler op s = op -> CompilerM op s () + +-- | Write a scalar to the given memory block with the given index and +-- in the given memory space. +type WriteScalar op s = VName -> CSExp -> PrimType -> Imp.SpaceId -> CSExp + -> CompilerM op s () + +-- | Read a scalar from the given memory block with the given index and +-- in the given memory space. +type ReadScalar op s = VName -> CSExp -> PrimType -> Imp.SpaceId + -> CompilerM op s CSExp + +-- | Allocate a memory block of the given size in the given memory +-- space, saving a reference in the given variable name. +type Allocate op s = VName -> CSExp -> Imp.SpaceId + -> CompilerM op s () + +-- | Copy from one memory block to another. +type Copy op s = VName -> CSExp -> Imp.Space -> + VName -> CSExp -> Imp.Space -> + CSExp -> PrimType -> + CompilerM op s () + +-- | Create a static array of values - initialised at load time. +type StaticArray op s = VName -> Imp.SpaceId -> PrimType -> [PrimValue] -> CompilerM op s () + +-- | Construct the C# array being returned from an entry point. +type EntryOutput op s = VName -> Imp.SpaceId -> + PrimType -> Imp.Signedness -> + [Imp.DimSize] -> + CompilerM op s CSExp + +-- | Unpack the array being passed to an entry point. +type EntryInput op s = VName -> Imp.MemSize -> Imp.SpaceId -> + PrimType -> Imp.Signedness -> + [Imp.DimSize] -> + CSExp -> + CompilerM op s () + +data Operations op s = Operations { opsWriteScalar :: WriteScalar op s + , opsReadScalar :: ReadScalar op s + , opsAllocate :: Allocate op s + , opsCopy :: Copy op s + , opsStaticArray :: StaticArray op s + , opsCompiler :: OpCompiler op s + , opsEntryOutput :: EntryOutput op s + , opsEntryInput :: EntryInput op s + , opsSyncRun :: CSStmt + } + +-- | A set of operations that fail for every operation involving +-- non-default memory spaces. Uses plain pointers and @malloc@ for +-- memory management. +defaultOperations :: Operations op s +defaultOperations = Operations { opsWriteScalar = defWriteScalar + , opsReadScalar = defReadScalar + , opsAllocate = defAllocate + , opsCopy = defCopy + , opsStaticArray = defStaticArray + , opsCompiler = defCompiler + , opsEntryOutput = defEntryOutput + , opsEntryInput = defEntryInput + , opsSyncRun = defSyncRun + } + where defWriteScalar _ _ _ _ _ = + fail "Cannot write to non-default memory space because I am dumb" + defReadScalar _ _ _ _ = + fail "Cannot read from non-default memory space" + defAllocate _ _ _ = + fail "Cannot allocate in non-default memory space" + defCopy _ _ _ _ _ _ _ _ = + fail "Cannot copy to or from non-default memory space" + defStaticArray _ _ _ _ = + fail "Cannot create static array in non-default memory space" + defCompiler _ = + fail "The default compiler cannot compile extended operations" + defEntryOutput _ _ _ _ = + fail "Cannot return array not in default memory space" + defEntryInput _ _ _ _ = + fail "Cannot accept array not in default memory space" + defSyncRun = + Pass + +data CompilerEnv op s = CompilerEnv { + envOperations :: Operations op s + , envFtable :: M.Map Name [Imp.Type] +} + +data CompilerAcc op s = CompilerAcc { + accItems :: [CSStmt] + , accFreedMem :: [VName] + } + +instance Sem.Semigroup (CompilerAcc op s) where + CompilerAcc items1 freed1 <> CompilerAcc items2 freed2 = + CompilerAcc (items1<>items2) (freed1<>freed2) + +instance Monoid (CompilerAcc op s) where + mempty = CompilerAcc mempty mempty + mappend = (Sem.<>) + +envOpCompiler :: CompilerEnv op s -> OpCompiler op s +envOpCompiler = opsCompiler . envOperations + +envReadScalar :: CompilerEnv op s -> ReadScalar op s +envReadScalar = opsReadScalar . envOperations + +envWriteScalar :: CompilerEnv op s -> WriteScalar op s +envWriteScalar = opsWriteScalar . envOperations + +envAllocate :: CompilerEnv op s -> Allocate op s +envAllocate = opsAllocate . envOperations + +envCopy :: CompilerEnv op s -> Copy op s +envCopy = opsCopy . envOperations + +envStaticArray :: CompilerEnv op s -> StaticArray op s +envStaticArray = opsStaticArray . envOperations + +envEntryOutput :: CompilerEnv op s -> EntryOutput op s +envEntryOutput = opsEntryOutput . envOperations + +envEntryInput :: CompilerEnv op s -> EntryInput op s +envEntryInput = opsEntryInput . envOperations + +envSyncFun :: CompilerEnv op s -> CSStmt +envSyncFun = opsSyncRun . envOperations + +newCompilerEnv :: Imp.Functions op -> Operations op s -> CompilerEnv op s +newCompilerEnv (Imp.Functions funs) ops = + CompilerEnv { envOperations = ops + , envFtable = ftable <> builtinFtable + } + where ftable = M.fromList $ map funReturn funs + funReturn (name, Imp.Function _ outparams _ _ _ _) = (name, paramsTypes outparams) + builtinFtable = M.map (map Imp.Scalar . snd) builtInFunctions + +data CompilerState s = CompilerState { + compNameSrc :: VNameSource + , compBeforeParse :: [CSStmt] + , compInit :: [CSStmt] + , compStaticMemDecls :: [CSStmt] + , compStaticMemAllocs :: [CSStmt] + , compDebugItems :: [CSStmt] + , compUserState :: s + , compMemberDecls :: [CSStmt] + , compAssignedVars :: [VName] + , compDeclaredMem :: [(VName, Space)] +} + +newCompilerState :: VNameSource -> s -> CompilerState s +newCompilerState src s = CompilerState { compNameSrc = src + , compBeforeParse = [] + , compInit = [] + , compStaticMemDecls = [] + , compStaticMemAllocs = [] + , compDebugItems = [] + , compMemberDecls = [] + , compUserState = s + , compAssignedVars = [] + , compDeclaredMem = [] + } + +newtype CompilerM op s a = CompilerM (RWS (CompilerEnv op s) (CompilerAcc op s) (CompilerState s) a) + deriving (Functor, Applicative, Monad, + MonadState (CompilerState s), + MonadReader (CompilerEnv op s), + MonadWriter (CompilerAcc op s)) + +instance MonadFreshNames (CompilerM op s) where + getNameSource = gets compNameSrc + putNameSource src = modify $ \s -> s { compNameSrc = src } + +collect :: CompilerM op s () -> CompilerM op s [CSStmt] +collect m = pass $ do + ((), w) <- listen m + return (accItems w, + const w { accItems = mempty} ) + +collect' :: CompilerM op s a -> CompilerM op s (a, [CSStmt]) +collect' m = pass $ do + (x, w) <- listen m + return ((x, accItems w), + const w { accItems = mempty}) + +beforeParse :: CSStmt -> CompilerM op s () +beforeParse x = modify $ \s -> + s { compBeforeParse = compBeforeParse s ++ [x] } + +atInit :: CSStmt -> CompilerM op s () +atInit x = modify $ \s -> + s { compInit = compInit s ++ [x] } + +staticMemDecl :: CSStmt -> CompilerM op s () +staticMemDecl x = modify $ \s -> + s { compStaticMemDecls = compStaticMemDecls s ++ [x] } + +staticMemAlloc :: CSStmt -> CompilerM op s () +staticMemAlloc x = modify $ \s -> + s { compStaticMemAllocs = compStaticMemAllocs s ++ [x] } + +addMemberDecl :: CSStmt -> CompilerM op s () +addMemberDecl x = modify $ \s -> + s { compMemberDecls = compMemberDecls s ++ [x] } + +contextFinalInits :: CompilerM op s [CSStmt] +contextFinalInits = gets compInit + +item :: CSStmt -> CompilerM op s () +item x = tell $ mempty { accItems = [x] } + +stm :: CSStmt -> CompilerM op s () +stm = item + +stms :: [CSStmt] -> CompilerM op s () +stms = mapM_ stm + +debugReport :: CSStmt -> CompilerM op s () +debugReport x = modify $ \s -> + s { compDebugItems = compDebugItems s ++ [x] } + +getVarAssigned :: VName -> CompilerM op s Bool +getVarAssigned vname = + elem vname <$> gets compAssignedVars + +setVarAssigned :: VName -> CompilerM op s () +setVarAssigned vname = modify $ \s -> + s { compAssignedVars = vname : compAssignedVars s} + +futharkFun :: String -> String +futharkFun s = "futhark_" ++ zEncodeString s + +paramsTypes :: [Imp.Param] -> [Imp.Type] +paramsTypes = map paramType + +paramType :: Imp.Param -> Imp.Type +paramType (Imp.MemParam _ space) = Imp.Mem (Imp.ConstSize 0) space +paramType (Imp.ScalarParam _ t) = Imp.Scalar t + +compileOutput :: Imp.Param -> (CSExp, CSType) +compileOutput = nameFun &&& typeFun + where nameFun = Var . compileName . Imp.paramName + typeFun = compileType . paramType + +getDefaultDecl :: Imp.Param -> CSStmt +getDefaultDecl (Imp.MemParam v DefaultSpace) = + Assign (Var $ compileName v) $ simpleCall "allocateMem" [Integer 0] +getDefaultDecl (Imp.MemParam v _) = + AssignTyped (CustomT "OpenCLMemblock") (Var $ compileName v) (Just $ simpleCall "EmptyMemblock" [Var "Ctx.EMPTY_MEM_HANDLE"]) +getDefaultDecl (Imp.ScalarParam v Cert) = + Assign (Var $ compileName v) $ Bool True +getDefaultDecl (Imp.ScalarParam v t) = + Assign (Var $ compileName v) $ simpleInitClass (compilePrimType t) [] + + +runCompilerM :: Imp.Functions op -> Operations op s + -> VNameSource + -> s + -> CompilerM op s a + -> a +runCompilerM prog ops src userstate (CompilerM m) = + fst $ evalRWS m (newCompilerEnv prog ops) (newCompilerState src userstate) + +standardOptions :: [Option] +standardOptions = [ + Option { optionLongName = "write-runtime-to" + , optionShortName = Just 't' + , optionArgument = RequiredArgument + , optionAction = + [ + If (BinOp "!=" (Var "RuntimeFile") Null) + [Exp $ simpleCall "RuntimeFile.Close" []] [] + , Reassign (Var "RuntimeFile") $ + simpleInitClass "FileStream" [Var "optarg", Var "FileMode.Create"] + , Reassign (Var "RuntimeFileWriter") $ + simpleInitClass "StreamWriter" [Var "RuntimeFile"] + ] + }, + Option { optionLongName = "runs" + , optionShortName = Just 'r' + , optionArgument = RequiredArgument + , optionAction = + [ Reassign (Var "NumRuns") $ simpleCall "Convert.ToInt32" [Var "optarg"] + , Reassign (Var "DoWarmupRun") $ Bool True + ] + }, + Option { optionLongName = "entry-point" + , optionShortName = Just 'e' + , optionArgument = RequiredArgument + , optionAction = + [ Reassign (Var "EntryPoint") $ Var "optarg" ] + } + ] + +-- | The class generated by the code generator must have a +-- constructor, although it can be vacuous. +data Constructor = Constructor [CSFunDefArg] [CSStmt] + +-- | A constructor that takes no arguments and does nothing. +emptyConstructor :: Constructor +emptyConstructor = Constructor [(Composite $ ArrayT $ Primitive StringT, "args")] [] + +constructorToConstructorDef :: Constructor -> String -> [CSStmt] -> CSStmt +constructorToConstructorDef (Constructor params body) name at_init = + ConstructorDef $ ClassConstructor name params $ body <> at_init + + +compileProg :: MonadFreshNames m => + Maybe String + -> Constructor + -> [CSStmt] + -> [CSStmt] + -> Operations op s + -> s + -> CompilerM op s () + -> [CSStmt] + -> [Space] + -> [Option] + -> Imp.Functions op + -> m String +compileProg module_name constructor imports defines ops userstate boilerplate pre_timing _ options prog@(Imp.Functions funs) = do + src <- getNameSource + let prog' = runCompilerM prog ops src userstate compileProg' + let imports' = [ Using Nothing "System" + , Using Nothing "System.Diagnostics" + , Using Nothing "System.Collections" + , Using Nothing "System.Collections.Generic" + , Using Nothing "System.IO" + , Using Nothing "System.Linq" + , Using Nothing "System.Runtime.InteropServices" + , Using Nothing "static System.ValueTuple" + , Using Nothing "static System.Convert" + , Using Nothing "static System.Math" + , Using Nothing "System.Numerics" + , Using Nothing "Mono.Options" ] ++ imports + + return $ pretty (CSProg $ imports' ++ prog') + where compileProg' = do + definitions <- mapM compileFunc funs + opencl_boilerplate <- collect boilerplate + compBeforeParses <- gets compBeforeParse + compInits <- gets compInit + staticDecls <- gets compStaticMemDecls + staticAllocs <- gets compStaticMemAllocs + extraMemberDecls <- gets compMemberDecls + let member_decls' = member_decls ++ extraMemberDecls ++ staticDecls + let at_inits' = at_inits ++ compBeforeParses ++ parse_options ++ compInits ++ staticAllocs + + + case module_name of + Just name -> do + entry_points <- mapM (compileEntryFun pre_timing) $ filter (Imp.functionEntry . snd) funs + let constructor' = constructorToConstructorDef constructor name at_inits' + return [ Namespace name [ClassDef $ PublicClass name $ member_decls' ++ + constructor' : defines' ++ opencl_boilerplate ++ + map PrivateFunDef definitions ++ + map PublicFunDef entry_points ]] + + + Nothing -> do + let name = "FutharkInternal" + let constructor' = constructorToConstructorDef constructor name at_inits' + (entry_point_defs, entry_point_names, entry_points) <- + unzip3 <$> mapM (callEntryFun pre_timing) + (filter (Imp.functionEntry . snd) funs) + + debug_ending <- gets compDebugItems + return [Namespace name ((ClassDef $ + PublicClass name $ + member_decls' ++ + constructor' : defines' ++ + opencl_boilerplate ++ + map PrivateFunDef (definitions ++ entry_point_defs) ++ + [PublicFunDef $ Def "InternalEntry" VoidT [] $ selectEntryPoint entry_point_names entry_points ++ debug_ending + ] + ) : + [ClassDef $ PublicClass "Program" + [StaticFunDef $ Def "Main" VoidT [(string_arrayT,"args")] main_entry]]) + ] + + + + string_arrayT = Composite $ ArrayT $ Primitive StringT + main_entry :: [CSStmt] + main_entry = [ Assign (Var "internalInstance") (simpleInitClass "FutharkInternal" [Var "args"]) + , Exp $ simpleCall "internalInstance.InternalEntry" [] + ] + + member_decls = + [ AssignTyped (CustomT "FileStream") (Var "RuntimeFile") Nothing + , AssignTyped (CustomT "StreamWriter") (Var "RuntimeFileWriter") Nothing + , AssignTyped (Primitive BoolT) (Var "DoWarmupRun") Nothing + , AssignTyped (Primitive $ CSInt Int32T) (Var "NumRuns") Nothing + , AssignTyped (Primitive StringT) (Var "EntryPoint") Nothing + ] + + at_inits = [ Reassign (Var "DoWarmupRun") (Bool False) + , Reassign (Var "NumRuns") (Integer 1) + , Reassign (Var "EntryPoint") (String "main") + , Exp $ simpleCall "ValueReader" [] + ] + + defines' = [ Escape csScalar + , Escape csMemory + , Escape csPanic + , Escape csExceptions + , Escape csReader] ++ defines + + parse_options = + generateOptionParser (standardOptions ++ options) + + selectEntryPoint entry_point_names entry_points = + [ Assign (Var "EntryPoints") $ + Collection "Dictionary" $ zipWith Pair (map String entry_point_names) entry_points, + If (simpleCall "!EntryPoints.ContainsKey" [Var "EntryPoint"]) + [ Exp $ simpleCall "Console.Error.WriteLine" + [simpleCall "string.Format" + [ String "No entry point '{0}'. Select another with --entry point. Options are:\n{1}" + , Var "EntryPoint" + , simpleCall "string.Join" + [ String "\n" + , Field (Var "EntryPoints") "Keys" ]]] + , Exp $ simpleCall "Environment.Exit" [Integer 1]] + [ Assign (Var "entryPointFun") $ + Index (Var "EntryPoints") (IdxExp $ Var "EntryPoint") + , Exp $ simpleCall "entryPointFun.Invoke" []] + ] + + +compileFunc :: (Name, Imp.Function op) -> CompilerM op s CSFunDef +compileFunc (fname, Imp.Function _ outputs inputs body _ _) = do + body' <- blockScope $ compileCode body + let inputs' = map compileTypedInput inputs + let outputs' = map compileOutput outputs + let outputDecls = map getDefaultDecl outputs + let (ret, retType) = unzip outputs' + let retType' = tupleOrSingleT retType + let ret' = [Return $ tupleOrSingle ret] + + case outputs of + [] -> return $ Def (futharkFun . nameToString $ fname) VoidT inputs' (outputDecls++body') + _ -> return $ Def (futharkFun . nameToString $ fname) retType' inputs' (outputDecls++body'++ret') + + +compileTypedInput :: Imp.Param -> (CSType, String) +compileTypedInput input = (typeFun input, nameFun input) + where nameFun = compileName . Imp.paramName + typeFun = compileType . paramType + +tupleOrSingleEntryT :: [CSType] -> CSType +tupleOrSingleEntryT [e] = e +tupleOrSingleEntryT es = Composite $ SystemTupleT es + +tupleOrSingleEntry :: [CSExp] -> CSExp +tupleOrSingleEntry [e] = e +tupleOrSingleEntry es = CreateSystemTuple es + +tupleOrSingleT :: [CSType] -> CSType +tupleOrSingleT [e] = e +tupleOrSingleT es = Composite $ TupleT es + +tupleOrSingle :: [CSExp] -> CSExp +tupleOrSingle [e] = e +tupleOrSingle es = Tuple es + +assignScalarPointer :: CSExp -> CSExp -> CSStmt +assignScalarPointer e ptr = + AssignTyped (PointerT VoidT) ptr (Just $ Addr e) + +-- | A 'Call' where the function is a variable and every argument is a +-- simple 'Arg'. +simpleCall :: String -> [CSExp] -> CSExp +simpleCall fname = Call (Var fname) . map simpleArg + +-- | A 'Call' where the function is a variable and every argument is a +-- simple 'Arg'. +parametrizedCall :: String -> String -> [CSExp] -> CSExp +parametrizedCall fname primtype = Call (Var fname') . map simpleArg + where fname' = concat [fname, "<", primtype, ">"] + +simpleArg :: CSExp -> CSArg +simpleArg = Arg Nothing + +-- | A CallMethod +callMethod :: CSExp -> String -> [CSExp] -> CSExp +callMethod object method = CallMethod object (Var method) . map simpleArg + +simpleInitClass :: String -> [CSExp] -> CSExp +simpleInitClass fname =CreateObject (Var fname) . map simpleArg + +compileName :: VName -> String +compileName = zEncodeString . pretty + +compileType :: Imp.Type -> CSType +compileType (Imp.Scalar p) = compilePrimTypeToAST p +compileType (Imp.Mem _ space) = rawMemCSType space + +compilePrimTypeToAST :: PrimType -> CSType +compilePrimTypeToAST (IntType Int8) = Primitive $ CSInt Int8T +compilePrimTypeToAST (IntType Int16) = Primitive $ CSInt Int16T +compilePrimTypeToAST (IntType Int32) = Primitive $ CSInt Int32T +compilePrimTypeToAST (IntType Int64) = Primitive $ CSInt Int64T +compilePrimTypeToAST (FloatType Float32) = Primitive $ CSFloat FloatT +compilePrimTypeToAST (FloatType Float64) = Primitive $ CSFloat DoubleT +compilePrimTypeToAST Imp.Bool = Primitive BoolT +compilePrimTypeToAST Imp.Cert = Primitive BoolT + +compilePrimTypeToASText :: PrimType -> Imp.Signedness -> CSType +compilePrimTypeToASText (IntType Int8) Imp.TypeUnsigned = Primitive $ CSUInt UInt8T +compilePrimTypeToASText (IntType Int16) Imp.TypeUnsigned = Primitive $ CSUInt UInt16T +compilePrimTypeToASText (IntType Int32) Imp.TypeUnsigned = Primitive $ CSUInt UInt32T +compilePrimTypeToASText (IntType Int64) Imp.TypeUnsigned = Primitive $ CSUInt UInt64T +compilePrimTypeToASText (IntType Int8) _ = Primitive $ CSInt Int8T +compilePrimTypeToASText (IntType Int16) _ = Primitive $ CSInt Int16T +compilePrimTypeToASText (IntType Int32) _ = Primitive $ CSInt Int32T +compilePrimTypeToASText (IntType Int64) _ = Primitive $ CSInt Int64T +compilePrimTypeToASText (FloatType Float32) _ = Primitive $ CSFloat FloatT +compilePrimTypeToASText (FloatType Float64) _ = Primitive $ CSFloat DoubleT +compilePrimTypeToASText Imp.Bool _ = Primitive BoolT +compilePrimTypeToASText Imp.Cert _ = Primitive BoolT + +compileDim :: Imp.DimSize -> CSExp +compileDim (Imp.ConstSize i) = Integer $ toInteger i +compileDim (Imp.VarSize v) = Var $ compileName v + +unpackDim :: CSExp -> Imp.DimSize -> Int32 -> CompilerM op s () +unpackDim arr_name (Imp.ConstSize c) i = do + let shape_name = Field arr_name "Item2" -- array tuples are currently (data array * dimension array) currently + let constant_c = Integer $ toInteger c + let constant_i = Integer $ toInteger i + stm $ Assert (BinOp "==" constant_c (Index shape_name $ IdxExp constant_i)) [String "constant dimension wrong"] + +unpackDim arr_name (Imp.VarSize var) i = do + let shape_name = Field arr_name "Item2" + let src = Index shape_name $ IdxExp $ Integer $ toInteger i + let dest = Var $ compileName var + isAssigned <- getVarAssigned var + if isAssigned + then + stm $ Reassign dest $ Cast (Primitive $ CSInt Int32T) src + else do + stm $ Assign dest $ Cast (Primitive $ CSInt Int32T) src + setVarAssigned var + +entryPointOutput :: Imp.ExternalValue -> CompilerM op s CSExp +entryPointOutput (Imp.OpaqueValue _ vs) = + CreateSystemTuple <$> mapM (entryPointOutput . Imp.TransparentValue) vs + +entryPointOutput (Imp.TransparentValue (Imp.ScalarValue bt ept name)) = + return $ cast $ Var $ compileName name + where cast = compileTypecastExt bt ept + +entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ Imp.DefaultSpace bt ept dims)) = do + let src = Var $ compileName mem + let createTuple = "createTuple_" ++ compilePrimTypeExt bt ept + return $ simpleCall createTuple [src, CreateArray (Primitive $ CSInt Int64T) $ map compileDim dims] + +entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ (Imp.Space sid) bt ept dims)) = do + unRefMem mem (Imp.Space sid) + pack_output <- asks envEntryOutput + pack_output mem sid bt ept dims + +entryPointInput :: (Int, Imp.ExternalValue, CSExp) -> CompilerM op s () +entryPointInput (i, Imp.OpaqueValue _ vs, e) = + mapM_ entryPointInput $ zip3 (repeat i) (map Imp.TransparentValue vs) $ + map (\idx -> Field e $ "Item" ++ show (idx :: Int)) [1..] + +entryPointInput (_, Imp.TransparentValue (Imp.ScalarValue bt _ name), e) = do + let vname' = Var $ compileName name + cast = compileTypecast bt + stm $ Assign vname' (cast e) + +entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem memsize Imp.DefaultSpace bt _ dims), e) = do + zipWithM_ (unpackDim e) dims [0..] + let arrayData = Field e "Item1" + let dest = Var $ compileName mem + unwrap_call = simpleCall "unwrapArray" [arrayData, sizeOf $ compilePrimTypeToAST bt] + case memsize of + Imp.VarSize sizevar -> + stm $ Assign (Var $ compileName sizevar) $ Field e "Item2.Length" + Imp.ConstSize _ -> + return () + stm $ Assign dest unwrap_call + +entryPointInput (_, Imp.TransparentValue (Imp.ArrayValue mem memsize (Imp.Space sid) bt ept dims), e) = do + unpack_input <- asks envEntryInput + unpack <- collect $ unpack_input mem memsize sid bt ept dims e + stms unpack + +extValueDescName :: Imp.ExternalValue -> String +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 +extName = (++"_ext") + +sizeOf :: CSType -> CSExp +sizeOf t = simpleCall "sizeof" [(Var . pretty) t] + +publicFunDef :: String -> CSType -> [(CSType, String)] -> [CSStmt] -> CSStmt +publicFunDef s t args stmts = PublicFunDef $ Def s t args stmts + +privateFunDef :: String -> CSType -> [(CSType, String)] -> [CSStmt] -> CSStmt +privateFunDef s t args stmts = PrivateFunDef $ Def s t args stmts + +valueDescName :: Imp.ValueDesc -> String +valueDescName = compileName . valueDescVName + +valueDescVName :: Imp.ValueDesc -> VName +valueDescVName (Imp.ScalarValue _ _ vname) = vname +valueDescVName (Imp.ArrayValue vname _ _ _ _ _) = vname + +consoleWrite :: String -> [CSExp] -> CSExp +consoleWrite str exps = simpleCall "Console.Write" $ String str:exps + +consoleWriteLine :: String -> [CSExp] -> CSExp +consoleWriteLine str exps = simpleCall "Console.WriteLine" $ String str:exps + +consoleErrorWrite :: String -> [CSExp] -> CSExp +consoleErrorWrite str exps = simpleCall "Console.Error.Write" $ String str:exps + +consoleErrorWriteLine :: String -> [CSExp] -> CSExp +consoleErrorWriteLine str exps = simpleCall "Console.Error.WriteLine" $ String str:exps + +readFun :: PrimType -> Imp.Signedness -> String +readFun (FloatType Float32) _ = "ReadF32" +readFun (FloatType Float64) _ = "ReadF64" +readFun (IntType Int8) Imp.TypeUnsigned = "ReadU8" +readFun (IntType Int16) Imp.TypeUnsigned = "ReadU16" +readFun (IntType Int32) Imp.TypeUnsigned = "ReadU32" +readFun (IntType Int64) Imp.TypeUnsigned = "ReadU64" +readFun (IntType Int8) Imp.TypeDirect = "ReadI8" +readFun (IntType Int16) Imp.TypeDirect = "ReadI16" +readFun (IntType Int32) Imp.TypeDirect = "ReadI32" +readFun (IntType Int64) Imp.TypeDirect = "ReadI64" +readFun Imp.Bool _ = "ReadBool" +readFun Cert _ = error "readFun: cert" + +readBinFun :: PrimType -> Imp.Signedness -> String +readBinFun (FloatType Float32) _bin_ = "ReadBinF32" +readBinFun (FloatType Float64) _bin_ = "ReadBinF64" +readBinFun (IntType Int8) Imp.TypeUnsigned = "ReadBinU8" +readBinFun (IntType Int16) Imp.TypeUnsigned = "ReadBinU16" +readBinFun (IntType Int32) Imp.TypeUnsigned = "ReadBinU32" +readBinFun (IntType Int64) Imp.TypeUnsigned = "ReadBinU64" +readBinFun (IntType Int8) Imp.TypeDirect = "ReadBinI8" +readBinFun (IntType Int16) Imp.TypeDirect = "ReadBinI16" +readBinFun (IntType Int32) Imp.TypeDirect = "ReadBinI32" +readBinFun (IntType Int64) Imp.TypeDirect = "ReadBinI64" +readBinFun Imp.Bool _ = "ReadBinBool" +readBinFun Cert _ = error "readFun: cert" + +-- The value returned will be used when reading binary arrays, to indicate what +-- the expected type is +-- Key into the FUTHARK_PRIMTYPES dict. +readTypeEnum :: PrimType -> Imp.Signedness -> String +readTypeEnum (IntType Int8) Imp.TypeUnsigned = "u8" +readTypeEnum (IntType Int16) Imp.TypeUnsigned = "u16" +readTypeEnum (IntType Int32) Imp.TypeUnsigned = "u32" +readTypeEnum (IntType Int64) Imp.TypeUnsigned = "u64" +readTypeEnum (IntType Int8) Imp.TypeDirect = "i8" +readTypeEnum (IntType Int16) Imp.TypeDirect = "i16" +readTypeEnum (IntType Int32) Imp.TypeDirect = "i32" +readTypeEnum (IntType Int64) Imp.TypeDirect = "i64" +readTypeEnum (FloatType Float32) _ = "f32" +readTypeEnum (FloatType Float64) _ = "f64" +readTypeEnum Imp.Bool _ = "bool" +readTypeEnum Cert _ = error "readTypeEnum: cert" + +readInput :: Imp.ExternalValue -> CSStmt +readInput (Imp.OpaqueValue desc _) = + Throw $ simpleInitClass "Exception" [String $ "Cannot read argument of type " ++ desc ++ "."] + +readInput decl@(Imp.TransparentValue (Imp.ScalarValue bt ept _)) = + let read_func = Var $ readFun bt ept + read_bin_func = Var $ readBinFun bt ept + type_enum = String $ readTypeEnum bt ept + bt' = compilePrimTypeExt bt ept + readScalar = initializeGenericFunction "ReadScalar" bt' + in Assign (Var $ extValueDescName decl) $ simpleCall readScalar [type_enum, read_func, read_bin_func] + +-- TODO: If the type identifier of 'Float32' is changed, currently the error +-- messages for reading binary input will not use this new name. This is also a +-- problem for the C runtime system. +readInput decl@(Imp.TransparentValue (Imp.ArrayValue _ _ _ bt ept dims)) = + let rank' = Var $ show $ length dims + type_enum = String $ readTypeEnum bt ept + bt' = compilePrimTypeExt bt ept + read_func = Var $ readFun bt ept + readArray = initializeGenericFunction "ReadArray" bt' + in Assign (Var $ extValueDescName decl) $ simpleCall readArray [rank', type_enum, read_func] + +initializeGenericFunction :: String -> String -> String +initializeGenericFunction fun tp = fun ++ "<" ++ tp ++ ">" + + +printPrimStm :: CSExp -> CSStmt +printPrimStm val = Exp $ simpleCall "WriteValue" [val] + +formatString :: String -> [CSExp] -> CSExp +formatString fmt contents = + simpleCall "String.Format" $ String fmt : contents + +printStm :: Imp.ValueDesc -> CSExp -> CSExp -> CompilerM op s CSStmt +printStm Imp.ScalarValue{} _ e = + return $ printPrimStm e +printStm (Imp.ArrayValue _ _ _ _ _ []) ind e = do + let e' = Index e (IdxExp (PostUnOp "++" ind)) + return $ printPrimStm e' + +printStm (Imp.ArrayValue mem memsize space bt ept (outer:shape)) ind e = do + ptr <- newVName "shapePtr" + first <- newVName "printFirst" + let size = callMethod (CreateArray (Primitive $ CSInt Int32T) $ map compileDim $ outer:shape) + "Aggregate" [ Integer 1 + , Lambda (Tuple [Var "acc", Var "val"]) + [Exp $ BinOp "*" (Var "acc") (Var "val")] + ] + emptystr = "empty(" ++ ppArrayType bt (length shape) ++ ")" + + printelem <- printStm (Imp.ArrayValue mem memsize space bt ept shape) ind e + return $ + If (BinOp "==" size (Integer 0)) + [puts emptystr] + [ Assign (Var $ pretty first) $ Var "true" + , puts "[" + , For (pretty ptr) (compileDim outer) + [ If (simpleCall "!" [Var $ pretty first]) [puts ", "] [] + , printelem + , Reassign (Var $ pretty first) $ Var "false" + ] + , puts "]" + ] + + where ppArrayType :: PrimType -> Int -> String + ppArrayType t 0 = prettyPrimType ept t + ppArrayType t n = "[]" ++ ppArrayType t (n-1) + + prettyPrimType Imp.TypeUnsigned (IntType Int8) = "u8" + prettyPrimType Imp.TypeUnsigned (IntType Int16) = "u16" + prettyPrimType Imp.TypeUnsigned (IntType Int32) = "u32" + prettyPrimType Imp.TypeUnsigned (IntType Int64) = "u64" + prettyPrimType _ t = pretty t + + puts s = Exp $ simpleCall "Console.Write" [String s] + +printValue :: [(Imp.ExternalValue, CSExp)] -> CompilerM op s [CSStmt] +printValue = fmap concat . mapM (uncurry printValue') + -- We copy non-host arrays to the host before printing. This is + -- done in a hacky way - we assume the value has a .get()-method + -- that returns an equivalent Numpy array. This works for CSOpenCL, + -- but we will probably need yet another plugin mechanism here in + -- the future. + where printValue' (Imp.OpaqueValue desc _) _ = + return [Exp $ simpleCall "Console.Write" + [String $ "#"]] + printValue' (Imp.TransparentValue r@Imp.ScalarValue{}) e = do + p <- printStm r (Integer 0) e + return [p, Exp $ simpleCall "Console.Write" [String "\n"]] + printValue' (Imp.TransparentValue r@Imp.ArrayValue{}) e = do + tuple <- newVName "resultArr" + i <- newVName "arrInd" + let i' = Var $ compileName i + p <- printStm r i' (Var $ compileName tuple) + let e' = Var $ pretty e + return [ Assign (Var $ compileName tuple) (Field e' "Item1") + , Assign i' (Integer 0) + , p + , Exp $ simpleCall "Console.Write" [String "\n"]] + +prepareEntry :: (Name, Imp.Function op) -> CompilerM op s + (String, [(CSType, String)], CSType, [CSStmt], [CSStmt], [CSStmt], [CSStmt], + [(Imp.ExternalValue, CSExp)], [CSStmt]) +prepareEntry (fname, Imp.Function _ outputs inputs _ results args) = do + let (output_types, output_paramNames) = unzip $ map compileTypedInput outputs + funTuple = tupleOrSingle $ fmap Var output_paramNames + + + (_, sizeDecls) <- collect' $ forM args declsfunction + + (argexps_mem_copies, prepare_run) <- collect' $ forM inputs $ \case + Imp.MemParam name space -> do + -- A program might write to its input parameters, so create a new memory + -- block and copy the source there. This way the program can be run more + -- than once. + name' <- newVName $ baseString name <> "_copy" + copy <- asks envCopy + allocate <- asks envAllocate + + let size = Var (compileName name ++ "_nbytes") + dest = name' + src = name + offset = Integer 0 + case space of + DefaultSpace -> + stm $ Reassign (Var (compileName name')) + (simpleCall "allocateMem" [size]) -- FIXME + Space sid -> + allocate name' size sid + copy dest offset space src offset space size (IntType Int64) -- FIXME + return $ Just (compileName name') + _ -> return Nothing + + prepareIn <- collect $ mapM_ entryPointInput $ zip3 [0..] args $ + map (Var . extValueDescName) args + (res, prepareOut) <- collect' $ mapM entryPointOutput results + + let mem_copies = mapMaybe liftMaybe $ zip argexps_mem_copies inputs + mem_copy_inits = map initCopy mem_copies + + argexps_lib = map (compileName . Imp.paramName) inputs + argexps_bin = zipWith fromMaybe argexps_lib argexps_mem_copies + fname' = futharkFun (nameToString fname) + arg_types = map (fst . compileTypedInput) inputs + inputs' = zip arg_types (map extValueDescName args) + output_type = tupleOrSingleEntryT output_types + call_lib = [Reassign funTuple $ simpleCall fname' (fmap Var argexps_lib)] + call_bin = [Reassign funTuple $ simpleCall fname' (fmap Var argexps_bin)] + prepareIn' = prepareIn ++ mem_copy_inits ++ sizeDecls + + return (nameToString fname, inputs', output_type, + prepareIn', call_lib, call_bin, prepareOut, + zip results res, prepare_run) + + where liftMaybe (Just a, b) = Just (a,b) + liftMaybe _ = Nothing + + initCopy (varName, Imp.MemParam _ space) = declMem' varName space + initCopy _ = Pass + + valueDescFun (Imp.ArrayValue mem _ Imp.DefaultSpace _ _ _) = + stm $ Assign (Var $ compileName mem ++ "_nbytes") (Var $ compileName mem ++ ".Length") + valueDescFun (Imp.ArrayValue mem _ (Imp.Space _) bt _ dims) = + stm $ Assign (Var $ compileName mem ++ "_nbytes") $ foldr (BinOp "*" . compileDim) (sizeOf $ compilePrimTypeToAST bt) dims + valueDescFun _ = stm Pass + + declsfunction (Imp.TransparentValue v) = valueDescFun v + declsfunction (Imp.OpaqueValue _ vs) = mapM_ valueDescFun vs + +copyMemoryDefaultSpace :: VName -> CSExp -> VName -> CSExp -> CSExp -> + CompilerM op s () +copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = + stm $ Exp $ simpleCall "Buffer.BlockCopy" [ Var (compileName srcmem), srcidx + , Var (compileName destmem), destidx, + nbytes] + +compileEntryFun :: [CSStmt] -> (Name, Imp.Function op) + -> CompilerM op s CSFunDef +compileEntryFun pre_timing entry@(_,Imp.Function _ outputs _ _ results args) = do + let params = map (getType &&& extValueDescName) args + let outputType = tupleOrSingleEntryT $ map getType results + + (fname', _, _, prepareIn, body_lib, _, prepareOut, res, _) <- prepareEntry entry + let ret = Return $ tupleOrSingleEntry $ map snd res + let outputDecls = map getDefaultDecl outputs + do_run = body_lib ++ pre_timing + (do_run_with_timing, close_runtime_file) <- addTiming do_run + + let do_warmup_run = If (Var "DoWarmupRun") do_run [] + do_num_runs = For "i" (Var "NumRuns") do_run_with_timing + + return $ Def fname' outputType params $ + prepareIn ++ outputDecls ++ [do_warmup_run, do_num_runs, close_runtime_file] ++ prepareOut ++ [ret] + + where getType :: Imp.ExternalValue -> CSType + getType (Imp.OpaqueValue _ valueDescs) = + let valueDescs' = map getType' valueDescs + in Composite $ SystemTupleT valueDescs' + getType (Imp.TransparentValue valueDesc) = + getType' valueDesc + + getType' :: Imp.ValueDesc -> CSType + getType' (Imp.ScalarValue primtype signedness _) = + compilePrimTypeToASText primtype signedness + getType' (Imp.ArrayValue _ _ _ primtype signedness _) = + let t = compilePrimTypeToASText primtype signedness + in Composite $ SystemTupleT [Composite $ ArrayT t, Composite $ ArrayT $ Primitive $ CSInt Int64T] + + +callEntryFun :: [CSStmt] -> (Name, Imp.Function op) + -> CompilerM op s (CSFunDef, String, CSExp) +callEntryFun pre_timing entry@(fname, Imp.Function _ outputs _ _ _ decl_args) = + if any isOpaque decl_args then + return (Def fname' VoidT [] [exitException], nameToString fname, Var fname') + else do + (_, _, _, prepareIn, _, body_bin, prepare_out, res, prepare_run) <- prepareEntry entry + let str_input = map readInput decl_args + + let outputDecls = map getDefaultDecl outputs + exitcall = [ + Exp $ simpleCall "Console.Error.WriteLine" [formatString "Assertion.{0} failed" [Var "e"]] + , Exp $ simpleCall "Environment.Exit" [Integer 1] + ] + except' = Catch (Var "Exception") exitcall + do_run = body_bin ++ pre_timing + (do_run_with_timing, close_runtime_file) <- addTiming do_run + + -- We ignore overflow errors and the like for executable entry + -- points. These are (somewhat) well-defined in Futhark. + + let maybe_free = + [If (BinOp "<" (Var "i") (BinOp "-" (Var "NumRuns") (Integer 1))) + prepare_out []] + + do_warmup_run = + If (Var "DoWarmupRun") (prepare_run ++ do_run ++ prepare_out) [] + + do_num_runs = + For "i" (Var "NumRuns") (prepare_run ++ do_run_with_timing ++ maybe_free) + + str_output <- printValue res + + return (Def fname' VoidT [] $ + str_input ++ prepareIn ++ outputDecls ++ + [Try [do_warmup_run, do_num_runs] [except']] ++ + [close_runtime_file] ++ + str_output, + + nameToString fname, + + Var fname') + + where fname' = "entry_" ++ nameToString fname + isOpaque Imp.TransparentValue{} = False + isOpaque _ = True + + exitException = Throw $ simpleInitClass "Exception" [String $ "The function " ++ nameToString fname ++ " is not available as an entry function."] + +addTiming :: [CSStmt] -> CompilerM s op ([CSStmt], CSStmt) +addTiming statements = do + syncFun <- asks envSyncFun + + return ([ Assign (Var "StopWatch") $ simpleInitClass "Stopwatch" [] + , syncFun + , Exp $ simpleCall "StopWatch.Start" [] ] ++ + statements ++ + [ syncFun + , Exp $ simpleCall "StopWatch.Stop" [] + , Assign (Var "timeElapsed") $ asMicroseconds (Var "StopWatch") + , If (not_null (Var "RuntimeFile")) [print_runtime] [] + ] + , If (not_null (Var "RuntimeFile")) [ + Exp $ simpleCall "RuntimeFileWriter.Close" [] , + Exp $ simpleCall "RuntimeFile.Close" [] + ] [] + ) + + where print_runtime = Exp $ simpleCall "RuntimeFileWriter.WriteLine" [ callMethod (Var "timeElapsed") "ToString" [] ] + not_null var = BinOp "!=" var Null + asMicroseconds watch = + BinOp "/" (Field watch "ElapsedTicks") + (BinOp "/" (Field (Var "TimeSpan") "TicksPerMillisecond") (Integer 1000)) + +compileUnOp :: Imp.UnOp -> String +compileUnOp op = + case op of + Not -> "!" + Complement{} -> "~" + Abs{} -> "Math.Abs" -- actually write these helpers + FAbs{} -> "Math.Abs" + SSignum{} -> "ssignum" + USignum{} -> "usignum" + +compileBinOpLike :: Monad m => + Imp.Exp -> Imp.Exp + -> CompilerM op s (CSExp, CSExp, String -> m CSExp) +compileBinOpLike x y = do + x' <- compileExp x + y' <- compileExp y + let simple s = return $ BinOp s x' y' + return (x', y', simple) + +-- | The ctypes type corresponding to a 'PrimType'. +compilePrimType :: PrimType -> String +compilePrimType t = + case t of + IntType Int8 -> "sbyte" + IntType Int16 -> "short" + IntType Int32 -> "int" + IntType Int64 -> "long" + FloatType Float32 -> "float" + FloatType Float64 -> "double" + Imp.Bool -> "bool" + Cert -> "bool" + +-- | The ctypes type corresponding to a 'PrimType', taking sign into account. +compilePrimTypeExt :: PrimType -> Imp.Signedness -> String +compilePrimTypeExt t ept = + case (t, ept) of + (IntType Int8, Imp.TypeUnsigned) -> "byte" + (IntType Int16, Imp.TypeUnsigned) -> "ushort" + (IntType Int32, Imp.TypeUnsigned) -> "uint" + (IntType Int64, Imp.TypeUnsigned) -> "ulong" + (IntType Int8, _) -> "sbyte" + (IntType Int16, _) -> "short" + (IntType Int32, _) -> "int" + (IntType Int64, _) -> "long" + (FloatType Float32, _) -> "float" + (FloatType Float64, _) -> "double" + (Imp.Bool, _) -> "bool" + (Cert, _) -> "byte" + +-- | Select function to retrieve bytes from byte array as specific data type +-- | The ctypes type corresponding to a 'PrimType'. +compileTypecastExt :: PrimType -> Imp.Signedness -> (CSExp -> CSExp) +compileTypecastExt t ept = + let t' = case (t, ept) of + (IntType Int8 , Imp.TypeUnsigned)-> Primitive $ CSUInt UInt8T + (IntType Int16 , Imp.TypeUnsigned)-> Primitive $ CSUInt UInt16T + (IntType Int32 , Imp.TypeUnsigned)-> Primitive $ CSUInt UInt32T + (IntType Int64 , Imp.TypeUnsigned)-> Primitive $ CSUInt UInt64T + (IntType Int8 , _)-> Primitive $ CSInt Int8T + (IntType Int16 , _)-> Primitive $ CSInt Int16T + (IntType Int32 , _)-> Primitive $ CSInt Int32T + (IntType Int64 , _)-> Primitive $ CSInt Int64T + (FloatType Float32, _)-> Primitive $ CSFloat FloatT + (FloatType Float64, _)-> Primitive $ CSFloat DoubleT + (Imp.Bool , _)-> Primitive BoolT + (Cert, _)-> Primitive $ CSInt Int8T + in Cast t' + +-- | The ctypes type corresponding to a 'PrimType'. +compileTypecast :: PrimType -> (CSExp -> CSExp) +compileTypecast t = + let t' = case t of + IntType Int8 -> Primitive $ CSInt Int8T + IntType Int16 -> Primitive $ CSInt Int16T + IntType Int32 -> Primitive $ CSInt Int32T + IntType Int64 -> Primitive $ CSInt Int64T + FloatType Float32 -> Primitive $ CSFloat FloatT + FloatType Float64 -> Primitive $ CSFloat DoubleT + Imp.Bool -> Primitive BoolT + Cert -> Primitive $ CSInt Int8T + in Cast t' + +-- | The ctypes type corresponding to a 'PrimType'. +compilePrimValue :: Imp.PrimValue -> CSExp +compilePrimValue (IntValue (Int8Value v)) = + Cast (Primitive $ CSInt Int8T) $ Integer $ toInteger v +compilePrimValue (IntValue (Int16Value v)) = + Cast (Primitive $ CSInt Int16T) $ Integer $ toInteger v +compilePrimValue (IntValue (Int32Value v)) = + Cast (Primitive $ CSInt Int32T) $ Integer $ toInteger v +compilePrimValue (IntValue (Int64Value v)) = + Cast (Primitive $ CSInt Int64T) $ Integer $ toInteger v +compilePrimValue (FloatValue (Float32Value v)) + | isInfinite v = + if v > 0 then Var "Single.PositiveInfinity" else Var "Single.NegativeInfinity" + | isNaN v = + Var "Single.NaN" + | otherwise = Cast (Primitive $ CSFloat FloatT) (Float $ fromRational $ toRational v) +compilePrimValue (FloatValue (Float64Value v)) + | isInfinite v = + if v > 0 then Var "Double.PositiveInfinity" else Var "Double.NegativeInfinity" + | isNaN v = + Var "Double.NaN" + | otherwise = Cast (Primitive $ CSFloat DoubleT) (Float $ fromRational $ toRational v) +compilePrimValue (BoolValue v) = Bool v +compilePrimValue Checked = Bool True + +compileExp :: Imp.Exp -> CompilerM op s CSExp + +compileExp (Imp.ValueExp v) = return $ compilePrimValue v + +compileExp (Imp.LeafExp (Imp.ScalarVar vname) _) = + return $ Var $ compileName vname + +compileExp (Imp.LeafExp (Imp.SizeOf t) _) = + return $ (compileTypecast $ IntType Int32) (Integer $ primByteSize t) + +compileExp (Imp.LeafExp (Imp.Index src (Imp.Count iexp) (IntType Int8) DefaultSpace _) _) = do + let src' = compileName src + iexp' <- compileExp iexp + return $ Cast (Primitive $ CSInt Int8T) (Index (Var src') (IdxExp iexp')) + +compileExp (Imp.LeafExp (Imp.Index src (Imp.Count iexp) bt DefaultSpace _) _) = do + iexp' <- compileExp iexp + let bt' = compilePrimType bt + return $ simpleCall ("indexArray_" ++ bt') [Var $ compileName src, iexp'] + +compileExp (Imp.LeafExp (Imp.Index src (Imp.Count iexp) restype (Imp.Space space) _) _) = + join $ asks envReadScalar + <*> pure src <*> compileExp iexp + <*> pure restype <*> pure space + +compileExp (Imp.BinOpExp op x y) = do + (x', y', simple) <- compileBinOpLike x y + case op of + FAdd{} -> simple "+" + FSub{} -> simple "-" + FMul{} -> simple "*" + FDiv{} -> simple "/" + LogAnd{} -> simple "&&" + LogOr{} -> simple "||" + _ -> return $ simpleCall (pretty op) [x', y'] + +compileExp (Imp.ConvOpExp conv x) = do + x' <- compileExp x + return $ simpleCall (pretty conv) [x'] + +compileExp (Imp.CmpOpExp cmp x y) = do + (x', y', simple) <- compileBinOpLike x y + case cmp of + CmpEq{} -> simple "==" + FCmpLt{} -> simple "<" + FCmpLe{} -> simple "<=" + _ -> return $ simpleCall (pretty cmp) [x', y'] + +compileExp (Imp.UnOpExp op exp1) = + PreUnOp (compileUnOp op) <$> compileExp exp1 + +compileExp (Imp.FunExp h args _) = + simpleCall (futharkFun (pretty h)) <$> mapM compileExp args + +compileCode :: Imp.Code op -> CompilerM op s () + +compileCode Imp.DebugPrint{} = + return () + +compileCode (Imp.Op op) = + join $ asks envOpCompiler <*> pure op + +compileCode (Imp.If cond tb fb) = do + cond' <- compileExp cond + tb' <- blockScope $ compileCode tb + fb' <- blockScope $ compileCode fb + stm $ If cond' tb' fb' + +compileCode (c1 Imp.:>>: c2) = do + compileCode c1 + compileCode c2 + +compileCode (Imp.While cond body) = do + cond' <- compileExp cond + body' <- blockScope $ compileCode body + stm $ While cond' body' + +compileCode (Imp.For i it bound body) = do + bound' <- compileExp bound + let i' = compileName i + body' <- blockScope $ compileCode body + counter <- pretty <$> newVName "counter" + one <- pretty <$> newVName "one" + stm $ Assign (Var i') $ compileTypecast (IntType it) (Integer 0) + stm $ Assign (Var one) $ compileTypecast (IntType it) (Integer 1) + stm $ For counter bound' $ body' ++ + [AssignOp "+" (Var i') (Var one)] + + +compileCode (Imp.SetScalar vname exp1) = do + let name' = Var $ compileName vname + exp1' <- compileExp exp1 + stm $ Reassign name' exp1' + +compileCode (Imp.DeclareMem v space) = declMem v space + +compileCode (Imp.DeclareScalar v Cert) = + stm $ Assign (Var $ compileName v) $ Bool True +compileCode (Imp.DeclareScalar v t) = + stm $ AssignTyped t' (Var $ compileName v) Nothing + where t' = compilePrimTypeToAST t + +compileCode (Imp.DeclareArray name DefaultSpace t vs) = + stms [Assign (Var $ "init_"++name') $ + simpleCall "unwrapArray" + [ + CreateArray (compilePrimTypeToAST t) (map compilePrimValue vs) + , simpleCall "sizeof" [Var $ compilePrimType t] + ] + , Assign (Var name') $ Var ("init_"++name') + ] + where name' = compileName name + + +compileCode (Imp.DeclareArray name (Space space) t vs) = + join $ asks envStaticArray <*> + pure name <*> pure space <*> pure t <*> pure vs + +compileCode (Imp.Comment s code) = do + code' <- blockScope $ compileCode code + stm $ Comment s code' + +compileCode (Imp.Assert e (Imp.ErrorMsg parts) (loc,locs)) = do + e' <- compileExp e + let onPart (i, Imp.ErrorString s) = return (printFormatArg i, String s) + onPart (i, Imp.ErrorInt32 x) = (printFormatArg i,) <$> compileExp x + (formatstrs, formatargs) <- unzip <$> mapM onPart (zip ([1..] :: [Integer]) parts) + stm $ Assert e' $ (String $ "Error at {0}:\n" <> concat formatstrs) : (String stacktrace : formatargs) + where stacktrace = intercalate " -> " (reverse $ map locStr $ loc:locs) + printFormatArg = printf "{%d}" + +compileCode (Imp.Call dests fname args) = do + args' <- mapM compileArg args + let dests' = tupleOrSingle $ fmap Var (map compileName dests) + fname' = futharkFun (pretty fname) + call' = simpleCall fname' args' + -- If the function returns nothing (is called only for side + -- effects), take care not to assign to an empty tuple. + stm $ if null dests + then Exp call' + else Reassign dests' call' + where compileArg (Imp.MemArg m) = return $ Var $ compileName m + compileArg (Imp.ExpArg e) = compileExp e + +compileCode (Imp.SetMem dest src DefaultSpace) = do + let src' = Var (compileName src) + let dest' = Var (compileName dest) + stm $ Reassign dest' src' + +compileCode (Imp.SetMem dest src _) = do + let src' = Var (compileName src) + let dest' = Var (compileName dest) + stm $ Exp $ simpleCall "MemblockSetDevice" [Ref $ Var "Ctx", Ref dest', Ref src', String (compileName src)] + +compileCode (Imp.Allocate name (Imp.Count e) DefaultSpace) = do + e' <- compileExp e + let allocate' = simpleCall "allocateMem" [e'] + let name' = Var (compileName name) + stm $ Reassign name' allocate' + +compileCode (Imp.Allocate name (Imp.Count e) (Imp.Space space)) = + join $ asks envAllocate + <*> pure name + <*> compileExp e + <*> pure space + +compileCode (Imp.Free name space) = do + unRefMem name space + tell $ mempty { accFreedMem = [name] } + +compileCode (Imp.Copy dest (Imp.Count destoffset) DefaultSpace src (Imp.Count srcoffset) DefaultSpace (Imp.Count size)) = do + destoffset' <- compileExp destoffset + srcoffset' <- compileExp srcoffset + let dest' = Var (compileName dest) + let src' = Var (compileName src) + size' <- compileExp size + stm $ Exp $ simpleCall "Buffer.BlockCopy" [src', srcoffset', dest', destoffset', size'] + +compileCode (Imp.Copy dest (Imp.Count destoffset) destspace src (Imp.Count srcoffset) srcspace (Imp.Count size)) = do + copy <- asks envCopy + join $ copy + <$> pure dest <*> compileExp destoffset <*> pure destspace + <*> pure src <*> compileExp srcoffset <*> pure srcspace + <*> compileExp size <*> pure (IntType Int64) -- FIXME + +compileCode (Imp.Write dest (Imp.Count idx) elemtype DefaultSpace _ elemexp) = do + idx' <- compileExp idx + elemexp' <- compileExp elemexp + let dest' = Var $ compileName dest + let elemtype' = compileTypecast elemtype + let ctype = elemtype' elemexp' + stm $ Exp $ simpleCall "writeScalarArray" [dest', idx', ctype] + +compileCode (Imp.Write dest (Imp.Count idx) elemtype (Imp.Space space) _ elemexp) = + join $ asks envWriteScalar + <*> pure dest + <*> compileExp idx + <*> pure elemtype + <*> pure space + <*> compileExp elemexp + +compileCode Imp.Skip = return () + +blockScope :: CompilerM op s () -> CompilerM op s [CSStmt] +blockScope = fmap snd . blockScope' + +blockScope' :: CompilerM op s a -> CompilerM op s (a, [CSStmt]) +blockScope' m = do + old_allocs <- gets compDeclaredMem + (x, items) <- pass $ do + (x, w) <- listen m + let items = accItems w + return ((x, items), const mempty) + new_allocs <- gets $ filter (`notElem` old_allocs) . compDeclaredMem + modify $ \s -> s { compDeclaredMem = old_allocs } + releases <- collect $ mapM_ (uncurry unRefMem) new_allocs + return (x, items <> releases) + +unRefMem :: VName -> Space -> CompilerM op s () +unRefMem mem (Space "device") = + (stm . Exp) $ simpleCall "MemblockUnrefDevice" [ Ref $ Var "Ctx" + , (Ref . Var . compileName) mem + , (String . compileName) mem] +unRefMem _ DefaultSpace = stm Pass +unRefMem _ (Space "local") = stm Pass +unRefMem _ (Space _) = fail "The default compiler cannot compile unRefMem for other spaces" + + +-- | Public names must have a consistent prefix. +publicName :: String -> String +publicName s = "Futhark" ++ s + +declMem :: VName -> Space -> CompilerM op s () +declMem name space = do + modify $ \s -> s { compDeclaredMem = (name, space) : compDeclaredMem s} + stm $ declMem' (compileName name) space + +declMem' :: String -> Space -> CSStmt +declMem' name DefaultSpace = + AssignTyped (Composite $ ArrayT $ Primitive ByteT) (Var name) Nothing +declMem' name (Space _) = + AssignTyped (CustomT "OpenCLMemblock") (Var name) (Just $ simpleCall "EmptyMemblock" [Var "Ctx.EMPTY_MEM_HANDLE"]) + +rawMemCSType :: Space -> CSType +rawMemCSType DefaultSpace = Composite $ ArrayT $ Primitive ByteT +rawMemCSType (Space _) = CustomT "OpenCLMemblock" + +toIntPtr :: CSExp -> CSExp +toIntPtr e = simpleInitClass "IntPtr" [e] diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs b/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs @@ -0,0 +1,411 @@ +{-# LANGUAGE PostfixOperators #-} + + +module Futhark.CodeGen.Backends.GenericCSharp.AST + ( CSExp(..) + , CSType(..) + , CSComp(..) + , CSPrim(..) + , CSInt(..) + , CSUInt(..) + , CSFloat(..) + , CSIdx (..) + , CSArg (..) + , CSStmt(..) + , module Language.Futhark.Core + , CSProg(..) + , CSExcept(..) + , CSFunDef(..) + , CSFunDefArg + , CSClassDef(..) + , CSConstructorDef(..) + ) + where + +import Language.Futhark.Core +import Data.List(intersperse) +import Futhark.Util.Pretty + +data MemT = Pointer + deriving (Eq, Show) + +data ArgMemType = ArgOut + | ArgRef + deriving (Eq, Show) + +instance Pretty ArgMemType where + ppr ArgOut = text "out" + ppr ArgRef = text "ref" + +instance Pretty CSComp where + ppr (ArrayT t) = ppr t <> text "[]" + ppr (TupleT ts) = parens(commasep $ map ppr ts) + ppr (SystemTupleT ts) = text "Tuple" <> angles(commasep $ map ppr ts) + +data CSInt = Int8T + | Int16T + | Int32T + | Int64T + deriving (Eq, Show) + +data CSUInt = UInt8T + | UInt16T + | UInt32T + | UInt64T + deriving (Eq, Show) + +data CSFloat = FloatT + | DoubleT + deriving (Eq, Show) + +data CSType = Composite CSComp + | PointerT CSType + | Primitive CSPrim + | CustomT String + | StaticT CSType + | OutT CSType + | RefT CSType + | VoidT + deriving (Eq, Show) + +data CSComp = ArrayT CSType + | TupleT [CSType] + | SystemTupleT [CSType] + deriving (Eq, Show) + +data CSPrim = CSInt CSInt + | CSUInt CSUInt + | CSFloat CSFloat + | BoolT + | ByteT + | StringT + | IntPtrT + deriving (Eq, Show) + +instance Pretty CSType where + ppr (Composite t) = ppr t + ppr (PointerT t) = ppr t <> text "*" + ppr (Primitive t) = ppr t + ppr (CustomT t) = text t + ppr (StaticT t) = text "static" <+> ppr t + ppr (OutT t) = text "out" <+> ppr t + ppr (RefT t) = text "ref" <+> ppr t + ppr VoidT = text "void" + +instance Pretty CSPrim where + ppr BoolT = text "bool" + ppr ByteT = text "byte" + ppr (CSInt t) = ppr t + ppr (CSUInt t) = ppr t + ppr (CSFloat t) = ppr t + ppr StringT = text "string" + ppr IntPtrT = text "IntPtr" + +instance Pretty CSInt where + ppr Int8T = text "sbyte" + ppr Int16T = text "short" + ppr Int32T = text "int" + ppr Int64T = text "long" + +instance Pretty CSUInt where + ppr UInt8T = text "byte" + ppr UInt16T = text "ushort" + ppr UInt32T = text "uint" + ppr UInt64T = text "ulong" + +instance Pretty CSFloat where + ppr FloatT = text "float" + ppr DoubleT = text "double" + +data UnOp = Not -- ^ Boolean negation. + | Complement -- ^ Bitwise complement. + | Negate -- ^ Numerical negation. + | Abs -- ^ Absolute/numerical value. + deriving (Eq, Show) + +data CSExp = Integer Integer + | Bool Bool + | Float Double + | String String + | RawStringLiteral String + | Var String + | Addr CSExp + | Ref CSExp + | Out CSExp + | Deref String + | BinOp String CSExp CSExp + | PreUnOp String CSExp + | PostUnOp String CSExp + | Ternary CSExp CSExp CSExp + | Cond CSExp CSExp CSExp + | Index CSExp CSIdx + | Pair CSExp CSExp + | Call CSExp [CSArg] + | CallMethod CSExp CSExp [CSArg] + | CreateObject CSExp [CSArg] + | CreateArray CSType [CSExp] + | CreateSystemTuple [CSExp] + | AllocArray CSType CSExp + | Cast CSType CSExp + | Tuple [CSExp] + | Array [CSExp] + | Field CSExp String + | Lambda CSExp [CSStmt] + | Collection String [CSExp] + | This CSExp + | Null + deriving (Eq, Show) + +instance Pretty CSExp where + ppr (Integer x) = ppr x + ppr (Float x) + | isInfinite x = text $ if x > 0 then "Double.PositiveInfinity" else "Double.NegativeInfinity" + | otherwise = ppr x + ppr (Bool True) = text "true" + ppr (Bool False) = text "false" + ppr (String x) = text $ show x + ppr (RawStringLiteral s) = text "@\"" <> text s <> text "\"" + ppr (Var n) = text $ map (\x -> if x == '\'' then 'm' else x) n + ppr (Addr e) = text "&" <> ppr e + ppr (Ref e) = text "ref" <+> ppr e + ppr (Out e) = text "out" <+> ppr e + ppr (Deref n) = text "*" <> text (map (\x -> if x == '\'' then 'm' else x) n) + ppr (BinOp s e1 e2) = parens(ppr e1 <+> text s <+> ppr e2) + ppr (PreUnOp s e) = text s <> parens (ppr e) + ppr (PostUnOp s e) = parens (ppr e) <> text s + ppr (Ternary b e1 e2) = ppr b <+> text "?" <+> ppr e1 <+> colon <+> ppr e2 + ppr (Cond e1 e2 e3) = text "if" <+> parens(ppr e1) <> braces(ppr e2) <+> text "else" <> braces(ppr e3) + ppr (Cast bt src) = parens(ppr bt) <+> ppr src + ppr (Index src (IdxExp idx)) = ppr src <> brackets(ppr idx) + ppr (Index src (IdxRange from to)) = text "MySlice" <> parens(commasep $ map ppr [src, from, to]) + ppr (Pair e1 e2) = braces(ppr e1 <> comma <> ppr e2) + ppr (Call fun args) = ppr fun <> parens(commasep $ map ppr args) + ppr (CallMethod obj method args) = ppr obj <> dot <> ppr method <> parens(commasep $ map ppr args) + ppr (CreateObject className args) = text "new" <+> ppr className <> parens(commasep $ map ppr args) + ppr (CreateArray t vs) = text "new" <+> ppr t <> text "[]" <+> braces(commasep $ map ppr vs) + ppr (CreateSystemTuple exps) = text "Tuple.Create" <> parens(commasep $ map ppr exps) + ppr (Tuple exps) = parens(commasep $ map ppr exps) + ppr (Array exps) = braces(commasep $ map ppr exps) -- uhoh is this right? + ppr (Field obj field) = ppr obj <> dot <> text field + ppr (Lambda expr [Exp e]) = ppr expr <+> text "=>" <+> ppr e + ppr (Lambda expr stmts) = ppr expr <+> text "=>" <+> braces(stack $ map ppr stmts) + ppr (Collection collection exps) = text "new" <+> text collection <> braces(commasep $ map ppr exps) + ppr (This e) = text "this" <> dot <> ppr e + ppr Null = text "null" + ppr (AllocArray t len) = text "new" <+> ppr t <> lbracket <> ppr len <> rbracket + +data CSIdx = IdxRange CSExp CSExp + | IdxExp CSExp + deriving (Eq, Show) + +data CSArg = ArgKeyword String CSArg -- please don't assign multiple keywords with the same argument + | Arg (Maybe ArgMemType) CSExp + deriving (Eq, Show) + +instance Pretty CSArg where + ppr (ArgKeyword kw arg) = text kw <> colon <+> ppr arg + ppr (Arg (Just mt) arg) = ppr mt <+> ppr arg + ppr (Arg Nothing arg) = ppr arg + +data CSStmt = If CSExp [CSStmt] [CSStmt] + | Try [CSStmt] [CSExcept] + | While CSExp [CSStmt] + | For String CSExp [CSStmt] + | ForEach String CSExp [CSStmt] + | UsingWith CSStmt [CSStmt] + | Unsafe [CSStmt] + | Fixed CSExp CSExp [CSStmt] + | Assign CSExp CSExp + | Reassign CSExp CSExp + | AssignOp String CSExp CSExp + | AssignTyped CSType CSExp (Maybe CSExp) + + | Comment String [CSStmt] + | Assert CSExp [CSExp] + | Throw CSExp + | Exp CSExp + | Return CSExp + | Pass + -- Definition-like statements. + | Using (Maybe String) String + | StaticFunDef CSFunDef + | PublicFunDef CSFunDef + | PrivateFunDef CSFunDef + | Namespace String [CSStmt] + | ClassDef CSClassDef + | ConstructorDef CSConstructorDef + | StructDef String [(CSType, String)] + + -- Some arbitrary string of CS code. + | Escape String + deriving (Eq, Show) + +instance Pretty CSStmt where + ppr (If cond tbranch []) = + text "if" <+> parens(ppr cond) + lbrace + indent 4 (stack $ map ppr tbranch) + rbrace + + ppr (If cond tbranch fbranch) = + text "if" <+> parens(ppr cond) + lbrace + indent 4 (stack $ map ppr tbranch) + rbrace + text "else" + lbrace + indent 4 (stack $ map ppr fbranch) + rbrace + + ppr (Try stmts excepts) = + text "try" + lbrace + indent 4 (stack $ map ppr stmts) + rbrace + stack (map ppr excepts) + + ppr (While cond body) = + text "while" <+> parens(ppr cond) + lbrace + indent 4 (stack $ map ppr body) + rbrace + + ppr (For i what body) = + text "for" <+> parens(initialize <> limit <> inc) + lbrace + indent 4 (stack $ map ppr body) + rbrace + where initialize = text "int" <+> text i <+> text "= 0" <+> semi + limit = text i <+> langle <+> ppr what <+> semi + inc = text i <> text "++" + + ppr (ForEach i what body) = + text "foreach" <+> parens initialize + lbrace + indent 4 (stack $ map ppr body) + rbrace + where initialize = text "var" <+> text i <+> text "in " <+> ppr what + + ppr (Using (Just as) from) = + text "using" <+> text as <+> text "=" <+> text from <> semi + + ppr (Using Nothing from) = + text "using" <+> text from <> semi + + ppr (Unsafe stmts) = + text "unsafe" + lbrace + indent 4 (stack $ map ppr stmts) + rbrace + + ppr (Fixed ptr e stmts) = + text "fixed" <+> parens(text "void*" <+> ppr ptr <+> text "=" <+> ppr e) + lbrace + indent 4 (stack $ map ppr stmts) + rbrace + + ppr (UsingWith assignment body) = + text "using" <+> parens(ppr assignment) + lbrace + indent 4 (stack $ map ppr body) + rbrace + + ppr (Assign e1 e2) = text "var" <+> ppr e1 <+> equals <+> ppr e2 <> semi + ppr (Reassign e1 e2) = ppr e1 <+> equals <+> ppr e2 <> semi + ppr (AssignTyped t e1 Nothing) = ppr t <+> ppr e1 <> semi + ppr (AssignTyped t e1 (Just e2)) = ppr t <+> ppr e1 <+> equals <+> ppr e2 <> semi + + ppr (AssignOp op e1 e2) = ppr e1 <+> text (op ++ "=") <+> ppr e2 <> semi + + ppr (Comment s body) = text "//" <> text s stack (map ppr body) + + ppr (Assert e []) = + text "FutharkAssert" <> parens(ppr e) <> semi + + ppr (Assert e exps) = + let exps' = stack $ intersperse (text ",") $ map ppr exps + formattedString = text "String.Format" <> parens exps' + in text "FutharkAssert" <> parens(ppr e <> text "," <+> formattedString) <> semi + + ppr (Throw e) = text "throw" <+> ppr e <> semi + + ppr (Exp e) = ppr e <> semi + + ppr (Return e) = text "return" <+> ppr e <> semi + + ppr (ClassDef d) = ppr d + + ppr (StaticFunDef d) = text "static" <+> ppr d + + ppr (PublicFunDef d) = text "public" <+> ppr d + + ppr (PrivateFunDef d) = text "private" <+> ppr d + + ppr (ConstructorDef d) = ppr d + + ppr (StructDef name assignments) = text "public struct" <+> text name <> braces(stack $ map (\(tp,field) -> text "public" <+> ppr tp <+> text field <> semi) assignments) + + ppr (Namespace name csstms) = text "namespace" <+> text name + lbrace + indent 4 (stack $ map ppr csstms) + rbrace + + ppr (Escape s) = stack $ map text $ lines s + + ppr Pass = empty + +instance Pretty CSFunDef where + ppr (Def fname retType args stmts) = + ppr retType <+> text fname <> parens( commasep(map ppr' args) ) + lbrace + indent 4 (stack (map ppr stmts)) + rbrace + where ppr' (tp, var) = ppr tp <+> text var + +instance Pretty CSClassDef where + ppr (Class cname body) = + text "class" <+> text cname + lbrace + indent 4 (stack (map ppr body)) + rbrace + + ppr (PublicClass cname body) = + text "public" <+> text "class" <+> text cname + lbrace + indent 4 (stack (map ppr body)) + rbrace + +instance Pretty CSConstructorDef where + ppr (ClassConstructor cname params body) = + text "public" <+> text cname <> parens(commasep $ map ppr' params) + lbrace + indent 4 (stack (map ppr body)) + rbrace + where ppr' (tp, var) = ppr tp <+> text var + +instance Pretty CSExcept where + ppr (Catch csexp stmts) = + text "catch" <+> parens(ppr csexp <+> text "e") + lbrace + indent 4 (stack (map ppr stmts)) + rbrace + +data CSExcept = Catch CSExp [CSStmt] + deriving (Eq, Show) + +type CSFunDefArg = (CSType, String) +data CSFunDef = Def String CSType [CSFunDefArg] [CSStmt] + deriving (Eq, Show) + +data CSClassDef = Class String [CSStmt] + | PublicClass String [CSStmt] + deriving (Eq, Show) + +data CSConstructorDef = ClassConstructor String [CSFunDefArg] [CSStmt] + deriving (Eq, Show) + +newtype CSProg = CSProg [CSStmt] + deriving (Eq, Show) + +instance Pretty CSProg where + ppr (CSProg stms) = stack (map ppr stms) diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp/Definitions.hs b/src/Futhark/CodeGen/Backends/GenericCSharp/Definitions.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericCSharp/Definitions.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE TemplateHaskell #-} +module Futhark.CodeGen.Backends.GenericCSharp.Definitions + ( csFunctions + , csReader + , csMemory + , csMemoryOpenCL + , csScalar + , csPanic + , csExceptions + , csOpenCL + ) where + +import Data.FileEmbed + +csFunctions :: String +csFunctions = $(embedStringFile "rts/csharp/functions.cs") + +csMemory :: String +csMemory = $(embedStringFile "rts/csharp/memory.cs") + +csScalar :: String +csScalar = $(embedStringFile "rts/csharp/scalar.cs") + +csReader :: String +csReader = $(embedStringFile "rts/csharp/reader.cs") + +csPanic :: String +csPanic = $(embedStringFile "rts/csharp/panic.cs") + +csExceptions :: String +csExceptions = $(embedStringFile "rts/csharp/exceptions.cs") + +csOpenCL :: String +csOpenCL = $(embedStringFile "rts/csharp/opencl.cs") + +csMemoryOpenCL :: String +csMemoryOpenCL = $(embedStringFile "rts/csharp/memory_opencl.cs") diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp/Options.hs b/src/Futhark/CodeGen/Backends/GenericCSharp/Options.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericCSharp/Options.hs @@ -0,0 +1,46 @@ +-- | This module defines a generator for @getopt@ based command +-- line argument parsing. Each option is associated with arbitrary +-- Python code that will perform side effects, usually by setting some +-- global variables. +module Futhark.CodeGen.Backends.GenericCSharp.Options + ( Option (..) + , OptionArgument (..) + , generateOptionParser + ) + where + +import Futhark.CodeGen.Backends.GenericCSharp.AST + +-- | Specification if a single command line option. The option must +-- have a long name, and may also have a short name. +-- +-- When the statement is being executed, the argument (if any) will be +-- stored in the variable @optarg@. +data Option = Option { optionLongName :: String + , optionShortName :: Maybe Char + , optionArgument :: OptionArgument + , optionAction :: [CSStmt] + } + +-- | Whether an option accepts an argument. +data OptionArgument = NoArgument + | RequiredArgument + | OptionalArgument + +-- | Generate option parsing code that accepts the given command line options. Will read from @sys.argv@. +-- +-- If option parsing fails for any reason, the entire process will +-- terminate with error code 1. +generateOptionParser :: [Option] -> [CSStmt] +generateOptionParser options = + [ Assign (Var "options") (Collection "OptionSet" $ map parseOption options) + , Assign (Var "extra") (Call (Var "options.Parse") [Arg Nothing (Var "args")]) + ] + where parseOption option = Array [ String $ option_string option + , Lambda (Var "optarg") $ optionAction option ] + option_string option = case optionArgument option of + RequiredArgument -> + concat [maybe "" prefix $ optionShortName option,optionLongName option,"="] + _ -> + maybe "" prefix (optionShortName option) ++ optionLongName option + prefix = flip (:) "|" diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericPython.hs @@ -0,0 +1,978 @@ +{-# LANGUAGE OverloadedStrings, GeneralizedNewtypeDeriving, LambdaCase #-} +{-# LANGUAGE TupleSections #-} +-- | A generic Python code generator which is polymorphic in the type +-- of the operations. Concretely, we use this to handle both +-- sequential and PyOpenCL Python code. +module Futhark.CodeGen.Backends.GenericPython + ( compileProg + , Constructor (..) + , emptyConstructor + + , compileName + , compileDim + , compileExp + , compileCode + , compilePrimValue + , compilePrimType + , compilePrimTypeExt + , compilePrimToNp + , compilePrimToExtNp + + , Operations (..) + , defaultOperations + + , unpackDim + + , CompilerM (..) + , OpCompiler + , WriteScalar + , ReadScalar + , Allocate + , Copy + , StaticArray + , EntryOutput + , EntryInput + + , CompilerEnv(..) + , CompilerState(..) + , stm + , stms + , atInit + , collect' + , collect + , simpleCall + + , copyMemoryDefaultSpace + ) where + +import Control.Monad.Identity +import Control.Monad.State +import Control.Monad.Reader +import Control.Monad.Writer +import Control.Monad.RWS +import Data.Maybe +import Data.List +import qualified Data.Map.Strict as M + +import Futhark.Representation.Primitive hiding (Bool) +import Futhark.MonadFreshNames +import Futhark.Representation.AST.Syntax (Space(..)) +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.CodeGen.Backends.GenericPython.AST +import Futhark.CodeGen.Backends.GenericPython.Options +import Futhark.CodeGen.Backends.GenericPython.Definitions +import Futhark.Util.Pretty(pretty) +import Futhark.Util (zEncodeString) +import Futhark.Representation.AST.Attributes (builtInFunctions, isBuiltInFunction) + +-- | A substitute expression compiler, tried before the main +-- compilation function. +type OpCompiler op s = op -> CompilerM op s () + +-- | Write a scalar to the given memory block with the given index and +-- in the given memory space. +type WriteScalar op s = VName -> PyExp -> PrimType -> Imp.SpaceId -> PyExp + -> CompilerM op s () + +-- | Read a scalar from the given memory block with the given index and +-- in the given memory space. +type ReadScalar op s = VName -> PyExp -> PrimType -> Imp.SpaceId + -> CompilerM op s PyExp + +-- | Allocate a memory block of the given size in the given memory +-- space, saving a reference in the given variable name. +type Allocate op s = VName -> PyExp -> Imp.SpaceId + -> CompilerM op s () + +-- | Copy from one memory block to another. +type Copy op s = VName -> PyExp -> Imp.Space -> + VName -> PyExp -> Imp.Space -> + PyExp -> PrimType -> + CompilerM op s () + +-- | Create a static array of values - initialised at load time. +type StaticArray op s = VName -> Imp.SpaceId -> PrimType -> [PrimValue] -> CompilerM op s () + +-- | Construct the Python array being returned from an entry point. +type EntryOutput op s = VName -> Imp.SpaceId -> + PrimType -> Imp.Signedness -> + [Imp.DimSize] -> + CompilerM op s PyExp + +-- | Unpack the array being passed to an entry point. +type EntryInput op s = VName -> Imp.MemSize -> Imp.SpaceId -> + PrimType -> Imp.Signedness -> + [Imp.DimSize] -> + PyExp -> + CompilerM op s () + + +data Operations op s = Operations { opsWriteScalar :: WriteScalar op s + , opsReadScalar :: ReadScalar op s + , opsAllocate :: Allocate op s + , opsCopy :: Copy op s + , opsStaticArray :: StaticArray op s + , opsCompiler :: OpCompiler op s + , opsEntryOutput :: EntryOutput op s + , opsEntryInput :: EntryInput op s + } + +-- | A set of operations that fail for every operation involving +-- non-default memory spaces. Uses plain pointers and @malloc@ for +-- memory management. +defaultOperations :: Operations op s +defaultOperations = Operations { opsWriteScalar = defWriteScalar + , opsReadScalar = defReadScalar + , opsAllocate = defAllocate + , opsCopy = defCopy + , opsStaticArray = defStaticArray + , opsCompiler = defCompiler + , opsEntryOutput = defEntryOutput + , opsEntryInput = defEntryInput + } + where defWriteScalar _ _ _ _ _ = + fail "Cannot write to non-default memory space because I am dumb" + defReadScalar _ _ _ _ = + fail "Cannot read from non-default memory space" + defAllocate _ _ _ = + fail "Cannot allocate in non-default memory space" + defCopy _ _ _ _ _ _ _ _ = + fail "Cannot copy to or from non-default memory space" + defStaticArray _ _ _ _ = + fail "Cannot create static array in non-default memory space" + defCompiler _ = + fail "The default compiler cannot compile extended operations" + defEntryOutput _ _ _ _ = + fail "Cannot return array not in default memory space" + defEntryInput _ _ _ _ = + fail "Cannot accept array not in default memory space" + +data CompilerEnv op s = CompilerEnv { + envOperations :: Operations op s + , envFtable :: M.Map Name [Imp.Type] +} + +envOpCompiler :: CompilerEnv op s -> OpCompiler op s +envOpCompiler = opsCompiler . envOperations + +envReadScalar :: CompilerEnv op s -> ReadScalar op s +envReadScalar = opsReadScalar . envOperations + +envWriteScalar :: CompilerEnv op s -> WriteScalar op s +envWriteScalar = opsWriteScalar . envOperations + +envAllocate :: CompilerEnv op s -> Allocate op s +envAllocate = opsAllocate . envOperations + +envCopy :: CompilerEnv op s -> Copy op s +envCopy = opsCopy . envOperations + +envStaticArray :: CompilerEnv op s -> StaticArray op s +envStaticArray = opsStaticArray . envOperations + +envEntryOutput :: CompilerEnv op s -> EntryOutput op s +envEntryOutput = opsEntryOutput . envOperations + +envEntryInput :: CompilerEnv op s -> EntryInput op s +envEntryInput = opsEntryInput . envOperations + +newCompilerEnv :: Imp.Functions op -> Operations op s -> CompilerEnv op s +newCompilerEnv (Imp.Functions funs) ops = + CompilerEnv { envOperations = ops + , envFtable = ftable <> builtinFtable + } + where ftable = M.fromList $ map funReturn funs + funReturn (name, Imp.Function _ outparams _ _ _ _) = (name, paramsTypes outparams) + builtinFtable = M.map (map Imp.Scalar . snd) builtInFunctions + +data CompilerState s = CompilerState { + compNameSrc :: VNameSource + , compInit :: [PyStmt] + , compUserState :: s +} + +newCompilerState :: VNameSource -> s -> CompilerState s +newCompilerState src s = CompilerState { compNameSrc = src + , compInit = [] + , compUserState = s } + +newtype CompilerM op s a = CompilerM (RWS (CompilerEnv op s) [PyStmt] (CompilerState s) a) + deriving (Functor, Applicative, Monad, + MonadState (CompilerState s), + MonadReader (CompilerEnv op s), + MonadWriter [PyStmt]) + +instance MonadFreshNames (CompilerM op s) where + getNameSource = gets compNameSrc + putNameSource src = modify $ \s -> s { compNameSrc = src } + +collect :: CompilerM op s () -> CompilerM op s [PyStmt] +collect m = pass $ do + ((), w) <- listen m + return (w, const mempty) + +collect' :: CompilerM op s a -> CompilerM op s (a, [PyStmt]) +collect' m = pass $ do + (x, w) <- listen m + return ((x, w), const mempty) + +atInit :: PyStmt -> CompilerM op s () +atInit x = modify $ \s -> + s { compInit = compInit s ++ [x] } + +stm :: PyStmt -> CompilerM op s () +stm x = tell [x] + +stms :: [PyStmt] -> CompilerM op s () +stms = mapM_ stm + +futharkFun :: String -> String +futharkFun s = "futhark_" ++ zEncodeString s + +paramsTypes :: [Imp.Param] -> [Imp.Type] +paramsTypes = map paramType + where paramType (Imp.MemParam _ space) = Imp.Mem (Imp.ConstSize 0) space + paramType (Imp.ScalarParam _ t) = Imp.Scalar t + +compileOutput :: [Imp.Param] -> [PyExp] +compileOutput = map (Var . compileName . Imp.paramName) + +runCompilerM :: Imp.Functions op -> Operations op s + -> VNameSource + -> s + -> CompilerM op s a + -> a +runCompilerM prog ops src userstate (CompilerM m) = + fst $ evalRWS m (newCompilerEnv prog ops) (newCompilerState src userstate) + +standardOptions :: [Option] +standardOptions = [ + Option { optionLongName = "write-runtime-to" + , optionShortName = Just 't' + , optionArgument = RequiredArgument + , optionAction = + [ + If (Var "runtime_file") + [Exp $ simpleCall "runtime_file.close" []] [] + , Assign (Var "runtime_file") $ + simpleCall "open" [Var "optarg", String "w"] + ] + }, + Option { optionLongName = "runs" + , optionShortName = Just 'r' + , optionArgument = RequiredArgument + , optionAction = + [ Assign (Var "num_runs") $ Var "optarg" + , Assign (Var "do_warmup_run") $ Bool True + ] + }, + Option { optionLongName = "entry-point" + , optionShortName = Just 'e' + , optionArgument = RequiredArgument + , optionAction = + [ Assign (Var "entry_point") $ Var "optarg" ] + }, + -- The -b option is just a dummy for now. + Option { optionLongName = "binary-output" + , optionShortName = Just 'b' + , optionArgument = NoArgument + , optionAction = [Pass] + } + ] + + +-- | The class generated by the code generator must have a +-- constructor, although it can be vacuous. +data Constructor = Constructor [String] [PyStmt] + +-- | A constructor that takes no arguments and does nothing. +emptyConstructor :: Constructor +emptyConstructor = Constructor ["self"] [Pass] + +constructorToFunDef :: Constructor -> [PyStmt] -> PyFunDef +constructorToFunDef (Constructor params body) at_init = + Def "__init__" params $ body <> at_init + +compileProg :: MonadFreshNames m => + Maybe String + -> Constructor + -> [PyStmt] + -> [PyStmt] + -> Operations op s + -> s + -> [PyStmt] + -> [Option] + -> Imp.Functions op + -> m String +compileProg module_name constructor imports defines ops userstate pre_timing options prog@(Imp.Functions funs) = do + src <- getNameSource + let prog' = runCompilerM prog ops src userstate compileProg' + maybe_shebang = + case module_name of Nothing -> "#!/usr/bin/env python\n" + Just _ -> "" + return $ maybe_shebang ++ + pretty (PyProg $ imports ++ + [Import "argparse" Nothing] ++ + defines ++ + [Escape pyUtility] ++ + prog') + where compileProg' = do + definitions <- mapM compileFunc funs + at_inits <- gets compInit + + let constructor' = constructorToFunDef constructor at_inits + + case module_name of + Just name -> do + (entry_points, entry_point_types) <- + unzip <$> mapM compileEntryFun (filter (Imp.functionEntry . snd) funs) + return [ClassDef $ Class name $ + Assign (Var "entry_points") (Dict entry_point_types) : + map FunDef (constructor' : definitions ++ entry_points)] + Nothing -> do + let classinst = Assign (Var "self") $ simpleCall "internal" [] + (entry_point_defs, entry_point_names, entry_points) <- + unzip3 <$> mapM (callEntryFun pre_timing) + (filter (Imp.functionEntry . snd) funs) + return (parse_options ++ + ClassDef (Class "internal" $ map FunDef $ + constructor' : definitions) : + classinst : + map FunDef entry_point_defs ++ + selectEntryPoint entry_point_names entry_points) + + parse_options = + Assign (Var "runtime_file") None : + Assign (Var "do_warmup_run") (Bool False) : + Assign (Var "num_runs") (Integer 1) : + Assign (Var "entry_point") (String "main") : + generateOptionParser (standardOptions ++ options) + + selectEntryPoint entry_point_names entry_points = + [ Assign (Var "entry_points") $ + Dict $ zip (map String entry_point_names) entry_points, + Assign (Var "entry_point_fun") $ + simpleCall "entry_points.get" [Var "entry_point"], + If (BinOp "==" (Var "entry_point_fun") None) + [Exp $ simpleCall "sys.exit" + [Call (Field + (String "No entry point '{}'. Select another with --entry point. Options are:\n{}") + "format") + [Arg $ Var "entry_point", + + Arg $ Call (Field (String "\n") "join") + [Arg $ simpleCall "entry_points.keys" []]]]] + [Exp $ simpleCall "entry_point_fun" []] + ] + +compileFunc :: (Name, Imp.Function op) -> CompilerM op s PyFunDef +compileFunc (fname, Imp.Function _ outputs inputs body _ _) = do + body' <- collect $ compileCode body + let inputs' = map (compileName . Imp.paramName) inputs + let ret = Return $ tupleOrSingle $ compileOutput outputs + return $ Def (futharkFun . nameToString $ fname) ("self" : inputs') (body'++[ret]) + +tupleOrSingle :: [PyExp] -> PyExp +tupleOrSingle [e] = e +tupleOrSingle es = Tuple es + +-- | A 'Call' where the function is a variable and every argument is a +-- simple 'Arg'. +simpleCall :: String -> [PyExp] -> PyExp +simpleCall fname = Call (Var fname) . map Arg + +compileName :: VName -> String +compileName = zEncodeString . pretty + +compileDim :: Imp.DimSize -> PyExp +compileDim (Imp.ConstSize i) = Integer $ toInteger i +compileDim (Imp.VarSize v) = Var $ compileName v + +unpackDim :: PyExp -> Imp.DimSize -> Int32 -> CompilerM op s () +unpackDim arr_name (Imp.ConstSize c) i = do + let shape_name = Field arr_name "shape" + let constant_c = Integer $ toInteger c + let constant_i = Integer $ toInteger i + stm $ Assert (BinOp "==" constant_c (Index shape_name $ IdxExp constant_i)) $ + String "constant dimension wrong" + +unpackDim arr_name (Imp.VarSize var) i = do + let shape_name = Field arr_name "shape" + src = Index shape_name $ IdxExp $ Integer $ toInteger i + stm $ Assign (Var $ compileName var) $ simpleCall "np.int32" [src] + +entryPointOutput :: Imp.ExternalValue -> CompilerM op s PyExp +entryPointOutput (Imp.OpaqueValue desc vs) = + simpleCall "opaque" . (String (pretty desc):) <$> + mapM (entryPointOutput . Imp.TransparentValue) vs +entryPointOutput (Imp.TransparentValue (Imp.ScalarValue bt ept name)) = + return $ simpleCall tf [Var $ compileName name] + where tf = compilePrimToExtNp bt ept +entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ Imp.DefaultSpace bt ept dims)) = do + let cast = Cast (Var $ compileName mem) (compilePrimTypeExt bt ept) + return $ simpleCall "createArray" [cast, Tuple $ map compileDim dims] +entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ (Imp.Space sid) bt ept dims)) = do + pack_output <- asks envEntryOutput + pack_output mem sid bt ept dims + +badInput :: Int -> PyExp -> String -> PyStmt +badInput i e t = + Raise $ simpleCall "TypeError" + [Call (Field (String err_msg) "format") + [Arg (String t), Arg $ simpleCall "type" [e], Arg e]] + where err_msg = unlines [ "Argument #" ++ show i ++ " has invalid value" + , "Futhark type: {}" + , "Argument has Python type {} and value: {}"] + + +entryPointInput :: (Int, Imp.ExternalValue, PyExp) -> CompilerM op s () +entryPointInput (i, Imp.OpaqueValue desc vs, e) = do + let type_is_ok = BinOp "and" (simpleCall "isinstance" [e, Var "opaque"]) + (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) $ + map (Index (Field e "data") . IdxExp . Integer) [0..] + +entryPointInput (i, Imp.TransparentValue (Imp.ScalarValue bt s name), e) = do + let vname' = Var $ compileName name + -- HACK: A Numpy int64 will signal an OverflowError if we pass + -- it a number bigger than 2**63. This does not happen if we + -- pass e.g. int8 a number bigger than 2**7. As a workaround, + -- we first go through the corresponding ctypes type, which does + -- not have this problem. + ctobject = compilePrimType bt + ctcall = simpleCall ctobject [e] + npobject = compilePrimToNp bt + npcall = simpleCall npobject [ctcall] + stm $ Try [Assign vname' npcall] + [Catch (Tuple [Var "TypeError", Var "AssertionError"]) + [badInput i e $ prettySigned (s==Imp.TypeUnsigned) bt]] + +entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem memsize Imp.DefaultSpace t s dims), e) = do + let type_is_wrong = + UnOp "not" $ + BinOp "and" + (BinOp "in" (simpleCall "type" [e]) (List [Var "np.ndarray"])) + (BinOp "==" (Field e "dtype") (Var (compilePrimToExtNp t s))) + stm $ If type_is_wrong + [badInput i e $ concat (replicate (length dims) "[]") ++ + prettySigned (s==Imp.TypeUnsigned) t] + [] + + zipWithM_ (unpackDim e) dims [0..] + let dest = Var $ compileName mem + unwrap_call = simpleCall "unwrapArray" [e] + + case memsize of + Imp.VarSize sizevar -> + stm $ Assign (Var $ compileName sizevar) $ + simpleCall "np.int32" [Field e "nbytes"] + Imp.ConstSize _ -> + return () + + stm $ Assign dest unwrap_call + +entryPointInput (i, Imp.TransparentValue (Imp.ArrayValue mem memsize (Imp.Space sid) bt ept dims), e) = do + unpack_input <- asks envEntryInput + unpack <- collect $ unpack_input mem memsize sid bt ept dims e + stm $ Try unpack + [Catch (Tuple [Var "TypeError", Var "AssertionError"]) + [badInput i e $ concat (replicate (length dims) "[]") ++ + prettySigned (ept==Imp.TypeUnsigned) bt]] + +extValueDescName :: Imp.ExternalValue -> String +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 +extName = (++"_ext") + +valueDescName :: Imp.ValueDesc -> String +valueDescName = compileName . valueDescVName + +valueDescVName :: Imp.ValueDesc -> VName +valueDescVName (Imp.ScalarValue _ _ vname) = vname +valueDescVName (Imp.ArrayValue vname _ _ _ _ _) = vname + +-- Key into the FUTHARK_PRIMTYPES dict. +readTypeEnum :: PrimType -> Imp.Signedness -> String +readTypeEnum (IntType Int8) Imp.TypeUnsigned = "u8" +readTypeEnum (IntType Int16) Imp.TypeUnsigned = "u16" +readTypeEnum (IntType Int32) Imp.TypeUnsigned = "u32" +readTypeEnum (IntType Int64) Imp.TypeUnsigned = "u64" +readTypeEnum (IntType Int8) Imp.TypeDirect = "i8" +readTypeEnum (IntType Int16) Imp.TypeDirect = "i16" +readTypeEnum (IntType Int32) Imp.TypeDirect = "i32" +readTypeEnum (IntType Int64) Imp.TypeDirect = "i64" +readTypeEnum (FloatType Float32) _ = "f32" +readTypeEnum (FloatType Float64) _ = "f64" +readTypeEnum Imp.Bool _ = "bool" +readTypeEnum Cert _ = error "readTypeEnum: cert" + +readInput :: Imp.ExternalValue -> PyStmt +readInput (Imp.OpaqueValue desc _) = + Raise $ simpleCall "Exception" + [String $ "Cannot read argument of type " ++ desc ++ "."] + +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)) = + let type_name = readTypeEnum bt ept + in Assign (Var $ extValueDescName decl) $ simpleCall "read_value" + [String $ concat (replicate (length dims) "[]") ++ type_name] + +printValue :: [(Imp.ExternalValue, PyExp)] -> CompilerM op s [PyStmt] +printValue = fmap concat . mapM (uncurry printValue') + -- We copy non-host arrays to the host before printing. This is + -- done in a hacky way - we assume the value has a .get()-method + -- that returns an equivalent Numpy array. This works for PyOpenCL, + -- but we will probably need yet another plugin mechanism here in + -- the future. + where printValue' (Imp.OpaqueValue desc _) _ = + return [Exp $ simpleCall "sys.stdout.write" + [String $ "#"]] + printValue' (Imp.TransparentValue (Imp.ArrayValue mem memsize (Space _) bt ept shape)) e = + printValue' (Imp.TransparentValue (Imp.ArrayValue mem memsize DefaultSpace bt ept shape)) $ + simpleCall (pretty e ++ ".get") [] + printValue' (Imp.TransparentValue _) e = + return [Exp $ simpleCall "write_value" [e], + Exp $ simpleCall "sys.stdout.write" [String "\n"]] + +prepareEntry :: (Name, Imp.Function op) -> CompilerM op s + (String, [String], [PyStmt], [PyStmt], [PyStmt], [PyStmt], + [(Imp.ExternalValue, PyExp)], [PyStmt]) +prepareEntry (fname, Imp.Function _ outputs inputs _ results args) = do + let output_paramNames = map (compileName . Imp.paramName) outputs + funTuple = tupleOrSingle $ fmap Var output_paramNames + + (argexps_mem_copies, prepare_run) <- collect' $ forM inputs $ \case + Imp.MemParam name space -> do + -- A program might write to its input parameters, so create a new memory + -- block and copy the source there. This way the program can be run more + -- than once. + name' <- newVName $ baseString name <> "_copy" + copy <- asks envCopy + allocate <- asks envAllocate + let size = Var (extName (compileName name) ++ ".nbytes") -- FIXME + dest = name' + src = name + offset = Integer 0 + case space of + DefaultSpace -> + stm $ Assign (Var (compileName name')) + (simpleCall "allocateMem" [size]) -- FIXME + Space sid -> + allocate name' size sid + copy dest offset space src offset space size (IntType Int32) -- FIXME + return $ Just $ compileName name' + _ -> return Nothing + + prepareIn <- collect $ mapM_ entryPointInput $ zip3 [0..] args $ + map (Var . extValueDescName) args + (res, prepareOut) <- collect' $ mapM entryPointOutput results + + let argexps_lib = map (compileName . Imp.paramName) inputs + argexps_bin = zipWith fromMaybe argexps_lib argexps_mem_copies + fname' = "self." ++ futharkFun (nameToString fname) + call_lib = [Assign funTuple $ simpleCall fname' (fmap Var argexps_lib)] + call_bin = [Assign funTuple $ simpleCall fname' (fmap Var argexps_bin)] + + return (nameToString fname, map extValueDescName args, + prepareIn, call_lib, call_bin, prepareOut, + zip results res, prepare_run) + +copyMemoryDefaultSpace :: VName -> PyExp -> VName -> PyExp -> PyExp -> + CompilerM op s () +copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = do + let offset_call1 = simpleCall "addressOffset" + [Var (compileName destmem), destidx, Var "ct.c_byte"] + let offset_call2 = simpleCall "addressOffset" + [Var (compileName srcmem), srcidx, Var "ct.c_byte"] + stm $ Exp $ simpleCall "ct.memmove" [offset_call1, offset_call2, nbytes] + +compileEntryFun :: (Name, Imp.Function op) + -> CompilerM op s (PyFunDef, (PyExp, PyExp)) +compileEntryFun entry = do + (fname', params, prepareIn, body_lib, _, prepareOut, res, _) <- prepareEntry entry + let ret = Return $ tupleOrSingle $ map snd res + (pts, rts) = entryTypes $ snd entry + return (Def fname' ("self" : params) $ + prepareIn ++ body_lib ++ prepareOut ++ [ret], + (String fname', Tuple [List (map String pts), List (map String rts)])) + +entryTypes :: Imp.Function op -> ([String], [String]) +entryTypes func = (map desc $ Imp.functionArgs func, + 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 + +callEntryFun :: [PyStmt] -> (Name, Imp.Function op) + -> CompilerM op s (PyFunDef, String, PyExp) +callEntryFun pre_timing entry@(fname, Imp.Function _ _ _ _ _ decl_args) = do + (_, _, prepareIn, _, body_bin, _, res, prepare_run) <- prepareEntry entry + + let str_input = map readInput decl_args + + exitcall = [Exp $ simpleCall "sys.exit" [Field (String "Assertion.{} failed") "format(e)"]] + except' = Catch (Var "AssertionError") exitcall + do_run = body_bin ++ pre_timing + (do_run_with_timing, close_runtime_file) = addTiming do_run + + -- We ignore overflow errors and the like for executable entry + -- points. These are (somewhat) well-defined in Futhark. + ignore s = ArgKeyword s $ String "ignore" + errstate = Call (Var "np.errstate") $ map ignore ["divide", "over", "under", "invalid"] + + do_warmup_run = + If (Var "do_warmup_run") (prepare_run ++ do_run) [] + + do_num_runs = + For "i" (simpleCall "range" [simpleCall "int" [Var "num_runs"]]) + (prepare_run ++ do_run_with_timing) + + str_output <- printValue res + + let fname' = "entry_" ++ nameToString fname + + return (Def fname' [] $ + str_input ++ prepareIn ++ + [Try [With errstate [do_warmup_run, do_num_runs]] [except']] ++ + [close_runtime_file] ++ + str_output, + + nameToString fname, + + Var fname') + +addTiming :: [PyStmt] -> ([PyStmt], PyStmt) +addTiming statements = + ([ Assign (Var "time_start") $ simpleCall "time.time" [] ] ++ + statements ++ + [ Assign (Var "time_end") $ simpleCall "time.time" [] + , If (Var "runtime_file") print_runtime [] ], + + If (Var "runtime_file") [Exp $ simpleCall "runtime_file.close" []] []) + where print_runtime = + [Exp $ simpleCall "runtime_file.write" + [simpleCall "str" + [BinOp "-" + (toMicroseconds (Var "time_end")) + (toMicroseconds (Var "time_start"))]], + Exp $ simpleCall "runtime_file.write" [String "\n"]] + toMicroseconds x = + simpleCall "int" [BinOp "*" x $ Integer 1000000] + +compileUnOp :: Imp.UnOp -> String +compileUnOp op = + case op of + Not -> "not" + Complement{} -> "~" + Abs{} -> "abs" + FAbs{} -> "abs" + SSignum{} -> "ssignum" + USignum{} -> "usignum" + +compileBinOpLike :: Monad m => + Imp.Exp -> Imp.Exp + -> CompilerM op s (PyExp, PyExp, String -> m PyExp) +compileBinOpLike x y = do + x' <- compileExp x + y' <- compileExp y + let simple s = return $ BinOp s x' y' + return (x', y', simple) + +-- | The ctypes type corresponding to a 'PrimType'. +compilePrimType :: PrimType -> String +compilePrimType t = + case t of + IntType Int8 -> "ct.c_int8" + IntType Int16 -> "ct.c_int16" + IntType Int32 -> "ct.c_int32" + IntType Int64 -> "ct.c_int64" + FloatType Float32 -> "ct.c_float" + FloatType Float64 -> "ct.c_double" + Imp.Bool -> "ct.c_bool" + Cert -> "ct.c_bool" + +-- | The ctypes type corresponding to a 'PrimType', taking sign into account. +compilePrimTypeExt :: PrimType -> Imp.Signedness -> String +compilePrimTypeExt t ept = + case (t, ept) of + (IntType Int8, Imp.TypeUnsigned) -> "ct.c_uint8" + (IntType Int16, Imp.TypeUnsigned) -> "ct.c_uint16" + (IntType Int32, Imp.TypeUnsigned) -> "ct.c_uint32" + (IntType Int64, Imp.TypeUnsigned) -> "ct.c_uint64" + (IntType Int8, _) -> "ct.c_int8" + (IntType Int16, _) -> "ct.c_int16" + (IntType Int32, _) -> "ct.c_int32" + (IntType Int64, _) -> "ct.c_int64" + (FloatType Float32, _) -> "ct.c_float" + (FloatType Float64, _) -> "ct.c_double" + (Imp.Bool, _) -> "ct.c_bool" + (Cert, _) -> "ct.c_byte" + +-- | The Numpy type corresponding to a 'PrimType'. +compilePrimToNp :: Imp.PrimType -> String +compilePrimToNp bt = + case bt of + IntType Int8 -> "np.int8" + IntType Int16 -> "np.int16" + IntType Int32 -> "np.int32" + IntType Int64 -> "np.int64" + FloatType Float32 -> "np.float32" + FloatType Float64 -> "np.float64" + Imp.Bool -> "np.byte" + Cert -> "np.byte" + +-- | The Numpy type corresponding to a 'PrimType', taking sign into account. +compilePrimToExtNp :: Imp.PrimType -> Imp.Signedness -> String +compilePrimToExtNp bt ept = + case (bt,ept) of + (IntType Int8, Imp.TypeUnsigned) -> "np.uint8" + (IntType Int16, Imp.TypeUnsigned) -> "np.uint16" + (IntType Int32, Imp.TypeUnsigned) -> "np.uint32" + (IntType Int64, Imp.TypeUnsigned) -> "np.uint64" + (IntType Int8, _) -> "np.int8" + (IntType Int16, _) -> "np.int16" + (IntType Int32, _) -> "np.int32" + (IntType Int64, _) -> "np.int64" + (FloatType Float32, _) -> "np.float32" + (FloatType Float64, _) -> "np.float64" + (Imp.Bool, _) -> "np.bool" + (Cert, _) -> "np.byte" + +compilePrimValue :: Imp.PrimValue -> PyExp +compilePrimValue (IntValue (Int8Value v)) = + simpleCall "np.int8" [Integer $ toInteger v] +compilePrimValue (IntValue (Int16Value v)) = + simpleCall "np.int16" [Integer $ toInteger v] +compilePrimValue (IntValue (Int32Value v)) = + simpleCall "np.int32" [Integer $ toInteger v] +compilePrimValue (IntValue (Int64Value v)) = + simpleCall "np.int64" [Integer $ toInteger v] +compilePrimValue (FloatValue (Float32Value v)) + | isInfinite v = + if v > 0 then Var "np.inf" else Var "-np.inf" + | isNaN v = + Var "np.nan" + | otherwise = simpleCall "np.float32" [Float $ fromRational $ toRational v] +compilePrimValue (FloatValue (Float64Value v)) + | isInfinite v = + if v > 0 then Var "np.inf" else Var "-np.inf" + | isNaN v = + Var "np.nan" + | otherwise = simpleCall "np.float64" [Float $ fromRational $ toRational v] +compilePrimValue (BoolValue v) = Bool v +compilePrimValue Checked = Var "True" + +compileExp :: Imp.Exp -> CompilerM op s PyExp + +compileExp (Imp.ValueExp v) = return $ compilePrimValue v + +compileExp (Imp.LeafExp (Imp.ScalarVar vname) _) = + return $ Var $ compileName vname + +compileExp (Imp.LeafExp (Imp.SizeOf t) _) = + return $ simpleCall (compilePrimToNp $ IntType Int32) [Integer $ primByteSize t] + +compileExp (Imp.LeafExp (Imp.Index src (Imp.Count iexp) bt DefaultSpace _) _) = do + iexp' <- compileExp iexp + let bt' = compilePrimType bt + let nptype = compilePrimToNp bt + return $ simpleCall "indexArray" [Var $ compileName src, iexp', Var bt', Var nptype] + +compileExp (Imp.LeafExp (Imp.Index src (Imp.Count iexp) restype (Imp.Space space) _) _) = + join $ asks envReadScalar + <*> pure src <*> compileExp iexp + <*> pure restype <*> pure space + +compileExp (Imp.BinOpExp op x y) = do + (x', y', simple) <- compileBinOpLike x y + case op of + Add{} -> simple "+" + Sub{} -> simple "-" + Mul{} -> simple "*" + FAdd{} -> simple "+" + FSub{} -> simple "-" + FMul{} -> simple "*" + FDiv{} -> simple "/" + Xor{} -> simple "^" + And{} -> simple "&" + Or{} -> simple "|" + Shl{} -> simple "<<" + LogAnd{} -> simple "and" + LogOr{} -> simple "or" + _ -> return $ simpleCall (pretty op) [x', y'] + +compileExp (Imp.ConvOpExp conv x) = do + x' <- compileExp x + return $ simpleCall (pretty conv) [x'] + +compileExp (Imp.CmpOpExp cmp x y) = do + (x', y', simple) <- compileBinOpLike x y + case cmp of + CmpEq{} -> simple "==" + FCmpLt{} -> simple "<" + FCmpLe{} -> simple "<=" + CmpLlt -> simple "<" + CmpLle -> simple "<=" + _ -> return $ simpleCall (pretty cmp) [x', y'] + +compileExp (Imp.UnOpExp op exp1) = + UnOp (compileUnOp op) <$> compileExp exp1 + +compileExp (Imp.FunExp h args _) = + simpleCall (futharkFun (pretty h)) <$> mapM compileExp args + +compileCode :: Imp.Code op -> CompilerM op s () + +compileCode Imp.DebugPrint{} = + return () + +compileCode (Imp.Op op) = + join $ asks envOpCompiler <*> pure op + +compileCode (Imp.If cond tb fb) = do + cond' <- compileExp cond + tb' <- collect $ compileCode tb + fb' <- collect $ compileCode fb + stm $ If cond' tb' fb' + +compileCode (c1 Imp.:>>: c2) = do + compileCode c1 + compileCode c2 + +compileCode (Imp.While cond body) = do + cond' <- compileExp cond + body' <- collect $ compileCode body + stm $ While cond' body' + +compileCode (Imp.For i it bound body) = do + bound' <- compileExp bound + let i' = compileName i + body' <- collect $ compileCode body + counter <- pretty <$> newVName "counter" + one <- pretty <$> newVName "one" + stm $ Assign (Var i') $ simpleCall (compilePrimToNp (IntType it)) [Integer 0] + stm $ Assign (Var one) $ simpleCall (compilePrimToNp (IntType it)) [Integer 1] + stm $ For counter (simpleCall "range" [bound']) $ + body' ++ [AssignOp "+" (Var i') (Var one)] + +compileCode (Imp.SetScalar vname exp1) = do + let name' = Var $ compileName vname + exp1' <- compileExp exp1 + stm $ Assign name' exp1' + +compileCode Imp.DeclareMem{} = return () +compileCode (Imp.DeclareScalar v Cert) = + stm $ Assign (Var $ compileName v) $ Var "True" +compileCode Imp.DeclareScalar{} = return () + +compileCode (Imp.DeclareArray name DefaultSpace t vs) = do + -- It is important to store the Numpy array in a temporary variable + -- to prevent it from going "out-of-scope" before calling + -- unwrapArray (which internally uses the .ctype method); see + -- https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.ctypes.html + atInit $ Assign (Field (Var "self") arr_name) $ Call (Var "np.array") + [Arg $ List $ map compilePrimValue vs, + ArgKeyword "dtype" $ Var $ compilePrimToNp t] + atInit $ + Assign (Field (Var "self") name') $ + simpleCall "unwrapArray" [Field (Var "self") arr_name] + stm $ Assign (Var name') $ Field (Var "self") name' + where name' = compileName name + arr_name = name' <> "_arr" + +compileCode (Imp.DeclareArray name (Space space) t vs) = + join $ asks envStaticArray <*> + pure name <*> pure space <*> pure t <*> pure vs + +compileCode (Imp.Comment s code) = do + code' <- collect $ compileCode code + stm $ Comment s code' + +compileCode (Imp.Assert e (Imp.ErrorMsg parts) (loc,locs)) = do + e' <- compileExp e + let onPart (Imp.ErrorString s) = return ("%s", String s) + onPart (Imp.ErrorInt32 x) = ("%d",) <$> compileExp x + (formatstrs, formatargs) <- unzip <$> mapM onPart parts + stm $ Assert e' (BinOp "%" + (String $ "Error at " ++ stacktrace ++ ": " ++ concat formatstrs) + (Tuple formatargs)) + where stacktrace = intercalate " -> " (reverse $ map locStr $ loc:locs) + +compileCode (Imp.Call dests fname args) = do + args' <- mapM compileArg args + let dests' = tupleOrSingle $ fmap Var (map compileName dests) + fname' + | isBuiltInFunction fname = futharkFun (pretty fname) + | otherwise = "self." ++ futharkFun (pretty fname) + call' = simpleCall fname' args' + -- If the function returns nothing (is called only for side + -- effects), take care not to assign to an empty tuple. + stm $ if null dests + then Exp call' + else Assign dests' call' + where compileArg (Imp.MemArg m) = return $ Var $ compileName m + compileArg (Imp.ExpArg e) = compileExp e + +compileCode (Imp.SetMem dest src _) = do + let src' = Var (compileName src) + let dest' = Var (compileName dest) + stm $ Assign dest' src' + +compileCode (Imp.Allocate name (Imp.Count e) DefaultSpace) = do + e' <- compileExp e + let allocate' = simpleCall "allocateMem" [e'] + let name' = Var (compileName name) + stm $ Assign name' allocate' + +compileCode (Imp.Free name _) = + stm $ Assign (Var (compileName name)) None + +compileCode (Imp.Allocate name (Imp.Count e) (Imp.Space space)) = + join $ asks envAllocate + <*> pure name + <*> compileExp e + <*> pure space + +compileCode (Imp.Copy dest (Imp.Count destoffset) DefaultSpace src (Imp.Count srcoffset) DefaultSpace (Imp.Count size)) = do + destoffset' <- compileExp destoffset + srcoffset' <- compileExp srcoffset + let dest' = Var (compileName dest) + let src' = Var (compileName src) + size' <- compileExp size + let offset_call1 = simpleCall "addressOffset" [dest', destoffset', Var "ct.c_byte"] + let offset_call2 = simpleCall "addressOffset" [src', srcoffset', Var "ct.c_byte"] + stm $ Exp $ simpleCall "ct.memmove" [offset_call1, offset_call2, size'] + +compileCode (Imp.Copy dest (Imp.Count destoffset) destspace src (Imp.Count srcoffset) srcspace (Imp.Count size)) = do + copy <- asks envCopy + join $ copy + <$> pure dest <*> compileExp destoffset <*> pure destspace + <*> pure src <*> compileExp srcoffset <*> pure srcspace + <*> compileExp size <*> pure (IntType Int32) -- FIXME + +compileCode (Imp.Write dest (Imp.Count idx) elemtype DefaultSpace _ elemexp) = do + idx' <- compileExp idx + elemexp' <- compileExp elemexp + let dest' = Var $ compileName dest + let elemtype' = compilePrimType elemtype + let ctype = simpleCall elemtype' [elemexp'] + stm $ Exp $ simpleCall "writeScalarArray" [dest', idx', ctype] + +compileCode (Imp.Write dest (Imp.Count idx) elemtype (Imp.Space space) _ elemexp) = + join $ asks envWriteScalar + <*> pure dest + <*> compileExp idx + <*> pure elemtype + <*> pure space + <*> compileExp elemexp + +compileCode Imp.Skip = return () diff --git a/src/Futhark/CodeGen/Backends/GenericPython/AST.hs b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs @@ -0,0 +1,203 @@ +module Futhark.CodeGen.Backends.GenericPython.AST + ( PyExp(..) + , PyIdx (..) + , PyArg (..) + , PyStmt(..) + , module Language.Futhark.Core + , PyProg(..) + , PyExcept(..) + , PyFunDef(..) + , PyClassDef(..) + ) + where + +import Language.Futhark.Core +import Futhark.Util.Pretty + + +data UnOp = Not -- ^ Boolean negation. + | Complement -- ^ Bitwise complement. + | Negate -- ^ Numerical negation. + | Abs -- ^ Absolute/numerical value. + deriving (Eq, Show) + +data PyExp = Integer Integer + | Bool Bool + | Float Double + | String String + | RawStringLiteral String + | Var String + | BinOp String PyExp PyExp + | UnOp String PyExp + | Cond PyExp PyExp PyExp + | Index PyExp PyIdx + | Call PyExp [PyArg] + | Cast PyExp String + | Tuple [PyExp] + | List [PyExp] + | Field PyExp String + | Dict [(PyExp, PyExp)] + | None + deriving (Eq, Show) + +data PyIdx = IdxRange PyExp PyExp + | IdxExp PyExp + deriving (Eq, Show) + +data PyArg = ArgKeyword String PyExp + | Arg PyExp + deriving (Eq, Show) + +data PyStmt = If PyExp [PyStmt] [PyStmt] + | Try [PyStmt] [PyExcept] + | While PyExp [PyStmt] + | For String PyExp [PyStmt] + | With PyExp [PyStmt] + | Assign PyExp PyExp + | AssignOp String PyExp PyExp + | Comment String [PyStmt] + | Assert PyExp PyExp + | Raise PyExp + | Exp PyExp + | Return PyExp + | Pass + + -- Definition-like statements. + | Import String (Maybe String) + | FunDef PyFunDef + | ClassDef PyClassDef + + -- Some arbitrary string of Python code. + | Escape String + deriving (Eq, Show) + +data PyExcept = Catch PyExp [PyStmt] + deriving (Eq, Show) + +data PyFunDef = Def String [String] [PyStmt] + deriving (Eq, Show) + +data PyClassDef = Class String [PyStmt] + deriving (Eq, Show) + +newtype PyProg = PyProg [PyStmt] + deriving (Eq, Show) + +instance Pretty PyIdx where + ppr (IdxExp e) = ppr e + ppr (IdxRange from to) = ppr from <> text ":" <> ppr to + +instance Pretty PyArg where + ppr (ArgKeyword k e) = text k <> equals <> ppr e + ppr (Arg e) = ppr e + + +instance Pretty PyExp where + ppr (Integer x) = ppr x + ppr (Bool x) = ppr x + ppr (Float x) + | isInfinite x = text $ if x > 0 then "float('inf')" else "float('-inf')" + | otherwise = ppr x + ppr (String x) = text $ show x + ppr (RawStringLiteral s) = text "\"\"\"" <> text s <> text "\"\"\"" + ppr (Var n) = text $ map (\x -> if x == '\'' then 'm' else x) n + ppr (Field e s) = ppr e <> text "." <> text s + ppr (BinOp s e1 e2) = parens(ppr e1 <+> text s <+> ppr e2) + ppr (UnOp s e) = text s <> parens (ppr e) + ppr (Cond e1 e2 e3) = ppr e2 <+> text "if" <+> ppr e1 <+> text "else" <+> ppr e3 + ppr (Cast src bt) = text "ct.cast" <> + parens (ppr src <> text "," <+> + text "ct.POINTER" <> parens(text bt)) + ppr (Index src idx) = ppr src <> brackets(ppr idx) + ppr (Call fun exps) = ppr fun <> parens(commasep $ map ppr exps) + ppr (Tuple [dim]) = parens(ppr dim <> text ",") + ppr (Tuple dims) = parens(commasep $ map ppr dims) + ppr (List es) = brackets $ commasep $ map ppr es + ppr (Dict kvs) = braces $ commasep $ map ppElem kvs + where ppElem (k, v) = ppr k <> colon <+> ppr v + + ppr None = text "None" + +instance Pretty PyStmt where + ppr (If cond [] []) = + text "if" <+> ppr cond <> text ":" + indent 2 (text "pass") + + ppr (If cond [] fbranch) = + text "if" <+> ppr cond <> text ":" + indent 2 (text "pass") + text "else:" + indent 2 (stack $ map ppr fbranch) + + ppr (If cond tbranch []) = + text "if" <+> ppr cond <> text ":" + indent 2 (stack $ map ppr tbranch) + + ppr (If cond tbranch fbranch) = + text "if" <+> ppr cond <> text ":" + indent 2 (stack $ map ppr tbranch) + text "else:" + indent 2 (stack $ map ppr fbranch) + + ppr (Try pystms pyexcepts) = + text "try:" + indent 2 (stack $ map ppr pystms) + stack (map ppr pyexcepts) + + ppr (While cond body) = + text "while" <+> ppr cond <> text ":" + indent 2 (stack $ map ppr body) + + ppr (For i what body) = + text "for" <+> ppr i <+> text "in" <+> ppr what <> text ":" + indent 2 (stack $ map ppr body) + + ppr (With what body) = + text "with" <+> ppr what <> text ":" + indent 2 (stack $ map ppr body) + + ppr (Assign e1 e2) = ppr e1 <+> text "=" <+> ppr e2 + + ppr (AssignOp op e1 e2) = ppr e1 <+> text (op ++ "=") <+> ppr e2 + + ppr (Comment s body) = text "#" <> text s stack (map ppr body) + + ppr (Assert e1 e2) = text "assert" <+> ppr e1 <> text "," <+> ppr e2 + + ppr (Raise e) = text "raise" <+> ppr e + + ppr (Exp c) = ppr c + + ppr (Return e) = text "return" <+> ppr e + + ppr Pass = text "pass" + + ppr (Import from (Just as)) = + text "import" <+> text from <+> text "as" <+> text as + + ppr (Import from Nothing) = + text "import" <+> text from + + ppr (FunDef d) = ppr d + + ppr (ClassDef d) = ppr d + + ppr (Escape s) = stack $ map text $ lines s + +instance Pretty PyFunDef where + ppr (Def fname params body) = + text "def" <+> text fname <> parens (commasep $ map ppr params) <> text ":" + indent 2 (stack (map ppr body)) + +instance Pretty PyClassDef where + ppr (Class cname body) = + text "class" <+> text cname <> text ":" + indent 2 (stack (map ppr body)) + +instance Pretty PyExcept where + ppr (Catch pyexp stms) = + text "except" <+> ppr pyexp <+> text "as e:" + indent 2 (stack $ map ppr stms) + +instance Pretty PyProg where + ppr (PyProg stms) = stack (map ppr stms) diff --git a/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs b/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs @@ -0,0 +1,21 @@ +{-# LANGUAGE TemplateHaskell #-} +module Futhark.CodeGen.Backends.GenericPython.Definitions + ( pyFunctions + , pyUtility + , pyValues + , pyPanic + ) where + +import Data.FileEmbed + +pyFunctions :: String +pyFunctions = $(embedStringFile "rts/python/memory.py") + +pyUtility :: String +pyUtility = $(embedStringFile "rts/python/scalar.py") + +pyValues :: String +pyValues = $(embedStringFile "rts/python/values.py") + +pyPanic :: String +pyPanic = $(embedStringFile "rts/python/panic.py") diff --git a/src/Futhark/CodeGen/Backends/GenericPython/Options.hs b/src/Futhark/CodeGen/Backends/GenericPython/Options.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericPython/Options.hs @@ -0,0 +1,71 @@ +-- | This module defines a generator for @getopt@ based command +-- line argument parsing. Each option is associated with arbitrary +-- Python code that will perform side effects, usually by setting some +-- global variables. +module Futhark.CodeGen.Backends.GenericPython.Options + ( Option (..) + , OptionArgument (..) + , generateOptionParser + ) + where + +import Futhark.CodeGen.Backends.GenericPython.AST + +-- | Specification if a single command line option. The option must +-- have a long name, and may also have a short name. +-- +-- When the statement is being executed, the argument (if any) will be +-- stored in the variable @optarg@. +data Option = Option { optionLongName :: String + , optionShortName :: Maybe Char + , optionArgument :: OptionArgument + , optionAction :: [PyStmt] + } + +-- | Whether an option accepts an argument. +data OptionArgument = NoArgument + | RequiredArgument + | OptionalArgument + +-- | Generate option parsing code that accepts the given command line options. Will read from @sys.argv@. +-- +-- If option parsing fails for any reason, the entire process will +-- terminate with error code 1. +generateOptionParser :: [Option] -> [PyStmt] +generateOptionParser options = + [Assign (Var "parser") + (Call (Var "argparse.ArgumentParser") + [ArgKeyword "description" $ + String "A compiled Futhark program."])] ++ + map parseOption options ++ + [Assign (Var "parser_result") $ + Call (Var "vars") [Arg $ Call (Var "parser.parse_args") [Arg $ Var "sys.argv[1:]"]]] ++ + map executeOption options + where parseOption option = + Exp $ Call (Var "parser.add_argument") $ + map (Arg . String) name_args ++ argument_args + where name_args = maybe id ((:) . ('-':) . (:[])) (optionShortName option) + ["--" ++ optionLongName option] + argument_args = case optionArgument option of + RequiredArgument -> + [ArgKeyword "action" (String "append"), + ArgKeyword "default" $ List []] + + NoArgument -> + [ArgKeyword "action" (String "append_const"), + ArgKeyword "default" $ List [], + ArgKeyword "const" None] + + OptionalArgument -> + [ArgKeyword "action" (String "append"), + ArgKeyword "default" $ List [], + ArgKeyword "nargs" $ String "?"] + + executeOption option = + For "optarg" (Index (Var "parser_result") $ + IdxExp $ String $ fieldName option) $ + optionAction option + + fieldName = map escape . optionLongName + where escape '-' = '_' + escape c = c diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs @@ -0,0 +1,295 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.CodeGen.Backends.PyOpenCL + ( compileProg + ) where + +import Control.Monad + +import Futhark.Error +import Futhark.Representation.ExplicitMemory (Prog, ExplicitMemory) +import Futhark.CodeGen.Backends.PyOpenCL.Boilerplate +import qualified Futhark.CodeGen.Backends.GenericPython as Py +import qualified Futhark.CodeGen.ImpCode.OpenCL as Imp +import qualified Futhark.CodeGen.ImpGen.OpenCL as ImpGen +import Futhark.CodeGen.Backends.GenericPython.AST +import Futhark.CodeGen.Backends.GenericPython.Options +import Futhark.CodeGen.Backends.GenericPython.Definitions +import Futhark.Util.Pretty(pretty) +import Futhark.MonadFreshNames + + +--maybe pass the config file rather than multiple arguments +compileProg :: MonadFreshNames m => + Maybe String -> Prog ExplicitMemory -> m (Either InternalError String) +compileProg module_name prog = do + res <- ImpGen.compileProg prog + --could probably be a better why do to this.. + case res of + Left err -> return $ Left err + Right (Imp.Program opencl_code opencl_prelude kernel_names types sizes prog') -> do + --prepare the strings for assigning the kernels and set them as global + let assign = unlines $ map (\x -> pretty $ Assign (Var ("self."++x++"_var")) (Var $ "program."++x)) kernel_names + + let defines = + [Assign (Var "synchronous") $ Bool False, + Assign (Var "preferred_platform") None, + Assign (Var "preferred_device") None, + Assign (Var "fut_opencl_src") $ RawStringLiteral $ opencl_prelude ++ opencl_code, + Escape pyValues, + Escape pyFunctions, + Escape pyPanic] + let imports = [Import "sys" Nothing, + Import "numpy" $ Just "np", + Import "ctypes" $ Just "ct", + Escape openClPrelude, + Import "pyopencl.array" Nothing, + Import "time" Nothing] + + let constructor = Py.Constructor [ "self" + , "command_queue=None" + , "interactive=False" + , "platform_pref=preferred_platform" + , "device_pref=preferred_device" + , "default_group_size=None" + , "default_num_groups=None" + , "default_tile_size=None" + , "sizes={}"] + [Escape $ openClInit types assign sizes] + options = [ Option { optionLongName = "platform" + , optionShortName = Just 'p' + , optionArgument = RequiredArgument + , optionAction = + [ Assign (Var "preferred_platform") $ Var "optarg" ] + } + , Option { optionLongName = "device" + , optionShortName = Just 'd' + , optionArgument = RequiredArgument + , optionAction = + [ Assign (Var "preferred_device") $ Var "optarg" ] + }] + + Right <$> Py.compileProg module_name constructor imports defines operations () + [Exp $ Py.simpleCall "self.queue.finish" []] options prog' + where operations :: Py.Operations Imp.OpenCL () + operations = Py.Operations + { Py.opsCompiler = callKernel + , Py.opsWriteScalar = writeOpenCLScalar + , Py.opsReadScalar = readOpenCLScalar + , Py.opsAllocate = allocateOpenCLBuffer + , Py.opsCopy = copyOpenCLMemory + , Py.opsStaticArray = staticOpenCLArray + , Py.opsEntryOutput = packArrayOutput + , Py.opsEntryInput = unpackArrayInput + } + +-- We have many casts to 'long', because PyOpenCL may get confused at +-- the 32-bit numbers that ImpCode uses for offsets and the like. +asLong :: PyExp -> PyExp +asLong x = Py.simpleCall "np.long" [x] + +callKernel :: Py.OpCompiler Imp.OpenCL () +callKernel (Imp.GetSize v key) = + Py.stm $ Assign (Var (Py.compileName v)) $ + Index (Var "self.sizes") (IdxExp $ String $ pretty key) +callKernel (Imp.CmpSizeLe v key x) = do + x' <- Py.compileExp x + Py.stm $ Assign (Var (Py.compileName v)) $ + BinOp "<=" (Index (Var "self.sizes") (IdxExp $ String $ pretty key)) x' +callKernel (Imp.GetSizeMax v size_class) = + Py.stm $ Assign (Var (Py.compileName v)) $ + Var $ "self.max_" ++ pretty size_class +callKernel (Imp.HostCode c) = + Py.compileCode c + +callKernel (Imp.LaunchKernel name args kernel_size workgroup_size) = do + kernel_size' <- mapM Py.compileExp kernel_size + let total_elements = foldl mult_exp (Integer 1) kernel_size' + let cond = BinOp "!=" total_elements (Integer 0) + workgroup_size' <- Tuple <$> mapM (fmap asLong . Py.compileExp) workgroup_size + body <- Py.collect $ launchKernel name kernel_size' workgroup_size' args + Py.stm $ If cond body [] + where mult_exp = BinOp "*" + +launchKernel :: String -> [PyExp] -> PyExp -> [Imp.KernelArg] -> Py.CompilerM op s () +launchKernel kernel_name kernel_dims workgroup_dims args = do + let kernel_dims' = Tuple $ map asLong kernel_dims + let kernel_name' = "self." ++ kernel_name ++ "_var" + args' <- mapM processKernelArg args + Py.stm $ Exp $ Py.simpleCall (kernel_name' ++ ".set_args") args' + Py.stm $ Exp $ Py.simpleCall "cl.enqueue_nd_range_kernel" + [Var "self.queue", Var kernel_name', kernel_dims', workgroup_dims] + finishIfSynchronous + where processKernelArg :: Imp.KernelArg -> Py.CompilerM op s PyExp + processKernelArg (Imp.ValueKArg e bt) = do + e' <- Py.compileExp e + return $ Py.simpleCall (Py.compilePrimToNp bt) [e'] + processKernelArg (Imp.MemKArg v) = return $ Var $ Py.compileName v + processKernelArg (Imp.SharedMemoryKArg (Imp.Count num_bytes)) = do + num_bytes' <- Py.compileExp num_bytes + return $ Py.simpleCall "cl.LocalMemory" [asLong num_bytes'] + +writeOpenCLScalar :: Py.WriteScalar Imp.OpenCL () +writeOpenCLScalar mem i bt "device" val = do + let mem' = Var $ Py.compileName mem + let nparr = Call (Var "np.array") + [Arg val, ArgKeyword "dtype" $ Var $ Py.compilePrimType bt] + Py.stm $ Exp $ Call (Var "cl.enqueue_copy") + [Arg $ Var "self.queue", Arg mem', Arg nparr, + ArgKeyword "device_offset" $ asLong i, + ArgKeyword "is_blocking" $ Var "synchronous"] + +writeOpenCLScalar _ _ _ space _ = + fail $ "Cannot write to '" ++ space ++ "' memory space." + +readOpenCLScalar :: Py.ReadScalar Imp.OpenCL () +readOpenCLScalar mem i bt "device" = do + val <- newVName "read_res" + let val' = Var $ pretty val + let mem' = Var $ Py.compileName mem + let nparr = Call (Var "np.empty") + [Arg $ Integer 1, + ArgKeyword "dtype" (Var $ Py.compilePrimType bt)] + Py.stm $ Assign val' nparr + Py.stm $ Exp $ Call (Var "cl.enqueue_copy") + [Arg $ Var "self.queue", Arg val', Arg mem', + ArgKeyword "device_offset" $ asLong i, + ArgKeyword "is_blocking" $ Bool True] + return $ Index val' $ IdxExp $ Integer 0 + +readOpenCLScalar _ _ _ space = + fail $ "Cannot read from '" ++ space ++ "' memory space." + +allocateOpenCLBuffer :: Py.Allocate Imp.OpenCL () +allocateOpenCLBuffer mem size "device" = + Py.stm $ Assign (Var $ Py.compileName mem) $ + Py.simpleCall "opencl_alloc" [Var "self", size, String $ pretty mem] + +allocateOpenCLBuffer _ _ space = + fail $ "Cannot allocate in '" ++ space ++ "' space" + +copyOpenCLMemory :: Py.Copy Imp.OpenCL () +copyOpenCLMemory destmem destidx Imp.DefaultSpace srcmem srcidx (Imp.Space "device") nbytes bt = do + let srcmem' = Var $ Py.compileName srcmem + let destmem' = Var $ Py.compileName destmem + let divide = BinOp "//" nbytes (Integer $ Imp.primByteSize bt) + let end = BinOp "+" destidx divide + let dest = Index destmem' (IdxRange destidx end) + Py.stm $ ifNotZeroSize nbytes $ + Exp $ Call (Var "cl.enqueue_copy") + [Arg $ Var "self.queue", Arg dest, Arg srcmem', + ArgKeyword "device_offset" $ asLong srcidx, + ArgKeyword "is_blocking" $ Var "synchronous"] + +copyOpenCLMemory destmem destidx (Imp.Space "device") srcmem srcidx Imp.DefaultSpace nbytes bt = do + let destmem' = Var $ Py.compileName destmem + let srcmem' = Var $ Py.compileName srcmem + let divide = BinOp "//" nbytes (Integer $ Imp.primByteSize bt) + let end = BinOp "+" srcidx divide + let src = Index srcmem' (IdxRange srcidx end) + Py.stm $ ifNotZeroSize nbytes $ + Exp $ Call (Var "cl.enqueue_copy") + [Arg $ Var "self.queue", Arg destmem', Arg src, + ArgKeyword "device_offset" $ asLong destidx, + ArgKeyword "is_blocking" $ Var "synchronous"] + +copyOpenCLMemory destmem destidx (Imp.Space "device") srcmem srcidx (Imp.Space "device") nbytes _ = do + let destmem' = Var $ Py.compileName destmem + let srcmem' = Var $ Py.compileName srcmem + Py.stm $ ifNotZeroSize nbytes $ + Exp $ Call (Var "cl.enqueue_copy") + [Arg $ Var "self.queue", Arg destmem', Arg srcmem', + ArgKeyword "dest_offset" $ asLong destidx, + ArgKeyword "src_offset" $ asLong srcidx, + ArgKeyword "byte_count" $ asLong nbytes] + finishIfSynchronous + +copyOpenCLMemory destmem destidx Imp.DefaultSpace srcmem srcidx Imp.DefaultSpace nbytes _ = + Py.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes + +copyOpenCLMemory _ _ destspace _ _ srcspace _ _= + error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace + +staticOpenCLArray :: Py.StaticArray Imp.OpenCL () +staticOpenCLArray name "device" t vs = do + mapM_ Py.atInit <=< Py.collect $ do + -- Create host-side Numpy array with intended values. + Py.stm $ Assign (Var name') $ + Call (Var "np.array") + [Arg $ List $ map Py.compilePrimValue vs, + ArgKeyword "dtype" $ Var $ Py.compilePrimToNp t] + + -- Create memory block on the device. + static_mem <- newVName "static_mem" + let size = Integer $ fromIntegral (length vs) * Imp.primByteSize t + allocateOpenCLBuffer static_mem size "device" + + -- Copy Numpy array to the device memory block. + Py.stm $ ifNotZeroSize size $ + Exp $ Call (Var "cl.enqueue_copy") + [Arg $ Var "self.queue", + Arg $ Var $ Py.compileName static_mem, + Arg $ Call (Var "normaliseArray") [Arg (Var name')], + ArgKeyword "is_blocking" $ Var "synchronous"] + + -- Store the memory block for later reference. + Py.stm $ Assign (Field (Var "self") name') $ + Var $ Py.compileName static_mem + + Py.stm $ Assign (Var name') (Field (Var "self") name') + where name' = Py.compileName name +staticOpenCLArray _ space _ _ = + fail $ "PyOpenCL backend cannot create static array in memory space '" ++ space ++ "'" + +packArrayOutput :: Py.EntryOutput Imp.OpenCL () +packArrayOutput mem "device" bt ept dims = + return $ Call (Var "cl.array.Array") + [Arg $ Var "self.queue", + Arg $ Tuple $ map Py.compileDim dims, + Arg $ Var $ Py.compilePrimTypeExt bt ept, + ArgKeyword "data" $ Var $ Py.compileName mem] +packArrayOutput _ sid _ _ _ = + fail $ "Cannot return array from " ++ sid ++ " space." + +unpackArrayInput :: Py.EntryInput Imp.OpenCL () +unpackArrayInput mem memsize "device" t s dims e = do + let type_is_ok = + BinOp "and" + (BinOp "in" (Py.simpleCall "type" [e]) (List [Var "np.ndarray", Var "cl.array.Array"])) + (BinOp "==" (Field e "dtype") (Var (Py.compilePrimToExtNp t s))) + Py.stm $ Assert type_is_ok $ String "Parameter has unexpected type" + + zipWithM_ (Py.unpackDim e) dims [0..] + + case memsize of + Imp.VarSize sizevar -> + Py.stm $ Assign (Var $ Py.compileName sizevar) $ + Py.simpleCall "np.int64" [Field e "nbytes"] + Imp.ConstSize _ -> + return () + + let memsize' = Py.compileDim memsize + pyOpenCLArrayCase = + [Assign mem_dest $ Field e "data"] + numpyArrayCase <- Py.collect $ do + allocateOpenCLBuffer mem memsize' "device" + Py.stm $ ifNotZeroSize memsize' $ + Exp $ Call (Var "cl.enqueue_copy") + [Arg $ Var "self.queue", + Arg $ Var $ Py.compileName mem, + Arg $ Call (Var "normaliseArray") [Arg e], + ArgKeyword "is_blocking" $ Var "synchronous"] + + Py.stm $ If (BinOp "==" (Py.simpleCall "type" [e]) (Var "cl.array.Array")) + pyOpenCLArrayCase + numpyArrayCase + where mem_dest = Var $ Py.compileName mem +unpackArrayInput _ _ sid _ _ _ _ = + fail $ "Cannot accept array from " ++ sid ++ " space." + +ifNotZeroSize :: PyExp -> PyStmt -> PyStmt +ifNotZeroSize e s = + If (BinOp "!=" e (Integer 0)) [s] [] + +finishIfSynchronous :: Py.CompilerM op s () +finishIfSynchronous = + Py.stm $ If (Var "synchronous") [Exp $ Py.simpleCall "self.queue.finish" []] [] diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} +module Futhark.CodeGen.Backends.PyOpenCL.Boilerplate + ( openClInit + , openClPrelude + ) where + +import Data.FileEmbed +import qualified Data.Map as M +import qualified Data.Text as T +import NeatInterpolation (text) + +import Futhark.CodeGen.ImpCode.OpenCL (PrimType(..), SizeClass(..)) +import Futhark.CodeGen.OpenCL.Kernels +import Futhark.CodeGen.Backends.GenericPython.AST +import Futhark.Util.Pretty (pretty, prettyText) + +-- | @rts/python/opencl.py@ embedded as a string. +openClPrelude :: String +openClPrelude = $(embedStringFile "rts/python/opencl.py") + +-- | Python code (as a string) that calls the +-- @initiatialize_opencl_object@ procedure. Should be put in the +-- class constructor. +openClInit :: [PrimType] -> String -> M.Map VName (SizeClass, Name) -> String +openClInit types assign sizes = T.unpack [text| +size_heuristics=$size_heuristics +program = initialise_opencl_object(self, + program_src=fut_opencl_src, + command_queue=command_queue, + interactive=interactive, + platform_pref=platform_pref, + device_pref=device_pref, + default_group_size=default_group_size, + default_num_groups=default_num_groups, + default_tile_size=default_tile_size, + size_heuristics=size_heuristics, + required_types=$types', + user_sizes=sizes, + all_sizes=$sizes') +$assign' +|] + where assign' = T.pack assign + size_heuristics = prettyText $ sizeHeuristicsToPython sizeHeuristicsTable + types' = prettyText $ map (show . pretty) types -- Looks enough like Python. + sizes' = prettyText $ sizeClassesToPython $ M.map fst sizes + +sizeClassesToPython :: M.Map VName SizeClass -> PyExp +sizeClassesToPython = Dict . map f . M.toList + where f (size_name, size_class) = + (String $ pretty size_name, + Dict [(String "class", String $ pretty size_class), + (String "value", None)]) + +sizeHeuristicsToPython :: [SizeHeuristic] -> PyExp +sizeHeuristicsToPython = List . map f + where f (SizeHeuristic platform_name device_type which what) = + Tuple [String platform_name, + clDeviceType device_type, + which', + what'] + + where clDeviceType DeviceGPU = Var "cl.device_type.GPU" + clDeviceType DeviceCPU = Var "cl.device_type.CPU" + + which' = case which of LockstepWidth -> String "lockstep_width" + NumGroups -> String "num_groups" + GroupSize -> String "group_size" + TileSize -> String "tile_size" + + what' = case what of + HeuristicConst x -> Integer $ toInteger x + HeuristicDeviceInfo s -> String s diff --git a/src/Futhark/CodeGen/Backends/SequentialC.hs b/src/Futhark/CodeGen/Backends/SequentialC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/SequentialC.hs @@ -0,0 +1,120 @@ +{-# LANGUAGE QuasiQuotes #-} +-- | C code generator. This module can convert a correct ImpCode +-- program to an equivalent C program. The C code is strictly +-- sequential, but can handle the full Futhark language. +module Futhark.CodeGen.Backends.SequentialC + ( compileProg + , GC.CParts(..) + , GC.asLibrary + , GC.asExecutable + ) where + +import Control.Monad + +import qualified Language.C.Quote.OpenCL as C + +import Futhark.Error +import Futhark.Representation.ExplicitMemory +import qualified Futhark.CodeGen.ImpCode.Sequential as Imp +import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGen +import qualified Futhark.CodeGen.Backends.GenericC as GC +import Futhark.MonadFreshNames + +compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError GC.CParts) +compileProg = + traverse (GC.compileProg operations generateContext "" [DefaultSpace] []) <=< + ImpGen.compileProg + where operations :: GC.Operations Imp.Sequential () + operations = GC.defaultOperations + { GC.opsCompiler = const $ return () + , GC.opsCopy = copySequentialMemory + } + + generateContext = do + cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { int debugging; };|]) + + GC.publicDef_ "context_config_new" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:cfg* $id:s();|], + [C.cedecl|struct $id:cfg* $id:s() { + struct $id:cfg *cfg = malloc(sizeof(struct $id:cfg)); + if (cfg == NULL) { + return NULL; + } + cfg->debugging = 0; + return cfg; + }|]) + + GC.publicDef_ "context_config_free" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { + free(cfg); + }|]) + + GC.publicDef_ "context_config_set_debugging" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) { + cfg->debugging = detail; + }|]) + + GC.publicDef_ "context_config_set_logging" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) { + /* Does nothing for this backend. */ + cfg = cfg; detail=detail; + }|]) + + (fields, init_fields) <- GC.contextContents + + ctx <- GC.publicDef "context" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { + int detail_memory; + int debugging; + typename lock_t lock; + char *error; + $sdecls:fields + };|]) + + GC.publicDef_ "context_new" GC.InitDecl $ \s -> + ([C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|], + [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) { + struct $id:ctx* ctx = malloc(sizeof(struct $id:ctx)); + if (ctx == NULL) { + return NULL; + } + ctx->detail_memory = cfg->debugging; + ctx->debugging = cfg->debugging; + ctx->error = NULL; + create_lock(&ctx->lock); + $stms:init_fields + return ctx; + }|]) + + GC.publicDef_ "context_free" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:ctx* ctx);|], + [C.cedecl|void $id:s(struct $id:ctx* ctx) { + free_lock(&ctx->lock); + free(ctx); + }|]) + + GC.publicDef_ "context_sync" GC.InitDecl $ \s -> + ([C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { + ctx=ctx; + return 0; + }|]) + GC.publicDef_ "context_get_error" GC.InitDecl $ \s -> + ([C.cedecl|char* $id:s(struct $id:ctx* ctx);|], + [C.cedecl|char* $id:s(struct $id:ctx* ctx) { + char* error = ctx->error; + ctx->error = NULL; + return error; + }|]) + +copySequentialMemory :: GC.Copy Imp.Sequential () +copySequentialMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes = + GC.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes +copySequentialMemory _ _ destspace _ _ srcspace _ = + error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace diff --git a/src/Futhark/CodeGen/Backends/SequentialCSharp.hs b/src/Futhark/CodeGen/Backends/SequentialCSharp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/SequentialCSharp.hs @@ -0,0 +1,40 @@ +module Futhark.CodeGen.Backends.SequentialCSharp + ( compileProg + ) where + +import Control.Monad +import Futhark.Error +import Futhark.Representation.ExplicitMemory +import qualified Futhark.CodeGen.ImpCode.Sequential as Imp +import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGen +import qualified Futhark.CodeGen.Backends.GenericCSharp as CS +import Futhark.CodeGen.Backends.GenericCSharp.AST () +import Futhark.MonadFreshNames + +compileProg :: MonadFreshNames m => + Maybe String -> Prog ExplicitMemory -> m (Either InternalError String) +compileProg module_name = + ImpGen.compileProg >=> + traverse (CS.compileProg + module_name + CS.emptyConstructor + [] + [] + operations + () + empty + [] + [] + []) + where operations :: CS.Operations Imp.Sequential () + operations = CS.defaultOperations + { CS.opsCompiler = const $ return () + , CS.opsCopy = copySequentialMemory + } + empty = return () + +copySequentialMemory :: CS.Copy Imp.Sequential () +copySequentialMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes _bt = + CS.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes +copySequentialMemory _ _ destspace _ _ srcspace _ _ = + error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace diff --git a/src/Futhark/CodeGen/Backends/SequentialPython.hs b/src/Futhark/CodeGen/Backends/SequentialPython.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/SequentialPython.hs @@ -0,0 +1,41 @@ +module Futhark.CodeGen.Backends.SequentialPython + ( compileProg + ) where + +import Control.Monad + +import Futhark.Error +import Futhark.Representation.ExplicitMemory +import qualified Futhark.CodeGen.ImpCode.Sequential as Imp +import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGen +import qualified Futhark.CodeGen.Backends.GenericPython as GenericPython +import Futhark.CodeGen.Backends.GenericPython.Definitions +import Futhark.CodeGen.Backends.GenericPython.AST +import Futhark.MonadFreshNames + +compileProg :: MonadFreshNames m => + Maybe String -> Prog ExplicitMemory -> m (Either InternalError String) +compileProg module_name = + ImpGen.compileProg >=> + traverse (GenericPython.compileProg + module_name + GenericPython.emptyConstructor + imports + defines + operations () [] []) + where imports = [Import "sys" Nothing, + Import "numpy" $ Just "np", + Import "ctypes" $ Just "ct", + Import "time" Nothing] + defines = [Escape pyValues, Escape pyFunctions, Escape pyPanic] + operations :: GenericPython.Operations Imp.Sequential () + operations = GenericPython.defaultOperations + { GenericPython.opsCompiler = const $ return () + , GenericPython.opsCopy = copySequentialMemory + } + +copySequentialMemory :: GenericPython.Copy Imp.Sequential () +copySequentialMemory destmem destidx DefaultSpace srcmem srcidx DefaultSpace nbytes _bt = + GenericPython.copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes +copySequentialMemory _ _ destspace _ _ srcspace _ _ = + error $ "Cannot copy to " ++ show destspace ++ " from " ++ show srcspace diff --git a/src/Futhark/CodeGen/Backends/SimpleRepresentation.hs b/src/Futhark/CodeGen/Backends/SimpleRepresentation.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/SimpleRepresentation.hs @@ -0,0 +1,428 @@ +{-# LANGUAGE QuasiQuotes #-} +-- | Simple C runtime representation. +module Futhark.CodeGen.Backends.SimpleRepresentation + ( sameRepresentation + , tupleField + , tupleFieldExp + , funName + , defaultMemBlockType + , intTypeToCType + , floatTypeToCType + , primTypeToCType + , signedPrimTypeToCType + + -- * Primitive value operations + , cIntOps + , cFloat32Ops, cFloat32Funs + , cFloat64Ops, cFloat64Funs + , cFloatConvOps + + ) + where + +import qualified Language.C.Syntax as C +import qualified Language.C.Quote.C as C + +import Futhark.CodeGen.ImpCode +import Futhark.Util.Pretty (pretty) +import Futhark.Util (zEncodeString) + +-- | The C type corresponding to a signed integer type. +intTypeToCType :: IntType -> C.Type +intTypeToCType Int8 = [C.cty|typename int8_t|] +intTypeToCType Int16 = [C.cty|typename int16_t|] +intTypeToCType Int32 = [C.cty|typename int32_t|] +intTypeToCType Int64 = [C.cty|typename int64_t|] + +-- | The C type corresponding to an unsigned integer type. +uintTypeToCType :: IntType -> C.Type +uintTypeToCType Int8 = [C.cty|typename uint8_t|] +uintTypeToCType Int16 = [C.cty|typename uint16_t|] +uintTypeToCType Int32 = [C.cty|typename uint32_t|] +uintTypeToCType Int64 = [C.cty|typename uint64_t|] + +-- | The C type corresponding to a float type. +floatTypeToCType :: FloatType -> C.Type +floatTypeToCType Float32 = [C.cty|float|] +floatTypeToCType Float64 = [C.cty|double|] + +-- | The C type corresponding to a primitive type. Integers are +-- assumed to be unsigned. +primTypeToCType :: PrimType -> C.Type +primTypeToCType (IntType t) = intTypeToCType t +primTypeToCType (FloatType t) = floatTypeToCType t +primTypeToCType Bool = [C.cty|typename bool|] +primTypeToCType Cert = [C.cty|typename bool|] + +-- | The C type corresponding to a primitive type. Integers are +-- assumed to have the specified sign. +signedPrimTypeToCType :: Signedness -> PrimType -> C.Type +signedPrimTypeToCType TypeUnsigned (IntType t) = uintTypeToCType t +signedPrimTypeToCType TypeDirect (IntType t) = intTypeToCType t +signedPrimTypeToCType _ t = primTypeToCType t + +-- | True if both types map to the same runtime representation. This +-- is the case if they are identical modulo uniqueness. +sameRepresentation :: [Type] -> [Type] -> Bool +sameRepresentation ets1 ets2 + | length ets1 == length ets2 = + and $ zipWith sameRepresentation' ets1 ets2 + | otherwise = False + +sameRepresentation' :: Type -> Type -> Bool +sameRepresentation' (Scalar t1) (Scalar t2) = + t1 == t2 +sameRepresentation' (Mem _ space1) (Mem _ space2) = space1 == space2 +sameRepresentation' _ _ = False + +-- | @tupleField i@ is the name of field number @i@ in a tuple. +tupleField :: Int -> String +tupleField i = "v" ++ show i + +-- | @tupleFieldExp e i@ is the expression for accesing field @i@ of +-- tuple @e@. If @e@ is an lvalue, so will the resulting expression +-- be. +tupleFieldExp :: C.ToExp a => a -> Int -> C.Exp +tupleFieldExp e i = [C.cexp|$exp:e.$id:(tupleField i)|] + +-- | @funName f@ is the name of the C function corresponding to +-- the Futhark function @f@. +funName :: Name -> String +funName = ("futrts_"++) . zEncodeString . nameToString + +funName' :: String -> String +funName' = funName . nameFromString + +-- | The type of memory blocks in the default memory space. +defaultMemBlockType :: C.Type +defaultMemBlockType = [C.cty|char*|] + +cIntOps :: [C.Definition] +cIntOps = concatMap (`map` [minBound..maxBound]) ops + where ops = [mkAdd, mkSub, mkMul, + mkUDiv, mkUMod, + mkSDiv, mkSMod, + mkSQuot, mkSRem, + mkSMin, mkUMin, + mkSMax, mkUMax, + mkShl, mkLShr, mkAShr, + mkAnd, mkOr, mkXor, + mkUlt, mkUle, mkSlt, mkSle, + mkPow, + mkIToB, mkBToI + ] ++ + map mkSExt [minBound..maxBound] ++ + map mkZExt [minBound..maxBound] + + taggedI s Int8 = s ++ "8" + taggedI s Int16 = s ++ "16" + taggedI s Int32 = s ++ "32" + taggedI s Int64 = s ++ "64" + + mkAdd = simpleIntOp "add" [C.cexp|x + y|] + mkSub = simpleIntOp "sub" [C.cexp|x - y|] + mkMul = simpleIntOp "mul" [C.cexp|x * y|] + mkUDiv = simpleUintOp "udiv" [C.cexp|x / y|] + mkUMod = simpleUintOp "umod" [C.cexp|x % y|] + mkUMax = simpleUintOp "umax" [C.cexp|x < y ? y : x|] + mkUMin = simpleUintOp "umin" [C.cexp|x < y ? x : y|] + + mkSDiv t = + let ct = intTypeToCType t + in [C.cedecl|static inline $ty:ct $id:(taggedI "sdiv" t)($ty:ct x, $ty:ct y) { + $ty:ct q = x / y; + $ty:ct r = x % y; + return q - + (((r != 0) && ((r < 0) != (y < 0))) ? 1 : 0); + }|] + mkSMod t = + let ct = intTypeToCType t + in [C.cedecl|static inline $ty:ct $id:(taggedI "smod" t)($ty:ct x, $ty:ct y) { + $ty:ct r = x % y; + return r + + ((r == 0 || (x > 0 && y > 0) || (x < 0 && y < 0)) ? 0 : y); + }|] + + mkSQuot = simpleIntOp "squot" [C.cexp|x / y|] + mkSRem = simpleIntOp "srem" [C.cexp|x % y|] + mkSMax = simpleIntOp "smax" [C.cexp|x < y ? y : x|] + mkSMin = simpleIntOp "smin" [C.cexp|x < y ? x : y|] + mkShl = simpleUintOp "shl" [C.cexp|x << y|] + mkLShr = simpleUintOp "lshr" [C.cexp|x >> y|] + mkAShr = simpleIntOp "ashr" [C.cexp|x >> y|] + mkAnd = simpleUintOp "and" [C.cexp|x & y|] + mkOr = simpleUintOp "or" [C.cexp|x | y|] + mkXor = simpleUintOp "xor" [C.cexp|x ^ y|] + mkUlt = uintCmpOp "ult" [C.cexp|x < y|] + mkUle = uintCmpOp "ule" [C.cexp|x <= y|] + mkSlt = intCmpOp "slt" [C.cexp|x < y|] + mkSle = intCmpOp "sle" [C.cexp|x <= y|] + + mkPow t = + let ct = intTypeToCType t + in [C.cedecl|static inline $ty:ct $id:(taggedI "pow" t)($ty:ct x, $ty:ct y) { + $ty:ct res = 1, rem = y; + while (rem != 0) { + if (rem & 1) { + res *= x; + } + rem >>= 1; + x *= x; + } + return res; + }|] + + mkSExt from_t to_t = + [C.cedecl|static inline $ty:to_ct + $id:name($ty:from_ct x) { return x;} |] + where name = "sext_"++pretty from_t++"_"++pretty to_t + from_ct = intTypeToCType from_t + to_ct = intTypeToCType to_t + + mkZExt from_t to_t = + [C.cedecl|static inline $ty:to_ct + $id:name($ty:from_ct x) { return x;} |] + where name = "zext_"++pretty from_t++"_"++pretty to_t + from_ct = uintTypeToCType from_t + to_ct = uintTypeToCType to_t + + mkBToI to_t = + [C.cedecl|static inline $ty:to_ct + $id:name($ty:from_ct x) { return x; } |] + where name = "btoi_bool_"++pretty to_t + from_ct = primTypeToCType Bool + to_ct = intTypeToCType to_t + + mkIToB from_t = + [C.cedecl|static inline $ty:to_ct + $id:name($ty:from_ct x) { return x; } |] + where name = "itob_"++pretty from_t++"_bool" + to_ct = primTypeToCType Bool + from_ct = intTypeToCType from_t + + simpleUintOp s e t = + [C.cedecl|static inline $ty:ct $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where ct = uintTypeToCType t + simpleIntOp s e t = + [C.cedecl|static inline $ty:ct $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where ct = intTypeToCType t + intCmpOp s e t = + [C.cedecl|static inline char $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where ct = intTypeToCType t + uintCmpOp s e t = + [C.cedecl|static inline char $id:(taggedI s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where ct = uintTypeToCType t + +cFloat32Ops :: [C.Definition] +cFloat64Ops :: [C.Definition] +cFloatConvOps :: [C.Definition] +(cFloat32Ops, cFloat64Ops, cFloatConvOps) = + ( map ($Float32) mkOps + , map ($Float64) mkOps + , [ mkFPConvFF "fpconv" from to | + from <- [minBound..maxBound], + to <- [minBound..maxBound] ]) + where taggedF s Float32 = s ++ "32" + taggedF s Float64 = s ++ "64" + convOp s from to = s ++ "_" ++ pretty from ++ "_" ++ pretty to + + mkOps = [mkFDiv, mkFAdd, mkFSub, mkFMul, mkFMin, mkFMax, mkPow, mkCmpLt, mkCmpLe] ++ + map (mkFPConvIF "sitofp") [minBound..maxBound] ++ + map (mkFPConvUF "uitofp") [minBound..maxBound] ++ + map (flip $ mkFPConvFI "fptosi") [minBound..maxBound] ++ + map (flip $ mkFPConvFU "fptoui") [minBound..maxBound] + + mkFDiv = simpleFloatOp "fdiv" [C.cexp|x / y|] + mkFAdd = simpleFloatOp "fadd" [C.cexp|x + y|] + mkFSub = simpleFloatOp "fsub" [C.cexp|x - y|] + mkFMul = simpleFloatOp "fmul" [C.cexp|x * y|] + mkFMin = simpleFloatOp "fmin" [C.cexp|x < y ? x : y|] + mkFMax = simpleFloatOp "fmax" [C.cexp|x < y ? y : x|] + mkCmpLt = floatCmpOp "cmplt" [C.cexp|x < y|] + mkCmpLe = floatCmpOp "cmple" [C.cexp|x <= y|] + + mkPow Float32 = + [C.cedecl|static inline float fpow32(float x, float y) { return pow(x, y); }|] + mkPow Float64 = + [C.cedecl|static inline double fpow64(double x, double y) { return pow(x, y); }|] + + mkFPConv from_f to_f s from_t to_t = + [C.cedecl|static inline $ty:to_ct + $id:(convOp s from_t to_t)($ty:from_ct x) { return x;} |] + where from_ct = from_f from_t + to_ct = to_f to_t + + mkFPConvFF = mkFPConv floatTypeToCType floatTypeToCType + mkFPConvFI = mkFPConv floatTypeToCType intTypeToCType + mkFPConvIF = mkFPConv intTypeToCType floatTypeToCType + mkFPConvFU = mkFPConv floatTypeToCType uintTypeToCType + mkFPConvUF = mkFPConv uintTypeToCType floatTypeToCType + + simpleFloatOp s e t = + [C.cedecl|static inline $ty:ct $id:(taggedF s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where ct = floatTypeToCType t + floatCmpOp s e t = + [C.cedecl|static inline char $id:(taggedF s t)($ty:ct x, $ty:ct y) { return $exp:e; }|] + where ct = floatTypeToCType t + +cFloat32Funs :: [C.Definition] +cFloat32Funs = [C.cunit| + static inline float $id:(funName' "log32")(float x) { + return log(x); + } + + static inline float $id:(funName' "log2_32")(float x) { + return log2(x); + } + + static inline float $id:(funName' "log10_32")(float x) { + return log10(x); + } + + static inline float $id:(funName' "sqrt32")(float x) { + return sqrt(x); + } + + static inline float $id:(funName' "exp32")(float x) { + return exp(x); + } + + static inline float $id:(funName' "cos32")(float x) { + return cos(x); + } + + static inline float $id:(funName' "sin32")(float x) { + return sin(x); + } + + static inline float $id:(funName' "tan32")(float x) { + return tan(x); + } + + static inline float $id:(funName' "acos32")(float x) { + return acos(x); + } + + static inline float $id:(funName' "asin32")(float x) { + return asin(x); + } + + static inline float $id:(funName' "atan32")(float x) { + return atan(x); + } + + static inline float $id:(funName' "atan2_32")(float x, float y) { + return atan2(x,y); + } + + static inline float $id:(funName' "round32")(float x) { + return rint(x); + } + + static inline char $id:(funName' "isnan32")(float x) { + return isnan(x); + } + + static inline char $id:(funName' "isinf32")(float x) { + return isinf(x); + } + + static inline typename int32_t $id:(funName' "to_bits32")(float x) { + union { + float f; + typename int32_t t; + } p; + p.f = x; + return p.t; + } + + static inline float $id:(funName' "from_bits32")(typename int32_t x) { + union { + typename int32_t f; + float t; + } p; + p.f = x; + return p.t; + } +|] + +cFloat64Funs :: [C.Definition] +cFloat64Funs = [C.cunit| + static inline double $id:(funName' "log64")(double x) { + return log(x); + } + + static inline double $id:(funName' "log2_64")(double x) { + return log2(x); + } + + static inline double $id:(funName' "log10_64")(double x) { + return log10(x); + } + + static inline double $id:(funName' "sqrt64")(double x) { + return sqrt(x); + } + + static inline double $id:(funName' "exp64")(double x) { + return exp(x); + } + + static inline double $id:(funName' "cos64")(double x) { + return cos(x); + } + + static inline double $id:(funName' "sin64")(double x) { + return sin(x); + } + + static inline double $id:(funName' "tan64")(double x) { + return tan(x); + } + + static inline double $id:(funName' "acos64")(double x) { + return acos(x); + } + + static inline double $id:(funName' "asin64")(double x) { + return asin(x); + } + + static inline double $id:(funName' "atan64")(double x) { + return atan(x); + } + + static inline double $id:(funName' "atan2_64")(double x, double y) { + return atan2(x,y); + } + + static inline double $id:(funName' "round64")(double x) { + return rint(x); + } + + static inline char $id:(funName' "isnan64")(double x) { + return isnan(x); + } + + static inline char $id:(funName' "isinf64")(double x) { + return isinf(x); + } + + static inline typename int64_t $id:(funName' "to_bits64")(double x) { + union { + double f; + typename int64_t t; + } p; + p.f = x; + return p.t; + } + + static inline double $id:(funName' "from_bits64")(typename int64_t x) { + union { + typename int64_t f; + double t; + } p; + p.f = x; + return p.t; + } +|] diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpCode.hs @@ -0,0 +1,498 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +-- | Imperative intermediate language used as a stepping stone in code generation. +-- +-- This is a generic representation parametrised on an extensible +-- arbitrary operation. +-- +-- Originally inspired by the paper "Defunctionalizing Push Arrays" +-- (FHPC '14). +module Futhark.CodeGen.ImpCode + ( Functions (..) + , Function + , FunctionT (..) + , ValueDesc (..) + , Signedness (..) + , ExternalValue (..) + , Param (..) + , paramName + , Size (..) + , MemSize + , DimSize + , Type (..) + , Space (..) + , SpaceId + , Code (..) + , PrimValue (..) + , ExpLeaf (..) + , Exp + , Volatility (..) + , Arg (..) + , var + , index + , ErrorMsg(..) + , ErrorMsgPart(..) + + -- * Typed enumerations + , Count (..) + , Bytes + , Elements + , elements + , bytes + , withElemType + + -- * Converting from sizes + , sizeToExp + , dimSizeToExp + , memSizeToExp + + -- * Analysis + + -- * Re-exports from other modules. + , module Language.Futhark.Core + , module Futhark.Representation.Primitive + , module Futhark.Analysis.PrimExp + ) + where + +import Data.Monoid ((<>)) +import Data.List +import Data.Loc +import Data.Traversable +import qualified Data.Set as S +import qualified Data.Semigroup as Sem + +import Language.Futhark.Core +import Futhark.Representation.Primitive +import Futhark.Representation.AST.Syntax + (Space(..), SpaceId, ErrorMsg(..), ErrorMsgPart(..)) +import Futhark.Representation.AST.Attributes.Names +import Futhark.Representation.AST.Pretty () +import Futhark.Util.IntegralExp +import Futhark.Analysis.PrimExp +import Futhark.Util.Pretty hiding (space) + +data Size = ConstSize Int64 + | VarSize VName + deriving (Eq, Show) + +type MemSize = Size +type DimSize = Size + +data Type = Scalar PrimType | Mem MemSize Space + +data Param = MemParam VName Space + | ScalarParam VName PrimType + deriving (Show) + +paramName :: Param -> VName +paramName (MemParam name _) = name +paramName (ScalarParam name _) = name + +-- | A collection of imperative functions. +newtype Functions a = Functions [(Name, Function a)] + +instance Sem.Semigroup (Functions a) where + Functions x <> Functions y = Functions $ x ++ y + +instance Monoid (Functions a) where + mempty = Functions [] + mappend = (Sem.<>) + +data Signedness = TypeUnsigned + | TypeDirect + deriving (Eq, Show) + +-- | A description of an externally meaningful value. +data ValueDesc = ArrayValue VName MemSize Space PrimType Signedness [DimSize] + -- ^ An array with memory block, memory block size, + -- memory space, element type, signedness of element + -- type (if applicable), and shape. + | ScalarValue PrimType Signedness VName + -- ^ A scalar value with signedness if applicable. + deriving (Eq, Show) + +-- | ^ An externally visible value. This can be an opaque value +-- (covering several physical internal values), or a single value that +-- can be used externally. +data ExternalValue = OpaqueValue String [ValueDesc] + -- ^ The string is a human-readable description + -- with no other semantics. + | TransparentValue ValueDesc + deriving (Show) + +-- | A imperative function, containing the body as well as its +-- low-level inputs and outputs, as well as its high-level arguments +-- and results. The latter are only used if the function is an entry +-- point. +data FunctionT a = Function { functionEntry :: Bool + , functionOutput :: [Param] + , functionInput :: [Param] + , functionbBody :: Code a + , functionResult :: [ExternalValue] + , functionArgs :: [ExternalValue] + } + deriving (Show) + +-- | Type alias for namespace control. +type Function = FunctionT + +data Code a = Skip + | Code a :>>: Code a + | For VName IntType Exp (Code a) + | While Exp (Code a) + | DeclareMem VName Space + | DeclareScalar VName PrimType + | DeclareArray VName Space PrimType [PrimValue] + -- ^ Create a read-only array containing the given values. + | Allocate VName (Count Bytes) Space + -- ^ Memory space must match the corresponding + -- 'DeclareMem'. + | Free VName Space + -- ^ Indicate that some memory block will never again be + -- referenced via the indicated variable. However, it + -- may still be accessed through aliases. It is only + -- safe to actually deallocate the memory block if this + -- is the last reference. There is no guarantee that + -- all memory blocks will be freed with this statement. + -- Backends are free to ignore it entirely. + | Copy VName (Count Bytes) Space VName (Count Bytes) Space (Count Bytes) + -- ^ Destination, offset in destination, destination + -- space, source, offset in source, offset space, number + -- of bytes. + | Write VName (Count Bytes) PrimType Space Volatility Exp + | SetScalar VName Exp + | SetMem VName VName Space + -- ^ Must be in same space. + | Call [VName] Name [Arg] + | If Exp (Code a) (Code a) + | Assert Exp (ErrorMsg Exp) (SrcLoc, [SrcLoc]) + | Comment String (Code a) + -- ^ Has the same semantics as the contained code, but + -- the comment should show up in generated code for ease + -- of inspection. + | DebugPrint String PrimType Exp + -- ^ Print the given value (of the given type) to the + -- screen, somehow annotated with the given string as a + -- description. This has no semantic meaning, but is + -- used entirely for debugging. Code generators are + -- free to ignore this statement. + | Op a + deriving (Show) + +-- | The volatility of a memory access. +data Volatility = Volatile | Nonvolatile + deriving (Eq, Ord, Show) + +instance Sem.Semigroup (Code a) where + Skip <> y = y + x <> Skip = x + x <> y = x :>>: y + +instance Monoid (Code a) where + mempty = Skip + mappend = (Sem.<>) + +data ExpLeaf = ScalarVar VName + | SizeOf PrimType + | Index VName (Count Bytes) PrimType Space Volatility + deriving (Eq, Show) + +type Exp = PrimExp ExpLeaf + +-- | A function call argument. +data Arg = ExpArg Exp + | MemArg VName + deriving (Show) + +-- | A wrapper around 'Imp.Exp' that maintains a unit as a phantom +-- type. +newtype Count u = Count { innerExp :: Exp } + deriving (Eq, Show, Num, IntegralExp, FreeIn, Pretty) + +-- | Phantom type for a count of elements. +data Elements + +-- | Phantom type for a count of bytes. +data Bytes + +elements :: Exp -> Count Elements +elements = Count + +bytes :: Exp -> Count Bytes +bytes = Count + +-- | Convert a count of elements into a count of bytes, given the +-- per-element size. +withElemType :: Count Elements -> PrimType -> Count Bytes +withElemType (Count e) t = bytes $ e * LeafExp (SizeOf t) (IntType Int32) + +dimSizeToExp :: DimSize -> Count Elements +dimSizeToExp = elements . sizeToExp + +memSizeToExp :: MemSize -> Count Bytes +memSizeToExp = bytes . sizeToExp + +sizeToExp :: Size -> Exp +sizeToExp (VarSize v) = LeafExp (ScalarVar v) (IntType Int32) +sizeToExp (ConstSize x) = ValueExp $ IntValue $ Int32Value $ fromIntegral x + +var :: VName -> PrimType -> Exp +var = LeafExp . ScalarVar + +index :: VName -> Count Bytes -> PrimType -> Space -> Volatility -> Exp +index arr i t s vol = LeafExp (Index arr i t s vol) t + +-- Prettyprinting definitions. + +instance Pretty op => Pretty (Functions op) where + ppr (Functions funs) = stack $ intersperse mempty $ map ppFun funs + where ppFun (name, fun) = + text "Function " <> ppr name <> colon indent 2 (ppr fun) + +instance Pretty op => Pretty (FunctionT op) where + ppr (Function _ outs ins body results args) = + text "Inputs:" block ins + text "Outputs:" block outs + text "Arguments:" block args + text "Result:" block results + text "Body:" indent 2 (ppr body) + where block :: Pretty a => [a] -> Doc + block = indent 2 . stack . map ppr + +instance Pretty Param where + ppr (ScalarParam name ptype) = + ppr ptype <+> ppr name + ppr (MemParam name space) = + text "mem" <> space' <+> ppr name + where space' = case space of Space s -> text "@" <> text s + DefaultSpace -> mempty + +instance Pretty ValueDesc where + ppr (ScalarValue t ept name) = + ppr t <+> ppr name <> ept' + where ept' = case ept of TypeUnsigned -> text " (unsigned)" + TypeDirect -> mempty + ppr (ArrayValue mem memsize space et ept shape) = + foldr f (ppr et) shape <+> text "at" <+> ppr mem <> parens (ppr memsize) <> space' <+> ept' + where f e s = brackets $ s <> comma <> ppr e + ept' = case ept of TypeUnsigned -> text " (unsigned)" + TypeDirect -> mempty + space' = case space of Space s -> text "@" <> text s + DefaultSpace -> mempty + + +instance Pretty ExternalValue where + ppr (TransparentValue v) = ppr v + ppr (OpaqueValue desc vs) = + text "opaque" <+> text desc <+> + nestedBlock "{" "}" (stack $ map ppr vs) + +instance Pretty Size where + ppr (ConstSize x) = ppr x + ppr (VarSize v) = ppr v + +instance Pretty op => Pretty (Code op) where + ppr (Op op) = ppr op + ppr Skip = text "skip" + ppr (c1 :>>: c2) = ppr c1 ppr c2 + ppr (For i it limit body) = + text "for" <+> ppr i <> text ":" <> ppr it <+> langle <+> ppr limit <+> text "{" + indent 2 (ppr body) + text "}" + ppr (While cond body) = + text "while" <+> ppr cond <+> text "{" + indent 2 (ppr body) + text "}" + ppr (DeclareMem name space) = + text "var" <+> ppr name <> text ": mem" <> parens (ppr space) + ppr (DeclareScalar name t) = + text "var" <+> ppr name <> text ":" <+> ppr t + ppr (DeclareArray name space t vs) = + text "array" <+> ppr name <> text "@" <> ppr space <+> text ":" <+> ppr t <+> + equals <+> braces (commasep $ map ppr vs) + ppr (Allocate name e space) = + ppr name <+> text "<-" <+> text "malloc" <> parens (ppr e) <> ppr space + ppr (Free name space) = + text "free" <> parens (ppr name) <> ppr space + ppr (Write name i bt space vol val) = + ppr name <> langle <> vol' <> ppr bt <> ppr space <> rangle <> brackets (ppr i) <+> + text "<-" <+> ppr val + where vol' = case vol of Volatile -> text "volatile " + Nonvolatile -> mempty + ppr (SetScalar name val) = + ppr name <+> text "<-" <+> ppr val + ppr (SetMem dest from space) = + ppr dest <+> text "<-" <+> ppr from <+> text "@" <> ppr space + ppr (Assert e msg _) = + text "assert" <> parens (commasep [text (show msg), ppr e]) + ppr (Copy dest destoffset destspace src srcoffset srcspace size) = + text "memcpy" <> + parens (ppMemLoc dest destoffset <> ppr destspace <> comma + ppMemLoc src srcoffset <> ppr srcspace <> comma + ppr size) + where ppMemLoc base offset = + ppr base <+> text "+" <+> ppr offset + ppr (If cond tbranch fbranch) = + text "if" <+> ppr cond <+> text "then {" + indent 2 (ppr tbranch) + text "} else {" + indent 2 (ppr fbranch) + text "}" + ppr (Call dests fname args) = + commasep (map ppr dests) <+> text "<-" <+> + ppr fname <> parens (commasep $ map ppr args) + ppr (Comment s code) = + text "--" <+> text s ppr code + ppr (DebugPrint desc pt e) = + text "debug" <+> parens (commasep [text (show desc), ppr pt, ppr e]) + +instance Pretty Arg where + ppr (MemArg m) = ppr m + ppr (ExpArg e) = ppr e + +instance Pretty ExpLeaf where + ppr (ScalarVar v) = + ppr v + ppr (Index v is bt space vol) = + ppr v <> langle <> vol' <> ppr bt <> space' <> rangle <> brackets (ppr is) + where space' = case space of DefaultSpace -> mempty + Space s -> text "@" <> text s + vol' = case vol of Volatile -> text "volatile " + Nonvolatile -> mempty + + ppr (SizeOf t) = + text "sizeof" <> parens (ppr t) + +instance Functor Functions where + fmap = fmapDefault + +instance Foldable Functions where + foldMap = foldMapDefault + +instance Traversable Functions where + traverse f (Functions funs) = + Functions <$> traverse f' funs + where f' (name, fun) = (name,) <$> traverse f fun + +instance Functor FunctionT where + fmap = fmapDefault + +instance Foldable FunctionT where + foldMap = foldMapDefault + +instance Traversable FunctionT where + traverse f (Function entry outs ins body results args) = + Function entry outs ins <$> traverse f body <*> pure results <*> pure args + +instance Functor Code where + fmap = fmapDefault + +instance Foldable Code where + foldMap = foldMapDefault + +instance Traversable Code where + traverse f (x :>>: y) = + (:>>:) <$> traverse f x <*> traverse f y + traverse f (For i it bound code) = + For i it bound <$> traverse f code + traverse f (While cond code) = + While cond <$> traverse f code + traverse f (If cond x y) = + If cond <$> traverse f x <*> traverse f y + traverse f (Op kernel) = + Op <$> f kernel + traverse _ Skip = + pure Skip + traverse _ (DeclareMem name space) = + pure $ DeclareMem name space + traverse _ (DeclareScalar name bt) = + pure $ DeclareScalar name bt + traverse _ (DeclareArray name space t vs) = + pure $ DeclareArray name space t vs + traverse _ (Allocate name size s) = + pure $ Allocate name size s + traverse _ (Free name space) = + pure $ Free name space + traverse _ (Copy dest destoffset destspace src srcoffset srcspace size) = + pure $ Copy dest destoffset destspace src srcoffset srcspace size + traverse _ (Write name i bt val space vol) = + pure $ Write name i bt val space vol + traverse _ (SetScalar name val) = + pure $ SetScalar name val + traverse _ (SetMem dest from space) = + pure $ SetMem dest from space + traverse _ (Assert e msg loc) = + pure $ Assert e msg loc + traverse _ (Call dests fname args) = + pure $ Call dests fname args + traverse f (Comment s code) = + Comment s <$> traverse f code + traverse _ (DebugPrint s t e) = + pure $ DebugPrint s t e + +declaredIn :: Code a -> Names +declaredIn (DeclareMem name _) = S.singleton name +declaredIn (DeclareScalar name _) = S.singleton name +declaredIn (DeclareArray name _ _ _) = S.singleton name +declaredIn (If _ t f) = declaredIn t <> declaredIn f +declaredIn (x :>>: y) = declaredIn x <> declaredIn y +declaredIn (For i _ _ body) = S.singleton i <> declaredIn body +declaredIn (While _ body) = declaredIn body +declaredIn (Comment _ body) = declaredIn body +declaredIn _ = mempty + +instance FreeIn a => FreeIn (Code a) where + freeIn (x :>>: y) = + freeIn x <> freeIn y `S.difference` declaredIn x + freeIn Skip = + mempty + freeIn (For i _ bound body) = + i `S.delete` (freeIn bound <> freeIn body) + freeIn (While cond body) = + freeIn cond <> freeIn body + freeIn DeclareMem{} = + mempty + freeIn DeclareScalar{} = + mempty + freeIn DeclareArray{} = + mempty + freeIn (Allocate name size _) = + freeIn name <> freeIn size + freeIn (Free name _) = + freeIn name + freeIn (Copy dest x _ src y _ n) = + freeIn dest <> freeIn x <> freeIn src <> freeIn y <> freeIn n + freeIn (SetMem x y _) = + freeIn x <> freeIn y + freeIn (Write v i _ _ _ e) = + freeIn v <> freeIn i <> freeIn e + freeIn (SetScalar x y) = + freeIn x <> freeIn y + freeIn (Call dests _ args) = + freeIn dests <> freeIn args + freeIn (If cond t f) = + freeIn cond <> freeIn t <> freeIn f + freeIn (Assert e _ _) = + freeIn e + freeIn (Op op) = + freeIn op + freeIn (Comment _ code) = + freeIn code + freeIn (DebugPrint _ _ e) = + freeIn e + +instance FreeIn ExpLeaf where + freeIn (Index v e _ _ _) = freeIn v <> freeIn e + freeIn (ScalarVar v) = freeIn v + freeIn (SizeOf _) = mempty + +instance FreeIn Arg where + freeIn (MemArg m) = freeIn m + freeIn (ExpArg e) = freeIn e + +instance FreeIn Size where + freeIn (VarSize name) = S.singleton name + freeIn (ConstSize _) = mempty diff --git a/src/Futhark/CodeGen/ImpCode/Kernels.hs b/src/Futhark/CodeGen/ImpCode/Kernels.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpCode/Kernels.hs @@ -0,0 +1,297 @@ +{-# LANGUAGE FlexibleContexts #-} +-- | 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 (..) + , AtomicOp (..) + , CallKernel (..) + , MapKernel (..) + , Kernel (..) + , LocalMemoryUse + , KernelUse (..) + , module Futhark.CodeGen.ImpCode + , module Futhark.Representation.Kernels.Sizes + -- * Utility functions + , getKernels + ) + where + +import Control.Monad.Writer +import Data.List +import qualified Data.Set as S + +import Futhark.CodeGen.ImpCode hiding (Function, Code) +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.Representation.Kernels.Sizes +import Futhark.Representation.AST.Attributes.Names +import Futhark.Representation.AST.Pretty () +import Futhark.Util.Pretty + +type Program = Functions HostOp +type Function = Imp.Function HostOp +-- | Host-level code that can call kernels. +type Code = Imp.Code CallKernel +-- | Code inside a kernel. +type KernelCode = Imp.Code KernelOp + +-- | A run-time constant related to kernels. +newtype KernelConst = SizeConst VName + deriving (Eq, Ord, Show) + +-- | An expression whose variables are kernel constants. +type KernelConstExp = PrimExp KernelConst + +data HostOp = CallKernel CallKernel + | GetSize VName VName SizeClass + | CmpSizeLe VName VName SizeClass Imp.Exp + | GetSizeMax VName SizeClass + deriving (Show) + +data CallKernel = Map MapKernel + | AnyKernel Kernel + | MapTranspose PrimType VName Exp VName Exp Exp Exp Exp Exp Exp + deriving (Show) + +-- | A generic kernel containing arbitrary kernel code. +data MapKernel = MapKernel { mapKernelThreadNum :: VName + -- ^ Stm position - also serves as a unique + -- name for the kernel. + , mapKernelDesc :: String + -- ^ Used to name the kernel for readability. + , mapKernelBody :: Imp.Code KernelOp + , mapKernelUses :: [KernelUse] + , mapKernelNumGroups :: DimSize + , mapKernelGroupSize :: DimSize + , mapKernelSize :: Imp.Exp + -- ^ Do not actually execute threads past this. + } + deriving (Show) + +data Kernel = Kernel + { kernelBody :: Imp.Code KernelOp + , kernelLocalMemory :: [LocalMemoryUse] + -- ^ The local memory used by this kernel. + + , kernelUses :: [KernelUse] + -- ^ The host variables referenced by the kernel. + + , kernelNumGroups :: DimSize + , kernelGroupSize :: DimSize + , kernelName :: VName + -- ^ Unique name for the kernel. + , kernelDesc :: String + -- ^ A short descriptive name - should be + -- alphanumeric and without spaces. + } + deriving (Show) + +-- ^ In-kernel name and per-workgroup size in bytes. +type LocalMemoryUse = (VName, Either MemSize KernelConstExp) + +data KernelUse = ScalarUse VName PrimType + | MemoryUse VName Imp.DimSize + | ConstUse VName KernelConstExp + deriving (Eq, Show) + +getKernels :: Program -> [CallKernel] +getKernels = nubBy sameKernel . execWriter . traverse getFunKernels + where getFunKernels (CallKernel kernel) = + tell [kernel] + getFunKernels _ = + return () + sameKernel (MapTranspose bt1 _ _ _ _ _ _ _ _ _) (MapTranspose bt2 _ _ _ _ _ _ _ _ _) = + bt1 == bt2 + sameKernel _ _ = False + +instance Pretty KernelConst where + ppr (SizeConst key) = text "get_size" <> parens (ppr key) + +instance Pretty KernelUse where + ppr (ScalarUse name t) = + text "scalar_copy" <> parens (commasep [ppr name, ppr t]) + ppr (MemoryUse name size) = + text "mem_copy" <> parens (commasep [ppr name, ppr size]) + ppr (ConstUse name e) = + 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 name _ x) = + freeIn dest <> freeIn name <> freeIn x + freeIn (GetSizeMax dest _) = + freeIn dest + freeIn (GetSize dest _ _) = + freeIn dest + +instance Pretty CallKernel where + ppr (Map k) = ppr k + ppr (AnyKernel k) = ppr k + ppr (MapTranspose bt dest destoffset src srcoffset num_arrays size_x size_y in_size out_size) = + text "mapTranspose" <> + parens (ppr bt <> comma + ppMemLoc dest destoffset <> comma + ppMemLoc src srcoffset <> comma + ppr num_arrays <> comma <+> + ppr size_x <> comma <+> + ppr size_y <> comma <+> + ppr in_size <> comma <+> + ppr out_size) + where ppMemLoc base offset = + ppr base <+> text "+" <+> ppr offset + +instance FreeIn CallKernel where + freeIn (Map k) = freeIn k + freeIn (AnyKernel k) = freeIn k + freeIn (MapTranspose _ dest destoffset src srcoffset num_arrays size_x size_y in_size out_size) = + freeIn [dest, src] <> freeIn [destoffset, srcoffset] <> freeIn num_arrays <> + freeIn [size_x, size_y] <> freeIn [in_size, out_size] + +instance FreeIn Kernel where + freeIn kernel = freeIn (kernelBody kernel) <> + freeIn [kernelNumGroups kernel, kernelGroupSize kernel] + +instance Pretty MapKernel where + ppr kernel = + text "mapKernel" <+> brace + (text "uses" <+> brace (commasep $ map ppr $ mapKernelUses kernel) + text "body" <+> brace (ppr (mapKernelThreadNum kernel) <+> + text "<- get_thread_number()" + ppr (mapKernelBody kernel))) + +instance Pretty Kernel where + ppr kernel = + text "kernel" <+> brace + (text "groups" <+> brace (ppr $ kernelNumGroups kernel) + text "group_size" <+> brace (ppr $ kernelGroupSize kernel) + text "local_memory" <+> brace (commasep $ + map ppLocalMemory $ + kernelLocalMemory kernel) + text "uses" <+> brace (commasep $ map ppr $ kernelUses kernel) + text "body" <+> brace (ppr $ kernelBody kernel)) + where ppLocalMemory (name, Left size) = + ppr name <+> parens (ppr size <+> text "bytes") + ppLocalMemory (name, Right size) = + ppr name <+> parens (ppr size <+> text "bytes (const)") + +instance FreeIn MapKernel where + freeIn kernel = + mapKernelThreadNum kernel `S.delete` freeIn (mapKernelBody kernel) + +data KernelOp = GetGroupId VName Int + | GetLocalId VName Int + | GetLocalSize VName Int + | GetGlobalSize VName Int + | GetGlobalId VName Int + | GetLockstepWidth VName + | Atomic AtomicOp + | Barrier + | MemFence + deriving (Show) + +-- Atomic operations return the value stored before the update. +-- This value is stored in the first VName. +data AtomicOp = AtomicAdd VName VName (Count Bytes) Exp + | AtomicSMax VName VName (Count Bytes) Exp + | AtomicSMin VName VName (Count Bytes) Exp + | AtomicUMax VName VName (Count Bytes) Exp + | AtomicUMin VName VName (Count Bytes) Exp + | AtomicAnd VName VName (Count Bytes) Exp + | AtomicOr VName VName (Count Bytes) Exp + | AtomicXor VName VName (Count Bytes) Exp + | AtomicCmpXchg VName VName (Count Bytes) Exp Exp + | AtomicXchg VName VName (Count Bytes) Exp + deriving (Show) + +instance FreeIn AtomicOp where + freeIn (AtomicAdd _ 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 <+> text "<-" <+> + text "get_group_id" <> parens (ppr i) + ppr (GetLocalId dest i) = + ppr dest <+> text "<-" <+> + text "get_local_id" <> parens (ppr i) + ppr (GetLocalSize dest i) = + ppr dest <+> text "<-" <+> + text "get_local_size" <> parens (ppr i) + ppr (GetGlobalSize dest i) = + ppr dest <+> text "<-" <+> + text "get_global_size" <> parens (ppr i) + ppr (GetGlobalId dest i) = + ppr dest <+> text "<-" <+> + text "get_global_id" <> parens (ppr i) + ppr (GetLockstepWidth dest) = + ppr dest <+> text "<-" <+> + text "get_lockstep_width()" + ppr Barrier = + text "barrier()" + ppr MemFence = + text "mem_fence()" + ppr (Atomic (AtomicAdd old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_add" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicSMax old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_smax" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicSMin old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_smin" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicUMax old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_umax" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicUMin old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_umin" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicAnd old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_and" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicOr old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_or" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicXor old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_xor" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x]) + ppr (Atomic (AtomicCmpXchg old arr ind x y)) = + ppr old <+> text "<-" <+> text "atomic_cmp_xchg" <> + parens (commasep [ppr arr <> brackets (ppr ind), ppr x, ppr y]) + ppr (Atomic (AtomicXchg old arr ind x)) = + ppr old <+> text "<-" <+> text "atomic_xchg" <> + 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 = text " {" indent 2 body text "}" diff --git a/src/Futhark/CodeGen/ImpCode/OpenCL.hs b/src/Futhark/CodeGen/ImpCode/OpenCL.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpCode/OpenCL.hs @@ -0,0 +1,74 @@ +-- | Imperative code with an OpenCL component. +-- +-- Apart from ordinary imperative code, this also carries around an +-- OpenCL program as a string, as well as a list of kernels defined by +-- the OpenCL program. +-- +-- The imperative code has been augmented with a 'LaunchKernel' +-- operation that allows one to execute an OpenCL kernel. +module Futhark.CodeGen.ImpCode.OpenCL + ( Program (..) + , Function + , FunctionT (Function) + , Code + , KernelName + , KernelArg (..) + , OpenCL (..) + , transposeBlockDim + , module Futhark.CodeGen.ImpCode + , module Futhark.Representation.Kernels.Sizes + ) + where + +import qualified Data.Map as M + +import Futhark.CodeGen.ImpCode hiding (Function, Code) +import Futhark.Representation.Kernels.Sizes +import qualified Futhark.CodeGen.ImpCode as Imp + +import Futhark.Util.Pretty + +-- | An program calling OpenCL kernels. +data Program = Program { openClProgram :: String + , openClPrelude :: String + -- ^ Must be prepended to the program. + , openClKernelNames :: [KernelName] + , openClUsedTypes :: [PrimType] + -- ^ So we can detect whether the device is capable. + , openClSizes :: M.Map VName (SizeClass, Name) + -- ^ Runtime-configurable constants. + , hostFunctions :: Functions OpenCL + } + +-- | A function calling OpenCL kernels. +type Function = Imp.Function OpenCL + +-- | A piece of code calling OpenCL. +type Code = Imp.Code OpenCL + +-- | The name of a kernel. +type KernelName = String + +-- | An argument to be passed to a kernel. +data KernelArg = ValueKArg Exp PrimType + -- ^ Pass the value of this scalar expression as argument. + | MemKArg VName + -- ^ Pass this pointer as argument. + | SharedMemoryKArg (Count Bytes) + -- ^ Create this much local memory per workgroup. + deriving (Show) + +-- | Host-level OpenCL operation. +data OpenCL = LaunchKernel KernelName [KernelArg] [Exp] [Exp] + | HostCode Code + | GetSize VName VName + | CmpSizeLe VName VName Exp + | GetSizeMax VName SizeClass + deriving (Show) + +-- | The block size when transposing. +transposeBlockDim :: Num a => a +transposeBlockDim = 16 + +instance Pretty OpenCL where + ppr = text . show diff --git a/src/Futhark/CodeGen/ImpCode/Sequential.hs b/src/Futhark/CodeGen/ImpCode/Sequential.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpCode/Sequential.hs @@ -0,0 +1,34 @@ +-- | Sequential imperative code. +module Futhark.CodeGen.ImpCode.Sequential + ( Program + , Function + , FunctionT (Function) + , Code + , Sequential + , module Futhark.CodeGen.ImpCode + ) + where + +import Futhark.CodeGen.ImpCode hiding (Function, Code) +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.Representation.AST.Attributes.Names + +import Futhark.Util.Pretty + +-- | An imperative program. +type Program = Imp.Functions Sequential + +-- | An imperative function. +type Function = Imp.Function Sequential + +-- | A piece of imperative code. +type Code = Imp.Code Sequential + +-- | Phantom type for identifying sequential imperative code. +data Sequential + +instance Pretty Sequential where + ppr _ = empty + +instance FreeIn Sequential where + freeIn _ = mempty diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -0,0 +1,1304 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, LambdaCase, TypeSynonymInstances, FlexibleInstances, MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +module Futhark.CodeGen.ImpGen + ( -- * Entry Points + compileProg + + -- * Pluggable Compiler + , OpCompiler + , ExpCompiler + , CopyCompiler + , BodyCompiler + , Operations (..) + , defaultOperations + , Destination (..) + , ValueDestination (..) + , MemLocation (..) + , MemEntry (..) + , ScalarEntry (..) + + -- * Monadic Compiler Interface + , ImpM + , Env (envVtable, envDefaultSpace) + , subImpM + , subImpM_ + , emit + , collect + , comment + , VarEntry (..) + , ArrayEntry (..) + + -- * Lookups + , lookupVar + , lookupArray + , arrayLocation + , lookupMemory + + -- * Building Blocks + , compileSubExp + , compileSubExpOfType + , compileSubExpTo + , compilePrimExp + , compileAlloc + , subExpToDimSize + , declaringLParams + , declaringFParams + , declaringVarEntry + , declaringScope + , declaringScopes + , declaringPrimVar + , declaringPrimVars + , withPrimVar + , everythingVolatile + , compileBody + , compileLoopBody + , defCompileBody + , compileStms + , compileExp + , defCompileExp + , sliceArray + , offsetArray + , strideArray + , fullyIndexArray + , fullyIndexArray' + , varIndex + , Imp.dimSizeToExp + , dimSizeToSubExp + , destinationFromParam + , destinationFromParams + , destinationFromPattern + , funcallTargets + , copy + , copyDWIM + , copyDWIMDest + , copyElementWise + ) + where + +import Control.Monad.RWS hiding (mapM, forM) +import Control.Monad.State hiding (mapM, forM) +import Control.Monad.Writer hiding (mapM, forM) +import Control.Monad.Except hiding (mapM, forM) +import qualified Control.Monad.Fail as Fail +import Data.Either +import Data.Traversable +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import Data.List + +import qualified Futhark.CodeGen.ImpCode as Imp +import Futhark.CodeGen.ImpCode + (Count (..), + Bytes, Elements, + bytes, withElemType) +import Futhark.Representation.ExplicitMemory +import Futhark.Representation.SOACS (SOACS) +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Construct (fullSliceNum) +import Futhark.MonadFreshNames +import Futhark.Error +import Futhark.Util + +-- | How to compile an 'Op'. +type OpCompiler lore op = Destination -> Op lore -> ImpM lore op () + +-- | How to compile a 'Body'. +type BodyCompiler lore op = Destination -> Body lore -> ImpM lore op () + +-- | How to compile an 'Exp'. +type ExpCompiler lore op = Destination -> Exp lore -> ImpM lore op () + +type CopyCompiler lore op = PrimType + -> MemLocation + -> MemLocation + -> Count Elements -- ^ Number of row elements of the source. + -> ImpM lore op () + +data Operations lore op = Operations { opsExpCompiler :: ExpCompiler lore op + , opsOpCompiler :: OpCompiler lore op + , opsBodyCompiler :: BodyCompiler lore op + , opsCopyCompiler :: CopyCompiler lore op + } + +-- | An operations set for which the expression compiler always +-- returns 'CompileExp'. +defaultOperations :: (ExplicitMemorish lore, FreeIn op) => + OpCompiler lore op -> Operations lore op +defaultOperations opc = Operations { opsExpCompiler = defCompileExp + , opsOpCompiler = opc + , opsBodyCompiler = defCompileBody + , opsCopyCompiler = defaultCopy + } + +-- | When an array is declared, this is where it is stored. +data MemLocation = MemLocation { memLocationName :: VName + , memLocationShape :: [Imp.DimSize] + , memLocationIxFun :: IxFun.IxFun Imp.Exp + } + deriving (Eq, Show) + +data ArrayEntry = ArrayEntry { + entryArrayLocation :: MemLocation + , entryArrayElemType :: PrimType + } + +entryArrayShape :: ArrayEntry -> [Imp.DimSize] +entryArrayShape = memLocationShape . entryArrayLocation + +data MemEntry = MemEntry { + entryMemSize :: Imp.MemSize + , entryMemSpace :: Imp.Space + } + +newtype ScalarEntry = ScalarEntry { + entryScalarType :: PrimType + } + +-- | 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 + +-- | When compiling an expression, this is a description of where the +-- result should end up. The integer is a reference to the construct +-- that gave rise to this destination (for patterns, this will be the +-- tag of the first name in the pattern). This can be used to make +-- the generated code easier to relate to the original code. +data Destination = Destination { destinationTag :: Maybe Int + , valueDestinations :: [ValueDestination] } + deriving (Show) + +data ValueDestination = ScalarDestination VName + | ArrayElemDestination VName PrimType Imp.Space (Count Bytes) + | MemoryDestination VName + | ArrayDestination (Maybe MemLocation) + -- ^ The 'MemLocation' is 'Just' if a copy if + -- required. If it is 'Nothing', then a + -- copy/assignment of a memory block somewhere + -- takes care of this array. + deriving (Show) + +-- | If the given value destination if a 'ScalarDestination', return +-- the variable name. Otherwise, 'Nothing'. +fromScalarDestination :: ValueDestination -> Maybe VName +fromScalarDestination (ScalarDestination name) = Just name +fromScalarDestination _ = Nothing + +data Env lore op = Env { + envVtable :: M.Map VName (VarEntry lore) + , envExpCompiler :: ExpCompiler lore op + , envBodyCompiler :: BodyCompiler lore op + , envOpCompiler :: OpCompiler lore op + , envCopyCompiler :: CopyCompiler lore op + , envDefaultSpace :: Imp.Space + , envVolatility :: Imp.Volatility + } + +newEnv :: Operations lore op -> Imp.Space -> Env lore op +newEnv ops ds = Env { envVtable = M.empty + , envExpCompiler = opsExpCompiler ops + , envBodyCompiler = opsBodyCompiler ops + , envOpCompiler = opsOpCompiler ops + , envCopyCompiler = opsCopyCompiler ops + , envDefaultSpace = ds + , envVolatility = Imp.Nonvolatile + } + +newtype ImpM lore op a = ImpM (RWST (Env lore op) (Imp.Code op) VNameSource (Either InternalError) a) + deriving (Functor, Applicative, Monad, + MonadState VNameSource, + MonadReader (Env lore op), + MonadWriter (Imp.Code op), + MonadError InternalError) + +instance Fail.MonadFail (ImpM lore op) where + fail = error . ("ImpM.fail: "++) + +instance MonadFreshNames (ImpM lore op) where + getNameSource = get + putNameSource = put + + +instance HasScope SOACS (ImpM lore op) where + askScope = M.map (LetInfo . entryType) <$> asks envVtable + where entryType (MemVar _ memEntry) = + Mem (dimSizeToSubExp $ entryMemSize memEntry) (entryMemSpace memEntry) + entryType (ArrayVar _ arrayEntry) = + Array + (entryArrayElemType arrayEntry) + (Shape $ map dimSizeToSubExp $ entryArrayShape arrayEntry) + NoUniqueness + entryType (ScalarVar _ scalarEntry) = + Prim $ entryScalarType scalarEntry + +runImpM :: ImpM lore op a + -> Operations lore op -> Imp.Space -> VNameSource + -> Either InternalError (a, VNameSource, Imp.Code op) +runImpM (ImpM m) comp = runRWST m . newEnv comp + +subImpM_ :: Operations lore' op' -> ImpM lore' op' a + -> ImpM lore op (Imp.Code op') +subImpM_ ops m = snd <$> subImpM ops m + +subImpM :: Operations lore' op' -> ImpM lore' op' a + -> ImpM lore op (a, Imp.Code op') +subImpM ops (ImpM m) = do + env <- ask + src <- getNameSource + case runRWST m env { envExpCompiler = opsExpCompiler ops + , envBodyCompiler = opsBodyCompiler ops + , envCopyCompiler = opsCopyCompiler ops + , envOpCompiler = opsOpCompiler ops + , envVtable = M.map scrubExps $ envVtable env + } + src of + Left err -> throwError err + Right (x, src', code) -> do + putNameSource src' + return (x, code) + where scrubExps (ArrayVar _ entry) = ArrayVar Nothing entry + scrubExps (MemVar _ entry) = MemVar Nothing entry + scrubExps (ScalarVar _ entry) = ScalarVar Nothing entry + +-- | Execute a code generation action, returning the code that was +-- emitted. +collect :: ImpM lore op () -> ImpM lore op (Imp.Code op) +collect m = pass $ do + ((), code) <- listen m + return (code, const mempty) + +collect' :: ImpM lore op a -> ImpM lore op (a, Imp.Code op) +collect' m = pass $ do + (x, code) <- listen m + return ((x, code), const mempty) + +-- | Execute a code generation action, wrapping the generated code +-- within a 'Imp.Comment' with the given description. +comment :: String -> ImpM lore op () -> ImpM lore op () +comment desc m = do code <- collect m + emit $ Imp.Comment desc code + +-- | Emit some generated imperative code. +emit :: Imp.Code op -> ImpM lore op () +emit = tell + +compileProg :: (ExplicitMemorish lore, MonadFreshNames m) => + Operations lore op -> Imp.Space + -> Prog lore -> m (Either InternalError (Imp.Functions op)) +compileProg ops ds prog = + modifyNameSource $ \src -> + case mapAccumLM (compileFunDef ops ds) src (progFunctions prog) of + Left err -> (Left err, src) + Right (src', funs) -> (Right $ Imp.Functions funs, src') + +compileInParam :: ExplicitMemorish lore => + FParam lore -> ImpM lore op (Either Imp.Param ArrayDecl) +compileInParam fparam = case paramAttr fparam of + MemPrim bt -> + return $ Left $ Imp.ScalarParam name bt + MemMem _ space -> + return $ Left $ Imp.MemParam name space + MemArray bt shape _ (ArrayIn mem ixfun) -> do + shape' <- mapM subExpToDimSize $ shapeDims shape + return $ Right $ ArrayDecl name bt $ + MemLocation mem shape' $ fmap compilePrimExp ixfun + where name = paramName fparam + +data ArrayDecl = ArrayDecl VName PrimType MemLocation + +fparamSizes :: Typed attr => Param attr -> S.Set VName +fparamSizes fparam + | Mem (Var size) _ <- paramType fparam = S.singleton size + | otherwise = S.fromList $ subExpVars $ arrayDims $ paramType fparam + +compileInParams :: ExplicitMemorish lore => + [FParam lore] -> [EntryPointType] + -> ImpM lore 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 + (inparams, arraydecls) <- partitionEithers <$> mapM compileInParam (ctx_params++val_params) + let findArray x = find (isArrayDecl x) arraydecls + sizes = mconcat $ map fparamSizes $ ctx_params++val_params + + summaries = M.fromList $ mapMaybe memSummary params + where memSummary param + | MemMem (Constant (IntValue (Int64Value size))) space <- paramAttr param = + Just (paramName param, (Imp.ConstSize size, space)) + | MemMem (Var size) space <- paramAttr param = + Just (paramName param, (Imp.VarSize size, space)) + | otherwise = + Nothing + + findMemInfo :: VName -> Maybe (Imp.MemSize, Space) + findMemInfo = flip M.lookup summaries + + mkValueDesc fparam signedness = + case (findArray $ paramName fparam, paramType fparam) of + (Just (ArrayDecl _ bt (MemLocation mem shape _)), _) -> do + (memsize, memspace) <- findMemInfo mem + Just $ Imp.ArrayValue mem memsize memspace bt signedness shape + (_, Prim bt) + | paramName fparam `S.member` sizes -> + Nothing + | otherwise -> + Just $ Imp.ScalarValue bt signedness $ paramName fparam + _ -> + Nothing + + mkExts (TypeOpaque desc n:epts) fparams = + let (fparams',rest) = splitAt n fparams + in Imp.OpaqueValue desc + (mapMaybe (`mkValueDesc` Imp.TypeDirect) fparams') : + mkExts epts rest + mkExts (TypeUnsigned:epts) (fparam:fparams) = + maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeUnsigned) ++ + mkExts epts fparams + mkExts (TypeDirect:epts) (fparam:fparams) = + maybeToList (Imp.TransparentValue <$> mkValueDesc fparam Imp.TypeDirect) ++ + mkExts epts fparams + mkExts _ _ = [] + + return (inparams, arraydecls, mkExts orig_epts val_params) + where isArrayDecl x (ArrayDecl y _ _) = x == y + +compileOutParams :: ExplicitMemorish lore => + [RetType lore] -> [EntryPointType] + -> ImpM lore 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) + let ctx_dests' = map snd $ sortOn fst $ M.toList ctx_dests + return (extvs, outparams, Destination Nothing $ ctx_dests' <> dests) + where imp = lift . lift + + mkExts (TypeOpaque 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, + dests ++ more_dests) + mkExts (TypeUnsigned:epts) (rt:rts) = do + (ev,dest) <- mkParam rt Imp.TypeUnsigned + (more_values, more_dests) <- mkExts epts rts + return (Imp.TransparentValue ev : more_values, + dest : more_dests) + mkExts (TypeDirect:epts) (rt:rts) = do + (ev,dest) <- mkParam rt Imp.TypeDirect + (more_values, more_dests) <- mkExts epts rts + return (Imp.TransparentValue ev : more_values, + dest : more_dests) + mkExts _ _ = return ([], []) + + mkParam MemMem{} _ = + compilerBugS "Functions may not explicitly return memory blocks." + mkParam (MemPrim t) ept = do + out <- imp $ newVName "scalar_out" + tell ([Imp.ScalarParam out t], mempty) + return (Imp.ScalarValue t ept out, ScalarDestination out) + mkParam (MemArray t shape _ attr) ept = do + space <- asks envDefaultSpace + (memout, memsize) <- case attr of + ReturnsNewBlock _ x x_size _ixfun -> do + memout <- imp $ newVName "out_mem" + sizeout <- ensureMemSizeOut x_size + tell ([Imp.MemParam memout space], + M.singleton x $ MemoryDestination memout) + return (memout, sizeout) + ReturnsInBlock memout _ -> do + memsize <- imp $ entryMemSize <$> lookupMemory memout + return (memout, memsize) + resultshape <- mapM inspectExtSize $ shapeDims shape + return (Imp.ArrayValue memout memsize space t ept resultshape, + ArrayDestination Nothing) + + inspectExtSize (Ext x) = do + (memseen,arrseen) <- get + case M.lookup x arrseen of + Nothing -> do + out <- imp $ newVName "out_arrsize" + tell ([Imp.ScalarParam out int32], + M.singleton x $ ScalarDestination out) + put (memseen, M.insert x out arrseen) + return $ Imp.VarSize out + Just out -> + return $ Imp.VarSize out + inspectExtSize (Free se) = + imp $ subExpToDimSize se + + -- | Return the name of the out-parameter for the memory size + -- 'x', creating it if it does not already exist. + ensureMemSizeOut (Ext x) = do + (memseen, arrseen) <- get + case M.lookup x memseen of + Nothing -> do sizeout <- imp $ newVName "out_memsize" + tell ([Imp.ScalarParam sizeout int64], + M.singleton x $ ScalarDestination sizeout) + put (M.insert x sizeout memseen, arrseen) + return $ Imp.VarSize sizeout + Just sizeout -> return $ Imp.VarSize sizeout + ensureMemSizeOut (Free v) = imp $ subExpToDimSize v + +compileFunDef :: ExplicitMemorish lore => + Operations lore op -> Imp.Space + -> VNameSource + -> FunDef lore + -> Either InternalError (VNameSource, (Name, Imp.Function op)) +compileFunDef ops ds src (FunDef entry fname rettype params body) = do + ((outparams, inparams, results, args), src', body') <- + runImpM compile ops ds src + return (src', + (fname, + Imp.Function (isJust entry) outparams inparams body' results args)) + where params_entry = maybe (replicate (length params) TypeDirect) fst entry + ret_entry = maybe (replicate (length rettype) TypeDirect) snd entry + compile = do + (inparams, arraydecls, args) <- compileInParams params params_entry + (results, outparams, dests) <- compileOutParams rettype ret_entry + withFParams params $ + withArrays arraydecls $ + compileBody dests body + return (outparams, inparams, results, args) + +compileBody :: Destination -> Body lore -> ImpM lore op () +compileBody dest body = do + cb <- asks envBodyCompiler + cb dest body + +defCompileBody :: (ExplicitMemorish lore, FreeIn op) => Destination -> Body lore -> ImpM lore op () +defCompileBody (Destination _ dest) (Body _ bnds ses) = + compileStms (freeIn ses) (stmsToList bnds) $ zipWithM_ compileSubExpTo dest ses + +compileLoopBody :: (ExplicitMemorish lore, FreeIn op) => + [VName] -> Body lore -> ImpM lore op (Imp.Code op) +compileLoopBody mergenames (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 + -- would thus be clobbered. Therefore, we first copy to new + -- variables mirroring the merge parameters, and then copy this + -- buffer to the merge parameters. This is efficient, because the + -- operations are all scalar operations. + tmpnames <- mapM (newVName . (++"_tmp") . baseString) mergenames + collect $ compileStms (freeIn ses) (stmsToList bnds) $ do + copy_to_merge_params <- forM (zip3 mergenames tmpnames ses) $ \(d,tmp,se) -> + subExpType se >>= \case + Prim bt -> do + se' <- compileSubExp se + emit $ Imp.DeclareScalar tmp bt + emit $ Imp.SetScalar tmp se' + return $ emit $ Imp.SetScalar d $ Imp.var tmp bt + Mem _ space | Var v <- se -> do + emit $ Imp.DeclareMem tmp space + emit $ Imp.SetMem tmp v space + return $ emit $ Imp.SetMem d tmp space + _ -> return $ return () + sequence_ copy_to_merge_params + +compileStms :: (ExplicitMemorish lore, FreeIn op) => + Names -> [Stm lore] -> ImpM lore op () -> ImpM lore op () +compileStms 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 + -- Free. This is very conservative, but can cut down on lifetimes + -- in some cases. + void $ compileStms' mempty all_stms + where compileStms' allocs (Let pat _ e:bs) = + declaringVars (Just e) (patternElements pat) $ do + dest <- destinationFromPattern pat + + e_code <- collect $ compileExp dest e + (live_after, bs_code) <- collect' $ compileStms' (patternAllocs pat <> allocs) bs + let dies_here v = not (v `S.member` live_after) && + v `S.member` freeIn e_code + to_free = S.filter (dies_here . fst) allocs + + emit e_code + mapM_ (emit . uncurry Imp.Free) to_free + emit bs_code + + return $ freeIn e_code <> live_after + compileStms' _ [] = do + code <- collect m + emit code + return $ freeIn code <> alive_after_stms + + patternAllocs = S.fromList . mapMaybe isMemPatElem . patternElements + isMemPatElem pe = case patElemType pe of + Mem _ space -> Just (patElemName pe, space) + _ -> Nothing + +compileExp :: Destination -> Exp lore -> ImpM lore op () +compileExp targets e = do + ec <- asks envExpCompiler + ec targets e + +defCompileExp :: (ExplicitMemorish lore, FreeIn op) => + Destination -> Exp lore -> ImpM lore op () + +defCompileExp dest (If cond tbranch fbranch _) = do + cond' <- compileSubExp cond + tcode <- collect $ compileBody dest tbranch + fcode <- collect $ compileBody dest fbranch + emit $ Imp.If cond' tcode fcode + +defCompileExp dest (Apply fname args _ _) = do + targets <- funcallTargets dest + args' <- catMaybes <$> mapM compileArg args + emit $ Imp.Call targets fname args' + where compileArg (se, _) = do + t <- subExpType se + case (se, t) of + (_, Prim pt) -> return $ Just $ Imp.ExpArg $ compileSubExpOfType pt se + (Var v, Mem{}) -> return $ Just $ Imp.MemArg v + _ -> return Nothing + +defCompileExp targets (BasicOp op) = defCompileBasicOp targets op + +defCompileExp (Destination _ dest) (DoLoop ctx val form body) = + declaringFParams mergepat $ do + forM_ merge $ \(p, se) -> do + na <- subExpNotArray se + when na $ + copyDWIM (paramName p) [] se [] + (bindForm, emitForm) <- + case form of + ForLoop i it bound loopvars -> do + bound' <- compileSubExp bound + let setLoopParam (p,a) + | Prim _ <- paramType p = + copyDWIM (paramName p) [] (Var a) [varIndex i] + | otherwise = + return () + + let emitForm body' = do + set_loop_params <- collect $ mapM_ setLoopParam loopvars + emit $ Imp.For i it bound' $ set_loop_params<>body' + return (declaringLParams (map fst loopvars) . + declaringLoopVar i it, + emitForm) + WhileLoop cond -> + return (id, emit . Imp.While (Imp.var cond Bool)) + + bindForm $ do + body' <- compileLoopBody mergenames body + emitForm body' + zipWithM_ compileSubExpTo dest $ map (Var . paramName . fst) merge + where merge = ctx ++ val + mergepat = map fst merge + mergenames = map paramName mergepat + +defCompileExp dest (Op op) = do + opc <- asks envOpCompiler + opc dest op + +defCompileBasicOp :: Destination -> BasicOp lore -> ImpM lore op () + +defCompileBasicOp (Destination _ [target]) (SubExp se) = + compileSubExpTo target se + +defCompileBasicOp (Destination _ [target]) (Opaque se) = + compileSubExpTo target se + +defCompileBasicOp (Destination _ [target]) (UnOp op e) = do + e' <- compileSubExp e + writeExp target $ Imp.UnOpExp op e' + +defCompileBasicOp (Destination _ [target]) (ConvOp conv e) = do + e' <- compileSubExp e + writeExp target $ Imp.ConvOpExp conv e' + +defCompileBasicOp (Destination _ [target]) (BinOp bop x y) = do + x' <- compileSubExp x + y' <- compileSubExp y + writeExp target $ Imp.BinOpExp bop x' y' + +defCompileBasicOp (Destination _ [target]) (CmpOp bop x y) = do + x' <- compileSubExp x + y' <- compileSubExp y + writeExp target $ Imp.CmpOpExp bop x' y' + +defCompileBasicOp (Destination _ [_]) (Assert e msg loc) = do + e' <- compileSubExp e + msg' <- traverse compileSubExp msg + emit $ Imp.Assert e' msg' loc + +defCompileBasicOp (Destination _ [target]) (Index src slice) + | Just idxs <- sliceIndices slice = + copyDWIMDest target [] (Var src) $ map (compileSubExpOfType int32) idxs + +defCompileBasicOp _ Index{} = + return () + +defCompileBasicOp (Destination _ [ArrayDestination (Just memloc)]) (Update _ slice se) + | MemLocation mem shape ixfun <- memloc = do + bt <- elemType <$> subExpType se + target' <- + case sliceIndices slice of + Just is -> do + (_, space, elemOffset) <- + fullyIndexArray' + (MemLocation mem shape ixfun) + (map (compileSubExpOfType int32) is) + bt + return $ ArrayElemDestination mem bt space elemOffset + Nothing -> + let memdest = sliceArray (MemLocation mem shape ixfun) $ + map (fmap (compileSubExpOfType int32)) slice + in return $ ArrayDestination $ Just memdest + + copyDWIMDest target' [] se [] + +defCompileBasicOp (Destination _ [dest]) (Replicate (Shape ds) se) = do + is <- replicateM (length ds) (newVName "i") + ds' <- mapM compileSubExp ds + declaringLoopVars Int32 is $ do + copy_elem <- collect $ copyDWIMDest dest (map varIndex is) se [] + emit $ foldl (.) id (zipWith (`Imp.For` Int32) is ds') copy_elem + +defCompileBasicOp (Destination _ [_]) Scratch{} = + return () + +defCompileBasicOp (Destination _ [dest]) (Iota n e s et) = do + i <- newVName "i" + x <- newVName "x" + n' <- compileSubExp n + e' <- compileSubExp e + s' <- compileSubExp s + emit $ Imp.DeclareScalar x $ IntType et + let i' = ConvOpExp (SExt Int32 et) $ Imp.var i $ IntType Int32 + declaringLoopVar i Int32 $ withPrimVar x (IntType et) $ + emit =<< (Imp.For i Int32 n' <$> + collect (do emit $ Imp.SetScalar x $ e' + i' * s' + copyDWIMDest dest [varIndex i] (Var x) [])) + +defCompileBasicOp (Destination _ [target]) (Copy src) = + compileSubExpTo target $ Var src + +defCompileBasicOp (Destination _ [target]) (Manifest _ src) = + compileSubExpTo target $ Var src + +defCompileBasicOp + (Destination _ [ArrayDestination (Just (MemLocation destmem destshape destixfun))]) + (Concat i x ys _) = do + xtype <- lookupType x + offs_glb <- newVName "tmp_offs" + withPrimVar offs_glb int32 $ do + emit $ Imp.DeclareScalar offs_glb int32 + emit $ Imp.SetScalar offs_glb 0 + let perm = [i] ++ [0..i-1] ++ [i+1..length destshape-1] + invperm = rearrangeInverse perm + destloc = MemLocation destmem destshape + (IxFun.permute (IxFun.offsetIndex (IxFun.permute destixfun perm) $ + varIndex offs_glb) + invperm) + + forM_ (x:ys) $ \y -> do + yentry <- lookupArray y + let srcloc = entryArrayLocation yentry + rows = case drop i $ entryArrayShape yentry of + [] -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y + r:_ -> innerExp $ Imp.dimSizeToExp r + copy (elemType xtype) destloc srcloc $ arrayOuterSize yentry + emit $ Imp.SetScalar offs_glb $ Imp.var offs_glb int32 + rows + +defCompileBasicOp (Destination _ [dest]) (ArrayLit es _) + | ArrayDestination (Just dest_mem) <- dest, + Just vs@(v:_) <- mapM isLiteral es = do + dest_space <- entryMemSpace <$> lookupMemory (memLocationName dest_mem) + let t = primValueType v + static_array <- newVName "static_array" + emit $ Imp.DeclareArray static_array dest_space t vs + let static_src = MemLocation static_array [Imp.ConstSize $ fromIntegral $ length es] $ + IxFun.iota [fromIntegral $ length es] + num_bytes = Imp.ConstSize $ fromIntegral (length es) * primByteSize t + entry = MemVar Nothing $ MemEntry num_bytes dest_space + local (insertInVtable static_array entry) $ + copy t dest_mem static_src $ fromIntegral $ length es + | otherwise = + forM_ (zip [0..] es) $ \(i,e) -> + copyDWIMDest dest [constIndex i] e [] + + where isLiteral (Constant v) = Just v + isLiteral _ = Nothing + +defCompileBasicOp _ Rearrange{} = + return () + +defCompileBasicOp _ Rotate{} = + return () + +defCompileBasicOp _ Reshape{} = + return () + +defCompileBasicOp _ Repeat{} = + return () + +defCompileBasicOp (Destination _ dests) (Partition n flags value_arrs) + | (sizedests, arrdest) <- splitAt n dests, + Just sizenames <- mapM fromScalarDestination sizedests, + Just destlocs <- mapM arrDestLoc arrdest = do + i <- newVName "i" + declaringLoopVar i Int32 $ do + outer_dim <- compileSubExp =<< (arraySize 0 <$> lookupType flags) + -- We will use 'i' to index the flag array and the value array. + -- Note that they have the same outer size ('outer_dim'). + read_flags_i <- readFromArray flags [varIndex i] + + -- First, for each of the 'n' output arrays, we compute the final + -- size. This is done by iterating through the flag array, but + -- first we declare scalars to hold the size. We do this by + -- creating a mapping from equivalence classes to the name of the + -- scalar holding the size. + let sizes = M.fromList $ zip [0..n-1] sizenames + + -- We initialise ecah size to zero. + forM_ sizenames $ \sizename -> + emit $ Imp.SetScalar sizename 0 + + -- Now iterate across the flag array, storing each element in + -- 'eqclass', then comparing it to the known classes and increasing + -- the appropriate size variable. + eqclass <- newVName "eqclass" + emit $ Imp.DeclareScalar eqclass int32 + let mkSizeLoopBody code c sizevar = + Imp.If (Imp.CmpOpExp (CmpEq int32) (Imp.var eqclass int32) (fromIntegral c)) + (Imp.SetScalar sizevar $ Imp.var sizevar int32 + 1) + code + sizeLoopBody = M.foldlWithKey' mkSizeLoopBody Imp.Skip sizes + emit $ Imp.For i Int32 outer_dim $ + Imp.SetScalar eqclass read_flags_i <> + sizeLoopBody + + -- We can now compute the starting offsets of each of the + -- partitions, creating a map from equivalence class to its + -- corresponding offset. + offsets <- flip evalStateT 0 $ forM sizes $ \size -> do + cur_offset <- get + partition_offset <- lift $ newVName "partition_offset" + lift $ emit $ Imp.DeclareScalar partition_offset int32 + lift $ emit $ Imp.SetScalar partition_offset cur_offset + put $ Imp.var partition_offset int32 + Imp.var size int32 + return partition_offset + + -- We create the memory location we use when writing a result + -- element. This is basically the index function of 'destloc', but + -- with a dynamic offset, stored in 'partition_cur_offset'. + partition_cur_offset <- newVName "partition_cur_offset" + emit $ Imp.DeclareScalar partition_cur_offset int32 + + -- Finally, we iterate through the data array and flag array in + -- parallel, and put each element where it is supposed to go. Note + -- that after writing to a partition, we increase the corresponding + -- offset. + ets <- mapM (fmap elemType . lookupType) value_arrs + srclocs <- mapM arrayLocation value_arrs + copy_elements <- forM (zip3 destlocs ets srclocs) $ \(destloc,et,srcloc) -> + copyArrayDWIM et + destloc [varIndex partition_cur_offset] + srcloc [varIndex i] + let mkWriteLoopBody code c offsetvar = + Imp.If (Imp.CmpOpExp (CmpEq int32) (Imp.var eqclass int32) (fromIntegral c)) + (Imp.SetScalar partition_cur_offset + (Imp.var offsetvar int32) + <> + mconcat copy_elements + <> + Imp.SetScalar offsetvar + (Imp.var offsetvar int32 + 1)) + code + writeLoopBody = M.foldlWithKey' mkWriteLoopBody Imp.Skip offsets + emit $ Imp.For i Int32 outer_dim $ + Imp.SetScalar eqclass read_flags_i <> + writeLoopBody + return () + where arrDestLoc (ArrayDestination destloc) = destloc + arrDestLoc _ = Nothing + +defCompileBasicOp (Destination _ []) _ = return () -- No arms, no cake. + +defCompileBasicOp target e = + compilerBugS $ "ImpGen.defCompileBasicOp: Invalid target\n " ++ + show target ++ "\nfor expression\n " ++ pretty e + +writeExp :: ValueDestination -> Imp.Exp -> ImpM lore op () +writeExp (ScalarDestination target) e = + emit $ Imp.SetScalar target e +writeExp (ArrayElemDestination destmem bt space elemoffset) e = do + vol <- asks envVolatility + emit $ Imp.Write destmem elemoffset bt space vol e +writeExp target e = + compilerBugS $ "Cannot write " ++ pretty e ++ " to " ++ show target + +insertInVtable :: VName -> VarEntry lore -> Env lore op -> Env lore op +insertInVtable name entry env = + env { envVtable = M.insert name entry $ envVtable env } + +withArray :: ArrayDecl -> ImpM lore op a -> ImpM lore op a +withArray (ArrayDecl name bt location) m = do + let entry = ArrayVar Nothing ArrayEntry + { entryArrayLocation = location + , entryArrayElemType = bt + } + local (insertInVtable name entry) m + +withArrays :: [ArrayDecl] -> ImpM lore op a -> ImpM lore op a +withArrays = flip $ foldr withArray + +-- | Like 'declaringFParams', but does not create new declarations. +withFParams :: ExplicitMemorish lore => [FParam lore] -> ImpM lore op a -> ImpM lore op a +withFParams = flip $ foldr withFParam + where withFParam fparam m = do + entry <- memBoundToVarEntry Nothing $ noUniquenessReturns $ paramAttr fparam + local (insertInVtable (paramName fparam) entry) m + +declaringVars :: ExplicitMemorish lore => + Maybe (Exp lore) -> [PatElem lore] -> ImpM lore op a -> ImpM lore op a +declaringVars e = flip $ foldr declaringVar + where declaringVar = declaringScope e . scopeOfPatElem + +declaringFParams :: ExplicitMemorish lore => [FParam lore] -> ImpM lore op a -> ImpM lore op a +declaringFParams = declaringScope Nothing . scopeOfFParams + +declaringLParams :: ExplicitMemorish lore => [LParam lore] -> ImpM lore op a -> ImpM lore op a +declaringLParams = declaringScope Nothing . scopeOfLParams + +declaringVarEntry :: VName -> VarEntry lore -> ImpM lore op a -> ImpM lore op a +declaringVarEntry name entry m = do + case entry of + MemVar _ entry' -> + emit $ Imp.DeclareMem name $ entryMemSpace entry' + ScalarVar _ entry' -> + emit $ Imp.DeclareScalar name $ entryScalarType entry' + ArrayVar _ _ -> + return () + local (insertInVtable name entry) m + +declaringPrimVar :: VName -> PrimType -> ImpM lore op a -> ImpM lore op a +declaringPrimVar name bt = + declaringVarEntry name $ ScalarVar Nothing $ ScalarEntry bt + +declaringPrimVars :: [(VName,PrimType)] -> ImpM lore op a -> ImpM lore op a +declaringPrimVars = flip $ foldr (uncurry declaringPrimVar) + +memBoundToVarEntry :: Maybe (Exp lore) -> MemBound NoUniqueness + -> ImpM lore op (VarEntry lore) +memBoundToVarEntry e (MemPrim bt) = + return $ ScalarVar e ScalarEntry { entryScalarType = bt } +memBoundToVarEntry e (MemMem size space) = do + size' <- subExpToDimSize size + return $ MemVar e MemEntry { entryMemSize = size' + , entryMemSpace = space + } +memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem ixfun)) = do + shape' <- mapM subExpToDimSize $ shapeDims shape + let location = MemLocation mem shape' $ fmap compilePrimExp ixfun + return $ ArrayVar e ArrayEntry { entryArrayLocation = location + , entryArrayElemType = bt + } + +declaringName :: Maybe (Exp lore) -> VName -> NameInfo ExplicitMemory + -> ImpM lore op a -> ImpM lore op a +declaringName e name info m = do + entry <- memBoundToVarEntry e $ infoAttr info + declaringVarEntry name entry m + where infoAttr (LetInfo attr) = attr + infoAttr (FParamInfo attr) = noUniquenessReturns attr + infoAttr (LParamInfo attr) = attr + infoAttr (IndexInfo it) = MemPrim $ IntType it + +declaringScope :: Maybe (Exp lore) -> Scope ExplicitMemory -> ImpM lore op a -> ImpM lore op a +declaringScope e scope m = foldr (uncurry $ declaringName e) m $ M.toList scope + +declaringScopes :: [(Maybe (Exp lore), Scope ExplicitMemory)] -> ImpM lore op a -> ImpM lore op a +declaringScopes es_and_scopes m = foldr (uncurry declaringScope) m es_and_scopes + +withPrimVar :: VName -> PrimType -> ImpM lore op a -> ImpM lore op a +withPrimVar name bt = + local (insertInVtable name $ ScalarVar Nothing $ ScalarEntry bt) + +declaringLoopVars :: IntType -> [VName] -> ImpM lore op a -> ImpM lore op a +declaringLoopVars it = flip $ foldr (`declaringLoopVar` it) + +declaringLoopVar :: VName -> IntType -> ImpM lore op a -> ImpM lore op a +declaringLoopVar name it = + withPrimVar name $ IntType it + +everythingVolatile :: ImpM lore op a -> ImpM lore op a +everythingVolatile = local $ \env -> env { envVolatility = Imp.Volatile } + +-- | Remove the array targets. +funcallTargets :: Destination -> ImpM lore op [VName] +funcallTargets (Destination _ dests) = + concat <$> mapM funcallTarget dests + where funcallTarget (ScalarDestination name) = + return [name] + funcallTarget ArrayElemDestination{} = + compilerBugS "Cannot put scalar function return in-place yet." -- FIXME + funcallTarget (ArrayDestination _) = + return [] + funcallTarget (MemoryDestination name) = + return [name] + +subExpToDimSize :: SubExp -> ImpM lore op Imp.DimSize +subExpToDimSize (Var v) = + return $ Imp.VarSize v +subExpToDimSize (Constant (IntValue (Int64Value i))) = + return $ Imp.ConstSize $ fromIntegral i +subExpToDimSize (Constant (IntValue (Int32Value i))) = + return $ Imp.ConstSize $ fromIntegral i +subExpToDimSize Constant{} = + compilerBugS "Size subexp is not an int32 or int64 constant." + +compileSubExpTo :: ValueDestination -> SubExp -> ImpM lore op () +compileSubExpTo dest se = copyDWIMDest dest [] se [] + +compileSubExp :: SubExp -> ImpM lore op Imp.Exp +compileSubExp (Constant v) = + return $ Imp.ValueExp v +compileSubExp (Var v) = do + t <- lookupType v + case t of + Prim pt -> return $ Imp.var v pt + _ -> compilerBugS $ "compileSubExp: SubExp is not a primitive type: " ++ pretty v + +compileSubExpOfType :: PrimType -> SubExp -> Imp.Exp +compileSubExpOfType _ (Constant v) = Imp.ValueExp v +compileSubExpOfType t (Var v) = Imp.var v t + +compilePrimExp :: PrimExp VName -> Imp.Exp +compilePrimExp = fmap Imp.ScalarVar + +varIndex :: VName -> Imp.Exp +varIndex name = LeafExp (Imp.ScalarVar name) int32 + +constIndex :: Int -> Imp.Exp +constIndex = fromIntegral + +lookupVar :: VName -> ImpM lore op (VarEntry lore) +lookupVar name = do + res <- asks $ M.lookup name . envVtable + case res of + Just entry -> return entry + _ -> compilerBugS $ "Unknown variable: " ++ pretty name + +lookupArray :: VName -> ImpM lore op ArrayEntry +lookupArray name = do + res <- lookupVar name + case res of + ArrayVar _ entry -> return entry + _ -> compilerBugS $ "ImpGen.lookupArray: not an array: " ++ pretty name + +arrayLocation :: VName -> ImpM lore op MemLocation +arrayLocation name = entryArrayLocation <$> lookupArray name + +lookupMemory :: VName -> ImpM lore op MemEntry +lookupMemory name = do + res <- lookupVar name + case res of + MemVar _ entry -> return entry + _ -> compilerBugS $ "Unknown memory block: " ++ pretty name + +destinationFromParam :: Param (MemBound u) -> ImpM lore op ValueDestination +destinationFromParam param + | MemArray _ shape _ (ArrayIn mem ixfun) <- paramAttr param = do + let dims = shapeDims shape + memloc <- MemLocation mem <$> mapM subExpToDimSize dims <*> + pure (fmap compilePrimExp ixfun) + return $ ArrayDestination $ Just memloc + | otherwise = + return $ ScalarDestination $ paramName param + +destinationFromParams :: [Param (MemBound u)] -> ImpM lore op Destination +destinationFromParams ps = fmap (Destination $ baseTag . paramName <$> maybeHead ps) . mapM destinationFromParam $ ps + +destinationFromPattern :: ExplicitMemorish lore => Pattern lore -> ImpM lore op Destination +destinationFromPattern pat = fmap (Destination (baseTag <$> maybeHead (patternNames pat))) . mapM inspect $ + patternElements pat + where ctx_names = patternContextNames pat + inspect patElem = do + let name = patElemName patElem + entry <- lookupVar name + case entry of + ArrayVar _ (ArrayEntry (MemLocation mem shape ixfun) _) -> + return $ ArrayDestination $ + if mem `elem` ctx_names + then Nothing + else Just $ MemLocation mem shape ixfun + MemVar{} -> + return $ MemoryDestination name + + ScalarVar{} -> + return $ ScalarDestination name + +fullyIndexArray :: VName -> [Imp.Exp] + -> ImpM lore op (VName, Imp.Space, Count Bytes) +fullyIndexArray name indices = do + arr <- lookupArray name + fullyIndexArray' (entryArrayLocation arr) indices $ entryArrayElemType arr + +fullyIndexArray' :: MemLocation -> [Imp.Exp] -> PrimType + -> ImpM lore op (VName, Imp.Space, Count Bytes) +fullyIndexArray' (MemLocation mem _ ixfun) indices bt = do + space <- entryMemSpace <$> lookupMemory mem + return (mem, space, + bytes $ IxFun.index ixfun indices $ primByteSize bt) + +readFromArray :: VName -> [Imp.Exp] + -> ImpM lore op Imp.Exp +readFromArray name indices = do + arr <- lookupArray name + (mem, space, i) <- + fullyIndexArray' (entryArrayLocation arr) indices $ entryArrayElemType arr + vol <- asks envVolatility + return $ Imp.index mem i (entryArrayElemType arr) space vol + +sliceArray :: MemLocation + -> Slice Imp.Exp + -> MemLocation +sliceArray (MemLocation mem shape ixfun) slice = + MemLocation mem (update shape slice) $ IxFun.slice ixfun slice + where update (d:ds) (DimSlice{}:is) = d : update ds is + update (_:ds) (DimFix{}:is) = update ds is + update _ _ = [] + +offsetArray :: MemLocation + -> Imp.Exp + -> MemLocation +offsetArray (MemLocation mem shape ixfun) offset = + MemLocation mem shape $ IxFun.offsetIndex ixfun offset + +strideArray :: MemLocation + -> Imp.Exp + -> MemLocation +strideArray (MemLocation mem shape ixfun) stride = + MemLocation mem shape $ IxFun.strideIndex ixfun stride + +subExpNotArray :: SubExp -> ImpM lore op Bool +subExpNotArray se = subExpType se >>= \case + Array {} -> return False + _ -> return True + +arrayOuterSize :: ArrayEntry -> Count Elements +arrayOuterSize = arrayDimSize 0 + +arrayDimSize :: Int -> ArrayEntry -> Count Elements +arrayDimSize i = + product . map Imp.dimSizeToExp . take 1 . drop i . entryArrayShape + +-- More complicated read/write operations that use index functions. + +copy :: CopyCompiler lore op +copy bt dest src n = do + cc <- asks envCopyCompiler + cc bt dest src n + +-- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'. +defaultCopy :: CopyCompiler lore op +defaultCopy bt dest src n + | ixFunMatchesInnerShape + (Shape $ map dimSizeToExp destshape) destIxFun, + ixFunMatchesInnerShape + (Shape $ map dimSizeToExp srcshape) srcIxFun, + Just destoffset <- + IxFun.linearWithOffset destIxFun bt_size, + Just srcoffset <- + IxFun.linearWithOffset srcIxFun bt_size = do + srcspace <- entryMemSpace <$> lookupMemory srcmem + destspace <- entryMemSpace <$> lookupMemory destmem + emit $ Imp.Copy + destmem (bytes destoffset) destspace + srcmem (bytes srcoffset) srcspace $ + (n * row_size) `withElemType` bt + | otherwise = + copyElementWise bt dest src n + where bt_size = primByteSize bt + row_size = product $ map Imp.dimSizeToExp $ drop 1 srcshape + MemLocation destmem destshape destIxFun = dest + MemLocation srcmem srcshape srcIxFun = src + +copyElementWise :: CopyCompiler lore op +copyElementWise bt (MemLocation destmem _ destIxFun) (MemLocation srcmem srcshape srcIxFun) n = do + is <- replicateM (IxFun.rank destIxFun) (newVName "i") + declaringLoopVars Int32 is $ do + let ivars = map varIndex is + destidx = IxFun.index destIxFun ivars bt_size + srcidx = IxFun.index srcIxFun ivars bt_size + bounds = map innerExp $ n : drop 1 (map Imp.dimSizeToExp srcshape) + srcspace <- entryMemSpace <$> lookupMemory srcmem + destspace <- entryMemSpace <$> lookupMemory destmem + vol <- asks envVolatility + emit $ foldl (.) id (zipWith (`Imp.For` Int32) is bounds) $ + Imp.Write destmem (bytes destidx) bt destspace vol $ + Imp.index srcmem (bytes srcidx) bt srcspace vol + where bt_size = primByteSize bt + +-- | Copy from here to there; both destination and source may be +-- indexeded. +copyArrayDWIM :: PrimType + -> MemLocation -> [Imp.Exp] + -> MemLocation -> [Imp.Exp] + -> ImpM lore op (Imp.Code op) +copyArrayDWIM bt + destlocation@(MemLocation _ destshape dest_ixfun) destis + srclocation@(MemLocation _ srcshape src_ixfun) srcis + + | length srcis == length srcshape, length destis == length destshape = do + (targetmem, destspace, targetoffset) <- + fullyIndexArray' destlocation destis bt + (srcmem, srcspace, srcoffset) <- + fullyIndexArray' srclocation srcis bt + vol <- asks envVolatility + return $ Imp.Write targetmem targetoffset bt destspace vol $ + Imp.index srcmem srcoffset bt srcspace vol + + | otherwise = do + let destlocation' = + sliceArray destlocation $ + fullSliceNum (IxFun.shape dest_ixfun) $ map DimFix destis + srclocation' = + sliceArray srclocation $ + fullSliceNum (IxFun.shape src_ixfun) $ map DimFix srcis + if destlocation' == srclocation' + then return mempty -- Copy would be no-op. + else collect $ copy bt destlocation' srclocation' $ + product $ map Imp.dimSizeToExp $ + take 1 $ drop (length srcis) srcshape + +-- | Like 'copyDWIM', but the target is a 'ValueDestination' +-- instead of a variable name. +copyDWIMDest :: ValueDestination -> [Imp.Exp] -> SubExp -> [Imp.Exp] + -> ImpM lore op () + +copyDWIMDest _ _ (Constant v) (_:_) = + compilerBugS $ + unwords ["copyDWIMDest: constant source", pretty v, "cannot be indexed."] +copyDWIMDest dest dest_is (Constant v) [] = + case dest of + ScalarDestination name -> + emit $ Imp.SetScalar name $ Imp.ValueExp v + ArrayElemDestination dest_mem _ dest_space dest_i -> do + vol <- asks envVolatility + emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.ValueExp v + MemoryDestination{} -> + compilerBugS $ + unwords ["copyDWIMDest: constant source", pretty v, "cannot be written to memory destination."] + ArrayDestination (Just dest_loc) -> do + (dest_mem, dest_space, dest_i) <- + fullyIndexArray' dest_loc dest_is bt + vol <- asks envVolatility + emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.ValueExp v + ArrayDestination Nothing -> + compilerBugS "copyDWIMDest: ArrayDestination Nothing" + where bt = primValueType v + +copyDWIMDest dest dest_is (Var src) src_is = do + src_entry <- lookupVar src + case (dest, src_entry) of + (MemoryDestination mem, MemVar _ (MemEntry _ space)) -> + emit $ Imp.SetMem mem src space + + (MemoryDestination{}, _) -> + compilerBugS $ + unwords ["copyDWIMDest: cannot write", pretty src, "to memory destination."] + + (_, MemVar{}) -> + compilerBugS $ + unwords ["copyDWIMDest: source", pretty src, "is a memory block."] + + (_, ScalarVar _ (ScalarEntry _)) | not $ null src_is -> + compilerBugS $ + unwords ["copyDWIMDest: prim-typed source", pretty src, "with nonzero indices."] + + + (ScalarDestination name, _) | not $ null dest_is -> + compilerBugS $ + unwords ["copyDWIMDest: prim-typed target", pretty name, "with nonzero indices."] + + (ScalarDestination name, ScalarVar _ (ScalarEntry pt)) -> + emit $ Imp.SetScalar name $ Imp.var src pt + + (ScalarDestination name, ArrayVar _ arr) -> do + let bt = entryArrayElemType arr + (mem, space, i) <- + fullyIndexArray' (entryArrayLocation arr) src_is bt + vol <- asks envVolatility + emit $ Imp.SetScalar name $ Imp.index mem i bt space vol + + (ArrayElemDestination{}, _) | not $ null dest_is-> + compilerBugS $ + unwords ["copyDWIMDest: array elemenent destination given indices:", pretty dest_is] + + (ArrayElemDestination dest_mem _ dest_space dest_i, + ScalarVar _ (ScalarEntry bt)) -> do + vol <- asks envVolatility + emit $ Imp.Write dest_mem dest_i bt dest_space vol $ Imp.var src bt + + (ArrayElemDestination dest_mem _ dest_space dest_i, ArrayVar _ src_arr) + | length (entryArrayShape src_arr) == length src_is -> do + let bt = entryArrayElemType src_arr + (src_mem, src_space, src_i) <- + fullyIndexArray' (entryArrayLocation src_arr) src_is bt + vol <- asks envVolatility + emit $ Imp.Write dest_mem dest_i bt dest_space vol $ + Imp.index src_mem src_i bt src_space vol + + (ArrayElemDestination{}, ArrayVar{}) -> + compilerBugS $ + unwords ["copyDWIMDest: array element destination, but array source", + pretty src, + "with incomplete indexing."] + + (ArrayDestination (Just dest_loc), ArrayVar _ src_arr) -> do + let src_loc = entryArrayLocation src_arr + bt = entryArrayElemType src_arr + emit =<< copyArrayDWIM bt dest_loc dest_is src_loc src_is + + (ArrayDestination (Just dest_loc), ScalarVar _ (ScalarEntry bt)) -> do + (dest_mem, dest_space, dest_i) <- + fullyIndexArray' dest_loc dest_is bt + vol <- asks envVolatility + emit $ Imp.Write dest_mem dest_i bt dest_space vol (Imp.var src bt) + + (ArrayDestination Nothing, _) -> + return () -- Nothing to do; something else set some memory + -- somewhere. + +-- | Copy from here to there; both destination and source be +-- indexeded. If so, they better be arrays of enough dimensions. +-- This function will generally just Do What I Mean, and Do The Right +-- Thing. Both destination and source must be in scope. +copyDWIM :: VName -> [Imp.Exp] -> SubExp -> [Imp.Exp] + -> ImpM lore op () +copyDWIM dest dest_is src src_is = do + dest_entry <- lookupVar dest + let dest_target = + case dest_entry of + ScalarVar _ _ -> + ScalarDestination dest + + ArrayVar _ (ArrayEntry (MemLocation mem shape ixfun) _) -> + ArrayDestination $ Just $ MemLocation mem shape ixfun + + MemVar _ _ -> + MemoryDestination dest + copyDWIMDest dest_target dest_is src src_is + +-- | @compileAlloc dest size space@ allocates @n@ bytes of memory in @space@, +-- writing the result to @dest@, which must be a single +-- 'MemoryDestination', +compileAlloc :: Destination -> SubExp -> Space + -> ImpM lore op () +compileAlloc (Destination _ [MemoryDestination mem]) e space = do + e' <- compileSubExp e + emit $ Imp.Allocate mem (Imp.bytes e') space +compileAlloc dest _ _ = + compilerBugS $ "compileAlloc: Invalid destination: " ++ show dest + +dimSizeToSubExp :: Imp.Size -> SubExp +dimSizeToSubExp (Imp.ConstSize n) = constant n +dimSizeToSubExp (Imp.VarSize v) = Var v + +dimSizeToExp :: Imp.Size -> Imp.Exp +dimSizeToExp = compilePrimExp . primExpFromSubExp int32 . dimSizeToSubExp diff --git a/src/Futhark/CodeGen/ImpGen/Kernels.hs b/src/Futhark/CodeGen/ImpGen/Kernels.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Kernels.hs @@ -0,0 +1,1390 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ConstraintKinds #-} +module Futhark.CodeGen.ImpGen.Kernels + ( compileProg + ) + where + +import Control.Arrow ((&&&)) +import Control.Monad.Except +import Control.Monad.Reader +import Data.Maybe +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.List + +import Prelude hiding (quot) + +import Futhark.Error +import Futhark.MonadFreshNames +import Futhark.Transform.Rename +import Futhark.Representation.ExplicitMemory +import qualified Futhark.CodeGen.ImpCode.Kernels as Imp +import Futhark.CodeGen.ImpCode.Kernels (bytes) +import qualified Futhark.CodeGen.ImpGen as ImpGen +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.CodeGen.SetDefaultSpace +import Futhark.Tools (partitionChunkedKernelLambdaParameters, fullSliceNum) +import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem, IntegralExp) +import Futhark.Util (splitAt3) + +type CallKernelGen = ImpGen.ImpM ExplicitMemory Imp.HostOp +type InKernelGen = ImpGen.ImpM InKernel Imp.KernelOp + +callKernelOperations :: ImpGen.Operations ExplicitMemory Imp.HostOp +callKernelOperations = + ImpGen.Operations { ImpGen.opsExpCompiler = expCompiler + , ImpGen.opsCopyCompiler = callKernelCopy + , ImpGen.opsOpCompiler = opCompiler + , ImpGen.opsBodyCompiler = ImpGen.defCompileBody + } + +inKernelOperations :: KernelConstants -> ImpGen.Operations InKernel Imp.KernelOp +inKernelOperations constants = (ImpGen.defaultOperations $ compileInKernelOp constants) + { ImpGen.opsCopyCompiler = inKernelCopy + , ImpGen.opsExpCompiler = inKernelExpCompiler + , ImpGen.opsBodyCompiler = compileNestedKernelBody constants + } + +compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError Imp.Program) +compileProg prog = + fmap (setDefaultSpace (Imp.Space "device")) <$> + ImpGen.compileProg callKernelOperations (Imp.Space "device") prog + +opCompiler :: ImpGen.Destination -> Op ExplicitMemory + -> CallKernelGen () +opCompiler dest (Alloc e space) = + ImpGen.compileAlloc dest e space +opCompiler dest (Inner kernel) = + kernelCompiler dest kernel + +compileInKernelOp :: KernelConstants -> ImpGen.Destination -> Op InKernel + -> InKernelGen () +compileInKernelOp _ (ImpGen.Destination _ [ImpGen.MemoryDestination mem]) Alloc{} = + compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." +compileInKernelOp _ dest Alloc{} = + compilerBugS $ "Invalid target for in-kernel allocation: " ++ show dest +compileInKernelOp constants dest (Inner op) = + compileKernelExp constants dest op + +-- | Recognise kernels (maps), give everything else back. +kernelCompiler :: ImpGen.Destination -> Kernel InKernel + -> CallKernelGen () + +kernelCompiler dest (GetSize key size_class) = do + [v] <- ImpGen.funcallTargets dest + ImpGen.emit $ Imp.Op $ Imp.GetSize v key size_class + +kernelCompiler dest (CmpSizeLe key size_class x) = do + [v] <- ImpGen.funcallTargets dest + ImpGen.emit =<< Imp.Op . Imp.CmpSizeLe v key size_class <$> ImpGen.compileSubExp x + +kernelCompiler dest (GetSizeMax size_class) = do + [v] <- ImpGen.funcallTargets dest + ImpGen.emit $ Imp.Op $ Imp.GetSizeMax v size_class + +kernelCompiler dest (Kernel desc space _ kernel_body) = do + + num_groups' <- ImpGen.subExpToDimSize $ spaceNumGroups space + group_size' <- ImpGen.subExpToDimSize $ spaceGroupSize space + num_threads' <- ImpGen.subExpToDimSize $ spaceNumThreads space + + let bound_in_kernel = + M.keys $ + scopeOfKernelSpace space <> + scopeOf (kernelBodyStms kernel_body) + + let global_tid = spaceGlobalId space + local_tid = spaceLocalId space + group_id = spaceGroupId space + wave_size <- newVName "wave_size" + inner_group_size <- newVName "group_size" + thread_active <- newVName "thread_active" + + let (space_is, space_dims) = unzip $ spaceDimensions space + space_dims' <- mapM ImpGen.compileSubExp space_dims + let constants = KernelConstants global_tid local_tid group_id + group_size' num_threads' + (Imp.VarSize wave_size) (zip space_is space_dims') + (Imp.var thread_active Bool) mempty + + kernel_body' <- + makeAllMemoryGlobal $ + ImpGen.subImpM_ (inKernelOperations constants) $ + ImpGen.declaringPrimVar wave_size int32 $ + ImpGen.declaringPrimVar inner_group_size int32 $ + ImpGen.declaringPrimVar thread_active Bool $ + ImpGen.declaringScope Nothing (scopeOfKernelSpace space) $ do + + ImpGen.emit $ + Imp.Op (Imp.GetGlobalId global_tid 0) <> + Imp.Op (Imp.GetLocalId local_tid 0) <> + Imp.Op (Imp.GetLocalSize inner_group_size 0) <> + Imp.Op (Imp.GetLockstepWidth wave_size) <> + Imp.Op (Imp.GetGroupId group_id 0) + + setSpaceIndices space + + ImpGen.emit $ Imp.SetScalar thread_active (isActive $ spaceDimensions space) + + compileKernelBody dest constants kernel_body + + (uses, local_memory) <- computeKernelUses kernel_body' bound_in_kernel + + forM_ (kernelHints desc) $ \(s,v) -> do + ty <- case v of + Constant pv -> return $ Prim $ primValueType pv + Var vn -> lookupType vn + unless (primType ty) $ fail $ concat [ "debugKernelHint '", s, "'" + , " in kernel '", kernelName desc, "'" + , " did not have primType value." ] + + ImpGen.compileSubExp v >>= ImpGen.emit . Imp.DebugPrint s (elemType ty) + + ImpGen.emit $ Imp.Op $ Imp.CallKernel $ Imp.AnyKernel Imp.Kernel + { Imp.kernelBody = kernel_body' + , Imp.kernelLocalMemory = local_memory + , Imp.kernelUses = uses + , Imp.kernelNumGroups = num_groups' + , Imp.kernelGroupSize = group_size' + , Imp.kernelName = global_tid + , Imp.kernelDesc = kernelName desc + } + +expCompiler :: ImpGen.ExpCompiler ExplicitMemory Imp.HostOp +-- We generate a simple kernel for itoa and replicate. +expCompiler + (ImpGen.Destination tag [ImpGen.ArrayDestination (Just destloc)]) + (BasicOp (Iota n x s et)) = do + thread_gid <- maybe (newVName "thread_gid") (return . VName (nameFromString "thread_gid")) tag + + makeAllMemoryGlobal $ do + (destmem, destspace, destidx) <- + ImpGen.fullyIndexArray' destloc [ImpGen.varIndex thread_gid] (IntType et) + + n' <- ImpGen.compileSubExp n + x' <- ImpGen.compileSubExp x + s' <- ImpGen.compileSubExp s + + let body = Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ + Imp.ConvOpExp (SExt Int32 et) (Imp.var thread_gid int32) * s' + x' + + (group_size, num_groups) <- computeMapKernelGroups n' + + (body_uses, _) <- computeKernelUses + (freeIn body <> freeIn [n',x',s']) + [thread_gid] + + ImpGen.emit $ Imp.Op $ Imp.CallKernel $ Imp.Map Imp.MapKernel + { Imp.mapKernelThreadNum = thread_gid + , Imp.mapKernelDesc = "iota" + , Imp.mapKernelNumGroups = Imp.VarSize num_groups + , Imp.mapKernelGroupSize = Imp.VarSize group_size + , Imp.mapKernelSize = n' + , Imp.mapKernelUses = body_uses + , Imp.mapKernelBody = body + } + +expCompiler + (ImpGen.Destination tag [dest]) (BasicOp (Replicate (Shape ds) se)) = do + constants <- simpleKernelConstants tag "replicate" + + t <- subExpType se + let thread_gid = kernelGlobalThreadId constants + row_dims = arrayDims t + dims = ds ++ row_dims + is' = unflattenIndex (map (ImpGen.compileSubExpOfType int32) dims) $ + ImpGen.varIndex thread_gid + ds' <- mapM ImpGen.compileSubExp ds + + makeAllMemoryGlobal $ do + body <- ImpGen.subImpM_ (inKernelOperations constants) $ + ImpGen.copyDWIMDest dest is' se $ drop (length ds) is' + + dims' <- mapM ImpGen.compileSubExp dims + (group_size, num_groups) <- computeMapKernelGroups $ product dims' + + (body_uses, _) <- computeKernelUses + (freeIn body <> freeIn ds') + [thread_gid] + + ImpGen.emit $ Imp.Op $ Imp.CallKernel $ Imp.Map Imp.MapKernel + { Imp.mapKernelThreadNum = thread_gid + , Imp.mapKernelDesc = "replicate" + , Imp.mapKernelNumGroups = Imp.VarSize num_groups + , Imp.mapKernelGroupSize = Imp.VarSize group_size + , Imp.mapKernelSize = product dims' + , Imp.mapKernelUses = body_uses + , Imp.mapKernelBody = body + } + +-- Allocation in the "local" space is just a placeholder. +expCompiler _ (Op (Alloc _ (Space "local"))) = + return () + +expCompiler dest e = + ImpGen.defCompileExp dest e + +callKernelCopy :: ImpGen.CopyCompiler ExplicitMemory Imp.HostOp +callKernelCopy bt + destloc@(ImpGen.MemLocation destmem destshape destIxFun) + srcloc@(ImpGen.MemLocation srcmem srcshape srcIxFun) + n + | Just (destoffset, srcoffset, + num_arrays, size_x, size_y, + src_elems, dest_elems) <- isMapTransposeKernel bt destloc srcloc = + ImpGen.emit $ Imp.Op $ Imp.CallKernel $ + Imp.MapTranspose bt + destmem destoffset + srcmem srcoffset + num_arrays size_x size_y + src_elems dest_elems + + | bt_size <- primByteSize bt, + ixFunMatchesInnerShape + (Shape $ map Imp.sizeToExp destshape) destIxFun, + ixFunMatchesInnerShape + (Shape $ map Imp.sizeToExp srcshape) srcIxFun, + Just destoffset <- + IxFun.linearWithOffset destIxFun bt_size, + Just srcoffset <- + IxFun.linearWithOffset srcIxFun bt_size = do + let row_size = product $ map ImpGen.dimSizeToExp $ drop 1 srcshape + srcspace <- ImpGen.entryMemSpace <$> ImpGen.lookupMemory srcmem + destspace <- ImpGen.entryMemSpace <$> ImpGen.lookupMemory destmem + ImpGen.emit $ Imp.Copy + destmem (bytes destoffset) destspace + srcmem (bytes srcoffset) srcspace $ + (n * row_size) `Imp.withElemType` bt + + | otherwise = do + global_thread_index <- newVName "copy_global_thread_index" + + -- Note that the shape of the destination and the source are + -- necessarily the same. + let shape = map Imp.sizeToExp srcshape + shape_se = map (Imp.innerExp . ImpGen.dimSizeToExp) srcshape + dest_is = unflattenIndex shape_se $ ImpGen.varIndex global_thread_index + src_is = dest_is + + makeAllMemoryGlobal $ do + (_, destspace, destidx) <- ImpGen.fullyIndexArray' destloc dest_is bt + (_, srcspace, srcidx) <- ImpGen.fullyIndexArray' srcloc src_is bt + + let body = Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ + Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile + + destmem_size <- ImpGen.entryMemSize <$> ImpGen.lookupMemory destmem + let writes_to = [Imp.MemoryUse destmem destmem_size] + + reads_from <- readsFromSet $ + S.singleton srcmem <> + freeIn destIxFun <> freeIn srcIxFun <> freeIn destshape + + let kernel_size = Imp.innerExp n * product (drop 1 shape) + (group_size, num_groups) <- computeMapKernelGroups kernel_size + + let bound_in_kernel = [global_thread_index] + (body_uses, _) <- computeKernelUses (kernel_size, body) bound_in_kernel + + ImpGen.emit $ Imp.Op $ Imp.CallKernel $ Imp.Map Imp.MapKernel + { Imp.mapKernelThreadNum = global_thread_index + , Imp.mapKernelDesc = "copy" + , Imp.mapKernelNumGroups = Imp.VarSize num_groups + , Imp.mapKernelGroupSize = Imp.VarSize group_size + , Imp.mapKernelSize = kernel_size + , Imp.mapKernelUses = nub $ body_uses ++ writes_to ++ reads_from + , Imp.mapKernelBody = body + } + +-- | We have no bulk copy operation (e.g. memmove) inside kernels, so +-- turn any copy into a loop. +inKernelCopy :: ImpGen.CopyCompiler InKernel Imp.KernelOp +inKernelCopy = ImpGen.copyElementWise + +inKernelExpCompiler :: ImpGen.ExpCompiler InKernel Imp.KernelOp +inKernelExpCompiler _ (BasicOp (Assert _ _ (loc, locs))) = + compilerLimitationS $ + unlines [ "Cannot compile assertion at " ++ + intercalate " -> " (reverse $ map locStr $ loc:locs) ++ + " inside parallel kernel." + , "As a workaround, surround the expression with 'unsafe'."] +-- The static arrays stuff does not work inside kernels. +inKernelExpCompiler (ImpGen.Destination _ [dest]) (BasicOp (ArrayLit es _)) = + forM_ (zip [0..] es) $ \(i,e) -> + ImpGen.copyDWIMDest dest [fromIntegral (i::Int32)] e [] +inKernelExpCompiler dest e = + ImpGen.defCompileExp dest e + +computeKernelUses :: FreeIn a => + a -> [VName] + -> CallKernelGen ([Imp.KernelUse], [Imp.LocalMemoryUse]) +computeKernelUses kernel_body bound_in_kernel = do + let actually_free = freeIn kernel_body `S.difference` S.fromList bound_in_kernel + + -- Compute the variables that we need to pass to the kernel. + reads_from <- readsFromSet actually_free + + -- Are we using any local memory? + local_memory <- computeLocalMemoryUse actually_free + return (nub reads_from, nub local_memory) + +readsFromSet :: Names -> CallKernelGen [Imp.KernelUse] +readsFromSet free = + fmap catMaybes $ + forM (S.toList free) $ \var -> do + t <- lookupType var + case t of + Array {} -> return Nothing + Mem _ (Space "local") -> return Nothing + Mem memsize _ -> Just <$> (Imp.MemoryUse var <$> + ImpGen.subExpToDimSize memsize) + Prim bt -> + isConstExp var >>= \case + Just ce -> return $ Just $ Imp.ConstUse var ce + Nothing | bt == Cert -> return Nothing + | otherwise -> return $ Just $ Imp.ScalarUse var bt + +computeLocalMemoryUse :: Names -> CallKernelGen [Imp.LocalMemoryUse] +computeLocalMemoryUse free = + fmap catMaybes $ + forM (S.toList free) $ \var -> do + t <- lookupType var + case t of + Mem memsize (Space "local") -> do + memsize' <- localMemSize =<< ImpGen.subExpToDimSize memsize + return $ Just (var, memsize') + _ -> return Nothing + +localMemSize :: Imp.MemSize -> CallKernelGen (Either Imp.MemSize Imp.KernelConstExp) +localMemSize (Imp.ConstSize x) = + return $ Right $ ValueExp $ IntValue $ Int64Value x +localMemSize (Imp.VarSize v) = isConstExp v >>= \case + Just e | isStaticExp e -> return $ Right e + _ -> return $ Left $ Imp.VarSize v + +-- | Only some constant expressions quality as *static* expressions, +-- which we can use for static memory allocation. This is a bit of a +-- hack, as it is primarly motivated by what you can put as the size +-- when declaring an array in C. +isStaticExp :: Imp.KernelConstExp -> Bool +isStaticExp LeafExp{} = True +isStaticExp ValueExp{} = True +isStaticExp (BinOpExp Add{} x y) = isStaticExp x && isStaticExp y +isStaticExp (BinOpExp Sub{} x y) = isStaticExp x && isStaticExp y +isStaticExp (BinOpExp Mul{} x y) = isStaticExp x && isStaticExp y +isStaticExp _ = False + +isConstExp :: VName -> CallKernelGen (Maybe Imp.KernelConstExp) +isConstExp v = do + vtable <- asks ImpGen.envVtable + let lookupConstExp name = constExp =<< hasExp =<< M.lookup name vtable + constExp (Op (Inner (GetSize key _))) = Just $ LeafExp (Imp.SizeConst key) int32 + constExp e = primExpFromExp lookupConstExp e + return $ lookupConstExp v + where hasExp (ImpGen.ArrayVar e _) = e + hasExp (ImpGen.ScalarVar e _) = e + hasExp (ImpGen.MemVar e _) = e + +-- | 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 declared as variables in the +-- kernel). +makeAllMemoryGlobal :: CallKernelGen a + -> CallKernelGen a +makeAllMemoryGlobal = + local $ \env -> env { ImpGen.envVtable = M.map globalMemory $ ImpGen.envVtable env + , ImpGen.envDefaultSpace = Imp.Space "global" + } + where globalMemory (ImpGen.MemVar _ entry) + | ImpGen.entryMemSpace entry /= Space "local" = + ImpGen.MemVar Nothing entry { ImpGen.entryMemSpace = Imp.Space "global" } + globalMemory entry = + entry + +computeMapKernelGroups :: Imp.Exp -> CallKernelGen (VName, VName) +computeMapKernelGroups kernel_size = do + group_size <- newVName "group_size" + num_groups <- newVName "num_groups" + let group_size_var = Imp.var group_size int32 + ImpGen.emit $ Imp.DeclareScalar group_size int32 + ImpGen.emit $ Imp.DeclareScalar num_groups int32 + ImpGen.emit $ Imp.Op $ Imp.GetSize group_size group_size Imp.SizeGroup + ImpGen.emit $ Imp.SetScalar num_groups $ + kernel_size `quotRoundingUp` Imp.ConvOpExp (SExt Int32 Int32) group_size_var + return (group_size, num_groups) + +isMapTransposeKernel :: PrimType -> ImpGen.MemLocation -> ImpGen.MemLocation + -> Maybe (Imp.Exp, Imp.Exp, + Imp.Exp, Imp.Exp, Imp.Exp, + Imp.Exp, Imp.Exp) +isMapTransposeKernel bt + (ImpGen.MemLocation _ _ destIxFun) + (ImpGen.MemLocation _ _ srcIxFun) + | Just (dest_offset, perm_and_destshape) <- IxFun.rearrangeWithOffset destIxFun bt_size, + (perm, destshape) <- unzip perm_and_destshape, + srcshape' <- IxFun.shape srcIxFun, + Just src_offset <- IxFun.linearWithOffset srcIxFun bt_size, + Just (r1, r2, _) <- isMapTranspose perm = + isOk (product srcshape') (product destshape) destshape swap r1 r2 dest_offset src_offset + | Just dest_offset <- IxFun.linearWithOffset destIxFun bt_size, + Just (src_offset, perm_and_srcshape) <- IxFun.rearrangeWithOffset srcIxFun bt_size, + (perm, srcshape) <- unzip perm_and_srcshape, + destshape' <- IxFun.shape destIxFun, + Just (r1, r2, _) <- isMapTranspose perm = + isOk (product srcshape) (product destshape') srcshape id r1 r2 dest_offset src_offset + | otherwise = + Nothing + where bt_size = primByteSize bt + swap (x,y) = (y,x) + + isOk src_elems dest_elems shape f r1 r2 dest_offset src_offset = do + let (num_arrays, size_x, size_y) = getSizes shape f r1 r2 + return (dest_offset, src_offset, + num_arrays, size_x, size_y, + src_elems, dest_elems) + + getSizes shape f r1 r2 = + let (mapped, notmapped) = splitAt r1 shape + (pretrans, posttrans) = f $ splitAt r2 notmapped + in (product mapped, product pretrans, product posttrans) + +writeParamToLocalMemory :: Typed (MemBound u) => + Imp.Exp -> (VName, t) -> Param (MemBound u) + -> ImpGen.ImpM lore op () +writeParamToLocalMemory i (mem, _) param + | Prim t <- paramType param = + ImpGen.emit $ + Imp.Write mem (bytes i') bt (Space "local") Imp.Volatile $ + Imp.var (paramName param) t + | otherwise = + return () + where i' = i * Imp.LeafExp (Imp.SizeOf bt) int32 + bt = elemType $ paramType param + +readParamFromLocalMemory :: Typed (MemBound u) => + VName -> Imp.Exp -> Param (MemBound u) -> (VName, t) + -> ImpGen.ImpM lore op () +readParamFromLocalMemory index i param (l_mem, _) + | Prim _ <- paramType param = + ImpGen.emit $ + Imp.SetScalar (paramName param) $ + Imp.index l_mem (bytes i') bt (Space "local") Imp.Volatile + | otherwise = + ImpGen.emit $ + Imp.SetScalar index i + where i' = i * Imp.LeafExp (Imp.SizeOf bt) int32 + bt = elemType $ paramType param + +computeThreadChunkSize :: SplitOrdering + -> Imp.Exp + -> Imp.Count Imp.Elements + -> Imp.Count Imp.Elements + -> VName + -> ImpGen.ImpM lore op () +computeThreadChunkSize (SplitStrided stride) thread_index elements_per_thread num_elements chunk_var = do + stride' <- ImpGen.compileSubExp stride + ImpGen.emit $ Imp.SetScalar chunk_var $ Imp.BinOpExp (SMin Int32) + (Imp.innerExp elements_per_thread) $ + (Imp.innerExp num_elements - thread_index) + `quotRoundingUp` + stride' + +computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do + starting_point <- newVName "starting_point" + remaining_elements <- newVName "remaining_elements" + + ImpGen.emit $ + Imp.DeclareScalar starting_point int32 + ImpGen.emit $ + Imp.SetScalar starting_point $ + thread_index * Imp.innerExp elements_per_thread + + ImpGen.emit $ + Imp.DeclareScalar remaining_elements int32 + ImpGen.emit $ + Imp.SetScalar remaining_elements $ + Imp.innerExp num_elements - Imp.var starting_point int32 + + let no_remaining_elements = Imp.CmpOpExp (CmpSle Int32) + (Imp.var remaining_elements int32) 0 + beyond_bounds = Imp.CmpOpExp (CmpSle Int32) + (Imp.innerExp num_elements) + (Imp.var starting_point int32) + + ImpGen.emit $ + Imp.If (Imp.BinOpExp LogOr no_remaining_elements beyond_bounds) + (Imp.SetScalar chunk_var 0) + (Imp.If is_last_thread + (Imp.SetScalar chunk_var $ Imp.innerExp last_thread_elements) + (Imp.SetScalar chunk_var $ Imp.innerExp elements_per_thread)) + where last_thread_elements = + num_elements - Imp.elements thread_index * elements_per_thread + is_last_thread = + Imp.CmpOpExp (CmpSlt Int32) + (Imp.innerExp num_elements) + ((thread_index + 1) * Imp.innerExp elements_per_thread) + +inBlockScan :: Imp.Exp + -> Imp.Exp + -> Imp.Exp + -> VName + -> [(VName, t)] + -> Lambda InKernel + -> InKernelGen () +inBlockScan lockstep_width block_size active local_id acc_local_mem scan_lam = ImpGen.everythingVolatile $ do + skip_threads <- newVName "skip_threads" + let in_block_thread_active = + Imp.CmpOpExp (CmpSle Int32) (Imp.var skip_threads int32) in_block_id + (scan_lam_i, other_index_param, actual_params) = + partitionChunkedKernelLambdaParameters $ lambdaParams scan_lam + (x_params, y_params) = + splitAt (length actual_params `div` 2) actual_params + read_operands <- + ImpGen.collect $ + zipWithM_ (readParamFromLocalMemory (paramName other_index_param) $ + Imp.var local_id int32 - Imp.var skip_threads int32) + x_params acc_local_mem + scan_y_dest <- ImpGen.destinationFromParams y_params + + -- Set initial y values + read_my_initial <- ImpGen.collect $ + zipWithM_ (readParamFromLocalMemory scan_lam_i $ Imp.var local_id int32) + y_params acc_local_mem + ImpGen.emit $ Imp.If active read_my_initial mempty + + op_to_y <- ImpGen.collect $ ImpGen.compileBody scan_y_dest $ lambdaBody scan_lam + write_operation_result <- + ImpGen.collect $ + zipWithM_ (writeParamToLocalMemory $ Imp.var local_id int32) + acc_local_mem y_params + let andBlockActive = Imp.BinOpExp LogAnd active + maybeBarrier = Imp.If (Imp.CmpOpExp (CmpSle Int32) lockstep_width (Imp.var skip_threads int32)) + (Imp.Op Imp.Barrier) mempty + + ImpGen.emit $ + Imp.Comment "in-block scan (hopefully no barriers needed)" $ + Imp.DeclareScalar skip_threads int32 <> + Imp.SetScalar skip_threads 1 <> + Imp.While (Imp.CmpOpExp (CmpSlt Int32) (Imp.var skip_threads int32) block_size) + (Imp.If (andBlockActive in_block_thread_active) + (Imp.Comment "read operands" read_operands <> + Imp.Comment "perform operation" op_to_y) mempty <> + + maybeBarrier <> + + Imp.If (andBlockActive in_block_thread_active) + (Imp.Comment "write result" write_operation_result) mempty <> + maybeBarrier <> + Imp.SetScalar skip_threads (Imp.var skip_threads int32 * 2)) + where block_id = Imp.BinOpExp (SQuot Int32) (Imp.var local_id int32) block_size + in_block_id = Imp.var local_id int32 - block_id * block_size + +data KernelConstants = KernelConstants + { kernelGlobalThreadId :: VName + , kernelLocalThreadId :: VName + , kernelGroupId :: VName + , kernelGroupSize :: Imp.DimSize + , _kernelNumThreads :: Imp.DimSize + , kernelWaveSize :: Imp.DimSize + , kernelDimensions :: [(VName, Imp.Exp)] + , kernelThreadActive :: Imp.Exp + , kernelStreamed :: [(VName, Imp.DimSize)] + -- ^ Chunk sizez and their maximum size. Hint + -- for unrolling. + } + +-- FIXME: wing a KernelConstants structure for use in Replicate +-- compilation. This cannot be the best way to do this... +simpleKernelConstants :: MonadFreshNames m => + Maybe Int -> String + -> m KernelConstants +simpleKernelConstants tag desc = do + thread_gtid <- maybe (newVName $ desc ++ "_gtid") + (return . VName (nameFromString $ desc ++ "_gtid")) tag + thread_ltid <- newVName $ desc ++ "_ltid" + thread_gid <- newVName $ desc ++ "_gid" + return $ KernelConstants + thread_gtid thread_ltid thread_gid + (Imp.ConstSize 0) (Imp.ConstSize 0) (Imp.ConstSize 0) + [] (Imp.ValueExp $ BoolValue True) mempty + +compileKernelBody :: ImpGen.Destination + -> KernelConstants + -> KernelBody InKernel + -> InKernelGen () +compileKernelBody (ImpGen.Destination _ dest) constants kbody = + compileKernelStms constants (stmsToList $ kernelBodyStms kbody) $ + zipWithM_ (compileKernelResult constants) dest $ + kernelBodyResult kbody + +compileNestedKernelBody :: KernelConstants + -> ImpGen.Destination + -> Body InKernel + -> InKernelGen () +compileNestedKernelBody constants (ImpGen.Destination _ dest) kbody = + compileKernelStms constants (stmsToList $ bodyStms kbody) $ + zipWithM_ ImpGen.compileSubExpTo dest $ bodyResult kbody + +compileKernelStms :: KernelConstants -> [Stm InKernel] + -> InKernelGen a + -> InKernelGen a +compileKernelStms constants ungrouped_bnds m = + compileGroupedKernelStms' $ groupStmsByGuard constants ungrouped_bnds + where compileGroupedKernelStms' [] = m + compileGroupedKernelStms' ((g, bnds):rest_bnds) = + ImpGen.declaringScopes + (map ((Just . stmExp) &&& (castScope . scopeOf)) bnds) $ do + protect g $ mapM_ compileKernelStm bnds + compileGroupedKernelStms' rest_bnds + + protect Nothing body_m = + body_m + protect (Just (Imp.ValueExp (BoolValue True))) body_m = + body_m + protect (Just g) body_m = do + body <- allThreads constants body_m + ImpGen.emit $ Imp.If g body mempty + + compileKernelStm (Let pat _ e) = do + dest <- ImpGen.destinationFromPattern pat + ImpGen.compileExp dest e + +groupStmsByGuard :: KernelConstants + -> [Stm InKernel] + -> [(Maybe Imp.Exp, [Stm InKernel])] +groupStmsByGuard constants bnds = + map collapse $ groupBy sameGuard $ zip (map bindingGuard bnds) bnds + where bindingGuard (Let _ _ Op{}) = Nothing + bindingGuard _ = Just $ kernelThreadActive constants + + sameGuard (g1, _) (g2, _) = g1 == g2 + + collapse [] = + (Nothing, []) + collapse l@((g,_):_) = + (g, map snd l) + +compileKernelExp :: KernelConstants -> ImpGen.Destination -> KernelExp InKernel + -> InKernelGen () + +compileKernelExp _ (ImpGen.Destination _ dests) (Barrier ses) = do + zipWithM_ ImpGen.compileSubExpTo dests ses + ImpGen.emit $ Imp.Op Imp.Barrier + +compileKernelExp _ dest (SplitSpace o w i elems_per_thread) + | ImpGen.Destination _ [ImpGen.ScalarDestination size] <- dest = do + num_elements <- Imp.elements <$> ImpGen.compileSubExp w + i' <- ImpGen.compileSubExp i + elems_per_thread' <- Imp.elements <$> ImpGen.compileSubExp elems_per_thread + computeThreadChunkSize o i' elems_per_thread' num_elements size + +compileKernelExp constants dest (Combine (CombineSpace scatter cspace) ts aspace body) = do + -- First we compute how many times we have to iterate to cover + -- cspace with our group size. It is a fairly common case that + -- we statically know that this requires 1 iteration, so we + -- could detect it and not generate a loop in that case. + -- However, it seems to have no impact on performance (an extra + -- conditional jump), so for simplicity we just always generate + -- the loop. + let cspace_dims = map (streamBounded . snd) cspace + num_iters = product cspace_dims `quotRoundingUp` + Imp.sizeToExp (kernelGroupSize constants) + + iter <- newVName "comb_iter" + cid <- newVName "flat_comb_id" + + one_iteration <- ImpGen.collect $ + ImpGen.declaringPrimVars (zip (map fst cspace) $ repeat int32) $ + ImpGen.declaringPrimVar cid int32 $ do + + -- Compute the *flat* array index. + ImpGen.emit $ Imp.SetScalar cid $ + Imp.var iter int32 * Imp.sizeToExp (kernelGroupSize constants) + + Imp.var (kernelLocalThreadId constants) int32 + + -- Turn it into a nested array index. + forM_ (zip (map fst cspace) $ unflattenIndex cspace_dims (Imp.var cid int32)) $ \(v, x) -> + ImpGen.emit $ Imp.SetScalar v x + + -- Construct the body. This is mostly about the book-keeping + -- for the scatter-like part. + let (scatter_ws, scatter_ns, _scatter_vs) = unzip3 scatter + scatter_ws_repl = concat $ zipWith replicate scatter_ns scatter_ws + (scatter_dests, normal_dests) = + splitAt (sum scatter_ns) $ ImpGen.valueDestinations dest + (res_is, res_vs, res_normal) = + splitAt3 (sum scatter_ns) (sum scatter_ns) $ bodyResult body + scatter_is = map (pure . DimFix . ImpGen.compileSubExpOfType int32) res_is + scatter_dests_repl = concat $ zipWith replicate scatter_ns scatter_dests + (scatter_dests', normal_dests') <- + case (sequence $ zipWith3 index scatter_is ts scatter_dests_repl, + zipWithM (index local_index) (drop (sum scatter_ns*2) ts) normal_dests) of + (Just x, Just y) -> return (x, y) + _ -> fail "compileKernelExp combine: invalid destination." + body' <- allThreads constants $ + ImpGen.compileStms (freeIn $ bodyResult body) (stmsToList $ bodyStms body) $ do + + forM_ (zip4 scatter_ws_repl res_is res_vs scatter_dests') $ + \(w, res_i, res_v, scatter_dest) -> do + let res_i' = ImpGen.compileSubExpOfType int32 res_i + w' = ImpGen.compileSubExpOfType int32 w + -- We have to check that 'res_i' is in-bounds wrt. an array of size 'w'. + in_bounds = BinOpExp LogAnd (CmpOpExp (CmpSle Int32) 0 res_i') + (CmpOpExp (CmpSlt Int32) res_i' w') + when_in_bounds <- ImpGen.collect $ ImpGen.compileSubExpTo scatter_dest res_v + ImpGen.emit $ Imp.If in_bounds when_in_bounds mempty + + zipWithM_ ImpGen.compileSubExpTo normal_dests' res_normal + + -- Execute the body if we are within bounds. + ImpGen.emit $ + Imp.If (Imp.BinOpExp LogAnd (isActive cspace) (isActive aspace)) body' mempty + + ImpGen.emit $ Imp.For iter Int32 num_iters one_iteration + ImpGen.emit $ Imp.Op Imp.Barrier + + where streamBounded (Var v) + | Just x <- lookup v $ kernelStreamed constants = + Imp.sizeToExp x + streamBounded se = ImpGen.compileSubExpOfType int32 se + + local_index = map (DimFix . ImpGen.varIndex . fst) cspace + + index i t (ImpGen.ArrayDestination (Just loc)) = + let space_dims = map (ImpGen.varIndex . fst) cspace + t_dims = map (ImpGen.compileSubExpOfType int32) $ arrayDims t + in Just $ ImpGen.ArrayDestination $ + Just $ ImpGen.sliceArray loc $ + fullSliceNum (space_dims++t_dims) i + index _ _ _ = Nothing + +compileKernelExp constants (ImpGen.Destination _ dests) (GroupReduce w lam input) = do + skip_waves <- newVName "skip_waves" + w' <- ImpGen.compileSubExp w + + let local_tid = kernelLocalThreadId constants + (_nes, arrs) = unzip input + (reduce_i, reduce_j_param, actual_reduce_params) = + partitionChunkedKernelLambdaParameters $ lambdaParams lam + (reduce_acc_params, reduce_arr_params) = + splitAt (length input) actual_reduce_params + reduce_j = paramName reduce_j_param + + offset <- newVName "offset" + ImpGen.emit $ Imp.DeclareScalar offset int32 + + ImpGen.Destination _ reduce_acc_targets <- + ImpGen.destinationFromParams reduce_acc_params + + ImpGen.declaringPrimVar skip_waves int32 $ + ImpGen.declaringLParams (lambdaParams lam) $ do + + ImpGen.emit $ Imp.SetScalar reduce_i $ Imp.var local_tid int32 + + let setOffset x = + Imp.SetScalar offset x <> + Imp.SetScalar reduce_j (Imp.var local_tid int32 + Imp.var offset int32) + ImpGen.emit $ setOffset 0 + + set_init_params <- ImpGen.collect $ + zipWithM_ (readReduceArgument offset) reduce_acc_params arrs + ImpGen.emit $ + Imp.If (Imp.CmpOpExp (CmpSlt Int32) (Imp.var local_tid int32) w') + set_init_params mempty + + let read_reduce_args = zipWithM_ (readReduceArgument offset) + reduce_arr_params arrs + reduce_acc_dest = ImpGen.Destination Nothing reduce_acc_targets + do_reduce = do ImpGen.comment "read array element" read_reduce_args + ImpGen.compileBody reduce_acc_dest $ lambdaBody lam + zipWithM_ (writeReduceOpResult local_tid) + reduce_acc_params arrs + + in_wave_reduce <- ImpGen.collect $ ImpGen.everythingVolatile do_reduce + cross_wave_reduce <- ImpGen.collect do_reduce + + let wave_size = Imp.sizeToExp $ kernelWaveSize constants + group_size = Imp.sizeToExp $ kernelGroupSize constants + wave_id = Imp.var local_tid int32 `quot` wave_size + in_wave_id = Imp.var local_tid int32 - wave_id * wave_size + num_waves = (group_size + wave_size - 1) `quot` wave_size + arg_in_bounds = Imp.CmpOpExp (CmpSlt Int32) + (Imp.var reduce_j int32) w' + + doing_in_wave_reductions = + Imp.CmpOpExp (CmpSlt Int32) (Imp.var offset int32) wave_size + apply_in_in_wave_iteration = + Imp.CmpOpExp (CmpEq int32) + (Imp.BinOpExp (And Int32) in_wave_id (2 * Imp.var offset int32 - 1)) 0 + in_wave_reductions = + setOffset 1 <> + Imp.While doing_in_wave_reductions + (Imp.If (Imp.BinOpExp LogAnd arg_in_bounds apply_in_in_wave_iteration) + in_wave_reduce mempty <> + setOffset (Imp.var offset int32 * 2)) + + doing_cross_wave_reductions = + Imp.CmpOpExp (CmpSlt Int32) (Imp.var skip_waves int32) num_waves + is_first_thread_in_wave = + Imp.CmpOpExp (CmpEq int32) in_wave_id 0 + wave_not_skipped = + Imp.CmpOpExp (CmpEq int32) + (Imp.BinOpExp (And Int32) wave_id (2 * Imp.var skip_waves int32 - 1)) + 0 + apply_in_cross_wave_iteration = + Imp.BinOpExp LogAnd arg_in_bounds $ + Imp.BinOpExp LogAnd is_first_thread_in_wave wave_not_skipped + cross_wave_reductions = + Imp.SetScalar skip_waves 1 <> + Imp.While doing_cross_wave_reductions + (Imp.Op Imp.Barrier <> + setOffset (Imp.var skip_waves int32 * wave_size) <> + Imp.If apply_in_cross_wave_iteration + cross_wave_reduce mempty <> + Imp.SetScalar skip_waves (Imp.var skip_waves int32 * 2)) + + ImpGen.emit $ + in_wave_reductions <> cross_wave_reductions + + forM_ (zip dests reduce_acc_params) $ \(dest, reduce_acc_param) -> + ImpGen.copyDWIMDest dest [] (Var $ paramName reduce_acc_param) [] + where readReduceArgument offset param arr + | Prim _ <- paramType param = + ImpGen.copyDWIM (paramName param) [] (Var arr) [i] + | otherwise = + return () + where i = ImpGen.varIndex (kernelLocalThreadId constants) + ImpGen.varIndex offset + + writeReduceOpResult i param arr + | Prim _ <- paramType param = + ImpGen.copyDWIM arr [ImpGen.varIndex i] (Var $ paramName param) [] + | otherwise = + return () + +compileKernelExp constants _ (GroupScan w lam input) = do + renamed_lam <- renameLambda lam + w' <- ImpGen.compileSubExp w + + when (any (not . primType . paramType) $ lambdaParams lam) $ + compilerLimitationS "Cannot compile parallel scans with array element type." + + let local_tid = kernelLocalThreadId constants + (_nes, arrs) = unzip input + (lam_i, other_index_param, actual_params) = + partitionChunkedKernelLambdaParameters $ lambdaParams lam + (x_params, y_params) = + splitAt (length input) actual_params + + ImpGen.declaringLParams (lambdaParams lam++lambdaParams renamed_lam) $ do + ImpGen.emit $ Imp.SetScalar lam_i $ Imp.var local_tid int32 + + acc_local_mem <- flip zip (repeat ()) <$> + mapM (fmap (ImpGen.memLocationName . ImpGen.entryArrayLocation) . + ImpGen.lookupArray) arrs + + -- 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 = Imp.ValueExp $ IntValue $ Int32Value 32 + simd_width = Imp.sizeToExp $ kernelWaveSize constants + block_id = Imp.var local_tid int32 `quot` block_size + in_block_id = Imp.var local_tid int32 - block_id * block_size + doInBlockScan active = inBlockScan simd_width block_size active local_tid acc_local_mem + lid_in_bounds = Imp.CmpOpExp (CmpSlt Int32) (Imp.var local_tid int32) w' + + doInBlockScan lid_in_bounds lam + ImpGen.emit $ Imp.Op Imp.Barrier + + pack_block_results <- + ImpGen.collect $ + zipWithM_ (writeParamToLocalMemory block_id) acc_local_mem y_params + + let last_in_block = + Imp.CmpOpExp (CmpEq int32) in_block_id $ block_size - 1 + ImpGen.comment + "last thread of block 'i' writes its result to offset 'i'" $ + ImpGen.emit $ Imp.If (Imp.BinOpExp LogAnd last_in_block lid_in_bounds) pack_block_results mempty + + ImpGen.emit $ Imp.Op Imp.Barrier + + let is_first_block = Imp.CmpOpExp (CmpEq int32) block_id 0 + ImpGen.comment + "scan the first block, after which offset 'i' contains carry-in for warp 'i+1'" $ + doInBlockScan (Imp.BinOpExp LogAnd is_first_block lid_in_bounds) renamed_lam + + ImpGen.emit $ Imp.Op Imp.Barrier + + read_carry_in <- + ImpGen.collect $ + zipWithM_ (readParamFromLocalMemory + (paramName other_index_param) (block_id - 1)) + x_params acc_local_mem + + y_dest <- ImpGen.destinationFromParams y_params + op_to_y <- ImpGen.collect $ ImpGen.compileBody y_dest $ lambdaBody lam + write_final_result <- ImpGen.collect $ + zipWithM_ (writeParamToLocalMemory $ Imp.var local_tid int32) acc_local_mem y_params + + ImpGen.comment "carry-in for every block except the first" $ + ImpGen.emit $ Imp.If (Imp.BinOpExp LogOr + is_first_block + (Imp.UnOpExp Not lid_in_bounds)) mempty $ + Imp.Comment "read operands" read_carry_in <> + Imp.Comment "perform operation" op_to_y <> + Imp.Comment "write final result" write_final_result + + ImpGen.emit $ Imp.Op Imp.Barrier + + ImpGen.comment "restore correct values for first block" $ + ImpGen.emit $ Imp.If is_first_block write_final_result mempty + + +compileKernelExp constants (ImpGen.Destination _ final_targets) (GroupStream w maxchunk lam accs _arrs) = do + let GroupStreamLambda block_size block_offset acc_params arr_params body = lam + block_offset' = Imp.var block_offset int32 + w' <- ImpGen.compileSubExp w + max_block_size <- ImpGen.compileSubExp maxchunk + acc_dest <- ImpGen.destinationFromParams acc_params + + ImpGen.declaringLParams (acc_params++arr_params) $ do + zipWithM_ ImpGen.compileSubExpTo (ImpGen.valueDestinations acc_dest) accs + ImpGen.declaringPrimVar block_size int32 $ + -- If the GroupStream is morally just a do-loop, generate simpler code. + case mapM isSimpleThreadInSpace $ stmsToList $ bodyStms body of + Just stms' | ValueExp x <- max_block_size, oneIsh x -> do + let body' = body { bodyStms = stmsFromList stms' } + body'' <- ImpGen.withPrimVar block_offset int32 $ + allThreads constants $ ImpGen.emit =<< + ImpGen.compileLoopBody (map paramName acc_params) body' + ImpGen.emit $ Imp.SetScalar block_size 1 + + -- Check if loop is candidate for unrolling. + let loop = + case w of + Var w_var | Just w_bound <- lookup w_var $ kernelStreamed constants, + w_bound /= Imp.ConstSize 1 -> + -- Candidate for unrolling, so generate two loops. + Imp.If (CmpOpExp (CmpEq int32) w' (Imp.sizeToExp w_bound)) + (Imp.For block_offset Int32 (Imp.sizeToExp w_bound) body'') + (Imp.For block_offset Int32 w' body'') + _ -> Imp.For block_offset Int32 w' body'' + + ImpGen.emit $ + if kernelThreadActive constants == Imp.ValueExp (BoolValue True) + then loop + else Imp.If (kernelThreadActive constants) loop mempty + + _ -> ImpGen.declaringPrimVar block_offset int32 $ do + body' <- streaming constants block_size maxchunk $ + ImpGen.compileBody acc_dest body + + ImpGen.emit $ Imp.SetScalar block_offset 0 + + let not_at_end = + Imp.CmpOpExp (CmpSlt Int32) block_offset' w' + set_block_size = + Imp.If (Imp.CmpOpExp (CmpSlt Int32) + (w' - block_offset') + max_block_size) + (Imp.SetScalar block_size (w' - block_offset')) + (Imp.SetScalar block_size max_block_size) + increase_offset = + Imp.SetScalar block_offset $ + block_offset' + max_block_size + + -- Three cases to consider for simpler generated code based + -- on max block size: (0) if full input size, do not + -- generate a loop; (1) if one, generate for-loop (2) + -- otherwise, generate chunked while-loop. + ImpGen.emit $ + if max_block_size == w' then + Imp.SetScalar block_size w' <> body' + else if max_block_size == Imp.ValueExp (value (1::Int32)) then + Imp.SetScalar block_size w' <> + Imp.For block_offset Int32 w' body' + else + Imp.While not_at_end $ + set_block_size <> body' <> increase_offset + + zipWithM_ ImpGen.compileSubExpTo final_targets $ + map (Var . paramName) acc_params + + where isSimpleThreadInSpace (Let _ _ Op{}) = Nothing + isSimpleThreadInSpace bnd = Just bnd + +compileKernelExp _ _ (GroupGenReduce w [a] op bucket [v] _) + | [Prim t] <- lambdaReturnType op, + primBitSize t == 32 = do + -- If we have only one array and one non-array value (this is a + -- one-to-one correspondance) then we need only one + -- update. If 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 bits. + + -- Common variables. + old <- newVName "old" + old_bits <- newVName "old_bits" + ImpGen.emit $ Imp.DeclareScalar old t + ImpGen.emit $ Imp.DeclareScalar old_bits int32 + bucket' <- mapM ImpGen.compileSubExp bucket + w' <- mapM ImpGen.compileSubExp w + + (arr', _a_space, bucket_offset) <- ImpGen.fullyIndexArray a bucket' + + case opHasAtomicSupport old arr' bucket_offset op of + Just f -> do + val' <- ImpGen.compileSubExp v + + ImpGen.emit $ + Imp.If (indexInBounds bucket' w') + (Imp.Op $ f val') + Imp.Skip + + Nothing -> do + -- Code generation target: + -- + -- old = d_his[idx]; + -- do { + -- assumed = old; + -- tmp = OP::apply(val, assumed); + -- old = atomicCAS(&d_his[idx], assumed, tmp); + -- } while(assumed != old); + assumed <- newVName "assumed" + run_loop <- newVName "run_loop" + ImpGen.emit $ Imp.DeclareScalar assumed t + ImpGen.emit $ Imp.DeclareScalar run_loop int32 + + read_old <- ImpGen.collect $ + ImpGen.copyDWIMDest (ImpGen.ScalarDestination old) [] (Var a) bucket' + + ImpGen.emit $ + Imp.If (indexInBounds bucket' w') + -- True branch: bucket in-bounds -> enter loop + (Imp.SetScalar run_loop 1 <> read_old) + -- False branch: bucket out-of-bounds -> skip loop + (Imp.SetScalar run_loop 0) + + -- Preparing parameters + let (acc_p:arr_p:_) = lambdaParams op + + -- Store result from operator in accumulators + dests <- ImpGen.destinationFromParams [acc_p] + + -- Critical section + ImpGen.declaringLParams (lambdaParams op) $ do + bind_acc_param <- ImpGen.collect $ + ImpGen.copyDWIMDest (ImpGen.ScalarDestination $ paramName acc_p) [] v [] + + let bind_arr_param = + Imp.SetScalar (paramName arr_p) $ Imp.var assumed t + + op_body <- ImpGen.collect $ + ImpGen.compileBody dests $ lambdaBody op + + -- While-loop: Try to insert your value + let (toBits, fromBits) = + case t of FloatType Float32 -> (\x -> Imp.FunExp "to_bits32" [x] int32, + \x -> Imp.FunExp "from_bits32" [x] t) + _ -> (id, id) + ImpGen.emit $ Imp.While (Imp.var run_loop int32) + (Imp.SetScalar assumed (Imp.var old t) <> + bind_acc_param <> bind_arr_param <> op_body + <> + (Imp.Op $ + Imp.Atomic $ + Imp.AtomicCmpXchg old_bits arr' bucket_offset + (toBits (Imp.var assumed int32)) (toBits (Imp.var (paramName acc_p) int32))) + <> + Imp.SetScalar old (fromBits (Imp.var old_bits int32)) + <> + Imp.If + (Imp.CmpOpExp + (CmpEq int32) (toBits $ Imp.var assumed t) (Imp.var old_bits int32)) + -- True branch: + (Imp.SetScalar run_loop 0) + -- False branch: + Imp.Skip + ) + + where opHasAtomicSupport old arr' bucket' lam = do + let atomic f = Imp.Atomic . f old arr' bucket' + atomics = [ (Add Int32, Imp.AtomicAdd) + , (SMax Int32, Imp.AtomicSMax) + , (SMin Int32, Imp.AtomicSMin) + , (UMax Int32, Imp.AtomicUMax) + , (UMin Int32, Imp.AtomicUMin) + , (And Int32, Imp.AtomicAnd) + , (Or Int32, Imp.AtomicOr) + , (Xor Int32, Imp.AtomicXor) + ] + [BasicOp (BinOp bop _ _)] <- + Just $ map stmExp $ stmsToList $ bodyStms $ lambdaBody lam + atomic <$> lookup bop atomics + +compileKernelExp _ _ (GroupGenReduce w arrs op bucket values locks) = do + old <- newVName "old" + tmp <- newVName "tmp" + loop_done <- newVName "loop_done" + ImpGen.emit $ + Imp.DeclareScalar old int32 <> + Imp.DeclareScalar tmp int32 <> + Imp.DeclareScalar loop_done int32 + + -- Check if bucket is in-bounds + bucket' <- mapM ImpGen.compileSubExp bucket + w' <- mapM ImpGen.compileSubExp w + + -- Correctly index into locks. + (locks', _locks_space, locks_offset) <- + ImpGen.fullyIndexArray locks bucket' + + ImpGen.emit $ + Imp.If (indexInBounds bucket' w') + -- True branch: bucket in-bounds -> enter loop + (Imp.SetScalar loop_done 0) + -- False branch: bucket out-of-bounds -> skip loop + (Imp.SetScalar loop_done 1) + + -- Preparing parameters + let (acc_params, arr_params) = + splitAt (length values) $ lambdaParams op + + -- Store result from operator in accumulators + dests <- ImpGen.destinationFromParams acc_params + + -- Critical section + ImpGen.declaringLParams (lambdaParams op) $ do + let try_acquire_lock = + Imp.Op $ Imp.Atomic $ + Imp.AtomicXchg old locks' locks_offset 1 + lock_acquired = + Imp.CmpOpExp (CmpEq int32) (Imp.var old int32) 0 + loop_cond = + Imp.CmpOpExp (CmpEq int32) (Imp.var loop_done int32) 0 + break_loop = + Imp.SetScalar loop_done 1 + + -- We copy the current value and the new value to the parameters + -- unless they are array-typed. If they are arrays, then the + -- index functions should already be set up correctly, so there is + -- nothing more to do. + bind_acc_params <- ImpGen.collect $ + forM_ (zip acc_params arrs) $ \(acc_p, arr) -> + when (primType (paramType acc_p)) $ + ImpGen.copyDWIMDest (ImpGen.ScalarDestination $ paramName acc_p) [] (Var arr) bucket' + + bind_arr_params <- ImpGen.collect $ + forM_ (zip arr_params values) $ \(arr_p, val) -> + when (primType (paramType arr_p)) $ + ImpGen.copyDWIMDest (ImpGen.ScalarDestination $ paramName arr_p) [] val [] + + op_body <- ImpGen.collect $ + ImpGen.compileBody dests $ lambdaBody op + + do_gen_reduce <- ImpGen.collect $ + zipWithM_ (writeArray bucket') arrs $ map (Var . paramName) acc_params + + release_lock <- ImpGen.collect $ + ImpGen.copyDWIM locks bucket' (intConst Int32 0) [] + + -- While-loop: Try to insert your value + ImpGen.emit $ Imp.While loop_cond + (try_acquire_lock <> + Imp.If lock_acquired + -- True branch + (bind_acc_params <> bind_arr_params <> op_body <> do_gen_reduce <> release_lock <> break_loop) + -- False branch + Imp.Skip + <> + Imp.Op Imp.MemFence + ) + where writeArray i arr val = + ImpGen.copyDWIM arr i val [] + +compileKernelExp _ dest e = + compilerBugS $ unlines ["Invalid target", " " ++ show dest, + "for kernel expression", " " ++ pretty e] + +-- Requires that the lists are of equal length, otherwise +-- zip with truncate the longer list. +indexInBounds :: [Imp.Exp] -> [Imp.Exp] -> Imp.Exp +indexInBounds inds bounds = + foldl1 (Imp.BinOpExp LogAnd) $ zipWith checkBound inds bounds + where checkBound ind bound = + Imp.BinOpExp LogAnd + (Imp.CmpOpExp (CmpSle Int32) 0 ind) + (Imp.CmpOpExp (CmpSlt Int32) ind bound) + +allThreads :: KernelConstants -> InKernelGen () -> InKernelGen Imp.KernelCode +allThreads constants = ImpGen.subImpM_ $ inKernelOperations constants' + where constants' = + constants { kernelThreadActive = Imp.ValueExp (BoolValue True) } + +streaming :: KernelConstants -> VName -> SubExp -> InKernelGen () -> InKernelGen Imp.KernelCode +streaming constants chunksize bound m = do + bound' <- ImpGen.subExpToDimSize bound + let constants' = + constants { kernelStreamed = (chunksize, bound') : kernelStreamed constants } + ImpGen.subImpM_ (inKernelOperations constants') m + +compileKernelResult :: KernelConstants -> ImpGen.ValueDestination -> KernelResult + -> InKernelGen () + +compileKernelResult constants dest (ThreadsReturn OneResultPerGroup what) = do + i <- newVName "i" + + in_local_memory <- arrayInLocalMemory what + let me = Imp.var (kernelLocalThreadId constants) int32 + + if not in_local_memory then do + write_result <- + ImpGen.collect $ + ImpGen.copyDWIMDest dest [ImpGen.varIndex $ kernelGroupId constants] what [] + + who' <- ImpGen.compileSubExp $ intConst Int32 0 + ImpGen.emit $ + Imp.If (Imp.CmpOpExp (CmpEq int32) me who') write_result mempty + else do + -- 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). + -- + -- We do the reads/writes multidimensionally, but the loop is + -- single-dimensional. + ws <- mapM ImpGen.compileSubExp . arrayDims =<< subExpType what + -- Compute how many elements this thread is responsible for. + -- Formula: (w - ltid) / group_size (rounded up). + let w = product ws + ltid = ImpGen.varIndex (kernelLocalThreadId constants) + group_size = Imp.sizeToExp (kernelGroupSize constants) + to_write = (w - ltid) `quotRoundingUp` group_size + is = unflattenIndex ws $ ImpGen.varIndex i * group_size + ltid + + write_result <- + ImpGen.collect $ + ImpGen.copyDWIMDest dest (ImpGen.varIndex (kernelGroupId constants) : is) + what is + + ImpGen.emit $ Imp.For i Int32 to_write write_result + +compileKernelResult constants dest (ThreadsReturn AllThreads what) = + ImpGen.copyDWIMDest dest [ImpGen.varIndex $ kernelGlobalThreadId constants] what [] + +compileKernelResult constants dest (ThreadsReturn (ThreadsPerGroup limit) what) = do + write_result <- + ImpGen.collect $ + ImpGen.copyDWIMDest dest [ImpGen.varIndex $ kernelGroupId constants] what [] + + ImpGen.emit $ Imp.If (isActive limit) write_result mempty + +compileKernelResult constants dest (ThreadsReturn ThreadsInSpace what) = do + let is = map (ImpGen.varIndex . fst) $ kernelDimensions constants + write_result <- ImpGen.collect $ ImpGen.copyDWIMDest dest is what [] + ImpGen.emit $ Imp.If (kernelThreadActive constants) + write_result mempty + +compileKernelResult constants dest (ConcatReturns SplitContiguous _ per_thread_elems moffset what) = do + ImpGen.ArrayDestination (Just dest_loc) <- return dest + let dest_loc_offset = ImpGen.offsetArray dest_loc offset + dest' = ImpGen.ArrayDestination $ Just dest_loc_offset + ImpGen.copyDWIMDest dest' [] (Var what) [] + where offset = case moffset of + Nothing -> ImpGen.compileSubExpOfType int32 per_thread_elems * + ImpGen.varIndex (kernelGlobalThreadId constants) + Just se -> ImpGen.compileSubExpOfType int32 se + +compileKernelResult constants dest (ConcatReturns (SplitStrided stride) _ _ moffset what) = do + ImpGen.ArrayDestination (Just dest_loc) <- return dest + let dest_loc' = ImpGen.strideArray + (ImpGen.offsetArray dest_loc offset) $ + ImpGen.compileSubExpOfType int32 stride + dest' = ImpGen.ArrayDestination $ Just dest_loc' + ImpGen.copyDWIMDest dest' [] (Var what) [] + where offset = case moffset of + Nothing -> ImpGen.varIndex (kernelGlobalThreadId constants) + Just se -> ImpGen.compileSubExpOfType int32 se + +compileKernelResult constants dest (WriteReturn rws _arr dests) = do + rws' <- mapM ImpGen.compileSubExp rws + forM_ dests $ \(is, e) -> do + is' <- mapM ImpGen.compileSubExp is + let condInBounds0 = Imp.CmpOpExp (Imp.CmpSle Int32) $ + Imp.ValueExp (IntValue (Int32Value 0)) + condInBounds1 = Imp.CmpOpExp (Imp.CmpSlt Int32) + condInBounds i rw = Imp.BinOpExp LogAnd (condInBounds0 i) (condInBounds1 i rw) + write = foldl (Imp.BinOpExp LogAnd) (kernelThreadActive constants) $ + zipWith condInBounds is' rws' + actual_body' <- ImpGen.collect $ + ImpGen.copyDWIMDest dest (map (ImpGen.compileSubExpOfType int32) is) e [] + ImpGen.emit $ Imp.If write actual_body' Imp.Skip + +compileKernelResult _ _ KernelInPlaceReturn{} = + -- Already in its place... said it was a hack. + return () + +isActive :: [(VName, SubExp)] -> Imp.Exp +isActive limit = case actives of + [] -> Imp.ValueExp $ BoolValue True + x:xs -> foldl (Imp.BinOpExp LogAnd) x xs + where (is, ws) = unzip limit + actives = zipWith active is $ map (ImpGen.compileSubExpOfType Bool) ws + active i = Imp.CmpOpExp (CmpSlt Int32) (Imp.var i Bool) + +setSpaceIndices :: KernelSpace -> InKernelGen () +setSpaceIndices space = + case spaceStructure space of + FlatThreadSpace is_and_dims -> + flatSpaceWith gtid is_and_dims + NestedThreadSpace is_and_dims -> do + let (gtids, gdims, ltids, ldims) = unzip4 is_and_dims + gdims' <- mapM ImpGen.compileSubExp gdims + ldims' <- mapM ImpGen.compileSubExp ldims + let (gtid_es, ltid_es) = unzip $ unflattenNestedIndex gdims' ldims' gtid + forM_ (zip gtids gtid_es) $ \(i,e) -> + ImpGen.emit $ Imp.SetScalar i e + forM_ (zip ltids ltid_es) $ \(i,e) -> + ImpGen.emit $ Imp.SetScalar i e + where gtid = Imp.var (spaceGlobalId space) int32 + + flatSpaceWith base is_and_dims = do + let (is, dims) = unzip is_and_dims + dims' <- mapM ImpGen.compileSubExp dims + let index_expressions = unflattenIndex dims' base + forM_ (zip is index_expressions) $ \(i, x) -> + ImpGen.emit $ Imp.SetScalar i x + +unflattenNestedIndex :: IntegralExp num => [num] -> [num] -> num -> [(num,num)] +unflattenNestedIndex global_dims group_dims global_id = + zip global_is local_is + where num_groups_dims = zipWith quotRoundingUp global_dims group_dims + group_size = product group_dims + group_id = global_id `Futhark.Util.IntegralExp.quot` group_size + local_id = global_id `Futhark.Util.IntegralExp.rem` group_size + + group_is = unflattenIndex num_groups_dims group_id + local_is = unflattenIndex group_dims local_id + global_is = zipWith (+) local_is $ zipWith (*) group_is group_dims + +arrayInLocalMemory :: SubExp -> InKernelGen Bool +arrayInLocalMemory (Var name) = do + res <- ImpGen.lookupVar name + case res of + ImpGen.ArrayVar _ entry -> + (Space "local"==) . ImpGen.entryMemSpace <$> + ImpGen.lookupMemory (ImpGen.memLocationName (ImpGen.entryArrayLocation entry)) + _ -> return False +arrayInLocalMemory Constant{} = return False diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs @@ -0,0 +1,640 @@ +{-# LANGUAGE QuasiQuotes #-} +{-# 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 + ) + where + +import Control.Monad.State +import Control.Monad.Identity +import Control.Monad.Writer +import Control.Monad.Reader +import Data.Maybe +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem + +import qualified Language.C.Syntax as C +import qualified Language.C.Quote.OpenCL as C + +import Futhark.Error +import Futhark.Representation.AST.Attributes.Types (int32) +import qualified Futhark.CodeGen.OpenCL.Kernels as Kernels +import qualified Futhark.CodeGen.Backends.GenericC as GenericC +import Futhark.CodeGen.Backends.SimpleRepresentation +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.MonadFreshNames +import Futhark.Util (zEncodeString) +import Futhark.Util.Pretty (pretty, prettyOneLine) +import Futhark.Util.IntegralExp (quotRoundingUp) + +-- | Translate a kernels-program to an OpenCL-program. +kernelsToOpenCL :: ImpKernels.Program + -> Either InternalError ImpOpenCL.Program +kernelsToOpenCL (ImpKernels.Functions funs) = do + (prog', ToOpenCL extra_funs kernels requirements sizes) <- + runWriterT $ fmap Functions $ forM funs $ \(fname, fun) -> + (fname,) <$> runReaderT (traverse onHostOp fun) fname + let kernel_names = M.keys kernels + opencl_code = openClCode $ M.elems kernels + opencl_prelude = pretty $ genOpenClPrelude requirements + return $ ImpOpenCL.Program opencl_code opencl_prelude kernel_names + (S.toList $ kernelUsedTypes requirements) sizes $ + ImpOpenCL.Functions (M.toList extra_funs) <> prog' + +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 = fail $ "'" ++ s ++ "' is not an OpenCL kernel address space." + +type UsedFunctions = [(String,C.Func)] -- The ordering is important! + +data OpenClRequirements = + OpenClRequirements { kernelUsedTypes :: S.Set PrimType + , _kernelConstants :: [(VName, KernelConstExp)] + } + +instance Sem.Semigroup OpenClRequirements where + OpenClRequirements ts1 consts1 <> OpenClRequirements ts2 consts2 = + OpenClRequirements (ts1 <> ts2) (consts1 <> consts2) + +instance Monoid OpenClRequirements where + mempty = OpenClRequirements mempty mempty + mappend = (Sem.<>) + +data ToOpenCL = ToOpenCL { clExtraFuns :: M.Map Name ImpOpenCL.Function + , clKernels :: M.Map KernelName C.Func + , clRequirements :: OpenClRequirements + , clSizes :: M.Map VName (SizeClass, Name) + } + +instance Sem.Semigroup ToOpenCL where + ToOpenCL f1 k1 r1 sz1 <> ToOpenCL f2 k2 r2 sz2 = + ToOpenCL (f1<>f2) (k1<>k2) (r1<>r2) (sz1<>sz2) + +instance Monoid ToOpenCL where + mempty = ToOpenCL mempty mempty mempty mempty + mappend = (Sem.<>) + +type OnKernelM = ReaderT Name (WriterT ToOpenCL (Either InternalError)) + +onHostOp :: HostOp -> OnKernelM OpenCL +onHostOp (CallKernel k) = onKernel k +onHostOp (ImpKernels.GetSize v key size_class) = do + fname <- ask + tell mempty { clSizes = M.singleton key (size_class, fname) } + return $ ImpOpenCL.GetSize v key +onHostOp (ImpKernels.CmpSizeLe v key size_class x) = do + fname <- ask + tell mempty { clSizes = M.singleton key (size_class, fname) } + return $ ImpOpenCL.CmpSizeLe v key x +onHostOp (ImpKernels.GetSizeMax v size_class) = + return $ ImpOpenCL.GetSizeMax v size_class + +onKernel :: CallKernel -> OnKernelM OpenCL + +onKernel called@(Map kernel) = do + let (funbody, _) = + GenericC.runCompilerM (Functions []) inKernelOperations blankNameSource mempty $ do + size <- GenericC.compileExp $ mapKernelSize kernel + let check = [C.citem|if ($id:(mapKernelThreadNum kernel) >= $exp:size) return;|] + body <- GenericC.blockScope $ GenericC.compileCode $ mapKernelBody kernel + return $ check : body + + params = mapMaybe useAsParam $ mapKernelUses kernel + + tell mempty + { clExtraFuns = mempty + , clKernels = M.singleton (mapKernelName kernel) + [C.cfun|__kernel void $id:(mapKernelName kernel) ($params:params) { + const uint $id:(mapKernelThreadNum kernel) = get_global_id(0); + $items:funbody + }|] + , clRequirements = OpenClRequirements + (typesInKernel called) + (mapMaybe useAsConst $ mapKernelUses kernel) + } + + return $ LaunchKernel + (calledKernelName called) (kernelArgs called) kernel_size workgroup_size + + where (kernel_size, workgroup_size) = kernelAndWorkgroupSize called + +onKernel called@(AnyKernel kernel) = do + let (kernel_body, _) = + GenericC.runCompilerM (Functions []) inKernelOperations blankNameSource mempty $ + GenericC.blockScope $ GenericC.compileCode $ kernelBody kernel + + use_params = mapMaybe useAsParam $ kernelUses kernel + + (local_memory_params, local_memory_init) = + unzip $ + flip evalState (blankNameSource :: VNameSource) $ + mapM prepareLocalMemory $ kernelLocalMemory kernel + + params = catMaybes local_memory_params ++ use_params + + tell mempty { clExtraFuns = mempty + , clKernels = M.singleton name + [C.cfun|__kernel void $id:name ($params:params) { + $items:local_memory_init + $items:kernel_body + }|] + , clRequirements = OpenClRequirements + (typesInKernel called) + (mapMaybe useAsConst $ kernelUses kernel) + } + + return $ LaunchKernel + (calledKernelName called) (kernelArgs called) kernel_size workgroup_size + + where prepareLocalMemory (mem, Left _) = do + mem_aligned <- newVName $ baseString mem ++ "_aligned" + return (Just [C.cparam|__local volatile typename int64_t* $id:mem_aligned|], + [C.citem|__local volatile char* restrict $id:mem = $id:mem_aligned;|]) + prepareLocalMemory (mem, Right size) = do + let size' = compilePrimExp size + return (Nothing, + [C.citem|ALIGNED_LOCAL_MEMORY($id:mem, $exp:size');|]) + name = calledKernelName called + (kernel_size, workgroup_size) = kernelAndWorkgroupSize called + +onKernel (MapTranspose bt + destmem destoffset + srcmem srcoffset + num_arrays x_elems y_elems in_elems out_elems) = do + generateTransposeFunction bt + return $ HostCode $ Call [] (transposeName bt) + [MemArg destmem, ExpArg destoffset, + MemArg srcmem, ExpArg srcoffset, + ExpArg num_arrays, ExpArg x_elems, ExpArg y_elems, + ExpArg in_elems, ExpArg out_elems] + +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|] + _ -> GenericC.primTypeToCType bt + in Just [C.cparam|$ty:ctp $id:name|] +useAsParam (MemoryUse name _) = + Just [C.cparam|__global unsigned char *$id:name|] +useAsParam ConstUse{} = + Nothing + +useAsConst :: KernelUse -> Maybe (VName, KernelConstExp) +useAsConst (ConstUse v e) = Just (v,e) +useAsConst _ = Nothing + +openClCode :: [C.Func] -> String +openClCode kernels = + pretty [C.cunit|$edecls:funcs|] + where funcs = + [[C.cedecl|$func:kernel_func|] | + kernel_func <- kernels ] + +genOpenClPrelude :: OpenClRequirements -> [C.Definition] +genOpenClPrelude (OpenClRequirements ts consts) = + -- Clang-based OpenCL implementations need this for 'static' to work. + [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_clang_storage_class_specifiers : enable")|] : + [[C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_fp64 : enable")|] | 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; +} + +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; + +$esc:("#define ALIGNED_LOCAL_MEMORY(m,size) __local unsigned char m[size] __attribute__ ((align))") +|] ++ + cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ + (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else []) ++ + [ [C.cedecl|$esc:def|] | def <- map constToDefine consts ] + where uses_float64 = FloatType Float64 `S.member` ts + constToDefine (name, e) = + let e' = compilePrimExp e + in unwords ["#define", zEncodeString (pretty name), "("++prettyOneLine e'++")"] + +compilePrimExp :: PrimExp KernelConst -> C.Exp +compilePrimExp e = runIdentity $ GenericC.compilePrimExp compileKernelConst e + where compileKernelConst (SizeConst key) = return [C.cexp|$id:(pretty key)|] + +mapKernelName :: MapKernel -> String +mapKernelName k = "kernel_"++ mapKernelDesc k ++ "_" ++ + show (baseTag $ mapKernelThreadNum k) + +calledKernelName :: CallKernel -> String +calledKernelName (Map k) = + mapKernelName k +calledKernelName (AnyKernel k) = + kernelDesc k ++ "_kernel_" ++ show (baseTag $ kernelName k) +calledKernelName (MapTranspose bt _ _ _ _ _ _ _ _ _) = + transposeKernelName bt Kernels.TransposeNormal + +kernelArgs :: CallKernel -> [KernelArg] +kernelArgs (Map kernel) = + mapMaybe useToArg $ mapKernelUses kernel +kernelArgs (AnyKernel kernel) = + mapMaybe (fmap (SharedMemoryKArg . memSizeToExp) . localMemorySize) + (kernelLocalMemory kernel) ++ + mapMaybe useToArg (kernelUses kernel) + where localMemorySize (_, Left size) = Just size + localMemorySize (_, Right{}) = Nothing +kernelArgs (MapTranspose bt destmem destoffset srcmem srcoffset _ x_elems y_elems in_elems out_elems) = + [ MemKArg destmem + , ValueKArg destoffset int32 + , MemKArg srcmem + , ValueKArg srcoffset int32 + , ValueKArg x_elems int32 + , ValueKArg y_elems int32 + , ValueKArg in_elems int32 + , ValueKArg out_elems int32 + , SharedMemoryKArg shared_memory + ] + where shared_memory = + bytes $ (transposeBlockDim + 1) * transposeBlockDim * + LeafExp (SizeOf bt) (IntType Int32) + +kernelAndWorkgroupSize :: CallKernel -> ([Exp], [Exp]) +kernelAndWorkgroupSize (Map kernel) = + ([sizeToExp (mapKernelNumGroups kernel) * + sizeToExp (mapKernelGroupSize kernel)], + [sizeToExp $ mapKernelGroupSize kernel]) +kernelAndWorkgroupSize (AnyKernel kernel) = + ([sizeToExp (kernelNumGroups kernel) * + sizeToExp (kernelGroupSize kernel)], + [sizeToExp $ kernelGroupSize kernel]) +kernelAndWorkgroupSize (MapTranspose _ _ _ _ _ num_arrays x_elems y_elems _ _) = + transposeKernelAndGroupSize num_arrays x_elems y_elems + +--- Generating C + +inKernelOperations :: GenericC.Operations KernelOp UsedFunctions +inKernelOperations = GenericC.Operations + { GenericC.opsCompiler = kernelOps + , GenericC.opsMemoryType = kernelMemoryType + , GenericC.opsWriteScalar = GenericC.writeScalarPointerWithQuals pointerQuals + , GenericC.opsReadScalar = GenericC.readScalarPointerWithQuals pointerQuals + , GenericC.opsAllocate = cannotAllocate + , GenericC.opsDeallocate = cannotDeallocate + , GenericC.opsCopy = copyInKernel + , GenericC.opsStaticArray = noStaticArrays + , GenericC.opsFatMemory = False + } + where kernelOps :: GenericC.OpCompiler KernelOp UsedFunctions + kernelOps (GetGroupId v i) = + GenericC.stm [C.cstm|$id:v = get_group_id($int:i);|] + kernelOps (GetLocalId v i) = + GenericC.stm [C.cstm|$id:v = get_local_id($int:i);|] + kernelOps (GetLocalSize v i) = + GenericC.stm [C.cstm|$id:v = get_local_size($int:i);|] + kernelOps (GetGlobalId v i) = + GenericC.stm [C.cstm|$id:v = get_global_id($int:i);|] + kernelOps (GetGlobalSize v i) = + GenericC.stm [C.cstm|$id:v = get_global_size($int:i);|] + kernelOps (GetLockstepWidth v) = + GenericC.stm [C.cstm|$id:v = LOCKSTEP_WIDTH;|] + kernelOps Barrier = + GenericC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|] + kernelOps MemFence = + GenericC.stm [C.cstm|mem_fence(CLK_GLOBAL_MEM_FENCE);|] + kernelOps (Atomic aop) = atomicOps aop + + atomicOps (AtomicAdd old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_add((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|] + + atomicOps (AtomicSMax old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_max((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|] + + atomicOps (AtomicSMin old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_min((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|] + + atomicOps (AtomicUMax old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_max((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|] + + atomicOps (AtomicUMin old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_min((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|] + + atomicOps (AtomicAnd old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_and((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|] + + atomicOps (AtomicOr old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_or((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|] + + atomicOps (AtomicXor old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_xor((volatile __global unsigned int *)&$id:arr[$exp:ind'], (unsigned int)$exp:val');|] + + atomicOps (AtomicCmpXchg old arr ind cmp val) = do + ind' <- GenericC.compileExp $ innerExp ind + cmp' <- GenericC.compileExp cmp + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_cmpxchg((volatile __global int *)&$id:arr[$exp:ind'], $exp:cmp', $exp:val');|] + + atomicOps (AtomicXchg old arr ind val) = do + ind' <- GenericC.compileExp $ innerExp ind + val' <- GenericC.compileExp val + GenericC.stm [C.cstm|$id:old = atomic_xchg((volatile __global int *)&$id:arr[$exp:ind'], $exp:val');|] + + cannotAllocate :: GenericC.Allocate KernelOp UsedFunctions + cannotAllocate _ = + fail "Cannot allocate memory in kernel" + + cannotDeallocate :: GenericC.Deallocate KernelOp UsedFunctions + cannotDeallocate _ _ = + fail "Cannot deallocate memory in kernel" + + copyInKernel :: GenericC.Copy KernelOp UsedFunctions + copyInKernel _ _ _ _ _ _ _ = + fail "Cannot bulk copy in kernel." + + noStaticArrays :: GenericC.StaticArray KernelOp UsedFunctions + noStaticArrays _ _ _ _ = + fail "Cannot create static array in kernel." + + kernelMemoryType space = do + quals <- pointerQuals space + return [C.cty|$tyquals:quals $ty:defaultMemBlockType|] + +--- Handling transpositions + +transposeKernelName :: PrimType -> Kernels.TransposeType -> String +transposeKernelName bt Kernels.TransposeNormal = + "fut_kernel_map_transpose_" ++ pretty bt +transposeKernelName bt Kernels.TransposeLowWidth = + "fut_kernel_map_transpose_lowwidth_" ++ pretty bt +transposeKernelName bt Kernels.TransposeLowHeight = + "fut_kernel_map_transpose_lowheight_" ++ pretty bt +transposeKernelName bt Kernels.TransposeSmall = + "fut_kernel_map_transpose_small_" ++ pretty bt + +transposeName :: PrimType -> Name +transposeName bt = nameFromString $ "map_transpose_opencl_" ++ pretty bt + +generateTransposeFunction :: PrimType -> OnKernelM () +generateTransposeFunction bt = + -- We have special functions to handle transposing an input array with low + -- width or low height, as this would cause very few threads to be active. See + -- comment in Futhark.CodeGen.OpenCL.OpenCL.Kernels.hs for more details. + + tell mempty + { clExtraFuns = M.singleton (transposeName bt) $ + ImpOpenCL.Function False [] params transpose_code [] [] + , clKernels = M.fromList $ + map (\tt -> let name = transposeKernelName bt tt + in (name, Kernels.mapTranspose name bt' tt)) + [Kernels.TransposeNormal, Kernels.TransposeLowWidth, + Kernels.TransposeLowHeight, Kernels.TransposeSmall] + + , clRequirements = mempty + } + + where bt' = GenericC.primTypeToCType bt + space = ImpOpenCL.Space "device" + memparam s i = MemParam (VName (nameFromString s) i) space + intparam s i = ScalarParam (VName (nameFromString s) i) $ IntType Int32 + + params = [destmem_p, destoffset_p, srcmem_p, srcoffset_p, + num_arrays_p, x_p, y_p, in_p, out_p] + + [destmem_p, destoffset_p, srcmem_p, srcoffset_p, + num_arrays_p, x_p, y_p, in_p, out_p, + muly, new_height, mulx, new_width] = + zipWith ($) [memparam "destmem", + intparam "destoffset", + memparam "srcmem", + intparam "srcoffset", + intparam "num_arrays", + intparam "x_elems", + intparam "y_elems", + intparam "in_elems", + intparam "out_elems", + -- The following is only used for low width/height + -- transpose kernels + intparam "muly", + intparam "new_height", + intparam "mulx", + intparam "new_width" + ] + [0..] + + asExp param = + ImpOpenCL.LeafExp (ImpOpenCL.ScalarVar (paramName param)) (IntType Int32) + + asArg (MemParam name _) = + MemKArg name + asArg (ScalarParam name t) = + ValueKArg (ImpOpenCL.LeafExp (ImpOpenCL.ScalarVar name) t) t + + normal_kernel_args = + map asArg [destmem_p, destoffset_p, srcmem_p, srcoffset_p, + x_p, y_p, in_p, out_p] ++ + [SharedMemoryKArg shared_memory] + + lowwidth_kernel_args = + map asArg [destmem_p, destoffset_p, srcmem_p, srcoffset_p, + x_p, y_p, in_p, out_p, muly] ++ + [SharedMemoryKArg shared_memory] + + lowheight_kernel_args = + map asArg [destmem_p, destoffset_p, srcmem_p, srcoffset_p, + x_p, y_p, in_p, out_p, mulx] ++ + [SharedMemoryKArg shared_memory] + + shared_memory = + bytes $ (transposeBlockDim + 1) * transposeBlockDim * + LeafExp (SizeOf bt) (IntType Int32) + + transposeBlockDimDivTwo = BinOpExp (SQuot Int32) transposeBlockDim 2 + + should_use_lowwidth = BinOpExp LogAnd + (CmpOpExp (CmpSle Int32) (asExp x_p) transposeBlockDimDivTwo) + (CmpOpExp (CmpSlt Int32) transposeBlockDim (asExp y_p)) + + should_use_lowheight = BinOpExp LogAnd + (CmpOpExp (CmpSle Int32) (asExp y_p) transposeBlockDimDivTwo) + (CmpOpExp (CmpSlt Int32) transposeBlockDim (asExp x_p)) + + should_use_small = BinOpExp LogAnd + (CmpOpExp (CmpSle Int32) (asExp x_p) transposeBlockDimDivTwo) + (CmpOpExp (CmpSle Int32) (asExp y_p) transposeBlockDimDivTwo) + + -- When an input array has either width==1 or height==1, performing a + -- transpose will be the same as performing a copy. If 'input_size' or + -- 'output_size' is not equal to width*height, then this trick will not + -- work when there are more than one array to process, as it is a per + -- array limit. We could copy each array individually, but currently we + -- do not. + can_use_copy = + let in_out_eq = CmpOpExp (CmpEq $ IntType Int32) (asExp in_p) (asExp out_p) + onearr = CmpOpExp (CmpEq $ IntType Int32) (asExp num_arrays_p) 1 + noprob_widthheight = CmpOpExp (CmpEq $ IntType Int32) + (asExp x_p * asExp y_p) + (asExp in_p) + height_is_one = CmpOpExp (CmpEq $ IntType Int32) (asExp y_p) 1 + width_is_one = CmpOpExp (CmpEq $ IntType Int32) (asExp x_p) 1 + in BinOpExp LogAnd + in_out_eq + (BinOpExp LogAnd + (BinOpExp LogOr onearr noprob_widthheight) + (BinOpExp LogOr width_is_one height_is_one)) + + input_is_empty = CmpOpExp (CmpEq $ IntType Int32) + (asExp num_arrays_p * asExp x_p * asExp y_p) 0 + + transpose_code = + ImpOpenCL.If input_is_empty mempty + (ImpOpenCL.If can_use_copy + copy_code + (ImpOpenCL.If should_use_lowwidth + lowwidth_transpose_code + (ImpOpenCL.If should_use_lowheight + lowheight_transpose_code + (ImpOpenCL.If should_use_small + small_transpose_code + normal_transpose_code)))) + + copy_code = + let num_bytes = + asExp in_p * ImpOpenCL.LeafExp (ImpOpenCL.SizeOf bt) (IntType Int32) + in ImpOpenCL.Copy + (paramName destmem_p) (Count $ asExp destoffset_p) space + (paramName srcmem_p) (Count $ asExp srcoffset_p) space + (Count num_bytes) + + normal_transpose_code = + let (kernel_size, workgroup_size) = + transposeKernelAndGroupSize (asExp num_arrays_p) (asExp x_p) (asExp y_p) + in ImpOpenCL.Op $ LaunchKernel + (transposeKernelName bt Kernels.TransposeNormal) normal_kernel_args kernel_size workgroup_size + + small_transpose_code = + let group_size = (transposeBlockDim * transposeBlockDim) + kernel_size = (asExp num_arrays_p * asExp x_p * asExp y_p) `roundUpTo` + group_size + in ImpOpenCL.Op $ LaunchKernel + (transposeKernelName bt Kernels.TransposeSmall) + (map asArg [destmem_p, destoffset_p, srcmem_p, srcoffset_p, + num_arrays_p, x_p, y_p, in_p, out_p]) + [kernel_size] [group_size] + + lowwidth_transpose_code = + let set_muly = DeclareScalar (paramName muly) (IntType Int32) + :>>: SetScalar (paramName muly) (BinOpExp (SQuot Int32) transposeBlockDim (asExp x_p)) + set_new_height = DeclareScalar (paramName new_height) (IntType Int32) + :>>: SetScalar (paramName new_height) (asExp y_p `quotRoundingUp` asExp muly) + (kernel_size, workgroup_size) = + transposeKernelAndGroupSize (asExp num_arrays_p) (asExp x_p) (asExp new_height) + launch = ImpOpenCL.Op $ LaunchKernel + (transposeKernelName bt Kernels.TransposeLowWidth) lowwidth_kernel_args kernel_size workgroup_size + in set_muly :>>: set_new_height :>>: launch + + lowheight_transpose_code = + let set_mulx = DeclareScalar (paramName mulx) (IntType Int32) + :>>: SetScalar (paramName mulx) (BinOpExp (SQuot Int32) transposeBlockDim (asExp y_p)) + set_new_width = DeclareScalar (paramName new_width) (IntType Int32) + :>>: SetScalar (paramName new_width) (asExp x_p `quotRoundingUp` asExp mulx) + (kernel_size, workgroup_size) = + transposeKernelAndGroupSize (asExp num_arrays_p) (asExp new_width) (asExp y_p) + launch = ImpOpenCL.Op $ LaunchKernel + (transposeKernelName bt Kernels.TransposeLowHeight) lowheight_kernel_args kernel_size workgroup_size + in set_mulx :>>: set_new_width :>>: launch + +transposeKernelAndGroupSize :: ImpOpenCL.Exp -> ImpOpenCL.Exp -> ImpOpenCL.Exp + -> ([ImpOpenCL.Exp], [ImpOpenCL.Exp]) +transposeKernelAndGroupSize num_arrays x_elems y_elems = + ([x_elems `roundUpTo` transposeBlockDim , + y_elems `roundUpTo` transposeBlockDim, + num_arrays], + [transposeBlockDim, transposeBlockDim, 1]) + +roundUpTo :: ImpOpenCL.Exp -> ImpOpenCL.Exp -> ImpOpenCL.Exp +roundUpTo x y = x + ((y - (x `impRem` y)) `impRem` y) + where impRem = BinOpExp $ SRem Int32 + +--- Checking requirements + +useToArg :: KernelUse -> Maybe KernelArg +useToArg (MemoryUse mem _) = Just $ MemKArg mem +useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt +useToArg ConstUse{} = Nothing + +typesInKernel :: CallKernel -> S.Set PrimType +typesInKernel (Map kernel) = typesInCode $ mapKernelBody kernel +typesInKernel (AnyKernel kernel) = typesInCode $ kernelBody kernel +typesInKernel MapTranspose{} = mempty + +typesInCode :: ImpKernels.KernelCode -> S.Set PrimType +typesInCode Skip = mempty +typesInCode (c1 :>>: c2) = typesInCode c1 <> typesInCode c2 +typesInCode (For _ it e c) = IntType it `S.insert` typesInExp e <> typesInCode c +typesInCode (While e c) = typesInExp e <> typesInCode c +typesInCode DeclareMem{} = mempty +typesInCode (DeclareScalar _ t) = S.singleton t +typesInCode (DeclareArray _ _ t _) = S.singleton t +typesInCode (Allocate _ (Count e) _) = typesInExp e +typesInCode Free{} = mempty +typesInCode (Copy _ (Count e1) _ _ (Count e2) _ (Count e3)) = + typesInExp e1 <> typesInExp e2 <> typesInExp e3 +typesInCode (Write _ (Count 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 e c1 c2) = + typesInExp e <> typesInCode c1 <> typesInCode c2 +typesInCode (Assert e _ _) = typesInExp e +typesInCode (Comment _ c) = typesInCode c +typesInCode (DebugPrint _ _ e) = typesInExp e +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 e) t _ _) _) = S.singleton t <> typesInExp e +typesInExp (LeafExp ScalarVar{} _) = mempty +typesInExp (LeafExp (SizeOf t) _) = S.singleton t diff --git a/src/Futhark/CodeGen/ImpGen/OpenCL.hs b/src/Futhark/CodeGen/ImpGen/OpenCL.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/OpenCL.hs @@ -0,0 +1,13 @@ +module Futhark.CodeGen.ImpGen.OpenCL + ( compileProg + ) where + +import Futhark.Error +import Futhark.Representation.ExplicitMemory +import qualified Futhark.CodeGen.ImpCode.OpenCL as OpenCL +import qualified Futhark.CodeGen.ImpGen.Kernels as ImpGenKernels +import Futhark.CodeGen.ImpGen.Kernels.ToOpenCL +import Futhark.MonadFreshNames + +compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError OpenCL.Program) +compileProg prog = either Left kernelsToOpenCL <$> ImpGenKernels.compileProg prog diff --git a/src/Futhark/CodeGen/ImpGen/Sequential.hs b/src/Futhark/CodeGen/ImpGen/Sequential.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Sequential.hs @@ -0,0 +1,20 @@ +{-# LANGUAGE TypeFamilies #-} +module Futhark.CodeGen.ImpGen.Sequential + ( compileProg + ) + where + +import Futhark.Error +import qualified Futhark.CodeGen.ImpCode.Sequential as Imp +import qualified Futhark.CodeGen.ImpGen as ImpGen +import Futhark.Representation.ExplicitMemory +import Futhark.MonadFreshNames + +compileProg :: MonadFreshNames m => Prog ExplicitMemory -> m (Either InternalError Imp.Program) +compileProg = ImpGen.compileProg ops Imp.DefaultSpace + where ops = ImpGen.defaultOperations opCompiler + opCompiler :: ImpGen.OpCompiler ExplicitMemory Imp.Sequential + opCompiler dest (Alloc e space) = + ImpGen.compileAlloc dest e space + opCompiler _ (Inner _) = + compilerBugS "Cannot handle kernels in sequential code generator." diff --git a/src/Futhark/CodeGen/OpenCL/Kernels.hs b/src/Futhark/CodeGen/OpenCL/Kernels.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/OpenCL/Kernels.hs @@ -0,0 +1,208 @@ +{-# LANGUAGE QuasiQuotes #-} +module Futhark.CodeGen.OpenCL.Kernels + ( SizeHeuristic (..) + , DeviceType (..) + , WhichSize (..) + , HeuristicValue (..) + , sizeHeuristicsTable + + , mapTranspose + , TransposeType(..) + ) + where + +import qualified Language.C.Syntax as C +import qualified Language.C.Quote.OpenCL as C + +-- Some OpenCL platforms have a SIMD/warp/wavefront-based execution +-- model that execute groups of threads in lockstep, permitting us to +-- perform cross-thread synchronisation within each such group without +-- the use of barriers. Unfortunately, there seems to be no reliable +-- way to query these sizes at runtime. Instead, we use this table to +-- figure out which size we should use for a specific platform and +-- device. If nothing matches here, the wave size should be set to +-- one. +-- +-- We also use this to select reasonable default group sizes and group +-- counts. + +-- | The type of OpenCL device that this heuristic applies to. +data DeviceType = DeviceCPU | DeviceGPU + +-- | The value supplies by a heuristic can be a constant, or inferred +-- from some device information. +data HeuristicValue = HeuristicConst Int + | HeuristicDeviceInfo String + +-- | A size that can be assigned a default. +data WhichSize = LockstepWidth | NumGroups | GroupSize | TileSize + +-- | A heuristic for setting the default value for something. +data SizeHeuristic = + SizeHeuristic { platformName :: String + , deviceType :: DeviceType + , heuristicSize :: WhichSize + , heuristicValue :: HeuristicValue + } + +-- | All of our heuristics. +sizeHeuristicsTable :: [SizeHeuristic] +sizeHeuristicsTable = + [ SizeHeuristic "NVIDIA CUDA" DeviceGPU LockstepWidth $ HeuristicConst 32 + , SizeHeuristic "AMD Accelerated Parallel Processing" DeviceGPU LockstepWidth $ HeuristicConst 64 + , SizeHeuristic "" DeviceGPU LockstepWidth $ HeuristicConst 1 + , SizeHeuristic "" DeviceGPU NumGroups $ HeuristicConst 128 + , SizeHeuristic "" DeviceGPU GroupSize $ HeuristicConst 256 + , SizeHeuristic "" DeviceGPU TileSize $ HeuristicConst 32 + + , SizeHeuristic "" DeviceCPU LockstepWidth $ HeuristicConst 1 + , SizeHeuristic "" DeviceCPU NumGroups $ HeuristicDeviceInfo "MAX_COMPUTE_UNITS" + , SizeHeuristic "" DeviceCPU GroupSize $ HeuristicConst 32 + , SizeHeuristic "" DeviceCPU TileSize $ HeuristicConst 4 + ] + +-- | Which form of transposition to generate code for. +data TransposeType = TransposeNormal + | TransposeLowWidth + | TransposeLowHeight + | TransposeSmall -- ^ For small arrays that do not + -- benefit from coalescing. + deriving (Eq, Ord, Show) + +-- | @mapTranspose name elem_type transpose_type@ Generate a transpose kernel +-- with requested @name@ for elements of type @elem_type@. There are special +-- support to handle input arrays with low width or low height, which can be +-- indicated by @transpose_type@. +-- +-- 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. +mapTranspose :: C.ToIdent a => a -> C.Type -> TransposeType -> C.Func +mapTranspose kernel_name elem_type transpose_type = + case transpose_type of + TransposeNormal -> + bigKernel [] + [C.cexp|get_global_id(0)|] + [C.cexp|get_global_id(1)|] + [C.cexp|get_group_id(1) * FUT_BLOCK_DIM + get_local_id(0)|] + [C.cexp|get_group_id(0) * FUT_BLOCK_DIM + get_local_id(1)|] + TransposeLowWidth -> + bigKernel [C.cparams|uint muly|] + [C.cexp|get_group_id(0) * FUT_BLOCK_DIM + (get_local_id(0) / muly)|] + [C.cexp|get_group_id(1) * FUT_BLOCK_DIM * muly + + get_local_id(1) + + (get_local_id(0) % muly) * FUT_BLOCK_DIM + |] + [C.cexp|get_group_id(1) * FUT_BLOCK_DIM * muly + + get_local_id(0) + + (get_local_id(1) % muly) * FUT_BLOCK_DIM|] + [C.cexp|get_group_id(0) * FUT_BLOCK_DIM + (get_local_id(1) / muly)|] + TransposeLowHeight -> + bigKernel [C.cparams|uint mulx|] + [C.cexp|get_group_id(0) * FUT_BLOCK_DIM * mulx + + get_local_id(0) + + (get_local_id(1) % mulx) * FUT_BLOCK_DIM + |] + [C.cexp|get_group_id(1) * FUT_BLOCK_DIM + (get_local_id(1) / mulx)|] + [C.cexp|get_group_id(1) * FUT_BLOCK_DIM + (get_local_id(0) / mulx)|] + [C.cexp|get_group_id(0) * FUT_BLOCK_DIM * mulx + + get_local_id(1) + + (get_local_id(0) % mulx) * FUT_BLOCK_DIM + |] + TransposeSmall -> + smallKernel + where + bigKernel extraparams x_in_index y_in_index x_out_index y_out_index = + [C.cfun| + // This kernel is optimized to ensure all global reads and writes are coalesced, + // and to avoid bank conflicts in shared memory. The shared memory array is sized + // to (BLOCK_DIM+1)*BLOCK_DIM. This pads each row of the 2D block in shared memory + // so that bank conflicts do not occur when threads address the array column-wise. + // + // Note that input_size/output_size may not equal width*height if we are dealing with + // a truncated array - this happens sometimes for coalescing optimisations. + __kernel void $id:kernel_name($params:params) { + uint x_index; + uint y_index; + uint our_array_offset; + + // Adjust the input and output arrays with the basic offset. + odata += odata_offset/sizeof($ty:elem_type); + idata += idata_offset/sizeof($ty:elem_type); + + // Adjust the input and output arrays for the third dimension. + our_array_offset = get_global_id(2) * width * height; + odata += our_array_offset; + idata += our_array_offset; + + // read the matrix tile into shared memory + x_index = $exp:x_in_index; + y_index = $exp:y_in_index; + + uint index_in = y_index * width + x_index; + + if(x_index < width && y_index < height && index_in < input_size) + { + block[get_local_id(1)*(FUT_BLOCK_DIM+1)+get_local_id(0)] = idata[index_in]; + } + + barrier(CLK_LOCAL_MEM_FENCE); + + // Scatter the transposed matrix tile to global memory. + x_index = $exp:x_out_index; + y_index = $exp:y_out_index; + + uint index_out = y_index * height + x_index; + + if(x_index < height && y_index < width && index_out < output_size) + { + odata[index_out] = block[get_local_id(0)*(FUT_BLOCK_DIM+1)+get_local_id(1)]; + } + }|] + where params = [C.cparams|__global $ty:elem_type *odata, + uint odata_offset, + __global $ty:elem_type *idata, + uint idata_offset, + uint width, + uint height, + uint input_size, + uint output_size|] ++ extraparams ++ + [C.cparams|__local $ty:elem_type* block|] + + smallKernel = + [C.cfun| + __kernel void $id:kernel_name(__global $ty:elem_type *odata, + uint odata_offset, + __global $ty:elem_type *idata, + uint idata_offset, + uint num_arrays, + uint width, + uint height, + uint input_size, + uint output_size) { + uint our_array_offset = get_global_id(0) / (height*width) * (height*width); + uint x_index = get_global_id(0) % (height*width) / height; + uint y_index = get_global_id(0) % height; + + // Adjust the input and output arrays with the basic offset. + odata += odata_offset/sizeof($ty:elem_type); + idata += idata_offset/sizeof($ty:elem_type); + + // Adjust the input and output arrays. + odata += our_array_offset; + idata += our_array_offset; + + // Read and write the element. + uint index_in = y_index * width + x_index; + uint index_out = x_index * height + y_index; + if (get_global_id(0) < input_size) { + odata[index_out] = idata[index_in]; + } + }|] diff --git a/src/Futhark/CodeGen/SetDefaultSpace.hs b/src/Futhark/CodeGen/SetDefaultSpace.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/SetDefaultSpace.hs @@ -0,0 +1,101 @@ +module Futhark.CodeGen.SetDefaultSpace + ( setDefaultSpace + ) + where + +import Futhark.CodeGen.ImpCode + +-- | Set all uses of 'DefaultSpace' in the given functions to another memory space. +setDefaultSpace :: Space -> Functions op -> Functions op +setDefaultSpace space (Functions fundecs) = + Functions [ (fname, setFunctionSpace space func) + | (fname, func) <- fundecs ] + +setFunctionSpace :: Space -> Function op -> Function op +setFunctionSpace space (Function entry outputs inputs body results args) = + Function entry + (map (setParamSpace space) outputs) + (map (setParamSpace space) inputs) + (setBodySpace space body) + (map (setExtValueSpace space) results) + (map (setExtValueSpace space) args) + +setParamSpace :: Space -> Param -> Param +setParamSpace space (MemParam name DefaultSpace) = + MemParam name space +setParamSpace _ param = + param + +setExtValueSpace :: Space -> ExternalValue -> ExternalValue +setExtValueSpace space (OpaqueValue desc vs) = + OpaqueValue desc $ map (setValueSpace space) vs +setExtValueSpace space (TransparentValue v) = + TransparentValue $ setValueSpace space v + +setValueSpace :: Space -> ValueDesc -> ValueDesc +setValueSpace space (ArrayValue mem memsize _ bt ept shape) = + ArrayValue mem memsize space bt ept shape +setValueSpace _ (ScalarValue bt ept v) = + ScalarValue bt ept v + +setBodySpace :: Space -> Code op -> Code op +setBodySpace space (Allocate v e old_space) = + Allocate v (setCountSpace space e) $ setSpace space old_space +setBodySpace space (Free v old_space) = + Free v $ setSpace space old_space +setBodySpace space (DeclareMem name old_space) = + DeclareMem name $ setSpace space old_space +setBodySpace space (DeclareArray name _ t vs) = + DeclareArray name space t vs +setBodySpace space (Copy dest dest_offset dest_space src src_offset src_space n) = + Copy + dest (setCountSpace space dest_offset) dest_space' + src (setCountSpace space src_offset) src_space' $ + setCountSpace space n + where dest_space' = setSpace space dest_space + src_space' = setSpace space src_space +setBodySpace space (Write dest dest_offset bt dest_space vol e) = + Write dest (setCountSpace space dest_offset) bt (setSpace space dest_space) + vol (setExpSpace space e) +setBodySpace space (c1 :>>: c2) = + setBodySpace space c1 :>>: setBodySpace space c2 +setBodySpace space (For i it e body) = + For i it (setExpSpace space e) $ setBodySpace space body +setBodySpace space (While e body) = + While (setExpSpace space e) $ setBodySpace space body +setBodySpace space (If e c1 c2) = + If (setExpSpace space e) (setBodySpace space c1) (setBodySpace space c2) +setBodySpace space (Comment s c) = + Comment s $ setBodySpace space c +setBodySpace _ Skip = + Skip +setBodySpace _ (DeclareScalar name bt) = + DeclareScalar name bt +setBodySpace space (SetScalar name e) = + SetScalar name $ setExpSpace space e +setBodySpace space (SetMem to from old_space) = + SetMem to from $ setSpace space old_space +setBodySpace space (Call dests fname args) = + Call dests fname $ map setArgSpace args + where setArgSpace (MemArg m) = MemArg m + setArgSpace (ExpArg e) = ExpArg $ setExpSpace space e +setBodySpace space (Assert e msg loc) = + Assert (setExpSpace space e) msg loc +setBodySpace space (DebugPrint s t e) = + DebugPrint s t (setExpSpace space e) +setBodySpace _ (Op op) = + Op op + +setCountSpace :: Space -> Count a -> Count a +setCountSpace space (Count e) = + Count $ setExpSpace space e + +setExpSpace :: Space -> Exp -> Exp +setExpSpace space = fmap setLeafSpace + where setLeafSpace (Index mem i bt DefaultSpace vol) = + Index mem i bt space vol + setLeafSpace e = e + +setSpace :: Space -> Space -> Space +setSpace space DefaultSpace = space +setSpace _ space = space diff --git a/src/Futhark/Compiler.hs b/src/Futhark/Compiler.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Compiler.hs @@ -0,0 +1,158 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Compiler + ( + runPipelineOnProgram + , runCompilerOnProgram + + , FutharkConfig (..) + , newFutharkConfig + , dumpError + , reportingIOErrors + + , module Futhark.Compiler.Program + , readProgram + , readLibrary + ) +where + +import Data.Semigroup ((<>)) +import Control.Exception +import Control.Monad +import Control.Monad.Reader +import Control.Monad.Except +import System.Exit (exitWith, ExitCode(..)) +import System.IO +import qualified Data.Text as T +import qualified Data.Text.IO as T + +import Futhark.Internalise +import Futhark.Pipeline +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import qualified Futhark.Representation.SOACS as I +import qualified Futhark.TypeCheck as I +import Futhark.Compiler.Program +import qualified Language.Futhark as E +import Futhark.Util.Log + +data FutharkConfig = FutharkConfig + { futharkVerbose :: (Verbosity, Maybe FilePath) + , futharkWarn :: Bool -- ^ Warn if True. + , futharkWerror :: Bool -- ^ If true, error on any warnings. + , futharkSafe :: Bool -- ^ If True, ignore @unsafe@. + } + +newFutharkConfig :: FutharkConfig +newFutharkConfig = FutharkConfig { futharkVerbose = (NotVerbose, Nothing) + , futharkWarn = True + , futharkWerror = False + , futharkSafe = False + } + +dumpError :: FutharkConfig -> CompilerError -> IO () +dumpError config err = + case err of + ExternalError s -> do + T.hPutStrLn stderr s + T.hPutStrLn stderr "If you find this error message confusing, uninformative, or wrong, please open an issue at https://github.com/diku-dk/futhark/issues." + InternalError s info CompilerBug -> do + T.hPutStrLn stderr "Internal compiler error." + T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues." + report s info + InternalError s info CompilerLimitation -> do + T.hPutStrLn stderr "Known compiler limitation encountered. Sorry." + T.hPutStrLn stderr "Revise your program or try a different Futhark compiler." + report s info + where report s info = do + T.hPutStrLn stderr s + when (fst (futharkVerbose config) > NotVerbose) $ + maybe (T.hPutStr stderr) T.writeFile + (snd (futharkVerbose config)) $ info <> "\n" + +-- | Catch all IO exceptions and print a better error message if they +-- happen. Use this at the top-level of all Futhark compiler +-- frontends. +reportingIOErrors :: IO () -> IO () +reportingIOErrors = flip catches [Handler onExit, Handler onError] + where onExit :: ExitCode -> IO () + onExit = throwIO + onError :: SomeException -> IO () + onError e + | Just UserInterrupt <- asyncExceptionFromException e = + return () -- This corresponds to CTRL-C, which is not an error. + | otherwise = do + T.hPutStrLn stderr "Internal compiler error (unhandled IO exception)." + T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues" + T.hPutStrLn stderr $ T.pack $ show e + exitWith $ ExitFailure 1 + +runCompilerOnProgram :: FutharkConfig + -> Pipeline I.SOACS lore + -> Action lore + -> FilePath + -> IO () +runCompilerOnProgram config pipeline action file = do + res <- runFutharkM compile $ fst $ futharkVerbose config + case res of + Left err -> liftIO $ do + dumpError config err + exitWith $ ExitFailure 2 + Right () -> + return () + where compile = do + prog <- runPipelineOnProgram config pipeline file + when ((>NotVerbose) . fst $ futharkVerbose config) $ + logMsg $ "Running action " ++ actionName action + actionProcedure action prog + when ((>NotVerbose) . fst $ futharkVerbose config) $ + logMsg ("Done." :: String) + +runPipelineOnProgram :: FutharkConfig + -> Pipeline I.SOACS tolore + -> FilePath + -> FutharkM (Prog tolore) +runPipelineOnProgram config pipeline file = do + when (pipelineVerbose pipeline_config) $ + logMsg ("Reading and type-checking source program" :: String) + (ws, prog_imports, namesrc) <- readProgram file + + when (futharkWarn config) $ do + liftIO $ hPutStr stderr $ show ws + when (futharkWerror config && ws /= mempty) $ + externalErrorS "Treating above warnings as errors due to --Werror." + + putNameSource namesrc + when (pipelineVerbose pipeline_config) $ + logMsg ("Internalising program" :: String) + res <- internaliseProg (futharkSafe config) prog_imports + case res of + Left err -> + internalErrorS ("During internalisation: " <> pretty err) $ E.Prog Nothing $ + concatMap (E.progDecs . fileProg . snd) prog_imports + Right int_prog -> do + when (pipelineVerbose pipeline_config) $ + logMsg ("Type-checking internalised program" :: String) + typeCheckInternalProgram int_prog + runPasses pipeline pipeline_config int_prog + where pipeline_config = + PipelineConfig { pipelineVerbose = fst (futharkVerbose config) > NotVerbose + , pipelineValidate = True + } + +typeCheckInternalProgram :: I.Prog -> FutharkM () +typeCheckInternalProgram prog = + case I.checkProg prog of + Left err -> internalErrorS ("After internalisation:\n" ++ show err) (Just prog) + Right () -> return () + +-- | Read and type-check a Futhark program, including all imports. +readProgram :: (MonadError CompilerError m, MonadIO m) => + FilePath -> m (Warnings, Imports, VNameSource) +readProgram = readLibrary . pure + +-- | Read and type-check a collection of Futhark files, including all +-- imports. +readLibrary :: (MonadError CompilerError m, MonadIO m) => + [FilePath] -> m (Warnings, Imports, VNameSource) +readLibrary = readLibraryWithBasis emptyBasis diff --git a/src/Futhark/Compiler/CLI.hs b/src/Futhark/Compiler/CLI.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Compiler/CLI.hs @@ -0,0 +1,127 @@ +-- | Convenient common interface for command line Futhark compilers. +-- Using this module ensures that all compilers take the same options. +-- A small amount of flexibility is provided for backend-specific +-- options. +module Futhark.Compiler.CLI + ( compilerMain + , CompilerOption + , CompilerMode(..) + ) +where + +import Control.Monad +import Data.Maybe +import System.FilePath +import System.Console.GetOpt + +import Futhark.Pipeline +import Futhark.Compiler +import Futhark.Representation.AST (Prog) +import Futhark.Representation.SOACS (SOACS) +import Futhark.Util.Options + +-- | Run a parameterised Futhark compiler, where @cfg@ is a user-given +-- configuration type. Call this from @main@. +compilerMain :: cfg -- ^ Initial configuration. + -> [CompilerOption cfg] -- ^ Options that affect the configuration. + -> String -- ^ The short action name (e.g. "compile to C"). + -> String -- ^ The longer action description. + -> Pipeline SOACS lore -- ^ The pipeline to use. + -> (cfg -> CompilerMode -> FilePath -> Prog lore -> FutharkM ()) + -- ^ The action to take on the result of the pipeline. + -> IO () +compilerMain cfg cfg_opts name desc pipeline doIt = + reportingIOErrors $ + mainWithOptions (newCompilerConfig cfg) (commandLineOptions ++ map wrapOption cfg_opts) + "options... program" inspectNonOptions + where inspectNonOptions [file] config = Just $ compile config file + inspectNonOptions _ _ = Nothing + + compile config filepath = + runCompilerOnProgram (futharkConfig config) + pipeline (action config filepath) filepath + + action config filepath = + Action { actionName = name + , actionDescription = desc + , actionProcedure = + doIt (compilerConfig config) (compilerMode config) $ + outputFilePath filepath config + } + +-- | An option that modifies the configuration of type @cfg@. +type CompilerOption cfg = OptDescr (Either (IO ()) (cfg -> cfg)) + +type CoreCompilerOption cfg = OptDescr (Either + (IO ()) + (CompilerConfig cfg -> CompilerConfig cfg)) + +commandLineOptions :: [CoreCompilerOption cfg] +commandLineOptions = + [ Option "o" [] + (ReqArg (\filename -> Right $ \config -> config { compilerOutput = Just filename }) + "FILE") + "Name of the compiled binary." + , Option "v" ["verbose"] + (OptArg (Right . incVerbosity) "FILE") + "Print verbose output on standard error; wrong program to FILE." + , Option [] ["library"] + (NoArg $ Right $ \config -> config { compilerMode = ToLibrary }) + "Generate a library instead of an executable." + , Option [] ["executable"] + (NoArg $ Right $ \config -> config { compilerMode = ToExecutable }) + "Generate an executable instead of a library (set by default)." + , Option [] ["Werror"] + (NoArg $ Right $ \config -> config { compilerWerror = True }) + "Treat warnings as errors." + , Option [] ["safe"] + (NoArg $ Right $ \config -> config { compilerSafe = True }) + "Ignore 'unsafe' in code." + ] + +wrapOption :: CompilerOption cfg -> CoreCompilerOption cfg +wrapOption = fmap wrap + where wrap f = do + g <- f + return $ \cfg -> cfg { compilerConfig = g (compilerConfig cfg) } + +incVerbosity :: Maybe FilePath -> CompilerConfig cfg -> CompilerConfig cfg +incVerbosity file cfg = + cfg { compilerVerbose = (v, file `mplus` snd (compilerVerbose cfg)) } + where v = case fst $ compilerVerbose cfg of + NotVerbose -> Verbose + Verbose -> VeryVerbose + VeryVerbose -> VeryVerbose + +data CompilerConfig cfg = + CompilerConfig { compilerOutput :: Maybe FilePath + , compilerVerbose :: (Verbosity, Maybe FilePath) + , compilerMode :: CompilerMode + , compilerWerror :: Bool + , compilerSafe :: Bool + , compilerConfig :: cfg + } + +-- | Are we compiling a library or an executable? +data CompilerMode = ToLibrary | ToExecutable deriving (Eq, Ord, Show) + +-- | The configuration of the compiler. +newCompilerConfig :: cfg -> CompilerConfig cfg +newCompilerConfig x = CompilerConfig { compilerOutput = Nothing + , compilerVerbose = (NotVerbose, Nothing) + , compilerMode = ToExecutable + , compilerWerror = False + , compilerSafe = False + , compilerConfig = x + } + +outputFilePath :: FilePath -> CompilerConfig cfg -> FilePath +outputFilePath srcfile = + fromMaybe (srcfile `replaceExtension` "") . compilerOutput + +futharkConfig :: CompilerConfig cfg -> FutharkConfig +futharkConfig config = + newFutharkConfig { futharkVerbose = compilerVerbose config + , futharkWerror = compilerWerror config + , futharkSafe = compilerSafe config + } diff --git a/src/Futhark/Compiler/Program.hs b/src/Futhark/Compiler/Program.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Compiler/Program.hs @@ -0,0 +1,191 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TupleSections #-} +-- | Low-level compilation parts. Look at "Futhark.Compiler" for a +-- more high-level API. +module Futhark.Compiler.Program + ( readLibraryWithBasis + , readImports + , Imports + , FileModule(..) + , E.Warnings + + , Basis(..) + , emptyBasis + ) +where + +import Data.Semigroup ((<>)) +import Data.Loc +import Control.Exception +import Control.Monad +import Control.Monad.Reader +import Control.Monad.State +import Control.Monad.Except +import qualified Data.Map.Strict as M +import Data.Maybe +import Data.List +import qualified System.FilePath.Posix as Posix +import System.IO.Error +import qualified Data.Text as T +import qualified Data.Text.IO as T + +import Futhark.Error +import Futhark.FreshNames +import Language.Futhark.Parser +import qualified Language.Futhark as E +import qualified Language.Futhark.TypeChecker as E +import Language.Futhark.Semantic +import Language.Futhark.Futlib + +-- | A little monad for reading and type-checking a Futhark program. +type CompilerM m = ReaderT [FilePath] (StateT ReaderState m) + +data ReaderState = ReaderState { alreadyImported :: Imports + , nameSource :: VNameSource + , warnings :: E.Warnings + } + +-- | Pre-typechecked imports, including a starting point for the name source. +data Basis = Basis { basisImports :: Imports + , basisNameSource :: VNameSource + , basisRoots :: [String] + -- ^ Files that should be implicitly opened. + } + +-- | A basis that contains no imports, and has a properly initialised +-- name source. +emptyBasis :: Basis +emptyBasis = Basis { basisImports = mempty + , basisNameSource = src + , basisRoots = mempty + } + where src = newNameSource $ succ $ maximum $ map E.baseTag $ M.keys E.intrinsics + +readImport :: (MonadError CompilerError m, MonadIO m) => + [ImportName] -> ImportName -> CompilerM m () +readImport steps include + | include `elem` steps = + throwError $ ExternalError $ T.pack $ + "Import cycle: " ++ intercalate " -> " + (map includeToString $ reverse $ include:steps) + | otherwise = do + already_done <- gets $ isJust . lookup (includeToString include) . alreadyImported + + unless already_done $ + uncurry (handleFile steps include) =<< readImportFile include + +handleFile :: (MonadIO m, MonadError CompilerError m) => + [ImportName] + -> ImportName + -> T.Text + -> FilePath + -> CompilerM m () +handleFile steps include file_contents file_name = do + prog <- case parseFuthark file_name file_contents of + Left err -> externalErrorS $ show err + Right prog -> return prog + + mapM_ (readImport steps' . uncurry (mkImportFrom include)) $ + E.progImports prog + + -- It is important to not read these before the above calls to + -- readImport. + imports <- gets alreadyImported + src <- gets nameSource + roots <- ask + + case E.checkProg imports src include $ prependRoots roots prog of + Left err -> + externalError $ T.pack $ show err + Right (m, ws, src') -> + modify $ \s -> + s { alreadyImported = (includeToString include,m) : imports + , nameSource = src' + , warnings = warnings s <> ws + } + where steps' = include:steps + +readFileSafely :: String -> IO (Maybe (Either String (String, T.Text))) +readFileSafely filepath = + (Just . Right . (filepath,) <$> T.readFile filepath) `catch` couldNotRead + where couldNotRead e + | isDoesNotExistError e = + return Nothing + | otherwise = + return $ Just $ Left $ show e + +readImportFile :: (MonadError CompilerError m, MonadIO m) => + ImportName -> m (T.Text, FilePath) +readImportFile include = do + -- First we try to find a file of the given name in the search path, + -- then we look at the builtin library if we have to. For the + -- builtins, we don't use the search path. + r <- liftIO $ readFileSafely $ includeToFilePath include + case (r, lookup futlib_str futlib) of + (Just (Right (filepath,s)), _) -> return (s, filepath) + (Just (Left e), _) -> externalErrorS e + (Nothing, Just t) -> return (t, futlib_str) + (Nothing, Nothing) -> externalErrorS not_found + where futlib_str = "/" Posix. includeToString include Posix.<.> "fut" + + not_found = + "Error at " ++ E.locStr (srclocOf include) ++ + ": could not find import '" ++ includeToString include ++ "'." + +-- | Read Futhark files from some basis, and printing log messages if +-- the first parameter is True. +readLibraryWithBasis :: (MonadError CompilerError m, MonadIO m) => + Basis -> [FilePath] + -> m (E.Warnings, + Imports, + VNameSource) +readLibraryWithBasis builtin fps = do + (_, imps, src) <- runCompilerM builtin $ + mapM (readImport [] . mkInitialImport) prelude + let basis = Basis imps src prelude + readLibrary' basis fps + +-- | Read and type-check a Futhark library (multiple files, relative +-- to the same search path), including all imports. +readLibrary' :: (MonadError CompilerError m, MonadIO m) => + Basis -> [FilePath] + -> m (E.Warnings, + Imports, + VNameSource) +readLibrary' basis fps = runCompilerM basis $ mapM onFile fps + where onFile fp = do + r <- liftIO $ readFileSafely fp + case r of + Just (Right (_, fs)) -> + handleFile [] (mkInitialImport fp_name) fs fp + Just (Left e) -> externalError $ T.pack e + Nothing -> externalErrorS $ fp ++ ": file not found." + where (fp_name, _) = Posix.splitExtension fp + +-- | Read and type-check Futhark imports (no @.fut@ extension; may +-- refer to baked-in futlib). This is an exotic operation that +-- probably only makes sense in an interactive environment. +readImports :: (MonadError CompilerError m, MonadIO m) => + Basis -> [ImportName] + -> m (E.Warnings, + Imports, + VNameSource) +readImports basis imps = + runCompilerM basis $ mapM (readImport []) imps + +runCompilerM :: Monad m => + Basis -> CompilerM m a + -> m (E.Warnings, [(String, FileModule)], VNameSource) +runCompilerM (Basis imports src roots) m = do + let s = ReaderState (reverse imports) src mempty + s' <- execStateT (runReaderT m roots) s + return (warnings s', + reverse $ alreadyImported s', + nameSource s') + +prependRoots :: [FilePath] -> E.UncheckedProg -> E.UncheckedProg +prependRoots roots (E.Prog doc ds) = + E.Prog doc $ map mkImport roots ++ ds + where mkImport fp = + E.LocalDec (E.OpenDec (E.ModImport fp E.NoInfo noLoc) E.NoInfo noLoc) noLoc diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Construct.hs @@ -0,0 +1,529 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE FlexibleInstances #-} +module Futhark.Construct + ( letSubExp + , letSubExps + , letExp + , letExps + , letTupExp + , letTupExp' + , letInPlace + + , eSubExp + , eIf + , eIf' + , eBinOp + , eCmpOp + , eConvOp + , eNegate + , eNot + , eAbs + , eSignum + , eCopy + , eAssert + , eBody + , eLambda + , eDivRoundingUp + , eRoundToMultipleOf + , eSliceArray + , eSplitArray + + , eWriteArray + + , asIntZ, asIntS + + , resultBody + , resultBodyM + , insertStmsM + , mapResult + + , foldBinOp + , binOpLambda + , cmpOpLambda + , fullSlice + , fullSliceNum + , isFullSlice + , ifCommon + + , module Futhark.Binder + + -- * Result types + , instantiateShapes + , instantiateShapes' + , instantiateShapesFromIdentList + , instantiateExtTypes + , instantiateIdents + , removeExistentials + + -- * Convenience + , simpleMkLetNames + + , ToExp(..) + ) +where + +import qualified Data.Map.Strict as M +import Data.Loc (SrcLoc) +import Data.List +import Control.Monad.Identity +import Control.Monad.State +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.MonadFreshNames +import Futhark.Binder +import Futhark.Util + +letSubExp :: MonadBinder m => + String -> Exp (Lore m) -> m SubExp +letSubExp _ (BasicOp (SubExp se)) = return se +letSubExp desc e = Var <$> letExp desc e + +letExp :: MonadBinder m => + String -> Exp (Lore m) -> m VName +letExp _ (BasicOp (SubExp (Var v))) = + return v +letExp desc e = do + n <- length <$> expExtType e + vs <- replicateM n $ newVName desc + idents <- letBindNames vs e + case idents of + [ident] -> return $ identName ident + _ -> fail $ "letExp: tuple-typed expression given:\n" ++ pretty e + +letInPlace :: MonadBinder m => + String -> VName -> Slice SubExp -> Exp (Lore m) + -> m VName +letInPlace desc src slice e = do + tmp <- letSubExp (desc ++ "_tmp") e + letExp desc $ BasicOp $ Update src slice tmp + +letSubExps :: MonadBinder m => + String -> [Exp (Lore m)] -> m [SubExp] +letSubExps desc = mapM $ letSubExp desc + +letExps :: MonadBinder m => + String -> [Exp (Lore m)] -> m [VName] +letExps desc = mapM $ letExp desc + +letTupExp :: (MonadBinder m) => + String -> Exp (Lore m) + -> m [VName] +letTupExp _ (BasicOp (SubExp (Var v))) = + return [v] +letTupExp name e = do + numValues <- length <$> expExtType e + names <- replicateM numValues $ newVName name + map identName <$> letBindNames names e + +letTupExp' :: (MonadBinder m) => + String -> Exp (Lore m) + -> m [SubExp] +letTupExp' _ (BasicOp (SubExp se)) = return [se] +letTupExp' name ses = map Var <$> letTupExp name ses + +eSubExp :: MonadBinder m => + SubExp -> m (Exp (Lore 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)) +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)) + -> IfSort + -> m (Exp (Lore m)) +eIf' ce te fe if_sort = do + ce' <- letSubExp "cond" =<< ce + te' <- insertStmsM te + fe' <- insertStmsM fe + -- We need to construct the context. + ts <- generaliseExtTypes <$> bodyExtType te' <*> bodyExtType fe' + te'' <- addContextForBranch ts te' + fe'' <- addContextForBranch ts fe' + return $ If ce' te'' fe'' $ IfAttr ts if_sort + where addContextForBranch ts (Body _ stms val_res) = do + body_ts <- extendedScope (traverse subExpType val_res) stmsscope + let ctx_res = map snd $ sortOn fst $ + M.toList $ shapeExtMapping ts body_ts + mkBodyM stms $ ctx_res++val_res + where stmsscope = scopeOf stms + +eBinOp :: MonadBinder m => + BinOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) + -> m (Exp (Lore m)) +eBinOp op x y = do + x' <- letSubExp "x" =<< x + y' <- letSubExp "y" =<< y + return $ BasicOp $ BinOp op x' y' + +eCmpOp :: MonadBinder m => + CmpOp -> m (Exp (Lore m)) -> m (Exp (Lore m)) + -> m (Exp (Lore m)) +eCmpOp op x y = do + x' <- letSubExp "x" =<< x + y' <- letSubExp "y" =<< y + return $ BasicOp $ CmpOp op x' y' + +eConvOp :: MonadBinder m => + ConvOp -> m (Exp (Lore m)) + -> m (Exp (Lore m)) +eConvOp op x = do + x' <- letSubExp "x" =<< x + return $ BasicOp $ ConvOp op x' + +eNegate :: MonadBinder m => + m (Exp (Lore m)) -> m (Exp (Lore m)) +eNegate em = do + e <- em + e' <- letSubExp "negate_arg" e + t <- subExpType e' + case t of + Prim (IntType int_t) -> + return $ BasicOp $ + BinOp (Sub int_t) (intConst int_t 0) e' + Prim (FloatType float_t) -> + return $ BasicOp $ + BinOp (FSub float_t) (floatConst float_t 0) e' + _ -> + fail $ "eNegate: operand " ++ pretty e ++ " has invalid type." + +eNot :: MonadBinder m => + m (Exp (Lore m)) -> m (Exp (Lore m)) +eNot e = BasicOp . UnOp Not <$> (letSubExp "not_arg" =<< e) + +eAbs :: MonadBinder m => + m (Exp (Lore m)) -> m (Exp (Lore m)) +eAbs em = do + e <- em + e' <- letSubExp "abs_arg" e + t <- subExpType e' + case t of + Prim (IntType int_t) -> + return $ BasicOp $ UnOp (Abs int_t) e' + Prim (FloatType float_t) -> + return $ BasicOp $ UnOp (FAbs float_t) e' + _ -> + fail $ "eAbs: operand " ++ pretty e ++ " has invalid type." + +eSignum :: MonadBinder m => + m (Exp (Lore m)) -> m (Exp (Lore m)) +eSignum em = do + e <- em + e' <- letSubExp "signum_arg" e + t <- subExpType e' + case t of + Prim (IntType int_t) -> + return $ BasicOp $ UnOp (SSignum int_t) e' + _ -> + fail $ "eSignum: operand " ++ pretty e ++ " has invalid type." + +eCopy :: MonadBinder m => + m (Exp (Lore m)) -> m (Exp (Lore m)) +eCopy e = BasicOp . Copy <$> (letExp "copy_arg" =<< e) + +eAssert :: MonadBinder m => + m (Exp (Lore m)) -> ErrorMsg SubExp -> SrcLoc -> m (Exp (Lore m)) +eAssert e msg loc = do e' <- letSubExp "assert_arg" =<< e + return $ BasicOp $ Assert e' msg (loc, mempty) + +eBody :: (MonadBinder m) => + [m (Exp (Lore m))] + -> m (Body (Lore m)) +eBody es = insertStmsM $ do + es' <- sequence es + xs <- mapM (letTupExp "x") es' + mkBodyM mempty $ map Var $ concat xs + +eLambda :: MonadBinder m => + Lambda (Lore m) -> [m (Exp (Lore m))] -> m [SubExp] +eLambda lam args = do zipWithM_ bindParam (lambdaParams lam) args + bodyBind $ lambdaBody lam + where bindParam param arg = letBindNames_ [paramName param] =<< arg + +-- | Note: unsigned division. +eDivRoundingUp :: MonadBinder m => + IntType -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m)) +eDivRoundingUp t x y = + eBinOp (SQuot t) (eBinOp (Add t) x (eBinOp (Sub t) y (eSubExp one))) y + where one = intConst t 1 + +eRoundToMultipleOf :: MonadBinder m => + IntType -> m (Exp (Lore m)) -> m (Exp (Lore m)) -> m (Exp (Lore m)) +eRoundToMultipleOf t x d = + ePlus x (eMod (eMinus d (eMod x d)) d) + where eMod = eBinOp (SMod t) + eMinus = eBinOp (Sub t) + ePlus = eBinOp (Add t) + +-- | Construct an 'Index' expressions that slices an array with unit stride. +eSliceArray :: MonadBinder m => + Int -> VName -> m (Exp (Lore m)) -> m (Exp (Lore m)) + -> m (Exp (Lore m)) +eSliceArray d arr i n = do + arr_t <- lookupType arr + let skips = map (slice (constant (0::Int32))) $ take d $ arrayDims arr_t + i' <- letSubExp "slice_i" =<< i + n' <- letSubExp "slice_n" =<< n + return $ BasicOp $ Index arr $ fullSlice arr_t $ skips ++ [slice i' n'] + where slice j m = DimSlice j m (constant (1::Int32)) + +-- | Construct an 'Index' expressions that splits an array in different parts along the outer dimension. +eSplitArray :: MonadBinder m => + VName -> [m (Exp (Lore m))] -> m [Exp (Lore m)] +eSplitArray arr sizes = do + sizes' <- mapM (letSubExp "split_size") =<< sequence sizes + -- Compute the starting offset for each slice. + (_, offsets) <- mapAccumLM increase (intConst Int32 0) sizes' + zipWithM (eSliceArray 0 arr) (map eSubExp offsets) (map eSubExp sizes') + where increase offset size = do + offset' <- letSubExp "offset" $ BasicOp $ BinOp (Add Int32) offset size + return (offset', offset) + +-- | Write to an index of the array, if within bounds. Otherwise, +-- nothing. Produces the updated array. +eWriteArray :: (MonadBinder m, BranchType (Lore m) ~ ExtType) => + VName -> [m (Exp (Lore m))] -> m (Exp (Lore m)) + -> m (Exp (Lore m)) +eWriteArray arr is v = do + arr_t <- lookupType arr + let ws = arrayDims arr_t + is' <- mapM (letSubExp "write_i") =<< sequence is + v' <- letSubExp "write_v" =<< v + let checkDim w i = do + less_than_zero <- letSubExp "less_than_zero" $ + BasicOp $ CmpOp (CmpSlt Int32) i (constant (0::Int32)) + greater_than_size <- letSubExp "greater_than_size" $ + BasicOp $ CmpOp (CmpSle Int32) w i + letSubExp "outside_bounds_dim" $ + BasicOp $ BinOp LogOr less_than_zero greater_than_size + + outside_bounds <- + letSubExp "outside_bounds" =<< + foldBinOp LogOr (constant False) =<< + zipWithM checkDim ws is' + + outside_bounds_branch <- insertStmsM $ resultBodyM [Var arr] + + in_bounds_branch <- insertStmsM $ do + res <- letInPlace "write_out_inside_bounds" arr + (fullSlice arr_t (map DimFix is')) $ BasicOp $ SubExp v' + resultBodyM [Var res] + + return $ + If outside_bounds outside_bounds_branch in_bounds_branch $ + ifCommon [arr_t] + +-- | Sign-extend to the given integer type. +asIntS :: MonadBinder m => IntType -> SubExp -> m SubExp +asIntS = asInt SExt + +-- | Zero-extend to the given integer type. +asIntZ :: MonadBinder m => IntType -> SubExp -> m SubExp +asIntZ = asInt ZExt + +asInt :: MonadBinder m => + (IntType -> IntType -> ConvOp) -> IntType -> SubExp -> m SubExp +asInt ext to_it e = do + e_t <- subExpType e + case e_t of + Prim (IntType from_it) + | to_it == from_it -> return e + | otherwise -> letSubExp s $ BasicOp $ ConvOp (ext from_it to_it) e + _ -> fail "asInt: wrong type" + where s = case e of Var v -> baseString v + _ -> "to_" ++ pretty to_it + + +-- | Apply a binary operator to several subexpressions. A left-fold. +foldBinOp :: MonadBinder m => + BinOp -> SubExp -> [SubExp] -> m (Exp (Lore m)) +foldBinOp _ ne [] = + return $ BasicOp $ SubExp ne +foldBinOp bop ne (e:es) = + eBinOp bop (pure $ BasicOp $ SubExp e) (foldBinOp bop ne es) + +-- | Create a two-parameter lambda whose body applies the given binary +-- operation to its arguments. It is assumed that both argument and +-- result types are the same. (This assumption should be fixed at +-- some point.) +binOpLambda :: (MonadBinder m, Bindable (Lore m)) => + BinOp -> PrimType -> m (Lambda (Lore m)) +binOpLambda bop t = binLambda (BinOp bop) t t + +-- | As 'binOpLambda', but for 'CmpOp's. +cmpOpLambda :: (MonadBinder m, Bindable (Lore m)) => + CmpOp -> PrimType -> m (Lambda (Lore m)) +cmpOpLambda cop t = binLambda (CmpOp cop) t Bool + +binLambda :: (MonadBinder m, Bindable (Lore m)) => + (SubExp -> SubExp -> BasicOp (Lore m)) -> PrimType -> PrimType + -> m (Lambda (Lore m)) +binLambda bop arg_t ret_t = do + x <- newVName "x" + y <- newVName "y" + body <- insertStmsM $ do + res <- letSubExp "res" $ BasicOp $ bop (Var x) (Var y) + return $ resultBody [res] + return Lambda { + lambdaParams = [Param x (Prim arg_t), + Param y (Prim arg_t)] + , lambdaReturnType = [Prim ret_t] + , lambdaBody = body + } + +-- | @fullSlice t slice@ returns @slice@, but with 'DimSlice's of +-- entire dimensions appended to the full dimensionality of @t@. This +-- function is used to turn incomplete indexing complete, as required +-- by 'Index'. +fullSlice :: Type -> [DimIndex SubExp] -> Slice SubExp +fullSlice t slice = + slice ++ + map (\d -> DimSlice (constant (0::Int32)) d (constant (1::Int32))) + (drop (length slice) $ arrayDims t) + +-- | Like 'fullSlice', but the dimensions are simply numeric. +fullSliceNum :: Num d => [d] -> [DimIndex d] -> Slice d +fullSliceNum dims slice = + slice ++ map (\d -> DimSlice 0 d 1) (drop (length slice) dims) + +-- | Does the slice describe the full size of the array? The most +-- obvious such slice is one that 'DimSlice's the full span of every +-- dimension, but also one that fixes all unit dimensions. +isFullSlice :: Shape -> Slice SubExp -> Bool +isFullSlice shape slice = and $ zipWith allOfIt (shapeDims shape) slice + where allOfIt (Constant v) DimFix{} = oneIsh v + allOfIt d (DimSlice _ n _) = d == n + allOfIt _ _ = False + +ifCommon :: [Type] -> IfAttr ExtType +ifCommon ts = IfAttr (staticShapes ts) IfNormal + +-- | Conveniently construct a body that contains no bindings. +resultBody :: Bindable lore => [SubExp] -> Body lore +resultBody = mkBody mempty + +-- | Conveniently construct a body that contains no bindings - but +-- this time, monadically! +resultBodyM :: MonadBinder m => + [SubExp] + -> m (Body (Lore 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)) +insertStmsM m = do + (Body _ bnds res, otherbnds) <- collectStms m + mkBodyM (otherbnds <> bnds) res + +-- | 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 +mapResult f (Body _ bnds res) = + let Body _ bnds2 newres = f res + in mkBody (bnds<>bnds2) newres + +-- | Instantiate all existential parts dimensions of the given +-- type, using a monadic action to create the necessary 'SubExp's. +-- You should call this function within some monad that allows you to +-- collect the actions performed (say, 'Writer'). +instantiateShapes :: Monad m => + (Int -> m SubExp) + -> [TypeBase ExtShape u] + -> m [TypeBase Shape u] +instantiateShapes f ts = evalStateT (mapM instantiate ts) M.empty + where instantiate t = do + shape <- mapM instantiate' $ shapeDims $ arrayShape t + return $ t `setArrayShape` Shape shape + instantiate' (Ext x) = do + m <- get + case M.lookup x m of + Just se -> return se + Nothing -> do se <- lift $ f x + put $ M.insert x se m + return se + instantiate' (Free se) = return se + +instantiateShapes' :: MonadFreshNames m => + [TypeBase ExtShape u] + -> m ([TypeBase Shape u], [Ident]) +instantiateShapes' ts = + runWriterT $ instantiateShapes instantiate ts + where instantiate _ = do v <- lift $ newIdent "size" $ Prim int32 + tell [v] + return $ Var $ identName v + +instantiateShapesFromIdentList :: [Ident] -> [ExtType] -> [Type] +instantiateShapesFromIdentList idents ts = + evalState (instantiateShapes instantiate ts) idents + where instantiate _ = do + idents' <- get + case idents' of + [] -> fail "instantiateShapesFromIdentList: insufficiently sized context" + ident:idents'' -> do put idents'' + return $ Var $ identName ident + +instantiateExtTypes :: [VName] -> [ExtType] -> [Ident] +instantiateExtTypes names rt = + let (shapenames,valnames) = splitAt (shapeContextSize rt) names + shapes = [ Ident name (Prim int32) | name <- shapenames ] + valts = instantiateShapesFromIdentList shapes rt + vals = [ Ident name t | (name,t) <- zip valnames valts ] + in shapes ++ vals + +instantiateIdents :: [VName] -> [ExtType] + -> Maybe ([Ident], [Ident]) +instantiateIdents names ts + | let n = shapeContextSize ts, + n + length ts == length names = do + let (context, vals) = splitAt n names + nextShape _ = do + (context', remaining) <- get + case remaining of [] -> lift Nothing + x:xs -> do let ident = Ident x (Prim int32) + put (context'++[ident], xs) + return $ Var x + (ts', (context', _)) <- + runStateT (instantiateShapes nextShape ts) ([],context) + return (context', zipWith Ident vals ts') + | otherwise = Nothing + +removeExistentials :: ExtType -> Type -> Type +removeExistentials t1 t2 = + t1 `setArrayDims` + zipWith nonExistential + (shapeDims $ arrayShape t1) + (arrayDims t2) + where nonExistential (Ext _) dim = dim + nonExistential (Free dim) _ = dim + +-- | Can be used as the definition of 'mkLetNames' for a 'Bindable' +-- instance for simple representations. +simpleMkLetNames :: (ExpAttr lore ~ (), LetAttr lore ~ Type, + MonadFreshNames m, TypedOp (Op lore), HasScope lore m) => + [VName] -> Exp lore -> m (Stm lore) +simpleMkLetNames names e = do + et <- expExtType e + (ts, shapes) <- instantiateShapes' et + let shapeElems = [ PatElem shape shapet | Ident shape shapet <- shapes ] + let valElems = zipWith PatElem names ts + return $ Let (Pattern shapeElems valElems) (StmAux mempty ()) e + +-- | 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)) + +instance ToExp SubExp where + toExp = return . BasicOp . SubExp + +instance ToExp VName where + toExp = return . BasicOp . SubExp . Var diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Doc/Generator.hs @@ -0,0 +1,745 @@ +{-# LANGUAGE OverloadedStrings #-} +module Futhark.Doc.Generator (renderFiles) where + +import Control.Arrow ((***)) +import Control.Monad +import Control.Monad.Reader +import Control.Monad.Writer +import Data.List (sort, sortOn, intersperse, inits, tails, isPrefixOf, find, groupBy, partition) +import Data.Char (isSpace, isAlpha, toUpper) +import Data.Loc +import Data.Maybe +import Data.Ord +import qualified Data.Map as M +import qualified Data.Set as S +import System.FilePath (splitPath, (), (-<.>), (<.>), makeRelative) +import Text.Blaze.Html5 (AttributeValue, Html, (!), toHtml) +import qualified Text.Blaze.Html5 as H +import qualified Text.Blaze.Html5.Attributes as A +import Data.String (fromString) +import Data.Version +import qualified Data.Text.Lazy as LT +import Text.Markdown + +import Prelude hiding (abs) + +import Language.Futhark.Semantic +import Language.Futhark.TypeChecker.Monad hiding (warn) +import Language.Futhark +import Futhark.Doc.Html +import Futhark.Version + +-- | A set of names that we should not generate links to, because they +-- are uninteresting. These are for example type parameters. +type NoLink = S.Set VName + +data Context = Context { ctxCurrent :: String + , ctxFileMod :: FileModule + , ctxImports :: Imports + , ctxNoLink :: NoLink + , ctxFileMap :: FileMap + , ctxVisibleMTys :: S.Set VName + -- ^ Local module types that show up in the + -- interface. These should be documented, + -- but clearly marked local. + } +type FileMap = M.Map VName (String, Namespace) +type DocM = ReaderT Context (WriterT Documented (Writer Warnings)) + +data IndexWhat = IndexValue | IndexFunction | IndexModule | IndexModuleType | IndexType + +-- | We keep a mapping of the names we have actually documented, so we +-- can generate an index. +type Documented = M.Map VName IndexWhat + +warn :: SrcLoc -> String -> DocM () +warn loc s = lift $ lift $ tell $ singleWarning loc s + +document :: VName -> IndexWhat -> DocM () +document v what = tell $ M.singleton v what + +noLink :: [VName] -> DocM a -> DocM a +noLink names = local $ \ctx -> + ctx { ctxNoLink = S.fromList names <> ctxNoLink ctx } + +selfLink :: AttributeValue -> Html -> Html +selfLink s = H.a ! A.id s ! A.href ("#" <> s) ! A.class_ "self_link" + +fullRow :: Html -> Html +fullRow = H.tr . (H.td ! A.colspan "3") + +emptyRow :: Html +emptyRow = H.tr $ H.td mempty <> H.td mempty <> H.td mempty + +specRow :: Html -> Html -> Html -> Html +specRow a b c = H.tr $ (H.td ! A.class_ "spec_lhs") a <> + (H.td ! A.class_ "spec_eql") b <> + (H.td ! A.class_ "spec_rhs") c + +vnameToFileMap :: Imports -> FileMap +vnameToFileMap = mconcat . map forFile + where forFile (file, FileModule abs file_env _prog) = + mconcat (map (vname Type) (M.keys abs)) <> + forEnv file_env + where vname ns v = M.singleton (qualLeaf v) (file, ns) + vname' ((ns, _), v) = vname ns v + + forEnv env = + mconcat (map vname' $ M.toList $ envNameMap env) <> + mconcat (map forMty $ M.elems $ envSigTable env) + forMod (ModEnv env) = forEnv env + forMod ModFun{} = mempty + forMty = forMod . mtyMod + +renderFiles :: [FilePath] -> Imports -> ([(FilePath, Html)], Warnings) +renderFiles important_imports imports = runWriter $ do + (import_pages, documented) <- runWriterT $ forM imports $ \(current, fm) -> + let ctx = Context current fm imports mempty file_map + (progModuleTypes $ fileProg fm) in + flip runReaderT ctx $ do + + (first_paragraph, maybe_abstract, maybe_sections) <- headerDoc $ fileProg fm + + synopsis <- (H.div ! A.id "module") <$> synopsisDecs (progDecs $ fileProg fm) + + description <- describeDecs $ progDecs $ fileProg fm + + return (current, + (H.docTypeHtml ! A.lang "en" $ + addBoilerplateWithNav important_imports imports ("doc" current) current $ + H.main $ maybe_abstract <> + selfLink "synopsis" (H.h2 "Synopsis") <> (H.div ! A.id "overview") synopsis <> + selfLink "description" (H.h2 "Description") <> description <> + maybe_sections, + first_paragraph)) + + return $ + [("index.html", contentsPage important_imports $ map (fmap snd) import_pages), + ("doc-index.html", indexPage important_imports imports documented file_map)] + ++ map (importHtml *** fst) import_pages + where file_map = vnameToFileMap imports + importHtml import_name = "doc" import_name <.> "html" + +-- | The header documentation (which need not be present) can contain +-- an abstract and further sections. +headerDoc :: Prog -> DocM (Html, Html, Html) +headerDoc prog = + case progDoc prog of + Just (DocComment doc loc) -> do + let (abstract, more_sections) = splitHeaderDoc doc + first_paragraph <- docHtml $ Just $ DocComment (firstParagraph abstract) loc + abstract' <- docHtml $ Just $ DocComment abstract loc + more_sections' <- docHtml $ Just $ DocComment more_sections loc + return (first_paragraph, + selfLink "abstract" (H.h2 "Abstract") <> abstract', + more_sections') + _ -> return mempty + where splitHeaderDoc s = fromMaybe (s, mempty) $ + find (("\n##" `isPrefixOf`) . snd) $ + zip (inits s) (tails s) + firstParagraph = unlines . takeWhile (not . paragraphSeparator) . lines + paragraphSeparator = all isSpace + + +contentsPage :: [FilePath] -> [(String, Html)] -> Html +contentsPage important_imports pages = + H.docTypeHtml $ addBoilerplate "index.html" "Futhark Library Documentation" $ + H.main $ H.h2 "Main libraries" <> + fileList important_pages <> + if null unimportant_pages then mempty else + H.h2 "Supporting libraries" <> + fileList unimportant_pages + where (important_pages, unimportant_pages) = + partition ((`elem` important_imports) . fst) pages + + fileList pages' = + H.dl ! A.class_ "file_list" $ + mconcat $ map linkTo $ sortOn fst pages' + + linkTo (name, maybe_abstract) = + H.div ! A.class_ "file_desc" $ + (H.dt ! A.class_ "desc_header") (importLink "index.html" name) <> + (H.dd ! A.class_ "desc_doc") maybe_abstract + +importLink :: FilePath -> String -> Html +importLink current name = + let file = relativise (makeRelative "/" $ "doc" name -<.> "html") current + in (H.a ! A.href (fromString file) $ fromString name) + +indexPage :: [FilePath] -> Imports -> Documented -> FileMap -> Html +indexPage important_imports imports documented fm = + H.docTypeHtml $ addBoilerplateWithNav important_imports imports "doc-index.html" "Index" $ + H.main $ + (H.ul ! A.id "doc_index_list" $ + mconcat $ map initialListEntry $ + letter_group_links ++ [symbol_group_link]) <> + (H.table ! A.id "doc_index" $ + H.thead (H.tr $ H.td "Who" <> H.td "What" <> H.td "Where") <> + mconcat (letter_groups ++ [symbol_group])) + where (letter_names, sym_names) = + partition (isLetterName . baseString . fst) $ + sortOn (map toUpper . baseString . fst) $ + mapMaybe isDocumented $ M.toList fm + + isDocumented (k, (file, _)) = do + what <- M.lookup k documented + Just (k, (file, what)) + + (letter_groups, letter_group_links) = + unzip $ map tbodyForNames $ groupBy sameInitial letter_names + (symbol_group, symbol_group_link) = + tbodyForInitial "Symbols" sym_names + + isLetterName [] = False + isLetterName (c:_) = isAlpha c + + sameInitial (x, _) (y, _) = + case (baseString x, baseString y) of + (x':_, y':_) -> toUpper x' == toUpper y' + _ -> False + + tbodyForNames names@((s,_):_) = + tbodyForInitial (map toUpper $ take 1 $ baseString s) names + tbodyForNames _ = mempty + + tbodyForInitial initial names = + (H.tbody $ mconcat $ initial' : map linkTo names, + initial) + where initial' = + H.tr $ H.td ! A.colspan "2" ! A.class_ "doc_index_initial" $ + H.a ! A.id (fromString initial) + ! A.href (fromString $ '#' : initial) + $ fromString initial + + initialListEntry initial = + H.li $ H.a ! A.href (fromString $ '#' : initial) $ fromString initial + + linkTo (name, (file, what)) = + let link = (H.a ! A.href (fromString (makeRelative "/" $ "doc" vnameLink' name "" file))) $ + fromString $ baseString name + what' = case what of IndexValue -> "value" + IndexFunction -> "function" + IndexType -> "type" + IndexModuleType -> "module type" + IndexModule -> "module" + html_file = makeRelative "/" $ "doc" file -<.> "html" + in H.tr $ + (H.td ! A.class_ "doc_index_name" $ link) <> + (H.td ! A.class_ "doc_index_namespace" $ what') <> + (H.td ! A.class_ "doc_index_file" $ + (H.a ! A.href (fromString html_file) $ fromString file)) + +addBoilerplate :: String -> String -> Html -> Html +addBoilerplate current titleText content = + let headHtml = H.head $ + H.meta ! A.charset "utf-8" <> + H.title (fromString titleText) <> + H.link ! A.href (fromString $ relativise "style.css" current) + ! A.rel "stylesheet" + ! A.type_ "text/css" + + navigation = H.ul ! A.id "navigation" $ + H.li (H.a ! A.href (fromString $ relativise "index.html" current) $ "Contents") <> + H.li (H.a ! A.href (fromString $ relativise "doc-index.html" current) $ "Index") + + madeByHtml = + "Generated by " <> (H.a ! A.href futhark_doc_url) "futhark-doc" + <> " " <> fromString (showVersion version) + in headHtml <> + H.body ((H.div ! A.id "header") (H.h1 (toHtml titleText) <> navigation) <> + (H.div ! A.id "content") content <> + (H.div ! A.id "footer") madeByHtml) + where futhark_doc_url = + "https://futhark.readthedocs.io/en/latest/man/futhark-doc.html" + +addBoilerplateWithNav :: [FilePath] -> Imports -> String -> String -> Html -> Html +addBoilerplateWithNav important_imports imports current titleText content = + addBoilerplate current titleText $ + (H.nav ! A.id "filenav" $ files) <> content + where files = H.ul $ mconcat $ map pp $ sort $ filter visible important_imports + pp name = H.li $ importLink current name + visible = (`elem` map fst imports) + +synopsisDecs :: [Dec] -> DocM Html +synopsisDecs decs = do + visible <- asks ctxVisibleMTys + fm <- asks ctxFileMod + -- We add an empty row to avoid generating invalid HTML in cases + -- where all rows are otherwise colspan=2. + (H.table ! A.class_ "specs") . (emptyRow<>) . mconcat <$> + sequence (mapMaybe (synopsisDec visible fm) decs) + +synopsisDec :: S.Set VName -> FileModule -> Dec -> Maybe (DocM Html) +synopsisDec visible fm dec = case dec of + SigDec s -> synopsisModType mempty s + ModDec m -> synopsisMod fm m + ValDec v -> synopsisValBind v + TypeDec t -> synopsisType t + OpenDec x (Info _names) _ + | Just opened <- synopsisOpened x -> Just $ do + opened' <- opened + return $ fullRow $ keyword "open " <> opened' + | otherwise -> + Just $ return $ fullRow $ + keyword "open" <> fromString (" <" <> pretty x <> ">") + LocalDec (SigDec s) _ + | sigName s `S.member` visible -> + synopsisModType (keyword "local" <> " ") s + LocalDec _ _ -> Nothing + +synopsisOpened :: ModExp -> Maybe (DocM Html) +synopsisOpened (ModVar qn _) = Just $ qualNameHtml qn +synopsisOpened (ModParens me _) = do me' <- synopsisOpened me + Just $ parens <$> me' +synopsisOpened (ModImport _ (Info file) _) = Just $ do + current <- asks ctxCurrent + let dest = fromString $ relativise file current <> ".html" + return $ keyword "import " <> (H.a ! A.href dest) (fromString $ show file) +synopsisOpened (ModAscript _ se _ _) = Just $ do + se' <- synopsisSigExp se + return $ "... : " <> se' +synopsisOpened _ = Nothing + +synopsisValBind :: ValBind -> Maybe (DocM Html) +synopsisValBind vb = Just $ do + let name' = vnameSynopsisDef $ valBindName vb + (lhs, mhs, rhs) <- valBindHtml name' vb + return $ specRow lhs (mhs <> " : ") rhs + +valBindHtml :: Html -> ValBind -> DocM (Html, Html, Html) +valBindHtml name (ValBind _ _ retdecl (Info rettype) tparams params _ _ _) = do + let tparams' = mconcat $ map ((" "<>) . typeParamHtml) tparams + noLink' = noLink $ map typeParamName tparams ++ + map identName (S.toList $ mconcat $ map patIdentSet params) + rettype' <- noLink' $ maybe (typeHtml rettype) typeExpHtml retdecl + params' <- noLink' $ mapM patternHtml params + return (keyword "val " <> (H.span ! A.class_ "decl_name") name, + tparams', + mconcat (intersperse " -> " $ params' ++ [rettype'])) + +synopsisModType :: Html -> SigBind -> Maybe (DocM Html) +synopsisModType prefix sb = Just $ do + let name' = vnameSynopsisDef $ sigName sb + fullRow <$> do + se' <- synopsisSigExp $ sigExp sb + return $ prefix <> keyword "module type " <> name' <> " = " <> se' + +synopsisMod :: FileModule -> ModBind -> Maybe (DocM Html) +synopsisMod fm (ModBind name ps sig _ _ _) = + case sig of Nothing -> (proceed <=< envSig) <$> M.lookup name modtable + Just (s,_) -> Just $ proceed =<< synopsisSigExp s + where proceed sig' = do + let name' = vnameSynopsisDef name + ps' <- modParamHtml ps + return $ specRow (keyword "module " <> name') ": " (ps' <> sig') + + FileModule _abs Env { envModTable = modtable} _ = fm + envSig (ModEnv e) = renderEnv e + envSig (ModFun (FunSig _ _ (MTy _ m))) = envSig m + +synopsisType :: TypeBind -> Maybe (DocM Html) +synopsisType tb = Just $ do + let name' = vnameSynopsisDef $ typeAlias tb + fullRow <$> typeBindHtml name' tb + +typeBindHtml :: Html -> TypeBind -> DocM Html +typeBindHtml name' (TypeBind _ tparams t _ _) = do + t' <- noLink (map typeParamName tparams) $ typeDeclHtml t + return $ typeAbbrevHtml Unlifted name' tparams <> " = " <> t' + +renderEnv :: Env -> DocM Html +renderEnv (Env vtable ttable sigtable modtable _) = do + typeBinds <- mapM renderTypeBind (M.toList ttable) + valBinds <- mapM renderValBind (M.toList vtable) + sigBinds <- mapM renderModType (M.toList sigtable) + modBinds <- mapM renderMod (M.toList modtable) + return $ braces $ mconcat $ typeBinds ++ valBinds ++ sigBinds ++ modBinds + +renderModType :: (VName, MTy) -> DocM Html +renderModType (name, _sig) = + (keyword "module type " <>) <$> qualNameHtml (qualName name) + +renderMod :: (VName, Mod) -> DocM Html +renderMod (name, _mod) = + (keyword "module " <>) <$> qualNameHtml (qualName name) + +renderValBind :: (VName, BoundV) -> DocM Html +renderValBind = fmap H.div . synopsisValBindBind + +renderTypeBind :: (VName, TypeBinding) -> DocM Html +renderTypeBind (name, TypeAbbr l tps tp) = do + tp' <- typeHtml tp + return $ H.div $ typeAbbrevHtml l (vnameHtml name) tps <> " = " <> tp' + +synopsisValBindBind :: (VName, BoundV) -> DocM Html +synopsisValBindBind (name, BoundV tps t) = do + let tps' = map typeParamHtml tps + t' <- typeHtml t + return $ keyword "val " <> vnameHtml name <> joinBy " " tps' <> ": " <> t' + +typeHtml :: StructType -> DocM Html +typeHtml t = case t of + Prim et -> return $ primTypeHtml et + Record fs + | Just ts <- areTupleFields fs -> + parens . commas <$> mapM typeHtml ts + | otherwise -> + braces . commas <$> mapM ppField (M.toList fs) + where ppField (name, tp) = do + tp' <- typeHtml tp + return $ toHtml (nameToString name) <> ": " <> tp' + TypeVar _ u et targs -> do + targs' <- mapM typeArgHtml targs + et' <- typeNameHtml et + return $ prettyU u <> et' <> joinBy " " targs' + Array et shape u -> do + shape' <- prettyShapeDecl shape + et' <- prettyElem et + return $ prettyU u <> shape' <> et' + Arrow _ pname t1 t2 -> do + t1' <- typeHtml t1 + t2' <- typeHtml t2 + return $ case pname of + Just v -> + parens (vnameHtml v <> ": " <> t1') <> " -> " <> t2' + Nothing -> + t1' <> " -> " <> t2' + +prettyElem :: ArrayElemTypeBase (DimDecl VName) () -> DocM Html +prettyElem (ArrayPrimElem et _) = return $ primTypeHtml et +prettyElem (ArrayPolyElem et targs _) = do + targs' <- mapM typeArgHtml targs + return $ prettyTypeName et <> joinBy " " targs' +prettyElem (ArrayRecordElem fs) + | Just ts <- areTupleFields fs = + parens . commas <$> mapM prettyRecordElem ts + | otherwise = + braces . commas <$> mapM ppField (M.toList fs) + where ppField (name, tp) = do + tp' <- prettyRecordElem tp + return $ toHtml (nameToString name) <> ": " <> tp' + +prettyRecordElem :: RecordArrayElemTypeBase (DimDecl VName) () -> DocM Html +prettyRecordElem (RecordArrayElem et) = prettyElem et +prettyRecordElem (RecordArrayArrayElem et shape u) = + typeHtml $ Array et shape u + +prettyShapeDecl :: ShapeDecl (DimDecl VName) -> DocM Html +prettyShapeDecl (ShapeDecl ds) = + mconcat <$> mapM (fmap brackets . dimDeclHtml) ds + +typeArgHtml :: TypeArg (DimDecl VName) () -> DocM Html +typeArgHtml (TypeArgDim d _) = brackets <$> dimDeclHtml d +typeArgHtml (TypeArgType t _) = typeHtml t + +modParamHtml :: [ModParamBase Info VName] -> DocM Html +modParamHtml [] = return mempty +modParamHtml (ModParam pname psig _ _ : mps) = + liftM2 f (synopsisSigExp psig) (modParamHtml mps) + where f se params = "(" <> vnameHtml pname <> + ": " <> se <> ") -> " <> params + +synopsisSigExp :: SigExpBase Info VName -> DocM Html +synopsisSigExp e = case e of + SigVar v _ -> qualNameHtml v + SigParens e' _ -> parens <$> synopsisSigExp e' + SigSpecs ss _ -> braces . (H.table ! A.class_ "specs") . mconcat <$> mapM synopsisSpec ss + SigWith s (TypeRef v ps t _) _ -> do + s' <- synopsisSigExp s + t' <- typeDeclHtml t + v' <- qualNameHtml v + let ps' = mconcat $ map ((" "<>) . typeParamHtml) ps + return $ s' <> keyword " with " <> v' <> ps' <> " = " <> t' + SigArrow Nothing e1 e2 _ -> + liftM2 f (synopsisSigExp e1) (synopsisSigExp e2) + where f e1' e2' = e1' <> " -> " <> e2' + SigArrow (Just v) e1 e2 _ -> + do let name = vnameHtml v + e1' <- synopsisSigExp e1 + e2' <- noLink [v] $ synopsisSigExp e2 + return $ "(" <> name <> ": " <> e1' <> ") -> " <> e2' + +keyword :: String -> Html +keyword = (H.span ! A.class_ "keyword") . fromString + +vnameHtml :: VName -> Html +vnameHtml (VName name tag) = + H.span ! A.id (fromString (show tag)) $ renderName name + +vnameDescDef :: VName -> IndexWhat -> DocM Html +vnameDescDef v what = do + document v what + return $ H.a ! A.id (fromString (show (baseTag v))) $ renderName (baseName v) + +vnameSynopsisDef :: VName -> Html +vnameSynopsisDef (VName name tag) = + H.span ! A.id (fromString (show tag ++ "s")) $ + H.a ! A.href (fromString ("#" ++ show tag)) $ renderName name + +vnameSynopsisRef :: VName -> Html +vnameSynopsisRef v = H.a ! A.class_ "synopsis_link" + ! A.href (fromString ("#" ++ show (baseTag v) ++ "s")) $ + "↑" + +synopsisSpec :: SpecBase Info VName -> DocM Html +synopsisSpec spec = case spec of + TypeAbbrSpec tpsig -> + fullRow <$> typeBindHtml (vnameSynopsisDef $ typeAlias tpsig) tpsig + TypeSpec Unlifted name ps _ _ -> + return $ fullRow $ keyword "type " <> vnameSynopsisDef name <> joinBy " " (map typeParamHtml ps) + TypeSpec Lifted name ps _ _ -> + return $ fullRow $ keyword "type" <> "^" <> vnameSynopsisDef name <> joinBy " " (map typeParamHtml ps) + ValSpec name tparams rettype _ _ -> do + let tparams' = map typeParamHtml tparams + rettype' <- noLink (map typeParamName tparams) $ + typeDeclHtml rettype + return $ specRow + (keyword "val " <> vnameSynopsisDef name) + (mconcat (map (" "<>) tparams') <> ": ") rettype' + ModSpec name sig _ _ -> + specRow (keyword "module " <> vnameSynopsisDef name) ": " <$> synopsisSigExp sig + IncludeSpec e _ -> fullRow . (keyword "include " <>) <$> synopsisSigExp e + +typeDeclHtml :: TypeDeclBase f VName -> DocM Html +typeDeclHtml = typeExpHtml . declaredType + +typeExpHtml :: TypeExp VName -> DocM Html +typeExpHtml e = case e of + TEUnique t _ -> ("*"<>) <$> typeExpHtml t + TEArray at d _ -> do + at' <- typeExpHtml at + d' <- dimDeclHtml d + return $ brackets d' <> at' + TETuple ts _ -> parens . commas <$> mapM typeExpHtml ts + TERecord fs _ -> braces . commas <$> mapM ppField fs + where ppField (name, t) = do + t' <- typeExpHtml t + return $ toHtml (nameToString name) <> ": " <> t' + TEVar name _ -> qualNameHtml name + TEApply t arg _ -> do + t' <- typeExpHtml t + arg' <- typeArgExpHtml arg + return $ t' <> " " <> arg' + TEArrow pname t1 t2 _ -> do + t1' <- case t1 of TEArrow{} -> parens <$> typeExpHtml t1 + _ -> typeExpHtml t1 + t2' <- typeExpHtml t2 + return $ case pname of + Just v -> + parens (vnameHtml v <> ": " <> t1') <> " -> " <> t2' + Nothing -> + t1' <> " -> " <> t2' + +qualNameHtml :: QualName VName -> DocM Html +qualNameHtml (QualName names vname@(VName name tag)) = + if tag <= maxIntrinsicTag + then return $ renderName name + else f <$> ref + where prefix :: Html + prefix = mapM_ ((<> ".") . renderName . baseName) names + f (Just s) = H.a ! A.href (fromString s) $ prefix <> renderName name + f Nothing = prefix <> renderName name + + ref = do boring <- asks $ S.member vname . ctxNoLink + if boring + then return Nothing + else Just <$> vnameLink vname + +vnameLink' :: VName -> String -> String -> String +vnameLink :: VName -> DocM String +vnameLink vname = do + current <- asks ctxCurrent + file <- maybe current fst <$> asks (M.lookup vname . ctxFileMap) + return $ vnameLink' vname current file + +vnameLink' (VName _ tag) current file = + if file == current + then "#" ++ show tag + else relativise file current ++ ".html#" ++ show tag + +typeNameHtml :: TypeName -> DocM Html +typeNameHtml = qualNameHtml . qualNameFromTypeName + +patternHtml :: Pattern -> DocM Html +patternHtml pat = do + let (pat_param, t) = patternParam pat + t' <- typeHtml t + return $ case pat_param of + Just v -> parens (vnameHtml v <> ": " <> t') + Nothing -> t' + +relativise :: FilePath -> FilePath -> FilePath +relativise dest src = + concat (replicate (length (splitPath src) - 1) "../") ++ dest + +dimDeclHtml :: DimDecl VName -> DocM Html +dimDeclHtml AnyDim = return mempty +dimDeclHtml (NamedDim v) = qualNameHtml v +dimDeclHtml (ConstDim n) = return $ toHtml (show n) + +typeArgExpHtml :: TypeArgExp VName -> DocM Html +typeArgExpHtml (TypeArgExpDim d _) = dimDeclHtml d +typeArgExpHtml (TypeArgExpType d) = typeExpHtml d + +typeParamHtml :: TypeParam -> Html +typeParamHtml (TypeParamDim name _) = brackets $ vnameHtml name +typeParamHtml (TypeParamType Unlifted name _) = "'" <> vnameHtml name +typeParamHtml (TypeParamType Lifted name _) = "'^" <> vnameHtml name + +typeAbbrevHtml :: Liftedness -> Html -> [TypeParam] -> Html +typeAbbrevHtml l name params = + what <> name <> mconcat (map ((" "<>) . typeParamHtml) params) + where what = case l of Lifted -> keyword "type " <> "^" + Unlifted -> keyword "type " + +docHtml :: Maybe DocComment -> DocM Html +docHtml (Just (DocComment doc loc)) = + markdown def { msAddHeadingId = True } . LT.pack <$> identifierLinks loc doc +docHtml Nothing = return mempty + +identifierLinks :: SrcLoc -> String -> DocM String +identifierLinks _ [] = return [] +identifierLinks loc s + | Just ((name, namespace, file), s') <- identifierReference s = do + let proceed x = (x<>) <$> identifierLinks loc s' + unknown = proceed $ "`" <> name <> "`" + case knownNamespace namespace of + Just namespace' -> do + maybe_v <- lookupName (namespace', name, file) + case maybe_v of + Nothing -> do + warn loc $ + "Identifier '" <> name <> "' not found in namespace '" <> + namespace <> "'" <> maybe "" (" in file "<>) file <> "." + unknown + Just v' -> do + link <- vnameLink v' + proceed $ "[`" <> name <> "`](" <> link <> ")" + _ -> do + warn loc $ "Unknown namespace '" <> namespace <> "'." + unknown + where knownNamespace "term" = Just Term + knownNamespace "mtype" = Just Signature + knownNamespace "type" = Just Type + knownNamespace _ = Nothing +identifierLinks loc (c:s') = (c:) <$> identifierLinks loc s' + +lookupName :: (Namespace, String, Maybe FilePath) -> DocM (Maybe VName) +lookupName (namespace, name, file) = do + current <- asks ctxCurrent + let file' = includeToString . flip (mkImportFrom (mkInitialImport current)) noLoc <$> file + env <- lookupEnvForFile file' + case M.lookup (namespace, nameFromString name) . envNameMap =<< env of + Nothing -> return Nothing + Just qn -> return $ Just $ qualLeaf qn + +lookupEnvForFile :: Maybe FilePath -> DocM (Maybe Env) +lookupEnvForFile Nothing = asks $ Just . fileEnv . ctxFileMod +lookupEnvForFile (Just file) = asks $ fmap fileEnv . lookup file . ctxImports + +describeGeneric :: VName + -> IndexWhat + -> Maybe DocComment + -> (Html -> DocM Html) + -> DocM Html +describeGeneric name what doc f = do + name' <- H.span ! A.class_ "decl_name" <$> vnameDescDef name what + decl_type <- f name' + doc' <- docHtml doc + let decl_doc = H.dd ! A.class_ "desc_doc" $ doc' + decl_header = (H.dt ! A.class_ "desc_header") $ + vnameSynopsisRef name <> decl_type + return $ decl_header <> decl_doc + +describeGenericMod :: VName + -> IndexWhat + -> SigExp + -> Maybe DocComment + -> (Html -> DocM Html) + -> DocM Html +describeGenericMod name what se doc f = do + name' <- H.span ! A.class_ "decl_name" <$> vnameDescDef name what + + decl_type <- f name' + + doc' <- case se of + SigSpecs specs _ -> (<>) <$> docHtml doc <*> describeSpecs specs + _ -> docHtml doc + + let decl_doc = H.dd ! A.class_ "desc_doc" $ doc' + decl_header = (H.dt ! A.class_ "desc_header") $ + vnameSynopsisRef name <> decl_type + return $ decl_header <> decl_doc + +describeDecs :: [Dec] -> DocM Html +describeDecs decs = do + visible <- asks ctxVisibleMTys + H.dl . mconcat <$> + mapM (fmap $ H.div ! A.class_ "decl_description") + (mapMaybe (describeDec visible) decs) + +describeDec :: S.Set VName -> Dec -> Maybe (DocM Html) +describeDec _ (ValDec vb) = Just $ + describeGeneric (valBindName vb) (valBindWhat vb) (valBindDoc vb) $ \name -> do + (lhs, mhs, rhs) <- valBindHtml name vb + return $ lhs <> mhs <> ": " <> rhs + +describeDec _ (TypeDec vb) = Just $ + describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb) + +describeDec _ (SigDec (SigBind name se doc _)) = Just $ + describeGenericMod name IndexModuleType se doc $ \name' -> + return $ keyword "module type " <> name' + +describeDec _ (ModDec mb) = Just $ + describeGeneric (modName mb) IndexModule (modDoc mb) $ \name' -> + return $ keyword "module " <> name' + +describeDec _ OpenDec{} = Nothing + +describeDec visible (LocalDec (SigDec (SigBind name se doc _)) _) + | name `S.member` visible = Just $ + describeGenericMod name IndexModuleType se doc $ \name' -> + return $ keyword "local module type " <> name' + +describeDec _ LocalDec{} = Nothing + +valBindWhat :: ValBind -> IndexWhat +valBindWhat vb = + if null (valBindParams vb) && orderZero (unInfo (valBindRetType vb)) + then IndexValue + else IndexFunction + +describeSpecs :: [Spec] -> DocM Html +describeSpecs specs = + H.dl . mconcat <$> mapM describeSpec specs + +describeSpec :: Spec -> DocM Html +describeSpec (ValSpec name tparams t doc _) = + describeGeneric name what doc $ \name' -> do + let tparams' = mconcat $ map ((" "<>) . typeParamHtml) tparams + t' <- noLink (map typeParamName tparams) $ + typeExpHtml $ declaredType t + return $ keyword "val " <> name' <> tparams' <> ": " <> t' + where what = if orderZero (unInfo $ expandedType t) + then IndexValue else IndexFunction +describeSpec (TypeAbbrSpec vb) = + describeGeneric (typeAlias vb) IndexType (typeDoc vb) (`typeBindHtml` vb) +describeSpec (TypeSpec l name tparams doc _) = + describeGeneric name IndexType doc $ + return . (\name' -> typeAbbrevHtml l name' tparams) +describeSpec (ModSpec name se doc _) = + describeGenericMod name IndexModule se doc $ \name' -> + case se of + SigSpecs{} -> return $ keyword "module " <> name' + _ -> do se' <- synopsisSigExp se + return $ keyword "module " <> name' <> ": " <> se' +describeSpec (IncludeSpec sig _) = do + sig' <- synopsisSigExp sig + doc' <- docHtml Nothing + let decl_header = (H.dt ! A.class_ "desc_header") $ + (H.span ! A.class_ "synopsis_link") mempty <> + keyword "include " <> + sig' + decl_doc = H.dd ! A.class_ "desc_doc" $ doc' + return $ decl_header <> decl_doc diff --git a/src/Futhark/Doc/Html.hs b/src/Futhark/Doc/Html.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Doc/Html.hs @@ -0,0 +1,50 @@ +module Futhark.Doc.Html + ( primTypeHtml + , prettyTypeName + , prettyU + , renderName + , joinBy + , commas + , brackets + , braces + , parens + ) +where + +import Data.Semigroup ((<>)) + +import Language.Futhark +import Futhark.Util.Pretty (Doc,ppr) + +import qualified Text.PrettyPrint.Mainland as PP (pretty) +import Text.Blaze.Html5 (toHtml, Html) + +docToHtml :: Doc -> Html +docToHtml = toHtml . PP.pretty 80 + +primTypeHtml :: PrimType -> Html +primTypeHtml = docToHtml . ppr + +prettyTypeName :: TypeName -> Html +prettyTypeName et = (docToHtml . ppr) (baseName <$> qualNameFromTypeName et) + +prettyU :: Uniqueness -> Html +prettyU = docToHtml . ppr + +renderName :: Name -> Html +renderName name = docToHtml (ppr name) + +joinBy :: Html -> [Html] -> Html +joinBy _ [] = mempty +joinBy _ [x] = x +joinBy sep (x:xs) = x <> foldMap (sep <>) xs + +commas :: [Html] -> Html +commas = joinBy (toHtml ", ") + +parens :: Html -> Html +parens x = toHtml "(" <> x <> toHtml ")" +braces :: Html -> Html +braces x = toHtml "{" <> x <> toHtml "}" +brackets :: Html -> Html +brackets x = toHtml "[" <> x <> toHtml "]" diff --git a/src/Futhark/Error.hs b/src/Futhark/Error.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Error.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE FlexibleContexts #-} +-- | Futhark error definitions. +module Futhark.Error + ( CompilerError(..) + , ErrorClass(..) + + , externalError + , externalErrorS + + , InternalError + , internalError + , compilerBug + , compilerBugS + , compilerLimitation + , compilerLimitationS + ) +where + +import Control.Monad.Error.Class +import qualified Data.Text as T + +-- | There are two classes of internal errors: actual bugs, and +-- implementation limitations. The latter are already known and need +-- not be reported. +data ErrorClass = CompilerBug + | CompilerLimitation + deriving (Eq, Ord, Show) + +data CompilerError = + ExternalError T.Text + -- ^ An error that happened due to something the user did, such as + -- provide incorrect code or options. + | InternalError T.Text T.Text ErrorClass + -- ^ An internal compiler error. The second text is extra data + -- for debugging, which can be written to a file. + +instance Show CompilerError where + show (ExternalError s) = T.unpack s + show (InternalError s _ _) = T.unpack s + +externalError :: MonadError CompilerError m => T.Text -> m a +externalError = throwError . ExternalError + +externalErrorS :: MonadError CompilerError m => String -> m a +externalErrorS = externalError . T.pack + +-- | An error that is not the users fault, but a bug (or limitation) +-- in the compiler. Compiler passes should only ever report this +-- error - any problems after the type checker are *our* fault, not +-- the users. +data InternalError = Error ErrorClass T.Text + +compilerBug :: MonadError InternalError m => T.Text -> m a +compilerBug = throwError . Error CompilerBug + +compilerLimitation :: MonadError InternalError m => T.Text -> m a +compilerLimitation = throwError . Error CompilerLimitation + +internalError :: MonadError CompilerError m => InternalError -> T.Text -> m a +internalError (Error c s) t = throwError $ InternalError s t c + +compilerBugS :: MonadError InternalError m => String -> m a +compilerBugS = compilerBug . T.pack + +compilerLimitationS :: MonadError InternalError m => String -> m a +compilerLimitationS = compilerLimitation . T.pack diff --git a/src/Futhark/FreshNames.hs b/src/Futhark/FreshNames.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/FreshNames.hs @@ -0,0 +1,52 @@ +{-# LANGUAGE DeriveLift #-} +-- | This module provides facilities for generating unique names. +module Futhark.FreshNames + ( VNameSource + , blankNameSource + , newNameSource + , newName + , newVName + , newVNameFromName + ) where + +import qualified Data.Semigroup as Sem +import Language.Haskell.TH.Syntax (Lift) + +import Language.Futhark.Core + +-- | A name source is conceptually an infinite sequence of names with +-- no repeating entries. In practice, when asked for a name, the name +-- source will return the name along with a new name source, which +-- should then be used in place of the original. +-- +-- The 'Ord' instance is based on how many names have been extracted +-- from the name source. +newtype VNameSource = VNameSource Int + deriving (Lift, Eq, Ord) + +instance Sem.Semigroup VNameSource where + VNameSource x <> VNameSource y = VNameSource (x `max` y) + +instance Monoid VNameSource where + mempty = blankNameSource + mappend = (Sem.<>) + +-- | Produce a fresh name, using the given name as a template. +newName :: VNameSource -> VName -> (VName, VNameSource) +newName (VNameSource i) k = (VName (baseName k) i, VNameSource (i+1)) + +-- | A blank name source. +blankNameSource :: VNameSource +blankNameSource = newNameSource 0 + +-- | A new name source that starts counting from the given number. +newNameSource :: Int -> VNameSource +newNameSource = VNameSource + +-- | Produce a fresh 'VName', using the given base name as a template. +newVName :: VNameSource -> String -> (VName, VNameSource) +newVName src = newVNameFromName src . nameFromString + +-- | Produce a fresh 'VName', using the given base name as a template. +newVNameFromName :: VNameSource -> Name -> (VName, VNameSource) +newVNameFromName src s = newName src $ VName s 0 diff --git a/src/Futhark/Internalise.hs b/src/Futhark/Internalise.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise.hs @@ -0,0 +1,1704 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +-- | +-- +-- This module implements a transformation from source to core +-- Futhark. +-- +module Futhark.Internalise (internaliseProg) where + +import Control.Monad.State +import Control.Monad.Reader +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Semigroup ((<>)) +import Data.List +import Data.Loc +import Data.Char (chr) + +import Language.Futhark as E hiding (TypeArg) +import Language.Futhark.Semantic (Imports) +import Futhark.Representation.SOACS as I hiding (stmPattern) +import Futhark.Transform.Rename as I +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Representation.AST.Attributes.Aliases +import qualified Futhark.Analysis.Alias as Alias +import Futhark.Util (splitAt3) + +import Futhark.Internalise.Monad as I +import Futhark.Internalise.AccurateSizes +import Futhark.Internalise.TypesValues +import Futhark.Internalise.Bindings +import Futhark.Internalise.Lambdas +import Futhark.Internalise.Defunctorise as Defunctorise +import Futhark.Internalise.Defunctionalise as Defunctionalise +import Futhark.Internalise.Monomorphise as Monomorphise + +-- | Convert a program in source Futhark to a program in the Futhark +-- core language. +internaliseProg :: MonadFreshNames m => + Bool -> Imports -> m (Either String I.Prog) +internaliseProg always_safe prog = do + prog_decs <- Defunctorise.transformProg prog + prog_decs' <- Monomorphise.transformProg prog_decs + prog_decs'' <- Defunctionalise.transformProg prog_decs' + prog' <- fmap (fmap I.Prog) $ runInternaliseM always_safe $ internaliseValBinds prog_decs'' + traverse I.renameProg prog' + +internaliseValBinds :: [E.ValBind] -> InternaliseM () +internaliseValBinds = mapM_ internaliseValBind + +internaliseFunName :: VName -> [E.Pattern] -> InternaliseM Name +internaliseFunName ofname [] = return $ nameFromString $ pretty ofname ++ "f" +internaliseFunName ofname _ = do + info <- lookupFunction' ofname + -- In some rare cases involving local functions, the same function + -- name may be re-used in multiple places. We check whether the + -- function name has already been used, and generate a new one if + -- so. + case info of + Just _ -> nameFromString . pretty <$> newNameFromString (baseString ofname) + Nothing -> return $ nameFromString $ pretty ofname + +internaliseValBind :: E.ValBind -> InternaliseM () +internaliseValBind fb@(E.ValBind entry fname retdecl (Info rettype) tparams params body _ loc) = do + info <- bindingParams tparams params $ \pcm shapeparams params' -> do + (rettype_bad, rcm) <- internaliseReturnType rettype + let rettype' = zeroExts rettype_bad + + let mkConstParam name = Param name $ I.Prim int32 + constparams = map (mkConstParam . snd) $ pcm<>rcm + constnames = map I.paramName constparams + constscope = M.fromList $ zip constnames $ repeat $ + FParamInfo $ I.Prim $ IntType Int32 + + shapenames = map I.paramName shapeparams + normal_params = map I.paramName constparams ++ shapenames ++ + map I.paramName (concat params') + normal_param_names = S.fromList normal_params + + fname' <- internaliseFunName fname params + + body' <- localScope constscope $ do + msg <- case retdecl of + Just dt -> ErrorMsg . + ("Function return value does not match shape of type ":) <$> + typeExpForError rcm dt + Nothing -> return $ ErrorMsg ["Function return value does not match shape of declared return type."] + internaliseBody body >>= + ensureResultExtShape asserting msg loc (map I.fromDecl rettype') + + let free_in_fun = freeInBody body' `S.difference` normal_param_names + + used_free_params <- forM (S.toList free_in_fun) $ \v -> do + v_t <- lookupType v + return $ Param v $ toDecl v_t Nonunique + + let free_shape_params = map (`Param` I.Prim int32) $ + concatMap (I.shapeVars . I.arrayShape . I.paramType) used_free_params + free_params = nub $ free_shape_params ++ used_free_params + all_params = constparams ++ free_params ++ shapeparams ++ concat params' + + addFunction $ I.FunDef Nothing fname' rettype' all_params body' + + return (fname', + pcm<>rcm, + map I.paramName free_params, + shapenames, + map declTypeOf $ concat params', + all_params, + applyRetType rettype' all_params) + + bindFunction fname info + when entry $ generateEntryPoint fb + + where + -- | Recompute existential sizes to start from zero. + -- Necessary because some convoluted constructions will start + -- them from somewhere else. + zeroExts ts = generaliseExtTypes ts ts + +generateEntryPoint :: E.ValBind -> InternaliseM () +generateEntryPoint (E.ValBind _ ofname retdecl (Info rettype) _ orig_params _ _ loc) = + -- We remove all shape annotations, so there should be no constant + -- parameters here. + bindingParams [] (map E.patternNoShapeAnnotations params) $ + \_ shapeparams params' -> do + (entry_rettype, _) <- internaliseEntryReturnType $ E.vacuousShapeAnnotations rettype + let entry' = entryPoint (zip params params') (retdecl, rettype, entry_rettype) + args = map (I.Var . I.paramName) $ concat params' + + entry_body <- insertStmsM $ do + vals <- fst <$> funcall "entry_result" (E.qualName ofname) args loc + ctx <- extractShapeContext (concat entry_rettype) <$> + mapM (fmap I.arrayDims . subExpType) vals + resultBodyM (ctx ++ vals) + + addFunction $ + I.FunDef (Just entry') (baseName ofname) + (concat entry_rettype) + (shapeparams ++ concat params') entry_body + + -- XXX: We massage the parameters a little bit to handle the case + -- where there is just a single parameter that is a tuple. This is + -- wide-spread in existing Futhark code, although I'd like to get + -- rid of it. + where params = case orig_params of + [TuplePattern ps _] -> ps + _ -> orig_params + +entryPoint :: [(E.Pattern,[I.FParam])] + -> (Maybe (E.TypeExp VName), E.StructType, [[I.TypeBase ExtShape Uniqueness]]) + -> EntryPoint +entryPoint params (retdecl, eret, crets) = + (concatMap (entryPointType . preParam) params, + case isTupleRecord eret of + Just ts -> concatMap entryPointType $ zip3 retdecls ts crets + _ -> entryPointType (retdecl, eret, concat crets)) + where preParam (p_pat, ps) = (paramOuterType p_pat, + E.patternStructType p_pat, + staticShapes $ map I.paramDeclType ps) + paramOuterType (E.PatternAscription _ tdecl _) = Just $ declaredType tdecl + paramOuterType (E.PatternParens p _) = paramOuterType p + paramOuterType _ = Nothing + + retdecls = case retdecl of Just (TETuple tes _) -> map Just tes + _ -> repeat Nothing + + entryPointType :: (Maybe (E.TypeExp VName), + E.StructType, + [I.TypeBase ExtShape Uniqueness]) + -> [EntryPointType] + entryPointType (_, E.Prim E.Unsigned{}, _) = + [I.TypeUnsigned] + entryPointType (_, E.Array (ArrayPrimElem Unsigned{} _) _ _, _) = + [I.TypeUnsigned] + entryPointType (_, E.Prim{}, _) = + [I.TypeDirect] + entryPointType (_, E.Array ArrayPrimElem{} _ _, _) = + [I.TypeDirect] + entryPointType (te, t, ts) = + [I.TypeOpaque desc $ length ts] + where desc = maybe (pretty t') pretty te + t' = removeShapeAnnotations t `E.setUniqueness` Nonunique + +internaliseIdent :: E.Ident -> InternaliseM I.VName +internaliseIdent (E.Ident name (Info tp) loc) = + case tp of + E.Prim{} -> return name + _ -> fail $ "Futhark.Internalise.internaliseIdent: asked to internalise non-prim-typed ident '" + ++ pretty name ++ " of type " ++ pretty tp ++ + " at " ++ locStr loc ++ "." + +internaliseBody :: E.Exp -> InternaliseM Body +internaliseBody e = insertStmsM $ resultBody <$> internaliseExp "res" e + +internaliseBodyStms :: E.Exp -> ([SubExp] -> InternaliseM (Body, a)) + -> InternaliseM (Body, a) +internaliseBodyStms e m = do + ((Body _ bnds res,x), otherbnds) <- + collectStms $ m =<< internaliseExp "res" e + (,x) <$> mkBodyM (otherbnds <> bnds) res + +internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp] + +internaliseExp desc (E.Parens e _) = + internaliseExp desc e + +internaliseExp desc (E.QualParens _ e _) = + internaliseExp desc e + +internaliseExp _ (E.Var (E.QualName _ name) (Info t) loc) = do + subst <- asks $ M.lookup name . envSubsts + case subst of + Just substs -> return substs + Nothing -> do + -- If this identifier is the name of a constant, we have to turn it + -- into a call to the corresponding function. + is_const <- lookupConstant loc name + case is_const of + Just ses -> return ses + Nothing -> (:[]) . I.Var <$> internaliseIdent (E.Ident name (Info t') loc) + where t' = removeShapeAnnotations t + +internaliseExp desc (E.Index e idxs _ loc) = do + vs <- internaliseExpToVars "indexed" e + dims <- case vs of [] -> return [] -- Will this happen? + v:_ -> I.arrayDims <$> lookupType v + (idxs', cs) <- internaliseSlice loc dims idxs + let index v = do v_t <- lookupType v + return $ I.BasicOp $ I.Index v $ fullSlice v_t idxs' + certifying cs $ letSubExps desc =<< mapM index vs + +internaliseExp desc (E.TupLit es _) = + concat <$> mapM (internaliseExp desc) es + +internaliseExp desc (E.RecordLit orig_fields _) = + concatMap snd . sortFields . M.unions . reverse <$> mapM internaliseField orig_fields + where internaliseField (E.RecordFieldExplicit name e _) = + M.singleton name <$> internaliseExp desc e + internaliseField (E.RecordFieldImplicit name t loc) = + internaliseField $ E.RecordFieldExplicit (baseName name) + (E.Var (E.qualName name) (vacuousShapeAnnotations <$> t) loc) loc + +internaliseExp desc (E.ArrayLit es (Info arr_t) loc) + -- If this is a multidimensional array literal of primitives, we + -- treat it specially by flattening it out followed by a reshape. + -- This cuts down on the amount of statements that are produced, and + -- thus allows us to efficiently handle huge array literals - a + -- corner case, but an important one. + | Just ((eshape,e'):es') <- mapM isArrayLiteral es, + not $ null eshape, + all ((eshape==) . fst) es', + Just basetype <- E.peelArray (length eshape) arr_t = do + let flat_lit = E.ArrayLit (e' ++ concatMap snd es') (Info basetype) loc + new_shape = length es:eshape + flat_arrs <- internaliseExpToVars "flat_literal" flat_lit + forM flat_arrs $ \flat_arr -> do + flat_arr_t <- lookupType flat_arr + let new_shape' = reshapeOuter (map (DimNew . constant) new_shape) + (length new_shape) $ arrayShape flat_arr_t + letSubExp desc $ I.BasicOp $ I.Reshape new_shape' flat_arr + + | otherwise = do + es' <- mapM (internaliseExp "arr_elem") es + case es' of + [] -> do + rowtypes <- internaliseType (rowtype `setAliases` ()) + let arraylit rt = I.BasicOp $ I.ArrayLit [] rt + letSubExps desc $ map (arraylit . zeroDim . fromDecl) rowtypes + e' : _ -> do + rowtypes <- mapM subExpType e' + let arraylit ks rt = do + ks' <- mapM (ensureShape asserting "shape of element differs from shape of first element" + loc rt "elem_reshaped") ks + return $ I.BasicOp $ I.ArrayLit ks' rt + letSubExps desc =<< zipWithM arraylit (transpose es') rowtypes + where rowtype = E.stripArray 1 arr_t + + zeroDim t = t `I.setArrayShape` + I.Shape (replicate (I.arrayRank t) (constant (0::Int32))) + + isArrayLiteral :: E.Exp -> Maybe ([Int],[E.Exp]) + isArrayLiteral (E.ArrayLit inner_es _ _) = do + (eshape,e):inner_es' <- mapM isArrayLiteral inner_es + guard $ all ((eshape==) . fst) inner_es' + return (length inner_es:eshape, e ++ concatMap snd inner_es') + isArrayLiteral e = + Just ([], [e]) + +internaliseExp desc (E.Range start maybe_second end _ _) = do + start' <- internaliseExp1 "range_start" start + end' <- internaliseExp1 "range_end" $ case end of + DownToExclusive e -> e + ToInclusive e -> e + UpToExclusive e -> e + + (it, le_op, lt_op) <- + case E.typeOf start of + E.Prim (E.Signed it) -> return (it, CmpSle it, CmpSlt it) + E.Prim (E.Unsigned it) -> return (it, CmpUle it, CmpUlt it) + start_t -> fail $ "Start value in range has type " ++ pretty start_t + + let one = intConst it 1 + negone = intConst it (-1) + default_step = case end of DownToExclusive{} -> negone + ToInclusive{} -> one + UpToExclusive{} -> one + + (step, step_zero) <- case maybe_second of + Just second -> do + second' <- internaliseExp1 "range_second" second + subtracted_step <- letSubExp "subtracted_step" $ I.BasicOp $ I.BinOp (I.Sub it) second' start' + step_zero <- letSubExp "step_zero" $ I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) start' second' + return (subtracted_step, step_zero) + Nothing -> + return (default_step, constant False) + + step_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum it) step + step_sign_i32 <- asIntS Int32 step_sign + + bounds_invalid_downwards <- letSubExp "bounds_invalid_downwards" $ + I.BasicOp $ I.CmpOp le_op start' end' + bounds_invalid_upwards <- letSubExp "bounds_invalid_upwards" $ + I.BasicOp $ I.CmpOp lt_op end' start' + + (distance, step_wrong_dir, bounds_invalid) <- case end of + DownToExclusive{} -> do + step_wrong_dir <- letSubExp "step_wrong_dir" $ + I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign one + distance <- letSubExp "distance" $ + I.BasicOp $ I.BinOp (Sub it) start' end' + distance_i32 <- asIntZ Int32 distance + return (distance_i32, step_wrong_dir, bounds_invalid_downwards) + UpToExclusive{} -> do + step_wrong_dir <- letSubExp "step_wrong_dir" $ + I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign negone + distance <- letSubExp "distance" $ I.BasicOp $ I.BinOp (Sub it) end' start' + distance_i32 <- asIntZ Int32 distance + return (distance_i32, step_wrong_dir, bounds_invalid_upwards) + ToInclusive{} -> do + downwards <- letSubExp "downwards" $ + I.BasicOp $ I.CmpOp (I.CmpEq $ IntType it) step_sign negone + distance_downwards_exclusive <- + letSubExp "distance_downwards_exclusive" $ + I.BasicOp $ I.BinOp (Sub it) start' end' + distance_upwards_exclusive <- + letSubExp "distance_upwards_exclusive" $ + I.BasicOp $ I.BinOp (Sub it) end' start' + + bounds_invalid <- letSubExp "bounds_invalid" $ + I.If downwards + (resultBody [bounds_invalid_downwards]) + (resultBody [bounds_invalid_upwards]) $ + ifCommon [I.Prim I.Bool] + distance_exclusive <- letSubExp "distance_exclusive" $ + I.If downwards + (resultBody [distance_downwards_exclusive]) + (resultBody [distance_upwards_exclusive]) $ + ifCommon [I.Prim $ IntType it] + distance_exclusive_i32 <- asIntZ Int32 distance_exclusive + distance <- letSubExp "distance" $ + I.BasicOp $ I.BinOp (Add Int32) + distance_exclusive_i32 (intConst Int32 1) + return (distance, constant False, bounds_invalid) + + step_invalid <- letSubExp "step_invalid" $ + I.BasicOp $ I.BinOp I.LogOr step_wrong_dir step_zero + invalid <- letSubExp "range_invalid" $ + I.BasicOp $ I.BinOp I.LogOr step_invalid bounds_invalid + + step_i32 <- asIntS Int32 step + pos_step <- letSubExp "pos_step" $ + I.BasicOp $ I.BinOp (Mul Int32) step_i32 step_sign_i32 + num_elems <- letSubExp "num_elems" =<< + eIf (eSubExp invalid) + (eBody [eSubExp $ intConst Int32 0]) + (eBody [eDivRoundingUp Int32 (eSubExp distance) (eSubExp pos_step)]) + pure <$> letSubExp desc (I.BasicOp $ I.Iota num_elems start' step it) + +internaliseExp desc (E.Ascript e (TypeDecl dt (Info et)) loc) = do + es <- internaliseExp desc e + (ts, cm) <- internaliseReturnType et + mapM_ (uncurry (internaliseDimConstant loc)) cm + dt' <- typeExpForError cm dt + forM (zip es ts) $ \(e',t') -> do + dims <- arrayDims <$> subExpType e' + let parts = ["Value of (core language) shape ("] ++ + intersperse ", " (map ErrorInt32 dims) ++ + [") cannot match shape of type `"] ++ dt' ++ ["`."] + ensureExtShape asserting (ErrorMsg parts) loc (I.fromDecl t') desc e' + +internaliseExp desc (E.Negate e _) = do + e' <- internaliseExp1 "negate_arg" e + et <- subExpType e' + case et of I.Prim (I.IntType t) -> + letTupExp' desc $ I.BasicOp $ I.BinOp (I.Sub t) (I.intConst t 0) e' + I.Prim (I.FloatType t) -> + letTupExp' desc $ I.BasicOp $ I.BinOp (I.FSub t) (I.floatConst t 0) e' + _ -> fail "Futhark.Internalise.internaliseExp: non-numeric type in Negate" + +internaliseExp desc e@E.Apply{} = do + (qfname, args, _) <- findFuncall e + let fname = nameFromString $ pretty $ baseName $ qualLeaf qfname + loc = srclocOf e + + -- Some functions are magical (overloaded) and we handle that here. + -- Note that polymorphic functions (which are not magical) are not + -- handled here. + case () of + () | Just internalise <- isOverloadedFunction qfname args loc -> + internalise desc + | Just (rettype, _) <- M.lookup fname I.builtInFunctions -> do + let tag ses = [ (se, I.Observe) | se <- ses ] + args' <- mapM (internaliseExp "arg") args + let args'' = concatMap tag args' + letTupExp' desc $ I.Apply fname args'' [I.Prim rettype] (Safe, loc, []) + | otherwise -> do + args' <- concat <$> mapM (internaliseExp "arg") args + fst <$> funcall desc qfname args' loc + +internaliseExp desc (E.LetPat tparams pat e body loc) = do + ses <- internaliseExp desc e + t <- I.staticShapes <$> mapM I.subExpType ses + stmPattern tparams pat t $ \cm pat_names match -> do + mapM_ (uncurry (internaliseDimConstant loc)) cm + ses' <- match loc ses + forM_ (zip pat_names ses') $ \(v,se) -> + letBindNames_ [v] $ I.BasicOp $ I.SubExp se + internaliseExp desc body + +internaliseExp desc (E.LetFun ofname (tparams, params, retdecl, Info rettype, body) letbody loc) = do + internaliseValBind $ E.ValBind False ofname retdecl (Info rettype) tparams params body Nothing loc + internaliseExp desc letbody + +internaliseExp desc (E.DoLoop tparams mergepat mergeexp form loopbody loc) = do + -- We pretend that we saw a let-binding first to ensure that the + -- initial values for the merge parameters match their annotated + -- sizes + ses <- internaliseExp "loop_init" mergeexp + t <- I.staticShapes <$> mapM I.subExpType ses + stmPattern tparams mergepat t $ \cm mergepat_names match -> do + mapM_ (uncurry (internaliseDimConstant loc)) cm + ses' <- match (srclocOf mergepat) ses + forM_ (zip mergepat_names ses') $ \(v,se) -> + letBindNames_ [v] $ I.BasicOp $ I.SubExp se + let mergeinit = map I.Var mergepat_names + + (loopbody', (form', shapepat, mergepat', mergeinit', pre_stms)) <- + handleForm mergeinit form + + addStms pre_stms + + mergeinit_ts' <- mapM subExpType mergeinit' + + let ctxinit = argShapes + (map I.paramName shapepat) + (map I.paramType mergepat') + mergeinit_ts' + ctxmerge = zip shapepat ctxinit + valmerge = zip mergepat' mergeinit' + merge = ctxmerge ++ valmerge + dropCond = case form of E.While{} -> drop 1 + _ -> id + + -- Ensure that the result of the loop matches the shapes of the + -- merge parameters, if any have been annotated by programmer. + let merge_names = map (I.paramName . fst) merge + merge_ts = existentialiseExtTypes merge_names $ + staticShapes $ map (I.paramType . fst) merge + loopbody'' <- localScope (scopeOfFParams $ map fst merge) $ + ensureResultExtShapeNoCtx asserting + "shape of loop result does not match shapes in loop parameters" + loc merge_ts loopbody' + + loop_res <- letTupExp desc $ I.DoLoop ctxmerge valmerge form' loopbody'' + return $ map I.Var $ dropCond loop_res + + where + forLoop nested_mergepat shapepat mergeinit form' = + inScopeOf form' $ internaliseBodyStms loopbody $ \ses -> do + sets <- mapM subExpType ses + let mergepat' = concat nested_mergepat + shapeargs = argShapes + (map I.paramName shapepat) + (map I.paramType mergepat') + sets + return (resultBody $ shapeargs ++ ses, + (form', + shapepat, + mergepat', + mergeinit, + mempty)) + + + handleForm mergeinit (E.ForIn x arr) = do + arr' <- internaliseExpToVars "for_in_arr" arr + arr_ts <- mapM lookupType arr' + let w = arraysSize 0 arr_ts + + i <- newVName "i" + + bindingParams tparams [mergepat] $ \mergecm shapepat nested_mergepat -> + bindingLambdaParams [] [x] (map rowType arr_ts) $ \x_cm x_params -> do + mapM_ (uncurry (internaliseDimConstant loc)) x_cm + mapM_ (uncurry (internaliseDimConstant loc)) mergecm + let loopvars = zip x_params arr' + forLoop nested_mergepat shapepat mergeinit $ I.ForLoop i Int32 w loopvars + + handleForm mergeinit (E.For i num_iterations) = do + num_iterations' <- internaliseExp1 "upper_bound" num_iterations + i' <- internaliseIdent i + num_iterations_t <- I.subExpType num_iterations' + it <- case num_iterations_t of + I.Prim (IntType it) -> return it + _ -> fail "internaliseExp DoLoop: invalid type" + + bindingParams tparams [mergepat] $ \mergecm shapepat nested_mergepat -> do + mapM_ (uncurry (internaliseDimConstant loc)) mergecm + forLoop nested_mergepat shapepat mergeinit $ I.ForLoop i' it num_iterations' [] + + handleForm mergeinit (E.While cond) = + bindingParams tparams [mergepat] $ \mergecm shapepat nested_mergepat -> do + mergeinit_ts <- mapM subExpType mergeinit + mapM_ (uncurry (internaliseDimConstant loc)) mergecm + let mergepat' = concat nested_mergepat + -- We need to insert 'cond' twice - once for the initial + -- condition (do we enter the loop at all?), and once with + -- the result values of the loop (do we continue into the + -- next iteration?). This is safe, as the type rules for + -- the external language guarantees that 'cond' does not + -- consume anything. + let shapeinit = argShapes + (map I.paramName shapepat) + (map I.paramType mergepat') + mergeinit_ts + + (loop_initial_cond, init_loop_cond_bnds) <- collectStms $ do + forM_ (zip shapepat shapeinit) $ \(p, se) -> + letBindNames_ [paramName p] $ BasicOp $ SubExp se + forM_ (zip (concat nested_mergepat) mergeinit) $ \(p, se) -> + unless (se == I.Var (paramName p)) $ + letBindNames_ [paramName p] $ BasicOp $ + case se of I.Var v | not $ primType $ paramType p -> + Reshape (map DimCoercion $ arrayDims $ paramType p) v + _ -> SubExp se + internaliseExp1 "loop_cond" cond + + internaliseBodyStms loopbody $ \ses -> do + sets <- mapM subExpType ses + loop_while <- newParam "loop_while" $ I.Prim I.Bool + let shapeargs = argShapes + (map I.paramName shapepat) + (map I.paramType mergepat') + sets + + -- Careful not to clobber anything. + loop_end_cond_body <- renameBody <=< insertStmsM $ do + forM_ (zip shapepat shapeargs) $ \(p, se) -> + unless (se == I.Var (paramName p)) $ + letBindNames_ [paramName p] $ BasicOp $ SubExp se + forM_ (zip (concat nested_mergepat) ses) $ \(p, se) -> + unless (se == I.Var (paramName p)) $ + letBindNames_ [paramName p] $ BasicOp $ + case se of I.Var v | not $ primType $ paramType p -> + Reshape (map DimCoercion $ arrayDims $ paramType p) v + _ -> SubExp se + resultBody <$> internaliseExp "loop_cond" cond + loop_end_cond <- bodyBind loop_end_cond_body + + return (resultBody $ shapeargs++loop_end_cond++ses, + (I.WhileLoop $ I.paramName loop_while, + shapepat, + loop_while : mergepat', + loop_initial_cond : mergeinit, + init_loop_cond_bnds)) + +internaliseExp desc (E.LetWith name src idxs ve body loc) = do + srcs <- internaliseExpToVars "src" $ + E.Var (qualName (E.identName src)) (vacuousShapeAnnotations <$> E.identType src) + (srclocOf src) + ves <- internaliseExp "lw_val" ve + dims <- case srcs of + [] -> return [] -- Will this happen? + v:_ -> I.arrayDims <$> lookupType v + (idxs', cs) <- internaliseSlice loc dims idxs + let comb sname ve' = do + sname_t <- lookupType sname + let slice = fullSlice sname_t idxs' + rowtype = sname_t `setArrayDims` sliceDims slice + ve'' <- ensureShape asserting "shape of value does not match shape of source array" + loc rowtype "lw_val_correct_shape" ve' + certifying cs $ + letInPlace "letwith_dst" sname (fullSlice sname_t idxs') $ BasicOp $ SubExp ve'' + dsts <- zipWithM comb srcs ves + dstt <- I.staticShapes <$> mapM lookupType dsts + let pat = E.Id (E.identName name) + (E.vacuousShapeAnnotations <$> E.identType name) + (srclocOf name) + stmPattern [] pat dstt $ \cm pat_names match -> do + mapM_ (uncurry (internaliseDimConstant loc)) cm + dsts' <- match loc $ map I.Var dsts + forM_ (zip pat_names dsts') $ \(v,dst) -> + letBindNames_ [v] $ I.BasicOp $ I.SubExp dst + internaliseExp desc body + +internaliseExp desc (E.Update src slice ve loc) = do + src_name <- newVName "update_src" + dest_name <- newVName "update_dest" + let src_t = E.typeOf src + src_ident = E.Ident src_name (E.Info src_t) loc + dest_ident = E.Ident dest_name (E.Info src_t) loc + + internaliseExp desc $ + E.LetPat [] (E.Id src_name (E.Info $ E.vacuousShapeAnnotations src_t) loc) src + (E.LetWith dest_ident src_ident slice ve + (E.Var (E.qualName dest_name) (E.Info (E.vacuousShapeAnnotations src_t)) loc) + loc) + loc + +internaliseExp desc (E.RecordUpdate src fields ve _ _) = do + src' <- internaliseExp desc src + ve' <- internaliseExp desc ve + replace (E.typeOf src `setAliases` ()) fields ve' src' + where replace (E.Record m) (f:fs) ve' src' + | Just t <- M.lookup f m = do + i <- fmap sum $ mapM (internalisedTypeSize . snd) $ + takeWhile ((/=f) . fst) $ sortFields m + k <- internalisedTypeSize t + let (bef, to_update, aft) = splitAt3 i k src' + src'' <- replace t fs ve' to_update + return $ bef ++ src'' ++ aft + replace _ _ ve' _ = return ve' + +internaliseExp desc (E.Unzip e _ _) = + internaliseExp desc e + +internaliseExp desc (E.Unsafe e _) = + local (\env -> env { envDoBoundsChecks = False }) $ + internaliseExp desc e + +internaliseExp desc (E.Assert e1 e2 (Info check) loc) = do + e1' <- internaliseExp1 "assert_cond" e1 + c <- assertingOne $ letExp "assert_c" $ + I.BasicOp $ I.Assert e1' (ErrorMsg [ErrorString check]) (loc, mempty) + -- Make sure there are some bindings to certify. + certifying c $ mapM rebind =<< internaliseExp desc e2 + where rebind v = do + v' <- newVName "assert_res" + letBindNames_ [v'] $ I.BasicOp $ I.SubExp v + return $ I.Var v' + +internaliseExp _ (E.Zip _ e es _ loc) = do + e' <- internaliseExpToVars "zip_arg" $ TupLit (e:es) loc + case e' of + e_key:es_unchecked -> do + -- We will reshape all of es_unchecked' to have the same outer + -- size as ts. We will not change any of the inner dimensions. + -- This will cause a runtime error if the outer sizes do not match, + -- thus preserving the semantics of zip(). + w <- arraySize 0 <$> lookupType e_key + let reshapeToOuter e_unchecked' = do + unchecked_t <- lookupType e_unchecked' + case I.arrayDims unchecked_t of + outer:inner | w /= outer -> do + cmp <- letSubExp "zip_cmp" $ I.BasicOp $ + I.CmpOp (I.CmpEq I.int32) w outer + c <- assertingOne $ + letExp "zip_assert" $ I.BasicOp $ + I.Assert cmp "arrays differ in length" (loc, mempty) + certifying c $ letExp (postfix e_unchecked' "_zip_res") $ + shapeCoerce (w:inner) e_unchecked' + _ -> return e_unchecked' + es' <- mapM reshapeToOuter es_unchecked + return $ map I.Var $ e_key : es' + [] -> return [] + + where postfix i s = baseString i ++ s + +internaliseExp desc (E.Map lam arr _ _) = do + arr' <- internaliseExpToVars "map_arr" arr + lam' <- internaliseMapLambda internaliseLambda lam $ map I.Var arr' + w <- arraysSize 0 <$> mapM lookupType arr' + letTupExp' desc $ I.Op $ + I.Screma w (I.mapSOAC lam') arr' + +internaliseExp desc (E.Reduce comm lam ne arr loc) = + internaliseScanOrReduce desc "reduce" reduce (lam, ne, arr, loc) + where reduce w red_lam nes arrs = + I.Screma w <$> I.reduceSOAC comm red_lam nes <*> pure arrs + +internaliseExp desc (E.GenReduce hist op ne buckets img loc) = do + ne' <- internaliseExp "gen_reduce_ne" ne + hist' <- internaliseExpToVars "gen_reduce_hist" hist + buckets' <- letExp "gen_reduce_buckets" . BasicOp . SubExp =<< + internaliseExp1 "gen_reduce_buckets" buckets + img' <- internaliseExpToVars "gen_reduce_img" img + + -- reshape neutral element to have same size as the destination array + ne_shp <- forM (zip ne' hist') $ \(n, h) -> do + rowtype <- I.stripArray 1 <$> lookupType h + ensureShape asserting + "Row shape of destination array does not match shape of neutral element" + loc rowtype "gen_reduce_ne_right_shape" n + ne_ts <- mapM I.subExpType ne_shp + his_ts <- mapM lookupType hist' + op' <- internaliseFoldLambda internaliseLambda op ne_ts his_ts + + -- reshape return type of bucket function to have same size as neutral element + -- (modulo the index) + bucket_param <- newParam "bucket_p" $ I.Prim int32 + img_params <- mapM (newParam "img_p" . rowType) =<< mapM lookupType img' + let params = bucket_param : img_params + rettype = I.Prim int32 : ne_ts + body = mkBody mempty $ map (I.Var . paramName) params + body' <- localScope (scopeOfLParams params) $ + ensureResultShape asserting + "Row shape of value array does not match row shape of gen_reduce target" + (srclocOf img) rettype body + + -- get sizes of histogram and image arrays + w_hist <- arraysSize 0 <$> mapM lookupType hist' + w_img <- arraysSize 0 <$> mapM lookupType img' + + -- Generate an assertion and reshapes to ensure that buckets' and + -- img' are the same size. + b_shape <- arrayShape <$> lookupType buckets' + let b_w = shapeSize 0 b_shape + cmp <- letSubExp "bucket_cmp" $ I.BasicOp $ I.CmpOp (I.CmpEq I.int32) b_w w_img + c <- assertingOne $ + letExp "bucket_cert" $ I.BasicOp $ + I.Assert cmp "length of index and value array does not match" (loc, mempty) + buckets'' <- certifying c $ letExp (baseString buckets') $ + I.BasicOp $ I.Reshape (reshapeOuter [DimCoercion w_img] 1 b_shape) buckets' + + letTupExp' desc $ I.Op $ + I.GenReduce w_img [GenReduceOp w_hist hist' ne_shp op'] (I.Lambda params body' rettype) $ buckets'' : img' + +internaliseExp desc (E.Scan lam ne arr loc) = + internaliseScanOrReduce desc "scan" scan (lam, ne, arr, loc) + where scan w scan_lam nes arrs = + I.Screma w <$> I.scanSOAC scan_lam nes <*> pure arrs + +internaliseExp _ (E.Filter lam arr _) = do + arrs <- internaliseExpToVars "filter_input" arr + lam' <- internalisePartitionLambda internaliseLambda 1 lam $ map I.Var arrs + uncurry (++) <$> partitionWithSOACS 1 lam' arrs + +internaliseExp _ (E.Partition k lam arr _) = do + arrs <- internaliseExpToVars "partition_input" arr + lam' <- internalisePartitionLambda internaliseLambda k lam $ map I.Var arrs + uncurry (++) <$> partitionWithSOACS k lam' arrs + +internaliseExp desc (E.Stream (E.MapLike o) lam arr _) = do + arrs <- internaliseExpToVars "stream_input" arr + lam' <- internaliseStreamMapLambda internaliseLambda lam $ map I.Var arrs + w <- arraysSize 0 <$> mapM lookupType arrs + let form = I.Parallel o Commutative (I.Lambda [] (mkBody mempty []) []) [] + letTupExp' desc $ I.Op $ I.Stream w form lam' arrs + +-- If the stream form is a reduce, we also have to fiddle with the +-- lambda to incorporate the reduce function. FIXME: can't we just +-- modify the internal representation of reduction streams? +internaliseExp desc (E.Stream (E.RedLike o comm lam0) lam arr _) = do + arrs <- internaliseExpToVars "stream_input" arr + rowts <- mapM (fmap I.rowType . lookupType) arrs + (lam_params, lam_body) <- + internaliseStreamLambda internaliseLambda lam rowts + let (chunk_param, _, lam_val_params) = + partitionChunkedFoldParameters 0 lam_params + + -- Synthesize neutral elements by applying the fold function + -- to an empty chunk. + letBindNames_ [I.paramName chunk_param] $ + I.BasicOp $ I.SubExp $ constant (0::Int32) + forM_ lam_val_params $ \p -> + letBindNames_ [I.paramName p] $ + I.BasicOp $ I.Scratch (I.elemType $ I.paramType p) $ + I.arrayDims $ I.paramType p + accs <- bodyBind =<< renameBody lam_body + + acctps <- mapM I.subExpType accs + outsz <- arraysSize 0 <$> mapM lookupType arrs + let acc_arr_tps = [ I.arrayOf t (I.Shape [outsz]) NoUniqueness | t <- acctps ] + lam0' <- internaliseFoldLambda internaliseLambda lam0 acctps acc_arr_tps + let lam0_acc_params = fst $ splitAt (length accs) $ I.lambdaParams lam0' + acc_params <- forM lam0_acc_params $ \p -> do + name <- newVName $ baseString $ I.paramName p + return p { I.paramName = name } + + body_with_lam0 <- + ensureResultShape asserting "shape of result does not match shape of initial value" + (srclocOf lam0) acctps <=< insertStmsM $ do + lam_res <- bodyBind lam_body + + let consumed = consumedByLambda $ Alias.analyseLambda lam0' + copyIfConsumed p (I.Var v) + | I.paramName p `S.member` consumed = + letSubExp "acc_copy" $ I.BasicOp $ I.Copy v + copyIfConsumed _ x = return x + + accs' <- zipWithM copyIfConsumed (I.lambdaParams lam0') accs + lam_res' <- ensureArgShapes asserting + "shape of chunk function result does not match shape of initial value" + (srclocOf lam) [] (map I.typeOf $ I.lambdaParams lam0') lam_res + new_lam_res <- eLambda lam0' $ map eSubExp $ accs' ++ lam_res' + return $ resultBody new_lam_res + + -- Make sure the chunk size parameter comes first. + let form = I.Parallel o comm lam0' accs + lam' = I.Lambda { lambdaParams = chunk_param : acc_params ++ lam_val_params + , lambdaBody = body_with_lam0 + , lambdaReturnType = acctps } + w <- arraysSize 0 <$> mapM lookupType arrs + letTupExp' desc $ I.Op $ I.Stream w form lam' arrs + +-- The "interesting" cases are over, now it's mostly boilerplate. + +internaliseExp _ (E.Literal v _) = + return [I.Constant $ internalisePrimValue v] + +internaliseExp _ (E.IntLit v (Info t) _) = + case t of + E.Prim (E.Signed it) -> + return [I.Constant $ I.IntValue $ intValue it v] + E.Prim (E.Unsigned it) -> + return [I.Constant $ I.IntValue $ intValue it v] + E.Prim (E.FloatType ft) -> + return [I.Constant $ I.FloatValue $ floatValue ft v] + _ -> fail $ "internaliseExp: nonsensical type for integer literal: " ++ pretty t + +internaliseExp _ (E.FloatLit v (Info t) _) = + case t of + E.Prim (E.FloatType ft) -> + return [I.Constant $ I.FloatValue $ floatValue ft v] + _ -> fail $ "internaliseExp: nonsensical type for float literal: " ++ pretty t + +internaliseExp desc (E.If ce te fe _ _) = + letTupExp' desc =<< eIf (BasicOp . SubExp <$> internaliseExp1 "cond" ce) + (internaliseBody te) (internaliseBody fe) + +-- Builtin operators are handled specially because they are +-- overloaded. +internaliseExp desc (E.BinOp op _ (xe,_) (ye,_) _ loc) + | Just internalise <- isOverloadedFunction op [xe, ye] loc = + internalise desc + +-- User-defined operators are just the same as a function call. +internaliseExp desc (E.BinOp op (Info t) (xarg, Info xt) (yarg, Info yt) _ loc) = + internaliseExp desc $ + E.Apply (E.Apply (E.Var op (Info t) loc) xarg (Info $ E.diet xt) + (Info $ foldFunType [E.fromStruct yt] t) loc) + yarg (Info $ E.diet yt) (Info t) loc + +internaliseExp desc (E.Project k e (Info rt) _) = do + n <- internalisedTypeSize $ rt `setAliases` () + i' <- fmap sum $ mapM internalisedTypeSize $ + case E.typeOf e `setAliases` () of + Record fs -> map snd $ takeWhile ((/=k) . fst) $ sortFields fs + t -> [t] + take n . drop i' <$> internaliseExp desc e + +internaliseExp _ e@E.Lambda{} = + fail $ "internaliseExp: Unexpected lambda at " ++ locStr (srclocOf e) + +internaliseExp _ e@E.OpSection{} = + fail $ "internaliseExp: Unexpected operator section at " ++ locStr (srclocOf e) + +internaliseExp _ e@E.OpSectionLeft{} = + fail $ "internaliseExp: Unexpected left operator section at " ++ locStr (srclocOf e) + +internaliseExp _ e@E.OpSectionRight{} = + fail $ "internaliseExp: Unexpected right operator section at " ++ locStr (srclocOf e) + +internaliseExp _ e@E.ProjectSection{} = + fail $ "internaliseExp: Unexpected projection section at " ++ locStr (srclocOf e) + +internaliseExp _ e@E.IndexSection{} = + fail $ "internaliseExp: Unexpected index section at " ++ locStr (srclocOf e) + +internaliseSlice :: SrcLoc + -> [SubExp] + -> [E.DimIndex] + -> InternaliseM ([I.DimIndex SubExp], Certificates) +internaliseSlice loc dims idxs = do + (idxs', oks, parts) <- unzip3 <$> zipWithM internaliseDimIndex dims idxs + c <- assertingOne $ do + ok <- letSubExp "index_ok" =<< foldBinOp I.LogAnd (constant True) oks + let msg = ErrorMsg $ ["Index ["] ++ intercalate [", "] parts ++ + ["] out of bounds for array of shape ["] ++ + intersperse "][" (map ErrorInt32 $ take (length idxs) dims) ++ ["]."] + letExp "index_certs" $ I.BasicOp $ I.Assert ok msg (loc, mempty) + return (idxs', c) + +internaliseDimIndex :: SubExp -> E.DimIndex + -> InternaliseM (I.DimIndex SubExp, SubExp, [ErrorMsgPart SubExp]) +internaliseDimIndex w (E.DimFix i) = do + (i', _) <- internaliseDimExp "i" i + let lowerBound = I.BasicOp $ + I.CmpOp (I.CmpSle I.Int32) (I.constant (0 :: I.Int32)) i' + upperBound = I.BasicOp $ + I.CmpOp (I.CmpSlt I.Int32) i' w + ok <- letSubExp "bounds_check" =<< eBinOp I.LogAnd (pure lowerBound) (pure upperBound) + return (I.DimFix i', ok, [ErrorInt32 i']) +internaliseDimIndex w (E.DimSlice i j s) = do + s' <- maybe (return one) (fmap fst . internaliseDimExp "s") s + s_sign <- letSubExp "s_sign" $ BasicOp $ I.UnOp (I.SSignum Int32) s' + backwards <- letSubExp "backwards" $ I.BasicOp $ I.CmpOp (I.CmpEq int32) s_sign negone + w_minus_1 <- letSubExp "w_minus_1" $ BasicOp $ I.BinOp (Sub Int32) w one + let i_def = letSubExp "i_def" $ I.If backwards + (resultBody [w_minus_1]) + (resultBody [zero]) $ ifCommon [I.Prim int32] + j_def = letSubExp "j_def" $ I.If backwards + (resultBody [negone]) + (resultBody [w]) $ ifCommon [I.Prim int32] + i' <- maybe i_def (fmap fst . internaliseDimExp "i") i + j' <- maybe j_def (fmap fst . internaliseDimExp "j") j + j_m_i <- letSubExp "j_m_i" $ BasicOp $ I.BinOp (Sub Int32) j' i' + -- Something like a division-rounding-up, but accomodating negative + -- operands. + let divRounding x y = + eBinOp (SQuot Int32) (eBinOp (Add Int32) x (eBinOp (Sub Int32) y (eSignum $ toExp s'))) y + n <- letSubExp "n" =<< divRounding (toExp j_m_i) (toExp s') + + -- Bounds checks depend on whether we are slicing forwards or + -- backwards. If forwards, we must check '0 <= i && i <= j'. If + -- backwards, '-1 <= j && j <= i'. In both cases, we check '0 <= + -- i+n*s && i+(n-1)*s < w'. We only check if the slice is nonempty. + empty_slice <- letSubExp "empty_slice" $ I.BasicOp $ I.CmpOp (CmpEq int32) n zero + + m <- letSubExp "m" $ I.BasicOp $ I.BinOp (Sub Int32) n one + m_t_s <- letSubExp "m_t_s" $ I.BasicOp $ I.BinOp (Mul Int32) m s' + i_p_m_t_s <- letSubExp "i_p_m_t_s" $ I.BasicOp $ I.BinOp (Add Int32) i' m_t_s + zero_leq_i_p_m_t_s <- letSubExp "zero_leq_i_p_m_t_s" $ + I.BasicOp $ I.CmpOp (I.CmpSle Int32) zero i_p_m_t_s + i_p_m_t_s_leq_w <- letSubExp "i_p_m_t_s_leq_w" $ + I.BasicOp $ I.CmpOp (I.CmpSle Int32) i_p_m_t_s w + i_p_m_t_s_lth_w <- letSubExp "i_p_m_t_s_leq_w" $ + I.BasicOp $ I.CmpOp (I.CmpSlt Int32) i_p_m_t_s w + + zero_lte_i <- letSubExp "zero_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) zero i' + i_lte_j <- letSubExp "i_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) i' j' + forwards_ok <- letSubExp "forwards_ok" =<< + foldBinOp I.LogAnd zero_lte_i + [zero_lte_i, i_lte_j, zero_leq_i_p_m_t_s, i_p_m_t_s_lth_w] + + negone_lte_j <- letSubExp "negone_lte_j" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) negone j' + j_lte_i <- letSubExp "j_lte_i" $ I.BasicOp $ I.CmpOp (I.CmpSle Int32) j' i' + backwards_ok <- letSubExp "backwards_ok" =<< + foldBinOp I.LogAnd negone_lte_j + [negone_lte_j, j_lte_i, zero_leq_i_p_m_t_s, i_p_m_t_s_leq_w] + + slice_ok <- letSubExp "slice_ok" $ I.If backwards + (resultBody [backwards_ok]) + (resultBody [forwards_ok]) $ + ifCommon [I.Prim I.Bool] + ok_or_empty <- letSubExp "ok_or_empty" $ + I.BasicOp $ I.BinOp I.LogOr empty_slice slice_ok + + let parts = case (i, j, s) of + (_, _, Just{}) -> + [maybe "" (const $ ErrorInt32 i') i, ":", + maybe "" (const $ ErrorInt32 j') j, ":", + ErrorInt32 s'] + (_, Just{}, _) -> + [maybe "" (const $ ErrorInt32 i') i, ":", + ErrorInt32 j'] ++ + maybe mempty (const [":", ErrorInt32 s']) s + (_, Nothing, Nothing) -> + [ErrorInt32 i'] + return (I.DimSlice i' n s', ok_or_empty, parts) + where zero = constant (0::Int32) + negone = constant (-1::Int32) + one = constant (1::Int32) + +internaliseScanOrReduce :: String -> String + -> (SubExp -> I.Lambda -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) + -> (E.Exp, E.Exp, E.Exp, SrcLoc) + -> InternaliseM [SubExp] +internaliseScanOrReduce desc what f (lam, ne, arr, loc) = do + arrs <- internaliseExpToVars (what++"_arr") arr + nes <- internaliseExp (what++"_ne") ne + nes' <- forM (zip nes arrs) $ \(ne', arr') -> do + rowtype <- I.stripArray 1 <$> lookupType arr' + ensureShape asserting + "Row shape of input array does not match shape of neutral element" + loc rowtype (what++"_ne_right_shape") ne' + nests <- mapM I.subExpType nes' + arrts <- mapM lookupType arrs + lam' <- internaliseFoldLambda internaliseLambda lam nests arrts + w <- arraysSize 0 <$> mapM lookupType arrs + letTupExp' desc . I.Op =<< f w lam' nes' arrs + +internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp +internaliseExp1 desc e = do + vs <- internaliseExp desc e + case vs of [se] -> return se + _ -> fail "Internalise.internaliseExp1: was passed not just a single subexpression" + +-- | Promote to dimension type as appropriate for the original type. +-- Also return original type. +internaliseDimExp :: String -> E.Exp -> InternaliseM (I.SubExp, IntType) +internaliseDimExp s e = do + e' <- internaliseExp1 s e + case E.typeOf e of + E.Prim (Signed it) -> (,it) <$> asIntS Int32 e' + E.Prim (Unsigned it) -> (,it) <$> asIntZ Int32 e' + _ -> fail "internaliseDimExp: bad type" + +internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName] +internaliseExpToVars desc e = + mapM asIdent =<< internaliseExp desc e + where asIdent (I.Var v) = return v + asIdent se = letExp desc $ I.BasicOp $ I.SubExp se + +internaliseOperation :: String + -> E.Exp + -> (I.VName -> InternaliseM I.BasicOp) + -> InternaliseM [I.SubExp] +internaliseOperation s e op = do + vs <- internaliseExpToVars s e + letSubExps s =<< mapM (fmap I.BasicOp . op) vs + +internaliseBinOp :: String + -> E.BinOp + -> I.SubExp -> I.SubExp + -> E.PrimType + -> E.PrimType + -> InternaliseM [I.SubExp] +internaliseBinOp desc E.Plus x y (E.Signed t) _ = + simpleBinOp desc (I.Add t) x y +internaliseBinOp desc E.Plus x y (E.Unsigned t) _ = + simpleBinOp desc (I.Add t) x y +internaliseBinOp desc E.Plus x y (E.FloatType t) _ = + simpleBinOp desc (I.FAdd t) x y +internaliseBinOp desc E.Minus x y (E.Signed t) _ = + simpleBinOp desc (I.Sub t) x y +internaliseBinOp desc E.Minus x y (E.Unsigned t) _ = + simpleBinOp desc (I.Sub t) x y +internaliseBinOp desc E.Minus x y (E.FloatType t) _ = + simpleBinOp desc (I.FSub t) x y +internaliseBinOp desc E.Times x y (E.Signed t) _ = + simpleBinOp desc (I.Mul t) x y +internaliseBinOp desc E.Times x y (E.Unsigned t) _ = + simpleBinOp desc (I.Mul t) x y +internaliseBinOp desc E.Times x y (E.FloatType t) _ = + simpleBinOp desc (I.FMul t) x y +internaliseBinOp desc E.Divide x y (E.Signed t) _ = + simpleBinOp desc (I.SDiv t) x y +internaliseBinOp desc E.Divide x y (E.Unsigned t) _ = + simpleBinOp desc (I.UDiv t) x y +internaliseBinOp desc E.Divide x y (E.FloatType t) _ = + simpleBinOp desc (I.FDiv t) x y +internaliseBinOp desc E.Pow x y (E.FloatType t) _ = + simpleBinOp desc (I.FPow t) x y +internaliseBinOp desc E.Pow x y (E.Signed t) _ = + simpleBinOp desc (I.Pow t) x y +internaliseBinOp desc E.Pow x y (E.Unsigned t) _ = + simpleBinOp desc (I.Pow t) x y +internaliseBinOp desc E.Mod x y (E.Signed t) _ = + simpleBinOp desc (I.SMod t) x y +internaliseBinOp desc E.Mod x y (E.Unsigned t) _ = + simpleBinOp desc (I.UMod t) x y +internaliseBinOp desc E.Quot x y (E.Signed t) _ = + simpleBinOp desc (I.SQuot t) x y +internaliseBinOp desc E.Quot x y (E.Unsigned t) _ = + simpleBinOp desc (I.UDiv t) x y +internaliseBinOp desc E.Rem x y (E.Signed t) _ = + simpleBinOp desc (I.SRem t) x y +internaliseBinOp desc E.Rem x y (E.Unsigned t) _ = + simpleBinOp desc (I.UMod t) x y +internaliseBinOp desc E.ShiftR x y (E.Signed t) _ = + simpleBinOp desc (I.AShr t) x y +internaliseBinOp desc E.ShiftR x y (E.Unsigned t) _ = + simpleBinOp desc (I.LShr t) x y +internaliseBinOp desc E.ShiftL x y (E.Signed t) _ = + simpleBinOp desc (I.Shl t) x y +internaliseBinOp desc E.ShiftL x y (E.Unsigned t) _ = + simpleBinOp desc (I.Shl t) x y +internaliseBinOp desc E.Band x y (E.Signed t) _ = + simpleBinOp desc (I.And t) x y +internaliseBinOp desc E.Band x y (E.Unsigned t) _ = + simpleBinOp desc (I.And t) x y +internaliseBinOp desc E.Xor x y (E.Signed t) _ = + simpleBinOp desc (I.Xor t) x y +internaliseBinOp desc E.Xor x y (E.Unsigned t) _ = + simpleBinOp desc (I.Xor t) x y +internaliseBinOp desc E.Bor x y (E.Signed t) _ = + simpleBinOp desc (I.Or t) x y +internaliseBinOp desc E.Bor x y (E.Unsigned t) _ = + simpleBinOp desc (I.Or t) x y + +internaliseBinOp desc E.Equal x y t _ = + simpleCmpOp desc (I.CmpEq $ internalisePrimType t) x y +internaliseBinOp desc E.NotEqual x y t _ = do + eq <- letSubExp (desc++"true") $ I.BasicOp $ I.CmpOp (I.CmpEq $ internalisePrimType t) x y + fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp I.Not eq +internaliseBinOp desc E.Less x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSlt t) x y +internaliseBinOp desc E.Less x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUlt t) x y +internaliseBinOp desc E.Leq x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSle t) x y +internaliseBinOp desc E.Leq x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUle t) x y +internaliseBinOp desc E.Greater x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSlt t) y x -- Note the swapped x and y +internaliseBinOp desc E.Greater x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUlt t) y x -- Note the swapped x and y +internaliseBinOp desc E.Geq x y (E.Signed t) _ = + simpleCmpOp desc (I.CmpSle t) y x -- Note the swapped x and y +internaliseBinOp desc E.Geq x y (E.Unsigned t) _ = + simpleCmpOp desc (I.CmpUle t) y x -- Note the swapped x and y +internaliseBinOp desc E.Less x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLt t) x y +internaliseBinOp desc E.Leq x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLe t) x y +internaliseBinOp desc E.Greater x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLt t) y x -- Note the swapped x and y +internaliseBinOp desc E.Geq x y (E.FloatType t) _ = + simpleCmpOp desc (I.FCmpLe t) y x -- Note the swapped x and y + +-- Relational operators for booleans. +internaliseBinOp desc E.Less x y E.Bool _ = + simpleCmpOp desc I.CmpLlt x y +internaliseBinOp desc E.Leq x y E.Bool _ = + simpleCmpOp desc I.CmpLle x y +internaliseBinOp desc E.Greater x y E.Bool _ = + simpleCmpOp desc I.CmpLlt y x -- Note the swapped x and y +internaliseBinOp desc E.Geq x y E.Bool _ = + simpleCmpOp desc I.CmpLle y x -- Note the swapped x and y + +internaliseBinOp _ op _ _ t1 t2 = + fail $ "Invalid binary operator " ++ pretty op ++ + " with operand types " ++ pretty t1 ++ ", " ++ pretty t2 + +simpleBinOp :: String + -> I.BinOp + -> I.SubExp -> I.SubExp + -> InternaliseM [I.SubExp] +simpleBinOp desc bop x y = + letTupExp' desc $ I.BasicOp $ I.BinOp bop x y + +simpleCmpOp :: String + -> I.CmpOp + -> I.SubExp -> I.SubExp + -> InternaliseM [I.SubExp] +simpleCmpOp desc op x y = + letTupExp' desc $ I.BasicOp $ I.CmpOp op x y + +findFuncall :: E.Exp -> InternaliseM (E.QualName VName, [E.Exp], [E.StructType]) +findFuncall (E.Var fname (Info t) _) = + let (remaining, _) = unfoldFunType t + in return (fname, [], map E.toStruct remaining) +findFuncall (E.Apply f arg _ (Info t) _) = do + let (remaining, _) = unfoldFunType t + (fname, args, _) <- findFuncall f + return (fname, args ++ [arg], map E.toStruct remaining) +findFuncall e = + fail $ "Invalid function expression in application: " ++ pretty e + +internaliseLambda :: InternaliseLambda + +internaliseLambda (E.Parens e _) rowtypes = + internaliseLambda e rowtypes + +internaliseLambda (E.Lambda tparams params body _ (Info (_, rettype)) loc) rowtypes = + bindingLambdaParams tparams params rowtypes $ \pcm params' -> do + (rettype', rcm) <- internaliseReturnType rettype + body' <- internaliseBody body + mapM_ (uncurry (internaliseDimConstant loc)) $ pcm<>rcm + return (params', body', map I.fromDecl rettype') + +internaliseLambda E.OpSection{} _ = fail "internaliseLambda: unexpected OpSection" + +internaliseLambda E.OpSectionLeft{} _ = fail "internaliseLambda: unexpected OpSectionLeft" + +internaliseLambda E.OpSectionRight{} _ = fail "internaliseLambda: unexpected OpSectionRight" + +internaliseLambda e rowtypes = do + (_, _, remaining_params_ts) <- findFuncall e + (params, param_args) <- fmap unzip $ forM remaining_params_ts $ \et -> do + name <- newVName "not_curried" + return (E.Id name (Info $ E.vacuousShapeAnnotations $ et `setAliases` mempty) loc, + E.Var (E.qualName name) + (Info (et `setAliases` mempty)) loc) + let rettype = E.typeOf e + body = foldl (\f arg -> E.Apply f arg (Info E.Observe) + (Info $ E.vacuousShapeAnnotations rettype) loc) + e + param_args + rettype' = E.vacuousShapeAnnotations $ rettype `E.setAliases` () + internaliseLambda (E.Lambda [] params body Nothing (Info (mempty, rettype')) loc) rowtypes + where loc = srclocOf e + +internaliseDimConstant :: SrcLoc -> Name -> VName -> InternaliseM () +internaliseDimConstant loc fname name = + letBind_ (basicPattern [] [I.Ident name $ I.Prim I.int32]) $ + I.Apply fname [] [I.Prim I.int32] (Safe, loc, mempty) + +-- | Some operators and functions are overloaded or otherwise special +-- - we detect and treat them here. +isOverloadedFunction :: E.QualName VName -> [E.Exp] -> SrcLoc + -> Maybe (String -> InternaliseM [SubExp]) +isOverloadedFunction qname args loc = do + guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag + handle args $ baseString $ qualLeaf qname + where + handle [x] "sign_i8" = Just $ toSigned I.Int8 x + handle [x] "sign_i16" = Just $ toSigned I.Int16 x + handle [x] "sign_i32" = Just $ toSigned I.Int32 x + handle [x] "sign_i64" = Just $ toSigned I.Int64 x + + handle [x] "unsign_i8" = Just $ toUnsigned I.Int8 x + handle [x] "unsign_i16" = Just $ toUnsigned I.Int16 x + handle [x] "unsign_i32" = Just $ toUnsigned I.Int32 x + handle [x] "unsign_i64" = Just $ toUnsigned I.Int64 x + + handle [x] "sgn" = Just $ signumF x + handle [x] "abs" = Just $ absF x + handle [x] "!" = Just $ notF x + handle [x] "~" = Just $ complementF x + + handle [x] "opaque" = Just $ \desc -> + mapM (letSubExp desc . BasicOp . Opaque) =<< internaliseExp "opaque_arg" x + + handle [x] s + | Just unop <- find ((==s) . pretty) allUnOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x' + + handle [x,y] s + | Just bop <- find ((==s) . pretty) allBinOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + y' <- internaliseExp1 "y" y + fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y' + | Just cmp <- find ((==s) . pretty) allCmpOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + y' <- internaliseExp1 "y" y + fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y' + handle [x] s + | Just conv <- find ((==s) . pretty) allConvOps = Just $ \desc -> do + x' <- internaliseExp1 "x" x + fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x' + + -- Short-circuiting operators are magical. + handle [x,y] "&&" = Just $ \desc -> + internaliseExp desc $ + E.If x y (E.Literal (E.BoolValue False) noLoc) (Info (E.Prim E.Bool)) noLoc + handle [x,y] "||" = Just $ \desc -> + internaliseExp desc $ + E.If x (E.Literal (E.BoolValue True) noLoc) y (Info (E.Prim E.Bool)) noLoc + + -- Handle equality and inequality specially, to treat the case of + -- arrays. + handle [xe,ye] op + | Just cmp_f <- isEqlOp op = Just $ \desc -> do + xe' <- internaliseExp "x" xe + ye' <- internaliseExp "y" ye + rs <- zipWithM (doComparison desc) xe' ye' + cmp_f desc =<< letSubExp "eq" =<< foldBinOp I.LogAnd (constant True) rs + where isEqlOp "!=" = Just $ \desc eq -> + letTupExp' desc $ I.BasicOp $ I.UnOp I.Not eq + isEqlOp "==" = Just $ \_ eq -> + return [eq] + isEqlOp _ = Nothing + + doComparison desc x y = do + x_t <- I.subExpType x + y_t <- I.subExpType y + case x_t of + I.Prim t -> letSubExp desc $ I.BasicOp $ I.CmpOp (I.CmpEq t) x y + _ -> do + let x_dims = I.arrayDims x_t + y_dims = I.arrayDims y_t + dims_match <- forM (zip x_dims y_dims) $ \(x_dim, y_dim) -> + letSubExp "dim_eq" $ I.BasicOp $ I.CmpOp (I.CmpEq int32) x_dim y_dim + shapes_match <- letSubExp "shapes_match" =<< + foldBinOp I.LogAnd (constant True) dims_match + compare_elems_body <- runBodyBinder $ do + -- Flatten both x and y. + x_num_elems <- letSubExp "x_num_elems" =<< + foldBinOp (I.Mul Int32) (constant (1::Int32)) x_dims + x' <- letExp "x" $ I.BasicOp $ I.SubExp x + y' <- letExp "x" $ I.BasicOp $ I.SubExp y + x_flat <- letExp "x_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] x' + y_flat <- letExp "y_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] y' + + -- Compare the elements. + cmp_lam <- cmpOpLambda (I.CmpEq (elemType x_t)) (elemType x_t) + cmps <- letExp "cmps" $ I.Op $ + I.Screma x_num_elems (I.mapSOAC cmp_lam) [x_flat, y_flat] + + -- Check that all were equal. + and_lam <- binOpLambda I.LogAnd I.Bool + reduce <- I.reduceSOAC Commutative and_lam [constant True] + all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems reduce [cmps] + return $ resultBody [all_equal] + + letSubExp "arrays_equal" $ + I.If shapes_match compare_elems_body (resultBody [constant False]) $ + ifCommon [I.Prim I.Bool] + + handle [x,y] name + | Just bop <- find ((name==) . pretty) [minBound..maxBound::E.BinOp] = + Just $ \desc -> do + x' <- internaliseExp1 "x" x + y' <- internaliseExp1 "y" y + case (E.typeOf x, E.typeOf y) of + (E.Prim t1, E.Prim t2) -> + internaliseBinOp desc bop x' y' t1 t2 + _ -> fail "Futhark.Internalise.internaliseExp: non-primitive type in BinOp." + + handle [E.TupLit [a, si, v] _] "scatter" = Just $ scatterF a si v + + handle [E.TupLit [e, E.ArrayLit vs _ _] _] "cmp_threshold" = do + s <- mapM isCharLit vs + Just $ \desc -> do + x <- internaliseExp1 "threshold_x" e + pure <$> letSubExp desc (Op $ CmpThreshold x s) + where isCharLit (Literal (SignedValue iv) _) = Just $ chr $ fromIntegral $ intToInt64 iv + isCharLit _ = Nothing + + handle [E.TupLit [n, m, arr] _] f + | f `elem` ["unflatten", "cosmin_unflatten"] = Just $ \desc -> do + arrs <- internaliseExpToVars "unflatten_arr" arr + n' <- internaliseExp1 "n" n + m' <- internaliseExp1 "m" m + -- The unflattened dimension needs to have the same number of elements + -- as the original dimension. + old_dim <- I.arraysSize 0 <$> mapM lookupType arrs + dim_ok <- assertingOne $ letExp "dim_ok" =<< + eAssert (eCmpOp (I.CmpEq I.int32) + (eBinOp (I.Mul Int32) (eSubExp n') (eSubExp m')) + (eSubExp old_dim)) + "new shape has different number of elements than old shape" loc + certifying dim_ok $ forM arrs $ \arr' -> do + arr_t <- lookupType arr' + letSubExp desc $ I.BasicOp $ + I.Reshape (reshapeOuter [DimNew n', DimNew m'] 1 $ arrayShape arr_t) arr' + + handle [arr] f + | f `elem` ["flatten", "cosmin_flatten"] = Just $ \desc -> do + arrs <- internaliseExpToVars "flatten_arr" arr + forM arrs $ \arr' -> do + arr_t <- lookupType arr' + let n = arraySize 0 arr_t + m = arraySize 1 arr_t + k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int32) n m + letSubExp desc $ I.BasicOp $ + I.Reshape (reshapeOuter [DimNew k] 2 $ arrayShape arr_t) arr' + + handle [TupLit [x, y] _] "concat" = Just $ \desc -> do + xs <- internaliseExpToVars "concat_x" x + ys <- internaliseExpToVars "concat_y" y + outer_size <- arraysSize 0 <$> mapM lookupType xs + let sumdims xsize ysize = letSubExp "conc_tmp" $ I.BasicOp $ + I.BinOp (I.Add I.Int32) xsize ysize + ressize <- foldM sumdims outer_size =<< + mapM (fmap (arraysSize 0) . mapM lookupType) [ys] + + let conc xarr yarr = do + -- All dimensions except for dimension 'i' must match. + xt <- lookupType xarr + yt <- lookupType yarr + let matches n m = + letExp "match" =<< + eAssert (pure $ I.BasicOp $ I.CmpOp (I.CmpEq I.int32) n m) + "arguments do not have the same row shape" loc + x_inner_dims = drop 1 $ I.arrayDims xt + y_inner_dims = drop 1 $ I.arrayDims yt + updims = zipWith3 updims' [(0::Int)..] (I.arrayDims xt) + updims' j xd yd | j == 0 = yd + | otherwise = xd + matchcs <- asserting $ Certificates <$> + zipWithM matches x_inner_dims y_inner_dims + yarr' <- certifying matchcs $ letExp "concat_y_reshaped" $ + shapeCoerce (updims $ I.arrayDims yt) yarr + return $ I.BasicOp $ I.Concat 0 xarr [yarr'] ressize + letSubExps desc =<< zipWithM conc xs ys + + handle [TupLit [offset, e] _] "rotate" = Just $ \desc -> do + offset' <- internaliseExp1 "rotation_offset" offset + internaliseOperation desc e $ \v -> do + r <- I.arrayRank <$> lookupType v + let zero = intConst Int32 0 + offsets = offset' : replicate (r-1) zero + return $ I.Rotate offsets v + + handle [e] "transpose" = Just $ \desc -> + internaliseOperation desc e $ \v -> do + r <- I.arrayRank <$> lookupType v + return $ I.Rearrange ([1,0] ++ [2..r-1]) v + + handle [TupLit [x, y] _] "zip" = Just $ \desc -> + (++) <$> internaliseExp (desc ++ "_zip_x") x + <*> internaliseExp (desc ++ "_zip_y") y + + handle [x] "unzip" = Just $ flip internaliseExp x + handle [x] "trace" = Just $ flip internaliseExp x + handle [x] "break" = Just $ flip internaliseExp x + + handle _ _ = Nothing + + toSigned int_to e desc = do + e' <- internaliseExp1 "trunc_arg" e + case E.typeOf e of + E.Prim E.Bool -> + letTupExp' desc $ I.If e' (resultBody [intConst int_to 1]) + (resultBody [intConst int_to 0]) $ + ifCommon [I.Prim $ I.IntType int_to] + E.Prim (E.Signed int_from) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.SExt int_from int_to) e' + E.Prim (E.Unsigned int_from) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' + E.Prim (E.FloatType float_from) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToSI float_from int_to) e' + _ -> fail "Futhark.Internalise.handle: non-numeric type in ToSigned" + + toUnsigned int_to e desc = do + e' <- internaliseExp1 "trunc_arg" e + case E.typeOf e of + E.Prim E.Bool -> + letTupExp' desc $ I.If e' (resultBody [intConst int_to 1]) + (resultBody [intConst int_to 0]) $ + ifCommon [I.Prim $ I.IntType int_to] + E.Prim (E.Signed int_from) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' + E.Prim (E.Unsigned int_from) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' + E.Prim (E.FloatType float_from) -> + letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToUI float_from int_to) e' + _ -> fail "Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned" + + signumF e desc = do + e' <- internaliseExp1 "signum_arg" e + case E.typeOf e of + E.Prim (E.Signed t) -> + letTupExp' desc $ I.BasicOp $ I.UnOp (I.SSignum t) e' + E.Prim (E.Unsigned t) -> + letTupExp' desc $ I.BasicOp $ I.UnOp (I.USignum t) e' + _ -> fail "Futhark.Internalise.internaliseExp: non-integer type in Signum" + + absF e desc = do + e' <- internaliseExp1 "abs_arg" e + case E.typeOf e of + E.Prim (E.Signed t) -> + letTupExp' desc $ I.BasicOp $ I.UnOp (I.Abs t) e' + E.Prim (E.Unsigned _) -> + return [e'] + E.Prim (E.FloatType t) -> + letTupExp' desc $ I.BasicOp $ I.UnOp (I.FAbs t) e' + _ -> fail "Futhark.Internalise.internaliseExp: non-integer type in Abs" + + notF e desc = do + e' <- internaliseExp1 "not_arg" e + letTupExp' desc $ I.BasicOp $ I.UnOp I.Not e' + + complementF e desc = do + e' <- internaliseExp1 "complement_arg" e + et <- subExpType e' + case et of I.Prim (I.IntType t) -> + letTupExp' desc $ I.BasicOp $ I.UnOp (I.Complement t) e' + _ -> + fail "Futhark.Internalise.internaliseExp: non-integer type in Complement" + + scatterF a si v desc = do + si' <- letExp "write_si" . BasicOp . SubExp =<< internaliseExp1 "write_arg_i" si + svs <- internaliseExpToVars "write_arg_v" v + sas <- internaliseExpToVars "write_arg_a" a + + si_shape <- I.arrayShape <$> lookupType si' + let si_w = shapeSize 0 si_shape + sv_ts <- mapM lookupType svs + + svs' <- forM (zip svs sv_ts) $ \(sv,sv_t) -> do + let sv_shape = I.arrayShape sv_t + sv_w = arraySize 0 sv_t + + -- Generate an assertion and reshapes to ensure that sv and si' are the same + -- size. + cmp <- letSubExp "write_cmp" $ I.BasicOp $ + I.CmpOp (I.CmpEq I.int32) si_w sv_w + c <- assertingOne $ + letExp "write_cert" $ I.BasicOp $ + I.Assert cmp "length of index and value array does not match" (loc, mempty) + certifying c $ letExp (baseString sv ++ "_write_sv") $ + I.BasicOp $ I.Reshape (reshapeOuter [DimCoercion si_w] 1 sv_shape) sv + + indexType <- rowType <$> lookupType si' + indexName <- newVName "write_index" + valueNames <- replicateM (length sv_ts) $ newVName "write_value" + + sa_ts <- mapM lookupType sas + let bodyTypes = replicate (length sv_ts) indexType ++ map rowType sa_ts + paramTypes = indexType : map rowType sv_ts + bodyNames = indexName : valueNames + bodyParams = zipWith I.Param bodyNames paramTypes + + -- This body is pretty boring right now, as every input is exactly the output. + -- But it can get funky later on if fused with something else. + body <- localScope (scopeOfLParams bodyParams) $ insertStmsM $ do + let outs = replicate (length valueNames) indexName ++ valueNames + results <- forM outs $ \name -> + letSubExp "write_res" $ I.BasicOp $ I.SubExp $ I.Var name + ensureResultShape asserting "scatter value has wrong size" loc + bodyTypes $ resultBody results + + let lam = I.Lambda { I.lambdaParams = bodyParams + , I.lambdaReturnType = bodyTypes + , I.lambdaBody = body + } + sivs = si' : svs' + + let sa_ws = map (arraySize 0) sa_ts + letTupExp' desc $ I.Op $ I.Scatter si_w lam sivs $ zip3 sa_ws (repeat 1) sas + +-- | Is the name a value constant? If so, create the necessary +-- function call and return the corresponding subexpressions. +lookupConstant :: SrcLoc -> VName -> InternaliseM (Maybe [SubExp]) +lookupConstant loc name = do + is_const <- lookupFunction' name + scope <- askScope + case is_const of + Just (fname, constparams, _, _, _, _, mk_rettype) + | name `M.notMember` scope -> do + (constargs, const_ds, const_ts) <- unzip3 <$> constFunctionArgs loc constparams + safety <- askSafety + case mk_rettype $ zip constargs $ map I.fromDecl const_ts of + Nothing -> fail $ "lookupConstant: " ++ + unwords (pretty name : zipWith (curry pretty) constargs const_ts) ++ + " failed" + Just rettype -> + fmap (Just . map I.Var) $ letTupExp (baseString name) $ + I.Apply fname (zip constargs const_ds) rettype (safety, loc, mempty) + _ -> return Nothing + +constFunctionArgs :: SrcLoc -> ConstParams -> InternaliseM [(SubExp, I.Diet, I.DeclType)] +constFunctionArgs loc = mapM arg + where arg (fname, name) = do + safety <- askSafety + se <- letSubExp (baseString name ++ "_arg") $ + I.Apply fname [] [I.Prim I.int32] (safety, loc, []) + return (se, I.Observe, I.Prim I.int32) + +funcall :: String -> QualName VName -> [SubExp] -> SrcLoc + -> InternaliseM ([SubExp], [I.ExtType]) +funcall desc (QualName _ fname) args loc = do + (fname', constparams, closure, shapes, value_paramts, fun_params, rettype_fun) <- + lookupFunction fname + (constargs, const_ds, _) <- unzip3 <$> constFunctionArgs loc constparams + argts <- mapM subExpType args + closure_ts <- mapM lookupType closure + let shapeargs = argShapes shapes value_paramts argts + diets = const_ds ++ replicate (length closure + length shapeargs) I.Observe ++ + map I.diet value_paramts + constOrShape = const $ I.Prim int32 + paramts = map constOrShape constargs ++ closure_ts ++ + map constOrShape shapeargs ++ map I.fromDecl value_paramts + args' <- ensureArgShapes asserting "function arguments of wrong shape" + loc (map I.paramName fun_params) + paramts (constargs ++ map I.Var closure ++ shapeargs ++ args) + argts' <- mapM subExpType args' + case rettype_fun $ zip args' argts' of + Nothing -> fail $ "Cannot apply " ++ pretty fname ++ " to arguments\n " ++ + pretty args' ++ "\nof types\n " ++ + pretty argts' ++ + "\nFunction has parameters\n " ++ pretty fun_params + Just ts -> do + safety <- askSafety + ses <- letTupExp' desc $ I.Apply fname' (zip args' diets) ts (safety, loc, mempty) + return (ses, map I.fromDecl ts) + +askSafety :: InternaliseM Safety +askSafety = do check <- asks envDoBoundsChecks + safe <- asks envSafe + return $ if check || safe then I.Safe else I.Unsafe + +-- Implement partitioning using maps, scans and writes. +partitionWithSOACS :: Int -> I.Lambda -> [I.VName] -> InternaliseM ([I.SubExp], [I.SubExp]) +partitionWithSOACS k lam arrs = do + arr_ts <- mapM lookupType arrs + let w = arraysSize 0 arr_ts + classes_and_increments <- letTupExp "increments" $ I.Op $ I.Screma w (mapSOAC lam) arrs + (classes, increments) <- case classes_and_increments of + classes : increments -> return (classes, take k increments) + _ -> fail "partitionWithSOACS" + + add_lam_x_params <- + replicateM k $ I.Param <$> newVName "x" <*> pure (I.Prim int32) + add_lam_y_params <- + replicateM k $ I.Param <$> newVName "y" <*> pure (I.Prim int32) + add_lam_body <- runBodyBinder $ + localScope (scopeOfLParams $ add_lam_x_params++add_lam_y_params) $ + fmap resultBody $ forM (zip add_lam_x_params add_lam_y_params) $ \(x,y) -> + letSubExp "z" $ I.BasicOp $ I.BinOp (I.Add Int32) + (I.Var $ I.paramName x) (I.Var $ I.paramName y) + let add_lam = I.Lambda { I.lambdaBody = add_lam_body + , I.lambdaParams = add_lam_x_params ++ add_lam_y_params + , I.lambdaReturnType = replicate k $ I.Prim int32 + } + nes = replicate (length increments) $ constant (0::Int32) + + scan <- I.scanSOAC add_lam nes + all_offsets <- letTupExp "offsets" $ I.Op $ I.Screma w scan increments + + -- We have the offsets for each of the partitions, but we also need + -- the total sizes, which are the last elements in the offests. We + -- just have to be careful in case the array is empty. + last_index <- letSubExp "last_index" $ I.BasicOp $ I.BinOp (I.Sub Int32) w $ constant (1::Int32) + nonempty_body <- runBodyBinder $ fmap resultBody $ forM all_offsets $ \offset_array -> + letSubExp "last_offset" $ I.BasicOp $ I.Index offset_array [I.DimFix last_index] + let empty_body = resultBody $ replicate k $ constant (0::Int32) + is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int32) w $ constant (0::Int32) + sizes <- letTupExp "partition_size" $ + I.If is_empty empty_body nonempty_body $ + ifCommon $ replicate k $ I.Prim int32 + + -- Compute total size of all partitions. + sum_of_partition_sizes <- letSubExp "sum_of_partition_sizes" =<< + foldBinOp (Add Int32) (constant (0::Int32)) (map I.Var sizes) + + -- Create scratch arrays for the result. + blanks <- forM arr_ts $ \arr_t -> + letExp "partition_dest" $ I.BasicOp $ + Scratch (elemType arr_t) (sum_of_partition_sizes : drop 1 (I.arrayDims arr_t)) + + -- Now write into the result. + write_lam <- do + c_param <- I.Param <$> newVName "c" <*> pure (I.Prim int32) + offset_params <- replicateM k $ I.Param <$> newVName "offset" <*> pure (I.Prim int32) + value_params <- forM arr_ts $ \arr_t -> + I.Param <$> newVName "v" <*> pure (I.rowType arr_t) + (offset, offset_stms) <- collectStms $ mkOffsetLambdaBody (map I.Var sizes) + (I.Var $ I.paramName c_param) 0 offset_params + return I.Lambda { I.lambdaParams = c_param : offset_params ++ value_params + , I.lambdaReturnType = replicate (length arr_ts) (I.Prim int32) ++ + map I.rowType arr_ts + , I.lambdaBody = mkBody offset_stms $ + replicate (length arr_ts) offset ++ + map (I.Var . I.paramName) value_params + } + results <- letTupExp "partition_res" $ I.Op $ I.Scatter w + write_lam (classes : all_offsets ++ arrs) $ + zip3 (repeat sum_of_partition_sizes) (repeat 1) blanks + sizes' <- letSubExp "partition_sizes" $ I.BasicOp $ + I.ArrayLit (map I.Var sizes) $ I.Prim int32 + return (map I.Var results, [sizes']) + where + mkOffsetLambdaBody :: [SubExp] + -> SubExp + -> Int + -> [I.LParam] + -> InternaliseM SubExp + mkOffsetLambdaBody _ _ _ [] = + return $ constant (-1::Int32) + mkOffsetLambdaBody sizes c i (p:ps) = do + is_this_one <- letSubExp "is_this_one" $ I.BasicOp $ I.CmpOp (CmpEq int32) c (constant i) + next_one <- mkOffsetLambdaBody sizes c (i+1) ps + this_one <- letSubExp "this_offset" =<< + foldBinOp (Add Int32) (constant (-1::Int32)) + (I.Var (I.paramName p) : take i sizes) + letSubExp "total_res" $ I.If is_this_one + (resultBody [this_one]) (resultBody [next_one]) $ ifCommon [I.Prim int32] + +typeExpForError :: ConstParams -> E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp] +typeExpForError _ (E.TEVar qn _) = + return [ErrorString $ pretty qn] +typeExpForError cm (E.TEUnique te _) = ("*":) <$> typeExpForError cm te +typeExpForError cm (E.TEArray te d _) = do + d' <- dimDeclForError cm d + te' <- typeExpForError cm te + return $ ["[", d', "]"] ++ te' +typeExpForError cm (E.TETuple tes _) = do + tes' <- mapM (typeExpForError cm) tes + return $ ["("] ++ intercalate [", "] tes' ++ [")"] +typeExpForError cm (E.TERecord fields _) = do + fields' <- mapM onField fields + return $ ["{"] ++ intercalate [", "] fields' ++ ["}"] + where onField (k, te) = (ErrorString (pretty k ++ ": "):) <$> typeExpForError cm te +typeExpForError cm (E.TEArrow _ t1 t2 _) = do + t1' <- typeExpForError cm t1 + t2' <- typeExpForError cm t2 + return $ t1' ++ [" -> "] ++ t2' +typeExpForError cm (E.TEApply t arg _) = do + t' <- typeExpForError cm t + arg' <- case arg of TypeArgExpType argt -> typeExpForError cm argt + TypeArgExpDim d _ -> pure <$> dimDeclForError cm d + return $ t' ++ [" "] ++ arg' + +dimDeclForError :: ConstParams -> E.DimDecl VName -> InternaliseM (ErrorMsgPart SubExp) +dimDeclForError cm (NamedDim d) = do + substs <- asks $ M.lookup (E.qualLeaf d) . envSubsts + let fname = nameFromString $ pretty (E.qualLeaf d) ++ "f" + d' <- case (substs, lookup fname cm) of + (Just [v], _) -> return v + (_, Just v) -> return $ I.Var v + _ -> return $ I.Var $ E.qualLeaf d + return $ ErrorInt32 d' +dimDeclForError _ (ConstDim d) = + return $ ErrorString $ pretty d +dimDeclForError _ AnyDim = return "" diff --git a/src/Futhark/Internalise/AccurateSizes.hs b/src/Futhark/Internalise/AccurateSizes.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/AccurateSizes.hs @@ -0,0 +1,138 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Internalise.AccurateSizes + ( shapeBody + , annotateArrayShape + , argShapes + , ensureResultShape + , ensureResultExtShape + , ensureResultExtShapeNoCtx + , ensureExtShape + , ensureShape + , ensureArgShapes + ) + where + +import Control.Monad +import Data.Loc +import qualified Data.Map.Strict as M + +import Futhark.Construct +import Futhark.Representation.AST + +shapeBody :: (HasScope lore m, MonadFreshNames m, BinderOps lore, Bindable lore) => + [VName] -> [Type] -> Body lore + -> m (Body lore) +shapeBody shapenames ts body = + runBodyBinder $ do + ses <- bodyBind body + sets <- mapM subExpType ses + return $ resultBody $ argShapes shapenames ts sets + +annotateArrayShape :: ArrayShape shape => + TypeBase shape u -> [Int] -> TypeBase Shape u +annotateArrayShape t newshape = + t `setArrayShape` Shape (take (arrayRank t) $ + map (intConst Int32 . toInteger) $ newshape ++ repeat 0) + +argShapes :: [VName] -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> [SubExp] +argShapes shapes valts valargts = + map addShape shapes + where mapping = shapeMapping valts valargts + addShape name + | Just se <- M.lookup name mapping = se + | otherwise = intConst Int32 0 + +ensureResultShape :: MonadBinder m => + (m Certificates -> m Certificates) + -> ErrorMsg SubExp -> SrcLoc -> [Type] -> Body (Lore m) + -> m (Body (Lore m)) +ensureResultShape asserting msg loc = + ensureResultExtShape asserting msg loc . staticShapes + +ensureResultExtShape :: MonadBinder m => + (m Certificates -> m Certificates) + -> ErrorMsg SubExp -> SrcLoc -> [ExtType] -> Body (Lore m) + -> m (Body (Lore m)) +ensureResultExtShape asserting msg loc rettype body = + insertStmsM $ do + reses <- bodyBind =<< + ensureResultExtShapeNoCtx asserting msg loc rettype body + ts <- mapM subExpType reses + let ctx = extractShapeContext rettype $ map arrayDims ts + mkBodyM mempty $ ctx ++ reses + +ensureResultExtShapeNoCtx :: MonadBinder m => + (m Certificates -> m Certificates) + -> ErrorMsg SubExp -> SrcLoc -> [ExtType] -> Body (Lore m) + -> m (Body (Lore m)) +ensureResultExtShapeNoCtx asserting msg loc rettype body = + insertStmsM $ do + es <- bodyBind body + es_ts <- mapM subExpType es + let ext_mapping = shapeExtMapping rettype es_ts + rettype' = foldr (uncurry fixExt) rettype $ M.toList ext_mapping + assertProperShape t se = + let name = "result_proper_shape" + in ensureExtShape asserting msg loc t name se + resultBodyM =<< zipWithM assertProperShape rettype' es + +ensureExtShape :: MonadBinder m => + (m Certificates -> m Certificates) + -> ErrorMsg SubExp -> SrcLoc -> ExtType -> String -> SubExp + -> m SubExp +ensureExtShape asserting msg loc t name orig + | Array{} <- t, Var v <- orig = + Var <$> ensureShapeVar asserting msg loc t name v + | otherwise = return orig + +ensureShape :: MonadBinder m => + (m Certificates -> m Certificates) + -> ErrorMsg SubExp -> SrcLoc -> Type -> String -> SubExp + -> m SubExp +ensureShape asserting msg loc = ensureExtShape asserting msg loc . staticShapes1 + +-- | Reshape the arguments to a function so that they fit the expected +-- shape declarations. Not used to change rank of arguments. Assumes +-- everything is otherwise type-correct. +ensureArgShapes :: (MonadBinder m, Typed (TypeBase Shape u)) => + (m Certificates -> m Certificates) + -> ErrorMsg SubExp -> SrcLoc -> [VName] -> [TypeBase Shape u] -> [SubExp] + -> m [SubExp] +ensureArgShapes asserting msg loc shapes paramts args = + zipWithM ensureArgShape (expectedTypes shapes paramts args) args + where ensureArgShape _ (Constant v) = return $ Constant v + ensureArgShape t (Var v) + | arrayRank t < 1 = return $ Var v + | otherwise = + ensureShape asserting msg loc t (baseString v) $ Var v + + +ensureShapeVar :: MonadBinder m => + (m Certificates -> m Certificates) + -> ErrorMsg SubExp -> SrcLoc -> ExtType -> String -> VName + -> m VName +ensureShapeVar asserting msg loc t name v + | Array{} <- t = do + newdims <- arrayDims . removeExistentials t <$> lookupType v + olddims <- arrayDims <$> lookupType v + if newdims == olddims + then return v + else do + certs <- asserting $ do + old_zero <- letSubExp "old_empty" =<< anyZero olddims + new_zero <- letSubExp "new_empty" =<< anyZero newdims + both_empty <- letSubExp "both_empty" $ BasicOp $ BinOp LogAnd old_zero new_zero + + matches <- zipWithM checkDim newdims olddims + all_match <- letSubExp "match" =<< foldBinOp LogAnd (constant True) matches + + empty_or_match <- letSubExp "empty_or_match" $ BasicOp $ BinOp LogOr both_empty all_match + Certificates . pure <$> letExp "empty_or_match_cert" + (BasicOp $ Assert empty_or_match msg (loc, [])) + certifying certs $ letExp name $ shapeCoerce newdims v + | otherwise = return v + where checkDim desired has = + letSubExp "dim_match" $ BasicOp $ CmpOp (CmpEq int32) desired has + anyZero = + foldBinOp LogOr (constant False) <=< + mapM (letSubExp "dim_zero" . BasicOp . CmpOp (CmpEq int32) (intConst Int32 0)) diff --git a/src/Futhark/Internalise/Bindings.hs b/src/Futhark/Internalise/Bindings.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/Bindings.hs @@ -0,0 +1,204 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Internalise.Bindings + ( + -- * Internalising bindings + bindingParams + , bindingLambdaParams + , stmPattern + , MatchPattern + ) + where + +import Control.Monad.State hiding (mapM) +import Control.Monad.Reader hiding (mapM) +import Control.Monad.Writer hiding (mapM) + +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Loc +import Data.Traversable (mapM) + +import Language.Futhark as E +import qualified Futhark.Representation.SOACS as I +import Futhark.MonadFreshNames + +import Futhark.Internalise.Monad +import Futhark.Internalise.TypesValues +import Futhark.Internalise.AccurateSizes +import Futhark.Util + +bindingParams :: [E.TypeParam] -> [E.Pattern] + -> (ConstParams -> [I.FParam] -> [[I.FParam]] -> InternaliseM a) + -> InternaliseM a +bindingParams tparams params m = do + flattened_params <- mapM flattenPattern params + let (params_idents, params_types) = unzip $ concat flattened_params + bound = boundInTypes tparams + param_names = M.fromList [ (E.identName x, y) | (x,y) <- params_idents ] + (params_ts, cm) <- internaliseParamTypes bound param_names params_types + let num_param_idents = map length flattened_params + num_param_ts = map (sum . map length) $ chunks num_param_idents params_ts + + (params_ts', unnamed_shape_params) <- + fmap unzip $ forM params_ts $ \param_ts -> do + (param_ts', param_unnamed_dims) <- instantiateShapesWithDecls mempty param_ts + + return (param_ts', + param_unnamed_dims) + + let named_shape_params = [ I.Param v $ I.Prim I.int32 | E.TypeParamDim v _ <- tparams ] + shape_params = named_shape_params ++ concat unnamed_shape_params + shape_subst = M.fromList [ (I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params ] + bindingFlatPattern params_idents (concat params_ts') $ \valueparams -> + I.localScope (I.scopeOfFParams $ shape_params++concat valueparams) $ + substitutingVars shape_subst $ m cm shape_params $ chunks num_param_ts (concat valueparams) + +bindingLambdaParams :: [E.TypeParam] -> [E.Pattern] -> [I.Type] + -> (ConstParams -> [I.LParam] -> InternaliseM a) + -> InternaliseM a +bindingLambdaParams tparams params ts m = do + (params_idents, params_types) <- + unzip . concat <$> mapM flattenPattern params + let bound = boundInTypes tparams + param_names = M.fromList [ (E.identName x, y) | (x,y) <- params_idents ] + (params_ts, cm) <- internaliseParamTypes bound param_names params_types + + let ascript_substs = lambdaShapeSubstitutions (concat params_ts) ts + + bindingFlatPattern params_idents ts $ \params' -> + local (\env -> env { envSubsts = ascript_substs `M.union` envSubsts env }) $ + I.localScope (I.scopeOfLParams $ concat params') $ m cm $ concat params' + +processFlatPattern :: Show t => [(E.Ident,VName)] -> [t] + -> InternaliseM ([[I.Param t]], VarSubstitutions) +processFlatPattern x y = processFlatPattern' [] x y + where + processFlatPattern' pat [] _ = do + let (vs, substs) = unzip pat + substs' = M.fromList substs + idents = reverse vs + return (idents, substs') + + processFlatPattern' pat ((p,name):rest) ts = do + (ps, subst, rest_ts) <- handleMapping ts <$> internaliseBindee (p, name) + processFlatPattern' ((ps, (E.identName p, map (I.Var . I.paramName) subst)) : pat) rest rest_ts + + handleMapping ts [] = + ([], [], ts) + handleMapping ts (r:rs) = + let (ps, reps, ts') = handleMapping' ts r + (pss, repss, ts'') = handleMapping ts' rs + in (ps++pss, reps:repss, ts'') + + handleMapping' (t:ts) (vname,_) = + let v' = I.Param vname t + in ([v'], v', ts) + handleMapping' [] _ = + error $ "processFlatPattern: insufficient identifiers in pattern." ++ show (x, y) + + internaliseBindee :: (E.Ident, VName) -> InternaliseM [(VName, I.DeclExtType)] + internaliseBindee (bindee, name) = do + -- XXX: we gotta be screwing up somehow by ignoring the extra + -- return values. If not, why not? + (tss, _) <- internaliseParamTypes nothing_bound mempty + [flip E.setAliases () $ E.vacuousShapeAnnotations $ + E.unInfo $ E.identType bindee] + case concat tss of + [t] -> return [(name, t)] + tss' -> forM tss' $ \t -> do + name' <- newVName $ baseString name + return (name', t) + + -- Fixed up later. + nothing_bound = boundInTypes [] + +bindingFlatPattern :: Show t => [(E.Ident, VName)] -> [t] + -> ([[I.Param t]] -> InternaliseM a) + -> InternaliseM a +bindingFlatPattern idents ts m = do + (ps, substs) <- processFlatPattern idents ts + local (\env -> env { envSubsts = substs `M.union` envSubsts env}) $ + m ps + +-- | Flatten a pattern. Returns a list of identifiers. The +-- structural type of each identifier is returned separately. +flattenPattern :: MonadFreshNames m => E.Pattern -> m [((E.Ident, VName), E.StructType)] +flattenPattern = flattenPattern' + where flattenPattern' (E.PatternParens p _) = + flattenPattern' p + flattenPattern' (E.Wildcard t loc) = do + name <- newVName "nameless" + flattenPattern' $ E.Id name t loc + flattenPattern' (E.Id v (Info t) loc) = do + new_name <- newVName $ baseString v + return [((E.Ident v (Info (E.removeShapeAnnotations t)) loc, + new_name), + t `E.setAliases` ())] + flattenPattern' (E.TuplePattern pats _) = + concat <$> mapM flattenPattern' pats + flattenPattern' (E.RecordPattern fs loc) = + flattenPattern' $ E.TuplePattern (map snd $ sortFields $ M.fromList fs) loc + flattenPattern' (E.PatternAscription p _ _) = + flattenPattern' p + +type MatchPattern = SrcLoc -> [I.SubExp] -> InternaliseM [I.SubExp] + +stmPattern :: [E.TypeParam] -> E.Pattern -> [I.ExtType] + -> (ConstParams -> [VName] -> MatchPattern -> InternaliseM a) + -> InternaliseM a +stmPattern tparams pat ts m = do + (pat', pat_types) <- unzip <$> flattenPattern pat + (ts',_) <- instantiateShapes' ts + (pat_types', cm) <- internaliseParamTypes (boundInTypes tparams) mempty pat_types + let pat_types'' = map I.fromDecl $ concat pat_types' + tparam_names = S.fromList $ map E.typeParamName tparams + let addShapeStms l = + m cm (map I.paramName $ concat l) (matchPattern tparam_names pat_types'') + bindingFlatPattern pat' ts' addShapeStms + +matchPattern :: S.Set VName -> [I.ExtType] -> MatchPattern +matchPattern tparam_names exts loc ses = + forM (zip exts ses) $ \(et, se) -> do + se_t <- I.subExpType se + et' <- unExistentialise tparam_names et se_t + ensureExtShape asserting (I.ErrorMsg [I.ErrorString "value cannot match pattern"]) + loc et' "correct_shape" se + +unExistentialise :: S.Set VName -> I.ExtType -> I.Type -> InternaliseM I.ExtType +unExistentialise tparam_names et t = do + new_dims <- zipWithM inspectDim (I.shapeDims $ I.arrayShape et) (I.arrayDims t) + return $ t `I.setArrayShape` I.Shape new_dims + where inspectDim (I.Free (I.Var v)) d + | v `S.member` tparam_names = do + letBindNames_ [v] $ I.BasicOp $ I.SubExp d + return $ I.Free $ I.Var v + inspectDim ed _ = return ed + +instantiateShapesWithDecls :: MonadFreshNames m => + M.Map Int I.Ident + -> [I.DeclExtType] + -> m ([I.DeclType], [I.FParam]) +instantiateShapesWithDecls ctx ts = + runWriterT $ instantiateShapes instantiate ts + where instantiate x + | Just v <- M.lookup x ctx = + return $ I.Var $ I.identName v + + | otherwise = do + v <- lift $ nonuniqueParamFromIdent <$> newIdent "size" (I.Prim I.int32) + tell [v] + return $ I.Var $ I.paramName v + +lambdaShapeSubstitutions :: [I.TypeBase I.ExtShape Uniqueness] + -> [I.Type] + -> VarSubstitutions +lambdaShapeSubstitutions param_ts ts = + mconcat $ zipWith matchTypes param_ts ts + where matchTypes pt t = + mconcat $ zipWith matchDims (I.shapeDims $ I.arrayShape pt) (I.arrayDims t) + matchDims (I.Free (I.Var v)) d = M.singleton v [d] + matchDims _ _ = mempty + +nonuniqueParamFromIdent :: I.Ident -> I.FParam +nonuniqueParamFromIdent (I.Ident name t) = + I.Param name $ I.toDecl t Nonunique diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/Defunctionalise.hs @@ -0,0 +1,935 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +-- | Defunctionalization of typed, monomorphic Futhark programs without modules. +module Futhark.Internalise.Defunctionalise + ( transformProg ) where + +import Control.Arrow (first, second) +import Control.Monad.RWS +import Data.Bifunctor hiding (first, second) +import Data.Foldable +import Data.List +import Data.Loc +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Data.Semigroup as Sem +import qualified Data.Sequence as Seq + +import Futhark.MonadFreshNames +import Language.Futhark +import Futhark.Representation.AST.Pretty () + +-- | A static value stores additional information about the result of +-- defunctionalization of an expression, aside from the residual expression. +data StaticVal = Dynamic CompType + | LambdaSV [VName] Pattern Exp Env + -- ^ The 'VName's are shape parameters that are bound + -- by the 'Pattern'. + | RecordSV [(Name, StaticVal)] + | DynamicFun (Exp, StaticVal) StaticVal + | IntrinsicSV + deriving (Show) + +-- | Environment mapping variable names to their associated static value. +type Env = M.Map VName StaticVal + +localEnv :: Env -> DefM a -> DefM a +localEnv env = local $ second (env<>) + +-- Even when using a "new" environment (for evaluating closures) we +-- still ram the global environment of DynamicFuns in there. +localNewEnv :: Env -> DefM a -> DefM a +localNewEnv env = local $ \(globals, old_env) -> + (globals, M.filterWithKey (\k _ -> k `S.member` globals) old_env <> env) + +extendEnv :: VName -> StaticVal -> DefM a -> DefM a +extendEnv vn sv = localEnv (M.singleton vn sv) + +askEnv :: DefM Env +askEnv = asks snd + +isGlobal :: VName -> DefM a -> DefM a +isGlobal v = local $ first (S.insert v) + +-- | Returns the defunctionalization environment restricted +-- to the given set of variable names and types. +restrictEnvTo :: NameSet -> DefM Env +restrictEnvTo (NameSet m) = restrict <$> ask + where restrict (globals, env) = M.mapMaybeWithKey keep env + where keep k sv = do guard $ not $ k `S.member` globals + u <- M.lookup k m + Just $ restrict' u sv + restrict' Nonunique (Dynamic t) = + Dynamic $ t `setUniqueness` Nonunique + restrict' _ (Dynamic t) = + Dynamic t + restrict' u (LambdaSV dims pat e env) = + LambdaSV dims pat e $ M.map (restrict' u) env + restrict' u (RecordSV fields) = + RecordSV $ map (fmap $ restrict' u) fields + restrict' u (DynamicFun (e, sv1) sv2) = + DynamicFun (e, restrict' u sv1) $ restrict' u sv2 + restrict' _ IntrinsicSV = IntrinsicSV + +-- | Defunctionalization monad. The Reader environment tracks both +-- the current Env as well as the set of globally defined dynamic +-- functions. This is used to avoid unnecessarily large closure +-- environments. +newtype DefM a = DefM (RWS (Names, Env) (Seq.Seq ValBind) VNameSource a) + deriving (Functor, Applicative, Monad, + MonadReader (Names, Env), + MonadWriter (Seq.Seq ValBind), + MonadFreshNames) + +-- | Run a computation in the defunctionalization monad. Returns the result of +-- the computation, a new name source, and a list of lifted function declations. +runDefM :: VNameSource -> DefM a -> (a, VNameSource, Seq.Seq ValBind) +runDefM src (DefM m) = runRWS m mempty src + +collectFuns :: DefM a -> DefM (a, Seq.Seq ValBind) +collectFuns m = pass $ do + (x, decs) <- listen m + return ((x, decs), const mempty) + +-- | Looks up the associated static value for a given name in the environment. +lookupVar :: SrcLoc -> VName -> DefM StaticVal +lookupVar loc x = do + env <- askEnv + case M.lookup x env of + Just sv -> return sv + Nothing -- If the variable is unknown, it may refer to the 'intrinsics' + -- module, which we will have to treat specially. + | baseTag x <= maxIntrinsicTag -> return IntrinsicSV + | otherwise -> error $ "Variable " ++ pretty x ++ " at " + ++ locStr loc ++ " is out of scope." + +-- | Defunctionalization of an expression. Returns the residual expression and +-- the associated static value in the defunctionalization monad. +defuncExp :: Exp -> DefM (Exp, StaticVal) + +defuncExp e@Literal{} = + return (e, Dynamic $ typeOf e) + +defuncExp e@IntLit{} = + return (e, Dynamic $ typeOf e) + +defuncExp e@FloatLit{} = + return (e, Dynamic $ typeOf e) + +defuncExp (Parens e loc) = do + (e', sv) <- defuncExp e + return (Parens e' loc, sv) + +defuncExp (QualParens qn e loc) = do + (e', sv) <- defuncExp e + return (QualParens qn e' loc, sv) + +defuncExp (TupLit es loc) = do + (es', svs) <- unzip <$> mapM defuncExp es + return (TupLit es' loc, RecordSV $ zip fields svs) + where fields = map (nameFromString . show) [(1 :: Int) ..] + +defuncExp (RecordLit fs loc) = do + (fs', names_svs) <- unzip <$> mapM defuncField fs + return (RecordLit fs' loc, RecordSV names_svs) + + where defuncField (RecordFieldExplicit vn e loc') = do + (e', sv) <- defuncExp e + return (RecordFieldExplicit vn e' loc', (vn, sv)) + defuncField (RecordFieldImplicit vn _ loc') = do + sv <- lookupVar loc' vn + case sv of + -- If the implicit field refers to a dynamic function, we + -- convert it to an explicit field with a record closing over + -- the environment and bind the corresponding static value. + DynamicFun (e, sv') _ -> let vn' = baseName vn + in return (RecordFieldExplicit vn' e loc', + (vn', sv')) + -- The field may refer to a functional expression, so we get the + -- type from the static value and not the one from the AST. + _ -> let tp = Info $ typeFromSV sv + in return (RecordFieldImplicit vn tp loc', (baseName vn, sv)) + +defuncExp (ArrayLit es t@(Info t') loc) = do + es' <- mapM defuncExp' es + return (ArrayLit es' t loc, Dynamic t') + +defuncExp (Range e1 me incl t@(Info t') loc) = do + e1' <- defuncExp' e1 + me' <- mapM defuncExp' me + incl' <- mapM defuncExp' incl + return (Range e1' me' incl' t loc, Dynamic t') + +defuncExp e@(Var qn _ loc) = do + sv <- lookupVar loc (qualLeaf qn) + case sv of + -- If the variable refers to a dynamic function, we return its closure + -- representation (i.e., a record expression capturing the free variables + -- and a 'LambdaSV' static value) instead of the variable itself. + DynamicFun closure _ -> return closure + -- Intrinsic functions used as variables are eta-expanded, so we + -- can get rid of them. + IntrinsicSV -> do + (pats, body, tp) <- etaExpand e + defuncExp $ Lambda [] pats body Nothing (Info (mempty, tp)) noLoc + _ -> let tp = typeFromSV sv + in return (Var qn (Info (vacuousShapeAnnotations tp)) loc, sv) + +defuncExp (Ascript e0 tydecl loc) + | orderZero (typeOf e0) = do (e0', sv) <- defuncExp e0 + return (Ascript e0' tydecl loc, sv) + | otherwise = defuncExp e0 + +defuncExp (LetPat tparams pat e1 e2 loc) = do + let env_dim = envFromShapeParams tparams + (e1', sv1) <- localEnv env_dim $ defuncExp e1 + let env = matchPatternSV pat sv1 + pat' = updatePattern pat sv1 + (e2', sv2) <- localEnv (env <> env_dim) $ defuncExp e2 + return (LetPat tparams pat' e1' e2' loc, sv2) + +defuncExp (LetFun vn (dims, pats, _, rettype@(Info ret), e1) e2 loc) = do + let env_dim = envFromShapeParams dims + (pats', e1', sv1) <- localEnv env_dim $ defuncLet dims pats e1 rettype + (e2', sv2) <- extendEnv vn sv1 $ defuncExp e2 + case pats' of + [] -> let t1 = combineTypeShapes (fromStruct ret) $ + vacuousShapeAnnotations $ typeOf e1' + in return (LetPat dims (Id vn (Info t1) noLoc) e1' e2' loc, sv2) + _:_ -> let t1 = combineTypeShapes ret $ + vacuousShapeAnnotations . toStruct $ typeOf e1' + in return (LetFun vn (dims, pats', Nothing, Info t1, e1') e2' loc, sv2) + +defuncExp (If e1 e2 e3 tp loc) = do + (e1', _ ) <- defuncExp e1 + (e2', sv) <- defuncExp e2 + (e3', _ ) <- defuncExp e3 + return (If e1' e2' e3' tp loc, sv) + +defuncExp e@Apply{} = defuncApply 0 e + +defuncExp (Negate e0 loc) = do + (e0', sv) <- defuncExp e0 + return (Negate e0' loc, sv) + +defuncExp e@(Lambda tparams pats e0 decl tp loc) = do + when (any isTypeParam tparams) $ + error $ "Received a lambda with type parameters at " ++ locStr loc + ++ ", but the defunctionalizer expects a monomorphic input program." + -- Extract the first parameter of the lambda and "push" the + -- remaining ones (if there are any) into the body of the lambda. + let (dims, pat, e0') = case pats of + [] -> error "Received a lambda with no parameters." + [pat'] -> (map typeParamName tparams, pat', e0) + (pat' : pats') -> + -- Split shape parameters into those that are determined by + -- the first pattern, and those that are determined by later + -- patterns. + let bound_by_pat = (`S.member` patternDimNames pat') . typeParamName + (pat_dims, rest_dims) = partition bound_by_pat tparams + in (map typeParamName pat_dims, pat', + Lambda rest_dims pats' e0 decl tp loc) + + -- Construct a record literal that closes over the environment of + -- the lambda. Closed-over 'DynamicFun's are converted to their + -- closure representation. + env <- restrictEnvTo (freeVars e) + let (fields, env') = unzip $ map closureFromDynamicFun $ M.toList env + return (RecordLit fields loc, LambdaSV dims pat e0' $ M.fromList env') + + where closureFromDynamicFun (vn, DynamicFun (clsr_env, sv) _) = + let name = nameFromString $ pretty vn + in (RecordFieldExplicit name clsr_env noLoc, (vn, sv)) + + closureFromDynamicFun (vn, sv) = + let name = nameFromString $ pretty vn + tp' = vacuousShapeAnnotations $ typeFromSV sv + in (RecordFieldExplicit name + (Var (qualName vn) (Info tp') noLoc) noLoc, (vn, sv)) + +-- Operator sections are expected to be converted to lambda-expressions +-- by the monomorphizer, so they should no longer occur at this point. +defuncExp OpSection{} = error "defuncExp: unexpected operator section." +defuncExp OpSectionLeft{} = error "defuncExp: unexpected operator section." +defuncExp OpSectionRight{} = error "defuncExp: unexpected operator section." +defuncExp ProjectSection{} = error "defuncExp: unexpected projection section." +defuncExp IndexSection{} = error "defuncExp: unexpected projection section." + +defuncExp (DoLoop tparams pat e1 form e3 loc) = do + let env_dim = envFromShapeParams tparams + (e1', sv1) <- defuncExp e1 + let env1 = matchPatternSV pat sv1 + (form', env2) <- case form of + For v e2 -> do e2' <- defuncExp' e2 + return (For v e2', envFromIdent v) + ForIn pat2 e2 -> do e2' <- defuncExp' e2 + return (ForIn pat2 e2', envFromPattern pat2) + While e2 -> do e2' <- localEnv (env1 <> env_dim) $ defuncExp' e2 + return (While e2', mempty) + (e3', sv) <- localEnv (env1 <> env2 <> env_dim) $ defuncExp e3 + return (DoLoop tparams pat e1' form' e3' loc, sv) + where envFromIdent (Ident vn (Info tp) _) = M.singleton vn $ Dynamic tp + +-- We handle BinOps by turning them into ordinary function applications. +defuncExp (BinOp qn (Info t) (e1, Info pt1) (e2, Info pt2) (Info ret) loc) = + defuncExp $ Apply (Apply (Var qn (Info t) loc) + e1 (Info (diet pt1)) (Info (Arrow mempty Nothing (fromStruct pt2) ret)) loc) + e2 (Info (diet pt2)) (Info ret) loc + +defuncExp (Project vn e0 tp@(Info tp') loc) = do + (e0', sv0) <- defuncExp e0 + case sv0 of + RecordSV svs -> case lookup vn svs of + Just sv -> return (Project vn e0' (Info $ typeFromSV sv) loc, sv) + Nothing -> error "Invalid record projection." + Dynamic _ -> return (Project vn e0' tp loc, Dynamic tp') + _ -> error $ "Projection of an expression with static value " ++ show sv0 + +defuncExp (LetWith id1 id2 idxs e1 body loc) = do + e1' <- defuncExp' e1 + sv1 <- lookupVar (identSrcLoc id2) $ identName id2 + idxs' <- mapM defuncDimIndex idxs + (body', sv) <- extendEnv (identName id1) sv1 $ defuncExp body + return (LetWith id1 id2 idxs' e1' body' loc, sv) + +defuncExp expr@(Index e0 idxs info loc) = do + e0' <- defuncExp' e0 + idxs' <- mapM defuncDimIndex idxs + return (Index e0' idxs' info loc, Dynamic $ typeOf expr) + +defuncExp (Update e1 idxs e2 loc) = do + (e1', sv) <- defuncExp e1 + idxs' <- mapM defuncDimIndex idxs + e2' <- defuncExp' e2 + return (Update e1' idxs' e2' loc, sv) + +-- Note that we might change the type of the record field here. This +-- is not permitted in the type checker due to problems with type +-- inference, but it actually works fine. +defuncExp (RecordUpdate e1 fs e2 _ loc) = do + (e1', sv1) <- defuncExp e1 + (e2', sv2) <- defuncExp e2 + let sv = staticField sv1 sv2 fs + return (RecordUpdate e1' fs e2' + (Info $ vacuousShapeAnnotations $ typeFromSV sv) loc, + sv) + where staticField (RecordSV svs) sv2 (f:fs') = + case lookup f svs of + Just sv -> RecordSV $ + (f, staticField sv sv2 fs') : filter ((/=f) . fst) svs + Nothing -> error "Invalid record projection." + staticField _ sv2 _ = sv2 + +defuncExp e@(Map fun arr t loc) = do + fun' <- defuncSoacExp fun + arr' <- defuncExp' arr + return (Map fun' arr' t loc, Dynamic $ typeOf e) + +defuncExp e@(Reduce comm fun ne arr loc) = do + fun' <- defuncSoacExp fun + ne' <- defuncExp' ne + arr' <- defuncExp' arr + return (Reduce comm fun' ne' arr' loc, Dynamic $ typeOf e) + +defuncExp e@(GenReduce hist op ne bfun img loc) = do + hist' <- defuncExp' hist + op' <- defuncSoacExp op + ne' <- defuncExp' ne + bfun' <- defuncSoacExp bfun + img' <- defuncExp' img + return (GenReduce hist' op' ne' bfun' img' loc, Dynamic $ typeOf e) + +defuncExp e@(Scan fun ne arr loc) = + (,) <$> (Scan <$> defuncSoacExp fun <*> defuncExp' ne <*> defuncExp' arr + <*> pure loc) + <*> pure (Dynamic $ typeOf e) + +defuncExp e@(Filter fun arr loc) = do + fun' <- defuncSoacExp fun + arr' <- defuncExp' arr + return (Filter fun' arr' loc, Dynamic $ typeOf e) + +defuncExp e@(Partition k fun arr loc) = do + fun' <- defuncSoacExp fun + arr' <- defuncExp' arr + return (Partition k fun' arr' loc, Dynamic $ typeOf e) + +defuncExp e@(Stream form lam arr loc) = do + form' <- case form of + MapLike _ -> return form + RedLike so comm e' -> RedLike so comm <$> defuncSoacExp e' + lam' <- defuncSoacExp lam + arr' <- defuncExp' arr + return (Stream form' lam' arr' loc, Dynamic $ typeOf e) + +defuncExp e@(Zip i e1 es t loc) = do + e1' <- defuncExp' e1 + es' <- mapM defuncExp' es + return (Zip i e1' es' t loc, Dynamic $ typeOf e) + +defuncExp e@(Unzip e0 tps loc) = do + e0' <- defuncExp' e0 + return (Unzip e0' tps loc, Dynamic $ typeOf e) + +defuncExp (Unsafe e1 loc) = do + (e1', sv) <- defuncExp e1 + return (Unsafe e1' loc, sv) + +defuncExp (Assert e1 e2 desc loc) = do + (e1', _) <- defuncExp e1 + (e2', sv) <- defuncExp e2 + return (Assert e1' e2' desc loc, sv) + +-- | Same as 'defuncExp', except it ignores the static value. +defuncExp' :: Exp -> DefM Exp +defuncExp' = fmap fst . defuncExp + +-- | Defunctionalize the function argument to a SOAC by eta-expanding if +-- necessary and then defunctionalizing the body of the introduced lambda. +defuncSoacExp :: Exp -> DefM Exp +defuncSoacExp e@OpSection{} = return e +defuncSoacExp e@OpSectionLeft{} = return e +defuncSoacExp e@OpSectionRight{} = return e +defuncSoacExp e@ProjectSection{} = return e + +defuncSoacExp (Parens e loc) = + Parens <$> defuncSoacExp e <*> pure loc + +defuncSoacExp (Lambda tparams params e0 decl tp loc) = do + let env_dim = envFromShapeParams tparams + env = foldMap envFromPattern params + e0' <- localEnv (env <> env_dim) $ defuncSoacExp e0 + return $ Lambda tparams params e0' decl tp loc + +defuncSoacExp e + | Arrow{} <- typeOf e = do + (pats, body, tp) <- etaExpand e + let env = foldMap envFromPattern pats + body' <- localEnv env $ defuncExp' body + return $ Lambda [] pats body' Nothing (Info (mempty, tp)) noLoc + | otherwise = defuncExp' e + +etaExpand :: Exp -> DefM ([Pattern], Exp, StructType) +etaExpand e = do + let (ps, ret) = getType $ typeOf e + (pats, vars) <- fmap unzip . forM ps $ \t -> do + x <- newNameFromString "x" + let t' = vacuousShapeAnnotations t + return (Id x (Info t') noLoc, + Var (qualName x) (Info t') noLoc) + let ps_st = map vacuousShapeAnnotations ps + e' = foldl' (\e1 (e2, t2, argtypes) -> + Apply e1 e2 (Info $ diet t2) + (Info (foldFunType argtypes (vacuousShapeAnnotations ret))) noLoc) + e $ zip3 vars ps (drop 1 $ tails ps_st) + return (pats, e', vacuousShapeAnnotations $ toStruct ret) + + where getType (Arrow _ _ t1 t2) = + let (ps, r) = getType t2 in (t1 : ps, r) + getType t = ([], t) + +-- | Defunctionalize an indexing of a single array dimension. +defuncDimIndex :: DimIndexBase Info VName -> DefM (DimIndexBase Info VName) +defuncDimIndex (DimFix e1) = DimFix . fst <$> defuncExp e1 +defuncDimIndex (DimSlice me1 me2 me3) = + DimSlice <$> defunc' me1 <*> defunc' me2 <*> defunc' me3 + where defunc' = mapM defuncExp' + +-- | Defunctionalize a let-bound function, while preserving parameters +-- that have order 0 types (i.e., non-functional). +defuncLet :: [TypeParam] -> [Pattern] -> Exp -> Info StructType + -> DefM ([Pattern], Exp, StaticVal) +defuncLet dims ps@(pat:pats) body (Info rettype) + | patternOrderZero pat = do + let env = envFromPattern pat + bound_by_pat = (`S.member` patternDimNames pat) . typeParamName + (_pat_dims, rest_dims) = partition bound_by_pat dims + (pats', body', sv) <- localEnv env $ defuncLet rest_dims pats body (Info rettype) + closure <- defuncExp $ Lambda dims ps body Nothing (Info (mempty, rettype)) noLoc + return (pat : pats', body', DynamicFun closure sv) + | otherwise = do + (e, sv) <- defuncExp $ Lambda dims ps body Nothing (Info (mempty, rettype)) noLoc + return ([], e, sv) +defuncLet _ [] body (Info rettype) = do + (body', sv) <- defuncExp body + case sv of + Dynamic _ -> return ([], body', Dynamic $ fromStruct $ removeShapeAnnotations rettype) + _ -> return ([], body', sv) + +-- | Defunctionalize an application expression at a given depth of application. +-- Calls to dynamic (first-order) functions are preserved at much as possible, +-- but a new lifted function is created if a dynamic function is only partially +-- applied. +defuncApply :: Int -> Exp -> DefM (Exp, StaticVal) +defuncApply depth e@(Apply e1 e2 d t@(Info ret) loc) = do + let (argtypes, _) = unfoldFunType ret + (e1', sv1) <- defuncApply (depth+1) e1 + (e2', sv2) <- defuncExp e2 + let e' = Apply e1' e2' d t loc + case sv1 of + LambdaSV dims pat e0 closure_env -> do + let env' = matchPatternSV pat sv2 + env_dim = envFromDimNames dims + (e0', sv) <- localNewEnv (env' <> closure_env <> env_dim) $ defuncExp e0 + + let closure_pat = buildEnvPattern closure_env + pat' = updatePattern pat sv2 + + -- Inline certain trivial lifted functions immediately. This is + -- purely an optimisation to avoid having the rest of the + -- compiler spend a lot if time processing them (they will end + -- up being inlined later anyway). We also try to simplify away + -- some let-bindings, to make the generated code look slightly + -- more comprehensible. + -- + -- If you are debugging the defunctionaliser, you may want to + -- turn this off to see the original structure of the generated + -- code instead. + let letPat (RecordPattern [] _) _ pbody = pbody + letPat (PatternParens pp _) pe pbody = letPat pp pe pbody + letPat (Id v1 _ _) pe (RecordLit [RecordFieldExplicit f (Var v2 _ _) floc] rloc) + | v1 == qualLeaf v2 = + RecordLit [RecordFieldExplicit f pe floc] rloc + letPat (RecordPattern [(pf,p)] _) (RecordLit [RecordFieldExplicit f pe _] _) pbody + | pf == f = + letPat p pe pbody + letPat pp pe pbody = LetPat [] pp pe pbody noLoc + + inline RecordLit{} = True + inline TupLit{} = True + inline (Apply x y _ _ _) = inline x && inline y + inline (BinOp _ _ (x, _) (y, _) _ _) = inline x && inline y + inline Var{} = True + inline Literal{} = True + inline (LetPat _ _ x y _) = inline x && inline y + inline Negate{} = True + inline _ = False + if inline e0' && null dims + then return (letPat closure_pat e1' $ letPat pat' e2' e0', sv) + else do + -- Lift lambda to top-level function definition. + let params = [closure_pat, pat'] + rettype = buildRetType closure_env params $ typeOf e0' + + -- Embed some information about the original function + -- into the name of the lifted function, to make the + -- result slightly more human-readable. + liftedName i (Var f _ _) = + "lifted_" ++ show i ++ "_" ++ baseString (qualLeaf f) + liftedName i (Apply f _ _ _ _) = + liftedName (i+1) f + liftedName _ _ = "lifted" + fname <- newNameFromString $ liftedName (0::Int) e1 + liftValDec fname rettype dims params e0' + + let t1 = vacuousShapeAnnotations . toStruct $ typeOf e1' + t2 = vacuousShapeAnnotations . toStruct $ typeOf e2' + fname' = qualName fname + return (Parens (Apply (Apply (Var fname' (Info (Arrow mempty Nothing (fromStruct t1) $ + Arrow mempty Nothing (fromStruct t2) rettype)) loc) + e1' (Info Observe) (Info $ Arrow mempty Nothing (fromStruct t2) rettype) loc) + e2' d (Info rettype) loc) noLoc, sv) + + -- If e1 is a dynamic function, we just leave the application in place, + -- but we update the types since it may be partially applied or return + -- a higher-order term. + DynamicFun _ sv -> + let (argtypes', rettype) = dynamicFunType sv argtypes + apply_e = Apply e1' e2' d (Info $ foldFunType argtypes' rettype) loc + in return (apply_e, sv) + + -- Propagate the 'IntrinsicsSV' until we reach the outermost application, + -- where we construct a dynamic static value with the appropriate type. + IntrinsicSV + | depth == 0 -> return (e', Dynamic $ typeOf e) + | otherwise -> return (e', IntrinsicSV) + + _ -> error $ "Application of an expression that is neither a static lambda " + ++ "nor a dynamic function, but has static value: " ++ show sv1 + +defuncApply depth e@(Var qn (Info t) loc) = do + let (argtypes, _) = unfoldFunType t + sv <- lookupVar loc (qualLeaf qn) + case sv of + DynamicFun _ _ + | fullyApplied sv depth -> + -- We still need to update the types in case the dynamic + -- function returns a higher-order term. + let (argtypes', rettype) = dynamicFunType sv argtypes + in return (Var qn (Info (foldFunType argtypes' rettype)) loc, sv) + + | otherwise -> do + fname <- newName $ qualLeaf qn + let (dims, pats, e0, sv') = liftDynFun sv depth + (argtypes', rettype) = dynamicFunType sv' argtypes + liftValDec fname rettype dims pats e0 + return (Var (qualName fname) + (Info (foldFunType argtypes' rettype)) loc, sv') + + IntrinsicSV -> return (e, IntrinsicSV) + + _ -> return (Var qn (Info (vacuousShapeAnnotations $ typeFromSV sv)) loc, sv) + +defuncApply _ expr = defuncExp expr + +-- | Check if a 'StaticVal' and a given application depth corresponds +-- to a fully applied dynamic function. +fullyApplied :: StaticVal -> Int -> Bool +fullyApplied (DynamicFun _ sv) depth + | depth == 0 = False + | depth > 0 = fullyApplied sv (depth-1) +fullyApplied _ _ = True + +-- | Converts a dynamic function 'StaticVal' into a list of +-- dimensions, a list of parameters, a function body, and the +-- appropriate static value for applying the function at the given +-- depth of partial application. +liftDynFun :: StaticVal -> Int -> ([VName], [Pattern], Exp, StaticVal) +liftDynFun (DynamicFun (e, sv) _) 0 = ([], [], e, sv) +liftDynFun (DynamicFun clsr@(_, LambdaSV dims pat _ _) sv) d + | d > 0 = let (dims', pats, e', sv') = liftDynFun sv (d-1) + in (dims ++ dims', pat : pats, e', DynamicFun clsr sv') +liftDynFun sv _ = error $ "Tried to lift a StaticVal " ++ show sv + ++ ", but expected a dynamic function." + +-- | Converts a pattern to an environment that binds the individual names of the +-- pattern to their corresponding types wrapped in a 'Dynamic' static value. +envFromPattern :: Pattern -> Env +envFromPattern pat = case pat of + TuplePattern ps _ -> foldMap envFromPattern ps + RecordPattern fs _ -> foldMap (envFromPattern . snd) fs + PatternParens p _ -> envFromPattern p + Id vn (Info t) _ -> M.singleton vn $ Dynamic $ removeShapeAnnotations t + Wildcard _ _ -> mempty + PatternAscription p _ _ -> envFromPattern p + +-- | Create an environment that binds the shape parameters. +envFromShapeParams :: [TypeParamBase VName] -> Env +envFromShapeParams = envFromDimNames . map dim + where dim (TypeParamDim vn _) = vn + dim tparam = error $ + "The defunctionalizer expects a monomorphic input program,\n" ++ + "but it received a type parameter " ++ pretty tparam ++ + " at " ++ locStr (srclocOf tparam) ++ "." + +envFromDimNames :: [VName] -> Env +envFromDimNames = M.fromList . flip zip (repeat $ Dynamic $ Prim $ Signed Int32) + +-- | Create a new top-level value declaration with the given function name, +-- return type, list of parameters, and body expression. +liftValDec :: VName -> PatternType -> [VName] -> [Pattern] -> Exp -> DefM () +liftValDec fname rettype dims pats body = tell $ Seq.singleton dec + where dims' = map (flip TypeParamDim noLoc) dims + rettype_st = vacuousShapeAnnotations $ toStruct rettype + dec = ValBind + { valBindEntryPoint = False + , valBindName = fname + , valBindRetDecl = Nothing + , valBindRetType = Info rettype_st + , valBindTypeParams = dims' + , valBindParams = pats + , valBindBody = body + , valBindDoc = Nothing + , valBindLocation = noLoc + } + +-- | Given a closure environment, construct a record pattern that +-- binds the closed over variables. +buildEnvPattern :: Env -> Pattern +buildEnvPattern env = RecordPattern (map buildField $ M.toList env) noLoc + where buildField (vn, sv) = let tp = vacuousShapeAnnotations (typeFromSV sv) + in (nameFromString (pretty vn), + Id vn (Info tp) noLoc) + +-- | Given a closure environment pattern and the type of a term, +-- construct the type of that term, where uniqueness is set to +-- `Nonunique` for those arrays that are bound in the environment or +-- pattern (except if they are unique there). This ensures that a +-- lifted function can create unique arrays as long as they do not +-- alias any of its parameters. XXX: it is not clear that this is a +-- sufficient property, unfortunately. +buildRetType :: Env -> [Pattern] -> CompType -> PatternType +buildRetType env pats = vacuousShapeAnnotations . descend + where bound = foldMap oneName (M.keys env) <> foldMap patternVars pats + boundAsUnique v = + maybe False (unique . unInfo . identType) $ + find ((==v) . identName) $ S.toList $ foldMap patIdentSet pats + problematic v = (v `member` bound) && not (boundAsUnique v) + descend t@Array{} + | any problematic (aliases t) = t `setUniqueness` Nonunique + descend (Record t) = Record $ fmap descend t + descend t = t + +-- | Compute the corresponding type for a given static value. +typeFromSV :: StaticVal -> CompType +typeFromSV (Dynamic tp) = tp +typeFromSV (LambdaSV _ _ _ env) = typeFromEnv env +typeFromSV (RecordSV ls) = Record $ M.fromList $ map (fmap typeFromSV) ls +typeFromSV (DynamicFun (_, sv) _) = typeFromSV sv +typeFromSV IntrinsicSV = error $ "Tried to get the type from the " + ++ "static value of an intrinsic." + +typeFromEnv :: Env -> CompType +typeFromEnv = Record . M.fromList . + map (bimap (nameFromString . pretty) typeFromSV) . M.toList + +-- | Construct the type for a fully-applied dynamic function from its +-- static value and the original types of its arguments. +dynamicFunType :: StaticVal -> [PatternType] -> ([PatternType], PatternType) +dynamicFunType (DynamicFun _ sv) (p:ps) = + let (ps', ret) = dynamicFunType sv ps in (p : ps', ret) +dynamicFunType sv _ = ([], vacuousShapeAnnotations $ typeFromSV sv) + +-- | Match a pattern with its static value. Returns an environment with +-- the identifier components of the pattern mapped to the corresponding +-- subcomponents of the static value. +matchPatternSV :: PatternBase Info VName -> StaticVal -> Env +matchPatternSV (TuplePattern ps _) (RecordSV ls) = + mconcat $ zipWith (\p (_, sv) -> matchPatternSV p sv) ps ls +matchPatternSV (RecordPattern ps _) (RecordSV ls) + | ps' <- sortOn fst ps, ls' <- sortOn fst ls, + map fst ps' == map fst ls' = + mconcat $ zipWith (\(_, p) (_, sv) -> matchPatternSV p sv) ps' ls' +matchPatternSV (PatternParens pat _) sv = matchPatternSV pat sv +matchPatternSV (Id vn (Info t) _) sv = + -- When matching a pattern with a zero-order STaticVal, the type of + -- the pattern wins out. This is important when matching a + -- nonunique pattern with a unique value. + if orderZeroSV sv + then M.singleton vn $ Dynamic $ removeShapeAnnotations t + else M.singleton vn sv +matchPatternSV (Wildcard _ _) _ = mempty +matchPatternSV (PatternAscription pat _ _) sv = matchPatternSV pat sv +matchPatternSV pat (Dynamic t) = matchPatternSV pat $ svFromType t +matchPatternSV pat sv = error $ "Tried to match pattern " ++ pretty pat + ++ " with static value " ++ show sv ++ "." + +orderZeroSV :: StaticVal -> Bool +orderZeroSV Dynamic{} = True +orderZeroSV (RecordSV fields) = all (orderZeroSV . snd) fields +orderZeroSV _ = False + +-- | Given a pattern and the static value for the defunctionalized argument, +-- update the pattern to reflect the changes in the types. +updatePattern :: Pattern -> StaticVal -> Pattern +updatePattern (TuplePattern ps loc) (RecordSV svs) = + TuplePattern (zipWith updatePattern ps $ map snd svs) loc +updatePattern (RecordPattern ps loc) (RecordSV svs) + | ps' <- sortOn fst ps, svs' <- sortOn fst svs = + RecordPattern (zipWith (\(n, p) (_, sv) -> + (n, updatePattern p sv)) ps' svs') loc +updatePattern (PatternParens pat loc) sv = + PatternParens (updatePattern pat sv) loc +updatePattern pat@(Id vn (Info tp) loc) sv + | orderZero tp = pat + | otherwise = Id vn (Info . vacuousShapeAnnotations $ + typeFromSV sv `setUniqueness` Nonunique) loc +updatePattern pat@(Wildcard (Info tp) loc) sv + | orderZero tp = pat + | otherwise = Wildcard (Info . vacuousShapeAnnotations $ typeFromSV sv) loc +updatePattern (PatternAscription pat tydecl loc) sv + | orderZero . unInfo $ expandedType tydecl = + PatternAscription (updatePattern pat sv) tydecl loc + | otherwise = updatePattern pat sv +updatePattern pat (Dynamic t) = updatePattern pat (svFromType t) +updatePattern pat sv = + error $ "Tried to update pattern " ++ pretty pat + ++ "to reflect the static value " ++ show sv + +-- | Convert a record (or tuple) type to a record static value. This is used for +-- "unwrapping" tuples and records that are nested in 'Dynamic' static values. +svFromType :: CompType -> StaticVal +svFromType (Record fs) = RecordSV . M.toList $ M.map svFromType fs +svFromType t = Dynamic t + +-- A set of names where we also track uniqueness. +newtype NameSet = NameSet (M.Map VName Uniqueness) + +instance Sem.Semigroup NameSet where + NameSet x <> NameSet y = NameSet $ M.unionWith max x y + +instance Monoid NameSet where + mempty = NameSet mempty + mappend = (Sem.<>) + +without :: NameSet -> NameSet -> NameSet +without (NameSet x) (NameSet y) = NameSet $ x `M.difference` y + +member :: VName -> NameSet -> Bool +member v (NameSet m) = v `M.member` m + +ident :: Ident -> NameSet +ident v = NameSet $ M.singleton (identName v) (uniqueness $ unInfo $ identType v) + +oneName :: VName -> NameSet +oneName v = NameSet $ M.singleton v Nonunique + +names :: Names -> NameSet +names = foldMap oneName + +-- | Compute the set of free variables of an expression. +freeVars :: Exp -> NameSet +freeVars expr = case expr of + Literal{} -> mempty + IntLit{} -> mempty + FloatLit{} -> mempty + Parens e _ -> freeVars e + QualParens _ e _ -> freeVars e + TupLit es _ -> foldMap freeVars es + + RecordLit fs _ -> foldMap freeVarsField fs + where freeVarsField (RecordFieldExplicit _ e _) = freeVars e + freeVarsField (RecordFieldImplicit vn t _) = ident $ Ident vn t noLoc + + ArrayLit es _ _ -> foldMap freeVars es + Range e me incl _ _ -> freeVars e <> foldMap freeVars me <> + foldMap freeVars incl + Var qn (Info t) _ -> NameSet $ M.singleton (qualLeaf qn) $ uniqueness t + Ascript e t _ -> freeVars e <> names (typeDimNames $ unInfo $ expandedType t) + LetPat _ pat e1 e2 _ -> freeVars e1 <> ((names (patternDimNames pat) <> freeVars e2) + `without` patternVars pat) + + LetFun vn (_, pats, _, _, e1) e2 _ -> + ((freeVars e1 <> names (foldMap patternDimNames pats)) + `without` foldMap patternVars pats) <> + (freeVars e2 `without` oneName vn) + + If e1 e2 e3 _ _ -> freeVars e1 <> freeVars e2 <> freeVars e3 + Apply e1 e2 _ _ _ -> freeVars e1 <> freeVars e2 + Negate e _ -> freeVars e + Lambda tps pats e0 _ _ _ -> (names (foldMap patternDimNames pats) <> freeVars e0) + `without` (foldMap patternVars pats <> + mconcat (map (oneName . typeParamName) tps)) + OpSection{} -> mempty + OpSectionLeft _ _ e _ _ _ -> freeVars e + OpSectionRight _ _ e _ _ _ -> freeVars e + ProjectSection{} -> mempty + IndexSection idxs _ _ -> foldMap freeDimIndex idxs + + DoLoop _ pat e1 form e3 _ -> let (e2fv, e2ident) = formVars form + in freeVars e1 <> e2fv <> + (freeVars e3 `without` (patternVars pat <> e2ident)) + where formVars (For v e2) = (freeVars e2, ident v) + formVars (ForIn p e2) = (freeVars e2, patternVars p) + formVars (While e2) = (freeVars e2, mempty) + + BinOp qn _ (e1, _) (e2, _) _ _ -> oneName (qualLeaf qn) <> + freeVars e1 <> freeVars e2 + Project _ e _ _ -> freeVars e + + LetWith id1 id2 idxs e1 e2 _ -> + ident id2 <> foldMap freeDimIndex idxs <> freeVars e1 <> + (freeVars e2 `without` ident id1) + + Index e idxs _ _ -> freeVars e <> foldMap freeDimIndex idxs + Update e1 idxs e2 _ -> freeVars e1 <> foldMap freeDimIndex idxs <> freeVars e2 + RecordUpdate e1 _ e2 _ _ -> freeVars e1 <> freeVars e2 + + Map e1 e2 _ _ -> freeVars e1 <> freeVars e2 + Reduce _ e1 e2 e3 _ -> freeVars e1 <> freeVars e2 <> freeVars e3 + GenReduce e1 e2 e3 e4 e5 _ -> freeVars e1 <> freeVars e2 <> freeVars e3 <> + freeVars e4 <> freeVars e5 + Scan e1 e2 e3 _ -> freeVars e1 <> freeVars e2 <> freeVars e3 + Filter e1 e2 _ -> freeVars e1 <> freeVars e2 + Partition _ e1 e2 _ -> freeVars e1 <> freeVars e2 + Stream form e1 e2 _ -> freeInForm form <> freeVars e1 <> freeVars e2 + where freeInForm (RedLike _ _ e) = freeVars e + freeInForm _ = mempty + + Zip _ e es _ _ -> freeVars e <> foldMap freeVars es + Unzip e _ _ -> freeVars e + Unsafe e _ -> freeVars e + Assert e1 e2 _ _ -> freeVars e1 <> freeVars e2 + +freeDimIndex :: DimIndexBase Info VName -> NameSet +freeDimIndex (DimFix e) = freeVars e +freeDimIndex (DimSlice me1 me2 me3) = + foldMap (foldMap freeVars) [me1, me2, me3] + +-- | Extract all the variable names bound in a pattern. +patternVars :: Pattern -> NameSet +patternVars = mconcat . map ident . S.toList . patIdentSet + +-- | Combine the shape information of types as much as possible. The first +-- argument is the orignal type and the second is the type of the transformed +-- expression. This is necessary since the original type may contain additional +-- information (e.g., shape restrictions) from the user given annotation. +combineTypeShapes :: ArrayDim dim => + TypeBase dim as -> TypeBase dim as -> TypeBase dim as +combineTypeShapes (Record ts1) (Record ts2) + | M.keys ts1 == M.keys ts2 = + Record $ M.map (uncurry combineTypeShapes) (M.intersectionWith (,) ts1 ts2) +combineTypeShapes (Array et1 shape1 u1) (Array et2 shape2 _u2) + | Just new_shape <- unifyShapes shape1 shape2 = + Array (combineElemTypeInfo et1 et2) new_shape u1 +combineTypeShapes _ new_tp = new_tp + +combineElemTypeInfo :: ArrayDim dim => + ArrayElemTypeBase dim as + -> ArrayElemTypeBase dim as -> ArrayElemTypeBase dim as +combineElemTypeInfo (ArrayRecordElem et1) (ArrayRecordElem et2) = + ArrayRecordElem $ M.map (uncurry combineRecordArrayTypeInfo) + (M.intersectionWith (,) et1 et2) +combineElemTypeInfo _ new_tp = new_tp + +combineRecordArrayTypeInfo :: ArrayDim dim => + RecordArrayElemTypeBase dim as + -> RecordArrayElemTypeBase dim as + -> RecordArrayElemTypeBase dim as +combineRecordArrayTypeInfo (RecordArrayElem et1) (RecordArrayElem et2) = + RecordArrayElem $ combineElemTypeInfo et1 et2 +combineRecordArrayTypeInfo (RecordArrayArrayElem et1 shape1 u1) + (RecordArrayArrayElem et2 shape2 u2) + | Just new_shape <- unifyShapes shape1 shape2 = + RecordArrayArrayElem (combineElemTypeInfo et1 et2) new_shape (u1 <> u2) +combineRecordArrayTypeInfo _ new_tp = new_tp + +-- | Defunctionalize a top-level value binding. Returns the +-- transformed result as well as an environment that binds the name of +-- the value binding to the static value of the transformed body. The +-- boolean is true if the function is a 'DynamicFun'. +defuncValBind :: ValBind -> DefM (ValBind, Env, Bool) + +-- Eta-expand entry points with a functional return type. +defuncValBind (ValBind True name retdecl (Info rettype) tparams params body _ loc) + | (rettype_ps, rettype') <- unfoldFunType rettype, + not $ null rettype_ps = do + (body_pats, body', _) <- etaExpand body + defuncValBind $ ValBind True name retdecl (Info rettype') + tparams (params <> body_pats) body' Nothing loc + +defuncValBind valbind@(ValBind _ name retdecl rettype tparams params body _ _) = do + let env = envFromShapeParams tparams + (params', body', sv) <- localEnv env $ defuncLet tparams params body rettype + -- Remove any shape parameters that no longer occur in the value parameters. + let dim_names = foldMap patternDimNames params' + tparams' = filter ((`S.member` dim_names) . typeParamName) tparams + + let rettype' = vacuousShapeAnnotations . toStruct $ typeOf body' + return ( valbind { valBindRetDecl = retdecl + , valBindRetType = Info $ combineTypeShapes + (unInfo rettype) rettype' + , valBindTypeParams = tparams' + , valBindParams = params' + , valBindBody = body' + } + , M.singleton name sv + , case sv of DynamicFun{} -> True + _ -> False) + +-- | Defunctionalize a list of top-level declarations. +defuncVals :: [ValBind] -> DefM (Seq.Seq ValBind) +defuncVals [] = return mempty +defuncVals (valbind : ds) = do + ((valbind', env, dyn), defs) <- collectFuns $ defuncValBind valbind + ds' <- localEnv env $ if dyn + then isGlobal (valBindName valbind') $ defuncVals ds + else defuncVals ds + return $ defs <> Seq.singleton valbind' <> ds' + +-- | Transform a list of top-level value bindings. May produce new +-- lifted function definitions, which are placed in front of the +-- resulting list of declarations. +transformProg :: MonadFreshNames m => [ValBind] -> m [ValBind] +transformProg decs = modifyNameSource $ \namesrc -> + let (decs', namesrc', liftedDecs) = runDefM namesrc $ defuncVals decs + in (toList $ liftedDecs <> decs', namesrc') diff --git a/src/Futhark/Internalise/Defunctorise.hs b/src/Futhark/Internalise/Defunctorise.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/Defunctorise.hs @@ -0,0 +1,299 @@ +-- | Partially evaluate all modules away from a source Futhark +-- program. This is implemented as a source-to-source transformation. +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Futhark.Internalise.Defunctorise (transformProg) where + +import Control.Monad.RWS.Strict +import Control.Monad.Identity +import qualified Data.DList as DL +import qualified Data.Map as M +import qualified Data.Set as S +import Data.Maybe +import Data.Loc +import qualified Data.Semigroup as Sem + +import Prelude hiding (mod, abs) + +import Futhark.MonadFreshNames +import Language.Futhark +import Language.Futhark.Traversals +import Language.Futhark.Semantic (Imports, FileModule(..)) + +-- | A substitution from names in the original program to names in the +-- generated/residual program. +type Substitutions = M.Map VName VName + +lookupSubst :: VName -> Substitutions -> VName +lookupSubst v substs = case M.lookup v substs of + Just v' | v' /= v -> lookupSubst v' substs + _ -> v + +data Mod = ModFun TySet Scope ModParam ModExp + -- ^ A pairing of a lexical closure and a module function. + | ModMod Scope + -- ^ A non-parametric module. + deriving (Show) + +modScope :: Mod -> Scope +modScope (ModMod scope) = scope +modScope ModFun{} = mempty + +data Scope = Scope { scopeSubsts :: Substitutions + , scopeMods :: M.Map VName Mod + } + deriving (Show) + +lookupSubstInScope :: QualName VName -> Scope -> (QualName VName, Scope) +lookupSubstInScope qn@(QualName quals name) scope@(Scope substs mods) = + case quals of + [] -> (qualName $ lookupSubst name substs, scope) + q:qs -> + let q' = lookupSubst q substs + in case M.lookup q' mods of + Just (ModMod mod_scope) -> lookupSubstInScope (QualName qs name) mod_scope + _ -> (qn, scope) + +instance Sem.Semigroup Scope where + Scope ss1 mt1 <> Scope ss2 mt2 = + Scope (ss1<>ss2) (mt1<>mt2) + +instance Monoid Scope where + mempty = Scope mempty mempty + mappend = (Sem.<>) + +type TySet = S.Set VName + +data Env = Env { envScope :: Scope + , envGenerating :: Bool + , envImports :: M.Map String Scope + , envAbs :: TySet + } + +newtype TransformM a = TransformM (RWS Env (DL.DList Dec) VNameSource a) + deriving (Applicative, Functor, Monad, + MonadFreshNames, + MonadReader Env, + MonadWriter (DL.DList Dec)) + +emit :: Dec -> TransformM () +emit = tell . DL.singleton + +askScope :: TransformM Scope +askScope = asks envScope + +localScope :: (Scope -> Scope) -> TransformM a -> TransformM a +localScope f = local $ \env -> env { envScope = f $ envScope env } + +extendScope :: Scope -> TransformM a -> TransformM a +extendScope (Scope substs mods) = localScope $ \scope -> + scope { scopeSubsts = M.map (forward (scopeSubsts scope)) substs <> scopeSubsts scope + , scopeMods = mods <> scopeMods scope } + where forward old_substs v = fromMaybe v $ M.lookup v old_substs + +substituting :: Substitutions -> TransformM a -> TransformM a +substituting substs = extendScope mempty { scopeSubsts = substs } + +boundName :: VName -> TransformM VName +boundName v = do g <- asks envGenerating + if g then newName v else return v + +bindingNames :: [VName] -> TransformM Scope -> TransformM Scope +bindingNames names m = do + names' <- mapM boundName names + let substs = M.fromList (zip names names') + substituting substs $ mappend <$> m <*> pure (Scope substs mempty) + +generating :: TransformM a -> TransformM a +generating = local $ \env -> env { envGenerating = True } + +bindingImport :: String -> Scope -> TransformM a -> TransformM a +bindingImport name scope = local $ \env -> + env { envImports = M.insert name scope $ envImports env } + +bindingAbs :: TySet -> TransformM a -> TransformM a +bindingAbs abs = local $ \env -> + env { envAbs = abs <> envAbs env } + +lookupImport :: String -> TransformM Scope +lookupImport name = maybe bad return =<< asks (M.lookup name . envImports) + where bad = fail $ "Unknown import: " ++ name + +lookupMod' :: QualName VName -> Scope -> Either String Mod +lookupMod' mname scope = + let (mname', scope') = lookupSubstInScope mname scope + in maybe (Left $ bad mname') Right $ M.lookup (qualLeaf mname') $ scopeMods scope' + where bad mname' = "Unknown module: " ++ pretty mname ++ " (" ++ pretty mname' ++ ")" + +lookupMod :: QualName VName -> TransformM Mod +lookupMod mname = either fail return . lookupMod' mname =<< askScope + +runTransformM :: VNameSource -> TransformM a -> (a, VNameSource, DL.DList Dec) +runTransformM src (TransformM m) = runRWS m env src + where env = Env mempty False mempty mempty + +maybeAscript :: SrcLoc -> Maybe (SigExp, Info (M.Map VName VName)) -> ModExp + -> ModExp +maybeAscript loc (Just (mtye, substs)) me = ModAscript me mtye substs loc +maybeAscript _ Nothing me = me + +substituteInMod :: Substitutions -> Mod -> Mod +substituteInMod substs (ModMod (Scope mod_substs mod_mods)) = + -- Forward all substitutions. + ModMod $ Scope substs' $ M.map (substituteInMod substs) mod_mods + where forward v = lookupSubst v $ mod_substs <> substs + substs' = M.map forward substs +substituteInMod substs (ModFun abs (Scope mod_substs mod_mods) mparam mbody) = + ModFun abs (Scope (substs'<>mod_substs) mod_mods) mparam mbody + where forward v = lookupSubst v mod_substs + substs' = M.map forward substs + +evalModExp :: ModExp -> TransformM Mod +evalModExp (ModVar qn _) = lookupMod qn +evalModExp (ModParens e _) = evalModExp e +evalModExp (ModDecs decs _) = ModMod <$> transformDecs decs +evalModExp (ModImport _ (Info fpath) _) = ModMod <$> lookupImport fpath +evalModExp (ModAscript me _ (Info ascript_substs) _) = + substituteInMod ascript_substs <$> evalModExp me +evalModExp (ModApply f arg (Info p_substs) (Info b_substs) loc) = do + f_mod <- evalModExp f + arg_mod <- evalModExp arg + case f_mod of + ModMod _ -> + fail $ "Cannot apply non-parametric module at " ++ locStr loc + ModFun f_abs f_closure f_p f_body -> + bindingAbs (f_abs <> S.fromList (unInfo (modParamAbs f_p))) $ + extendScope f_closure $ generating $ do + outer_substs <- scopeSubsts <$> askScope + abs <- asks envAbs + let forward (k,v) = (lookupSubst k outer_substs, v) + p_substs' = M.fromList $ map forward $ M.toList p_substs + abs_substs = M.filterWithKey (const . flip S.member abs) $ + p_substs' <> + scopeSubsts f_closure <> + scopeSubsts (modScope arg_mod) + extendScope (Scope abs_substs (M.singleton (modParamName f_p) $ + substituteInMod p_substs' arg_mod)) $ do + substs <- scopeSubsts <$> askScope + x <- evalModExp f_body + return $ addSubsts abs abs_substs $ substituteInMod (b_substs <> substs) x + where addSubsts abs substs (ModFun mabs (Scope msubsts mods) mp me) = + ModFun (abs<>mabs) (Scope (substs<>msubsts) mods) mp me + addSubsts _ substs (ModMod (Scope msubsts mods)) = + ModMod $ Scope (substs<>msubsts) mods +evalModExp (ModLambda p ascript e loc) = do + scope <- askScope + abs <- asks envAbs + return $ ModFun abs scope p $ maybeAscript loc ascript e + +transformName :: VName -> TransformM VName +transformName v = lookupSubst v . scopeSubsts <$> askScope + +-- | A general-purpose substitution of names. +transformNames :: ASTMappable x => x -> TransformM x +transformNames x = do + scope <- askScope + return $ runIdentity $ astMap (substituter scope) x + where substituter scope = + ASTMapper { mapOnExp = onExp scope + , mapOnName = \v -> + return $ qualLeaf $ fst $ lookupSubstInScope (qualName v) scope + , mapOnQualName = \v -> + return $ fst $ lookupSubstInScope v scope + , mapOnType = astMap (substituter scope) + , mapOnCompType = astMap (substituter scope) + , mapOnStructType = astMap (substituter scope) + , mapOnPatternType = astMap (substituter scope) + } + onExp scope e = + -- One expression is tricky, because it interacts with scoping rules. + case e of + QualParens mn e' _ -> + case lookupMod' mn scope of + Left err -> fail err + Right mod -> + astMap (substituter $ modScope mod<>scope) e' + _ -> astMap (substituter scope) e + +transformTypeExp :: TypeExp VName -> TransformM (TypeExp VName) +transformTypeExp = transformNames + +transformStructType :: StructType -> TransformM StructType +transformStructType = transformNames + +transformExp :: Exp -> TransformM Exp +transformExp = transformNames + +transformValBind :: ValBind -> TransformM () +transformValBind (ValBind entry name tdecl (Info t) tparams params e doc loc) = do + name' <- transformName name + tdecl' <- traverse transformTypeExp tdecl + t' <- transformStructType t + e' <- transformExp e + tparams' <- traverse transformNames tparams + params' <- traverse transformNames params + emit $ ValDec $ ValBind entry name' tdecl' (Info t') tparams' params' e' doc loc + +transformTypeDecl :: TypeDecl -> TransformM TypeDecl +transformTypeDecl (TypeDecl dt (Info et)) = + TypeDecl <$> transformTypeExp dt <*> (Info <$> transformStructType et) + +transformTypeBind :: TypeBind -> TransformM () +transformTypeBind (TypeBind name tparams te doc loc) = do + name' <- transformName name + emit =<< TypeDec <$> (TypeBind name' <$> traverse transformNames tparams + <*> transformTypeDecl te <*> pure doc <*> pure loc) + +transformModBind :: ModBind -> TransformM Scope +transformModBind mb = do + let addParam p me = ModLambda p Nothing me $ srclocOf me + mod <- evalModExp $ foldr addParam + (maybeAscript (srclocOf mb) (modSignature mb) $ modExp mb) $ + modParams mb + mname <- transformName $ modName mb + return $ Scope mempty $ M.singleton mname mod + +transformDecs :: [Dec] -> TransformM Scope +transformDecs ds = + case ds of + [] -> + return mempty + LocalDec d _ : ds' -> + transformDecs $ d : ds' + ValDec fdec : ds' -> + bindingNames [valBindName fdec] $ do + transformValBind fdec + transformDecs ds' + TypeDec tb : ds' -> + bindingNames [typeAlias tb] $ do + transformTypeBind tb + transformDecs ds' + SigDec {} : ds' -> + transformDecs ds' + ModDec mb : ds' -> + bindingNames [modName mb] $ do + mod_scope <- transformModBind mb + extendScope mod_scope $ mappend <$> transformDecs ds' <*> pure mod_scope + OpenDec e _ _ : ds' -> do + scope <- modScope <$> evalModExp e + extendScope scope $ mappend <$> transformDecs ds' <*> pure scope + +transformImports :: Imports -> TransformM () +transformImports [] = return () +transformImports ((name,imp):imps) = do + let abs = S.fromList $ map qualLeaf $ M.keys $ fileAbs imp + scope <- censor (fmap maybeHideEntryPoint) $ + bindingAbs abs $ transformDecs $ progDecs $ fileProg imp + bindingAbs abs $ bindingImport name scope $ transformImports imps + where + -- Only the "main" file (last import) is allowed to have entry points. + permit_entry_points = null imps + + maybeHideEntryPoint (ValDec vdec) = + ValDec vdec { valBindEntryPoint = + valBindEntryPoint vdec && permit_entry_points } + maybeHideEntryPoint d = d + +transformProg :: MonadFreshNames m => Imports -> m [Dec] +transformProg prog = modifyNameSource $ \namesrc -> + let ((), namesrc', prog') = runTransformM namesrc $ transformImports prog + in (DL.toList prog', namesrc') diff --git a/src/Futhark/Internalise/Lambdas.hs b/src/Futhark/Internalise/Lambdas.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/Lambdas.hs @@ -0,0 +1,184 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Internalise.Lambdas + ( InternaliseLambda + , internaliseMapLambda + , internaliseStreamMapLambda + , internaliseFoldLambda + , internaliseStreamLambda + , internalisePartitionLambda + ) + where + +import Control.Monad +import Data.Loc +import qualified Data.Set as S + +import Language.Futhark as E +import Futhark.Representation.SOACS as I +import Futhark.MonadFreshNames + +import Futhark.Internalise.Monad +import Futhark.Internalise.AccurateSizes +import Futhark.Representation.SOACS.Simplify (simplifyLambda) + +-- | A function for internalising lambdas. +type InternaliseLambda = + E.Exp -> [I.Type] -> InternaliseM ([I.LParam], I.Body, [I.ExtType]) + +internaliseMapLambda :: InternaliseLambda + -> E.Exp + -> [I.SubExp] + -> InternaliseM I.Lambda +internaliseMapLambda internaliseLambda lam args = do + argtypes <- mapM I.subExpType args + let rowtypes = map I.rowType argtypes + (params, body, rettype) <- internaliseLambda lam rowtypes + (rettype', inner_shapes) <- instantiateShapes' rettype + let outer_shape = arraysSize 0 argtypes + shapefun <- makeShapeFun params body rettype' inner_shapes + bindMapShapes index0 [] inner_shapes shapefun args outer_shape + body' <- localScope (scopeOfLParams params) $ + ensureResultShape asserting + (ErrorMsg [ErrorString "not all iterations produce same shape"]) + (srclocOf lam) rettype' body + return $ I.Lambda params body' rettype' + where index0 arg = do + arg' <- letExp "arg" $ I.BasicOp $ I.SubExp arg + arg_t <- lookupType arg' + return $ I.BasicOp $ I.Index arg' $ fullSlice arg_t [I.DimFix zero] + zero = constant (0::I.Int32) + +internaliseStreamMapLambda :: InternaliseLambda + -> E.Exp + -> [I.SubExp] + -> InternaliseM I.Lambda +internaliseStreamMapLambda internaliseLambda lam args = do + chunk_size <- newVName "chunk_size" + let chunk_param = I.Param chunk_size (I.Prim int32) + outer = (`setOuterSize` I.Var chunk_size) + localScope (scopeOfLParams [chunk_param]) $ do + argtypes <- mapM I.subExpType args + (params, body, rettype) <- internaliseLambda lam $ map outer argtypes + (rettype', inner_shapes) <- instantiateShapes' rettype + let outer_shape = arraysSize 0 argtypes + shapefun <- makeShapeFun (chunk_param:params) body rettype' inner_shapes + bindMapShapes (slice0 chunk_size) [zero] inner_shapes shapefun args outer_shape + body' <- localScope (scopeOfLParams params) $ + ensureResultShape asserting + (ErrorMsg [ErrorString "not all iterations produce same shape"]) + (srclocOf lam) (map outer rettype') body + return $ I.Lambda (chunk_param:params) body' (map outer rettype') + where slice0 chunk_size arg = do + arg' <- letExp "arg" $ I.BasicOp $ I.SubExp arg + arg_t <- lookupType arg' + return $ I.BasicOp $ I.Index arg' $ + fullSlice arg_t [I.DimSlice zero (I.Var chunk_size) one] + zero = constant (0::I.Int32) + one = constant (1::I.Int32) + +makeShapeFun :: [I.LParam] -> I.Body -> [Type] -> [I.Ident] + -> InternaliseM I.Lambda +makeShapeFun params body val_rettype inner_shapes = do + -- Some of 'params' may be unique, which means that the shape slice + -- would consume its input. This is not acceptable - that input is + -- needed for the value function! Hence, for all unique parameters, + -- we create a substitute non-unique parameter, and insert a + -- copy-binding in the body of the function. + (params', copystms) <- nonuniqueParams params + shape_body <- runBodyBinder $ localScope (scopeOfLParams params') $ do + mapM_ addStm copystms + shapeBody (map I.identName inner_shapes) val_rettype body + return $ I.Lambda params' shape_body rettype + where rettype = replicate (length inner_shapes) $ I.Prim int32 + +bindMapShapes :: (SubExp -> InternaliseM I.Exp) -> [SubExp] + -> [I.Ident] -> I.Lambda -> [I.SubExp] -> SubExp + -> InternaliseM () +bindMapShapes indexArg extra_args inner_shapes sizefun args outer_shape + | null $ I.lambdaReturnType sizefun = return () + | otherwise = do + let size_args = replicate (length $ lambdaParams sizefun) Nothing + sizefun' <- simplifyLambda sizefun size_args + let sizefun_safe = + all (I.safeExp . I.stmExp) $ I.bodyStms $ I.lambdaBody sizefun' + sizefun_arg_invariant = + not $ any (`S.member` freeInBody (I.lambdaBody sizefun')) $ + map I.paramName $ lambdaParams sizefun' + if sizefun_safe && sizefun_arg_invariant + then do ses <- bodyBind $ lambdaBody sizefun' + forM_ (zip inner_shapes ses) $ \(v, se) -> + letBind_ (basicPattern [] [v]) $ I.BasicOp $ I.SubExp se + else letBind_ (basicPattern [] inner_shapes) =<< + eIf' isnonempty nonemptybranch emptybranch IfFallback + + where emptybranch = + pure $ resultBody (map (const zero) $ I.lambdaReturnType sizefun) + nonemptybranch = insertStmsM $ + resultBody <$> (eLambda sizefun . (map eSubExp extra_args++) $ map indexArg args) + + isnonempty = eNot $ eCmpOp (I.CmpEq I.int32) + (pure $ I.BasicOp $ I.SubExp outer_shape) + (pure $ I.BasicOp $ SubExp zero) + zero = constant (0::I.Int32) + +internaliseFoldLambda :: InternaliseLambda + -> E.Exp + -> [I.Type] -> [I.Type] + -> InternaliseM I.Lambda +internaliseFoldLambda internaliseLambda lam acctypes arrtypes = do + let rowtypes = map I.rowType arrtypes + (params, body, rettype) <- internaliseLambda lam $ acctypes ++ rowtypes + let rettype' = [ t `I.setArrayShape` I.arrayShape shape + | (t,shape) <- zip rettype acctypes ] + -- The result of the body must have the exact same shape as the + -- initial accumulator. We accomplish this with an assertion and + -- reshape(). + body' <- localScope (scopeOfLParams params) $ + ensureResultShape asserting + (ErrorMsg [ErrorString "shape of result does not match shape of initial value"]) + (srclocOf lam) rettype' body + return $ I.Lambda params body' rettype' + +internaliseStreamLambda :: InternaliseLambda + -> E.Exp + -> [I.Type] + -> InternaliseM ([LParam], Body) +internaliseStreamLambda internaliseLambda lam rowts = do + chunk_size <- newVName "chunk_size" + let chunk_param = I.Param chunk_size $ I.Prim int32 + chunktypes = map (`arrayOfRow` I.Var chunk_size) rowts + (params, body, _) <- localScope (scopeOfLParams [chunk_param]) $ + internaliseLambda lam chunktypes + return (chunk_param:params, body) + +-- Given @k@ lambdas, this will return a lambda that returns an +-- (k+2)-element tuple of integers. The first element is the +-- equivalence class ID in the range [0,k]. The remaining are all zero +-- except for possibly one element. +internalisePartitionLambda :: InternaliseLambda + -> Int + -> E.Exp + -> [I.SubExp] + -> InternaliseM I.Lambda +internalisePartitionLambda internaliseLambda k lam args = do + argtypes <- mapM I.subExpType args + let rowtypes = map I.rowType argtypes + (params, body, _) <- internaliseLambda lam rowtypes + body' <- localScope (scopeOfLParams params) $ + lambdaWithIncrement body + return $ I.Lambda params body' rettype + where rettype = replicate (k+2) $ I.Prim int32 + result i = map constant $ (fromIntegral i :: Int32) : + (replicate i 0 ++ [1::Int32] ++ replicate (k-i) 0) + + mkResult _ i | i >= k = return $ result i + mkResult eq_class i = do + is_i <- letSubExp "is_i" $ BasicOp $ CmpOp (CmpEq int32) eq_class (constant i) + fmap (map I.Var) . letTupExp "part_res" =<< + eIf (eSubExp is_i) (pure $ resultBody $ result i) + (resultBody <$> mkResult eq_class (i+1)) + + lambdaWithIncrement :: I.Body -> InternaliseM I.Body + lambdaWithIncrement lam_body = runBodyBinder $ do + [eq_class] <- bodyBind lam_body + resultBody <$> mkResult eq_class 0 diff --git a/src/Futhark/Internalise/Monad.hs b/src/Futhark/Internalise/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/Monad.hs @@ -0,0 +1,194 @@ +{-# LANGUAGE FlexibleInstances, TypeFamilies, GeneralizedNewtypeDeriving, MultiParamTypeClasses #-} +module Futhark.Internalise.Monad + ( InternaliseM + , runInternaliseM + , throwError + , VarSubstitutions + , InternaliseEnv (..) + , ConstParams + , Closure + , FunInfo + + , substitutingVars + , addFunction + + , lookupFunction + , lookupFunction' + + , bindFunction + + , asserting + , assertingOne + + -- * Type Handling + , InternaliseTypeM + , liftInternaliseM + , runInternaliseTypeM + , lookupDim + , withDims + , DimTable + + -- * Convenient reexports + , module Futhark.Tools + ) + where + +import Control.Monad.Except +import Control.Monad.State +import Control.Monad.Reader +import Control.Monad.Writer +import Control.Monad.RWS +import qualified Control.Monad.Fail as Fail +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem + +import Futhark.Representation.SOACS +import Futhark.MonadFreshNames +import Futhark.Tools + +type ConstParams = [(Name,VName)] + +-- | Extra parameters to pass when calling this function. This +-- corresponds to the closure of a locally defined function. +type Closure = [VName] + +type FunInfo = (Name, ConstParams, Closure, + [VName], [DeclType], + [FParam], + [(SubExp,Type)] -> Maybe [DeclExtType]) + +type FunTable = M.Map VName FunInfo + +-- | A mapping from external variable names to the corresponding +-- internalised subexpressions. +type VarSubstitutions = M.Map VName [SubExp] + +data InternaliseEnv = InternaliseEnv { + envSubsts :: VarSubstitutions + , envDoBoundsChecks :: Bool + , envSafe :: Bool + } + +data InternaliseState = InternaliseState { + stateNameSource :: VNameSource + , stateFunTable :: FunTable + } + +newtype InternaliseResult = InternaliseResult [FunDef] + deriving (Sem.Semigroup, Monoid) + +newtype InternaliseM a = InternaliseM (BinderT SOACS + (RWST + InternaliseEnv + InternaliseResult + InternaliseState + (Except String)) + a) + deriving (Functor, Applicative, Monad, + MonadReader InternaliseEnv, + MonadState InternaliseState, + MonadFreshNames, + MonadError String, + HasScope SOACS, + LocalScope SOACS) + +instance (Monoid w, Monad m) => MonadFreshNames (RWST r w InternaliseState m) where + getNameSource = gets stateNameSource + putNameSource src = modify $ \s -> s { stateNameSource = src } + +instance Fail.MonadFail InternaliseM where + fail = InternaliseM . throwError + +instance MonadBinder InternaliseM where + type Lore InternaliseM = SOACS + mkExpAttrM pat e = InternaliseM $ mkExpAttrM pat e + mkBodyM bnds res = InternaliseM $ mkBodyM bnds res + mkLetNamesM pat e = InternaliseM $ mkLetNamesM pat e + + addStms = InternaliseM . addStms + collectStms (InternaliseM m) = InternaliseM $ collectStms m + certifying cs (InternaliseM m) = InternaliseM $ certifying cs m + +runInternaliseM :: MonadFreshNames m => + Bool -> InternaliseM () + -> m (Either String [FunDef]) +runInternaliseM safe (InternaliseM m) = + modifyNameSource $ \src -> do + let onError e = (Left e, src) + onSuccess (funs,src') = (Right funs, src') + either onError onSuccess $ runExcept $ do + (_, s, InternaliseResult funs) <- runRWST (runBinderT m mempty) newEnv (newState src) + return (funs, stateNameSource s) + where newEnv = InternaliseEnv { + envSubsts = mempty + , envDoBoundsChecks = True + , envSafe = safe + } + newState src = + InternaliseState { stateNameSource = src + , stateFunTable = mempty + } + +substitutingVars :: VarSubstitutions -> InternaliseM a -> InternaliseM a +substitutingVars substs = local $ \env -> env { envSubsts = substs <> envSubsts env } + +-- | Add a function definition to the program being constructed. +addFunction :: FunDef -> InternaliseM () +addFunction = InternaliseM . lift . tell . InternaliseResult . pure + +lookupFunction' :: VName -> InternaliseM (Maybe FunInfo) +lookupFunction' fname = gets $ M.lookup fname . stateFunTable + +lookupFunction :: VName -> InternaliseM FunInfo +lookupFunction fname = maybe bad return =<< lookupFunction' fname + where bad = fail $ "Internalise.lookupFunction: Function '" ++ pretty fname ++ "' not found." + +bindFunction :: VName -> FunInfo -> InternaliseM () +bindFunction fname info = + modify $ \s -> s { stateFunTable = M.insert fname info $ stateFunTable s } + +-- | Execute the given action if 'envDoBoundsChecks' is true, otherwise +-- just return an empty list. +asserting :: InternaliseM Certificates + -> InternaliseM Certificates +asserting m = do + doBoundsChecks <- asks envDoBoundsChecks + if doBoundsChecks + then m + else return mempty + +-- | Execute the given action if 'envDoBoundsChecks' is true, otherwise +-- just return an empty list. +assertingOne :: InternaliseM VName + -> InternaliseM Certificates +assertingOne m = asserting $ Certificates . pure <$> m + +type DimTable = M.Map VName ExtSize + +newtype TypeEnv = TypeEnv { typeEnvDims :: DimTable } + +type TypeState = (Int, ConstParams) + +newtype InternaliseTypeM a = + InternaliseTypeM (ReaderT TypeEnv (StateT TypeState InternaliseM) a) + deriving (Functor, Applicative, Monad, + MonadReader TypeEnv, + MonadState TypeState, + MonadError String) + +liftInternaliseM :: InternaliseM a -> InternaliseTypeM a +liftInternaliseM = InternaliseTypeM . lift . lift + +runInternaliseTypeM :: InternaliseTypeM a + -> InternaliseM (a, ConstParams) +runInternaliseTypeM (InternaliseTypeM m) = do + let new_env = TypeEnv mempty + new_state = (0, mempty) + (x, (_, cm)) <- runStateT (runReaderT m new_env) new_state + return (x, cm) + +withDims :: DimTable -> InternaliseTypeM a -> InternaliseTypeM a +withDims dtable = local $ \env -> env { typeEnvDims = dtable <> typeEnvDims env } + +lookupDim :: VName -> InternaliseTypeM (Maybe ExtSize) +lookupDim name = M.lookup name <$> asks typeEnvDims diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -0,0 +1,604 @@ +-- | This monomorphization module converts a well-typed, polymorphic, +-- module-free Futhark program into an equivalent monomorphic program. +-- +-- This pass also does a few other simplifications to make the job of +-- subsequent passes easier. Specifically, it does the following: +-- +-- * Turn operator sections into explicit lambdas. +-- +-- * Converts identifiers of record type into record patterns (and +-- similarly for tuples). +-- +-- * Converts applications of intrinsic SOACs into SOAC AST nodes +-- (Map, Reduce, etc). +-- +-- * Elide functions that are not reachable from an entry point (this +-- is a side effect of the monomorphisation algorithm, which uses +-- the entry points as roots). +-- +-- * Turns implicit record fields into explicit record fields. +-- +-- Note that these changes are unfortunately not visible in the AST +-- representation. +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Futhark.Internalise.Monomorphise + ( transformProg + , transformDecs + , runMonoM + ) where + +import Control.Monad.RWS +import Control.Monad.State +import Data.Loc +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem +import qualified Data.Sequence as Seq +import Data.Foldable + +import Futhark.MonadFreshNames +import Language.Futhark +import Language.Futhark.Traversals +import Language.Futhark.TypeChecker.Monad (TypeBinding(..)) +import Language.Futhark.TypeChecker.Types + +-- | The monomorphization monad reads 'PolyBinding's and writes 'ValBinding's. +-- The 'TypeParam's in a 'ValBinding' can only be shape parameters. +newtype PolyBinding = PolyBinding (VName, [TypeParam], [Pattern], + Maybe (TypeExp VName), StructType, Exp, SrcLoc) + +-- | Mapping from record names to the variable names that contain the +-- fields. This is used because the monomorphiser also expands all +-- record patterns. +type RecordReplacements = M.Map VName RecordReplacement + +type RecordReplacement = M.Map Name (VName, PatternType) + +-- | Monomorphization environment mapping names of polymorphic functions to a +-- representation of their corresponding function bindings. +data Env = Env { envPolyBindings :: M.Map VName PolyBinding + , envTypeBindings :: M.Map VName TypeBinding + , envRecordReplacements :: RecordReplacements + } + +instance Sem.Semigroup Env where + Env tb1 pb1 rr1 <> Env tb2 pb2 rr2 = Env (tb1 <> tb2) (pb1 <> pb2) (rr1 <> rr2) + +instance Monoid Env where + mempty = Env mempty mempty mempty + mappend = (Sem.<>) + +localEnv :: Env -> MonoM a -> MonoM a +localEnv env = local (env <>) + +extendEnv :: VName -> PolyBinding -> MonoM a -> MonoM a +extendEnv vn binding = localEnv + mempty { envPolyBindings = M.singleton vn binding } + +withRecordReplacements :: RecordReplacements -> MonoM a -> MonoM a +withRecordReplacements rr = localEnv mempty { envRecordReplacements = rr} + +noRecordReplacements :: MonoM a -> MonoM a +noRecordReplacements = local $ \env -> env { envRecordReplacements = mempty } + +-- | The monomorphization monad. +newtype MonoM a = MonoM (RWST Env (Seq.Seq (VName, ValBind)) VNameSource + (State Lifts) a) + deriving (Functor, Applicative, Monad, + MonadReader Env, + MonadWriter (Seq.Seq (VName, ValBind)), + MonadFreshNames) + +runMonoM :: VNameSource -> MonoM a -> ((a, Seq.Seq (VName, ValBind)), VNameSource) +runMonoM src (MonoM m) = ((a, defs), src') + where (a, src', defs) = evalState (runRWST m mempty src) mempty + +lookupFun :: VName -> MonoM (Maybe PolyBinding) +lookupFun vn = do + env <- asks envPolyBindings + case M.lookup vn env of + Just valbind -> return $ Just valbind + Nothing -> return Nothing + +lookupRecordReplacement :: VName -> MonoM (Maybe RecordReplacement) +lookupRecordReplacement v = asks $ M.lookup v . envRecordReplacements + +-- | Mapping from function name and instance list to a new function name in case +-- the function has already been instantiated with those concrete types. +type Lifts = [((VName, TypeBase () ()), VName)] + +getLifts :: MonoM Lifts +getLifts = MonoM $ lift get + +modifyLifts :: (Lifts -> Lifts) -> MonoM () +modifyLifts = MonoM . lift . modify + +addLifted :: VName -> TypeBase () () -> VName -> MonoM () +addLifted fname il lifted_fname = + modifyLifts (((fname, il), lifted_fname) :) + +lookupLifted :: VName -> TypeBase () () -> MonoM (Maybe VName) +lookupLifted fname t = lookup (fname, t) <$> getLifts + +transformFName :: VName -> TypeBase () () -> MonoM VName +transformFName fname t + | baseTag fname <= maxIntrinsicTag = return fname + | otherwise = do + maybe_fname <- lookupLifted fname t + maybe_funbind <- lookupFun fname + case (maybe_fname, maybe_funbind) of + -- The function has already been monomorphized. + (Just fname', _) -> return fname' + -- An intrinsic function. + (Nothing, Nothing) -> return fname + -- A polymorphic function. + (Nothing, Just funbind) -> do + (fname', funbind') <- monomorphizeBinding funbind t + tell $ Seq.singleton (fname, funbind') + addLifted fname t fname' + return fname' + +-- | Monomorphization of expressions. +transformExp :: Exp -> MonoM Exp +transformExp e@Literal{} = return e +transformExp e@IntLit{} = return e +transformExp e@FloatLit{} = return e + +transformExp (Parens e loc) = + Parens <$> transformExp e <*> pure loc + +transformExp (QualParens qn e loc) = + QualParens qn <$> transformExp e <*> pure loc + +transformExp (TupLit es loc) = + TupLit <$> mapM transformExp es <*> pure loc + +transformExp (RecordLit fs loc) = + RecordLit <$> mapM transformField fs <*> pure loc + where transformField (RecordFieldExplicit name e loc') = + RecordFieldExplicit name <$> transformExp e <*> pure loc' + transformField (RecordFieldImplicit v t _) = + transformField $ RecordFieldExplicit (baseName v) + (Var (qualName v) (vacuousShapeAnnotations <$> t) loc) loc + +transformExp (ArrayLit es tp loc) = + ArrayLit <$> mapM transformExp es <*> pure tp <*> pure loc + +transformExp (Range e1 me incl tp loc) = do + e1' <- transformExp e1 + me' <- mapM transformExp me + incl' <- mapM transformExp incl + return $ Range e1' me' incl' tp loc + +transformExp (Var (QualName qs fname) (Info t) loc) = do + maybe_fs <- lookupRecordReplacement fname + case maybe_fs of + Just fs -> do + let toField (f, (f_v, f_t)) = + let f_v' = Var (qualName f_v) (Info $ vacuousShapeAnnotations f_t) loc + in RecordFieldExplicit f f_v' loc + return $ RecordLit (map toField $ M.toList fs) loc + Nothing -> do + fname' <- transformFName fname (toStructural t) + return $ Var (QualName qs fname') (Info t) loc + +transformExp (Ascript e tp loc) = + Ascript <$> transformExp e <*> pure tp <*> pure loc + +transformExp (LetPat tparams pat e1 e2 loc) = do + (pat', rr) <- expandRecordPattern pat + LetPat tparams pat' <$> transformExp e1 <*> + withRecordReplacements rr (transformExp e2) <*> pure loc + +transformExp (LetFun fname (tparams, params, retdecl, Info ret, body) e loc) + | any isTypeParam tparams = do + -- Retrieve the lifted monomorphic function bindings that are produced, + -- filter those that are monomorphic versions of the current let-bound + -- function and insert them at this point, and propagate the rest. + let funbind = PolyBinding (fname, tparams, params, retdecl, ret, body, loc) + pass $ do + (e', bs) <- listen $ extendEnv fname funbind $ transformExp e + let (bs_local, bs_prop) = Seq.partition ((== fname) . fst) bs + return (unfoldLetFuns (map snd $ toList bs_local) e', const bs_prop) + + | otherwise = + transformExp $ LetPat [] (Id fname (Info ft) loc) lam e loc + where lam = Lambda tparams params body Nothing (Info (mempty, ret)) loc + ft = foldFunType (map (vacuousShapeAnnotations . patternType) params) $ fromStruct ret + +transformExp (If e1 e2 e3 tp loc) = do + e1' <- transformExp e1 + e2' <- transformExp e2 + e3' <- transformExp e3 + return $ If e1' e2' e3' tp loc + +transformExp (Apply e1 e2 d tp loc) = + -- We handle on an ad-hoc basis certain polymorphic higher-order + -- intrinsics here. They can only be used in very particular ways, + -- or the compiler will fail. In practice they will only be used + -- once, in the basis library, to define normal functions. + case (e1, e2) of + (Var v _ _, TupLit [op, ne, arr] _) + | intrinsic "reduce" v -> + transformExp $ Reduce Noncommutative op ne arr loc + | intrinsic "reduce_comm" v -> + transformExp $ Reduce Commutative op ne arr loc + | intrinsic "scan" v -> + transformExp $ Scan op ne arr loc + (Var v _ _, TupLit [f, arr] _) + | intrinsic "map" v -> + transformExp $ Map f arr (removeShapeAnnotations <$> tp) loc + | intrinsic "filter" v -> + transformExp $ Filter f arr loc + (Var v _ _, TupLit [k, f, arr] _) + | intrinsic "partition" v, + Just k' <- isInt32 k -> + transformExp $ Partition (fromIntegral k') f arr loc + (Var v _ _, TupLit [op, f, arr] _) + | intrinsic "stream_red" v -> + transformExp $ Stream (RedLike InOrder Noncommutative op) f arr loc + | intrinsic "stream_red_per" v -> + transformExp $ Stream (RedLike Disorder Commutative op) f arr loc + (Var v _ _, TupLit [f, arr] _) + | intrinsic "stream_map" v -> + transformExp $ Stream (MapLike InOrder) f arr loc + | intrinsic "stream_map_per" v -> + transformExp $ Stream (MapLike Disorder) f arr loc + (Var v _ _, TupLit [dest, op, ne, buckets, img] _) + | intrinsic "gen_reduce" v -> + transformExp $ GenReduce dest op ne buckets img loc + + _ -> do + e1' <- transformExp e1 + e2' <- transformExp e2 + return $ Apply e1' e2' d tp loc + where intrinsic s (QualName _ v) = + baseTag v <= maxIntrinsicTag && baseName v == nameFromString s + + isInt32 (Literal (SignedValue (Int32Value k)) _) = Just k + isInt32 (IntLit k (Info (Prim (Signed Int32))) _) = Just $ fromInteger k + isInt32 _ = Nothing + +transformExp (Negate e loc) = + Negate <$> transformExp e <*> pure loc + +transformExp (Lambda tparams params e0 decl tp loc) = do + e0' <- transformExp e0 + return $ Lambda tparams params e0' decl tp loc + +transformExp (OpSection qn t loc) = + transformExp $ Var qn t loc + +transformExp (OpSectionLeft (QualName qs fname) (Info t) e + (Info xtype, Info ytype) (Info rettype) loc) = do + fname' <- transformFName fname (toStructural t) + e' <- transformExp e + desugarBinOpSection (QualName qs fname') (Just e') Nothing t xtype ytype rettype loc + +transformExp (OpSectionRight (QualName qs fname) (Info t) e + (Info xtype, Info ytype) (Info rettype) loc) = do + fname' <- transformFName fname (toStructural t) + e' <- transformExp e + desugarBinOpSection (QualName qs fname') Nothing (Just e') t xtype ytype rettype loc + +transformExp (ProjectSection fields (Info t) loc) = + desugarProjectSection fields t loc + +transformExp (IndexSection idxs (Info t) loc) = + desugarIndexSection idxs t loc + +transformExp (DoLoop tparams pat e1 form e3 loc) = do + e1' <- transformExp e1 + form' <- case form of + For ident e2 -> For ident <$> transformExp e2 + ForIn pat2 e2 -> ForIn pat2 <$> transformExp e2 + While e2 -> While <$> transformExp e2 + e3' <- transformExp e3 + return $ DoLoop tparams pat e1' form' e3' loc + +transformExp (BinOp (QualName qs fname) (Info t) (e1, d1) (e2, d2) tp loc) = do + fname' <- transformFName fname (toStructural t) + e1' <- transformExp e1 + e2' <- transformExp e2 + return $ BinOp (QualName qs fname') (Info t) (e1', d1) (e2', d2) tp loc + +transformExp (Project n e tp loc) = do + maybe_fs <- case e of + Var qn _ _ -> lookupRecordReplacement (qualLeaf qn) + _ -> return Nothing + case maybe_fs of + Just m | Just (v, _) <- M.lookup n m -> + return $ Var (qualName v) (vacuousShapeAnnotations <$> tp) loc + _ -> do + e' <- transformExp e + return $ Project n e' tp loc + +transformExp (LetWith id1 id2 idxs e1 body loc) = do + idxs' <- mapM transformDimIndex idxs + e1' <- transformExp e1 + body' <- transformExp body + return $ LetWith id1 id2 idxs' e1' body' loc + +transformExp (Index e0 idxs info loc) = + Index <$> transformExp e0 <*> mapM transformDimIndex idxs <*> pure info <*> pure loc + +transformExp (Update e1 idxs e2 loc) = + Update <$> transformExp e1 <*> mapM transformDimIndex idxs + <*> transformExp e2 <*> pure loc + +transformExp (RecordUpdate e1 fs e2 t loc) = + RecordUpdate <$> transformExp e1 <*> pure fs + <*> transformExp e2 <*> pure t <*> pure loc + +transformExp (Map e1 es t loc) = + Map <$> transformExp e1 <*> transformExp es <*> pure t <*> pure loc + +transformExp (Reduce comm e1 e2 e3 loc) = + Reduce comm <$> transformExp e1 <*> transformExp e2 + <*> transformExp e3 <*> pure loc + +transformExp (Scan e1 e2 e3 loc) = + Scan <$> transformExp e1 <*> transformExp e2 <*> transformExp e3 <*> pure loc + +transformExp (Filter e1 e2 loc) = + Filter <$> transformExp e1 <*> transformExp e2 <*> pure loc + +transformExp (Partition k f e0 loc) = + Partition k <$> transformExp f <*> transformExp e0 <*> pure loc + +transformExp (Stream form e1 e2 loc) = do + form' <- case form of + MapLike _ -> return form + RedLike so comm e -> RedLike so comm <$> transformExp e + Stream form' <$> transformExp e1 <*> transformExp e2 <*> pure loc + +transformExp (GenReduce e1 e2 e3 e4 e5 loc) = + GenReduce + <$> transformExp e1 -- hist + <*> transformExp e2 -- operator + <*> transformExp e3 -- neutral element + <*> transformExp e4 -- buckets + <*> transformExp e5 -- input image + <*> pure loc + +transformExp (Zip i e1 es t loc) = do + e1' <- transformExp e1 + es' <- mapM transformExp es + return $ Zip i e1' es' t loc + +transformExp (Unzip e0 tps loc) = + Unzip <$> transformExp e0 <*> pure tps <*> pure loc + +transformExp (Unsafe e1 loc) = + Unsafe <$> transformExp e1 <*> pure loc + +transformExp (Assert e1 e2 desc loc) = + Assert <$> transformExp e1 <*> transformExp e2 <*> pure desc <*> pure loc + +transformDimIndex :: DimIndexBase Info VName -> MonoM (DimIndexBase Info VName) +transformDimIndex (DimFix e) = DimFix <$> transformExp e +transformDimIndex (DimSlice me1 me2 me3) = + DimSlice <$> trans me1 <*> trans me2 <*> trans me3 + where trans = mapM transformExp + +-- | Transform an operator section into a lambda. +desugarBinOpSection :: QualName VName -> Maybe Exp -> Maybe Exp + -> PatternType -> StructType -> StructType -> PatternType -> SrcLoc -> MonoM Exp +desugarBinOpSection qn e_left e_right t xtype ytype rettype loc = do + (e1, p1) <- makeVarParam e_left $ fromStruct xtype + (e2, p2) <- makeVarParam e_right $ fromStruct ytype + let body = BinOp qn (Info t) (e1, Info xtype) (e2, Info ytype) (Info rettype) loc + rettype' = vacuousShapeAnnotations $ toStruct rettype + return $ Lambda [] (p1 ++ p2) body Nothing (Info (mempty, rettype')) loc + + where makeVarParam (Just e) _ = return (e, []) + makeVarParam Nothing argtype = do + x <- newNameFromString "x" + return (Var (qualName x) (Info argtype) noLoc, + [Id x (Info $ fromStruct argtype) noLoc]) + +desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp +desugarProjectSection fields (Arrow _ _ t1 t2) loc = do + p <- newVName "project_p" + let body = foldl project (Var (qualName p) (Info t1) noLoc) fields + return $ Lambda [] [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc + where project e field = + case typeOf e of + Record fs | Just t <- M.lookup field fs -> + Project field e (Info t) noLoc + t -> error $ "desugarOpSection: type " ++ pretty t ++ + " does not have field " ++ pretty field +desugarProjectSection _ t _ = error $ "desugarOpSection: not a function type: " ++ pretty t + +desugarIndexSection :: [DimIndex] -> PatternType -> SrcLoc -> MonoM Exp +desugarIndexSection idxs (Arrow _ _ t1 t2) loc = do + p <- newVName "index_i" + let body = Index (Var (qualName p) (Info t1) loc) idxs (Info t2') loc + return $ Lambda [] [Id p (Info t1) noLoc] body Nothing (Info (mempty, toStruct t2)) loc + where t2' = removeShapeAnnotations t2 +desugarIndexSection _ t _ = error $ "desugarIndexSection: not a function type: " ++ pretty t + +noticeDims :: TypeBase (DimDecl VName) as -> MonoM () +noticeDims = mapM_ notice . nestedDims + where notice (NamedDim v) = void $ transformFName (qualLeaf v) $ Prim $ Signed Int32 + notice _ = return () + +-- | Convert a collection of 'ValBind's to a nested sequence of let-bound, +-- monomorphic functions with the given expression at the bottom. +unfoldLetFuns :: [ValBind] -> Exp -> Exp +unfoldLetFuns [] e = e +unfoldLetFuns (ValBind _ fname _ rettype dim_params params body _ loc : rest) e = + LetFun fname (dim_params, params, Nothing, rettype, body) e' loc + where e' = unfoldLetFuns rest e + +expandRecordPattern :: Pattern -> MonoM (Pattern, RecordReplacements) +expandRecordPattern (Id v (Info (Record fs)) loc) = do + let fs' = M.toList fs + (fs_ks, fs_ts) <- fmap unzip $ forM fs' $ \(f, ft) -> + (,) <$> newVName (nameToString f) <*> pure ft + return (RecordPattern (zip (map fst fs') + (zipWith3 Id fs_ks (map Info fs_ts) $ repeat loc)) + loc, + M.singleton v $ M.fromList $ zip (map fst fs') $ zip fs_ks fs_ts) +expandRecordPattern (Id v t loc) = return (Id v t loc, mempty) +expandRecordPattern (TuplePattern pats loc) = do + (pats', rrs) <- unzip <$> mapM expandRecordPattern pats + return (TuplePattern pats' loc, mconcat rrs) +expandRecordPattern (RecordPattern fields loc) = do + let (field_names, field_pats) = unzip fields + (field_pats', rrs) <- unzip <$> mapM expandRecordPattern field_pats + return (RecordPattern (zip field_names field_pats') loc, mconcat rrs) +expandRecordPattern (PatternParens pat loc) = do + (pat', rr) <- expandRecordPattern pat + return (PatternParens pat' loc, rr) +expandRecordPattern (Wildcard t loc) = return (Wildcard t loc, mempty) +expandRecordPattern (PatternAscription pat td loc) = do + (pat', rr) <- expandRecordPattern pat + return (PatternAscription pat' td loc, rr) + +-- | Monomorphize a polymorphic function at the types given in the instance +-- list. Monomorphizes the body of the function as well. Returns the fresh name +-- of the generated monomorphic function and its 'ValBind' representation. +monomorphizeBinding :: PolyBinding -> TypeBase () () -> MonoM (VName, ValBind) +monomorphizeBinding (PolyBinding (name, tparams, params, retdecl, rettype, body, loc)) t = + noRecordReplacements $ do + t' <- removeTypeVariablesInType t + let bind_t = foldFunType (map (toStructural . patternType) params) $ + toStructural rettype + substs = typeSubsts bind_t t' + rettype' = applySubst (`M.lookup` substs) rettype + params' = map (substPattern $ applySubst (`M.lookup` substs)) params + + (params'', rrs) <- unzip <$> mapM expandRecordPattern params' + + mapM_ noticeDims $ rettype : map patternStructType params'' + + body' <- updateExpTypes (`M.lookup` substs) body + body'' <- withRecordReplacements (mconcat rrs) $ transformExp body' + name' <- if null tparams then return name else newName name + return (name', toValBinding name' params'' rettype' body'') + + where shape_params = filter (not . isTypeParam) tparams + + updateExpTypes substs = astMap $ mapper substs + mapper substs = ASTMapper { mapOnExp = astMap $ mapper substs + , mapOnName = pure + , mapOnQualName = pure + , mapOnType = pure . applySubst substs + , mapOnCompType = pure . applySubst substs + , mapOnStructType = pure . applySubst substs + , mapOnPatternType = pure . applySubst substs + } + + toValBinding name' params'' rettype' body'' = + ValBind { valBindEntryPoint = False + , valBindName = name' + , valBindRetDecl = retdecl + , valBindRetType = Info rettype' + , valBindTypeParams = shape_params + , valBindParams = params'' + , valBindBody = body'' + , valBindDoc = Nothing + , valBindLocation = loc + } + +typeSubsts :: TypeBase () () -> TypeBase () () + -> M.Map VName (TypeBase () ()) +typeSubsts (Record fields1) (Record fields2) = + mconcat $ zipWith typeSubsts + (map snd $ sortFields fields1) (map snd $ sortFields fields2) +typeSubsts (TypeVar _ _ v _) t = + M.singleton (typeLeaf v) t +typeSubsts Prim{} Prim{} = mempty +typeSubsts (Arrow _ _ t1a t1b) (Arrow _ _ t2a t2b) = + typeSubsts t1a t2a <> typeSubsts t1b t2b +typeSubsts t1@Array{} t2@Array{} + | Just t1' <- peelArray (arrayRank t1) t1, + Just t2' <- peelArray (arrayRank t1) t2 = + typeSubsts t1' t2' +typeSubsts t1 t2 = error $ unlines ["typeSubsts: mismatched types:", pretty t1, pretty t2] + +-- | Perform a given substitution on the types in a pattern. +substPattern :: (PatternType -> PatternType) -> Pattern -> Pattern +substPattern f pat = case pat of + TuplePattern pats loc -> TuplePattern (map (substPattern f) pats) loc + RecordPattern fs loc -> RecordPattern (map substField fs) loc + where substField (n, p) = (n, substPattern f p) + PatternParens p loc -> PatternParens (substPattern f p) loc + Id vn (Info tp) loc -> Id vn (Info $ f tp) loc + Wildcard (Info tp) loc -> Wildcard (Info $ f tp) loc + PatternAscription p td loc -> PatternAscription (substPattern f p) td loc + +toPolyBinding :: ValBind -> PolyBinding +toPolyBinding (ValBind _ name retdecl (Info rettype) tparams params body _ loc) = + PolyBinding (name, tparams, params, retdecl, rettype, body, loc) + +-- | Remove all type variables and type abbreviations from a value binding. +removeTypeVariables :: ValBind -> MonoM ValBind +removeTypeVariables valbind@(ValBind _ _ _ (Info rettype) _ pats body _ _) = do + subs <- asks $ M.map TypeSub . envTypeBindings + let substPatternType = fromStruct . substituteTypes subs . toStruct + mapper = ASTMapper { + mapOnExp = astMap mapper + , mapOnName = pure + , mapOnQualName = pure + , mapOnType = pure . removeShapeAnnotations . + substituteTypes subs . vacuousShapeAnnotations + , mapOnCompType = pure . fromStruct . removeShapeAnnotations . + substituteTypes subs . + vacuousShapeAnnotations . toStruct + , mapOnStructType = pure . substituteTypes subs + , mapOnPatternType = pure . substPatternType + } + + body' <- astMap mapper body + + return valbind { valBindRetType = Info $ substituteTypes subs rettype + , valBindParams = map (substPattern substPatternType) pats + , valBindBody = body' + } + +removeTypeVariablesInType :: TypeBase dim () -> MonoM (TypeBase () ()) +removeTypeVariablesInType t = do + subs <- asks $ M.map TypeSub . envTypeBindings + return $ removeShapeAnnotations $ substituteTypes subs $ vacuousShapeAnnotations t + +transformValBind :: ValBind -> MonoM Env +transformValBind valbind = do + valbind' <- toPolyBinding <$> removeTypeVariables valbind + when (valBindEntryPoint valbind) $ do + t <- removeTypeVariablesInType $ removeShapeAnnotations $ foldFunType + (map patternStructType (valBindParams valbind)) $ + unInfo $ valBindRetType valbind + (name, valbind'') <- monomorphizeBinding valbind' t + tell $ Seq.singleton (name, valbind'' { valBindEntryPoint = True}) + addLifted (valBindName valbind) t name + return mempty { envPolyBindings = M.singleton (valBindName valbind) valbind' } + +transformTypeBind :: TypeBind -> MonoM Env +transformTypeBind (TypeBind name tparams tydecl _ _) = do + subs <- asks $ M.map TypeSub . envTypeBindings + noticeDims $ unInfo $ expandedType tydecl + let tp = substituteTypes subs . unInfo $ expandedType tydecl + tbinding = TypeAbbr Lifted tparams tp -- The Lifted is arbitrary. + return mempty { envTypeBindings = M.singleton name tbinding } + +-- | Monomorphize a list of top-level declarations. A module-free input program +-- is expected, so only value declarations and type declaration are accepted. +transformDecs :: [Dec] -> MonoM () +transformDecs [] = return () +transformDecs (ValDec valbind : ds) = do + env <- transformValBind valbind + localEnv env $ transformDecs ds + +transformDecs (TypeDec typebind : ds) = do + env <- transformTypeBind typebind + localEnv env $ transformDecs ds + +transformDecs (dec : _) = + error $ "The monomorphization module expects a module-free " ++ + "input program, but received: " ++ pretty dec + +transformProg :: MonadFreshNames m => [Dec] -> m [ValBind] +transformProg decs = + fmap (toList . fmap snd . snd) $ modifyNameSource $ \namesrc -> + runMonoM namesrc $ transformDecs decs diff --git a/src/Futhark/Internalise/TypesValues.hs b/src/Futhark/Internalise/TypesValues.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Internalise/TypesValues.hs @@ -0,0 +1,154 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Futhark.Internalise.TypesValues + ( + -- * Internalising types + BoundInTypes + , boundInTypes + , internaliseReturnType + , internaliseEntryReturnType + , internaliseParamTypes + , internaliseType + , internalisePrimType + , internalisedTypeSize + + -- * Internalising values + , internalisePrimValue + ) + where + +import Control.Monad.State +import Control.Monad.Reader +import Data.List +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import Data.Monoid ((<>)) +import Data.Semigroup (Semigroup) + +import qualified Language.Futhark as E +import Futhark.Representation.SOACS as I +import Futhark.Internalise.Monad + +internaliseUniqueness :: E.Uniqueness -> I.Uniqueness +internaliseUniqueness E.Nonunique = I.Nonunique +internaliseUniqueness E.Unique = I.Unique + +-- | The names that are bound for some types, either implicitly or +-- explicitly. +newtype BoundInTypes = BoundInTypes (S.Set VName) + deriving (Semigroup, Monoid) + +-- | Determine the names bound for some types. +boundInTypes :: [E.TypeParam] -> BoundInTypes +boundInTypes = BoundInTypes . S.fromList . mapMaybe isTypeParam + where isTypeParam (E.TypeParamDim v _) = Just v + isTypeParam _ = Nothing + +internaliseParamTypes :: BoundInTypes + -> M.Map VName VName + -> [E.TypeBase (E.DimDecl VName) ()] + -> InternaliseM ([[I.TypeBase ExtShape Uniqueness]], + ConstParams) +internaliseParamTypes (BoundInTypes bound) pnames ts = + runInternaliseTypeM $ withDims (bound' <> M.map (Free . Var) pnames) $ + mapM internaliseTypeM ts + where bound' = M.fromList (zip (S.toList bound) + (map (Free . Var) $ S.toList bound)) + +internaliseReturnType :: E.TypeBase (E.DimDecl VName) () + -> InternaliseM ([I.TypeBase ExtShape Uniqueness], + ConstParams) +internaliseReturnType t = do + (ts', cm') <- internaliseEntryReturnType t + return (concat ts', cm') + +-- | As 'internaliseReturnType', but returns components of a top-level +-- tuple type piecemeal. +internaliseEntryReturnType :: E.TypeBase (E.DimDecl VName) () + -> InternaliseM ([[I.TypeBase ExtShape Uniqueness]], + ConstParams) +internaliseEntryReturnType t = do + let ts = case E.isTupleRecord t of Just tts -> tts + _ -> [t] + runInternaliseTypeM $ mapM internaliseTypeM ts + +internaliseType :: E.TypeBase () () + -> InternaliseM [I.TypeBase I.ExtShape Uniqueness] +internaliseType = + fmap fst . runInternaliseTypeM . internaliseTypeM . E.vacuousShapeAnnotations + +newId :: InternaliseTypeM Int +newId = do (i,cm) <- get + put (i + 1, cm) + return i + +internaliseDim :: E.DimDecl VName + -> InternaliseTypeM ExtSize +internaliseDim d = + case d of + E.AnyDim -> Ext <$> newId + E.ConstDim n -> return $ Free $ intConst I.Int32 $ toInteger n + E.NamedDim name -> namedDim name + where namedDim (E.QualName _ name) = do + subst <- liftInternaliseM $ asks $ M.lookup name . envSubsts + is_dim <- lookupDim name + case (is_dim, subst) of + (Just dim, _) -> return dim + (Nothing, Just [v]) -> return $ I.Free v + _ -> do -- Then it must be a constant. + let fname = nameFromString $ pretty name ++ "f" + (i,cm) <- get + case find ((==fname) . fst) cm of + Just (_, known) -> return $ I.Free $ I.Var known + Nothing -> do new <- liftInternaliseM $ newVName $ baseString name + put (i, (fname,new):cm) + return $ I.Free $ I.Var new + +internaliseTypeM :: E.StructType + -> InternaliseTypeM [I.TypeBase ExtShape Uniqueness] +internaliseTypeM orig_t = + case orig_t of + E.Prim bt -> return [I.Prim $ internalisePrimType bt] + E.TypeVar{} -> + fail "internaliseTypeM: cannot handle type variable." + E.Record ets -> + concat <$> mapM (internaliseTypeM . snd) (E.sortFields ets) + E.Array et shape u -> do + dims <- internaliseShape shape + ets <- internaliseElemType et + return [I.arrayOf et' (Shape dims) $ internaliseUniqueness u | et' <- ets ] + E.Arrow{} -> fail $ "internaliseTypeM: cannot handle function type: " ++ pretty orig_t + + where internaliseElemType E.ArrayPolyElem{} = + fail "internaliseElemType: cannot handle type variable." + internaliseElemType (E.ArrayPrimElem bt _) = + return [I.Prim $ internalisePrimType bt] + internaliseElemType (E.ArrayRecordElem elemts) = + concat <$> mapM (internaliseRecordElem . snd) (E.sortFields elemts) + + internaliseRecordElem (E.RecordArrayElem et) = + internaliseElemType et + internaliseRecordElem (E.RecordArrayArrayElem et shape u) = + internaliseTypeM $ E.Array et shape u + + internaliseShape = mapM internaliseDim . E.shapeDims + +-- | How many core language values are needed to represent one source +-- language value of the given type? +internalisedTypeSize :: E.TypeBase dim () -> InternaliseM Int +internalisedTypeSize = fmap length . internaliseType . E.removeShapeAnnotations + +-- | Convert an external primitive to an internal primitive. +internalisePrimType :: E.PrimType -> I.PrimType +internalisePrimType (E.Signed t) = I.IntType t +internalisePrimType (E.Unsigned t) = I.IntType t +internalisePrimType (E.FloatType t) = I.FloatType t +internalisePrimType E.Bool = I.Bool + +-- | Convert an external primitive value to an internal primitive value. +internalisePrimValue :: E.PrimValue -> I.PrimValue +internalisePrimValue (E.SignedValue v) = I.IntValue v +internalisePrimValue (E.UnsignedValue v) = I.IntValue v +internalisePrimValue (E.FloatValue v) = I.FloatValue v +internalisePrimValue (E.BoolValue b) = I.BoolValue b diff --git a/src/Futhark/MonadFreshNames.hs b/src/Futhark/MonadFreshNames.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/MonadFreshNames.hs @@ -0,0 +1,166 @@ +{-# LANGUAGE FlexibleInstances, UndecidableInstances #-} +-- | This module provides a monadic facility similar (and built on top +-- of) "Futhark.FreshNames". The removes the need for a (small) amount of +-- boilerplate, at the cost of using some GHC extensions. The idea is +-- that if your compiler pass runs in a monad that is an instance of +-- 'MonadFreshNames', you can automatically use the name generation +-- functions exported by this module. +module Futhark.MonadFreshNames + ( MonadFreshNames (..) + , modifyNameSource + , newName + , newNameFromString + , newID + , newIDFromString + , newVName + , newVName' + , newIdent + , newIdent' + , newIdents + , newParam + , newParam' + , module Futhark.FreshNames + ) where + +import Control.Monad.Except +import qualified Control.Monad.State.Lazy +import qualified Control.Monad.State.Strict +import qualified Control.Monad.Writer.Lazy +import qualified Control.Monad.Writer.Strict +import qualified Control.Monad.RWS.Lazy +import qualified Control.Monad.RWS.Strict +import qualified Control.Monad.Trans.Maybe +import Control.Monad.Reader + +import Futhark.Representation.AST.Syntax +import qualified Futhark.FreshNames as FreshNames +import Futhark.FreshNames hiding (newName, newVName) + +-- | A monad that stores a name source. The following is a good +-- instance for a monad in which the only state is a @NameSource vn@: +-- +-- @ +-- instance MonadFreshNames vn MyMonad where +-- getNameSource = get +-- putNameSource = put +-- @ +class (Applicative m, Monad m) => MonadFreshNames m where + getNameSource :: m VNameSource + putNameSource :: VNameSource -> m () + +instance (Applicative im, Monad im) => MonadFreshNames (Control.Monad.State.Lazy.StateT VNameSource im) where + getNameSource = Control.Monad.State.Lazy.get + putNameSource = Control.Monad.State.Lazy.put + +instance (Applicative im, Monad im) => MonadFreshNames (Control.Monad.State.Strict.StateT VNameSource im) where + getNameSource = Control.Monad.State.Strict.get + putNameSource = Control.Monad.State.Strict.put + +instance (Applicative im, Monad im, Monoid w) => + MonadFreshNames (Control.Monad.RWS.Lazy.RWST r w VNameSource im) where + getNameSource = Control.Monad.RWS.Lazy.get + putNameSource = Control.Monad.RWS.Lazy.put + +instance (Applicative im, Monad im, Monoid w) => + MonadFreshNames (Control.Monad.RWS.Strict.RWST r w VNameSource im) where + getNameSource = Control.Monad.RWS.Strict.get + putNameSource = Control.Monad.RWS.Strict.put + +-- | Run a computation needing a fresh name source and returning a new +-- one, using 'getNameSource' and 'putNameSource' before and after the +-- computation. +modifyNameSource :: MonadFreshNames m => (VNameSource -> (a, VNameSource)) -> m a +modifyNameSource m = do src <- getNameSource + let (x,src') = m src + putNameSource src' + return x + +-- | Produce a fresh name, using the given name as a template. +newName :: MonadFreshNames m => VName -> m VName +newName = modifyNameSource . flip FreshNames.newName + +-- | As @newName@, but takes a 'String' for the name template. +newNameFromString :: MonadFreshNames m => String -> m VName +newNameFromString s = newName $ VName (nameFromString s) 0 + +-- | Produce a fresh 'ID', using the given base name as a template. +newID :: MonadFreshNames m => Name -> m VName +newID s = newName $ VName s 0 + +-- | As 'newID', but takes a 'String' for the name template. +newIDFromString :: MonadFreshNames m => String -> m VName +newIDFromString = newID . nameFromString + +-- | Produce a fresh 'VName', using the given base name as a template. +newVName :: MonadFreshNames m => String -> m VName +newVName = newID . nameFromString + +-- | Produce a fresh 'VName', using the given name as a template, but +-- possibly appending something more.. +newVName' :: MonadFreshNames m => (String -> String) -> String -> m VName +newVName' f = newID . nameFromString . f + +-- | Produce a fresh 'Ident', using the given name as a template. +newIdent :: MonadFreshNames m => + String -> Type -> m Ident +newIdent s t = do + s' <- newID $ nameFromString s + return $ Ident s' t + +-- | Produce a fresh 'Ident', using the given 'Ident' as a template, +-- but possibly modifying the name. +newIdent' :: MonadFreshNames m => + (String -> String) + -> Ident -> m Ident +newIdent' f ident = + newIdent (f $ nameToString $ baseName $ identName ident) + (identType ident) + +-- | Produce several 'Ident's, using the given name as a template, +-- based on a list of types. +newIdents :: MonadFreshNames m => + String -> [Type] -> m [Ident] +newIdents = mapM . newIdent + +-- | Produce a fresh 'Param', using the given name as a template. +newParam :: MonadFreshNames m => + String -> attr -> m (Param attr) +newParam s t = do + s' <- newID $ nameFromString s + return $ Param s' t + +-- | Produce a fresh 'Param', using the given 'Param' as a template, +-- but possibly modifying the name. +newParam' :: MonadFreshNames m => + (String -> String) + -> Param attr -> m (Param attr) +newParam' f param = + newParam (f $ nameToString $ baseName $ paramName param) + (paramAttr param) + +-- Utility instance defintions for MTL classes. This requires +-- UndecidableInstances, but saves on typing elsewhere. + +instance MonadFreshNames m => MonadFreshNames (ReaderT s m) where + getNameSource = lift getNameSource + putNameSource = lift . putNameSource + +instance (MonadFreshNames m, Monoid s) => + MonadFreshNames (Control.Monad.Writer.Lazy.WriterT s m) where + getNameSource = lift getNameSource + putNameSource = lift . putNameSource + +instance (MonadFreshNames m, Monoid s) => + MonadFreshNames (Control.Monad.Writer.Strict.WriterT s m) where + getNameSource = lift getNameSource + putNameSource = lift . putNameSource + +instance MonadFreshNames m => + MonadFreshNames (Control.Monad.Trans.Maybe.MaybeT m) where + getNameSource = lift getNameSource + putNameSource = lift . putNameSource + +instance MonadFreshNames m => + MonadFreshNames (ExceptT e m) where + getNameSource = lift getNameSource + putNameSource = lift . putNameSource diff --git a/src/Futhark/Optimise/CSE.hs b/src/Futhark/Optimise/CSE.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/CSE.hs @@ -0,0 +1,207 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +-- | This module implements common-subexpression elimination. This +-- module does not actually remove the duplicate, but only replaces +-- one with a diference to the other. E.g: +-- +-- @ +-- let a = x + y +-- let b = x + y +-- @ +-- +-- becomes: +-- +-- @ +-- let a = x + y +-- let b = a +-- @ +-- +-- After which copy propagation in the simplifier will actually remove +-- the definition of @b@. +-- +-- Our CSE is still rather stupid. No normalisation is performed, so +-- the expressions @x+y@ and @y+x@ will be considered distinct. +-- Furthermore, no expression with its own binding will be considered +-- equal to any other, since the variable names will be distinct. +-- This affects SOACs in particular. +module Futhark.Optimise.CSE + ( performCSE + , CSEInOp + ) + where + +import Control.Monad.Reader +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import Data.Semigroup ((<>)) + +import Futhark.Analysis.Alias +import Futhark.Representation.AST +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.Aliases + (removeFunDefAliases, Aliases, consumedInStms) +import qualified Futhark.Representation.Kernels.Kernel as Kernel +import qualified Futhark.Representation.Kernels.KernelExp as KernelExp +import qualified Futhark.Representation.SOACS.SOAC as SOAC +import qualified Futhark.Representation.ExplicitMemory as ExplicitMemory +import Futhark.Transform.Substitute +import Futhark.Pass + +-- | Perform CSE on every functioon in a program. +performCSE :: (Attributes lore, CanBeAliased (Op lore), + CSEInOp (OpWithAliases (Op lore))) => + Bool -> Pass lore lore +performCSE cse_arrays = + Pass "CSE" "Combine common subexpressions." $ + intraproceduralTransformation $ + return . removeFunDefAliases . cseInFunDef cse_arrays . analyseFun + +cseInFunDef :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => + Bool -> FunDef lore -> FunDef lore +cseInFunDef cse_arrays fundec = + fundec { funDefBody = + runReader (cseInBody $ funDefBody fundec) $ newCSEState cse_arrays + } + +type CSEM lore = Reader (CSEState lore) + +cseInBody :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => + Body lore -> CSEM lore (Body lore) +cseInBody (Body bodyattr bnds res) = + cseInStms (consumedInStms bnds res) (stmsToList bnds) $ do + CSEState (_, nsubsts) _ <- ask + return $ Body bodyattr mempty $ substituteNames nsubsts res + +cseInLambda :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => + Lambda lore -> CSEM lore (Lambda lore) +cseInLambda lam = do + body' <- cseInBody $ lambdaBody lam + return lam { lambdaBody = body' } + +cseInStms :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => + Names -> [Stm lore] + -> CSEM lore (Body lore) + -> CSEM lore (Body lore) +cseInStms _ [] m = m +cseInStms consumed (bnd:bnds) m = + cseInStm consumed bnd $ \bnd' -> do + Body bodyattr bnds' es <- cseInStms consumed bnds m + bnd'' <- mapM nestedCSE bnd' + return $ Body bodyattr (stmsFromList bnd''<>bnds') es + where nestedCSE bnd' = do + e <- mapExpM cse $ stmExp bnd' + return bnd' { stmExp = e } + cse = identityMapper { mapOnBody = const cseInBody + , mapOnOp = cseInOp + } + +cseInStm :: Attributes lore => + Names -> Stm lore + -> ([Stm lore] -> CSEM lore a) + -> CSEM lore a +cseInStm consumed (Let pat (StmAux cs eattr) e) m = do + CSEState (esubsts, nsubsts) cse_arrays <- ask + let e' = substituteNames nsubsts e + pat' = substituteNames nsubsts pat + if any (bad cse_arrays) $ patternValueElements pat then + m [Let pat' (StmAux cs eattr) e'] + else + case M.lookup (eattr, e') esubsts of + Just subpat -> + local (addNameSubst pat' subpat) $ do + let lets = + [ Let (Pattern [] [patElem']) (StmAux cs eattr) $ + BasicOp $ SubExp $ Var $ patElemName patElem + | (name,patElem) <- zip (patternNames pat') $ patternElements subpat , + let patElem' = patElem { patElemName = name } + ] + m lets + _ -> local (addExpSubst pat' eattr e') $ + m [Let pat' (StmAux cs eattr) e'] + + where bad cse_arrays pe + | Mem{} <- patElemType pe = True + | Array{} <- patElemType pe, not cse_arrays = True + | patElemName pe `S.member` consumed = True + | otherwise = False + +type ExpressionSubstitutions lore = M.Map + (ExpAttr lore, Exp lore) + (Pattern lore) +type NameSubstitutions = M.Map VName VName + +data CSEState lore = CSEState + { _cseSubstitutions :: (ExpressionSubstitutions lore, NameSubstitutions) + , _cseArrays :: Bool + } + +newCSEState :: Bool -> CSEState lore +newCSEState = CSEState (M.empty, M.empty) + +mkSubsts :: PatternT attr -> PatternT attr -> M.Map VName VName +mkSubsts pat vs = M.fromList $ zip (patternNames pat) (patternNames vs) + +addNameSubst :: PatternT attr -> PatternT attr -> CSEState lore -> CSEState lore +addNameSubst pat subpat (CSEState (esubsts, nsubsts) cse_arrays) = + CSEState (esubsts, mkSubsts pat subpat `M.union` nsubsts) cse_arrays + +addExpSubst :: Attributes lore => + Pattern lore -> ExpAttr lore -> Exp lore + -> CSEState lore + -> CSEState lore +addExpSubst pat eattr e (CSEState (esubsts, nsubsts) cse_arrays) = + CSEState (M.insert (eattr,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 + +instance CSEInOp () where + cseInOp () = return () + +subCSE :: CSEM lore r -> CSEM otherlore r +subCSE m = do + CSEState _ cse_arrays <- ask + return $ runReader m $ newCSEState cse_arrays + +instance (Attributes lore, Aliased lore, CSEInOp (Op lore)) => CSEInOp (Kernel.Kernel lore) where + cseInOp = subCSE . + Kernel.mapKernelM + (Kernel.KernelMapper return cseInLambda cseInBody + return return cseInKernelBody) + +cseInKernelBody :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => + Kernel.KernelBody lore -> CSEM lore (Kernel.KernelBody lore) +cseInKernelBody (Kernel.KernelBody bodyattr bnds res) = do + Body _ bnds' _ <- cseInBody $ Body bodyattr bnds [] + return $ Kernel.KernelBody bodyattr bnds' res + +instance (Attributes lore, Aliased lore, CSEInOp (Op lore)) => CSEInOp (KernelExp.KernelExp lore) where + cseInOp (KernelExp.Combine cspace ts active body) = + subCSE $ KernelExp.Combine cspace ts active <$> cseInBody body + cseInOp (KernelExp.GroupReduce w lam input) = + subCSE $ KernelExp.GroupReduce w <$> cseInLambda lam <*> pure input + cseInOp (KernelExp.GroupStream w max_chunk lam nes arrs) = + subCSE $ KernelExp.GroupStream w max_chunk <$> cseInGroupStreamLambda lam <*> pure nes <*> pure arrs + cseInOp op = return op + +cseInGroupStreamLambda :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => + KernelExp.GroupStreamLambda lore + -> CSEM lore (KernelExp.GroupStreamLambda lore) +cseInGroupStreamLambda lam = do + body' <- cseInBody $ KernelExp.groupStreamLambdaBody lam + return lam { KernelExp.groupStreamLambdaBody = body' } + + +instance CSEInOp op => CSEInOp (ExplicitMemory.MemOp op) where + cseInOp o@ExplicitMemory.Alloc{} = return o + cseInOp (ExplicitMemory.Inner k) = ExplicitMemory.Inner <$> subCSE (cseInOp k) + +instance (Attributes lore, + CanBeAliased (Op lore), + CSEInOp (OpWithAliases (Op lore))) => + CSEInOp (SOAC.SOAC (Aliases lore)) where + cseInOp = subCSE . SOAC.mapSOACM (SOAC.SOACMapper return cseInLambda return) diff --git a/src/Futhark/Optimise/DoubleBuffer.hs b/src/Futhark/Optimise/DoubleBuffer.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/DoubleBuffer.hs @@ -0,0 +1,284 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | The simplification engine is only willing to hoist allocations +-- out of loops if the memory block resulting from the allocation is +-- dead at the end of the loop. If it is not, we may cause data +-- hazards. +-- +-- This module rewrites loops with memory block merge parameters such +-- that each memory block is copied at the end of the iteration, thus +-- ensuring that any allocation inside the loop is dead at the end of +-- the loop. This is only possible for allocations whose size is +-- loop-invariant, although the initial size may differ from the size +-- produced by the loop result. +-- +-- Additionally, inside parallel kernels we also copy the initial +-- 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 + ( doubleBuffer ) + where + +import Control.Monad.State +import Control.Monad.Writer +import Control.Monad.Reader +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import Data.List + +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + hiding (Prog, Body, Stm, Pattern, PatElem, + BasicOp, Exp, Lambda, FunDef, FParam, LParam, RetType) +import Futhark.Pass + +doubleBuffer :: Pass ExplicitMemory ExplicitMemory +doubleBuffer = + Pass { passName = "Double buffer" + , passDescription = "Perform double buffering for merge parameters of sequential loops." + , passFunction = intraproceduralTransformation optimiseFunDef + } + +-- This pass is written in a slightly weird way because we want to +-- apply essentially the same transformation both outside and inside +-- kernel bodies, which are different (but similar) representations. +-- Thus, the environment is parametrised by the lore and contains the +-- function used to transform 'Op's for the lore. + +optimiseFunDef :: FunDef ExplicitMemory -> PassM (FunDef ExplicitMemory) +optimiseFunDef fundec = modifyNameSource $ \src -> + let m = runDoubleBufferM $ inScopeOf fundec $ optimiseBody $ funDefBody fundec + (body', src') = runState (runReaderT m env) src + in (fundec { funDefBody = body' }, src') + where env = Env mempty optimiseKernelOp doNotTouchLoop + + optimiseKernelOp (Inner k) = do + scope <- castScope <$> askScope + modifyNameSource $ + runState (runReaderT (runDoubleBufferM $ Inner <$> optimiseKernel k) $ + Env scope optimiseInKernelOp optimiseLoop) + where optimiseKernel = + mapKernelM identityKernelMapper + { mapOnKernelBody = optimiseBody + , mapOnKernelKernelBody = optimiseKernelBody + , mapOnKernelLambda = optimiseLambda + } + optimiseKernelOp op = return op + + optimiseInKernelOp (Inner (GroupStream w maxchunk lam accs arrs)) = do + lam' <- optimiseGroupStreamLambda lam + return $ Inner $ GroupStream w maxchunk lam' accs arrs + optimiseInKernelOp op = return op + + doNotTouchLoop ctx val body = return (mempty, ctx, val, body) + +data Env lore = Env { envScope :: Scope lore + , envOptimiseOp :: Op lore -> DoubleBufferM lore (Op lore) + , envOptimiseLoop :: OptimiseLoop lore + } + +newtype DoubleBufferM lore a = + DoubleBufferM { runDoubleBufferM :: ReaderT (Env lore) (State VNameSource) a } + deriving (Functor, Applicative, Monad, MonadReader (Env lore), MonadFreshNames) + +instance Annotations lore => HasScope lore (DoubleBufferM lore) where + askScope = asks envScope + +instance Annotations lore => LocalScope lore (DoubleBufferM lore) where + localScope scope = local $ \env -> env { envScope = envScope env <> scope } + +-- | Bunch up all the constraints for less typing. +type LoreConstraints lore inner = + (ExpAttr lore ~ (), BodyAttr lore ~ (), + ExplicitMemorish lore, Op lore ~ MemOp inner) + +optimiseBody :: LoreConstraints lore inner => + Body lore -> DoubleBufferM lore (Body lore) +optimiseBody body = do + bnds' <- optimiseStms $ stmsToList $ bodyStms body + return $ body { bodyStms = stmsFromList bnds' } + +optimiseStms :: LoreConstraints lore inner => + [Stm lore] -> DoubleBufferM lore [Stm lore] +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 inner. + LoreConstraints lore inner => + Stm lore -> DoubleBufferM lore [Stm lore] +optimiseStm (Let pat aux (DoLoop ctx val form body)) = do + body' <- localScope (scopeOf form <> scopeOfFParams (map fst $ ctx++val)) $ + optimiseBody body + opt_loop <- asks envOptimiseLoop + (bnds, ctx', val', body'') <- opt_loop ctx val body' + return $ bnds ++ [Let pat aux $ DoLoop ctx' val' form body''] +optimiseStm (Let pat aux e) = + pure . Let pat aux <$> mapExpM optimise e + where optimise = identityMapper { mapOnBody = \_ x -> + -- This type annotation is + -- necessary to prevent the GHC + -- 8.4 type checker from going + -- nuts. + (optimiseBody x :: DoubleBufferM lore (Body lore)) + , mapOnOp = optimiseOp + } + +optimiseOp :: Op lore -> DoubleBufferM lore (Op lore) +optimiseOp op = do f <- asks envOptimiseOp + f op + +optimiseKernelBody :: KernelBody InKernel + -> DoubleBufferM InKernel (KernelBody InKernel) +optimiseKernelBody kbody = do + stms' <- optimiseStms $ stmsToList $ kernelBodyStms kbody + return $ kbody { kernelBodyStms = stmsFromList stms' } + +optimiseLambda :: Lambda InKernel -> DoubleBufferM InKernel (Lambda InKernel) +optimiseLambda lam = do + body <- localScope (castScope $ scopeOf lam) $ optimiseBody $ lambdaBody lam + return lam { lambdaBody = body } + +optimiseGroupStreamLambda :: GroupStreamLambda InKernel + -> DoubleBufferM InKernel (GroupStreamLambda InKernel) +optimiseGroupStreamLambda lam = do + body <- localScope (scopeOf lam) $ + optimiseBody $ groupStreamLambdaBody lam + return lam { groupStreamLambdaBody = body } + +type OptimiseLoop lore = + [(FParam lore, SubExp)] -> [(FParam lore, SubExp)] -> Body lore + -> DoubleBufferM lore ([Stm lore], + [(FParam lore, SubExp)], + [(FParam lore, SubExp)], + Body lore) + +optimiseLoop :: LoreConstraints lore inner => OptimiseLoop lore +optimiseLoop ctx val body = do + -- We start out by figuring out which of the merge variables should + -- be double-buffered. + buffered <- doubleBufferMergeParams + (zip (map fst ctx) (bodyResult body)) (map fst merge) + (boundInBody body) + -- Then create the allocations of the buffers and copies of the + -- initial values. + (merge', allocs) <- allocStms merge buffered + -- Modify the loop body to copy buffered result arrays. + let body' = doubleBufferResult (map fst merge) buffered body + (ctx', val') = splitAt (length ctx) merge' + -- Modify the initial merge p + return (allocs, ctx', val', body') + where merge = ctx ++ val + +-- | The booleans indicate whether we should also play with the +-- initial merge values. +data DoubleBuffer lore = BufferAlloc VName SubExp Space Bool + | BufferCopy VName IxFun VName Bool + -- ^ First name is the memory block to copy to, + -- second is the name of the array copy. + | NoBuffer + deriving (Show) + +doubleBufferMergeParams :: (ExplicitMemorish lore, MonadFreshNames m) => + [(FParam lore,SubExp)] + -> [FParam lore] -> Names + -> m [DoubleBuffer lore] +doubleBufferMergeParams ctx_and_res val_params bound_in_loop = + evalStateT (mapM buffer val_params) M.empty + where loopVariant v = v `S.member` bound_in_loop || + v `elem` map (paramName . fst) ctx_and_res + + loopInvariantSize (Constant v) = + Just (Constant v, True) + loopInvariantSize (Var v) = + case find ((==v) . paramName . fst) ctx_and_res of + Just (_, Constant val) -> + Just (Constant val, False) + Just (_, Var v') | not $ loopVariant v' -> + Just (Var v', False) + Just _ -> + Nothing + Nothing -> + Just (Var v, True) + + buffer fparam = case paramType fparam of + Mem size space + | Just (size', b) <- loopInvariantSize size -> do + -- Let us double buffer this! + bufname <- lift $ newVName "double_buffer_mem" + modify $ M.insert (paramName fparam) (bufname, b) + return $ BufferAlloc bufname size' space b + Array {} + | MemArray _ _ _ (ArrayIn mem ixfun) <- paramAttr fparam -> do + buffered <- gets $ M.lookup mem + case buffered of + Just (bufname, b) -> do + copyname <- lift $ newVName "double_buffer_array" + return $ BufferCopy bufname ixfun copyname b + Nothing -> + return NoBuffer + _ -> return NoBuffer + +allocStms :: LoreConstraints lore inner => + [(FParam lore,SubExp)] -> [DoubleBuffer lore] + -> DoubleBufferM lore ([(FParam lore, SubExp)], [Stm lore]) +allocStms merge = runWriterT . zipWithM allocation merge + where allocation m@(Param pname _, _) (BufferAlloc name size space b) = do + tell [Let (Pattern [] [PatElem name $ MemMem size space]) (defAux ()) $ + Op $ Alloc size space] + if b then return (Param pname $ MemMem size space, Var name) + else return m + allocation (f, Var v) (BufferCopy mem _ _ b) | b = do + v_copy <- lift $ newVName $ baseString v ++ "_double_buffer_copy" + (_v_mem, v_ixfun) <- lift $ lookupArraySummary v + let bt = elemType $ paramType f + shape = arrayShape $ paramType f + bound = MemArray bt shape NoUniqueness $ ArrayIn mem v_ixfun + tell [Let (Pattern [] [PatElem v_copy bound]) (defAux ()) $ + BasicOp $ Copy v] + return (f, Var v_copy) + allocation (f, se) _ = + return (f, se) + +doubleBufferResult :: (ExplicitMemorish lore, + ExpAttr lore ~ (), BodyAttr lore ~ ()) => + [FParam lore] -> [DoubleBuffer lore] + -> Body lore -> Body lore +doubleBufferResult valparams buffered (Body () bnds res) = + let (ctx_res, val_res) = splitAt (length res - length valparams) res + (copybnds,val_res') = + unzip $ zipWith3 buffer valparams buffered val_res + in Body () (bnds<>stmsFromList (catMaybes copybnds)) $ ctx_res ++ val_res' + where buffer _ (BufferAlloc bufname _ _ _) _ = + (Nothing, Var bufname) + + buffer fparam (BufferCopy bufname ixfun copyname _) (Var v) = + -- To construct the copy we will need to figure out its type + -- based on the type of the function parameter. + let t = resultType $ paramType fparam + summary = MemArray (elemType t) (arrayShape t) NoUniqueness $ ArrayIn bufname ixfun + copybnd = Let (Pattern [] [PatElem copyname summary]) (defAux ()) $ + BasicOp $ Copy v + in (Just copybnd, Var copyname) + + buffer _ _ se = + (Nothing, se) + + parammap = M.fromList $ zip (map paramName valparams) res + + resultType t = t `setArrayDims` map substitute (arrayDims t) + + substitute (Var v) + | Just replacement <- M.lookup v parammap = replacement + substitute se = + se diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Fusion.hs @@ -0,0 +1,962 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ViewPatterns #-} +-- | Perform horizontal and vertical fusion of SOACs. +module Futhark.Optimise.Fusion ( fuseSOACs ) + where + +import Control.Monad.State +import Control.Monad.Reader +import Control.Monad.Except +import qualified Data.Semigroup as Sem +import Data.Maybe +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Data.List as L + +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.SOACS hiding (SOAC(..)) +import qualified Futhark.Representation.Aliases as Aliases +import qualified Futhark.Representation.SOACS as Futhark +import Futhark.MonadFreshNames +import Futhark.Representation.SOACS.Simplify +import Futhark.Optimise.Fusion.LoopKernel +import Futhark.Construct +import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC +import qualified Futhark.Analysis.Alias as Alias +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import Futhark.Pass + +data VarEntry = IsArray VName (NameInfo SOACS) Names SOAC.Input + | IsNotArray VName (NameInfo SOACS) + +varEntryType :: VarEntry -> NameInfo SOACS +varEntryType (IsArray _ attr _ _) = + attr +varEntryType (IsNotArray _ attr) = + attr + +varEntryAliases :: VarEntry -> Names +varEntryAliases (IsArray _ _ x _) = x +varEntryAliases _ = mempty + +data FusionGEnv = FusionGEnv { + soacs :: M.Map VName [VName] + -- ^ Mapping from variable name to its entire family. + , varsInScope:: M.Map VName VarEntry + , fusedRes :: FusedRes + } + +lookupArr :: VName -> FusionGEnv -> Maybe SOAC.Input +lookupArr v env = asArray =<< M.lookup v (varsInScope env) + where asArray (IsArray _ _ _ input) = Just input + asArray IsNotArray{} = Nothing + +newtype Error = Error String + +instance Show Error where + show (Error msg) = "Fusion error:\n" ++ msg + +newtype FusionGM a = FusionGM (ExceptT Error (StateT VNameSource (Reader FusionGEnv)) a) + deriving (Monad, Applicative, Functor, + MonadError Error, + MonadState VNameSource, + MonadReader FusionGEnv) + +instance MonadFreshNames FusionGM where + getNameSource = get + putNameSource = put + +instance HasScope SOACS FusionGM where + askScope = toScope <$> asks varsInScope + where toScope = M.map varEntryType + +------------------------------------------------------------------------ +--- Monadic Helpers: bind/new/runFusionGatherM, etc +------------------------------------------------------------------------ + +-- | Binds an array name to the set of used-array vars +bindVar :: FusionGEnv -> (Ident, Names) -> FusionGEnv +bindVar env (Ident name t, aliases) = + env { varsInScope = M.insert name entry $ varsInScope env } + where entry = case t of + Array {} -> IsArray name (LetInfo t) aliases' $ SOAC.identInput $ Ident name t + _ -> IsNotArray name $ LetInfo t + expand = maybe mempty varEntryAliases . flip M.lookup (varsInScope env) + aliases' = aliases <> mconcat (map expand $ S.toList aliases) + +bindVars :: FusionGEnv -> [(Ident, Names)] -> FusionGEnv +bindVars = foldl bindVar + +binding :: [(Ident, Names)] -> FusionGM a -> FusionGM a +binding vs = local (`bindVars` vs) + +gatherStmPattern :: Pattern -> Exp -> FusionGM FusedRes -> FusionGM FusedRes +gatherStmPattern pat e = binding $ zip idents aliases + where idents = patternIdents pat + aliases = replicate (length (patternContextNames pat)) mempty ++ + expAliases (Alias.analyseExp e) + +bindingPat :: Pattern -> FusionGM a -> FusionGM a +bindingPat = binding . (`zip` repeat mempty) . patternIdents + +bindingParams :: Typed t => [Param t] -> FusionGM a -> FusionGM a +bindingParams = binding . (`zip` repeat mempty) . map paramIdent + +-- | Binds an array name to the set of soac-produced vars +bindingFamilyVar :: [VName] -> FusionGEnv -> Ident -> FusionGEnv +bindingFamilyVar faml env (Ident nm t) = + env { soacs = M.insert nm faml $ soacs env + , varsInScope = M.insert nm (IsArray nm (LetInfo t) mempty $ + SOAC.identInput $ Ident nm t) $ + varsInScope env + } + +varAliases :: VName -> FusionGM Names +varAliases v = asks $ S.insert v . maybe mempty varEntryAliases . + M.lookup v . varsInScope + +varsAliases :: Names -> FusionGM Names +varsAliases = fmap mconcat . mapM varAliases . S.toList + +checkForUpdates :: FusedRes -> Exp -> FusionGM FusedRes +checkForUpdates res (BasicOp (Update src is _)) = do + res' <- foldM addVarToInfusible res $ + src : S.toList (mconcat $ map freeIn is) + aliases <- varAliases src + let inspectKer k = k { inplace = aliases <> inplace k } + return res' { kernels = M.map inspectKer $ kernels res' } +checkForUpdates res _ = return res + +-- | Updates the environment: (i) the @soacs@ (map) by binding each pattern +-- element identifier to all pattern elements (identifiers) and (ii) the +-- variables in scope (map) by inserting each (pattern-array) name. +-- Finally, if the binding is an in-place update, then the @inplace@ field +-- of each (result) kernel is updated with the new in-place updates. +bindingFamily :: Pattern -> FusionGM FusedRes -> FusionGM FusedRes +bindingFamily pat = local bind + where idents = patternIdents pat + family = patternNames pat + bind env = foldl (bindingFamilyVar family) env idents + +bindingTransform :: PatElem -> VName -> SOAC.ArrayTransform -> FusionGM a -> FusionGM a +bindingTransform pe srcname trns = local $ \env -> + case M.lookup srcname $ varsInScope env of + Just (IsArray src' _ aliases input) -> + env { varsInScope = + M.insert vname + (IsArray src' (LetInfo attr) (srcname `S.insert` aliases) $ + trns `SOAC.addTransform` input) $ + varsInScope env + } + _ -> bindVar env (patElemIdent pe, S.singleton vname) + where vname = patElemName pe + attr = patElemAttr pe + +-- | Binds the fusion result to the environment. +bindRes :: FusedRes -> FusionGM a -> FusionGM a +bindRes rrr = local (\x -> x { fusedRes = rrr }) + +-- | The fusion transformation runs in this monad. The mutable +-- state refers to the fresh-names engine. +-- The reader hides the vtable that associates ... to ... (fill in, please). +-- The 'Either' monad is used for error handling. +runFusionGatherM :: MonadFreshNames m => + FusionGM a -> FusionGEnv -> m (Either Error a) +runFusionGatherM (FusionGM a) env = + modifyNameSource $ \src -> runReader (runStateT (runExceptT a) src) env + +------------------------------------------------------------------------ +--- Fusion Entry Points: gather the to-be-fused kernels@pgm level --- +--- and fuse them in a second pass! --- +------------------------------------------------------------------------ + +fuseSOACs :: Pass SOACS SOACS +fuseSOACs = + Pass { passName = "Fuse SOACs" + , passDescription = "Perform higher-order optimisation, i.e., fusion." + , passFunction = simplifySOACS <=< renameProg <=< intraproceduralTransformation fuseFun + } + +fuseFun :: FunDef -> PassM FunDef +fuseFun fun = do + let env = FusionGEnv { soacs = M.empty + , varsInScope = M.empty + , fusedRes = mempty + } + k <- cleanFusionResult <$> + liftEitherM (runFusionGatherM (fusionGatherFun fun) env) + if not $ rsucc k + then return fun + else liftEitherM $ runFusionGatherM (fuseInFun k fun) env + +fusionGatherFun :: FunDef -> FusionGM FusedRes +fusionGatherFun fundec = + bindingParams (funDefParams fundec) $ + fusionGatherBody mempty $ funDefBody fundec + +fuseInFun :: FusedRes -> FunDef -> FusionGM FunDef +fuseInFun res fundec = do + body' <- bindingParams (funDefParams fundec) $ + bindRes res $ + fuseInBody $ funDefBody fundec + return $ fundec { funDefBody = body' } + +--------------------------------------------------- +--------------------------------------------------- +---- RESULT's Data Structure +--------------------------------------------------- +--------------------------------------------------- + +-- | A type used for (hopefully) uniquely referring a producer SOAC. +-- The uniquely identifying value is the name of the first array +-- returned from the SOAC. +newtype KernName = KernName { unKernName :: VName } + deriving (Eq, Ord, Show) + +data FusedRes = FusedRes { + rsucc :: Bool + -- ^ Whether we have fused something anywhere. + + , outArr :: M.Map VName KernName + -- ^ Associates an array to the name of the + -- SOAC kernel that has produced it. + + , inpArr :: M.Map VName (S.Set KernName) + -- ^ Associates an array to the names of the + -- SOAC kernels that uses it. These sets include + -- only the SOAC input arrays used as full variables, i.e., no `a[i]'. + + , infusible :: Names + -- ^ the (names of) arrays that are not fusible, i.e., + -- + -- 1. they are either used other than input to SOAC kernels, or + -- + -- 2. are used as input to at least two different kernels that + -- are not located on disjoint control-flow branches, or + -- + -- 3. are used in the lambda expression of SOACs + + , kernels :: M.Map KernName FusedKer + -- ^ The map recording the uses + } + +instance Sem.Semigroup FusedRes where + res1 <> res2 = + FusedRes (rsucc res1 || rsucc res2) + (outArr res1 `M.union` outArr res2) + (M.unionWith S.union (inpArr res1) (inpArr res2) ) + (infusible res1 `S.union` infusible res2) + (kernels res1 `M.union` kernels res2) + +instance Monoid FusedRes where + mempty = FusedRes { rsucc = False, outArr = M.empty, inpArr = M.empty, + infusible = S.empty, kernels = M.empty } + mappend = (Sem.<>) + +isInpArrInResModKers :: FusedRes -> S.Set KernName -> VName -> Bool +isInpArrInResModKers ress kers nm = + case M.lookup nm (inpArr ress) of + Nothing -> False + Just s -> not $ S.null $ s `S.difference` kers + +getKersWithInpArrs :: FusedRes -> [VName] -> S.Set KernName +getKersWithInpArrs ress = + S.unions . mapMaybe (`M.lookup` inpArr ress) + +-- | extend the set of names to include all the names +-- produced via SOACs (by querring the vtable's soac) +expandSoacInpArr :: [VName] -> FusionGM [VName] +expandSoacInpArr = + foldM (\y nm -> do bnd <- asks $ M.lookup nm . soacs + case bnd of + Nothing -> return (y++[nm]) + Just nns -> return (y++nns ) + ) [] + +---------------------------------------------------------------------- +---------------------------------------------------------------------- + +soacInputs :: SOAC -> FusionGM ([VName], [VName]) +soacInputs soac = do + let (inp_idds, other_idds) = getIdentArr $ SOAC.inputs soac + (inp_nms0, other_nms0) = (inp_idds, other_idds) + inp_nms <- expandSoacInpArr inp_nms0 + other_nms <- expandSoacInpArr other_nms0 + return (inp_nms, other_nms) + +addNewKerWithInfusible :: FusedRes -> ([Ident], Certificates, SOAC, Names) -> Names -> FusionGM FusedRes +addNewKerWithInfusible res (idd, cs, soac, consumed) ufs = do + nm_ker <- KernName <$> newVName "ker" + scope <- askScope + let out_nms = map identName idd + new_ker = newKernel cs soac consumed out_nms scope + comb = M.unionWith S.union + os' = M.fromList [(arr,nm_ker) | arr <- out_nms] + `M.union` outArr res + is' = M.fromList [(arr,S.singleton nm_ker) + | arr <- map SOAC.inputArray $ SOAC.inputs soac] + `comb` inpArr res + return $ FusedRes (rsucc res) os' is' ufs + (M.insert nm_ker new_ker (kernels res)) + +lookupInput :: VName -> FusionGM (Maybe SOAC.Input) +lookupInput name = asks $ lookupArr name + +inlineSOACInput :: SOAC.Input -> FusionGM SOAC.Input +inlineSOACInput (SOAC.Input ts v t) = do + maybe_inp <- lookupInput v + case maybe_inp of + Nothing -> + return $ SOAC.Input ts v t + Just (SOAC.Input ts2 v2 t2) -> + return $ SOAC.Input (ts2<>ts) v2 t2 + +inlineSOACInputs :: SOAC -> FusionGM SOAC +inlineSOACInputs soac = do + inputs' <- mapM inlineSOACInput $ SOAC.inputs soac + return $ inputs' `SOAC.setInputs` soac + + +-- | Attempts to fuse between SOACs. Input: +-- @rem_bnds@ are the bindings remaining in the current body after @orig_soac@. +-- @lam_used_nms@ the infusible names +-- @res@ the fusion result (before processing the current soac) +-- @orig_soac@ and @out_idds@ the current SOAC and its binding pattern +-- @consumed@ is the set of names consumed by the SOAC. +-- Output: a new Fusion Result (after processing the current SOAC binding) +greedyFuse :: [Stm] -> Names -> FusedRes -> (Pattern, Certificates, SOAC, Names) + -> FusionGM FusedRes +greedyFuse rem_bnds lam_used_nms res (out_idds, cs, orig_soac, consumed) = do + soac <- inlineSOACInputs orig_soac + (inp_nms, other_nms) <- soacInputs soac + -- Assumption: the free vars in lambda are already in @infusible res@. + let out_nms = patternNames out_idds + isInfusible = (`S.member` infusible res) + is_screma = case orig_soac of + SOAC.Screma _ form _ -> + (isJust (isRedomapSOAC form) || isJust (isScanomapSOAC form)) && + not (isJust (isReduceSOAC form) || isJust (isScanSOAC form)) + _ -> False + -- + -- Conditions for fusion: + -- If current soac is a replicate OR (current soac a redomap/scanomap AND + -- (i) none of @out_idds@ belongs to the infusible set) + -- THEN try applying producer-consumer fusion + -- ELSE try applying horizontal fusion + -- (without duplicating computation in both cases) + + (ok_kers_compat, fused_kers, fused_nms, old_kers, oldker_nms) <- + if is_screma || any isInfusible out_nms + then horizontGreedyFuse rem_bnds res (out_idds, cs, soac, consumed) + else prodconsGreedyFuse res (out_idds, cs, soac, consumed) + -- + -- (ii) check whether fusing @soac@ will violate any in-place update + -- restriction, e.g., would move an input array past its in-place update. + let all_used_names = S.toList $ S.unions [lam_used_nms, S.fromList inp_nms, S.fromList other_nms] + has_inplace ker = any (`S.member` inplace ker) all_used_names + ok_inplace = not $ any has_inplace old_kers + -- + -- (iii) there are some kernels that use some of `out_idds' as inputs + -- (iv) and producer-consumer or horizontal fusion succeeds with those. + let fusible_ker = not (null old_kers) && ok_inplace && ok_kers_compat + -- + -- Start constructing the fusion's result: + -- (i) inparr ids other than vars will be added to infusible list, + -- (ii) will also become part of the infusible set the inparr vars + -- that also appear as inparr of another kernel, + -- BUT which said kernel is not the one we are fusing with (now)! + let mod_kerS = if fusible_ker then S.fromList oldker_nms else S.empty + let used_inps = filter (isInpArrInResModKers res mod_kerS) inp_nms + let ufs = S.unions [infusible res, S.fromList used_inps, + S.fromList other_nms `S.difference` + S.fromList (map SOAC.inputArray $ SOAC.inputs soac)] + let comb = M.unionWith S.union + + if not fusible_ker then + addNewKerWithInfusible res (patternIdents out_idds, cs, soac, consumed) ufs + else do + -- Need to suitably update `inpArr': + -- (i) first remove the inpArr bindings of the old kernel + let inpArr' = + foldl (\inpa (kold, knm) -> + S.foldl' + (\inpp nm -> + case M.lookup nm inpp of + Nothing -> inpp + Just s -> let new_set = S.delete knm s + in if S.null new_set + then M.delete nm inpp + else M.insert nm new_set inpp + ) + inpa $ arrInputs kold + ) + (inpArr res) (zip old_kers oldker_nms) + -- (ii) then add the inpArr bindings of the new kernel + let fused_ker_nms = zip fused_nms fused_kers + inpArr''= foldl (\inpa' (knm, knew) -> + M.fromList [ (k, S.singleton knm) + | k <- S.toList $ arrInputs knew ] + `comb` inpa' + ) + inpArr' fused_ker_nms + -- Update the kernels map (why not delete the ones that have been fused?) + let kernels' = M.fromList fused_ker_nms `M.union` kernels res + -- nothing to do for `outArr' (since we have not added a new kernel) + -- DO IMPROVEMENT: attempt to fuse the resulting kernel AGAIN until it fails, + -- but make sure NOT to add a new kernel! + return $ FusedRes True (outArr res) inpArr'' ufs kernels' + +prodconsGreedyFuse :: FusedRes -> (Pattern, Certificates, SOAC, Names) + -> FusionGM (Bool, [FusedKer], [KernName], [FusedKer], [KernName]) +prodconsGreedyFuse res (out_idds, cs, soac, consumed) = do + let out_nms = patternNames out_idds -- Extract VNames from output patterns + to_fuse_knmSet = getKersWithInpArrs res out_nms -- Find kernels which consume outputs + to_fuse_knms = S.toList to_fuse_knmSet + lookup_kern k = case M.lookup k (kernels res) of + Nothing -> throwError $ Error + ("In Fusion.hs, greedyFuse, comp of to_fuse_kers: " + ++ "kernel name not found in kernels field!") + Just ker -> return ker + to_fuse_kers <- mapM lookup_kern to_fuse_knms -- Get all consumer kernels + -- try producer-consumer fusion + (ok_kers_compat, fused_kers) <- do + kers <- forM to_fuse_kers $ + attemptFusion S.empty (patternNames out_idds) soac consumed + case sequence kers of + Nothing -> return (False, []) + Just kers' -> return (True, map certifyKer kers') + return (ok_kers_compat, fused_kers, to_fuse_knms, to_fuse_kers, to_fuse_knms) + where certifyKer k = k { certificates = certificates k <> cs } + +horizontGreedyFuse :: [Stm] -> FusedRes -> (Pattern, Certificates, SOAC, Names) + -> FusionGM (Bool, [FusedKer], [KernName], [FusedKer], [KernName]) +horizontGreedyFuse rem_bnds res (out_idds, cs, soac, consumed) = do + (inp_nms, _) <- soacInputs soac + let out_nms = patternNames out_idds + infusible_nms = S.fromList $ filter (`S.member` infusible res) out_nms + out_arr_nms = case soac of + -- the accumulator result cannot be fused! + SOAC.Screma _ (ScremaForm (_, scan_nes) (_, _, red_nes) _) _ -> + drop (length scan_nes + length red_nes) out_nms + SOAC.Stream _ frm _ _ -> drop (length $ getStreamAccums frm) out_nms + _ -> out_nms + to_fuse_knms1 = S.toList $ getKersWithInpArrs res (out_arr_nms++inp_nms) + to_fuse_knms2 = getKersWithSameInpSize (SOAC.width soac) res + to_fuse_knms = S.toList $ S.fromList $ to_fuse_knms1 ++ to_fuse_knms2 + lookupKernel k = case M.lookup k (kernels res) of + Nothing -> throwError $ Error + ("In Fusion.hs, greedyFuse, comp of to_fuse_kers: " + ++ "kernel name not found in kernels field!") + Just ker -> return ker + + -- for each kernel get the index in the bindings where the kernel is located + -- and sort based on the index so that partial fusion may succeed. + let bnd_nms = map (patternNames . stmPattern) rem_bnds + kernminds <- forM to_fuse_knms $ \ker_nm -> do + ker <- lookupKernel ker_nm + let out_nm = case fsoac ker of + SOAC.Stream _ frm _ _ + | x:_ <- drop (length $ getStreamAccums frm) $ outNames ker -> + x + SOAC.Screma _ (ScremaForm (_, scan_nes) (_, _, red_nes) _) _ + | x:_ <- drop (length scan_nes + length red_nes) $ outNames ker -> + x + _ -> head $ outNames ker + case L.findIndex (elem out_nm) bnd_nms of + Nothing -> return Nothing + Just i -> return $ Just (ker,ker_nm,i) + + scope <- askScope + let kernminds' = L.sortBy (\(_,_,i1) (_,_,i2)->compare i1 i2) $ catMaybes kernminds + soac_kernel = newKernel cs soac consumed out_nms scope + -- now try to fuse kernels one by one (in a fold); @ok_ind@ is the index of the + -- kernel until which fusion succeded, and @fused_ker@ is the resulted kernel. + (_,ok_ind,_,fused_ker,_) <- + foldM (\(cur_ok,n,prev_ind,cur_ker,ufus_nms) (ker, _ker_nm, bnd_ind) -> do + -- check that we still try fusion and that the intermediate + -- bindings do not use the results of cur_ker + let curker_outnms = outNames cur_ker + curker_outset = S.fromList curker_outnms + new_ufus_nms = S.fromList $ outNames ker ++ S.toList ufus_nms + -- disable horizontal fusion in the case when an output array of + -- producer SOAC is a non-trivially transformed input of the consumer + out_transf_ok = let ker_inp = SOAC.inputs $ fsoac ker + unfuse1 = S.fromList (map SOAC.inputArray ker_inp) `S.difference` + S.fromList (mapMaybe SOAC.isVarInput ker_inp) + unfuse2 = S.intersection curker_outset ufus_nms + in S.null $ S.intersection unfuse1 unfuse2 + -- Disable horizontal fusion if consumer has any + -- output transforms. + cons_no_out_transf = SOAC.nullTransforms $ outputTransform ker + + consumer_ok <- do let consumer_bnd = rem_bnds !! bnd_ind + maybesoac <- SOAC.fromExp $ stmExp consumer_bnd + case maybesoac of + -- check that consumer's lambda body does not use + -- directly the produced arrays (e.g., see noFusion3.fut). + Right conssoac -> return $ S.null $ S.intersection curker_outset $ + freeInBody $ lambdaBody $ SOAC.lambda conssoac + Left _ -> return True + + let interm_bnds_ok = cur_ok && consumer_ok && out_transf_ok && cons_no_out_transf && + foldl (\ok bnd-> ok && -- hardwired to False after first fail + -- (i) check that the in-between bindings do + -- not use the result of current kernel OR + S.null ( S.intersection curker_outset $ + freeInExp (stmExp bnd) ) || + --(ii) that the pattern-binding corresponds to + -- the result of the consumer kernel; in the + -- latter case it means it corresponds to a + -- kernel that has been fused in the consumer, + -- hence it should be ignored + not ( null $ curker_outnms `L.intersect` + patternNames (stmPattern bnd)) + ) True (drop (prev_ind+1) $ take bnd_ind rem_bnds) + if not interm_bnds_ok then return (False,n,bnd_ind,cur_ker,S.empty) + else do new_ker <- attemptFusion ufus_nms (outNames cur_ker) + (fsoac cur_ker) (fusedConsumed cur_ker) ker + case new_ker of + Nothing -> return (False, n,bnd_ind,cur_ker,S.empty) + Just krn-> return (True,n+1,bnd_ind,krn,new_ufus_nms) + ) (True,0,0,soac_kernel,infusible_nms) kernminds' + + -- Find the kernels we have fused into and the name of the last such + -- kernel (if any). + let (to_fuse_kers', to_fuse_knms',_) = unzip3 $ take ok_ind kernminds' + new_kernms = drop (ok_ind-1) to_fuse_knms' + + return (ok_ind>0, [fused_ker], new_kernms, to_fuse_kers', to_fuse_knms') + + where getKersWithSameInpSize :: SubExp -> FusedRes -> [KernName] + getKersWithSameInpSize sz ress = + map fst $ filter (\ (_,ker) -> sz == SOAC.width (fsoac ker)) $ M.toList $ kernels ress + +------------------------------------------------------------------------ +------------------------------------------------------------------------ +------------------------------------------------------------------------ +--- Fusion Gather for EXPRESSIONS and BODIES, --- +--- i.e., where work is being done: --- +--- i) bottom-up AbSyn traversal (backward analysis) --- +--- ii) soacs are fused greedily iff does not duplicate computation--- +--- E.g., (y1, y2, y3) = mapT(f, x1, x2[i]) --- +--- (z1, z2) = mapT(g1, y1, y2) --- +--- (q1, q2) = mapT(g2, y3, z1, a, y3) --- +--- res = reduce(op, ne, q1, q2, z2, y1, y3) --- +--- can be fused if y1,y2,y3, z1,z2, q1,q2 are not used elsewhere: --- +--- res = redomap(op, \(x1,x2i,a)-> --- +--- let (y1,y2,y3) = f (x1, x2i) in--- +--- let (z1,z2) = g1(y1, y2) in--- +--- let (q1,q2) = g2(y3, z1, a, y3) in--- +--- (q1, q2, z2, y1, y3) --- +--- x1, x2[i], a) --- +------------------------------------------------------------------------ +------------------------------------------------------------------------ +------------------------------------------------------------------------ + +fusionGatherBody :: FusedRes -> Body -> FusionGM FusedRes + +-- Some forms of do-loops can profitably be considered streamSeqs. We +-- are careful to ensure that the generated nested loop cannot itself +-- be considered a stream, to avoid infinite recursion. +fusionGatherBody fres (Body blore (stmsToList -> + Let (Pattern [] pes) bndtp + (DoLoop [] merge (ForLoop i it w loop_vars) body) + :bnds) res) | not $ null loop_vars = do + let (merge_params,merge_init) = unzip merge + (loop_params,loop_arrs) = unzip loop_vars + chunk_size <- newVName "chunk_size" + offset <- newVName "offset" + let chunk_param = Param chunk_size $ Prim int32 + offset_param = Param offset $ Prim $ IntType it + + acc_params <- forM merge_params $ \p -> + Param <$> newVName (baseString (paramName p) ++ "_outer") <*> + pure (paramType p) + + chunked_params <- forM loop_vars $ \(p,arr) -> + Param <$> newVName (baseString arr ++ "_chunk") <*> + pure (paramType p `arrayOfRow` Futhark.Var chunk_size) + + let lam_params = chunk_param : acc_params ++ [offset_param] ++ chunked_params + + lam_body <- runBodyBinder $ localScope (scopeOfLParams lam_params) $ do + let merge' = zip merge_params $ map (Futhark.Var . paramName) acc_params + j <- newVName "j" + loop_body <- runBodyBinder $ do + forM_ (zip loop_params chunked_params) $ \(p,a_p) -> + letBindNames_ [paramName p] $ BasicOp $ Index (paramName a_p) $ + fullSlice (paramType a_p) [DimFix $ Futhark.Var j] + letBindNames_ [i] $ BasicOp $ BinOp (Add it) (Futhark.Var offset) (Futhark.Var j) + return body + eBody [pure $ + DoLoop [] merge' (ForLoop j it (Futhark.Var chunk_size) []) loop_body, + pure $ + BasicOp $ BinOp (Add Int32) (Futhark.Var offset) (Futhark.Var chunk_size)] + let lam = Lambda { lambdaParams = lam_params + , lambdaBody = lam_body + , lambdaReturnType = map paramType $ acc_params ++ [offset_param] + } + stream = Futhark.Stream w (Sequential $ merge_init ++ [intConst it 0]) lam loop_arrs + + -- It is important that the (discarded) final-offset is not the + -- first element in the pattern, as we use the first element to + -- identify the SOAC in the second phase of fusion. + discard <- newVName "discard" + let discard_pe = PatElem discard $ Prim int32 + + fusionGatherBody fres $ Body blore + (oneStm (Let (Pattern [] (pes<>[discard_pe])) bndtp (Op stream))<>stmsFromList bnds) res + +fusionGatherBody fres (Body _ (stmsToList -> (bnd@(Let pat _ e):bnds)) res) = do + maybesoac <- SOAC.fromExp e + case maybesoac of + Right soac@(SOAC.Scatter _len lam _ivs _as) -> do + -- We put the variables produced by Scatter into the infusible + -- set to force horizontal fusion. It is not possible to + -- producer/consumer-fuse Scatter anyway. + fres' <- addNamesToInfusible fres $ S.fromList $ patternNames pat + mapLike fres' soac lam + + Right soac@(SOAC.GenReduce _ _ lam _) -> do + -- We put the variables produced by GenReduce into the infusible + -- set to force horizontal fusion. It is not possible to + -- producer/consumer-fuse GenReduce anyway. + fres' <- addNamesToInfusible fres $ S.fromList $ patternNames pat + mapLike fres' soac lam + + Right soac@(SOAC.Screma _ (ScremaForm (scan_lam, scan_nes) + (_, reduce_lam, reduce_nes) + map_lam) _) -> + reduceLike soac [scan_lam, reduce_lam, map_lam] $ scan_nes <> reduce_nes + + Right soac@(SOAC.Stream _ form lam _) -> do + -- a redomap does not neccessarily start a new kernel, e.g., + -- @let a= reduce(+,0,A) in ... bnds ... in let B = map(f,A)@ + -- can be fused into a redomap that replaces the @map@, if @a@ + -- and @B@ are defined in the same scope and @bnds@ does not uses @a@. + -- a redomap always starts a new kernel + let lambdas = case form of + Parallel _ _ lout _ -> [lout, lam] + _ -> [lam] + reduceLike soac lambdas $ getStreamAccums form + + _ | [pe] <- patternValueElements pat, + Just (src,trns) <- SOAC.transformFromExp (stmCerts bnd) e -> + bindingTransform pe src trns $ fusionGatherBody fres body + | otherwise -> do + let pat_vars = map (BasicOp . SubExp . Var) $ patternNames pat + bres <- gatherStmPattern pat e $ fusionGatherBody fres body + bres' <- checkForUpdates bres e + foldM fusionGatherExp bres' (e:pat_vars) + + where body = mkBody (stmsFromList bnds) res + cs = stmCerts bnd + rem_bnds = bnd : bnds + consumed = consumedInExp $ Alias.analyseExp e + + reduceLike soac lambdas nes = do + (used_lam, lres) <- foldM fusionGatherLam (S.empty, fres) lambdas + bres <- bindingFamily pat $ fusionGatherBody lres body + bres' <- foldM fusionGatherSubExp bres nes + consumed' <- varsAliases consumed + greedyFuse rem_bnds used_lam bres' (pat, cs, soac, consumed') + + mapLike fres' soac lambda = do + bres <- bindingFamily pat $ fusionGatherBody fres' body + (used_lam, blres) <- fusionGatherLam (S.empty, bres) lambda + consumed' <- varsAliases consumed + greedyFuse rem_bnds used_lam blres (pat, cs, soac, consumed') + +fusionGatherBody fres (Body _ _ res) = + foldM fusionGatherExp fres $ map (BasicOp . SubExp) res + +fusionGatherExp :: FusedRes -> Exp -> FusionGM FusedRes + +----------------------------------------- +---- Index/If ---- +----------------------------------------- + +fusionGatherExp fres (DoLoop ctx val form loop_body) = do + fres' <- addNamesToInfusible fres $ freeIn form <> freeIn ctx <> freeIn val + let form_idents = + case form of + ForLoop i _ _ loopvars -> + Ident i (Prim int32) : map (paramIdent . fst) loopvars + WhileLoop{} -> [] + + new_res <- binding (zip (form_idents ++ map (paramIdent . fst) (ctx<>val)) $ + repeat mempty) $ + fusionGatherBody mempty loop_body + -- make the inpArr infusible, so that they + -- cannot be fused from outside the loop: + let (inp_arrs, _) = unzip $ M.toList $ inpArr new_res + let new_res' = new_res { infusible = foldl (flip S.insert) (infusible new_res) inp_arrs } + -- merge new_res with fres' + return $ new_res' <> fres' + +fusionGatherExp fres (If cond e_then e_else _) = do + then_res <- fusionGatherBody mempty e_then + else_res <- fusionGatherBody mempty e_else + let both_res = then_res <> else_res + fres' <- fusionGatherSubExp fres cond + mergeFusionRes fres' both_res + +----------------------------------------------------------------------------------- +--- Errors: all SOACs, (because normalization ensures they appear +--- directly in let exp, i.e., let x = e) +----------------------------------------------------------------------------------- + +fusionGatherExp _ (Op Futhark.Screma{}) = errorIllegal "screma" +fusionGatherExp _ (Op Futhark.Scatter{}) = errorIllegal "write" + +----------------------------------- +---- Generic Traversal ---- +----------------------------------- + +fusionGatherExp fres e = + addNamesToInfusible fres $ freeInExp e + +fusionGatherSubExp :: FusedRes -> SubExp -> FusionGM FusedRes +fusionGatherSubExp fres (Var idd) = addVarToInfusible fres idd +fusionGatherSubExp fres _ = return fres + +addNamesToInfusible :: FusedRes -> Names -> FusionGM FusedRes +addNamesToInfusible fres = foldM addVarToInfusible fres . S.toList + +addVarToInfusible :: FusedRes -> VName -> FusionGM FusedRes +addVarToInfusible fres name = do + trns <- asks $ lookupArr name + let name' = case trns of + Nothing -> name + Just (SOAC.Input _ orig _) -> orig + return fres { infusible = S.insert name' $ infusible fres } + +-- Lambdas create a new scope. Disallow fusing from outside lambda by +-- adding inp_arrs to the infusible set. +fusionGatherLam :: (Names, FusedRes) -> Lambda -> FusionGM (S.Set VName, FusedRes) +fusionGatherLam (u_set,fres) (Lambda idds body _) = do + new_res <- bindingParams idds $ fusionGatherBody mempty body + -- make the inpArr infusible, so that they + -- cannot be fused from outside the lambda: + let inp_arrs = S.fromList $ M.keys $ inpArr new_res + let unfus = infusible new_res `S.union` inp_arrs + bnds <- M.keys <$> asks varsInScope + let unfus' = unfus `S.intersection` S.fromList bnds + -- merge fres with new_res' + let new_res' = new_res { infusible = unfus' } + -- merge new_res with fres' + return (u_set `S.union` unfus', new_res' <> fres) + +------------------------------------------------------------- +------------------------------------------------------------- +--- FINALLY, Substitute the kernels in function +------------------------------------------------------------- +------------------------------------------------------------- + +fuseInBody :: Body -> FusionGM Body + +fuseInBody (Body _ stms res) + | Let pat aux e:bnds <- stmsToList stms = do + body' <- bindingPat pat $ fuseInBody $ mkBody (stmsFromList bnds) res + soac_bnds <- replaceSOAC pat aux e + return $ insertStms soac_bnds body' + | otherwise = return $ Body () mempty res + +fuseInExp :: Exp -> FusionGM Exp + +-- Handle loop specially because we need to bind the types of the +-- merge variables. +fuseInExp (DoLoop ctx val form loopbody) = + binding (zip form_idents $ repeat mempty) $ + bindingParams (map fst $ ctx ++ val) $ + DoLoop ctx val form <$> fuseInBody loopbody + where form_idents = case form of + WhileLoop{} -> [] + ForLoop i it _ loopvars -> + Ident i (Prim $ IntType it) : + map (paramIdent . fst) loopvars + +fuseInExp e = mapExpM fuseIn e + +fuseIn :: Mapper SOACS SOACS FusionGM +fuseIn = identityMapper { + mapOnBody = const fuseInBody + , mapOnOp = mapSOACM identitySOACMapper { mapOnSOACLambda = fuseInLambda } + } + +fuseInLambda :: Lambda -> FusionGM Lambda +fuseInLambda (Lambda params body rtp) = do + body' <- bindingParams params $ fuseInBody body + return $ Lambda params body' rtp + +replaceSOAC :: Pattern -> StmAux () -> Exp -> FusionGM (Stms SOACS) +replaceSOAC (Pattern _ []) _ _ = return mempty +replaceSOAC pat@(Pattern _ (patElem : _)) aux e = do + fres <- asks fusedRes + let pat_nm = patElemName patElem + names = patternIdents pat + case M.lookup pat_nm (outArr fres) of + Nothing -> + oneStm . Let pat aux <$> fuseInExp e + Just knm -> + case M.lookup knm (kernels fres) of + Nothing -> throwError $ Error + ("In Fusion.hs, replaceSOAC, outArr in ker_name " + ++"which is not in Res: "++pretty (unKernName knm)) + Just ker -> do + when (null $ fusedVars ker) $ + throwError $ Error + ("In Fusion.hs, replaceSOAC, unfused kernel " + ++"still in result: "++pretty names) + insertKerSOAC (outNames ker) ker + +insertKerSOAC :: [VName] -> FusedKer -> FusionGM (Stms SOACS) +insertKerSOAC names ker = do + new_soac' <- finaliseSOAC $ fsoac ker + runBinder_ $ do + f_soac <- SOAC.toSOAC new_soac' + -- The fused kernel may consume more than the original SOACs (see + -- issue #224). We insert copy expressions to fix it. + f_soac' <- copyNewlyConsumed (fusedConsumed ker) $ addOpAliases f_soac + validents <- zipWithM newIdent (map baseString names) $ SOAC.typeOf new_soac' + letBind_ (basicPattern [] validents) $ Op f_soac' + transformOutput (outputTransform ker) names validents + +-- | Perform simplification and fusion inside the lambda(s) of a SOAC. +finaliseSOAC :: SOAC.SOAC SOACS -> FusionGM (SOAC.SOAC SOACS) +finaliseSOAC new_soac = + case new_soac of + SOAC.Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs -> do + scan_lam' <- simplifyAndFuseInLambda scan_lam + red_lam' <- simplifyAndFuseInLambda red_lam + map_lam' <- simplifyAndFuseInLambda map_lam + return $ SOAC.Screma w (ScremaForm (scan_lam', scan_nes) + (comm, red_lam', red_nes) + map_lam') + arrs + SOAC.Scatter w lam inps dests -> do + lam' <- simplifyAndFuseInLambda lam + return $ SOAC.Scatter w lam' inps dests + SOAC.GenReduce w ops lam arrs -> do + lam' <- simplifyAndFuseInLambda lam + return $ SOAC.GenReduce w ops lam' arrs + SOAC.Stream w form lam inps -> do + lam' <- simplifyAndFuseInLambda lam + return $ SOAC.Stream w form lam' inps + +simplifyAndFuseInLambda :: Lambda -> FusionGM Lambda +simplifyAndFuseInLambda lam = do + let args = replicate (length $ lambdaParams lam) Nothing + lam' <- simplifyLambda lam args + (_, nfres) <- fusionGatherLam (S.empty, mkFreshFusionRes) lam' + let nfres' = cleanFusionResult nfres + bindRes nfres' $ fuseInLambda lam' + +copyNewlyConsumed :: Names + -> Futhark.SOAC (Aliases.Aliases SOACS) + -> Binder SOACS (Futhark.SOAC SOACS) +copyNewlyConsumed was_consumed soac = + case soac of + Futhark.Screma w (Futhark.ScremaForm + (scan_lam, scan_nes) + (comm, reduce_lam, reduce_nes) + map_lam) arrs -> do + -- Copy any arrays that are consumed now, but were not in the + -- constituents. + arrs' <- mapM copyConsumedArr arrs + -- Any consumed free variables will have to be copied inside the + -- lambda, and we have to substitute the name of the copy for + -- the original. + map_lam' <- copyFreeInLambda map_lam + return $ Futhark.Screma w + (Futhark.ScremaForm + (Aliases.removeLambdaAliases scan_lam, scan_nes) + (comm, Aliases.removeLambdaAliases reduce_lam, reduce_nes) + map_lam') arrs' + + _ -> return $ removeOpAliases soac + where consumed = consumedInOp soac + newly_consumed = consumed `S.difference` was_consumed + + copyConsumedArr a + | a `S.member` newly_consumed = + letExp (baseString a <> "_copy") $ BasicOp $ Copy a + | otherwise = return a + + copyFreeInLambda lam = do + let free_consumed = consumedByLambda lam `S.difference` + S.fromList (map paramName $ lambdaParams lam) + (bnds, subst) <- + foldM copyFree (mempty, mempty) $ S.toList free_consumed + let lam' = Aliases.removeLambdaAliases lam + return $ if null bnds + then lam' + else lam' { lambdaBody = + insertStms bnds $ + substituteNames subst $ lambdaBody lam' + } + + copyFree (bnds, subst) v = do + v_copy <- newVName $ baseString v <> "_copy" + copy <- mkLetNamesM [v_copy] $ BasicOp $ Copy v + return (oneStm copy<>bnds, M.insert v v_copy subst) + +--------------------------------------------------- +--------------------------------------------------- +---- HELPERS +--------------------------------------------------- +--------------------------------------------------- + +-- | Get a new fusion result, i.e., for when entering a new scope, +-- e.g., a new lambda or a new loop. +mkFreshFusionRes :: FusedRes +mkFreshFusionRes = + FusedRes { rsucc = False, outArr = M.empty, inpArr = M.empty, + infusible = S.empty, kernels = M.empty } + +mergeFusionRes :: FusedRes -> FusedRes -> FusionGM FusedRes +mergeFusionRes res1 res2 = do + let ufus_mres = infusible res1 `S.union` infusible res2 + inp_both <- expandSoacInpArr $ M.keys $ inpArr res1 `M.intersection` inpArr res2 + let m_unfus = foldl (flip S.insert) ufus_mres inp_both + return $ FusedRes (rsucc res1 || rsucc res2) + (outArr res1 `M.union` outArr res2) + (M.unionWith S.union (inpArr res1) (inpArr res2) ) + m_unfus + (kernels res1 `M.union` kernels res2) + + +-- | The expression arguments are supposed to be array-type exps. +-- Returns a tuple, in which the arrays that are vars are in the +-- first element of the tuple, and the one which are indexed or +-- transposes (or otherwise transformed) should be in the second. +-- +-- E.g., for expression `mapT(f, a, b[i])', the result should be +-- `([a],[b])' +getIdentArr :: [SOAC.Input] -> ([VName], [VName]) +getIdentArr = foldl comb ([],[]) + where comb (vs,os) (SOAC.Input ts idd _) + | SOAC.nullTransforms ts = (idd:vs, os) + comb (vs, os) inp = + (vs, SOAC.inputArray inp : os) + +cleanFusionResult :: FusedRes -> FusedRes +cleanFusionResult fres = + let newks = M.filter (not . null . fusedVars) (kernels fres) + newoa = M.filter (`M.member` newks) (outArr fres) + newia = M.map (S.filter (`M.member` newks)) (inpArr fres) + in fres { outArr = newoa, inpArr = newia, kernels = newks } + +-------------- +--- Errors --- +-------------- + +errorIllegal :: String -> FusionGM FusedRes +errorIllegal soac_name = + throwError $ Error + ("In Fusion.hs, soac "++soac_name++" appears illegally in pgm!") diff --git a/src/Futhark/Optimise/Fusion/Composing.hs b/src/Futhark/Optimise/Fusion/Composing.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Fusion/Composing.hs @@ -0,0 +1,213 @@ +-- | Facilities for composing SOAC functions. Mostly intended for use +-- by the fusion module, but factored into a separate module for ease +-- of testing, debugging and development. Of course, there is nothing +-- preventing you from using the exported functions whereever you +-- want. +-- +-- Important: this module is \"dumb\" in the sense that it does not +-- check the validity of its inputs, and does not have any +-- functionality for massaging SOACs to be fusible. It is assumed +-- that the given SOACs are immediately compatible. +-- +-- The module will, however, remove duplicate inputs after fusion. +module Futhark.Optimise.Fusion.Composing + ( fuseMaps + , fuseRedomap + , mergeReduceOps + ) + where + +import Data.List +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe + +import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC + +import Futhark.Representation.AST +import Futhark.Binder (Bindable(..), insertStm, insertStms, mkLet) +import Futhark.Construct (mapResult) +import Futhark.Util (splitAt3, takeLast, dropLast) + +-- | @fuseMaps lam1 inp1 out1 lam2 inp2@ fuses the function @lam1@ into +-- @lam2@. Both functions must be mapping functions, although @lam2@ +-- may have leading reduction parameters. @inp1@ and @inp2@ are the +-- array inputs to the SOACs containing @lam1@ and @lam2@ +-- respectively. @out1@ are the identifiers to which the output of +-- the SOAC containing @lam1@ is bound. It is nonsensical to call +-- this function unless the intersection of @out1@ and @inp2@ is +-- non-empty. +-- +-- If @lam2@ accepts more parameters than there are elements in +-- @inp2@, it is assumed that the surplus (which are positioned at the +-- beginning of the parameter list) are reduction (accumulator) +-- parameters, that do not correspond to array elements, and they are +-- thus not modified. +-- +-- The result is the fused function, and a list of the array inputs +-- expected by the SOAC containing the fused function. +fuseMaps :: Bindable lore => + Names -- ^ The producer var names that still need to be returned + -> Lambda lore -- ^ Function of SOAC to be fused. + -> [SOAC.Input] -- ^ Input of SOAC to be fused. + -> [(VName,Ident)] -- ^ Output of SOAC to be fused. The + -- first identifier is the name of the + -- actual output, where the second output + -- is an identifier that can be used to + -- bind a single element of that output. + -> Lambda lore -- ^ Function to be fused with. + -> [SOAC.Input] -- ^ Input of SOAC to be fused with. + -> (Lambda lore, [SOAC.Input]) -- ^ The fused lambda and the inputs of + -- the resulting SOAC. +fuseMaps unfus_nms lam1 inp1 out1 lam2 inp2 = (lam2', M.elems inputmap) + where lam2' = + lam2 { lambdaParams = [ Param name t + | Ident name t <- lam2redparams ++ M.keys inputmap ] + , lambdaBody = new_body2' + } + new_body2 = let bnds res = [ mkLet [] [p] $ BasicOp $ SubExp e + | (p,e) <- zip pat res] + bindLambda res = + stmsFromList (bnds res) `insertStms` makeCopiesInner (lambdaBody lam2) + in makeCopies $ mapResult bindLambda (lambdaBody lam1) + new_body2_rses = bodyResult new_body2 + new_body2'= new_body2 { bodyResult = new_body2_rses ++ + map (Var . identName) unfus_pat } + -- infusible variables are added at the end of the result/pattern/type + (lam2redparams, unfus_pat, pat, inputmap, makeCopies, makeCopiesInner) = + fuseInputs unfus_nms lam1 inp1 out1 lam2 inp2 + --(unfus_accpat, unfus_arrpat) = splitAt (length unfus_accs) unfus_pat + +fuseInputs :: Bindable lore => + Names + -> Lambda lore -> [SOAC.Input] -> [(VName,Ident)] + -> Lambda lore -> [SOAC.Input] + -> ([Ident], [Ident], [Ident], + M.Map Ident SOAC.Input, + Body lore -> Body lore, Body lore -> Body lore) +fuseInputs unfus_nms lam1 inp1 out1 lam2 inp2 = + (lam2redparams, unfus_vars, outbnds, inputmap, makeCopies, makeCopiesInner) + where (lam2redparams, lam2arrparams) = + splitAt (length lam2params - length inp2) lam2params + lam1params = map paramIdent $ lambdaParams lam1 + lam2params = map paramIdent $ lambdaParams lam2 + lam1inputmap = M.fromList $ zip lam1params inp1 + lam2inputmap = M.fromList $ zip lam2arrparams inp2 + (lam2inputmap', makeCopiesInner) = removeDuplicateInputs lam2inputmap + originputmap = lam1inputmap `M.union` lam2inputmap' + outins = uncurry (outParams $ map fst out1) $ + unzip $ M.toList lam2inputmap' + outbnds= filterOutParams out1 outins + (inputmap, makeCopies) = + removeDuplicateInputs $ originputmap `M.difference` outins + -- Cosmin: @unfus_vars@ is supposed to be the lam2 vars corresponding to unfus_nms (?) + getVarParPair x = case SOAC.isVarInput (snd x) of + Just nm -> Just (nm, fst x) + Nothing -> Nothing --should not be reached! + outinsrev = M.fromList $ mapMaybe getVarParPair $ M.toList outins + unfusible outname + | outname `S.member` unfus_nms = + outname `M.lookup` M.union outinsrev (M.fromList out1) + unfusible _ = Nothing + unfus_vars= mapMaybe (unfusible . fst) out1 + +outParams :: [VName] -> [Ident] -> [SOAC.Input] + -> M.Map Ident SOAC.Input +outParams out1 lam2arrparams inp2 = + M.fromList $ mapMaybe isOutParam $ zip lam2arrparams inp2 + where isOutParam (p, inp) + | Just a <- SOAC.isVarInput inp, + a `elem` out1 = Just (p, inp) + isOutParam _ = Nothing + +filterOutParams :: [(VName,Ident)] + -> M.Map Ident SOAC.Input + -> [Ident] +filterOutParams out1 outins = + snd $ mapAccumL checkUsed outUsage out1 + where outUsage = M.foldlWithKey' add M.empty outins + where add m p inp = + case SOAC.isVarInput inp of + Just v -> M.insertWith (++) v [p] m + Nothing -> m + + checkUsed m (a,ra) = + case M.lookup a m of + Just (p:ps) -> (M.insert a ps m, p) + _ -> (m, ra) + +removeDuplicateInputs :: Bindable lore => + M.Map Ident SOAC.Input + -> (M.Map Ident SOAC.Input, Body lore -> Body lore) +removeDuplicateInputs = fst . M.foldlWithKey' comb ((M.empty, id), M.empty) + where comb ((parmap, inner), arrmap) par arr = + case M.lookup arr arrmap of + Nothing -> ((M.insert par arr parmap, inner), + M.insert arr (identName par) arrmap) + Just par' -> ((parmap, inner . forward par par'), + arrmap) + forward to from b = + mkLet [] [to] (BasicOp $ SubExp $ Var from) + `insertStm` b + +fuseRedomap :: Bindable lore => + Names -> [VName] + -> Lambda lore -> [SubExp] -> [SubExp] -> [SOAC.Input] + -> [(VName,Ident)] + -> Lambda lore -> [SubExp] -> [SubExp] -> [SOAC.Input] + -> (Lambda lore, [SOAC.Input]) +fuseRedomap unfus_nms outVars p_lam p_scan_nes p_red_nes p_inparr outPairs + c_lam c_scan_nes c_red_nes c_inparr = + -- We hack the implementation of map o redomap to handle this case: + -- (i) we remove the accumulator formal paramter and corresponding + -- (body) result from from redomap's fold-lambda body + let p_num_nes = length p_scan_nes + length p_red_nes + unfus_arrs = filter (`S.member` unfus_nms) outVars + p_lam_body = lambdaBody p_lam + (p_lam_scan_ts, p_lam_red_ts, p_lam_map_ts) = + splitAt3 (length p_scan_nes) (length p_red_nes) $ lambdaReturnType p_lam + (p_lam_scan_res, p_lam_red_res, p_lam_map_res) = + splitAt3 (length p_scan_nes) (length p_red_nes) $ bodyResult p_lam_body + p_lam_hacked = p_lam { lambdaParams = takeLast (length p_inparr) $ lambdaParams p_lam + , lambdaBody = p_lam_body { bodyResult = p_lam_map_res } + , lambdaReturnType = p_lam_map_ts } + + -- (ii) we remove the accumulator's (global) output result from + -- @outPairs@, then ``map o redomap'' fuse the two lambdas + -- (in the usual way), and construct the extra return types + -- for the arrays that fall through. + (res_lam, new_inp) = fuseMaps (S.fromList unfus_arrs) p_lam_hacked p_inparr + (drop p_num_nes outPairs) c_lam c_inparr + (res_lam_scan_ts, res_lam_red_ts, res_lam_map_ts) = + splitAt3 (length c_scan_nes) (length c_red_nes) $ lambdaReturnType res_lam + (_,extra_map_ts) = unzip $ filter (\(nm,_)->elem nm unfus_arrs) $ + zip (drop p_num_nes outVars) $ drop p_num_nes $ + lambdaReturnType p_lam + + -- (iii) Finally, we put back the accumulator's formal parameter and + -- (body) result in the first position of the obtained lambda. + accpars = dropLast (length p_inparr) $ lambdaParams p_lam + res_body = lambdaBody res_lam + (res_lam_scan_res, res_lam_red_res, res_lam_map_res) = + splitAt3 (length c_scan_nes) (length c_red_nes) $ bodyResult res_body + res_body'= res_body { bodyResult = p_lam_scan_res ++ res_lam_scan_res ++ + p_lam_red_res ++ res_lam_red_res ++ + res_lam_map_res } + res_lam' = res_lam { lambdaParams = accpars ++ lambdaParams res_lam + , lambdaBody = res_body' + , lambdaReturnType = p_lam_scan_ts ++ res_lam_scan_ts ++ + p_lam_red_ts ++ res_lam_red_ts ++ + res_lam_map_ts ++ extra_map_ts + } + in (res_lam', new_inp) + + +mergeReduceOps :: Lambda lore -> Lambda lore -> Lambda lore +mergeReduceOps (Lambda par1 bdy1 rtp1) (Lambda par2 bdy2 rtp2) = + let body' = Body (bodyAttr bdy1) + (bodyStms bdy1 <> bodyStms bdy2) + (bodyResult bdy1 ++ bodyResult bdy2) + (len1, len2) = (length rtp1, length rtp2) + par' = take len1 par1 ++ take len2 par2 ++ drop len1 par1 ++ drop len2 par2 + in Lambda par' body' (rtp1++rtp2) diff --git a/src/Futhark/Optimise/Fusion/LoopKernel.hs b/src/Futhark/Optimise/Fusion/LoopKernel.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Fusion/LoopKernel.hs @@ -0,0 +1,786 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +module Futhark.Optimise.Fusion.LoopKernel + ( FusedKer(..) + , newKernel + , inputs + , setInputs + , arrInputs + , kernelType + , transformOutput + , attemptFusion + , SOAC + , MapNest + ) + where + +import Control.Applicative +import Control.Arrow (first) +import Control.Monad +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import Data.Maybe +import Data.Semigroup ((<>)) +import Data.List + +import Futhark.Representation.SOACS hiding (SOAC(..)) +import qualified Futhark.Representation.SOACS as Futhark +import Futhark.Transform.Rename (renameLambda) +import Futhark.Transform.Substitute +import Futhark.MonadFreshNames +import qualified Futhark.Analysis.HORepresentation.SOAC as SOAC +import qualified Futhark.Analysis.HORepresentation.MapNest as MapNest +import Futhark.Pass.ExtractKernels.ISRWIM (rwimPossible) +import Futhark.Optimise.Fusion.TryFusion +import Futhark.Optimise.Fusion.Composing +import Futhark.Construct +import Futhark.Util (splitAt3) + +type SOAC = SOAC.SOAC SOACS +type MapNest = MapNest.MapNest SOACS + +-- XXX: This function is very gross. +transformOutput :: SOAC.ArrayTransforms -> [VName] -> [Ident] + -> Binder SOACS () +transformOutput ts names = descend ts + where descend ts' validents = + case SOAC.viewf ts' of + SOAC.EmptyF -> + forM_ (zip names validents) $ \(k, valident) -> + letBindNames [k] $ BasicOp $ SubExp $ Var $ identName valident + t SOAC.:< ts'' -> do + let (es,css) = unzip $ map (applyTransform t) validents + mkPat (Ident nm tp) = Pattern [] [PatElem nm tp] + opts <- concat <$> mapM primOpType es + newIds <- forM (zip names opts) $ \(k, opt) -> + newIdent (baseString k) opt + forM_ (zip3 css newIds es) $ \(cs,ids,e) -> + certifying cs $ letBind (mkPat ids) (BasicOp e) + descend ts'' newIds + +applyTransform :: SOAC.ArrayTransform -> Ident -> (BasicOp, Certificates) +applyTransform (SOAC.Rearrange cs perm) v = + (Rearrange perm $ identName v, cs) +applyTransform (SOAC.Reshape cs shape) v = + (Reshape shape $ identName v, cs) +applyTransform (SOAC.ReshapeOuter cs shape) v = + let shapes = reshapeOuter shape 1 $ arrayShape $ identType v + in (Reshape shapes $ identName v, cs) +applyTransform (SOAC.ReshapeInner cs shape) v = + let shapes = reshapeInner shape 1 $ arrayShape $ identType v + in (Reshape shapes $ identName v, cs) +applyTransform (SOAC.Replicate cs n) v = + (Replicate n $ Var $ identName v, cs) + +inputToOutput :: SOAC.Input -> Maybe (SOAC.ArrayTransform, SOAC.Input) +inputToOutput (SOAC.Input ts ia iat) = + case SOAC.viewf ts of + t SOAC.:< ts' -> Just (t, SOAC.Input ts' ia iat) + SOAC.EmptyF -> Nothing + +data FusedKer = FusedKer { + fsoac :: SOAC + -- ^ the SOAC expression, e.g., mapT( f(a,b), x, y ) + + , inplace :: Names + -- ^ Variables used in in-place updates in the kernel itself, as + -- well as on the path to the kernel from the current position. + -- This is used to avoid fusion that would violate in-place + -- restrictions. + + , fusedVars :: [VName] + -- ^ whether at least a fusion has been performed. + + , fusedConsumed :: Names + -- ^ The set of variables that were consumed by the SOACs + -- contributing to this kernel. Note that, by the type rules, the + -- final SOAC may actually consume _more_ than its original + -- contributors, which implies the need for 'Copy' expressions. + + , kernelScope :: Scope SOACS + -- ^ The names in scope at the kernel. + + , outputTransform :: SOAC.ArrayTransforms + , outNames :: [VName] + , certificates :: Certificates + } + deriving (Show) + +newKernel :: Certificates -> SOAC -> Names -> [VName] -> Scope SOACS -> FusedKer +newKernel cs soac consumed out_nms scope = + FusedKer { fsoac = soac + , inplace = consumed + , fusedVars = [] + , fusedConsumed = consumed + , outputTransform = SOAC.noTransforms + , outNames = out_nms + , kernelScope = scope + , certificates = cs + } + +arrInputs :: FusedKer -> S.Set VName +arrInputs = S.fromList . map SOAC.inputArray . inputs + +inputs :: FusedKer -> [SOAC.Input] +inputs = SOAC.inputs . fsoac + +setInputs :: [SOAC.Input] -> FusedKer -> FusedKer +setInputs inps ker = ker { fsoac = inps `SOAC.setInputs` fsoac ker } + +kernelType :: FusedKer -> [Type] +kernelType = SOAC.typeOf . fsoac + +tryOptimizeSOAC :: Names -> [VName] -> SOAC -> Names -> FusedKer + -> TryFusion FusedKer +tryOptimizeSOAC unfus_nms outVars soac consumed ker = do + (soac', ots) <- optimizeSOAC Nothing soac mempty + let ker' = map (addInitialTransformIfRelevant ots) (inputs ker) `setInputs` ker + outIdents = zipWith Ident outVars $ SOAC.typeOf soac' + ker'' = fixInputTypes outIdents ker' + applyFusionRules unfus_nms outVars soac' consumed ker'' + where addInitialTransformIfRelevant ots inp + | SOAC.inputArray inp `elem` outVars = + SOAC.addInitialTransforms ots inp + | otherwise = + inp + +tryOptimizeKernel :: Names -> [VName] -> SOAC -> Names -> FusedKer + -> TryFusion FusedKer +tryOptimizeKernel unfus_nms outVars soac consumed ker = do + ker' <- optimizeKernel (Just outVars) ker + applyFusionRules unfus_nms outVars soac consumed ker' + +tryExposeInputs :: Names -> [VName] -> SOAC -> Names -> FusedKer + -> TryFusion FusedKer +tryExposeInputs unfus_nms outVars soac consumed ker = do + (ker', ots) <- exposeInputs outVars ker + if SOAC.nullTransforms ots + then fuseSOACwithKer unfus_nms outVars soac consumed ker' + else do + (soac', ots') <- pullOutputTransforms soac ots + let outIdents = zipWith Ident outVars $ SOAC.typeOf soac' + ker'' = fixInputTypes outIdents ker' + if SOAC.nullTransforms ots' + then applyFusionRules unfus_nms outVars soac' consumed ker'' + else fail "tryExposeInputs could not pull SOAC transforms" + +fixInputTypes :: [Ident] -> FusedKer -> FusedKer +fixInputTypes outIdents ker = + ker { fsoac = fixInputTypes' $ fsoac ker } + where fixInputTypes' soac = + map fixInputType (SOAC.inputs soac) `SOAC.setInputs` soac + fixInputType (SOAC.Input ts v _) + | Just v' <- find ((==v) . identName) outIdents = + SOAC.Input ts v $ identType v' + fixInputType inp = inp + +applyFusionRules :: Names -> [VName] -> SOAC -> Names -> FusedKer + -> TryFusion FusedKer +applyFusionRules unfus_nms outVars soac consumed ker = + tryOptimizeSOAC unfus_nms outVars soac consumed ker <|> + tryOptimizeKernel unfus_nms outVars soac consumed ker <|> + fuseSOACwithKer unfus_nms outVars soac consumed ker <|> + tryExposeInputs unfus_nms outVars soac consumed ker + +attemptFusion :: MonadFreshNames m => + Names -> [VName] -> SOAC -> Names -> FusedKer + -> m (Maybe FusedKer) +attemptFusion unfus_nms outVars soac consumed ker = + fmap removeUnusedParamsFromKer <$> + tryFusion (applyFusionRules unfus_nms outVars soac consumed ker) + (kernelScope ker) + +removeUnusedParamsFromKer :: FusedKer -> FusedKer +removeUnusedParamsFromKer ker = + case soac of SOAC.Screma {} -> ker { fsoac = soac' } + _ -> ker + where soac = fsoac ker + l = SOAC.lambda soac + inps = SOAC.inputs soac + (l', inps') = removeUnusedParams l inps + soac' = l' `SOAC.setLambda` + (inps' `SOAC.setInputs` soac) + +removeUnusedParams :: Lambda -> [SOAC.Input] -> (Lambda, [SOAC.Input]) +removeUnusedParams l inps = + (l { lambdaParams = ps' }, inps') + where pInps = zip (lambdaParams l) inps + (ps', inps') = case (unzip $ filter (used . fst) pInps, pInps) of + (([], []), (p,inp):_) -> ([p], [inp]) + ((ps_, inps_), _) -> (ps_, inps_) + used p = paramName p `S.member` freeVars + freeVars = freeInBody $ lambdaBody l + +-- | Check that the consumer uses at least one output of the producer +-- unmodified. +mapFusionOK :: [VName] -> FusedKer -> Bool +mapFusionOK outVars ker = any (`elem` inpIds) outVars + where inpIds = mapMaybe SOAC.isVarishInput (inputs ker) + +-- | Check that the consumer uses all the outputs of the producer unmodified. +mapWriteFusionOK :: [VName] -> FusedKer -> Bool +mapWriteFusionOK outVars ker = all (`elem` inpIds) outVars + where inpIds = mapMaybe SOAC.isVarishInput (inputs ker) + +-- | The brain of this module: Fusing a SOAC with a Kernel. +fuseSOACwithKer :: Names -> [VName] -> SOAC -> Names -> FusedKer + -> TryFusion FusedKer +fuseSOACwithKer unfus_set outVars soac_p soac_p_consumed ker = do + -- We are fusing soac_p into soac_c, i.e, the output of soac_p is going + -- into soac_c. + let soac_c = fsoac ker + inp_p_arr = SOAC.inputs soac_p + horizFuse= not (S.null unfus_set) && + SOAC.width soac_p == SOAC.width soac_c + inp_c_arr = SOAC.inputs soac_c + lam_p = SOAC.lambda soac_p + lam_c = SOAC.lambda soac_c + w = SOAC.width soac_p + returned_outvars = filter (`S.member` unfus_set) outVars + success res_outnms res_soac = do + let fusedVars_new = fusedVars ker++outVars + -- Avoid name duplication, because the producer lambda is not + -- removed from the program until much later. + uniq_lam <- renameLambda $ SOAC.lambda res_soac + return $ ker { fsoac = uniq_lam `SOAC.setLambda` res_soac + , fusedVars = fusedVars_new + , inplace = inplace ker <> soac_p_consumed + , fusedConsumed = fusedConsumed ker <> soac_p_consumed + , outNames = res_outnms + } + + outPairs <- forM (zip outVars $ map rowType $ SOAC.typeOf soac_p) $ \(outVar, t) -> do + outVar' <- newVName $ baseString outVar ++ "_elem" + return (outVar, Ident outVar' t) + + let mapLikeFusionCheck = + let (res_lam, new_inp) = fuseMaps unfus_set lam_p inp_p_arr outPairs lam_c inp_c_arr + (extra_nms,extra_rtps) = unzip $ filter ((`S.member` unfus_set) . fst) $ + zip outVars $ map (stripArray 1) $ SOAC.typeOf soac_p + res_lam' = res_lam { lambdaReturnType = lambdaReturnType res_lam ++ extra_rtps } + in (extra_nms, res_lam', new_inp) + + when (horizFuse && not (SOAC.nullTransforms $ outputTransform ker)) $ + fail "Horizontal fusion is invalid in the presence of output transforms." + + case (soac_c, soac_p) of + _ | SOAC.width soac_p /= SOAC.width soac_c -> fail "SOAC widths must match." + + (SOAC.Screma _ (ScremaForm (scan_lam_c, scan_nes_c) (comm_c, red_lam_c, red_nes_c) _) _, + SOAC.Screma _ (ScremaForm (scan_lam_p, scan_nes_p) (comm_p, red_lam_p, red_nes_p) _) _) + | mapFusionOK (drop (length $ scan_nes_p++red_nes_p) outVars) ker || horizFuse -> do + let (res_lam', new_inp) = fuseRedomap unfus_set outVars + lam_p scan_nes_p red_nes_p inp_p_arr + outPairs + lam_c scan_nes_c red_nes_c inp_c_arr + (soac_p_scanout, soac_p_redout, _soac_p_mapout) = + splitAt3 (length scan_nes_p) (length red_nes_p) outVars + (soac_c_scanout, soac_c_redout, soac_c_mapout) = + splitAt3 (length scan_nes_c) (length red_nes_c) $ outNames ker + unfus_arrs = returned_outvars \\ (soac_p_scanout++soac_p_redout) + scan_lam' = mergeReduceOps scan_lam_p scan_lam_c + red_lam' = mergeReduceOps red_lam_p red_lam_c + success (soac_p_scanout ++ soac_c_scanout ++ + soac_p_redout ++ soac_c_redout ++ + soac_c_mapout ++ unfus_arrs) $ + SOAC.Screma w (ScremaForm (scan_lam', scan_nes_p++scan_nes_c) + (comm_p<>comm_c, red_lam', red_nes_p++red_nes_c) + res_lam') + new_inp + + ------------------ + -- Scatter fusion -- + ------------------ + + -- Map-write fusion. + -- + -- The 'inplace' mechanism for kernels already takes care of + -- checking that the Scatter is not writing to any array used in + -- the Map. + (SOAC.Scatter _len _lam _ivs dests, + SOAC.Screma _ form _) + | isJust $ isMapSOAC form, + -- 1. all arrays produced by the map are ONLY used (consumed) + -- by the scatter, i.e., not used elsewhere. + not (any (`S.member` unfus_set) outVars), + -- 2. all arrays produced by the map are input to the scatter. + mapWriteFusionOK outVars ker -> do + let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck + success (outNames ker ++ extra_nms) $ + SOAC.Scatter w res_lam' new_inp dests + + -- Map-genreduce fusion. + -- + -- The 'inplace' mechanism for kernels already takes care of + -- checking that the GenReduce is not writing to any array used in + -- the Map. + (SOAC.GenReduce _ ops _ _, + SOAC.Screma _ form _) + | isJust $ isMapSOAC form, + -- 1. all arrays produced by the map are ONLY used (consumed) + -- by the genreduce, i.e., not used elsewhere. + not (any (`S.member` unfus_set) outVars), + -- 2. all arrays produced by the map are input to the scatter. + mapWriteFusionOK outVars ker -> do + let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck + success (outNames ker ++ extra_nms) $ + SOAC.GenReduce w ops res_lam' new_inp + + -- Genreduce-Genreduce fusion + (SOAC.GenReduce _ ops_c _ _, + SOAC.GenReduce _ ops_p _ _) + | horizFuse -> do + let p_num_buckets = length ops_p + c_num_buckets = length ops_c + (body_p, body_c) = (lambdaBody lam_p, lambdaBody lam_c) + body' = + Body { bodyAttr = bodyAttr body_p -- body_p and body_c have the same lores + , bodyStms = bodyStms body_p <> bodyStms body_c + , bodyResult = take c_num_buckets (bodyResult body_c) ++ + take p_num_buckets (bodyResult body_p) ++ + drop c_num_buckets (bodyResult body_c) ++ + drop p_num_buckets (bodyResult body_p) + } + lam' = + Lambda { lambdaParams = lambdaParams lam_c ++ lambdaParams lam_p + , lambdaBody = body' + , lambdaReturnType = replicate (c_num_buckets+p_num_buckets) (Prim int32) ++ + drop c_num_buckets (lambdaReturnType lam_c) ++ + drop p_num_buckets (lambdaReturnType lam_p) + } + success (outNames ker ++ returned_outvars) $ + SOAC.GenReduce w (ops_c <> ops_p) lam' (inp_c_arr <> inp_p_arr) + + -- Scatter-write fusion. + (SOAC.Scatter _len2 _lam_c ivs2 as2, + SOAC.Scatter _len_p _lam_p ivs_p as_p) + | horizFuse -> do + let zipW xs ys = ys_p ++ xs_p ++ ys2 ++ xs2 + where lenx = length xs `div` 2 + xs_p = take lenx xs + xs2 = drop lenx xs + leny = length ys `div` 2 + ys_p = take leny ys + ys2 = drop leny ys + let (body_p, body2) = (lambdaBody lam_p, lambdaBody lam_c) + let body' = Body { bodyAttr = bodyAttr body_p -- body_p and body2 have the same lores + , bodyStms = bodyStms body_p <> bodyStms body2 + , bodyResult = zipW (bodyResult body_p) (bodyResult body2) + } + let lam' = Lambda { lambdaParams = lambdaParams lam_p ++ lambdaParams lam_c + , lambdaBody = body' + , lambdaReturnType = zipW (lambdaReturnType lam_p) (lambdaReturnType lam_c) + } + success (outNames ker ++ returned_outvars) $ + SOAC.Scatter w lam' (ivs_p ++ ivs2) (as2 ++ as_p) + + (SOAC.Scatter {}, _) -> + fail "Cannot fuse a write with anything else than a write or a map" + (_, SOAC.Scatter {}) -> + fail "Cannot fuse a write with anything else than a write or a map" + + ---------------------------- + -- Stream-Stream Fusions: -- + ---------------------------- + (SOAC.Stream _ Sequential{} _ _, SOAC.Stream _ form_p@Sequential{} _ _) + | mapFusionOK (drop (length $ getStreamAccums form_p) outVars) ker || horizFuse -> do + -- fuse two SEQUENTIAL streams + (res_nms, res_stream) <- fuseStreamHelper (outNames ker) unfus_set outVars outPairs soac_c soac_p + success res_nms res_stream + + (SOAC.Stream _ Sequential{} _ _, SOAC.Stream _ Sequential{} _ _) -> + fail "Fusion conditions not met for two SEQ streams!" + + (SOAC.Stream _ Sequential{} _ _, SOAC.Stream{}) -> + fail "Cannot fuse a parallel with a sequential Stream!" + + (SOAC.Stream{}, SOAC.Stream _ Sequential{} _ _) -> + fail "Cannot fuse a parallel with a sequential Stream!" + + (SOAC.Stream{}, SOAC.Stream _ form_p _ _) + | mapFusionOK (drop (length $ getStreamAccums form_p) outVars) ker || horizFuse -> do + -- fuse two PARALLEL streams + (res_nms, res_stream) <- fuseStreamHelper (outNames ker) unfus_set outVars outPairs soac_c soac_p + success res_nms res_stream + + (SOAC.Stream{}, SOAC.Stream {}) -> + fail "Fusion conditions not met for two PAR streams!" + + ------------------------------------------------------------------- + --- If one is a stream, translate the other to a stream as well.--- + --- This does not get in trouble (infinite computation) because --- + --- scan's translation to Stream introduces a hindrance to --- + --- (horizontal fusion), hence repeated application is for the--- + --- moment impossible. However, if with a dependence-graph rep--- + --- we could run in an infinite recursion, i.e., repeatedly --- + --- fusing map o scan into an infinity of Stream levels! --- + ------------------------------------------------------------------- + (SOAC.Stream _ form2 _ _, _) -> do + -- If this rule is matched then soac_p is NOT a stream. + -- To fuse a stream kernel, we transform soac_p to a stream, which + -- borrows the sequential/parallel property of the soac_c Stream, + -- and recursively perform stream-stream fusion. + (soac_p', newacc_ids) <- SOAC.soacToStream soac_p + soac_p'' <- case form2 of + Sequential{} -> toSeqStream soac_p' + _ -> return soac_p' + if soac_p' == soac_p + then fail "SOAC could not be turned into stream." + else fuseSOACwithKer unfus_set (map identName newacc_ids++outVars) soac_p'' soac_p_consumed ker + + (_, SOAC.Screma _ form _) | Just _ <- Futhark.isScanSOAC form -> do + -- A Scan soac can be currently only fused as a (sequential) stream, + -- hence it is first translated to a (sequential) Stream and then + -- fusion with a kernel is attempted. + (soac_p', newacc_ids) <- SOAC.soacToStream soac_p + if soac_p' /= soac_p then + fuseSOACwithKer unfus_set (map identName newacc_ids++outVars) soac_p' soac_p_consumed ker + else fail "SOAC could not be turned into stream." + + (_, SOAC.Stream _ form_p _ _) -> do + -- If it reached this case then soac_c is NOT a Stream kernel, + -- hence transform the kernel's soac to a stream and attempt + -- stream-stream fusion recursivelly. + -- The newly created stream corresponding to soac_c borrows the + -- sequential/parallel property of the soac_p stream. + (soac_c', newacc_ids) <- SOAC.soacToStream soac_c + when (soac_c' == soac_c) $ fail "SOAC could not be turned into stream." + soac_c'' <- case form_p of + Sequential _ -> toSeqStream soac_c' + _ -> return soac_c' + + fuseSOACwithKer unfus_set outVars soac_p soac_p_consumed $ + ker { fsoac = soac_c'', outNames = map identName newacc_ids ++ outNames ker } + + --------------------------------- + --- DEFAULT, CANNOT FUSE CASE --- + --------------------------------- + _ -> fail "Cannot fuse" + +fuseStreamHelper :: [VName] -> Names -> [VName] -> [(VName,Ident)] + -> SOAC -> SOAC -> TryFusion ([VName], SOAC) +fuseStreamHelper out_kernms unfus_set outVars outPairs + (SOAC.Stream w2 form2 lam2 inp2_arr) + (SOAC.Stream _ form1 lam1 inp1_arr) = + if getStreamOrder form2 /= getStreamOrder form1 + then fail "fusion conditions not met!" + else do -- very similar to redomap o redomap composition, but need + -- to remove first the `chunk' parameters of streams' + -- lambdas and put them in the resulting stream lambda. + let nes1 = getStreamAccums form1 + chunk1 = head $ lambdaParams lam1 + chunk2 = head $ lambdaParams lam2 + hmnms = M.fromList [(paramName chunk2, paramName chunk1)] + lam20 = substituteNames hmnms lam2 + lam1' = lam1 { lambdaParams = tail $ lambdaParams lam1 } + lam2' = lam20 { lambdaParams = tail $ lambdaParams lam20 } + (res_lam', new_inp) = fuseRedomap unfus_set outVars + lam1' [] nes1 + inp1_arr outPairs + lam2' [] (getStreamAccums form2) + inp2_arr + res_lam'' = res_lam' { lambdaParams = chunk1 : lambdaParams res_lam' } + unfus_accs = take (length nes1) outVars + unfus_arrs = filter (`S.member` unfus_set) outVars + res_form <- mergeForms form2 form1 + return (unfus_accs ++ out_kernms ++ unfus_arrs, + SOAC.Stream w2 res_form res_lam'' new_inp ) + where mergeForms (Sequential acc2) (Sequential acc1) = return $ Sequential (acc1++acc2) + mergeForms (Parallel _ comm2 lam2r acc2) (Parallel o1 comm1 lam1r acc1) = + return $ Parallel o1 (comm1<>comm2) (mergeReduceOps lam1r lam2r) (acc1++acc2) + mergeForms _ _ = fail "Fusing sequential to parallel stream disallowed!" +fuseStreamHelper _ _ _ _ _ _ = fail "Cannot Fuse Streams!" + +-- | If a Stream is passed as argument then it converts it to a +-- Sequential Stream; Otherwise it FAILS! +toSeqStream :: SOAC -> TryFusion SOAC +toSeqStream s@(SOAC.Stream _ (Sequential _) _ _) = return s +toSeqStream (SOAC.Stream w (Parallel _ _ _ acc) l inps) = + return $ SOAC.Stream w (Sequential acc) l inps +toSeqStream _ = fail "toSeqStream expects a stream, but given a SOAC." + +-- Here follows optimizations and transforms to expose fusability. + +optimizeKernel :: Maybe [VName] -> FusedKer -> TryFusion FusedKer +optimizeKernel inp ker = do + (soac, resTrans) <- optimizeSOAC inp (fsoac ker) startTrans + return $ ker { fsoac = soac + , outputTransform = resTrans + } + where startTrans = outputTransform ker + +optimizeSOAC :: Maybe [VName] -> SOAC -> SOAC.ArrayTransforms + -> TryFusion (SOAC, SOAC.ArrayTransforms) +optimizeSOAC inp soac os = do + res <- foldM comb (False, soac, os) optimizations + case res of + (False, _, _) -> fail "No optimisation applied" + (True, soac', os') -> return (soac', os') + where comb (changed, soac', os') f = do + (soac'', os'') <- f inp soac' os + return (True, soac'', os'') + <|> return (changed, soac', os') + +type Optimization = Maybe [VName] + -> SOAC + -> SOAC.ArrayTransforms + -> TryFusion (SOAC, SOAC.ArrayTransforms) + +optimizations :: [Optimization] +optimizations = [iswim] + +iswim :: Maybe [VName] -> SOAC -> SOAC.ArrayTransforms + -> TryFusion (SOAC, SOAC.ArrayTransforms) +iswim _ (SOAC.Screma w form arrs) ots + | Just (scan_fun, nes) <- Futhark.isScanSOAC form, + Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible scan_fun, + Just nes_names <- mapM subExpVar nes = do + + let nes_idents = zipWith Ident nes_names $ lambdaReturnType scan_fun + map_nes = map SOAC.identInput nes_idents + map_arrs' = map_nes ++ map (SOAC.transposeInput 0 1) arrs + (scan_acc_params, scan_elem_params) = + splitAt (length arrs) $ lambdaParams scan_fun + map_params = map removeParamOuterDim scan_acc_params ++ + map (setParamOuterDimTo w) scan_elem_params + map_rettype = map (`setOuterSize` w) $ lambdaReturnType scan_fun + + scan_params = lambdaParams map_fun + scan_body = lambdaBody map_fun + scan_rettype = lambdaReturnType map_fun + scan_fun' = Lambda scan_params scan_body scan_rettype + nes' = map Var $ take (length map_nes) $ map paramName map_params + arrs' = drop (length map_nes) $ map paramName map_params + + id_map_lam <- mkIdentityLambda $ lambdaReturnType scan_fun' + + let map_body = mkBody (oneStm $ + Let (setPatternOuterDimTo w map_pat) (defAux ()) $ + Op $ Futhark.Screma w (ScremaForm (scan_fun', nes') + (mempty, nilFn, mempty) + id_map_lam) arrs') $ + map Var $ patternNames map_pat + map_fun' = Lambda map_params map_body map_rettype + perm = case lambdaReturnType map_fun of + [] -> [] + t:_ -> 1 : 0 : [2..arrayRank t] + + return (SOAC.Screma map_w + (ScremaForm (nilFn, mempty) (mempty, nilFn, mempty) map_fun') + map_arrs', + ots SOAC.|> SOAC.Rearrange map_cs perm) + +iswim _ _ _ = + fail "ISWIM does not apply." + +removeParamOuterDim :: LParam -> LParam +removeParamOuterDim param = + let t = rowType $ paramType param + in param { paramAttr = t } + +setParamOuterDimTo :: SubExp -> LParam -> LParam +setParamOuterDimTo w param = + let t = paramType param `setOuterSize` w + in param { paramAttr = t } + +setPatternOuterDimTo :: SubExp -> Pattern -> Pattern +setPatternOuterDimTo w = fmap (`setOuterSize` w) + +-- Now for fiddling with transpositions... + +commonTransforms :: [VName] -> [SOAC.Input] + -> (SOAC.ArrayTransforms, [SOAC.Input]) +commonTransforms interesting inps = commonTransforms' inps' + where inps' = [ (SOAC.inputArray inp `elem` interesting, inp) + | inp <- inps ] + +commonTransforms' :: [(Bool, SOAC.Input)] -> (SOAC.ArrayTransforms, [SOAC.Input]) +commonTransforms' inps = + case foldM inspect (Nothing, []) inps of + Just (Just mot, inps') -> first (mot SOAC.<|) $ commonTransforms' $ reverse inps' + _ -> (SOAC.noTransforms, map snd inps) + where inspect (mot, prev) (True, inp) = + case (mot, inputToOutput inp) of + (Nothing, Just (ot, inp')) -> Just (Just ot, (True, inp') : prev) + (Just ot1, Just (ot2, inp')) + | ot1 == ot2 -> Just (Just ot2, (True, inp') : prev) + _ -> Nothing + inspect (mot, prev) inp = Just (mot,inp:prev) + +mapDepth :: MapNest -> Int +mapDepth (MapNest.MapNest _ lam levels _) = + min resDims (length levels) + 1 + where resDims = minDim $ case levels of + [] -> lambdaReturnType lam + nest:_ -> MapNest.nestingReturnType nest + minDim [] = 0 + minDim (t:ts) = foldl min (arrayRank t) $ map arrayRank ts + +pullRearrange :: SOAC -> SOAC.ArrayTransforms + -> TryFusion (SOAC, SOAC.ArrayTransforms) +pullRearrange soac ots = do + nest <- liftMaybe =<< MapNest.fromSOAC soac + SOAC.Rearrange cs perm SOAC.:< ots' <- return $ SOAC.viewf ots + if rearrangeReach perm <= mapDepth nest then do + let -- Expand perm to cover the full extent of the input dimensionality + perm' inp = take r perm ++ [length perm..r-1] + where r = SOAC.inputRank inp + addPerm inp = SOAC.addTransform (SOAC.Rearrange cs $ perm' inp) inp + inputs' = map addPerm $ MapNest.inputs nest + soac' <- MapNest.toSOAC $ + inputs' `MapNest.setInputs` rearrangeReturnTypes nest perm + return (soac', ots') + else fail "Cannot pull transpose" + +pushRearrange :: [VName] -> SOAC -> SOAC.ArrayTransforms + -> TryFusion (SOAC, SOAC.ArrayTransforms) +pushRearrange inpIds soac ots = do + nest <- liftMaybe =<< MapNest.fromSOAC soac + (perm, inputs') <- liftMaybe $ fixupInputs inpIds $ MapNest.inputs nest + if rearrangeReach perm <= mapDepth nest then do + let invertRearrange = SOAC.Rearrange mempty $ rearrangeInverse perm + soac' <- MapNest.toSOAC $ + inputs' `MapNest.setInputs` + rearrangeReturnTypes nest perm + return (soac', invertRearrange SOAC.<| ots) + else fail "Cannot push transpose" + +-- | Actually also rearranges indices. +rearrangeReturnTypes :: MapNest -> [Int] -> MapNest +rearrangeReturnTypes nest@(MapNest.MapNest w body nestings inps) perm = + MapNest.MapNest w + body + (zipWith setReturnType + nestings $ + drop 1 $ iterate (map rowType) ts) + inps + where origts = MapNest.typeOf nest + -- The permutation may be deeper than the rank of the type, + -- but it is required that it is an identity permutation + -- beyond that. This is supposed to be checked as an + -- invariant by whoever calls rearrangeReturnTypes. + rearrangeType' t = rearrangeType (take (arrayRank t) perm) t + ts = map rearrangeType' origts + + setReturnType nesting t' = + nesting { MapNest.nestingReturnType = t' } + +fixupInputs :: [VName] -> [SOAC.Input] -> Maybe ([Int], [SOAC.Input]) +fixupInputs inpIds inps = + case mapMaybe inputRearrange $ filter exposable inps of + perm:_ -> do inps' <- mapM (fixupInput (rearrangeReach perm) perm) inps + return (perm, inps') + _ -> Nothing + where exposable = (`elem` inpIds) . SOAC.inputArray + + inputRearrange (SOAC.Input ts _ _) + | _ SOAC.:> SOAC.Rearrange _ perm <- SOAC.viewl ts = Just perm + inputRearrange _ = Nothing + + fixupInput d perm inp + | r <- SOAC.inputRank inp, + r >= d = + Just $ SOAC.addTransform (SOAC.Rearrange mempty $ take r perm) inp + | otherwise = Nothing + +pullReshape :: SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms) +pullReshape (SOAC.Screma _ form inps) ots + | Just maplam <- Futhark.isMapSOAC form, + SOAC.Reshape cs shape SOAC.:< ots' <- SOAC.viewf ots, + all primType $ lambdaReturnType maplam = do + let mapw' = case reverse $ newDims shape of + [] -> intConst Int32 0 + d:_ -> d + inputs' = map (SOAC.addTransform $ SOAC.ReshapeOuter cs shape) inps + inputTypes = map SOAC.inputType inputs' + + let outersoac :: ([SOAC.Input] -> SOAC) -> (SubExp, [SubExp]) + -> TryFusion ([SOAC.Input] -> SOAC) + outersoac inner (w, outershape) = do + let addDims t = arrayOf t (Shape outershape) NoUniqueness + retTypes = map addDims $ lambdaReturnType maplam + + ps <- forM inputTypes $ \inpt -> + newParam "pullReshape_param" $ + stripArray (length shape-length outershape) inpt + + inner_body <- runBodyBinder $ + eBody [SOAC.toExp $ inner $ map (SOAC.identInput . paramIdent) ps] + let inner_fun = Lambda { lambdaParams = ps + , lambdaReturnType = retTypes + , lambdaBody = inner_body + } + return $ SOAC.Screma w $ Futhark.mapSOAC inner_fun + + op' <- foldM outersoac (SOAC.Screma mapw' $ Futhark.mapSOAC maplam) $ + zip (drop 1 $ reverse $ newDims shape) $ + drop 1 $ reverse $ drop 1 $ tails $ newDims shape + return (op' inputs', ots') +pullReshape _ _ = fail "Cannot pull reshape" + +-- We can make a Replicate output-transform part of a map SOAC simply +-- by adding another dimension to the SOAC. +pullReplicate :: SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms) +pullReplicate soac@(SOAC.Screma _ form _) ots + | Just _ <- isMapSOAC form, + SOAC.Replicate cs (Shape [n]) SOAC.:< ots' <- SOAC.viewf ots = do + let rettype = SOAC.typeOf soac + body <- runBodyBinder $ do + names <- certifying cs $ + letTupExp "pull_replicate" =<< SOAC.toExp soac + resultBodyM $ map Var names + let lam = Lambda { lambdaReturnType = rettype + , lambdaBody = body + , lambdaParams = [] + } + return (SOAC.Screma n (Futhark.mapSOAC lam) [], ots') +pullReplicate _ _ = fail "Cannot pull replicate" + +-- Tie it all together in exposeInputs (for making inputs to a +-- consumer available) and pullOutputTransforms (for moving +-- output-transforms of a producer to its inputs instead). + +exposeInputs :: [VName] -> FusedKer + -> TryFusion (FusedKer, SOAC.ArrayTransforms) +exposeInputs inpIds ker = + (exposeInputs' =<< pushRearrange') <|> + (exposeInputs' =<< pullRearrange') <|> + exposeInputs' ker + where ot = outputTransform ker + + pushRearrange' = do + (soac', ot') <- pushRearrange inpIds (fsoac ker) ot + return ker { fsoac = soac' + , outputTransform = ot' + } + + pullRearrange' = do + (soac',ot') <- pullRearrange (fsoac ker) ot + unless (SOAC.nullTransforms ot') $ + fail "pullRearrange was not enough" + return ker { fsoac = soac' + , outputTransform = SOAC.noTransforms + } + + exposeInputs' ker' = + case commonTransforms inpIds $ inputs ker' of + (ot', inps') | all exposed inps' -> + return (ker' { fsoac = inps' `SOAC.setInputs` fsoac ker'}, ot') + _ -> fail "Cannot expose" + + exposed (SOAC.Input ts _ _) + | SOAC.nullTransforms ts = True + exposed inp = SOAC.inputArray inp `notElem` inpIds + +outputTransformPullers :: [SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms)] +outputTransformPullers = [pullRearrange, pullReshape, pullReplicate] + +pullOutputTransforms :: SOAC -> SOAC.ArrayTransforms + -> TryFusion (SOAC, SOAC.ArrayTransforms) +pullOutputTransforms = attempt outputTransformPullers + where attempt [] _ _ = fail "Cannot pull anything" + attempt (p:ps) soac ots = do + (soac',ots') <- p soac ots + if SOAC.nullTransforms ots' then return (soac', SOAC.noTransforms) + else pullOutputTransforms soac' ots' <|> return (soac', ots') + <|> attempt ps soac ots diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -0,0 +1,34 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Futhark.Optimise.Fusion.TryFusion + ( TryFusion + , tryFusion + , liftMaybe + ) + where + +import Control.Applicative +import Control.Monad.State +import Control.Monad.Reader +import qualified Control.Monad.Fail as Fail + +import Futhark.Representation.SOACS +import Futhark.MonadFreshNames + +newtype TryFusion a = TryFusion (ReaderT (Scope SOACS) + (StateT VNameSource Maybe) + a) + deriving (Functor, Applicative, Alternative, Monad, Fail.MonadFail, + MonadFreshNames, + HasScope SOACS, + LocalScope SOACS) + +tryFusion :: MonadFreshNames m => + TryFusion a -> Scope SOACS -> m (Maybe a) +tryFusion (TryFusion m) types = modifyNameSource $ \src -> + case runStateT (runReaderT m types) src of + Just (x, src') -> (Just x, src') + Nothing -> (Nothing, src) + +liftMaybe :: Maybe a -> TryFusion a +liftMaybe Nothing = fail "Nothing" +liftMaybe (Just x) = return x diff --git a/src/Futhark/Optimise/InPlaceLowering.hs b/src/Futhark/Optimise/InPlaceLowering.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/InPlaceLowering.hs @@ -0,0 +1,335 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE UndecidableInstances #-} +-- | This module implements an optimisation that moves in-place +-- updates into/before loops where possible, with the end goal of +-- minimising memory copies. As an example, consider this program: +-- +-- @ +-- loop (r = r0) = for i < n do +-- let a = r[i] in +-- let r[i] = a * i in +-- r +-- in +-- ... +-- let x = y with [k] <- r in +-- ... +-- @ +-- +-- We want to turn this into the following: +-- +-- @ +-- let x0 = y with [k] <- r0 +-- loop (x = x0) = for i < n do +-- let a = a[k,i] in +-- let x[k,i] = a * i in +-- x +-- in +-- let r = x[y] in +-- ... +-- @ +-- +-- The intent is that we are also going to optimise the new data +-- movement (in the @x0@-binding), possibly by changing how @r0@ is +-- defined. For the above transformation to be valid, a number of +-- conditions must be fulfilled: +-- +-- (1) @r@ must not be consumed after the original in-place update. +-- +-- (2) @k@ and @y@ must be available at the beginning of the loop. +-- +-- (3) @x@ must be visible whenever @r@ is visible. (This means +-- that both @x@ and @r@ must be bound in the same 'Body'.) +-- +-- (4) If @x@ is consumed at a point after the loop, @r@ must not +-- be used after that point. +-- +-- (5) The size of @r@ is invariant inside the loop. +-- +-- (6) The value @r@ must come from something that we can actually +-- optimise (e.g. not a function parameter). +-- +-- (7) @y@ (or its aliases) may not be used inside the body of the +-- loop. +-- +-- FIXME: the implementation is not finished yet. Specifically, the +-- above conditions are not really checked. +module Futhark.Optimise.InPlaceLowering + ( + inPlaceLowering + ) where + +import Control.Monad.RWS +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Data.Semigroup as Sem + +import Futhark.Analysis.Alias +import Futhark.Representation.Aliases +import Futhark.Representation.Kernels +import Futhark.Optimise.InPlaceLowering.LowerIntoStm +import Futhark.MonadFreshNames +import Futhark.Binder +import Futhark.Pass +import Futhark.Tools (fullSlice) + +-- | Apply the in-place lowering optimisation to the given program. +inPlaceLowering :: Pass Kernels Kernels +inPlaceLowering = + Pass "In-place lowering" "Lower in-place updates into loops" $ + fmap removeProgAliases . + intraproceduralTransformation optimiseFunDef . + aliasAnalysis + +optimiseFunDef :: MonadFreshNames m => FunDef (Aliases Kernels) + -> m (FunDef (Aliases Kernels)) +optimiseFunDef fundec = + modifyNameSource $ runForwardingM lowerUpdateKernels onKernelOp $ + bindingFParams (funDefParams fundec) $ do + body <- optimiseBody $ funDefBody fundec + return $ fundec { funDefBody = body } + +type Constraints lore = (Bindable lore, CanBeAliased (Op lore)) + +optimiseBody :: Constraints lore => + Body (Aliases lore) -> ForwardingM lore (Body (Aliases lore)) +optimiseBody (Body als bnds res) = do + bnds' <- deepen $ optimiseStms (stmsToList bnds) $ + mapM_ seen res + return $ Body als (stmsFromList bnds') res + where seen Constant{} = return () + seen (Var v) = seenVar v + +optimiseStms :: Constraints lore => + [Stm (Aliases lore)] -> ForwardingM lore () + -> ForwardingM lore [Stm (Aliases lore)] +optimiseStms [] m = m >> return [] + +optimiseStms (bnd:bnds) m = do + (bnds', bup) <- tapBottomUp $ bindingStm bnd $ optimiseStms bnds m + bnd' <- optimiseInStm bnd + case filter ((`elem` boundHere) . updateValue) $ + forwardThese bup of + [] -> checkIfForwardableUpdate bnd' bnds' + updates -> do + let updateStms = map updateStm updates + lower <- asks lowerUpdate + -- Condition (5) and (7) are assumed to be checked by + -- lowerUpdate. + case lower bnd' updates of + Just lowering -> do new_bnds <- lowering + new_bnds' <- optimiseStms new_bnds $ + tell bup { forwardThese = [] } + return $ new_bnds' ++ bnds' + Nothing -> checkIfForwardableUpdate bnd' $ + updateStms ++ bnds' + + where boundHere = patternNames $ stmPattern bnd + + checkIfForwardableUpdate bnd'@(Let (Pattern [] [PatElem v attr]) + (StmAux cs _) e) bnds' + | BasicOp (Update src (DimFix i:slice) (Var ve)) <- e, + slice == drop 1 (fullSlice (typeOf attr) [DimFix i]) = do + forwarded <- maybeForward ve v attr cs src i + return $ if forwarded + then bnds' + else bnd' : bnds' + checkIfForwardableUpdate bnd' bnds' = + return $ bnd' : bnds' + +optimiseInStm :: Constraints lore => Stm (Aliases lore) -> ForwardingM lore (Stm (Aliases lore)) +optimiseInStm (Let pat attr e) = + Let pat attr <$> optimiseExp e + +optimiseExp :: Constraints lore => Exp (Aliases lore) -> ForwardingM lore (Exp (Aliases lore)) +optimiseExp (DoLoop ctx val form body) = + bindingScope (scopeOf form) $ + bindingFParams (map fst $ ctx ++ val) $ + DoLoop ctx val form <$> optimiseBody body +optimiseExp (Op op) = do + f <- asks onOp + Op <$> f op +optimiseExp e = mapExpM optimise e + where optimise = identityMapper { mapOnBody = const optimiseBody + } +onKernelOp :: OnOp Kernels +onKernelOp (Kernel debug kspace ts kbody) = do + old_scope <- askScope + modifyNameSource $ runForwardingM lowerUpdateInKernel onKernelExp $ + bindingScope (castScope old_scope <> scopeOfKernelSpace kspace) $ do + stms <- deepen $ optimiseStms (stmsToList (kernelBodyStms kbody)) $ + mapM_ seenVar $ freeIn $ kernelBodyResult kbody + return $ Kernel debug kspace ts $ kbody { kernelBodyStms = stmsFromList stms } +onKernelOp op = return op + +onKernelExp :: OnOp InKernel +onKernelExp (GroupStream w maxchunk lam accs arrs) = do + lam_body <- bindingScope (scopeOf lam) $ + optimiseBody $ groupStreamLambdaBody lam + let lam' = lam { groupStreamLambdaBody = lam_body } + return $ GroupStream w maxchunk lam' accs arrs +onKernelExp op = return op + +data Entry lore = Entry { entryNumber :: Int + , entryAliases :: Names + , entryDepth :: Int + , entryOptimisable :: Bool + , entryType :: NameInfo (Aliases lore) + } + +type VTable lore = M.Map VName (Entry lore) + +type OnOp lore = Op (Aliases lore) -> ForwardingM lore (Op (Aliases lore)) + +data TopDown lore = TopDown { topDownCounter :: Int + , topDownTable :: VTable lore + , topDownDepth :: Int + , lowerUpdate :: LowerUpdate lore (ForwardingM lore) + , onOp :: OnOp lore + } + +data BottomUp lore = BottomUp { bottomUpSeen :: Names + , forwardThese :: [DesiredUpdate (LetAttr (Aliases lore))] + } + +instance Sem.Semigroup (BottomUp lore) where + BottomUp seen1 forward1 <> BottomUp seen2 forward2 = + BottomUp (seen1 <> seen2) (forward1 <> forward2) + +instance Monoid (BottomUp lore) where + mempty = BottomUp mempty mempty + mappend = (Sem.<>) + +updateStm :: Constraints lore => DesiredUpdate (LetAttr (Aliases lore)) -> Stm (Aliases lore) +updateStm fwd = + mkLet [] [Ident (updateName fwd) $ typeOf $ updateType fwd] $ + BasicOp $ Update (updateSource fwd) + (fullSlice (typeOf $ updateType fwd) $ updateIndices fwd) $ + Var $ updateValue fwd + +newtype ForwardingM lore a = ForwardingM (RWS (TopDown lore) (BottomUp lore) VNameSource a) + deriving (Monad, Applicative, Functor, + MonadReader (TopDown lore), + MonadWriter (BottomUp lore), + MonadState VNameSource) + +instance MonadFreshNames (ForwardingM lore) where + getNameSource = get + putNameSource = put + +instance Constraints lore => HasScope (Aliases lore) (ForwardingM lore) where + askScope = M.map entryType <$> asks topDownTable + +runForwardingM :: LowerUpdate lore (ForwardingM lore) -> OnOp lore -> ForwardingM lore a + -> VNameSource -> (a, VNameSource) +runForwardingM f g (ForwardingM m) src = let (x, src', _) = runRWS m emptyTopDown src + in (x, src') + where emptyTopDown = TopDown { topDownCounter = 0 + , topDownTable = M.empty + , topDownDepth = 0 + , lowerUpdate = f + , onOp = g + } + +bindingParams :: (attr -> NameInfo (Aliases lore)) + -> [Param attr] + -> ForwardingM lore a + -> ForwardingM lore a +bindingParams f params = local $ \(TopDown n vtable d x y) -> + let entry fparam = + (paramName fparam, + Entry n mempty d False $ f $ paramAttr fparam) + entries = M.fromList $ map entry params + in TopDown (n+1) (M.union entries vtable) d x y + +bindingFParams :: [FParam (Aliases lore)] + -> ForwardingM lore a + -> ForwardingM lore a +bindingFParams = bindingParams FParamInfo + +bindingScope :: Scope (Aliases lore) + -> ForwardingM lore a + -> ForwardingM lore a +bindingScope scope = local $ \(TopDown n vtable d x y) -> + let entries = M.map entry scope + infoAliases (LetInfo (aliases, _)) = unNames aliases + infoAliases _ = mempty + entry info = Entry n (infoAliases info) d False info + in TopDown (n+1) (entries<>vtable) d x y + +bindingStm :: Stm (Aliases lore) + -> ForwardingM lore a + -> ForwardingM lore a +bindingStm (Let pat _ _) = local $ \(TopDown n vtable d x y) -> + let entries = M.fromList $ map entry $ patternElements pat + entry patElem = + let (aliases, _) = patElemAttr patElem + in (patElemName patElem, + Entry n (unNames aliases) d True $ LetInfo $ patElemAttr patElem) + in TopDown (n+1) (M.union entries vtable) d x y + +bindingNumber :: VName -> ForwardingM lore Int +bindingNumber name = do + res <- asks $ fmap entryNumber . M.lookup name . topDownTable + case res of Just n -> return n + Nothing -> fail $ "bindingNumber: variable " ++ + pretty name ++ " not found." + +deepen :: ForwardingM lore a -> ForwardingM lore a +deepen = local $ \env -> env { topDownDepth = topDownDepth env + 1 } + +areAvailableBefore :: [SubExp] -> VName -> ForwardingM lore Bool +areAvailableBefore ses point = do + pointN <- bindingNumber point + nameNs <- mapM bindingNumber $ subExpVars ses + return $ all (< pointN) nameNs + +isInCurrentBody :: VName -> ForwardingM lore Bool +isInCurrentBody name = do + current <- asks topDownDepth + res <- asks $ fmap entryDepth . M.lookup name . topDownTable + case res of Just d -> return $ d == current + Nothing -> fail $ "isInCurrentBody: variable " ++ + pretty name ++ " not found." + +isOptimisable :: VName -> ForwardingM lore Bool +isOptimisable name = do + res <- asks $ fmap entryOptimisable . M.lookup name . topDownTable + case res of Just b -> return b + Nothing -> fail $ "isOptimisable: variable " ++ + pretty name ++ " not found." + +seenVar :: VName -> ForwardingM lore () +seenVar name = do + aliases <- asks $ + maybe mempty entryAliases . + M.lookup name . topDownTable + tell $ mempty { bottomUpSeen = S.insert name aliases } + +tapBottomUp :: ForwardingM lore a -> ForwardingM lore (a, BottomUp lore) +tapBottomUp m = do (x,bup) <- listen m + return (x, bup) + +maybeForward :: Constraints lore => + VName + -> VName -> LetAttr (Aliases lore) -> Certificates -> VName -> SubExp + -> ForwardingM lore Bool +maybeForward v dest_nm dest_attr cs src i = do + -- Checks condition (2) + available <- [i,Var src] `areAvailableBefore` v + -- ...subcondition, the certificates must also. + certs_available <- map Var (S.toList $ freeIn cs) `areAvailableBefore` v + -- Check condition (3) + samebody <- isInCurrentBody v + -- Check condition (6) + optimisable <- isOptimisable v + not_prim <- not . primType <$> lookupType v + if available && certs_available && samebody && optimisable && not_prim then do + let fwd = DesiredUpdate dest_nm dest_attr cs src [DimFix i] v + tell mempty { forwardThese = [fwd] } + return True + else return False diff --git a/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/InPlaceLowering/LowerIntoStm.hs @@ -0,0 +1,251 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +module Futhark.Optimise.InPlaceLowering.LowerIntoStm + ( + lowerUpdateInKernel + , lowerUpdateKernels + , LowerUpdate + , DesiredUpdate (..) + ) where + +import Control.Monad +import Control.Monad.Writer +import Data.List (find) +import Data.Maybe (mapMaybe) +import Data.Either +import qualified Data.Set as S + +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.Aliases +import Futhark.Representation.Kernels +import Futhark.Construct +import Futhark.Optimise.InPlaceLowering.SubstituteIndices +import Futhark.Tools (fullSlice) + +data DesiredUpdate attr = + DesiredUpdate { updateName :: VName -- ^ Name of result. + , updateType :: attr -- ^ Type of result. + , updateCertificates :: Certificates + , updateSource :: VName + , updateIndices :: Slice SubExp + , updateValue :: VName + } + deriving (Show) + +instance Functor DesiredUpdate where + f `fmap` u = u { updateType = f $ updateType u } + +updateHasValue :: VName -> DesiredUpdate attr -> Bool +updateHasValue name = (name==) . updateValue + +type LowerUpdate lore m = Stm (Aliases lore) + -> [DesiredUpdate (LetAttr (Aliases lore))] + -> Maybe (m [Stm (Aliases lore)]) + +lowerUpdate :: (MonadFreshNames m, Bindable lore, + LetAttr lore ~ Type, CanBeAliased (Op lore)) => LowerUpdate lore m +lowerUpdate (Let pat aux (DoLoop ctx val form body)) updates = do + canDo <- lowerUpdateIntoLoop updates pat ctx val body + Just $ do + (prebnds, postbnds, ctxpat, valpat, ctx', val', body') <- canDo + return $ + prebnds ++ [certify (stmAuxCerts aux) $ + mkLet ctxpat valpat $ DoLoop ctx' val' form body'] ++ postbnds +lowerUpdate + (Let pat aux (BasicOp (SubExp (Var v)))) + [DesiredUpdate bindee_nm bindee_attr cs src is val] + | patternNames pat == [src] = + let is' = fullSlice (typeOf bindee_attr) is + in Just $ + return [certify (stmAuxCerts aux <> cs) $ + mkLet [] [Ident bindee_nm $ typeOf bindee_attr] $ + BasicOp $ Update v is' $ Var val] +lowerUpdate _ _ = + Nothing + +lowerUpdateKernels :: MonadFreshNames m => LowerUpdate Kernels m +lowerUpdateKernels + (Let (Pattern [] [PatElem v v_attr]) aux (Op (Kernel debug kspace ts kbody))) + [update@(DesiredUpdate bindee_nm bindee_attr cs _src is val)] + | v == val = do + kbody' <- lowerUpdateIntoKernel update kspace kbody + let is' = fullSlice (typeOf bindee_attr) is + Just $ return [certify (stmAuxCerts aux <> cs) $ + mkLet [] [Ident bindee_nm $ typeOf bindee_attr] $ + Op $ Kernel debug kspace ts kbody', + mkLet [] [Ident v $ typeOf v_attr] $ BasicOp $ Index bindee_nm is'] +lowerUpdateKernels stm updates = lowerUpdate stm updates + +lowerUpdateInKernel :: MonadFreshNames m => LowerUpdate InKernel m +lowerUpdateInKernel = lowerUpdate + +lowerUpdateIntoKernel :: DesiredUpdate (LetAttr (Aliases Kernels)) + -> KernelSpace -> KernelBody (Aliases InKernel) + -> Maybe (KernelBody (Aliases InKernel)) +lowerUpdateIntoKernel update kspace kbody = do + [ThreadsReturn ThreadsInSpace se] <- Just $ kernelBodyResult kbody + is' <- mapM dimFix is + let ret = WriteReturn (arrayDims $ snd bindee_attr) src [(is'++map Var gtids, se)] + return kbody { kernelBodyResult = [ret] } + where DesiredUpdate _bindee_nm bindee_attr _cs src is _val = update + gtids = map fst $ spaceDimensions kspace + +lowerUpdateIntoLoop :: (Bindable lore, BinderOps lore, + Aliased lore, LetAttr lore ~ (als, Type), + MonadFreshNames m) => + [DesiredUpdate (LetAttr lore)] + -> Pattern lore + -> [(FParam lore, SubExp)] + -> [(FParam lore, SubExp)] + -> Body lore + -> Maybe (m ([Stm lore], + [Stm lore], + [Ident], + [Ident], + [(FParam lore, SubExp)], + [(FParam lore, SubExp)], + Body lore)) +lowerUpdateIntoLoop updates pat ctx val body = do + -- Algorithm: + -- + -- 0) Map each result of the loop body to a corresponding in-place + -- update, if one exists. + -- + -- 1) Create new merge variables corresponding to the arrays being + -- updated; extend the pattern and the @res@ list with these, + -- and remove the parts of the result list that have a + -- corresponding in-place update. + -- + -- (The creation of the new merge variable identifiers is + -- actually done at the same time as step (0)). + -- + -- 2) Create in-place updates at the end of the loop body. + -- + -- 3) Create index expressions that read back the values written + -- in (2). If the merge parameter corresponding to this value + -- is unique, also @copy@ this value. + -- + -- 4) Update the result of the loop body to properly pass the new + -- arrays and indexed elements to the next iteration of the + -- loop. + -- + -- We also check that the merge parameters we work with have + -- loop-invariant shapes. + mk_in_place_map <- summariseLoop updates usedInBody resmap val + Just $ do + in_place_map <- mk_in_place_map + (val',prebnds,postbnds) <- mkMerges in_place_map + let (ctxpat,valpat) = mkResAndPat in_place_map + idxsubsts = indexSubstitutions in_place_map + (idxsubsts', newbnds) <- substituteIndices idxsubsts $ bodyStms body + (body_res, res_bnds) <- manipulateResult in_place_map idxsubsts' + let body' = mkBody (newbnds<>res_bnds) body_res + return (prebnds, postbnds, ctxpat, valpat, ctx, val', body') + where usedInBody = freeInBody body + resmap = zip (bodyResult body) $ patternValueIdents pat + + mkMerges :: (MonadFreshNames m, Bindable lore) => + [LoopResultSummary (als, Type)] + -> m ([(Param DeclType, SubExp)], [Stm lore], [Stm lore]) + mkMerges summaries = do + ((origmerge, extramerge), (prebnds, postbnds)) <- + runWriterT $ partitionEithers <$> mapM mkMerge summaries + return (origmerge ++ extramerge, prebnds, postbnds) + + mkMerge summary + | Just (update, mergename, mergeattr) <- relatedUpdate summary = do + source <- newVName "modified_source" + let source_t = snd $ updateType update + elmident = Ident (updateValue update) $ rowType source_t + tell ([mkLet [] [Ident source source_t] $ BasicOp $ Update + (updateSource update) + (fullSlice source_t $ updateIndices update) $ + snd $ mergeParam summary], + [mkLet [] [elmident] $ BasicOp $ Index + (updateName update) (fullSlice (typeOf $ updateType update) $ updateIndices update)]) + return $ Right (Param + mergename + (toDecl (typeOf mergeattr) Unique), + Var source) + | otherwise = return $ Left $ mergeParam summary + + mkResAndPat summaries = + let (origpat,extrapat) = partitionEithers $ map mkResAndPat' summaries + in (patternContextIdents pat, + origpat ++ extrapat) + + mkResAndPat' summary + | Just (update, _, _) <- relatedUpdate summary = + Right (Ident (updateName update) (snd $ updateType update)) + | otherwise = + Left (inPatternAs summary) + +summariseLoop :: MonadFreshNames m => + [DesiredUpdate (als, Type)] + -> Names + -> [(SubExp, Ident)] + -> [(Param DeclType, SubExp)] + -> Maybe (m [LoopResultSummary (als, Type)]) +summariseLoop updates usedInBody resmap merge = + sequence <$> zipWithM summariseLoopResult resmap merge + where summariseLoopResult (se, v) (fparam, mergeinit) + | Just update <- find (updateHasValue $ identName v) updates = + if updateSource update `S.member` usedInBody + then Nothing + else if hasLoopInvariantShape fparam then Just $ do + lowered_array <- newVName "lowered_array" + return LoopResultSummary { resultSubExp = se + , inPatternAs = v + , mergeParam = (fparam, mergeinit) + , relatedUpdate = Just (update, + lowered_array, + updateType update) + } + else Nothing + summariseLoopResult _ _ = + Nothing -- XXX: conservative; but this entire pass is going away. + + hasLoopInvariantShape = all loopInvariant . arrayDims . paramType + + merge_param_names = map (paramName . fst) merge + + loopInvariant (Var v) = v `notElem` merge_param_names + loopInvariant Constant{} = True + +data LoopResultSummary attr = + LoopResultSummary { resultSubExp :: SubExp + , inPatternAs :: Ident + , mergeParam :: (Param DeclType, SubExp) + , relatedUpdate :: Maybe (DesiredUpdate attr, VName, attr) + } + deriving (Show) + +indexSubstitutions :: [LoopResultSummary attr] + -> IndexSubstitutions attr +indexSubstitutions = mapMaybe getSubstitution + where getSubstitution res = do + (DesiredUpdate _ _ cs _ is _, nm, attr) <- relatedUpdate res + let name = paramName $ fst $ mergeParam res + return (name, (cs, nm, attr, is)) + +manipulateResult :: (Bindable lore, MonadFreshNames m) => + [LoopResultSummary (LetAttr lore)] + -> IndexSubstitutions (LetAttr lore) + -> m (Result, Stms lore) +manipulateResult summaries substs = do + let (orig_ses,updated_ses) = partitionEithers $ map unchangedRes summaries + (subst_ses, res_bnds) <- runWriterT $ zipWithM substRes updated_ses substs + return (orig_ses ++ subst_ses, stmsFromList res_bnds) + where + unchangedRes summary = + case relatedUpdate summary of + Nothing -> Left $ resultSubExp summary + Just _ -> Right $ resultSubExp summary + substRes (Var res_v) (subst_v, (_, nm, _, _)) + | res_v == subst_v = + return $ Var nm + substRes res_se (_, (cs, nm, attr, is)) = do + v' <- newIdent' (++"_updated") $ Ident nm $ typeOf attr + tell [certify cs $ mkLet [] [v'] $ BasicOp $ + Update nm (fullSlice (typeOf attr) is) res_se] + return $ Var $ identName v' diff --git a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs @@ -0,0 +1,135 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | This module exports facilities for transforming array accesses in +-- a list of 'Stm's (intended to be the bindings in a body). The +-- idea is that you can state that some variable @x@ is in fact an +-- array indexing @v[i0,i1,...]@. +module Futhark.Optimise.InPlaceLowering.SubstituteIndices + ( + substituteIndices + , IndexSubstitution + , IndexSubstitutions + ) where + +import Data.Semigroup ((<>)) +import Control.Monad +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.AST +import Futhark.Construct +import Futhark.Tools (fullSlice) +import Futhark.Util + +type IndexSubstitution attr = (Certificates, VName, attr, Slice SubExp) +type IndexSubstitutions attr = [(VName, IndexSubstitution attr)] + +typeEnvFromSubstitutions :: LetAttr lore ~ attr => + IndexSubstitutions attr -> Scope lore +typeEnvFromSubstitutions = M.fromList . map (fromSubstitution . snd) + where fromSubstitution (_, name, t, _) = + (name, LetInfo t) + +substituteIndices :: (MonadFreshNames m, BinderOps lore, Bindable lore, + Aliased lore, LetAttr lore ~ attr) => + IndexSubstitutions attr -> Stms lore + -> m (IndexSubstitutions attr, Stms lore) +substituteIndices substs bnds = + runBinderT (substituteIndicesInStms substs bnds) types + where types = typeEnvFromSubstitutions substs + +substituteIndicesInStms :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => + IndexSubstitutions (LetAttr (Lore m)) + -> Stms (Lore m) + -> m (IndexSubstitutions (LetAttr (Lore m))) +substituteIndicesInStms = foldM substituteIndicesInStm + +substituteIndicesInStm :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => + IndexSubstitutions (LetAttr (Lore m)) + -> Stm (Lore m) + -> m (IndexSubstitutions (LetAttr (Lore m))) +substituteIndicesInStm substs (Let pat lore e) = do + e' <- substituteIndicesInExp substs e + (substs', pat') <- substituteIndicesInPattern substs pat + addStm $ Let pat' lore e' + return substs' + +substituteIndicesInPattern :: (MonadBinder m, LetAttr (Lore m) ~ attr) => + IndexSubstitutions (LetAttr (Lore m)) + -> PatternT attr + -> m (IndexSubstitutions (LetAttr (Lore m)), PatternT attr) +substituteIndicesInPattern substs pat = do + (substs', context) <- mapAccumLM sub substs $ patternContextElements pat + (substs'', values) <- mapAccumLM sub substs' $ patternValueElements pat + return (substs'', Pattern context values) + where sub substs' patElem = return (substs', patElem) + +substituteIndicesInExp :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m), + LetAttr (Lore m) ~ attr) => + IndexSubstitutions (LetAttr (Lore m)) + -> Exp (Lore m) + -> m (Exp (Lore m)) +substituteIndicesInExp substs e = do + substs' <- copyAnyConsumed e + let substitute = identityMapper { mapOnSubExp = substituteIndicesInSubExp substs' + , mapOnVName = substituteIndicesInVar substs' + , mapOnBody = const $ substituteIndicesInBody substs' + } + + mapExpM substitute e + where copyAnyConsumed = + let consumingSubst substs' v + | Just (cs2, src2, src2attr, is2) <- lookup v substs = do + row <- certifying cs2 $ + letExp (baseString v ++ "_row") $ + BasicOp $ Index src2 $ fullSlice (typeOf src2attr) is2 + row_copy <- letExp (baseString v ++ "_row_copy") $ + BasicOp $ Copy row + return $ update v v (mempty, + row_copy, + src2attr `setType` + stripArray (length is2) (typeOf src2attr), + []) substs' + consumingSubst substs' _ = + return substs' + in foldM consumingSubst substs . S.toList . consumedInExp + +substituteIndicesInSubExp :: MonadBinder m => + IndexSubstitutions (LetAttr (Lore m)) + -> SubExp + -> m SubExp +substituteIndicesInSubExp substs (Var v) = Var <$> substituteIndicesInVar substs v +substituteIndicesInSubExp _ se = return se + +substituteIndicesInVar :: MonadBinder m => + IndexSubstitutions (LetAttr (Lore m)) + -> VName + -> m VName +substituteIndicesInVar substs v + | Just (cs2, src2, _, []) <- lookup v substs = + certifying cs2 $ letExp (baseString src2) $ BasicOp $ SubExp $ Var src2 + | Just (cs2, src2, src2_attr, is2) <- lookup v substs = + certifying cs2 $ + letExp "idx" $ BasicOp $ Index src2 $ fullSlice (typeOf src2_attr) is2 + | otherwise = + return v + +substituteIndicesInBody :: (MonadBinder m, Bindable (Lore m), Aliased (Lore m)) => + IndexSubstitutions (LetAttr (Lore m)) + -> Body (Lore m) + -> m (Body (Lore m)) +substituteIndicesInBody substs body = do + (substs', bnds') <- inScopeOf bnds $ + collectStms $ substituteIndicesInStms substs bnds + (ses, ses_bnds) <- inScopeOf bnds' $ + collectStms $ mapM (substituteIndicesInSubExp substs') $ bodyResult body + mkBodyM (bnds'<>ses_bnds) ses + where bnds = bodyStms body + +update :: VName -> VName -> IndexSubstitution attr -> IndexSubstitutions attr + -> IndexSubstitutions attr +update needle name subst ((othername, othersubst) : substs) + | needle == othername = (name, subst) : substs + | otherwise = (othername, othersubst) : update needle name subst substs +update needle _ _ [] = error $ "Cannot find substitution for " ++ pretty needle diff --git a/src/Futhark/Optimise/InliningDeadFun.hs b/src/Futhark/Optimise/InliningDeadFun.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/InliningDeadFun.hs @@ -0,0 +1,147 @@ +{-# LANGUAGE FlexibleContexts #-} +-- | This module implements a compiler pass for inlining functions, +-- then removing those that have become dead. +module Futhark.Optimise.InliningDeadFun + ( inlineAndRemoveDeadFunctions + , removeDeadFunctions + ) + where + +import Control.Monad.Identity +import Data.List +import Data.Loc +import Data.Maybe +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Futhark.Representation.SOACS +import Futhark.Transform.Rename +import Futhark.Analysis.CallGraph +import Futhark.Binder +import Futhark.Pass + +aggInlining :: CallGraph -> [FunDef] -> [FunDef] +aggInlining cg = filter keep . recurse + where noInterestingCalls :: S.Set Name -> FunDef -> Bool + noInterestingCalls interesting fundec = + case M.lookup (funDefName fundec) cg of + Just calls | not $ any (`elem` interesting') calls -> True + _ -> False + where interesting' = funDefName fundec `S.insert` interesting + + recurse funs = + let interesting = S.fromList $ map funDefName funs + (to_be_inlined, to_inline_in) = + partition (noInterestingCalls interesting) funs + inlined_but_entry_points = + filter (isJust . funDefEntryPoint) to_be_inlined + in if null to_be_inlined then funs + else inlined_but_entry_points ++ + recurse (map (`doInlineInCaller` to_be_inlined) to_inline_in) + + keep fundec = isJust (funDefEntryPoint fundec) || callsRecursive fundec + + callsRecursive fundec = maybe False (any recursive) $ + M.lookup (funDefName fundec) cg + + recursive fname = case M.lookup fname cg of + Just calls -> fname `elem` calls + Nothing -> False + +-- | @doInlineInCaller caller inlcallees@ inlines in @calleer@ the functions +-- in @inlcallees@. At this point the preconditions are that if @inlcallees@ +-- is not empty, and, more importantly, the functions in @inlcallees@ do +-- not call any other functions. Further extensions that transform a +-- tail-recursive function to a do or while loop, should do the transformation +-- first and then do the inlining. +doInlineInCaller :: FunDef -> [FunDef] -> FunDef +doInlineInCaller (FunDef entry name rtp args body) inlcallees = + let body' = inlineInBody inlcallees body + in FunDef entry name rtp args body' + +inlineInBody :: [FunDef] -> Body -> Body +inlineInBody inlcallees (Body attr stms res) = Body attr stms' res + where stms' = stmsFromList (concatMap inline $ stmsToList stms) + + inline (Let pat _ (Apply fname args _ (safety,loc,locs))) + | fun:_ <- filter ((== fname) . funDefName) inlcallees = + let param_stms = zipWith reshapeIfNecessary (map paramIdent $ funDefParams fun) (map fst args) + body_stms = stmsToList $ addLocations safety + (filter notNoLoc (loc:locs)) $ bodyStms $ funDefBody fun + res_stms = zipWith reshapeIfNecessary (patternIdents pat) + (bodyResult $ funDefBody fun) + in param_stms ++ body_stms ++ res_stms + inline stm = [inlineInStm inlcallees stm] + + reshapeIfNecessary ident se + | t@Array{} <- identType ident, + Var v <- se = + mkLet [] [ident] $ shapeCoerce (arrayDims t) v + | otherwise = + mkLet [] [ident] $ BasicOp $ SubExp se + +notNoLoc :: SrcLoc -> Bool +notNoLoc = (/=NoLoc) . locOf + +inliner :: Monad m => [FunDef] -> Mapper SOACS SOACS m +inliner funs = identityMapper { mapOnBody = const $ return . inlineInBody funs + , mapOnOp = return . inlineInSOAC funs + } + +inlineInSOAC :: [FunDef] -> SOAC SOACS -> SOAC SOACS +inlineInSOAC inlcallees = runIdentity . mapSOACM identitySOACMapper + { mapOnSOACLambda = return . inlineInLambda inlcallees + } + +inlineInStm :: [FunDef] -> Stm -> Stm +inlineInStm inlcallees (Let pat aux e) = + Let pat aux $ mapExp (inliner inlcallees) e + +inlineInLambda :: [FunDef] -> Lambda -> Lambda +inlineInLambda inlcallees (Lambda params body ret) = + Lambda params (inlineInBody inlcallees body) ret + +addLocations :: Safety -> [SrcLoc] -> Stms SOACS -> Stms SOACS +addLocations caller_safety more_locs = fmap onStm + where onStm stm = stm { stmExp = onExp $ stmExp stm } + onExp (Apply fname args t (safety, loc,locs)) = + Apply fname args t (min caller_safety safety, loc,locs++more_locs) + onExp (BasicOp (Assert cond desc (loc,locs))) = + case caller_safety of + Safe -> BasicOp $ Assert cond desc (loc,locs++more_locs) + Unsafe -> BasicOp $ SubExp $ Constant Checked + onExp (Op soac) = Op $ runIdentity $ mapSOACM + identitySOACMapper { mapOnSOACLambda = return . onLambda + } soac + onExp e = mapExp identityMapper { mapOnBody = const $ return . onBody + } e + onBody body = + body { bodyStms = addLocations caller_safety more_locs $ bodyStms body } + onLambda :: Lambda -> Lambda + onLambda lam = lam { lambdaBody = onBody $ lambdaBody lam } + +-- | A composition of 'inlineAggressively' and 'removeDeadFunctions', +-- to avoid the cost of type-checking the intermediate stage. +inlineAndRemoveDeadFunctions :: Pass SOACS SOACS +inlineAndRemoveDeadFunctions = + Pass { passName = "Inline and remove dead functions" + , passDescription = "Inline and remove resulting dead functions." + , passFunction = pass + } + where pass prog = do + let cg = buildCallGraph prog + renameProg . Prog . aggInlining cg . progFunctions =<< renameProg prog + +-- | @removeDeadFunctions prog@ removes the functions that are unreachable from +-- the main function from the program. +removeDeadFunctions :: Pass SOACS SOACS +removeDeadFunctions = + Pass { passName = "Remove dead functions" + , passDescription = "Remove the functions that are unreachable from the main function" + , passFunction = return . pass + } + where pass prog = + let cg = buildCallGraph prog + live_funs = filter (isFunInCallGraph cg) (progFunctions prog) + in Prog live_funs + isFunInCallGraph cg fundec = isJust $ M.lookup (funDefName fundec) cg diff --git a/src/Futhark/Optimise/MemoryBlockMerging.hs b/src/Futhark/Optimise/MemoryBlockMerging.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging.hs @@ -0,0 +1,28 @@ +-- | Merge memory blocks. +module Futhark.Optimise.MemoryBlockMerging + ( memoryBlockMergingCoalescing + , memoryBlockMergingReuse + ) where + +import Futhark.Pass +import Futhark.Representation.ExplicitMemory (ExplicitMemory) + +import Futhark.Optimise.MemoryBlockMerging.Coalescing (coalesceInProg) +import Futhark.Optimise.MemoryBlockMerging.Reuse (reuseInProg) + + +-- | Apply the coalescing part of the memory block merging optimisation. +memoryBlockMergingCoalescing :: Pass ExplicitMemory ExplicitMemory +memoryBlockMergingCoalescing = + Pass + "Memory block merging (coalescing)" + "Coalesce the memory blocks of arrays" + coalesceInProg + +-- | Apply the reuse part of the memory block merging optimisation. +memoryBlockMergingReuse :: Pass ExplicitMemory ExplicitMemory +memoryBlockMergingReuse = + Pass + "Memory block merging (reuse)" + "Reuse the memory blocks of arrays" + reuseInProg diff --git a/src/Futhark/Optimise/MemoryBlockMerging/ActualVariables.hs b/src/Futhark/Optimise/MemoryBlockMerging/ActualVariables.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/ActualVariables.hs @@ -0,0 +1,358 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find the actual variables that need updating when a variable attribute +-- needs updating. This is different than variable aliasing: Variable aliasing +-- is a theoretical concept, while this module has the practical purpose of +-- finding any extra variables that also need a change when a variable has a +-- change of memory block. +-- +-- If and DoLoop statements have special requirements, as do some aliasing +-- expressions. We don't want to (just) use the obvious statement variable; +-- sometimes updating the memory block of one variable actually means updating +-- the memory block of other variables as well. + +module Futhark.Optimise.MemoryBlockMerging.ActualVariables + ( findActualVariables + ) where + +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import qualified Data.List as L +import Data.Maybe (fromMaybe, mapMaybe, catMaybes) +import Control.Monad +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory ( + ExplicitMemorish, ExplicitMemory, InKernel) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types +import Futhark.Optimise.MemoryBlockMerging.AllExpVars + + +data Context = Context + { ctxVarToMem :: VarMemMappings MemorySrc + , ctxFirstUses :: FirstUses + } + deriving (Show) + +newtype FindM lore a = FindM { unFindM :: RWS Context () ActualVariables a } + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadState ActualVariables) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore, + LookInKernelExp lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +recordActuals :: VName -> Names -> FindM lore () +recordActuals stmt_var more_actuals = do + -- If S.empty has already been recorded, keep it at that. This is because the + -- ActualVariables system is currently also used for disabling memory block + -- optimisations -- if a variables resolves to the empty set, don't touch it. + -- This keeps some edge cases simple. FIXME at some point. + current_actuals <- M.lookup stmt_var <$> get + case S.null <$> current_actuals of + Just True -> return () + _ -> modify (insertOrUpdateMany stmt_var more_actuals) + +-- Find all the actual variables in a function definition. +findActualVariables :: VarMemMappings MemorySrc -> FirstUses -> + FunDef ExplicitMemory -> ActualVariables +findActualVariables var_mem_mappings first_uses fundef = + let context = Context var_mem_mappings first_uses + m = unFindM $ lookInBody $ funDefBody fundef + actual_variables = fst $ execRWS m context M.empty + in actual_variables + +lookInFParam :: FParam lore -> FindM lore () +lookInFParam (Param v _) = + recordActuals v $ S.singleton v + +lookInLParam :: LParam lore -> FindM lore () +lookInLParam (Param v _) = + recordActuals v $ S.singleton v + +lookInLambda :: LoreConstraints lore => Lambda lore -> FindM lore () +lookInLambda (Lambda params body _) = do + forM_ params lookInLParam + lookInBody body + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm stm@(Let (Pattern patctxelems patvalelems) _ e) = do + case (patvalelems, e) of + ([PatElem var _], BasicOp (Update orig _ _)) -> do + let actuals = S.fromList [var, orig] + -- When coalescing an in-place update statement, also look at the original + -- array. + recordActuals var actuals + -- When reusing a previous memory block, make sure to also update related + -- in-place updates. + recordActuals orig actuals + _ -> return () + + -- Ignore the existential memory blocks. + let bodyResult' = drop (length patctxelems) . bodyResult + + -- Special handling of loops, ifs, etc. + case e of + DoLoop _mergectxparams mergevalparams loopform body -> do + let body_vars0 = mapMaybe (subExpVar . snd) mergevalparams + body_vars1 = map (paramName . fst) mergevalparams + body_vars2 = S.toList $ findAllExpVars e + body_vars = body_vars0 ++ body_vars1 ++ body_vars2 + forM_ patvalelems $ \(PatElem var membound) -> do + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _) -> do + -- If mem is existential, we need to find the return memory that it + -- refers to. We cannot just look at its memory aliases, since it + -- likely aliases both the initial memory and the final memory. + + let zipped = zip patctxelems (bodyResult body) + mem_search = case L.find ((== mem) . patElemName . fst) zipped of + Just (_, Var res_mem) -> res_mem + _ -> mem + -- Find the ones using the same memory as the result of the loop + -- expression. + body_vars' <- filterM (lookupGivesMem mem_search) body_vars + -- Not only the result variable needs to change its memory block in + -- case of a future memory merging with it; also the variables + -- extracted above. + let actuals = var : body_vars' + forM_ actuals $ \a -> recordActuals a (S.fromList actuals) + -- Some of these can be changed later on to have an actual variable + -- set of S.empty, e.g. if one of the variables using the memory is + -- a rearrange operation. This is fine, and will occur in the walk + -- later on. + + -- If you extend this loop handling, make sure not to target existential + -- memory blocks. We want those to stay. + _ -> return () + + -- It seems wrong to change the memory of merge variables, so we disable + -- it. If we were to accept it, we would need to record what other + -- variables to change as well. Seems hard. + recordActuals var S.empty + + case loopform of + ForLoop _ _ _ loop_vars -> + -- Link 'array' to 'lvar' in 'for lvar in array' loop expressions. + forM_ loop_vars $ \(Param lvar _, array) -> + aliasOpHandleVar array lvar + WhileLoop _ -> return () + + If _se body_then body_else _types -> + -- We don't want to coalesce the existiential memory block of the if. + -- However, if a branch result has a memory block that is firstly used + -- inside the branch, it is okay to coalesce that in a future statement. + forM_ (zip3 patvalelems (bodyResult' body_then) (bodyResult' body_else)) + $ \(PatElem var membound, res_then, res_else) -> do + let body_vars = S.toList $ findAllExpVars e + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _) -> + if mem `L.elem` map patElemName patctxelems + then + -- If the memory block is existential, we say that the If result + -- refers to all results in the If. + recordActuals var + $ S.fromList (var : catMaybes [subExpVar res_then, subExpVar res_else]) + + else do + -- If the memory block is not existential, we need to find all the + -- variables in any sub-bodies using the same memory block (like + -- with loops). + body_vars' <- filterM (lookupGivesMem mem) body_vars + + first_uses <- asks ctxFirstUses + case filter ((mem `S.member`) . (`lookupEmptyable` first_uses)) body_vars' of + [] -> + -- Not just the result variable needs to change its memory + -- block in case of a future memory block merging with it; + -- also the variables extracted above. + recordActuals var $ S.fromList (var : body_vars') + _ -> + -- If we come across a non-existential If which can be said to + -- create a new array *and* which has one or more bodies which + -- can also be said to create a new array *in the same memory* + -- (i.e. has first memory uses), then we disable it. This is + -- not at all an impossible case to handle, but such an If is + -- weird, since it would make more sense if it had existential + -- memory, so maybe something needs to be done somewhere else + -- in the compiler? If this is naively enabled, we can get an + -- error because the sub-body results are first uses while the + -- main result is not. This can be "fixed" by stating that + -- the If as a whole is also a first use of the memory, but + -- this seems too conservative. FIXME. + forM_ (var : body_vars') $ \v -> recordActuals v S.empty + + _ -> return () + + BasicOp (Index orig _) -> do + let ielem = head patvalelems -- Should be okay. + var = patElemName ielem + case patElemAttr ielem of + ExpMem.MemArray{} -> + -- Disable merging for index expressions that return arrays. Maybe + -- too restrictive. Make sure the source also updates the memory of + -- the index when updated. The array might be an aliasing operation, + -- in which case we try to find the original array. + aliasOpHandleVar orig var + _ -> return () + + -- Support reusing the memory of reshape operations by recording the origin + -- array that is being reshaped. Only partial support for reshape + -- operations: If the shape is more than one-dimensional, mark the statement + -- as disabled for memory merging operations. + BasicOp (Reshape shapechange_var orig) -> + forM_ (map patElemName patvalelems) $ \var -> do + orig' <- aliasOpRoot' orig + mem_orig <- M.lookup orig' <$> asks ctxVarToMem + case (shapechange_var, mem_orig) of + ([_], Just (MemorySrc _ _ (Shape [_]))) -> + recordActuals var $ S.fromList [var, orig] + -- Works, but only in limited cases where the reshape is not even + -- that useful to begin with; mostly cases where a reshape was + -- inserted by the compiler in an assert-like manner. + _ -> + recordActuals var S.empty + -- FIXME: The problem with these more complex cases with more than + -- one dimension is that a slice is relative to the shape of the + -- reshaped array, and not the original array. Disabled for now. + recordActuals orig' $ S.fromList [orig', var] + + -- For the other aliasing operations, disable their use for now. If the + -- source has a change of memory block, make sure to change this as well. + BasicOp (Rearrange _ orig) -> + aliasOpHandle orig patvalelems + + BasicOp (Rotate _ orig) -> + aliasOpHandle orig patvalelems + + BasicOp (Opaque (Var orig)) -> + aliasOpHandle orig patvalelems + + _ -> forM_ patvalelems $ \(PatElem var membound) -> do + let body_vars = S.toList $ findAllExpVars e + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _) -> do + body_vars' <- filterM (lookupGivesMem mem) body_vars + recordActuals var $ S.fromList (var : body_vars') + _ -> return () + + -- If we are inside a kernel, check for actual variables in the KernelExp of + -- the statement. + lookInKernelExp stm + + -- Recurse over any sub-bodies. + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + , walkOnFParam = lookInFParam + , walkOnLParam = lookInLParam + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInLambda + , walkOnKernelLParam = lookInLParam + } + +-- If we have a rotate or similar, we want to find the original array and +-- associate *that* with this aliasing array, so that changes to the original +-- array will affect this one as well. +aliasOpHandle :: VName -> [PatElem lore] -> FindM lore () +aliasOpHandle orig patvalelems = + forM_ (map patElemName patvalelems) $ aliasOpHandleVar orig + +aliasOpHandleVar :: VName -> VName -> FindM lore () +aliasOpHandleVar orig var = do + recordActuals var S.empty + orig' <- aliasOpRoot' orig + recordActuals orig' $ S.fromList [orig', var] + +aliasOpRoot :: VName -> FindM lore (Maybe VName) +aliasOpRoot orig = do + current_actuals <- get + return $ case S.null <$> M.lookup orig current_actuals of + -- If the original array is itself an aliasing operation, find the *actual* + -- original array. There can be more than one reference. We just pick the + -- first one -- any one should do, since there is a transitive closure + -- calculation later on. + Just True -> case M.keys (M.filter (orig `S.member`) current_actuals) of + orig' : _ -> Just orig' + _ -> Nothing + -- Else, just return orig. + _ -> Just orig + +aliasOpRoot' :: VName -> FindM lore VName +aliasOpRoot' orig = + fromJust ("at some point there will have been a proper statement: " + ++ pretty orig) <$> aliasOpRoot orig + +-- Is the memory block of 'v' the same as 'mem'? +lookupGivesMem :: MName -> VName -> FindM lore Bool +lookupGivesMem mem v = do + m <- M.lookup v <$> asks ctxVarToMem + return (Just mem == (memSrcName <$> m)) + +class LookInKernelExp lore where + -- Find actual vars in 'KernelExp's. + lookInKernelExp :: Stm lore -> FindM lore () + +instance LookInKernelExp ExplicitMemory where + lookInKernelExp (Let (Pattern _ patvalelems) _ e) = case e of + Op (ExpMem.Inner (Kernel _ _ _ (KernelBody _ _ ress))) -> + zipWithM_ (\(PatElem var _) res -> case res of + WriteReturn _ arr _ -> + recordActuals arr $ S.singleton var + _ -> return () + ) patvalelems ress + _ -> return () + +instance LookInKernelExp InKernel where + lookInKernelExp (Let _ _ e) = case e of + Op (ExpMem.Inner ke) -> case ke of + ExpMem.GroupReduce _ _ input -> do + let arrs = map snd input + extendActualVarsInKernel e arrs + ExpMem.GroupScan _ _ input -> do + let arrs = map snd input + extendActualVarsInKernel e arrs + ExpMem.GroupStream _ _ _ _ arrs -> + extendActualVarsInKernel e arrs + _ -> return () + _ -> return () + +-- Record actual variables for input arrays to 'KernelExp's. +extendActualVarsInKernel :: Exp InKernel -> [VName] -> FindM InKernel () +extendActualVarsInKernel e arrs = forM_ arrs $ \var -> do + -- The array might be an aliasing operation, in which case we try to find the + -- original array. + var' <- fromMaybe var <$> aliasOpRoot var + varmem <- M.lookup var <$> asks ctxVarToMem + case varmem of + Just mem -> do + let body_vars = findAllExpVars e + body_vars' <- filterSetM (lookupGivesMem $ memSrcName mem) body_vars + let actuals = S.insert var' body_vars' + recordActuals var' actuals + Nothing -> return () diff --git a/src/Futhark/Optimise/MemoryBlockMerging/AllExpVars.hs b/src/Futhark/Optimise/MemoryBlockMerging/AllExpVars.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/AllExpVars.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE LambdaCase #-} +-- | Find all variables in a statement. +module Futhark.Optimise.MemoryBlockMerging.AllExpVars + ( findAllExpVars + ) where + +import qualified Data.Set as S +import Control.Monad +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemorish) +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +newtype FindM lore a = FindM { unFindM :: Writer Names a } + deriving (Monad, Functor, Applicative, + MonadWriter Names) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- Find all the variables (both free and bound) that occur in a statement and +-- any nested bodies. We use this to record which extra variables need to have +-- their memory blocks updated when some variable needs updating. The result +-- might be an empty set, but in the case of If, DoLoop, and kernels, the result +-- might be nonempty. We cannot just find all variables in the program and look +-- through them every time we need to, since a memory block can (at least in +-- theory) be present in two different places (which also means by two different +-- variable sets) in a program, so we should limit ourselves to looking in the +-- statement declaring a new current use of the memory. +findAllExpVars :: LoreConstraints lore => + Exp lore -> Names +findAllExpVars e = + let m = unFindM $ lookInExp e + in execWriter m + +lookInExp :: LoreConstraints lore => + Exp lore -> FindM lore () +lookInExp = fullWalkExpM walker walker_kernel + where walker = identityWalker + { walkOnBody = lookInBody + , walkOnFParam = lookInFParam + , walkOnLParam = lookInLParam + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInLambda + , walkOnKernelLParam = lookInLParam + } + +lookInFParam :: FParam lore -> FindM lore () +lookInFParam (Param x _) = + tell $ S.singleton x + +lookInLParam :: LParam lore -> FindM lore () +lookInLParam (Param x _) = + tell $ S.singleton x + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds res) = do + mapM_ lookInStm bnds + forM_ res $ \case + ThreadsReturn{} -> return () + WriteReturn _ arr _ -> tell $ S.singleton arr + ConcatReturns{} -> return () + KernelInPlaceReturn v -> tell $ S.singleton v + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern _ patvalelems) _ e) = do + forM_ patvalelems $ \(PatElem x _) -> + tell $ S.singleton x + lookInExp e + +lookInLambda :: LoreConstraints lore => + Lambda lore -> FindM lore () +lookInLambda (Lambda params body _) = do + forM_ params lookInLParam + lookInBody body diff --git a/src/Futhark/Optimise/MemoryBlockMerging/AuxiliaryInfo.hs b/src/Futhark/Optimise/MemoryBlockMerging/AuxiliaryInfo.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/AuxiliaryInfo.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +-- | Helper information for the main optimisation passes. +module Futhark.Optimise.MemoryBlockMerging.AuxiliaryInfo + ( AuxiliaryInfo(..), getAuxiliaryInfo) +where + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemory) + +import Futhark.Optimise.MemoryBlockMerging.Types + +import Futhark.Optimise.MemoryBlockMerging.VariableMemory +import Futhark.Optimise.MemoryBlockMerging.MemoryAliases +import Futhark.Optimise.MemoryBlockMerging.VariableAliases +import Futhark.Optimise.MemoryBlockMerging.Liveness.FirstUse +import Futhark.Optimise.MemoryBlockMerging.Liveness.LastUse +import Futhark.Optimise.MemoryBlockMerging.Liveness.Interference +import Futhark.Optimise.MemoryBlockMerging.ActualVariables +import Futhark.Optimise.MemoryBlockMerging.Existentials + +-- Information needed by multiple transformations. +data AuxiliaryInfo = AuxiliaryInfo + { auxName :: Name -- For debugging. + , auxVarMemMappings :: VarMemMappings MemorySrc + , auxMemAliases :: MemAliases + , auxVarAliases :: VarAliases + , auxFirstUses :: FirstUses + , auxLastUses :: LastUses + , auxInterferences :: Interferences + , auxPotentialKernelDataRaceInterferences + :: PotentialKernelDataRaceInterferences + , auxActualVariables :: ActualVariables + , auxExistentials :: Names + } + deriving (Show) + +getAuxiliaryInfo :: FunDef ExplicitMemory -> AuxiliaryInfo +getAuxiliaryInfo fundef = + let name = funDefName fundef + var_to_mem = findVarMemMappings fundef + mem_aliases = findMemAliases fundef var_to_mem + var_aliases = findVarAliases fundef + first_uses = findFirstUses var_to_mem mem_aliases fundef + last_uses = findLastUses var_to_mem mem_aliases first_uses existentials + fundef + (interferences, potential_kernel_interferences) = + findInterferences var_to_mem mem_aliases first_uses last_uses + existentials fundef + actual_variables = findActualVariables var_to_mem first_uses fundef + existentials = findExistentials fundef + in AuxiliaryInfo + { auxName = name + , auxVarMemMappings = var_to_mem + , auxMemAliases = mem_aliases + , auxVarAliases = var_aliases + , auxFirstUses = first_uses + , auxLastUses = last_uses + , auxInterferences = interferences + , auxPotentialKernelDataRaceInterferences = potential_kernel_interferences + , auxActualVariables = actual_variables + , auxExistentials = existentials + } diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing.hs b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing.hs @@ -0,0 +1,31 @@ +-- | Coalesce the memory blocks of arrays. +-- +-- Enable by setting the environment variable MEMORY_BLOCK_MERGING_COALESCING=1. +module Futhark.Optimise.MemoryBlockMerging.Coalescing + ( coalesceInProg + ) where + +import Futhark.Pass + +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemory) + +import Futhark.Optimise.MemoryBlockMerging.AuxiliaryInfo +import Futhark.Optimise.MemoryBlockMerging.Coalescing.AllocationMovingUp +import Futhark.Optimise.MemoryBlockMerging.Coalescing.Core + + +coalesceInProg :: Prog ExplicitMemory -> PassM (Prog ExplicitMemory) +coalesceInProg = intraproceduralTransformation coalesceInFunDef + +coalesceInFunDef :: MonadFreshNames m + => FunDef ExplicitMemory + -> m (FunDef ExplicitMemory) +coalesceInFunDef fundef0 = do + let fundef1 = moveUpAllocsFunDef fundef0 + aux1 = getAuxiliaryInfo fundef1 + coreCoalesceFunDef fundef1 + (auxVarMemMappings aux1) (auxMemAliases aux1) + (auxVarAliases aux1) (auxFirstUses aux1) (auxLastUses aux1) + (auxActualVariables aux1) (auxExistentials aux1) diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/AllocationMovingUp.hs b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/AllocationMovingUp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/AllocationMovingUp.hs @@ -0,0 +1,94 @@ +-- | Move allocation statements upwards in the bodies of a program to enable +-- more memory block coalescings. +-- +-- This should be run *before* the coalescing pass, as it enables more +-- optimisations. +module Futhark.Optimise.MemoryBlockMerging.Coalescing.AllocationMovingUp + ( moveUpAllocsFunDef + ) where + +import qualified Data.Set as S +import Data.Maybe (mapMaybe) + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem + +import Futhark.Optimise.MemoryBlockMerging.CrudeMovingUp + + +findAllocHoistees :: Body ExplicitMemory -> Maybe [FParam ExplicitMemory] + -> [VName] +findAllocHoistees body params = + let all_found = mapMaybe findThemStm stms + ++ maybe [] (mapMaybe findThemFParam) params + extras = concatMap snd all_found + allocs = map fst all_found + -- We must hoist the alloc expressions in the end. If we hoist an alloc + -- before we hoist one of its array creations (in case of in-place + -- updates), that array creation might in turn hoist something depending + -- on another memory block mem_y further up than the allocation of memory + -- block mem_x. This will become a problem if mem_y can get coalesced + -- into mem_x. + -- + -- Maybe there is a nicer way to guarantee that this does not happen, but + -- this seems to work for now. + -- + -- We reverse the non-alloc dependencies to ensure (sloppily) that they do + -- not change positions internally compared to the original program: For + -- example, if a statement x is located before a statement y, and both x + -- and y need to be hoisted, then we need to hoist x in the end, so that + -- it can be hoisted further than y, which might have been hoisted to + -- before x. A better solution is welcome! + in reverse extras ++ reverse allocs + + where stms :: [Stm ExplicitMemory] + stms = stmsToList $ bodyStms body + + findThemStm :: Stm ExplicitMemory -> Maybe (VName, [VName]) + findThemStm (Let (Pattern _ [PatElem xmem _]) _ (Op ExpMem.Alloc{})) = + usedByCopyOrConcat xmem + findThemStm _ = Nothing + + -- A function paramater can be a unique memory block. While we cannot + -- hoist that, we may have to hoist an index in an in-place update that + -- uses the memory. + findThemFParam :: FParam ExplicitMemory -> Maybe (VName, [VName]) + findThemFParam (Param xmem ExpMem.MemMem{}) = usedByCopyOrConcat xmem + findThemFParam _ = Nothing + + -- Is the allocated memory used by either Copy or Concat in the function + -- body? Those are the only kinds of memory we care about, since those + -- are the cases handled by coalescing. Also find the names used by + -- in-place updates, since those also need to be hoisted (as an example + -- of this, consider the 'copy/pos1.fut' test where the replicate + -- expression needs to be hoisted as well as its memory allocation). + usedByCopyOrConcat :: VName -> Maybe (VName, [VName]) + usedByCopyOrConcat xmem_alloc = + let vs = mapMaybe checkStm stms + vs' = if null vs then Nothing else Just (xmem_alloc, concat vs) + in vs' + + where checkStm :: Stm ExplicitMemory -> Maybe [VName] + checkStm (Let + (Pattern _ + [PatElem _ (ExpMem.MemArray _ _ _ (ExpMem.ArrayIn xmem_pat _))]) + _ + (BasicOp bop)) + | xmem_pat == xmem_alloc = + case bop of + Update v slice _ -> + -- The source array must also be hoisted so that it + -- is initialized before it is used by the + -- coalesced party. Any index variables are also + -- hoisted. + Just $ v : S.toList (freeIn slice) + Copy{} -> Just [] + Concat{} -> Just [] + _ -> Nothing + checkStm _ = Nothing + +moveUpAllocsFunDef :: FunDef ExplicitMemory + -> FunDef ExplicitMemory +moveUpAllocsFunDef fundef = + moveUpInFunDef fundef findAllocHoistees diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Core.hs b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Core.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Core.hs @@ -0,0 +1,624 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +module Futhark.Optimise.MemoryBlockMerging.Coalescing.Core + ( coreCoalesceFunDef + ) where + +import qualified Data.Set as S +import qualified Data.List as L +import qualified Data.Map.Strict as M +import Data.Maybe (maybe, fromMaybe, mapMaybe, isJust) +import Control.Monad +import Control.Monad.RWS + +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory ( + ExplicitMemory, ExplicitMemorish) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Tools + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types +import Futhark.Optimise.MemoryBlockMerging.MemoryUpdater + +import Futhark.Optimise.MemoryBlockMerging.PrimExps (findPrimExpsFunDef) +import Futhark.Optimise.MemoryBlockMerging.Coalescing.Exps +import Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition2 +import Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition3 +import Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition5 +import Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizes + + +-- Some of these attributes could be split into separate Coalescing helper +-- modules if it becomes confusing. Their computations are fairly independent. +data Current = Current + { -- Coalescings state. Also save offsets and slices in the case that an + -- optimistic coalescing later becomes part of a chain of coalescings, where + -- it is offset yet again, and where it should maintain its old relative + -- offset. FIXME: This works, but is inefficient in the long run, as we + -- need to update it whenever we come across a coalescing that also affects + -- previous coalescings. The directions of the coalescings is inherently + -- bottom-up, but our algorithm is top-down. It should be possible to + -- rewrite it. + curCoalescedIntos :: CoalescedIntos + , curMemsCoalesced :: MemsCoalesced + } + deriving (Show) + +type CoalescedIntos = M.Map VName (S.Set (VName, PrimExp VName, + [Slice (PrimExp VName)])) +type MemsCoalesced = M.Map VName MemoryLoc + +emptyCurrent :: Current +emptyCurrent = Current + { curCoalescedIntos = M.empty + , curMemsCoalesced = M.empty + } + +data Context = Context + { ctxFunDef :: FunDef ExplicitMemory + -- ^ Keep the entire function definition around for lookup purposes. + , ctxVarToMem :: VarMemMappings MemorySrc + -- ^ From the module VariableMemory. + , ctxMemAliases :: MemAliases + -- ^ From the module MemoryAliases. + , ctxVarAliases :: VarAliases + -- ^ From the module VariableAliases. + , ctxFirstUses :: FirstUses + -- ^ From the module FirstUses. + , ctxLastUses :: LastUses + -- ^ From the module LastUses. + , ctxActualVars :: M.Map VName Names + -- ^ From the module ActualVariables. + , ctxExistentials :: Names + -- ^ From the module Existentials. + , ctxVarPrimExps :: M.Map VName (PrimExp VName) + -- ^ From the module PrimExps. + , ctxVarExps :: M.Map VName Exp' + -- ^ Statement-name-to-expression mappins for the entire function. + , ctxAllocatedBlocksBeforeCreation :: M.Map VName MNames + -- ^ Safety condition 2. + , ctxVarsInUseBeforeMem :: M.Map MName Names + -- ^ Safety condition 5. + , ctxCurSnapshot :: Current + -- ^ Keep a snapshot (used in 'tryCoalesce' for Concat). + } + deriving (Show) + +newtype FindM lore a = FindM { unFindM :: RWS Context () Current a } + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadState Current) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +modifyCurCoalescedIntos :: (CoalescedIntos -> CoalescedIntos) -> FindM lore () +modifyCurCoalescedIntos f = + modify $ \c -> c { curCoalescedIntos = f $ curCoalescedIntos c } + +modifyCurMemsCoalesced :: (MemsCoalesced -> MemsCoalesced) -> FindM lore () +modifyCurMemsCoalesced f = + modify $ \c -> c { curMemsCoalesced = f $ curMemsCoalesced c } + +ifExp :: MonadReader Context m => + VName -> m (Maybe Exp') +ifExp var = do + var_exp <- M.lookup var <$> asks ctxVarExps + return $ case var_exp of + Just e@(Exp _ _ If{}) -> Just e + _ -> Nothing + +isIfExp :: MonadReader Context m => + VName -> m Bool +isIfExp var = isJust <$> ifExp var + +isLoopExp :: MonadReader Context m => + VName -> m Bool +isLoopExp var = do + var_exp <- M.lookup var <$> asks ctxVarExps + return $ case var_exp of + Just (Exp _ _ DoLoop{}) -> True + _ -> False + +isReshapeExp :: MonadReader Context m => + VName -> m Bool +isReshapeExp var = do + var_exp <- M.lookup var <$> asks ctxVarExps + return $ case var_exp of + Just (Exp _ _ (BasicOp Reshape{})) -> True + _ -> False + +-- Lookup the memory block statically associated with a variable. +lookupVarMem :: MonadReader Context m => + VName -> m MemorySrc +lookupVarMem var = + -- This should always be called from a place where it is certain that 'var' + -- refers to a statement with an array expression. + fromJust ("lookup memory block from " ++ pretty var) . M.lookup var + <$> asks ctxVarToMem + +lookupActualVars :: MonadReader Context m => + VName -> m Names +lookupActualVars var = do + actual_vars <- asks ctxActualVars + -- Do this recursively. + let actual_vars' = expandWithAliases actual_vars actual_vars + return $ fromMaybe (S.singleton var) $ M.lookup var actual_vars' + +-- Lookup the memory block currenty associated with a variable. In most cases +-- (maybe all) this could probably replace 'lookupVarMem', though it would not +-- always be necessary. +lookupCurrentVarMem :: VName -> FindM lore (Maybe VName) +lookupCurrentVarMem var = do + -- Current result... + mem_cur <- M.lookup var . curMemsCoalesced <$> asks ctxCurSnapshot + -- ... or original result. + -- + -- This is why we save the variables after creation, not the memory + -- blocks: Variables stay the same, but memory blocks may change, which + -- is relevant in the case of a chain of coalescings. + mem_orig <- M.lookup var <$> asks ctxVarToMem + return $ case (mem_cur, mem_orig) of + (Just m, _) -> Just (memLocName m) -- priority choice + (_, Just m) -> Just (memSrcName m) + _ -> Nothing + +withMemAliases :: MonadReader Context m => + VName -> m Names +withMemAliases mem = + -- The only memory blocks with memory aliases are the existiential ones, so + -- using a static ctxMemAliases should be okay, as they will not change during + -- the transformation in this module. + S.union (S.singleton mem) . lookupEmptyable mem + <$> asks ctxMemAliases + +data Bindage = BindInPlace VName (Slice SubExp) + | BindVar + +recordOptimisticCoalescing :: VName -> PrimExp VName + -> [Slice (PrimExp VName)] + -> VName -> MemoryLoc -> Bindage -> FindM lore () +recordOptimisticCoalescing src offset ixfun_slices dst dst_memloc bindage = do + modifyCurCoalescedIntos $ insertOrUpdate dst (src, offset, ixfun_slices) + + -- If this is an in-place operation, we future-proof future coalescings by + -- recording that they also need to take a look at the original array, not + -- just the result of an in-place update into it. + case bindage of + BindVar -> return () + BindInPlace orig _ -> + modifyCurCoalescedIntos $ insertOrUpdate dst (orig, zeroOffset, []) + + modifyCurMemsCoalesced $ M.insert src dst_memloc + +coreCoalesceFunDef :: MonadFreshNames m => + FunDef ExplicitMemory -> VarMemMappings MemorySrc + -> MemAliases -> VarAliases -> FirstUses -> LastUses + -> ActualVariables -> Names -> m (FunDef ExplicitMemory) +coreCoalesceFunDef fundef var_to_mem mem_aliases var_aliases first_uses + last_uses actual_vars existentials = do + let primexps = findPrimExpsFunDef fundef + exps = findExpsFunDef fundef + cond2 = findSafetyCondition2FunDef fundef + cond5 = findSafetyCondition5FunDef fundef first_uses + context = Context { ctxFunDef = fundef + , ctxVarToMem = var_to_mem + , ctxMemAliases = mem_aliases + , ctxVarAliases = var_aliases + , ctxFirstUses = first_uses + , ctxLastUses = last_uses + , ctxActualVars = actual_vars + , ctxExistentials = existentials + , ctxVarPrimExps = primexps + , ctxVarExps = exps + , ctxAllocatedBlocksBeforeCreation = cond2 + , ctxVarsInUseBeforeMem = cond5 + , ctxCurSnapshot = emptyCurrent + } + m = unFindM $ lookInBody $ funDefBody fundef + var_to_mem_res = curMemsCoalesced $ fst $ execRWS m context emptyCurrent + sizes = memBlockSizesFunDef fundef + transformFromVarMemMappings var_to_mem_res (M.map memSrcName var_to_mem) (M.map fst sizes) (M.map fst sizes) False fundef + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +zeroOffset :: PrimExp VName +zeroOffset = primExpFromSubExp (IntType Int32) (constant (0 :: Int32)) + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern _patctxelems patvalelems) _ e) = do + -- COALESCING-SPECIFIC HANDLING for Copy and Concat. + case patvalelems of + [PatElem dst ExpMem.MemArray{}] -> do + -- We create a function and pass it around instead of just applying it to + -- the memory of the MemBound. We do this, since any source variables + -- might have more actual variables with different index functions that + -- also need to be fixed -- e.g. in the case of reshape, where both the + -- reshaped array and the original array need to get their index functions + -- updated. + -- + -- We take a snapshot of the current state of the curCoalescedIntos state + -- field. We need this feature to avoid having fewer coalescings just + -- because of the placement of the sources. For example, for + -- + -- let b = ... + -- let a = ... + -- let c = concat a b + -- + -- the coalescing pass will first coalesce m_a into m_c, which will + -- succeed. Then it will to coalesce m_b into m_c, which will (naively) + -- fail because of safety condition 3 arguing that m_c is now in use after + -- the creation of 'b' and before its use, since 'a' now uses m_c. + -- + -- (Alternatively, we could do some more general index function analysis + -- to check for things that will never overlap in merged memory, but this + -- seems easier.) + cur_snapshot <- get + var_to_mem <- asks ctxVarToMem + local (\ctx -> ctx { ctxCurSnapshot = cur_snapshot }) + $ case e of + -- In-place update. + BasicOp (Update orig slice (Var src)) -> + case M.lookup src var_to_mem of + Just _ -> + let ixfun_slices = + let slice' = map (primExpFromSubExp (IntType Int32) <$>) slice + in [slice'] + bindage = BindInPlace orig slice + in tryCoalesce dst ixfun_slices bindage src zeroOffset + Nothing -> + return () + + -- Copy. + BasicOp (Copy src) -> + tryCoalesce dst [] BindVar src zeroOffset + + -- Concat. + BasicOp (Concat 0 src0 src0s _) -> do + let srcs = src0 : src0s + shapes <- mapM ((memSrcShape <$>) . lookupVarMem) srcs + let getOffsets offset_prev shape = + let se = head (shapeDims shape) -- Should work. + len = primExpFromSubExp (IntType Int32) se + offset_new = offset_prev + len + in offset_new + offsets = init (scanl getOffsets zeroOffset shapes) + zipWithM_ (tryCoalesce dst [] BindVar) srcs offsets + _ -> return () + _ -> return () + + + -- RECURSIVE BODY WALK. + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + +tryCoalesce :: VName -> [Slice (PrimExp VName)] -> Bindage -> + VName -> PrimExp VName -> FindM lore () +tryCoalesce dst ixfun_slices bindage src offset = do + mem_dst <- lookupVarMem dst + + -- For ifs and loops and some aliasing expressions (e.g. reshape), this tells + -- us what non-existential source variables actually need to have assigned the + -- new memory block. + src's <- S.toList <$> lookupActualVars src + + -- From earlier optimistic coalescings. Remember to also get the coalescings + -- from the actual variables in e.g. loops. + coalesced_intos <- curCoalescedIntos <$> asks ctxCurSnapshot + let (src0s, offset0s, ixfun_slice0ss) = + unzip3 $ S.toList $ S.unions + $ map (`lookupEmptyable` coalesced_intos) (src : src's) + + var_to_pe <- asks ctxVarPrimExps + + let srcs = src's ++ src0s + -- The same number of base offsets as in src's. + offsets = replicate (length src's) offset + -- The offsets of any previously optimistically coalesced src0s must be + -- re-offset relative to the offset of the newest coalescing. + ++ map (\o0 -> if o0 == zeroOffset && offset == zeroOffset + -- This should not be necessary, and maybe it + -- is not (but there were some problems). + then zeroOffset + else offset + o0) offset0s + ixfun_slicess = replicate (length src's) ixfun_slices + -- Same as above, kind of. + ++ map (\slices0 -> ixfun_slices ++ slices0) ixfun_slice0ss + + let ixfuns' = zipWith (\offset_local islices -> + let ixfun0 = memSrcIxFun mem_dst + ixfun1 = foldl IxFun.slice ixfun0 islices + + -- 'ixfun_slices' contain the slices that are the + -- result of a new coalescing, contrary to the + -- slices in 'ixfun_slice0ss' which contain + -- previously registered slices. + -- 'offsetIndexDWIM' handles the case that we + -- want to offset a DimFix if it is the result of + -- a previous coalescing, and not the current + -- one. We do that by counting the number of + -- 'DimFix'es that originate in the new + -- coalescing, and then ignore those for our + -- heuristic. This is a hack. + initial_dimfixes = L.takeWhile (isJust . dimFix) (concat ixfun_slices) + ixfun2 = if offset_local == zeroOffset + then ixfun1 -- Should not be necessary, + -- but it makes the type + -- checker happy for now. + else IxFun.offsetIndexDWIM (length initial_dimfixes) ixfun1 offset_local + ixfun3 = expandIxFun var_to_pe ixfun2 + in ixfun3 + ) offsets ixfun_slicess + + -- Not everything supported yet. This dials back the optimisation on areas + -- where it fails. + existentials <- asks ctxExistentials + let currentlyDisabled src_local = do + -- This case covers the problem described in several programs in + -- tests/coalescing/wip/loop/ (for programs where it is overly + -- conservative) and tests/coalescing/loop/replicate-in-loop.fut (where + -- it is absolutely needed to keep the program correct). It is a + -- conservative requirement and could likely be loosened up. + + src_local_is_loop <- isLoopExp src_local + + -- if the source contains the result a loop expression, and that result + -- is an array with existential memory, don't coalesce. Since memory + -- can be allocated inside loops, coalescing with no further rules might + -- end up having the same arrays use memory allocated outside the loop, + -- which is not always okay. + let res = src_local_is_loop + && src_local `L.elem` existentials + return res + + safe0 <- not . or <$> mapM currentlyDisabled srcs + + -- Safety condition 1 is the same for all eventual previous arrays from srcs + -- that also need to be coalesced into dst, so we check it here instead of + -- checking it independently for every sub src. This also ensures that we + -- check that the destination memory is lastly used in *just* this statement, + -- not also in any previous statement that uses the same memory block, which + -- could very well fail. + mem_src_base <- lookupVarMem src + safe1 <- safetyCond1 dst mem_src_base + + when (safe0 && safe1) $ do + safes <- zipWithM (canBeCoalesced dst) srcs ixfuns' + when (and safes) $ do + -- Any previous src0s coalescings must be deleted. + modifyCurCoalescedIntos $ M.delete src + -- The rest will be overwritten below. + + -- We then need to record that, from what we currently know, src and any + -- nested src0s can all use the memory of dst with the new index functions. + forM_ (L.zip4 srcs offsets ixfun_slicess ixfuns') + $ \(src_local, offset_local, ixfun_slices_local, ixfun_local) -> do + denotes_existential <- S.member src_local <$> asks ctxExistentials + is_if <- isIfExp src_local + dst_memloc <- + if denotes_existential && not is_if + then do + -- Only use the new index function. Keep the existential memory + -- block. This means we have to make fewer changes to the program. + -- + -- FIXME: However, if we are at an If expression with an existential + -- memory block, we ignore it. This is due to some special handling + -- of If in MemoryUpdater, which is again due to branches having + -- explicit return types. This might not be correct. + mem_src <- lookupVarMem src_local + return $ MemoryLoc (memSrcName mem_src) ixfun_local + else + -- Use both the new memory block and the new index function. + return $ MemoryLoc (memSrcName mem_dst) ixfun_local + recordOptimisticCoalescing + src_local offset_local ixfun_slices_local + dst dst_memloc bindage + +canBeCoalesced :: VName -> VName -> ExpMem.IxFun -> FindM lore Bool +canBeCoalesced dst src ixfun = do + mem_dst <- lookupVarMem dst + mem_src <- lookupVarMem src + + safe2 <- safetyCond2 src mem_dst + safe3 <- safetyCond3 src dst mem_dst + safe4 <- safetyCond4 src + safe5 <- safetyCond5 mem_src ixfun + + safe_if <- safetyIf src dst + + let safe_all = safe2 && safe3 && safe4 && safe5 && safe_if + return safe_all + +-- Safety conditions for each statement with a Copy or Concat: +-- +-- 1. mem_src is not used beyond the statement. Handle by checking LastUses for +-- the statement. +-- +-- 2. The allocation of mem_dst occurs before the creation of src, i.e. the +-- first use of mem_src. Handle by checking +-- ctxAllocatedBlocksBeforeCreation. +-- +-- 3. There is no use or creation of mem_dst after the creation of src and +-- before the current statement. Handle by calling getVarUsesBetween and +-- looking at both the original var-mem mappings *and* the new, temporary +-- ones. +-- +-- 4. src (the variable, not the memory) does not alias anything. Handle by +-- checking VarAliases. +-- +-- 5. The new index function of src only uses variables declared prior to the +-- first use of mem_src. Handle by first using curVarPrimExps and +-- ExpMem.substituteInIxFun to create a (possibly larger) index function that +-- uses earlier variables. Then use ctxVarsInUseBeforeMem to check that all +-- the variables in the new index function are available before the creation +-- of mem_src. +-- +-- If an array src0 has been coalesced into mem_src, handle that by *also* +-- checking src0 and mem_src0 where src and mem_src are checked. We choose to +-- coalesce in a top-down fashion, even though that might exclude some potential +-- coalescings -- however, doing it differently might exclude some other +-- potentials, so we just make a choice. +-- +-- We only coalesce src into dst if all eventual src0 can also be coalesced into +-- dst. It does not make sense to coalesce only part of them, since in that +-- case both memory blocks and related allocations will still be around. + +safetyCond1 :: MonadReader Context m => + VName -> MemorySrc -> m Bool +safetyCond1 dst mem_src = do + last_uses <- lookupEmptyable (FromStm dst) <$> asks ctxLastUses + let res = S.member (memSrcName mem_src) last_uses + return res + +safetyCond2 :: MonadReader Context m => + VName -> MemorySrc -> m Bool +safetyCond2 src mem_dst = do + allocs_before_src <- lookupEmptyable src + <$> asks ctxAllocatedBlocksBeforeCreation + let res = S.member (memSrcName mem_dst) allocs_before_src + return res + +safetyCond3 :: VName -> VName -> MemorySrc -> FindM lore Bool +safetyCond3 src dst mem_dst = do + fundef <- asks ctxFunDef + let uses_after_src_vars = S.toList $ getVarUsesBetween fundef src dst + uses_after_src <- mapM (maybe (return S.empty) withMemAliases + <=< lookupCurrentVarMem) uses_after_src_vars + return $ not $ S.member (memSrcName mem_dst) (S.unions uses_after_src) + +safetyCond4 :: MonadReader Context m => + VName -> m Bool +safetyCond4 src = do + -- Special If handling: An If can have aliases, but that can be okay and is + -- checked in safe If: It is okay for it to have one alias (one of the + -- branches), while two aliases are wrong. + if_handling <- isIfExp src + + -- Special Reshape handling: If a reshape has variables associated with it, it + -- is okay to use it. + src_actuals <- lookupEmptyable src <$> asks ctxActualVars + reshape_handling <- isReshapeExp src <&&> pure (not (S.null src_actuals)) + + -- This needs to be extended if support for e.g. reshape coalescing is wanted: + -- Some operations can be aliasing, but still be okay to coalesce if you also + -- coalesce their aliased sources. + src_aliases <- lookupEmptyable src <$> asks ctxVarAliases + let res = if_handling || reshape_handling || S.null src_aliases + return res + +safetyCond5 :: MonadReader Context m => + MemorySrc -> ExpMem.IxFun -> m Bool +safetyCond5 mem_src ixfun = do + in_use_before_mem_src <- lookupEmptyable (memSrcName mem_src) + <$> asks ctxVarsInUseBeforeMem + let used_vars = freeIn ixfun + res = all (`S.member` in_use_before_mem_src) $ S.toList used_vars + return res + +safetyIf :: VName -> VName -> FindM lore Bool +safetyIf src dst = do + -- Special handling: If src refers to an If expression, we need to check that + -- not just is mem_dst not used after src and before dst, but neither is any + -- other memory that will be merged after the coalescing. Normally this is + -- not an issue, since a coalescing means changing just one memory block -- + -- but in the case of an If expression, each branch can have its own memory + -- block, and both of them will try to be coalesced. This extra test only + -- applies to the actual memory blocks in the branches, not any existential + -- memory block in the If, which in any case will be "used" in both branches. + -- + -- See tests/coalescing/if/if-neg-3.fut for an example of where this should + -- fail. + mem_src <- lookupVarMem src + actual_srcs <- S.toList <$> lookupActualVars src + existentials <- asks ctxExistentials + var_to_mem <- asks ctxVarToMem + first_uses_all <- asks ctxFirstUses + + -- Find all variables that have 'src' as an actual var, and then check if one + -- of those is an If expression. + reverse_actual_srcs <- + S.toList . S.unions . M.elems . M.filter (src `S.member`) + <$> asks ctxActualVars + outer <- mapMaybeM ifExp reverse_actual_srcs + let (is_in_if, + if_branch_results_from_outer, + at_least_one_creation_inside) = case outer of + -- This is the if expression of which we are currently looking at one of + -- its branch results. + [Exp nctx nthpat (If _ body0 body1 _)] -> + let results_from_outer = S.fromList $ mapMaybe subExpVar + $ concatMap (drop nctx . bodyResult) + $ filter (null . bodyStms) [body0, body1] + + resultCreatedInside body se = fromMaybe False $ do + res <- subExpVar se + res_mem <- memSrcName <$> M.lookup res var_to_mem + let body_vars = concatMap (map patElemName . patternValueElements + . stmPattern) $ bodyStms body + body_first_uses = S.unions $ map (`lookupEmptyable` first_uses_all) + body_vars + return $ S.member res_mem body_first_uses + + at_least = resultCreatedInside body0 (bodyResult body0 !! (nctx + nthpat)) + || resultCreatedInside body1 (bodyResult body1 !! (nctx + nthpat)) + in (True, results_from_outer, at_least) + _ -> (False, S.empty, False) + + -- This success requirement is independent of whichever branch we are in right + -- now. We say that the results of an if-expression can be coalesced if the + -- branch-specific requirements hold *and* this general rule holds: Either the + -- If has no existentials (e.g. if it does in-place updates), or it has + -- existentials and at least one of the branches returns an array that was + -- created inside the branch. + let res_general = not is_in_if || (not (any (`S.member` existentials) actual_srcs) + || at_least_one_creation_inside) + + -- Check if the branch described by 'src' needs special handling. + let if_handling = + -- We are sure this is an if. This might not actually be necessary. + is_in_if + -- This does not refer to the result of a branch where the array is + -- created outside the if. It is a requirement that there is at most + -- one such branch. The extra safety here only relates to branches + -- whose result arrays are created inside. + && not (any (`S.member` if_branch_results_from_outer) actual_srcs) + -- Ignore existentials as well. + && not (src `S.member` existentials) + + -- This success requirement is part is specific to this branch. + res_current <- + if if_handling + then do + -- Get the memory used in the other branch. Use a reverse lookup. + mem_actual_srcs <- L.nub <$> mapM lookupVarMem reverse_actual_srcs + let mem_actual_srcs_cur = L.delete mem_src mem_actual_srcs + and <$> mapM (safetyCond3 src dst) mem_actual_srcs_cur + else return True + + -- The full result. + let res = res_general && res_current + return res diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Exps.hs b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Exps.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Exps.hs @@ -0,0 +1,70 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE ExistentialQuantification #-} +-- | Get a mapping from statement patterns to statement expression for all +-- statements. +module Futhark.Optimise.MemoryBlockMerging.Coalescing.Exps + ( Exp'(..) + , findExpsFunDef + ) where + +import qualified Data.Map.Strict as M +import Control.Monad +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemorish) +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + +-- | Describes the nth pattern and the statement expression. +data Exp' = forall lore. Annotations lore => Exp Int Int (Exp lore) +instance Show Exp' where + show (Exp _nctxpatters _nthvalpattern e) = show e + +type Exps = M.Map VName Exp' + +newtype FindM lore a = FindM { unFindM :: Writer Exps a } + deriving (Monad, Functor, Applicative, + MonadWriter Exps) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +findExpsFunDef :: LoreConstraints lore => + FunDef lore -> Exps +findExpsFunDef fundef = + let m = unFindM $ lookInBody $ funDefBody fundef + res = execWriter m + in res + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern patctxelems patvalelems) _ e) = do + forM_ (zip patvalelems [0..]) $ \(PatElem var _, i) -> + tell $ M.singleton var $ Exp (length patctxelems) i e + + -- Recursive body walk. + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + } diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition2.hs b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition2.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition2.hs @@ -0,0 +1,110 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find safety condition 2 for all statements. +module Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition2 + ( findSafetyCondition2FunDef + ) where + +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Control.Monad +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory ( + ExplicitMemory, InKernel, ExplicitMemorish) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Types +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +type CurrentAllocatedBlocks = MNames +type AllocatedBlocksBeforeCreation = M.Map VName MNames + +newtype FindM lore a = FindM { unFindM :: RWS () + AllocatedBlocksBeforeCreation CurrentAllocatedBlocks a } + deriving (Monad, Functor, Applicative, + MonadWriter AllocatedBlocksBeforeCreation, + MonadState CurrentAllocatedBlocks) + +type LoreConstraints lore = (ExplicitMemorish lore, + IsAlloc lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +findSafetyCondition2FunDef :: FunDef ExplicitMemory + -> AllocatedBlocksBeforeCreation +findSafetyCondition2FunDef fundef = + let m = unFindM $ do + forM_ (funDefParams fundef) lookInFParam + lookInBody $ funDefBody fundef + res = snd $ evalRWS m () S.empty + in res + +lookInFParam :: FParam ExplicitMemory -> FindM lore () +lookInFParam (Param _ membound) = + -- Unique array function parameters also count as "allocations" in which + -- memory can be coalesced. + case membound of + ExpMem.MemArray _ _ Unique (ExpMem.ArrayIn mem _) -> + modify $ S.insert mem + _ -> return () + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern patctxelems patvalelems) _ e) = do + let new_decls0 = map patElemName (patctxelems ++ patvalelems) + new_decls1 = case e of + DoLoop _mergectxparams mergevalparams _loopform _body -> + -- Technically not a declaration for the current expression, but very + -- close, and hopefully okay to consider it as one. + map (paramName . fst) mergevalparams + _ -> [] + new_decls = new_decls0 ++ new_decls1 + + cur_allocated_blocks <- get + forM_ new_decls $ \x -> + tell $ M.singleton x cur_allocated_blocks + + case patvalelems of + [PatElem mem _] -> + when (isAlloc e) $ modify $ S.insert mem + _ -> return () + + -- RECURSIVE BODY WALK. + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + , walkOnFParam = lookInFParam + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + +class IsAlloc lore where + isAlloc :: Exp lore -> Bool + +instance IsAlloc ExplicitMemory where + isAlloc (Op ExpMem.Alloc{}) = True + isAlloc _ = False + +instance IsAlloc InKernel where + isAlloc _ = False diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition3.hs b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition3.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition3.hs @@ -0,0 +1,136 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Safety condition 3 verification. +module Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition3 + ( getVarUsesBetween + ) where + +import qualified Data.Set as S +import qualified Data.List as L +import Control.Monad +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory ( + ExplicitMemory, ExplicitMemorish) +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +data Context = Context + { ctxSource :: VName + , ctxDestination :: VName + } + deriving (Show) + +data Current = Current + { curHasReachedSource :: Bool + , curHasReachedDestination :: Bool + , curVars :: Names + } + deriving (Show) + +newtype FindM lore a = FindM { unFindM :: RWS Context () Current a } + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadState Current) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +modifyCurVars :: (Names -> Names) -> FindM lore () +modifyCurVars f = modify $ \c -> c { curVars = f $ curVars c } + +-- Find all the variables present between the creations of two variables (not +-- inclusive). +getVarUsesBetween :: FunDef ExplicitMemory + -> VName -> VName + -> Names +getVarUsesBetween fundef src dst = + let context = Context src dst + m = unFindM $ lookInBody $ funDefBody fundef + res = curVars $ fst $ execRWS m context (Current False False S.empty) + in res + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm stm@(Let _ _ e) = do + let new_decls = newDeclarationsStm stm + + dst <- asks ctxDestination + when (dst `L.elem` new_decls) + $ modify $ \c -> c { curHasReachedDestination = True } + + is_after_source <- gets curHasReachedSource + is_before_destination <- gets curHasReachedDestination + + unless is_before_destination $ do + let e_free_vars = freeInExp e + e_used_vars = S.union e_free_vars (S.fromList new_decls) + + -- If the source has been created, add the newly used variables. + -- + -- Note that "used after creation" refers both to used in subsequent + -- statements AND any statements in any sub-bodies (if and loop). + when is_after_source + $ modifyCurVars $ S.union e_used_vars + + -- If the source is present in the declarations, state that it has been + -- created. + src <- asks ctxSource + when (src `L.elem` new_decls) + $ modify $ \c -> c { curHasReachedSource = True } + + -- RECURSIVE BODY WALK. + case e of + If _ body0 body1 _ -> do + -- This is not very If-specific, but rather specific to expressions with + -- multiple, independent bodies, where If is just the only such + -- expression. + -- + -- We do not want the state (for safety condition 3) after traversing + -- the first branch to be present when traversing the second branch, + -- since they really will never both be run, so we compute them + -- independently and then merge them at the end. + before <- get + lookInBody body0 + after0 <- get + put Current { curHasReachedSource = curHasReachedSource before + , curHasReachedDestination = curHasReachedDestination after0 + , curVars = curVars before + } + lookInBody body1 + after1 <- get + put Current { curHasReachedSource = + curHasReachedSource after0 || curHasReachedSource after1 + , curHasReachedDestination = + curHasReachedDestination after0 || curHasReachedDestination after1 + , curVars = + S.union (curVars after0) (curVars after1) + } + _ -> do + -- In the general case, just look through any 'Body' you can find. (This + -- is the case for loops.) + let walker = identityWalker { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + fullWalkExpM walker walker_kernel e diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition5.hs b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition5.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition5.hs @@ -0,0 +1,120 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find safety condition 5 for all statements. +module Futhark.Optimise.MemoryBlockMerging.Coalescing.SafetyCondition5 + ( findSafetyCondition5FunDef + ) where + +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Control.Monad +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory ( + InKernel, ExplicitMemory, ExplicitMemorish) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Types +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +type DeclarationsSoFar = Names +type VarsInUseBeforeMem = M.Map MName Names + +newtype FindM lore a = FindM { unFindM :: RWS FirstUses + VarsInUseBeforeMem DeclarationsSoFar a } + deriving (Monad, Functor, Applicative, + MonadReader FirstUses, + MonadWriter VarsInUseBeforeMem, + MonadState DeclarationsSoFar) + +type LoreConstraints lore = (ExplicitMemorish lore, + ExtractKernelDefVars lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +findSafetyCondition5FunDef :: FunDef ExplicitMemory -> FirstUses + -> VarsInUseBeforeMem +findSafetyCondition5FunDef fundef first_uses = + let m = unFindM $ do + forM_ (funDefParams fundef) lookInFParam + lookInBody $ funDefBody fundef + res = snd $ evalRWS m first_uses S.empty + in res + +lookInFParam :: FParam lore -> FindM lore () +lookInFParam (Param x _) = + modify $ S.insert x + +lookInLParam :: LParam lore -> FindM lore () +lookInLParam (Param x _) = + modify $ S.insert x + +lookInBody :: LoreConstraints lore => Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => Stm lore -> FindM lore () +lookInStm stm@(Let _ _ e) = do + let new_decls = newDeclarationsStm stm + + first_uses <- ask + declarations_so_far <- get + forM_ (S.toList $ S.unions $ map (`lookupEmptyable` first_uses) new_decls) $ \mem -> + tell $ M.singleton mem declarations_so_far + + forM_ new_decls $ \x -> + modify $ S.insert x + + -- Special loop handling: Extract useful variables that are in use. + case e of + DoLoop _ _ loopform _ -> + case loopform of + ForLoop i _ _ _ -> modify $ S.insert i + WhileLoop c -> modify $ S.insert c + _ -> return () + + modify $ S.union (extractKernelDefVars e) + + -- Recursive body walk. + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + , walkOnFParam = lookInFParam + , walkOnLParam = lookInLParam + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInLambda + , walkOnKernelLParam = lookInLParam + } + +lookInLambda :: LoreConstraints lore => + Lambda lore -> FindM lore () +lookInLambda (Lambda params body _) = do + forM_ params lookInLParam + lookInBody body + +class ExtractKernelDefVars lore where + -- Extract variables from a kernel definition. + extractKernelDefVars :: Exp lore -> Names + +instance ExtractKernelDefVars ExplicitMemory where + extractKernelDefVars (Op (ExpMem.Inner (Kernel _ kernelspace _ _))) = + S.fromList $ map ($ kernelspace) + [spaceGlobalId, spaceLocalId, spaceGroupId] + extractKernelDefVars _ = S.empty + +instance ExtractKernelDefVars InKernel where + extractKernelDefVars _ = S.empty diff --git a/src/Futhark/Optimise/MemoryBlockMerging/CrudeMovingUp.hs b/src/Futhark/Optimise/MemoryBlockMerging/CrudeMovingUp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/CrudeMovingUp.hs @@ -0,0 +1,263 @@ +-- | Move variables as much as possible upwards in a program. +module Futhark.Optimise.MemoryBlockMerging.CrudeMovingUp + ( moveUpInFunDef + ) where + +import qualified Data.Set as S +import qualified Data.List as L +import qualified Data.Map.Strict as M +import Data.Maybe (mapMaybe) +import Control.Monad +import Control.Monad.RWS +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + +import Control.Monad.State +import Control.Monad.Identity + + +type Line = Int +data Origin = FromFParam + | FromLine Line (Exp ExplicitMemory) + deriving (Eq, Ord, Show) + +-- The dependencies and the location. +data PrimBinding = PrimBinding { pbFrees :: Names + , _pbConsumed :: Names + , pbOrigin :: Origin + } + deriving (Show) + +-- A mapping from names to PrimBinding. The key is a collection of names, since +-- a statement can have multiple patterns. +type BindingMap = [(Names, PrimBinding)] + +-- | Call 'findHoistees' for every body, and then hoist every one of the found +-- hoistees (variables). +moveUpInFunDef :: FunDef ExplicitMemory + -> (Body ExplicitMemory -> Maybe [FParam ExplicitMemory] -> [VName]) + -> FunDef ExplicitMemory +moveUpInFunDef fundef findHoistees = + let scope_new = scopeOf fundef + bindingmap_cur = [] + body' = hoistInBody scope_new bindingmap_cur + (Just (funDefParams fundef)) findHoistees (funDefBody fundef) + fundef' = fundef { funDefBody = body' } + in fundef' + +lookupPrimBinding :: VName -> State BindingMap PrimBinding +lookupPrimBinding vname = + gets $ snd . fromJust (pretty vname ++ " was not found in BindingMap." + ++ " This should not happen!") + . L.find ((vname `S.member`) . fst) + +namesDependingOn :: VName -> State BindingMap Names +namesDependingOn v = + gets $ S.unions . map fst . filter (\(_, pb) -> v `S.member` pbFrees pb) + +scopeBindingMap :: (VName, NameInfo ExplicitMemory) + -> BindingMap +scopeBindingMap (x, _) = [(S.singleton x, PrimBinding S.empty S.empty FromFParam)] + +-- Find all variables bound in a KernelSpace. +boundInKernelSpace :: ExpMem.KernelSpace -> Names +boundInKernelSpace space = + -- This might do too much. + S.fromList ([ ExpMem.spaceGlobalId space + , ExpMem.spaceLocalId space + , ExpMem.spaceGroupId space] + ++ (case ExpMem.spaceStructure space of + ExpMem.FlatThreadSpace ts -> + map fst ts ++ mapMaybe (subExpVar . snd) ts + ExpMem.NestedThreadSpace ts -> + map (\(x, _, _, _) -> x) ts + ++ mapMaybe (subExpVar . (\(_, x, _, _) -> x)) ts + ++ map (\(_, _, x, _) -> x) ts + ++ mapMaybe (subExpVar . (\(_, _, _, x) -> x)) ts + )) + +-- FIXME: The results of this should maybe go in the core 'freeIn' function, or +-- perhaps the ExplicitMemory module, instead of this arbitrary module. +boundInExpExtra :: Exp ExplicitMemory -> Names +boundInExpExtra = execWriter . inExp + where inExp :: Exp ExplicitMemory -> Writer Names () + inExp e = case e of + Op (ExpMem.Inner (ExpMem.Kernel _ space _ _)) -> + tell $ boundInKernelSpace space + _ -> walkExpM walker e + + walker = identityWalker { + walkOnBody = mapM_ (inExp . stmExp) . bodyStms + } + +bodyBindingMap :: [Stm ExplicitMemory] -> BindingMap +bodyBindingMap stms = + concatMap createBindingStmt $ zip [0..] stms + -- We do not need to run this recursively on any sub-bodies, since this will + -- be run for every call to hoistInBody, which *does* run recursively on + -- sub-bodies. + + where createBindingStmt :: (Line, Stm ExplicitMemory) + -> BindingMap + createBindingStmt (line, stmt@(Let (Pattern patctxelems patvalelems) _ e)) = + let stmt_vars = S.fromList (map patElemName (patctxelems ++ patvalelems)) + frees = freeInStm stmt + consumed = case e of BasicOp (Update src _ _) -> S.singleton src + _ -> mempty + bound_extra = boundInExpExtra e + frees' = frees `S.difference` bound_extra + vars_binding = (stmt_vars, PrimBinding frees' consumed (FromLine line e)) + + -- Some variables exist only in a shape declaration. + shape_sizes = S.fromList $ concatMap shapeSizes (patctxelems ++ patvalelems) + sizes_binding = (shape_sizes, PrimBinding frees' consumed (FromLine line e)) + + -- Some expressions contain special identifiers that are used in a + -- body. This should go somewhere else than here. + param_vars = case e of + Op (ExpMem.Inner (ExpMem.Kernel _ space _ _)) -> + boundInKernelSpace space + _ -> S.empty + params_binding = (param_vars, PrimBinding S.empty S.empty FromFParam) + + bmap = [vars_binding, sizes_binding, params_binding] + in bmap + + shapeSizes (PatElem _ (ExpMem.MemArray _ shape _ _)) = + mapMaybe subExpVar $ shapeDims shape + shapeSizes _ = [] + +hoistInBody :: Scope ExplicitMemory + -> BindingMap + -> Maybe [FParam ExplicitMemory] + -> (Body ExplicitMemory -> Maybe [FParam ExplicitMemory] -> [VName]) + -> Body ExplicitMemory + -> Body ExplicitMemory +hoistInBody scope_new bindingmap_old params findHoistees body = + let hoistees = findHoistees body params + + -- We use the possibly non-empty scope to extend our BindingMap. + bindingmap_fromscope = concatMap scopeBindingMap $ M.toList scope_new + bindingmap_body = bodyBindingMap $ stmsToList $ bodyStms body + bindingmap = bindingmap_old ++ bindingmap_fromscope ++ bindingmap_body + + -- Create a new body where all hoistees have been moved as much upwards in + -- the statement list as possible. + (Body () bnds res, bindingmap') = + foldl (\(body0, lbindingmap) -> hoist lbindingmap body0) + (body, bindingmap) hoistees + + -- Touch upon any subbodies. + bnds' = fmap (hoistRecursivelyStm bindingmap' findHoistees) bnds + body' = Body () bnds' res + + in body' + +hoistRecursivelyStm :: BindingMap + -> (Body ExplicitMemory -> Maybe [FParam ExplicitMemory] -> [VName]) + -> Stm ExplicitMemory + -> Stm ExplicitMemory +hoistRecursivelyStm bindingmap findHoistees (Let pat aux e) = + runIdentity (Let pat aux <$> mapExpM transform e) + + where transform = identityMapper { mapOnBody = mapper } + mapper scope_new = return . hoistInBody scope_new bindingmap' Nothing findHoistees + -- The nested body cannot move to any of its locations of its parent's + -- body, so we say that all its parent's bindings are parameters. + bindingmap' = map (\(ns, PrimBinding frees consumed _) -> + (ns, PrimBinding frees consumed FromFParam)) + bindingmap + +-- Hoist the statement denoted by 'hoistee' as much upwards as possible in +-- 'body', and return the new body. +hoist :: BindingMap + -> Body ExplicitMemory + -> VName + -> (Body ExplicitMemory, BindingMap) +hoist bindingmap_cur body hoistee = + let bindingmap = bindingmap_cur <> bodyBindingMap (stmsToList $ bodyStms body) + + body' = runState (moveLetUpwards hoistee body) bindingmap + + in body' + +-- Move a statement as much up as possible. +moveLetUpwards :: VName -> Body ExplicitMemory + -> State BindingMap (Body ExplicitMemory) +moveLetUpwards letname body = do + PrimBinding deps consumed letorig <- lookupPrimBinding letname + + -- Extend the dependencies with all those statements that use the consumed + -- variables of this statement, except the current statement. + deps' <- S.delete letname + <$> (S.union deps + <$> (S.unions <$> mapM namesDependingOn (S.toList consumed))) + + case letorig of + FromFParam -> return body + FromLine line_cur exp_cur -> + case exp_cur of + -- We do not want to change the structure of the program too much, so we + -- restrict the aggressive hoister to *stop* and not hoist loops and + -- kernels, as hoisting these expressions might actually make a + -- hoisting-dependent optimisation *poorer* because of some assumptions + -- about the structure. FIXME: Do this nicer in a way where it is easy + -- to argue for it. + DoLoop{} -> return body + Op ExpMem.Inner{} -> return body + _ -> do + -- Sort by how close they are to the beginning of the body. The closest + -- one should be the first one to hoist, so that the other ones can maybe + -- exploit it. + deps'' <- sortByKeyM (fmap pbOrigin . lookupPrimBinding) + $ S.toList deps' + body' <- foldM (flip moveLetUpwards) body deps'' + origins <- mapM (fmap pbOrigin . lookupPrimBinding) deps'' + let line_dest = case foldl max FromFParam origins of + FromFParam -> 0 + FromLine n _e -> n + 1 + + PrimBinding _ _ letorig' <- lookupPrimBinding letname + when (letorig' /= letorig) $ error "Assertion: This should not happen." + + stms' <- moveLetToLine letname line_cur line_dest $ stmsToList $ bodyStms body' + + return body' { bodyStms = stmsFromList stms' } + +-- Both move the statement to the new line and update the BindingMap. +moveLetToLine :: VName -> Line -> Line -> [Stm ExplicitMemory] + -> State BindingMap [Stm ExplicitMemory] +moveLetToLine stm_cur_name line_cur line_dest stms + | line_cur == line_dest = return stms + | otherwise = do + + let stm_cur = stms !! line_cur + stms1 = take line_cur stms ++ drop (line_cur + 1) stms + stms2 = take line_dest stms1 ++ [stm_cur] ++ drop line_dest stms1 + + modify $ map (\t@(ns, PrimBinding frees consumed orig) -> + case orig of + FromFParam -> t + FromLine l e -> if l >= line_dest && l < line_cur + then (ns, PrimBinding frees consumed + (FromLine (l + 1) e)) + else t) + + r <- lookupPrimBinding stm_cur_name + case r of + PrimBinding frees consumed (FromLine _ exp_cur) -> + modify $ replaceWhere stm_cur_name (PrimBinding frees consumed + (FromLine line_dest exp_cur)) + _ -> error "moveLetToLine: unhandled case" -- fixme + return stms2 + +replaceWhere :: VName -> PrimBinding -> BindingMap -> BindingMap +replaceWhere n pb1 = + map (\(ns, pb) -> (ns, if n `S.member` ns + then pb1 + else pb)) diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Existentials.hs b/src/Futhark/Optimise/MemoryBlockMerging/Existentials.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Existentials.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find all existential variables. +module Futhark.Optimise.MemoryBlockMerging.Existentials + ( findExistentials + ) where + +import qualified Data.Set as S +import qualified Data.List as L +import Control.Monad +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemorish) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +newtype FindM lore a = FindM { unFindM :: Writer Names a } + deriving (Monad, Functor, Applicative, + MonadWriter Names) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +record :: VName -> FindM lore () +record = tell . S.singleton + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +findExistentials :: LoreConstraints lore => + FunDef lore -> Names +findExistentials fundef = + let m = unFindM $ lookInBody $ funDefBody fundef + existentials = execWriter m + in existentials + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern patctxelems patvalelems) _ e) = do + forM_ patvalelems $ \(PatElem var membound) -> + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _) -> + when (mem `L.elem` map patElemName patctxelems) + $ record var + _ -> return () + + case e of + DoLoop mergectxparams mergevalparams _loopform _body -> + forM_ mergevalparams $ \(Param var membound, _) -> + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _) -> + when (mem `L.elem` map (paramName . fst) mergectxparams) + $ record var + _ -> return () + _ -> return () + + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Liveness/FirstUse.hs b/src/Futhark/Optimise/MemoryBlockMerging/Liveness/FirstUse.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Liveness/FirstUse.hs @@ -0,0 +1,199 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find first uses for all memory blocks. +-- +-- Array creation points. Maps statements to memory block names. +-- +-- A memory block can have more than one first use. +module Futhark.Optimise.MemoryBlockMerging.Liveness.FirstUse + ( findFirstUses + , createsNewArrayBase + ) where + +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) +import Control.Monad +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + (ExplicitMemory, InKernel, ExplicitMemorish) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types + + +data Context = Context + { ctxVarToMem :: VarMemMappings MemorySrc + , ctxMemAliases :: MemAliases + , ctxCurOuterFirstUses :: Names + -- ^ First uses found in outer bodies. + } + deriving (Show) + +newtype FindM lore a = FindM { unFindM :: RWS Context () FirstUses a } + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadWriter (), + MonadState FirstUses) + +type LoreConstraints lore = (ExplicitMemorish lore, + ArrayUtils lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- Find the memory blocks used or aliased by a variable. +varMems :: VName -> FindM lore MNames +varMems var = do + var_to_mem <- asks ctxVarToMem + mem_aliases <- asks ctxMemAliases + return $ fromMaybe S.empty $ do + mem <- memSrcName <$> M.lookup var var_to_mem + return $ S.union (S.singleton mem) $ lookupEmptyable mem mem_aliases + +recordMapping :: VName -> MName -> FindM lore () +recordMapping stmt_var mem = + modify $ M.unionWith S.union (M.singleton stmt_var $ S.singleton mem) + +-- | Find all first uses of *memory blocks* in a function definition. +findFirstUses :: VarMemMappings MemorySrc -> MemAliases + -> FunDef ExplicitMemory -> FirstUses +findFirstUses var_to_mem mem_aliases fundef = + let context = Context { ctxVarToMem = var_to_mem + , ctxMemAliases = mem_aliases + , ctxCurOuterFirstUses = S.empty + } + m = unFindM $ do + forM_ (funDefParams fundef) lookInFunDefFParam + lookInBody $ funDefBody fundef + first_uses = removeEmptyMaps $ expandWithAliases mem_aliases + $ fst $ execRWS m context M.empty + in first_uses + +lookInFunDefFParam :: LoreConstraints lore => + FParam lore -> FindM lore () +lookInFunDefFParam (Param x (ExpMem.MemArray _ _ _ (ExpMem.ArrayIn xmem _))) = + recordMapping x xmem +lookInFunDefFParam _ = return () + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern patctxelems patvalelems) _ e) = do + outer_first_uses <- asks ctxCurOuterFirstUses + when (createsNewArray e) $ do + let e_free_vars = freeInExp e + e_mems <- S.unions <$> mapM varMems (S.toList e_free_vars) + forM_ patvalelems $ \(PatElem x membound) -> + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn xmem _) -> do + x_mems <- varMems xmem + + -- For the first use to be a proper first use, it must write to + -- the memory, but not read from it. We need to check this to + -- support multiple liveness intervals. If we don't check this, + -- the last use analysis and the interference analysis might end + -- up wrong. + when (S.null $ S.intersection x_mems e_mems) + -- We only record the mapping between the statement and the + -- memory block, not any of its aliased memory blocks. They + -- would not be aliased unless they are themselves created at + -- some point, so they will get their own FirstUses. Putting + -- them into first use here would probably also be too + -- conservative. + -- + -- If it is a first use of a memory inside a loop or a kernel, and + -- that memory already has a first use outside the loop, ignore it, + -- since it is not a proper first use. This can be an issue after + -- the coalescing transformation, where multidimensional maps are + -- first-order-transformed into nested loops, each loop having its + -- own Scratch expression. FIXME: This might be too conservative + -- for multiple liveness intervals, but it does not seem to be a + -- problem with our tests. It is quite possible that this case only + -- occurs because the coalescing pass does not remove the inner + -- scratches, so maybe it should be fixed there. + $ unless (xmem `S.member` outer_first_uses) + $ recordMapping x xmem + _ -> return () + + -- Find first uses of existential memory blocks. Fairly conservative. + -- Covers the case where a loop uses multiple arrays by saying every + -- existential memory block overlaps with every result memory block. Fine + -- for now. + forM_ patctxelems + $ \p -> forM_ patvalelems + $ \el -> lookInPatCtxElem (patElemName el) p + case e of + DoLoop mergectxparams _mergevalparams _loopform _body -> + forM_ mergectxparams + $ \p -> forM_ patvalelems + $ \el -> lookInMergeCtxParam (patElemName el) p + _ -> return () + + cur_first_uses <- get + local (\ctx -> ctx { ctxCurOuterFirstUses = S.unions $ M.elems cur_first_uses }) + $ fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + +lookInPatCtxElem :: LoreConstraints lore => + VName -> PatElem lore -> FindM lore () +lookInPatCtxElem x (PatElem xmem ExpMem.MemMem{}) = + recordMapping x xmem +lookInPatCtxElem _ _ = return () + +lookInMergeCtxParam :: LoreConstraints lore => + VName -> (FParam lore, SubExp) -> FindM lore () +lookInMergeCtxParam x (Param xmem ExpMem.MemMem{}, _) = + recordMapping x xmem +lookInMergeCtxParam _ _ = return () + +class ArrayUtils lore where + -- Does an expression constitute a new array? + createsNewArray :: Exp lore -> Bool + +createsNewArrayBase :: Exp lore -> Bool +createsNewArrayBase e = case e of + BasicOp Partition{} -> True + BasicOp Replicate{} -> True + BasicOp Iota{} -> True + BasicOp Manifest{} -> True + BasicOp Copy{} -> True + BasicOp Concat{} -> True + BasicOp ArrayLit{} -> True + BasicOp Scratch{} -> True + _ -> False + +instance ArrayUtils ExplicitMemory where + createsNewArray e = case e of + Op (ExpMem.Inner ExpMem.Kernel{}) -> True + _ -> createsNewArrayBase e + +instance ArrayUtils InKernel where + createsNewArray e = case e of + Op (ExpMem.Inner ExpMem.GroupReduce{}) -> True + Op (ExpMem.Inner ExpMem.GroupScan{}) -> True + Op (ExpMem.Inner ExpMem.GroupStream{}) -> True + Op (ExpMem.Inner ExpMem.Combine{}) -> True + _ -> createsNewArrayBase e diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Liveness/Interference.hs b/src/Futhark/Optimise/MemoryBlockMerging/Liveness/Interference.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Liveness/Interference.hs @@ -0,0 +1,520 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE LambdaCase #-} +-- | Find memory block interferences. Maps a memory block to its interference +-- set. + +module Futhark.Optimise.MemoryBlockMerging.Liveness.Interference + ( findInterferences + ) where + +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import qualified Data.List as L +import Data.Maybe (mapMaybe, fromMaybe, catMaybes) +import Control.Monad +import Control.Monad.RWS +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory ( + ExplicitMemorish, ExplicitMemory, InKernel) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types + + +data Context = Context { ctxVarToMem :: VarMemMappings MemorySrc + , ctxMemAliases :: MemAliases + , ctxFirstUses :: FirstUses + , ctxLastUses :: LastUses + , ctxExistentials :: Names + , ctxLoopCorrespondingVar :: M.Map VName (VName, SubExp) + } + deriving (Show) + +type InterferencesList = [(MName, MNames)] + +getInterferencesMap :: InterferencesList -> Interferences +getInterferencesMap = M.unionsWith S.union . map (uncurry M.singleton) + +data Current = Current { curAlive :: MNames + + , curResPotentialKernelInterferences + :: PotentialKernelDataRaceInterferences + } + deriving (Show) + +newtype FindM lore a = FindM + { unFindM :: RWS Context InterferencesList Current a } + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadWriter InterferencesList, + MonadState Current) + +type LoreConstraints lore = (ExplicitMemorish lore, + KernelInterferences lore, + SpecialBodyExceptions lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +awaken :: MName -> FindM lore () +awaken mem = modifyCurAlive $ S.insert mem + +kill :: MName -> FindM lore () +kill mem = modifyCurAlive $ S.delete mem + +modifyCurAlive :: (MNames -> MNames) -> FindM lore () +modifyCurAlive f = modify $ \c -> c { curAlive = f $ curAlive c } + +addPotentialKernelInterferenceGroup :: + PotentialKernelDataRaceInterferenceGroup -> FindM lore () +addPotentialKernelInterferenceGroup set = + modify $ \c -> c { curResPotentialKernelInterferences = + curResPotentialKernelInterferences c ++ [set] } + +recordCurrentInterferences :: FindM lore () +recordCurrentInterferences = do + current <- gets curAlive + -- Interferences are commutative. Reflect that in the resulting data. + forM_ (S.toList current) $ \mem -> + tell [(mem, current)] + +recordNewInterferences :: MNames -> FindM lore () +recordNewInterferences mems_in_stm = do + current <- gets curAlive + -- Interferences are commutative. Reflect that in the resulting data. + forM_ (S.toList current) $ \mem -> + tell [(mem, mems_in_stm)] + forM_ (S.toList mems_in_stm) $ \mem -> + tell [(mem, current)] + +-- | Find all memory block interferences in a function definition. +findInterferences :: VarMemMappings MemorySrc -> MemAliases -> + FirstUses -> LastUses -> Names -> FunDef ExplicitMemory + -> (Interferences, PotentialKernelDataRaceInterferences) +findInterferences var_to_mem mem_aliases first_uses last_uses existentials fundef = + let context = Context { ctxVarToMem = var_to_mem + , ctxMemAliases = mem_aliases + , ctxFirstUses = first_uses + , ctxLastUses = last_uses + , ctxExistentials = existentials + , ctxLoopCorrespondingVar = M.empty + } + m = unFindM $ do + forM_ (funDefParams fundef) lookInFunDefFParam + lookInBody $ funDefBody fundef + (cur, interferences_list) = execRWS m context (Current S.empty []) + interferences = removeEmptyMaps $ removeKeyFromMapElems $ makeCommutativeMap + $ getInterferencesMap interferences_list + potential_kernel_interferences = curResPotentialKernelInterferences cur + in (interferences, potential_kernel_interferences) + +lookInFunDefFParam :: FParam lore -> FindM lore () +lookInFunDefFParam (Param var _) = do + first_uses_var <- lookupEmptyable var <$> asks ctxFirstUses + mapM_ awaken $ S.toList first_uses_var + recordCurrentInterferences + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds res) = do + mapM_ lookInStm bnds + lookInRes res + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds res) = do + mapM_ lookInStm bnds + lookInRes $ map kernelResultSubExp res + +awakenFirstUses :: [PatElem lore] -> FindM lore () +awakenFirstUses patvalelems = + forM_ patvalelems $ \(PatElem var _) -> do + first_uses_var <- lookupEmptyable var <$> asks ctxFirstUses + mapM_ awaken $ S.toList first_uses_var + +isNoOp :: Exp lore -> Bool +isNoOp (BasicOp bop) = case bop of + Scratch{} -> True + _ -> False +isNoOp _ = False + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm stm@(Let (Pattern _patctxelems patvalelems) _ e) + | isNoOp e = + awakenFirstUses patvalelems + -- There is no reason to record interferences if the current statement will + -- not generate any code in the end. We have this check to use the result + -- index sharing analysis on loop bodies and not get bogged down by the + -- result of a Scratch statement hanging around. + | otherwise = do + awakenFirstUses patvalelems + ctx <- ask + let ctx' = ctx { ctxLoopCorrespondingVar = + M.union (ctxLoopCorrespondingVar ctx) + (findLoopCorrespondingVar ctx stm) + } + let stm_exceptions = fromMaybe [] $ do + indices <- specialBodyIndices e + let walker_exc = + identityWalker + { walkOnBody = \body -> let (body', lcv) = innermostLoopNestBody ctx body + ctx'' = ctx' { ctxLoopCorrespondingVar = + M.union (ctxLoopCorrespondingVar ctx') lcv } + in tell $ interferenceExceptions ctx'' + (bodyStms body') (bodyResult body') + indices Nothing } + walker_kernel_exc = + identityKernelWalker + { walkOnKernelBody = \body -> let (body', lcv) = innermostLoopNestBody ctx body + ctx'' = ctx' { ctxLoopCorrespondingVar = + M.union (ctxLoopCorrespondingVar ctx') lcv } + in tell $ interferenceExceptions ctx'' + (bodyStms body') (bodyResult body') + indices Nothing + , walkOnKernelKernelBody = \kbody -> tell $ interferenceExceptions ctx' + (kernelBodyStms kbody) + (mapMaybe (\case + ThreadsReturn _ se -> Just se + _ -> Nothing) + $ kernelBodyResult kbody) + indices + (specialBodyWriteMems stm) + } + return $ execWriter $ fullWalkExpM walker_exc walker_kernel_exc e + + first_uses <- asks ctxFirstUses + last_uses <- asks ctxLastUses + let stm_mems = + S.unions $ map (\pelem -> + let v = patElemName pelem + in S.union + (lookupEmptyable v first_uses) + (lookupEmptyable (FromStm v) last_uses)) patvalelems + + ((), stm_interferences) <- censor (const []) $ listen $ do + recordNewInterferences stm_mems + local (const ctx') $ fullWalkExpM walker walker_kernel e + let stm_interferences' = + map (\(k, vs) -> + (k, S.fromList + $ filter (\v -> not ((k, v) `L.elem` stm_exceptions + || (v, k) `L.elem` stm_exceptions)) + $ S.toList vs)) + stm_interferences + tell stm_interferences' + + potential_kernel_interferences <- findKernelDataRaceInterferences e + forM_ potential_kernel_interferences addPotentialKernelInterferenceGroup + + forM_ patvalelems $ \(PatElem var _) -> do + last_uses_var <- lookupEmptyable (FromStm var) <$> asks ctxLastUses + mapM_ kill last_uses_var + + where walker = identityWalker + { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + +-- For perfectly nested loops. Make it possible to find the index function for +-- the outer loop. +findLoopCorrespondingVar :: LoreConstraints lore => + Context -> Stm lore -> M.Map VName (VName, SubExp) +findLoopCorrespondingVar ctx (Let (Pattern _patctxelems patvalelems) _ + (DoLoop _ _ _ (Body _ stms res))) = + M.fromList $ catMaybes $ zipWith findIt patvalelems res + where findIt (PatElem pat_v (ExpMem.MemArray _ _ _ (ExpMem.ArrayIn pat_mem _))) (Var res_v) + | not (null stms) = case L.last $ stmsToList stms of + -- This is how the program looks after coalescing. + Let (Pattern _ [PatElem _last_v + (ExpMem.MemArray _ _ _ (ExpMem.ArrayIn last_stm_mem _))]) _ + (BasicOp (Update _ (DimFix slice_part : _) (Var copy_v))) -> + if pat_mem == last_stm_mem + then let res_v' = + if (memSrcName <$> M.lookup copy_v (ctxVarToMem ctx)) + == Just last_stm_mem + then Just copy_v + else Just res_v + in res_v' >>= \t -> Just (t, (pat_v, slice_part)) + -- Fix this mess. + else Nothing + _ -> Nothing + | otherwise = Nothing + findIt _ _ = Nothing +findLoopCorrespondingVar _ _ = M.empty + +innermostLoopNestBody :: LoreConstraints lore => + Context -> Body lore -> (Body lore, M.Map VName (VName, SubExp)) +innermostLoopNestBody ctx body = case stmsToList $ bodyStms body of + -- This checks for how perfect nested loops looks like after coalescing. This + -- is very brittle. If it detects such a nesting, it will ask the + -- interference exception algorithm to look in the innermost body. + Let _ _ (BasicOp Scratch{}) : loopstm@(Let _ _ (DoLoop _ _ _ body')) : _ -> + let (body'', loop_corresponding_var) = innermostLoopNestBody ctx body' + in (body'', M.union + (findLoopCorrespondingVar ctx loopstm) + loop_corresponding_var) + _ -> (body, M.empty) + +lookInRes :: [SubExp] -> FindM lore () +lookInRes ses = do + let vs = subExpVars ses + last_uses <- asks ctxLastUses + let last_uses_v = + S.unions $ map (\v -> lookupEmptyable (FromRes v) last_uses) vs + recordNewInterferences last_uses_v + mapM_ kill $ S.toList last_uses_v + +firstUsesInStm :: LoreConstraints lore => FirstUses -> + Stm lore -> [KernelFirstUse] +firstUsesInStm first_uses stm = + let m = lookFUInStm stm + in snd $ evalRWS m first_uses () + +firstUsesInExp :: LoreConstraints lore => + Exp lore -> FindM lore [KernelFirstUse] +firstUsesInExp e = do + let m = lookFUInExp e + first_uses <- asks ctxFirstUses + return $ snd $ evalRWS m first_uses () + +lookFUInStm :: LoreConstraints lore => + Stm lore -> RWS FirstUses [KernelFirstUse] () () +lookFUInStm (Let (Pattern _patctxelems patvalelems) _ e_stm) = do + forM_ patvalelems $ \(PatElem patname membound) -> + case membound of + ExpMem.MemArray pt _ _ (ExpMem.ArrayIn _ ixfun) -> do + fus <- lookupEmptyable patname <$> ask + forM_ fus $ \fu -> tell [(fu, patname, pt, ixfun)] + _ -> return () + lookFUInExp e_stm + +lookFUInExp :: LoreConstraints lore => + Exp lore -> RWS FirstUses [KernelFirstUse] () () +lookFUInExp = fullWalkExpM fu_walker fu_walker_kernel + where fu_walker = identityWalker + { walkOnBody = mapM_ lookFUInStm . bodyStms } + fu_walker_kernel = identityKernelWalker + { walkOnKernelBody = mapM_ lookFUInStm . bodyStms + , walkOnKernelKernelBody = mapM_ lookFUInStm . kernelBodyStms + , walkOnKernelLambda = mapM_ lookFUInStm . bodyStms . lambdaBody + } + +class KernelInterferences lore where + findKernelDataRaceInterferences :: + Exp lore -> FindM lore (Maybe PotentialKernelDataRaceInterferenceGroup) + +instance KernelInterferences ExplicitMemory where + findKernelDataRaceInterferences e = case e of + Op (ExpMem.Inner Kernel{}) -> Just <$> firstUsesInExp e + _ -> return Nothing + +instance KernelInterferences InKernel where + findKernelDataRaceInterferences _ = return Nothing + +-- Base info for kernel bodies. +class SpecialBodyExceptions lore where + specialBodyIndices :: Exp lore -> Maybe [MName] + specialBodyWriteMems :: Stm lore -> Maybe [(MName, ExpMem.IxFun, PrimType)] + +instance SpecialBodyExceptions ExplicitMemory where + specialBodyIndices (Op (ExpMem.Inner (Kernel _ kernelspace _ _))) = + Just $ map fst $ spaceDimensions kernelspace + specialBodyIndices e = specialBodyIndicesBase e + + specialBodyWriteMems (Let (Pattern _patctxelems patvalelems) _ + (Op (ExpMem.Inner Kernel{}))) = + Just $ mapMaybe (\p -> case patElemAttr p of + ExpMem.MemArray t _ _ (ExpMem.ArrayIn mem ixfun) -> Just (mem, ixfun, t) + _ -> Nothing) patvalelems + specialBodyWriteMems _ = Nothing + +instance SpecialBodyExceptions InKernel where + specialBodyIndices = specialBodyIndicesBase + specialBodyWriteMems = const Nothing + +specialBodyIndicesBase :: Exp lore -> Maybe [MName] +specialBodyIndicesBase (DoLoop _ _ (ForLoop i _ _ _) _) = Just [i] +specialBodyIndicesBase _ = Nothing + +-- Use first use analysis and last use analysis to find any exceptions to the +-- naive interference recorded for a statement. +interferenceExceptions :: LoreConstraints lore => + Context -> Stms lore -> [SubExp] -> [MName] -> + Maybe [(MName, ExpMem.IxFun, PrimType)] -> [(MName, MName)] +interferenceExceptions ctx stms res indices output_mems_may = + let output_vars = subExpVars res + indices_slice = map (DimFix . Var) indices + stms_first_uses = map (\(mem, _, _, _) -> mem) + $ concatMap (firstUsesInStm (ctxFirstUses ctx)) stms + results = + concat $ flip map (stmsToList stms) $ \(Let (Pattern _patctxelems patvalelems) _ e) -> + flip map patvalelems $ \(PatElem v membound) -> + let fromread = case e of + BasicOp (Index orig slice) -> do + orig_mem <- M.lookup orig $ ctxVarToMem ctx + if + -- These two extra requirements might be superfluous. + memSrcName orig_mem `L.notElem` stms_first_uses && + not (memSrcName orig_mem `S.member` ctxExistentials ctx) + then return (v, typeOf membound, orig_mem, slice) + else Nothing + _ -> Nothing + fromwrite = case e of + BasicOp Update{} + | ExpMem.MemArray pt _ _ _ <- membound -> do + -- The coalescing pass can have created a program where some + -- dependencies are a bit indirect. We find the core index function. + let (orig', slice') = + fixpointIterateMay + (\(v0, ss0) -> do + (v1, s1) <- M.lookup v0 (ctxLoopCorrespondingVar ctx) + return (v1, DimFix s1 : ss0)) + (v, []) + + orig_mem <- M.lookup orig' $ ctxVarToMem ctx + if + -- These two extra requirements might be superfluous. + memSrcName orig_mem `L.notElem` stms_first_uses && + not (memSrcName orig_mem `S.member` ctxExistentials ctx) + then return (v, Prim pt, orig_mem, slice') + else Nothing + _ -> Nothing + in (fromread, fromwrite) + fromreads = mapMaybe fst results + fromwrites = mapMaybe snd results + fromwrites' = filter (\(v, _, _, _) -> v `L.elem` output_vars) fromwrites + + fus_input_vars = M.fromList $ map (\(v, _, mem, _) -> + (v, S.singleton $ memSrcName mem)) fromreads + lus_input_vars = mapFromListSetUnion $ mapMaybe + (\(v, typ, mem, _) -> + let check e_pat = + let frees = freeInExp e_pat + + -- We need to handle scalars and arrays differently: A last + -- use of a scalar variable is the definite last use of the + -- memory it represents, while the last use of an array can + -- be distorted by reshapes and other aliasing operations, + -- so in that case we need to find the last use of the + -- memory block. + b = case typ of + Prim _ -> + v `S.member` frees + _ -> + memSrcName mem `L.elem` + mapMaybe ((memSrcName <$>) . (`M.lookup` ctxVarToMem ctx)) + (S.toList frees) + + in b + check' (Let _ _ e) = check e + in (\stm -> (FromStm $ patElemName $ head $ patternValueElements $ stmPattern stm, + S.singleton $ memSrcName mem)) <$> + L.find check' (reverse $ stmsToList stms)) fromreads + + -- 'Just' if in kernel, 'Nothing' otherwise. + fus_output_vars = mapFromListSetUnion $ case output_mems_may of + Just _ -> [] + _ -> map (\(v, _, mem, _) -> (v, S.singleton $ memSrcName mem)) fromwrites' + fus_result = mapFromListSetUnion $ case output_mems_may of + Just mems -> zip output_vars $ map (S.singleton . (\(mem, _, _) -> mem)) mems + _ -> [] + + -- Extended first uses and last uses. + fus = M.unionsWith S.union [ctxFirstUses ctx, fus_input_vars, fus_output_vars] + lus = M.unionsWith S.union [ctxLastUses ctx, lus_input_vars] + + -- Memory-to-slice mappings. + input_mem_slices = M.fromList $ map (\(_, _, mem, slice) -> + (memSrcName mem, slice)) fromreads + output_mem_slices = M.fromList $ case output_mems_may of + Just mems -> + map (\(mem, _, _) -> (mem, indices_slice)) mems + _ -> + map (\(_, _, mem, slice) -> (memSrcName mem, slice)) fromwrites' + mem_slices = M.union input_mem_slices output_mem_slices + + -- Memory-to-ixfun mappings. + input_mem_ixfuns = M.fromList $ map (\(_, _, mem, _) -> + (memSrcName mem, memSrcIxFun mem)) fromreads + output_mem_ixfuns = M.fromList $ case output_mems_may of + Just mems -> map (\(mem, ixfun, _) -> (mem, ixfun)) mems + _ -> map (\(_, _, mem, _) -> (memSrcName mem, memSrcIxFun mem)) fromwrites' + mem_ixfuns = M.union input_mem_ixfuns output_mem_ixfuns + + -- Memory-to-primtype-size mappings. + input_mem_primtypes = M.fromList + $ map (\(_, t, mem, _) -> (memSrcName mem, elemType t)) fromreads + output_mem_primtypes = M.fromList $ case output_mems_may of + Just mems -> map (\(mem, _, pt) -> (mem, pt)) mems + _ -> map (\(_, t, mem, _) -> (memSrcName mem, elemType t)) fromwrites' + mem_primtypes = M.union input_mem_primtypes output_mem_primtypes + + -- Separation of input memory blocks and output memory blocks. + mem_ins0 = S.fromList $ map (\(_, _, mem, _) -> memSrcName mem) fromreads + mem_outs0 = S.fromList $ case output_mems_may of + Just mems -> map (\(mem, _, _) -> mem) mems + _ -> map (\(_, _, mem, _) -> memSrcName mem) fromwrites' + -- An input memory must not be an output memory, and vice versa. + mem_ins = S.difference mem_ins0 mem_outs0 + mem_outs = S.difference mem_outs0 mem_ins0 + + exceptions = snd $ evalRWS (findExceptions fus fus_result lus + mem_ins mem_outs mem_slices mem_ixfuns + mem_primtypes output_vars) () S.empty + in exceptions + + where findExceptions :: FirstUses -> FirstUses -> LastUses -> Names -> Names -> + M.Map VName (Slice SubExp) -> M.Map VName ExpMem.IxFun -> + M.Map VName PrimType -> [VName] -> + RWS () [(VName, VName)] LocalDeaths () + findExceptions fus fus_result lus mem_ins mem_outs mem_slices mem_ixfuns mem_primtypes output_vars = do + forM_ stms $ \(Let (Pattern _patctxelems patvalelems) _ _) -> do + let vs = map patElemName patvalelems + fus_stm = S.unions $ map (`lookupEmptyable` fus) vs + lus_stm = S.unions $ map ((`lookupEmptyable` lus) . FromStm) vs + recordNewExceptions mem_ins mem_outs mem_slices mem_ixfuns mem_primtypes fus_stm + modify $ S.union lus_stm + forM_ output_vars $ \ov -> do + let fus_ov = lookupEmptyable ov fus_result + recordNewExceptions mem_ins mem_outs mem_slices mem_ixfuns mem_primtypes fus_ov + + recordNewExceptions :: Names -> Names -> + M.Map VName (Slice SubExp) -> M.Map VName ExpMem.IxFun -> + M.Map VName PrimType -> Names -> + RWS () [(VName, VName)] LocalDeaths () + recordNewExceptions mem_ins mem_outs mem_slices mem_ixfuns mem_primtypes fus_cur = do + deaths <- get + forM_ (S.toList fus_cur) $ \mem_fu -> forM_ deaths $ \mem_killed -> + fromMaybe (return ()) $ do + slice_fu <- M.lookup mem_fu mem_slices + slice_killed <- M.lookup mem_killed mem_slices + ixfun_fu <- M.lookup mem_fu mem_ixfuns + ixfun_killed <- M.lookup mem_killed mem_ixfuns + pt_fu <- M.lookup mem_fu mem_primtypes + pt_killed <- M.lookup mem_killed mem_primtypes + return $ when + ( -- Is the killed memory read from and the first use memory + -- written to? + mem_fu `S.member` mem_outs && mem_killed `S.member` mem_ins && + -- Same index functions? + ixfun_fu == ixfun_killed && -- too conservative? + -- Same slices? + slice_fu == slice_killed && + -- Same primitive type byte sizes? + (primByteSize pt_fu :: Int) == primByteSize pt_killed + ) $ tell [(mem_fu, mem_killed)] + +-- Memory blocks that have had their last use locally in the body. +type LocalDeaths = Names diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Liveness/LastUse.hs b/src/Futhark/Optimise/MemoryBlockMerging/Liveness/LastUse.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Liveness/LastUse.hs @@ -0,0 +1,281 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find last uses for all memory blocks. +-- +-- A memory block can have more than one last use. +module Futhark.Optimise.MemoryBlockMerging.Liveness.LastUse + ( findLastUses + ) where + +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import Control.Monad +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + (ExplicitMemorish, ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types + + +type LastUsesList = [LastUses] + +getLastUsesMap :: LastUsesList -> LastUses +getLastUsesMap = M.unionsWith S.union + +-- Mapping from a memory block to its currently assumed last use statement +-- variable. +type OptimisticLastUses = M.Map VName (StmOrRes, Bool) + +data Context = Context + { ctxVarToMem :: VarMemMappings MemorySrc + , ctxMemAliases :: MemAliases + , ctxFirstUses :: FirstUses + , ctxExistentials :: Names + , ctxCurFirstUsesOuter :: Names + } + deriving (Show) + +data Current = Current + { curOptimisticLastUses :: OptimisticLastUses + , curFirstUses :: Names + } + deriving (Show) + +newtype FindM lore a = FindM { unFindM :: RWS Context LastUsesList Current a } + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadWriter LastUsesList, + MonadState Current) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- Find the memory blocks used or aliased by a variable. +varMems :: VName -> FindM lore MNames +varMems var = + maybe S.empty (S.singleton . memSrcName) <$> asks (M.lookup var . ctxVarToMem) + +modifyCurOptimisticLastUses :: (OptimisticLastUses -> OptimisticLastUses) -> FindM lore () +modifyCurOptimisticLastUses f = + modify $ \c -> c { curOptimisticLastUses = f $ curOptimisticLastUses c } + +modifyCurFirstUses :: (Names -> Names) -> FindM lore () +modifyCurFirstUses f = modify $ \c -> c { curFirstUses = f $ curFirstUses c } + +withLocalCurFirstUses :: FindM lore a -> FindM lore a +withLocalCurFirstUses m = do + cur_first_uses <- gets curFirstUses + res <- m + modifyCurFirstUses $ const cur_first_uses + return res + +recordMapping :: StmOrRes -> MName -> FindM lore () +recordMapping var mem = tell [M.singleton var (S.singleton mem)] + +-- | Find all last uses of *memory blocks* in a function definition. +findLastUses :: VarMemMappings MemorySrc -> MemAliases -> FirstUses -> Names + -> FunDef ExplicitMemory -> LastUses +findLastUses var_to_mem mem_aliases first_uses existentials fundef = + let context = Context + { ctxVarToMem = var_to_mem + , ctxMemAliases = mem_aliases + , ctxFirstUses = first_uses + , ctxExistentials = existentials + , ctxCurFirstUsesOuter = S.empty + } + m = unFindM $ do + forM_ (funDefParams fundef) lookInFunDefFParam + lookInBody $ funDefBody fundef + mapM_ lookInRes $ bodyResult $ funDefBody fundef + optimistics <- gets curOptimisticLastUses + forM_ (M.keys optimistics) $ \mem -> + commitOptimistic mem + + last_uses = removeEmptyMaps $ getLastUsesMap + $ snd $ evalRWS m context (Current M.empty S.empty) + in last_uses + +-- Optimistically say that the last use of 'mem' and all its memory aliases is +-- at 'x_lu'. Exclude 'exclude' from the memory aliases (necessary in a few +-- edge cases). +setOptimistic :: MName -> StmOrRes -> MNames -> FindM lore () +setOptimistic mem x_lu exclude = do + -- Will override any previous optimistic last use. + mem_aliases <- asks ctxMemAliases + let mems = S.difference (S.union (S.singleton mem) + $ lookupEmptyable mem mem_aliases) exclude + + forM_ mems $ \mem' -> do + let is_indirect = mem' /= mem + modifyCurOptimisticLastUses $ M.insert mem' (x_lu, is_indirect) + +-- If an optimistic last use 'mem' was added through a memory alias, forget +-- about it. +removeIndirectOptimistic :: MName -> FindM lore () +removeIndirectOptimistic mem = do + res <- M.lookup mem <$> gets curOptimisticLastUses + case res of + Just (_, True) -> -- Means that is was added indirectly. + modifyCurOptimisticLastUses $ M.delete mem + _ -> return () + +-- Set the optimistic last use in stone. +commitOptimistic :: MName -> FindM lore () +commitOptimistic mem = do + res <- M.lookup mem <$> gets curOptimisticLastUses + case res of + Just (x_lu, _) -> recordMapping x_lu mem + Nothing -> return () + +lookInFunDefFParam :: FParam lore -> FindM lore () +lookInFunDefFParam (Param x _) = do + first_uses_x <- lookupEmptyable x <$> asks ctxFirstUses + modifyCurFirstUses $ S.union first_uses_x + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern _patctxelems patvalelems) _ e) = do + -- When an loop, a scan, a reduce, or a stream contains a use of an array that + -- is created before the expression body, it should not get a last use in a + -- statement inside the inner body, since loops can have cycles, and so its + -- proper last use should really be in the statement declaring the sub-body, + -- and not in some statement in the sub-body. See + -- 'tests/reuse/loop/copy-from-outside.fut for an example of this. + cur_first_uses <- gets curFirstUses + let mMod = case e of + If{} -> id -- If is the only other expression with a body. + _ -> local $ \ctx -> ctx { ctxCurFirstUsesOuter = cur_first_uses } + + -- First handle all pattern elements by themselves. + forM_ patvalelems $ \(PatElem x membound) -> + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn xmem _) -> do + first_uses_x <- lookupEmptyable x <$> asks ctxFirstUses + modifyCurFirstUses $ S.union first_uses_x + -- When this is a new first use of a memory block, commit the previous + -- optimistic last use of it, so that it can be considered unused in + -- the statements inbetween. + when (S.member xmem first_uses_x) $ commitOptimistic xmem + _ -> return () + + -- Then find the new memory blocks. + let e_free_vars = freeInExp e `S.difference` S.fromList (freeExcludes e) + e_mems <- S.unions <$> mapM varMems (S.toList e_free_vars) + + mem_aliases <- asks ctxMemAliases + first_uses_outer <- asks ctxCurFirstUsesOuter + -- Then handle the pattern elements by themselves again. + forM_ patvalelems $ \(PatElem x _) -> + -- Set all memory blocks being used as optimistic last uses. + forM_ (S.toList e_mems) $ \mem -> do + -- If the memory has its first use outside the current body, it is + -- dangerous to set its last use to be in a statement inside the body, + -- since the body can be run multiple times in cases of loops or kernels, + -- so we only set the last use of a memory to this statement if it also + -- has its first use inside the current body. + -- + -- If it (or any aliased memory) does have its first use outside the body, + -- we remove any existing optimistic last use, although only if such an + -- optimistic last use was added as a side effect of adding an existential + -- optimistic last use (i.e. it was aliased by the existential memory + -- which had a last use). + let from_outer = any (`S.member` first_uses_outer) + (mem : S.toList (lookupEmptyable mem mem_aliases)) + if from_outer + then removeIndirectOptimistic mem + else setOptimistic mem (FromStm x) S.empty + + if S.null (lookupEmptyable mem mem_aliases) + then + -- If not existential, update the potential last use of any existential + -- memory aliasing it, but do not set the potential last use of the + -- memory itself, since there are cycles in loops, and it must also + -- contain the same data in the next iteration, so it can never be + -- reused inside the loop body, and must therefore always have its last + -- use outside the body. But since the existential memory might in the + -- current iteration refer to it, its last use needs to be updated. + + -- Note that while it is not wrong to run the code below also when the + -- memory has its first use inside the body, in that case it should not + -- be necessary, since we would be outside the body by then, and it + -- would result in a too conservative analysis. As an example, see + -- tests/mix/loop-interference-use.fut. + when from_outer $ do + -- If the memory has its first use outside the current body, we need + -- to find its actual last use (if it occurs in the body) through + -- memory aliases. + -- + -- If memory block t aliases memory block u (meaning that the memory of + -- t *can* be the memory of u), and u has a potential last use here, + -- then t also has a potential last use here (the relation is not + -- commutative, so it does not work the other way round). + let reverse_mem_aliases = M.keys $ M.filter (mem `S.member`) mem_aliases + exclude = S.singleton mem + forM_ reverse_mem_aliases $ \mem' -> + setOptimistic mem' (FromStm x) exclude + else + -- Just set the last use. + unless from_outer $ setOptimistic mem (FromStm x) S.empty + + withLocalCurFirstUses $ mMod $ fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + +-- Look in body results. +lookInRes :: SubExp -> FindM lore () +lookInRes (Var v) = do + exis <- asks ctxExistentials + -- If v is a existential variable, there is no reason to record its last use, + -- as existential memory cannot be reused (this is also the case for other + -- setOptimistic calls, but not in a clear way). + unless (v `S.member` exis) $ do + mem_v <- M.lookup v <$> asks ctxVarToMem + case mem_v of + Just mem -> + setOptimistic (memSrcName mem) (FromRes v) S.empty + Nothing -> + return () +lookInRes _ = return () + +-- Some freeInExp results are too limiting and give us too conservative last use +-- results (especially in the CPU pipeline). We only care about a free variable +-- if we *read* from it. If it only exists for *writing*, then we don't have to +-- look at its memory, since whatever is there we overwrite, and so there cannot +-- be any last *use*. +freeExcludes :: Exp lore -> [VName] +freeExcludes e = case e of + DoLoop _ _mergevalparams _ _ -> + -- FIXME: If the returned memory block-associated mergevalparams do not come + -- directly from a Scratch creation, we should be able to ignore them and + -- thereby become less conservative. + [] + + BasicOp (Update orig _ _) -> + [orig] + + _ -> [] diff --git a/src/Futhark/Optimise/MemoryBlockMerging/MemoryAliases.hs b/src/Futhark/Optimise/MemoryBlockMerging/MemoryAliases.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/MemoryAliases.hs @@ -0,0 +1,162 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE LambdaCase #-} +-- | Find memory block aliases. The conceptual difference from variable aliases +-- is that if a variable x has an alias y, it means that x and y use the same +-- memory block, but if a memory block xmem has an alias ymem, it means that +-- xmem and ymem refer to the same *memory*. This is not commutative. +module Futhark.Optimise.MemoryBlockMerging.MemoryAliases + ( findMemAliases + ) where + +import Data.Maybe (mapMaybe) +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Data.List as L +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.Aliases +import Futhark.Representation.ExplicitMemory + (ExplicitMemorish, ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel +import Futhark.Analysis.Alias (analyseFun) + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types + + +newtype FindM lore a = FindM { unFindM :: RWS (VarMemMappings MemorySrc) [MemAliases] () a } + deriving (Monad, Functor, Applicative, + MonadReader (VarMemMappings MemorySrc), + MonadWriter [MemAliases]) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalkAliases lore) + +recordMapping :: MName -> MNames -> FindM lore () +recordMapping mem mems = tell [M.singleton mem (S.delete mem mems)] + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +lookupMems :: Names -> FindM lore MNames +lookupMems var_aliases = do + var_to_mem <- ask + return $ S.fromList $ mapMaybe ((memSrcName <$>) . flip M.lookup var_to_mem) + $ S.toList var_aliases + +-- | Find all memory aliases in a function definition. +findMemAliases :: FunDef ExplicitMemory -> VarMemMappings MemorySrc -> MemAliases +findMemAliases fundef var_to_mem = + let fundef' = analyseFun fundef + m = unFindM $ lookInBody $ funDefBody fundef' + mem_aliases = M.unionsWith S.union $ snd $ evalRWS m var_to_mem () + mem_aliases' = removeEmptyMaps $ expandWithAliases mem_aliases mem_aliases + in mem_aliases' + +lookInBody :: LoreConstraints lore => + Body (Aliases lore) -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody (Aliases lore) -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm (Aliases lore) -> FindM lore () +lookInStm (Let (Pattern patctxelems patvalelems) _ e) = do + forM_ (patctxelems ++ patvalelems) lookInPatElem + + case e of + DoLoop mergectxparams mergevalparams _loopform body -> do + -- There are most likely more body results than + -- mergectxparams, but we are only interested in the first + -- body results anyway (those that have a matching location + -- with the mergectxparams). + zipWithM_ lookInMergeCtxParam mergectxparams (bodyResult body) + zipWithM_ lookInCtx patctxelems mergectxparams + mapM_ (lookInMergeValParam body) mergevalparams + mapM_ (lookInBodyTuples patctxelems (map snd mergectxparams) (bodyResult body)) + patvalelems + If _ body_then body_else _ -> do + -- Alias everything. FIXME: This is maybe more conservative than + -- necessary if the If works on tuples of arrays. + let ress = mapMaybe subExpVar + (bodyResult body_then ++ bodyResult body_else) + var_to_mem <- ask + let mems = map memSrcName $ mapMaybe (`M.lookup` var_to_mem) ress + forM_ patctxelems $ \case + (PatElem patmem (_, ExpMem.MemMem{})) -> + recordMapping patmem $ S.fromList mems + _ -> return () + _ -> return () + + fullWalkAliasesExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + +lookInCtx :: LoreConstraints lore => + PatElem (Aliases lore) -> (FParam (Aliases lore), SubExp) + -> FindM lore () +lookInCtx (PatElem patmem (_, ExpMem.MemMem{})) (Param parammem ExpMem.MemMem{}, _) = do + recordMapping patmem (S.singleton parammem) + recordMapping parammem (S.singleton patmem) +lookInCtx _ _ = return () + +lookInMergeCtxParam :: LoreConstraints lore => + (FParam (Aliases lore), SubExp) -> SubExp -> FindM lore () +lookInMergeCtxParam (Param xmem ExpMem.MemMem{}, Var param_mem) (Var body_mem_res) = do + let aliases = S.fromList [param_mem, body_mem_res] + recordMapping xmem aliases +lookInMergeCtxParam _ _ = return () + +lookInMergeValParam :: LoreConstraints lore => + Body (Aliases lore) -> (FParam (Aliases lore), SubExp) + -> FindM lore () +lookInMergeValParam body (Param _ (ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _)), _t) = do + -- FIXME: This is maybe more conservative than necessary in case you have more + -- than one loop array. Fixing this would require either changing the Aliases + -- representation, or building something on top of it. + aliases <- S.unions + <$> mapM (lookupMems . unNames) (fst $ fst $ bodyAttr body) + recordMapping mem aliases +lookInMergeValParam _ _ = return () + +lookInBodyTuples :: LoreConstraints lore => + [PatElem (Aliases lore)] + -> [SubExp] -> [SubExp] + -> PatElem (Aliases lore) + -> FindM lore () +-- When a parameter refers to a existential memory, we want to find +-- which return memory in the loop that the existential memory refers +-- to. +lookInBodyTuples patctxelems body_params body_results + (PatElem _ (_, ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _))) = do + let zipped = zip3 patctxelems body_params body_results + case L.find ((== mem) . patElemName . (\(x, _, _) -> x)) zipped of + Just (_, Var param_mem, Var res_mem) -> + recordMapping mem (S.fromList [param_mem, res_mem]) + _ -> return () +lookInBodyTuples _ _ _ _ = return () + +lookInPatElem :: LoreConstraints lore => + PatElem (Aliases lore) -> FindM lore () +lookInPatElem (PatElem _ (names', ExpMem.MemArray _ _ _ (ExpMem.ArrayIn xmem _))) = do + aliases <- lookupMems $ unNames names' + recordMapping xmem aliases +lookInPatElem (PatElem xmem (names', ExpMem.MemMem {})) = do + aliases <- lookupMems $ unNames names' + recordMapping xmem aliases +lookInPatElem _ = return () diff --git a/src/Futhark/Optimise/MemoryBlockMerging/MemoryUpdater.hs b/src/Futhark/Optimise/MemoryBlockMerging/MemoryUpdater.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/MemoryUpdater.hs @@ -0,0 +1,418 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE LambdaCase #-} + +-- | Transform a function based on a mapping from variable to memory and index +-- function: Change every variable in the mapping to its possibly new memory +-- block. +module Futhark.Optimise.MemoryBlockMerging.MemoryUpdater + ( transformFromVarMemMappings + ) where + +import qualified Data.Map.Strict as M +import qualified Data.List as L +import Data.Maybe (mapMaybe, fromMaybe) +import Control.Applicative ((<|>)) +import Control.Arrow (second) +import Control.Monad +import Control.Monad.RWS + +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + (ExplicitMemorish, ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Types +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +data Context = Context { ctxVarToMem :: VarMemMappings MemoryLoc + , ctxVarToMemOrig :: VarMemMappings MName + , ctxAllocSizes :: M.Map MName SubExp + , ctxAllocSizesOrig :: M.Map MName SubExp + , ctxHasMaxedSize :: Bool + } + deriving (Show) + +newtype FindM lore a = FindM { unFindM :: RWS Context () (VNameSource, [(MName, VName)]) a } + deriving (Monad, Functor, Applicative, + MonadReader Context, MonadState (VNameSource, [(MName, VName)])) + +instance MonadFreshNames (FindM lore) where + getNameSource = gets fst + putNameSource s = modify $ \(_, m) -> (s, m) + +modifyMemSizeMapping :: ([(MName, VName)] -> [(MName, VName)]) -> FindM lore () +modifyMemSizeMapping f = modify $ second f + +type LoreConstraints lore = (ExplicitMemorish lore, + FullMap lore, + BodyAttr lore ~ (), + ExpAttr lore ~ ()) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- | Transform a function to use new memory blocks. +transformFromVarMemMappings :: MonadFreshNames m => + VarMemMappings MemoryLoc -> + VarMemMappings MName -> + M.Map MName SubExp -> M.Map MName SubExp -> Bool -> + FunDef ExplicitMemory -> + m (FunDef ExplicitMemory) +transformFromVarMemMappings var_to_mem var_to_mem_orig alloc_sizes alloc_sizes_orig has_maxed_size fundef = + let m = unFindM $ transformFunDefBody $ funDefBody fundef + ctx = Context { ctxVarToMem = var_to_mem + , ctxVarToMemOrig = var_to_mem_orig + , ctxAllocSizes = alloc_sizes + , ctxAllocSizesOrig = alloc_sizes_orig + , ctxHasMaxedSize = has_maxed_size + } + in modifyNameSource (\src -> + let (body', (src', _), ()) = runRWS m ctx (src, []) + in (fundef { funDefBody = body' }, src') + ) + +transformFunDefBody :: LoreConstraints lore => + Body lore -> FindM lore (Body lore) +transformFunDefBody (Body () bnds res) = do + bnds' <- mapM transformStm $ stmsToList bnds + res' <- transformFunDefBodyResult res + return $ Body () (stmsFromList bnds') res' + +transformFunDefBodyResult :: [SubExp] -> FindM lore [SubExp] +transformFunDefBodyResult ses = do + var_to_mem_orig <- asks ctxVarToMemOrig + var_to_mem <- asks ctxVarToMem + mem_to_size_orig <- asks ctxAllocSizesOrig + mem_to_size <- asks ctxAllocSizes + mem_to_new_size <- gets snd + + let check se + | Var v <- se + , Just orig <- M.lookup v var_to_mem_orig + , Just new <- memLocName <$> M.lookup v var_to_mem + = ((Var orig, Nothing), Var new) : case (M.lookup orig mem_to_size_orig, + (Var <$> L.lookup new mem_to_new_size) <|> M.lookup new mem_to_size) of + (Just size_orig, Just size_new) -> + [((size_orig, Just (Var orig)), size_new)] + _ -> [] + | otherwise = [] + + check_size_only se + | Var v <- se + , Just orig <- M.lookup v mem_to_size_orig + , Just new <- (Var <$> L.lookup v mem_to_new_size) <|> M.lookup v mem_to_size + , orig /= new + = [((orig, Just (Var v)), new)] + | otherwise = [] + mem_orig_to_new1 = concatMap check ses + mem_orig_to_new2 = concatMap check_size_only ses + mem_orig_to_new = mem_orig_to_new1 ++ mem_orig_to_new2 + + return $ zipWith ( + \se ts -> fromMaybe se ( + -- FIXME: This assumes that a memory block always + -- comes just after its size variable. We ought + -- to instead properly find this information from + -- the funDefRetType 'ExtSize's. + (se, Nothing) `L.lookup` mem_orig_to_new + <|> case ts of + (ts0 : _) -> + (se, Just ts0) `L.lookup` mem_orig_to_new + _ -> Nothing + ) + ) ses (L.tail $ L.tails ses) + +transformBody :: LoreConstraints lore => + Body lore -> FindM lore (Body lore) +transformBody (Body () bnds res) = do + bnds' <- mapM transformStm $ stmsToList bnds + return $ Body () (stmsFromList bnds') res + +transformKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore (KernelBody lore) +transformKernelBody (KernelBody () bnds res) = do + bnds' <- mapM transformStm $ stmsToList bnds + return $ KernelBody () (stmsFromList bnds') res + +transformMemInfo :: ExpMem.MemInfo d u ExpMem.MemReturn -> MemoryLoc -> + ExpMem.MemInfo d u ExpMem.MemReturn +transformMemInfo meminfo memloc = case meminfo of + ExpMem.MemArray pt shape u _memreturn -> + let extixfun = ExpMem.existentialiseIxFun [] $ memLocIxFun memloc + in ExpMem.MemArray pt shape u + (ExpMem.ReturnsInBlock (memLocName memloc) extixfun) + _ -> meminfo + +data BranchReturn = ExistingBranchReturn ExpMem.BodyReturns + | NewBranchReturn (Int -> ExpMem.BodyReturns) + VName VName VName + +transformStm :: LoreConstraints lore => + Stm lore -> FindM lore (Stm lore) +transformStm (Let (Pattern patctxelems patvalelems) aux e) = do + patvalelems' <- mapM transformPatValElem patvalelems + + e' <- fullMapExpM mapper mapper_kernel e + var_to_mem <- asks ctxVarToMem + var_to_mem_orig <- asks ctxVarToMemOrig + mem_to_size <- asks ctxAllocSizes + mem_to_new_size <- gets snd + (e'', patctxelems') <- case e' of + If cond body_then body_else (IfAttr rets sort) -> do + let bodyVarMemLocs body = + map (flip M.lookup var_to_mem <=< subExpVar) + $ drop (length patctxelems) $ bodyResult body + + -- FIXME: This is a mess. We try to "reverse-engineer" the origin of + -- how the If results came to look as they do, so that we can produce + -- a correct IfAttr. + findBodyResMem i body_results = + let imem = patElemName (patctxelems L.!! i) + matching_var = mapMaybe ( + \(p, p_i) -> + case patElemAttr p of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn vmem _) -> + if imem == vmem + then Just p_i + else Nothing + _ -> + Nothing + ) (zip patvalelems [0..]) + in do + j <- case matching_var of + [t] -> Just t + _ -> Nothing + body_res_var <- subExpVar (body_results L.!! (length patctxelems + j)) + MemoryLoc mem _ixfun <- M.lookup body_res_var var_to_mem + return mem + + fixBodyExistentials body = + body { bodyResult = + zipWith (\res i -> if i < length patctxelems + then maybe res Var $ findBodyResMem i (bodyResult body) + else res) + (bodyResult body) [0..] } + + let ms_then = bodyVarMemLocs body_then + ms_else = bodyVarMemLocs body_else + + -- Fix values. + let rets' = + if ms_then == ms_else + then zipWith (\r m -> case m of + Nothing -> r + Just m' -> + transformMemInfo r m' + ) rets ms_then + else rets + + let body_then' = fixBodyExistentials body_then + body_else' = fixBodyExistentials body_else + + + -- Fix existential memory blocks. + let mem_size mem = L.lookup mem mem_to_new_size <|> (subExpVar =<< M.lookup mem mem_to_size) + v_size v = do + mem <- M.lookup v (M.map memLocName var_to_mem) <|> M.lookup v var_to_mem_orig + mem_size mem + + has_maxed_size <- asks ctxHasMaxedSize + let rets_branch_returns = + L.zipWith4 (\r pat th el -> case (r, pat, th, el) of + (ExpMem.MemArray pt shape u + (ExpMem.ReturnsNewBlock space n + (Free (Var _size)) extixfun), + PatElem _ + (ExpMem.MemArray _ _ _ + (ExpMem.ArrayIn patmem _)), + Var v_th, Var v_el) -> + case (v_size v_th, v_size v_el) of + (Just s_th, Just s_el) -> + if not has_maxed_size --s_th == s_el || not has_maxed_size + then ExistingBranchReturn r + else NewBranchReturn + (\nth_ctxelem -> + ExpMem.MemArray pt shape u + (ExpMem.ReturnsNewBlock space n + (Ext nth_ctxelem) extixfun)) + s_th s_el patmem + _ -> error ("both branch return arrays should use a memory block with a size: " ++ show v_th ++ " and " ++ show v_el) + _ -> ExistingBranchReturn r + ) + rets' + patvalelems + (drop (length patctxelems) (bodyResult body_then')) + (drop (length patctxelems) (bodyResult body_else')) + + patctxelems_new <- + replicateM + (length (filter (\case + NewBranchReturn{} -> True + ExistingBranchReturn{} -> False + ) rets_branch_returns)) + (newVName "new_memory_size") + let (rets'', _, body_ext_new, _, patmem_to_new_size) = + foldl (\(prev, i, ext, patctxelems_new', mapping) rb -> case rb of + ExistingBranchReturn r -> + (prev ++ [r], i, ext, patctxelems_new', mapping) + NewBranchReturn rf s_th s_el patmem -> + (prev ++ [rf i], i + 1, ext ++ [(s_th, s_el)], + tail patctxelems_new', + mapping ++ [(patmem, head patctxelems_new')]) + ) ([], length patctxelems, [], patctxelems_new, []) rets_branch_returns + modifyMemSizeMapping (++ patmem_to_new_size) + let (th_ext_new, el_ext_new) = unzip body_ext_new + body_then'' = body_then' { bodyResult = + take (length patctxelems) (bodyResult body_then') ++ + map Var th_ext_new ++ + drop (length patctxelems) (bodyResult body_then') + } + body_else'' = body_else' { bodyResult = + take (length patctxelems) (bodyResult body_else') ++ + map Var el_ext_new ++ + drop (length patctxelems) (bodyResult body_else') + } + patctxelems_replaced = map (\pe -> case pe of + PatElem name (ExpMem.MemMem _size space) -> + case L.lookup name patmem_to_new_size of + Just size_new -> + PatElem name (ExpMem.MemMem (Var size_new) space) + Nothing -> pe + _ -> pe + ) patctxelems + patctxelems' = patctxelems_replaced ++ map (\v -> PatElem v (ExpMem.MemPrim (IntType Int64))) patctxelems_new + + return (If cond body_then'' body_else'' (IfAttr rets'' sort), + patctxelems') + + DoLoop mergectxparams mergevalparams loopform body -> do + -- More special loop handling because of its extra + -- pattern-like info. + mergectxparams' <- mapM (transformMergeCtxParam mergevalparams) mergectxparams + mergevalparams' <- mapM transformMergeValParam mergevalparams + + -- The body of a loop can return a memory block in its results. This is + -- the memory block used by a variable which is also part of the results. + -- If the memory block of that variable is changed, we need a way to + -- record that the memory block in the body result also needs to change. + let zipped = zip [(0::Int)..] (patctxelems ++ patvalelems) + + findMemLinks (i, PatElem _x (ExpMem.MemArray _ _ _ (ExpMem.ArrayIn xmem _))) = + case L.find (\(_, PatElem ymem _) -> ymem == xmem) zipped of + Just (j, _) -> Just (j, i) + Nothing -> Nothing + findMemLinks _ = Nothing + + mem_links = mapMaybe findMemLinks zipped + + res = bodyResult body + + fixResRecord i se + | Var _mem <- se + , Just j <- L.lookup i mem_links + , Var related_var <- res L.!! j + , Just mem_new <- M.lookup related_var var_to_mem = + Var $ memLocName mem_new + | otherwise = se + + res' = zipWith fixResRecord [(0::Int)..] res + body' = body { bodyResult = res' } + + loopform' <- case loopform of + ForLoop i it bound loop_vars -> + ForLoop i it bound <$> mapM transformForLoopVar loop_vars + WhileLoop _ -> return loopform + return (DoLoop mergectxparams' mergevalparams' loopform' body', + patctxelems) + _ -> return (e', patctxelems) + return (Let (Pattern patctxelems' patvalelems') aux e'') + where mapper = identityMapper + { mapOnBody = const transformBody + , mapOnFParam = transformFParam + , mapOnLParam = transformLParam + } + mapper_kernel = identityKernelMapper + { mapOnKernelBody = coerce . transformBody + , mapOnKernelKernelBody = coerce . transformKernelBody + , mapOnKernelLambda = coerce . transformLambda + , mapOnKernelLParam = transformLParam + } + + + +-- Update the actual memory block referred to by a context (existential) memory +-- block in a loop. +transformMergeCtxParam :: [(FParam ExplicitMemory, SubExp)] -> + (FParam ExplicitMemory, SubExp) + -> FindM lore (FParam ExplicitMemory, SubExp) +transformMergeCtxParam mergevalparams (param@(Param ctxmem ExpMem.MemMem{}), mem) = do + var_to_mem <- asks ctxVarToMem + + let usesCtxMem (Param _ (ExpMem.MemArray _ _ _ (ExpMem.ArrayIn pmem _))) = ctxmem == pmem + usesCtxMem _ = False + + -- If the initial value of a loop merge parameter is a memory block name, + -- we may have to update that. If the context memory block is used in an + -- array in one of the value merge parameters, see if that array variable + -- refers to an array that has been set to reuse a memory block. + mem' = fromMaybe mem $ do + (_, Var orig_var) <- L.find (usesCtxMem . fst) mergevalparams + orig_mem <- M.lookup orig_var var_to_mem + return $ Var $ memLocName orig_mem + return (param, mem') +transformMergeCtxParam _ t = return t + +transformMergeValParam :: (FParam ExplicitMemory, SubExp) + -> FindM lore (FParam ExplicitMemory, SubExp) +transformMergeValParam (Param x membound, se) = do + membound' <- newMemBound membound x + return (Param x membound', se) + +transformPatValElem :: PatElem ExplicitMemory -> FindM lore (PatElem ExplicitMemory) +transformPatValElem (PatElem x membound) = + PatElem x <$> newMemBound membound x + +transformFParam :: LoreConstraints lore => + FParam lore -> FindM lore (FParam lore) +transformFParam (Param x membound) = + Param x <$> newMemBound membound x + +transformLParam :: LoreConstraints lore => + LParam lore -> FindM lore (LParam lore) +transformLParam (Param x membound) = + Param x <$> newMemBound membound x + +transformLambda :: LoreConstraints lore => + Lambda lore -> FindM lore (Lambda lore) +transformLambda (Lambda params body types) = do + params' <- mapM transformLParam params + body' <- transformBody body + return $ Lambda params' body' types + +transformForLoopVar :: LoreConstraints lore => + (LParam lore, VName) -> + FindM lore (LParam lore, VName) +transformForLoopVar (Param x membound, array) = do + membound' <- newMemBound membound x + return (Param x membound', array) + +-- Find a new memory block and index function if they exist. +newMemBound :: ExpMem.MemBound u -> VName -> FindM lore (ExpMem.MemBound u) +newMemBound membound var = do + var_to_mem <- asks ctxVarToMem + + let membound' + | ExpMem.MemArray pt shape u _ <- membound + , Just (MemoryLoc mem ixfun) <- M.lookup var var_to_mem = + Just $ ExpMem.MemArray pt shape u $ ExpMem.ArrayIn mem ixfun + | otherwise = Nothing + + return $ fromMaybe membound membound' diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Miscellaneous.hs b/src/Futhark/Optimise/MemoryBlockMerging/Miscellaneous.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Miscellaneous.hs @@ -0,0 +1,263 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | Miscellaneous helper functions. Perpetually in need of a cleanup. +module Futhark.Optimise.MemoryBlockMerging.Miscellaneous + ( makeCommutativeMap + , insertOrUpdate + , insertOrUpdateMany + , insertOrNew + , removeEmptyMaps + , removeKeyFromMapElems + , newDeclarationsStm + , lookupEmptyable + , fromJust + , maybeFromBoolM + , sortByKeyM + , mapMaybeM + , anyM + , whenM + , expandPrimExp + , expandIxFun + , mapFromListSetUnion + , fixpointIterateMay + , filterSetM + , (<&&>), (<||>) + + , expandWithAliases + , FullWalk(..) + , fullWalkAliasesExpM + , FullWalkAliases + , FullMap + , fullMapExpM + ) where + +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Data.List as L +import Control.Monad +import Data.Maybe (fromMaybe, catMaybes) +import Data.Function (on) + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + (ExplicitMemory, InKernel) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel +import Futhark.Representation.Kernels.KernelExp +import Futhark.Representation.Aliases +import Futhark.Analysis.PrimExp.Convert + +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Optimise.MemoryBlockMerging.Types + + +-- If a property is commutative in a map, build a map that reflects it. A bit +-- crude. We could also just use a function that calculates this whenever +-- needed. +makeCommutativeMap :: Ord v => M.Map v (S.Set v) -> M.Map v (S.Set v) +makeCommutativeMap m = + let names = S.toList (S.union (M.keysSet m) (S.unions (M.elems m))) + assocs = map (\n -> + let existing = lookupEmptyable n m + newly_found = S.unions $ map (\(k, v) -> + if S.member n v + then S.singleton k + else S.empty) $ M.assocs m + ns = S.union existing newly_found + in (n, ns)) names + in M.fromList assocs + +insertOrUpdate :: (Ord k, Ord v) => k -> v -> + M.Map k (S.Set v) -> M.Map k (S.Set v) +insertOrUpdate k v = M.alter (insertOrNew (S.singleton v)) k + +insertOrUpdateMany :: (Ord k, Ord v) => k -> S.Set v -> + M.Map k (S.Set v) -> M.Map k (S.Set v) +insertOrUpdateMany k vs = M.alter (insertOrNew vs) k + +insertOrNew :: Ord a => S.Set a -> Maybe (S.Set a) -> Maybe (S.Set a) +insertOrNew xs m = Just $ case m of + Just s -> S.union xs s + Nothing -> xs + +removeEmptyMaps :: M.Map k (S.Set v) -> M.Map k (S.Set v) +removeEmptyMaps = M.filter (not . S.null) + +removeKeyFromMapElems :: (Ord k) => M.Map k (S.Set k) -> M.Map k (S.Set k) +removeKeyFromMapElems = M.mapWithKey S.delete + +newDeclarationsStm :: Stm lore -> [VName] +newDeclarationsStm (Let (Pattern patctxelems patvalelems) _ e) = + let new_decls0 = map patElemName (patctxelems ++ patvalelems) + new_decls1 = case e of + DoLoop mergectxparams mergevalparams _loopform _body -> + -- Technically not a declaration for the current expression, but very + -- close. + map (paramName . fst) (mergectxparams ++ mergevalparams) + _ -> [] + new_decls = new_decls0 ++ new_decls1 + in new_decls + +lookupEmptyable :: (Ord a, Monoid b) => a -> M.Map a b -> b +lookupEmptyable x m = fromMaybe mempty $ M.lookup x m + +fromJust :: String -> Maybe a -> a +fromJust _ (Just x) = x +fromJust mistake Nothing = error ("error: " ++ mistake) + +maybeFromBoolM :: Monad m => (a -> m Bool) -> (a -> m (Maybe a)) +maybeFromBoolM f a = do + res <- f a + return $ if res + then Just a + else Nothing + +expandWithAliases :: forall v. Ord v => MemAliases -> M.Map v Names -> M.Map v Names +expandWithAliases mem_aliases = fixpointIterate expand + where expand :: M.Map v Names -> M.Map v Names + expand mems_map = + M.fromList (map (\(v, mems) -> + (v, S.unions (mems : map (`lookupEmptyable` mem_aliases) + (S.toList mems)))) + (M.assocs mems_map)) + +fixpointIterate :: Eq a => (a -> a) -> a -> a +fixpointIterate f x + | f x == x = x + | otherwise = fixpointIterate f (f x) + +fixpointIterateMay :: (a -> Maybe a) -> a -> a +fixpointIterateMay f x = maybe x (fixpointIterateMay f) (f x) + +mapFromListSetUnion :: (Ord k, Ord v) => [(k, S.Set v)] -> M.Map k (S.Set v) +mapFromListSetUnion = M.unionsWith S.union . map (uncurry M.singleton) + +-- Replace variables with subtrees of their constituents wherever possible. It +-- naively expands a PrimExp as much as the input map allows, and can enable +-- more expressions to have it in scope, since it will likely consist of fewer +-- variables. +expandPrimExp :: M.Map VName (ExpMem.PrimExp VName) -> ExpMem.PrimExp VName + -> ExpMem.PrimExp VName +expandPrimExp var_to_pe = fixpointIterate (substituteInPrimExp var_to_pe) + +expandIxFun :: M.Map VName (ExpMem.PrimExp VName) -> ExpMem.IxFun -> ExpMem.IxFun +expandIxFun var_to_pe = fixpointIterate (IxFun.substituteInIxFun var_to_pe) + +(<&&>) :: Monad m => m Bool -> m Bool -> m Bool +m <&&> n = (&&) <$> m <*> n + +(<||>) :: Monad m => m Bool -> m Bool -> m Bool +m <||> n = (||) <$> m <*> n + +anyM :: Monad m => (a -> m Bool) -> [a] -> m Bool +anyM f xs = or <$> mapM f xs + +whenM :: Monad m => m Bool -> m () -> m () +whenM b m = do + b' <- b + when b' m + +mapMaybeM :: Monad m => (a -> m (Maybe b)) -> [a] -> m [b] +mapMaybeM f xs = catMaybes <$> mapM f xs + +sortByKeyM :: (Ord t, Monad m) => (a -> m t) -> [a] -> m [a] +sortByKeyM f xs = + map fst . L.sortBy (compare `on` snd) . zip xs <$> mapM f xs + +filterSetM :: (Ord a, Monad m) => (a -> m Bool) -> S.Set a -> m (S.Set a) +filterSetM f xs = S.fromList <$> filterM f (S.toList xs) + +-- Map on both ExplicitMemory and InKernel. +class FullMap lore where + fullMapExpM :: Monad m => Mapper lore lore m -> KernelMapper InKernel InKernel m + -> Exp lore -> m (Exp lore) + +instance FullMap ExplicitMemory where + fullMapExpM mapper mapper_kernel e = + case e of + Op (ExpMem.Inner kernel) -> + Op . ExpMem.Inner <$> mapKernelM mapper_kernel kernel + _ -> mapExpM mapper e + +instance FullMap InKernel where + fullMapExpM mapper mapper_kernel e = case e of + Op (ExpMem.Inner ke) -> Op . ExpMem.Inner <$> case ke of + ExpMem.Combine a b c body -> + ExpMem.Combine a b c <$> mapOnKernelBody mapper_kernel body + ExpMem.GroupReduce a lambda b -> + ExpMem.GroupReduce a + <$> mapOnKernelLambda mapper_kernel lambda + <*> pure b + ExpMem.GroupScan a lambda b -> + ExpMem.GroupScan a + <$> mapOnKernelLambda mapper_kernel lambda + <*> pure b + ExpMem.GroupStream a b (ExpMem.GroupStreamLambda a1 b1 params0 params1 gsbody) c d -> + ExpMem.GroupStream a b + <$> (ExpMem.GroupStreamLambda a1 b1 + <$> mapM (mapOnKernelLParam mapper_kernel) params0 + <*> mapM (mapOnKernelLParam mapper_kernel) params1 + <*> mapOnKernelBody mapper_kernel gsbody + ) + <*> pure c <*> pure d + _ -> return ke + _ -> mapExpM mapper e + +-- Walk on both ExplicitMemory and InKernel. +class FullWalk lore where + fullWalkExpM :: Monad m => Walker lore m -> KernelWalker InKernel m + -> Exp lore -> m () + +-- FIXME: This can maybe be integrated into the above typeclass. +class FullWalkAliases lore where + fullWalkAliasesExpM :: Monad m => Walker (Aliases lore) m + -> KernelWalker (Aliases InKernel) m + -> Exp (Aliases lore) -> m () + +instance FullWalk ExplicitMemory where + fullWalkExpM walker walker_kernel e = do + walkExpM walker e + case e of + Op (ExpMem.Inner kernel) -> + walkKernelM walker_kernel kernel + _ -> return () + +instance FullWalkAliases ExplicitMemory where + fullWalkAliasesExpM walker walker_kernel e = do + walkExpM walker e + case e of + Op (ExpMem.Inner kernel) -> + walkKernelM walker_kernel kernel + _ -> return () + +instance FullWalk InKernel where + fullWalkExpM walker walker_kernel e = case e of + Op (ExpMem.Inner ke) -> walkOnKernelExpM walker_kernel ke + _ -> walkExpM walker e + +instance FullWalkAliases InKernel where + fullWalkAliasesExpM walker walker_kernel e = case e of + Op (ExpMem.Inner ke) -> walkOnKernelExpM walker_kernel ke + _ -> walkExpM walker e + +walkOnKernelExpM :: Monad m => KernelWalker lore m -> + KernelExp lore -> m () +walkOnKernelExpM walker_kernel ke = case ke of + ExpMem.Combine _ _ _ body -> + walkOnKernelBody walker_kernel body + ExpMem.GroupReduce _ lambda _ -> + walkOnKernelLambda walker_kernel lambda + ExpMem.GroupScan _ lambda _ -> + walkOnKernelLambda walker_kernel lambda + ExpMem.GroupStream _ _ gslambda _ _ -> + walkOnGroupStreamLambdaM walker_kernel gslambda + _ -> return () + +walkOnGroupStreamLambdaM :: Monad m => KernelWalker lore m -> + GroupStreamLambda lore -> m () +walkOnGroupStreamLambdaM walker_kernel (GroupStreamLambda _ _ + params0 params1 gsbody) = do + mapM_ (walkOnKernelLParam walker_kernel) params0 + mapM_ (walkOnKernelLParam walker_kernel) params1 + walkOnKernelBody walker_kernel gsbody diff --git a/src/Futhark/Optimise/MemoryBlockMerging/PrimExps.hs b/src/Futhark/Optimise/MemoryBlockMerging/PrimExps.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/PrimExps.hs @@ -0,0 +1,105 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Get a mapping from statement name to PrimExp (if the statement has a +-- primitive expression) for all statements. +module Futhark.Optimise.MemoryBlockMerging.PrimExps + ( findPrimExpsFunDef + ) where + +import qualified Data.Map.Strict as M +import Data.Maybe (mapMaybe) +import Control.Monad +import Control.Monad.RWS + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + (ExplicitMemorish, ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel +import Futhark.Tools + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +type CurrentTypes = M.Map VName PrimType +type PrimExps = M.Map VName (PrimExp VName) + +newtype FindM lore a = FindM { unFindM :: RWS () PrimExps CurrentTypes a } + deriving (Monad, Functor, Applicative, + MonadWriter PrimExps, + MonadState CurrentTypes) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- Find/construct all 'PrimExp's in a function definition. +findPrimExpsFunDef :: FunDef ExplicitMemory -> PrimExps +findPrimExpsFunDef fundef = + let m = unFindM $ do + lookInFParams $ funDefParams fundef + lookInBody $ funDefBody fundef + res = snd $ evalRWS m () M.empty + in res + +lookInFParams :: LoreConstraints lore => + [FParam lore] -> FindM lore () +lookInFParams params = forM_ params $ \(Param var membound) -> do + case typeOf membound of + Prim pt -> modify $ M.insert var pt + _ -> return () + + case membound of + ExpMem.MemArray pt shape _ (ExpMem.ArrayIn mem _) -> do + let matchingSizeVar (Param mem1 (ExpMem.MemMem (Var mem_size) _)) + | mem1 == mem = Just mem_size + matchingSizeVar _ = Nothing + case mapMaybe matchingSizeVar params of + [mem_size] -> do + let prod_i32 = product (map (primExpFromSubExp (IntType Int32)) (shapeDims shape)) + let prod_i64 = ConvOpExp (SExt Int32 Int64) prod_i32 + let pe = prod_i64 * primByteSize pt + tell $ M.singleton mem_size pe + _ -> return () + _ -> return () + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern _patctxelems patvalelems) _ e) = do + prim_types <- get + let varUse v = ExpMem.LeafExp v <$> M.lookup v prim_types + + case patvalelems of + [PatElem dst _] -> + forM_ (primExpFromExp varUse e) $ tell . M.singleton dst + _ -> return () + + forM_ patvalelems $ \(PatElem var membound) -> + case typeOf membound of + Prim pt -> + modify $ M.insert var pt + _ -> return () + + -- Recursive body walk. + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Reuse.hs b/src/Futhark/Optimise/MemoryBlockMerging/Reuse.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Reuse.hs @@ -0,0 +1,30 @@ +-- | Reuse the memory blocks of arrays. +-- +-- Enable by setting the environment variable MEMORY_BLOCK_MERGING_REUSE=1. +module Futhark.Optimise.MemoryBlockMerging.Reuse + ( reuseInProg + ) where + +import Futhark.Pass + +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemory) + +import Futhark.Optimise.MemoryBlockMerging.AuxiliaryInfo +import Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizeMovingUp +import Futhark.Optimise.MemoryBlockMerging.Reuse.Core + +reuseInProg :: Prog ExplicitMemory -> PassM (Prog ExplicitMemory) +reuseInProg = intraproceduralTransformation reuseInFunDef + +reuseInFunDef :: MonadFreshNames m + => FunDef ExplicitMemory + -> m (FunDef ExplicitMemory) +reuseInFunDef fundef0 = do + let fundef1 = moveUpAllocSizesFunDef fundef0 + aux1 = getAuxiliaryInfo fundef1 + coreReuseFunDef fundef1 + (auxFirstUses aux1) (auxInterferences aux1) + (auxPotentialKernelDataRaceInterferences aux1) (auxVarMemMappings aux1) + (auxActualVariables aux1) (auxExistentials aux1) diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeMovingUp.hs b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeMovingUp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeMovingUp.hs @@ -0,0 +1,32 @@ +-- | Move size variables used in allocation statements upwards in the bodies of +-- a program to enable more memory block reuses. +-- +-- This should be run *before* the reuse pass, as it enables more optimisations. +-- Specifically, it helps with reusing memory whose size needs to be changed to +-- be the maximum of itself and another size -- and so, that other size needs to +-- have been hoisted so that is in scope at that point. This module hoists all +-- sizes as much as possible. +module Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizeMovingUp + ( moveUpAllocSizesFunDef + ) where + +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory (ExplicitMemory) + +import Futhark.Optimise.MemoryBlockMerging.CrudeMovingUp +import Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizes + +findAllocSizeHoistees :: Body ExplicitMemory -> Maybe [FParam ExplicitMemory] + -> [VName] +findAllocSizeHoistees body params = + let subexps = map fst $ M.elems + $ memBlockSizesParamsBodyNonRec (fromMaybe [] params) body + in subExpVars subexps + +moveUpAllocSizesFunDef :: FunDef ExplicitMemory + -> FunDef ExplicitMemory +moveUpAllocSizesFunDef fundef = + moveUpInFunDef fundef findAllocSizeHoistees diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeUses.hs b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeUses.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeUses.hs @@ -0,0 +1,127 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find out where allocation sizes are used. For each statement, which sizes +-- are in scope? +module Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizeUses + ( findSizeUsesFunDef + ) where + +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe (mapMaybe) +import Control.Monad +import Control.Monad.RWS +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory ( + ExplicitMemory, ExplicitMemorish) +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizes +import Futhark.Optimise.MemoryBlockMerging.PrimExps + + +type SizeVars = Names +type DeclarationsSoFar = Names + +-- The final return value. Describes which size variables are in scope at the +-- creation of the key size variable. +type UsesBefore = M.Map VName Names + +newtype FindM lore a = FindM { unFindM :: RWS SizeVars + UsesBefore DeclarationsSoFar a } + deriving (Monad, Functor, Applicative, + MonadReader SizeVars, + MonadWriter UsesBefore, + MonadState DeclarationsSoFar) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +addDeclarations :: Names -> FindM lore () +addDeclarations = modify . S.union + +addUsesBefore :: VName -> Names -> FindM lore () +addUsesBefore var declarations_so_far = + tell $ M.singleton var declarations_so_far + +findSizeUsesFunDef :: FunDef ExplicitMemory -> UsesBefore +findSizeUsesFunDef fundef = + let size_vars = mapMaybe (subExpVar . fst) $ M.elems $ memBlockSizesFunDef fundef + var_to_pe = findPrimExpsFunDef fundef + -- We want to find 'uses before' for all size vars *and* which variables + -- they depend on. This is a compromise between recording the + -- relationship for only size variables and all variables. We need this + -- compromise for 'sizesCanBeMaxedKernelArray' in Reuse.Core. + find_pe_vars v0 = maybe S.empty + (S.insert v0 . execWriter . traverse + (\v -> do + tell $ S.singleton v + tell $ find_pe_vars v + return v)) $ M.lookup v0 var_to_pe + size_vars' = S.unions $ map find_pe_vars size_vars + m = unFindM $ do + forM_ (funDefParams fundef) lookInFParam + lookInBody $ funDefBody fundef + res = snd $ evalRWS m size_vars' S.empty + in res + +lookInFParam :: FParam lore -> FindM lore () +lookInFParam (Param x _) = + lookAtNewDecls $ S.singleton x + +lookInLParam :: LParam lore -> FindM lore () +lookInLParam (Param x _) = + lookAtNewDecls $ S.singleton x + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm stm@(Let _ _ e) = do + let new_decls = S.fromList $ newDeclarationsStm stm + lookAtNewDecls new_decls + + -- Recursive body walk. + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + , walkOnFParam = lookInFParam + , walkOnLParam = lookInLParam + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInLambda + , walkOnKernelLParam = lookInLParam + } + +lookInLambda :: LoreConstraints lore => + Lambda lore -> FindM lore () +lookInLambda (Lambda params body _) = do + forM_ params lookInLParam + lookInBody body + +lookAtNewDecls :: Names -> FindM lore () +lookAtNewDecls new_decls = do + all_size_vars <- ask + declarations_so_far <- get + let new_size_vars = S.intersection all_size_vars new_decls + forM_ new_size_vars $ \var -> + addUsesBefore var declarations_so_far + addDeclarations new_size_vars diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizes.hs b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizes.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizes.hs @@ -0,0 +1,132 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find all Alloc statements and associate their memory blocks with the +-- allocation size. +module Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizes + ( memBlockSizesFunDef, memBlockSizesParamsBodyNonRec + , Sizes + ) where + +import qualified Data.Map.Strict as M +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + (ExplicitMemorish, ExplicitMemory, InKernel) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Types +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous + + +-- | maps memory blocks to its size and space/type +type Sizes = M.Map MName (SubExp, Space) -- Also Space information + +newtype FindM lore a = FindM { unFindM :: Writer Sizes a } + deriving (Monad, Functor, Applicative, + MonadWriter Sizes) + +type LoreConstraints lore = (ExplicitMemorish lore, + AllocSizeUtils lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +recordMapping :: VName -> (SubExp, Space) -> FindM lore () +recordMapping var (size, space) = tell $ M.singleton var (size, space) + +memBlockSizesFunDef :: LoreConstraints lore => + FunDef lore -> Sizes +memBlockSizesFunDef fundef = + let m = unFindM $ do + mapM_ lookInFParam $ funDefParams fundef + lookInBody $ funDefBody fundef + mem_sizes = execWriter m + in mem_sizes + +memBlockSizesParamsBodyNonRec :: LoreConstraints lore => + [FParam lore] -> Body lore -> Sizes +memBlockSizesParamsBodyNonRec params body = + let m = unFindM $ do + mapM_ lookInFParam params + mapM_ lookInStm $ bodyStms body + mem_sizes = execWriter m + in mem_sizes + +lookInFParam :: LoreConstraints lore => + FParam lore -> FindM lore () +lookInFParam (Param mem (ExpMem.MemMem size space)) = + recordMapping mem (size, space) +lookInFParam _ = return () + +lookInLParam :: LoreConstraints lore => + LParam lore -> FindM lore () +lookInLParam (Param mem (ExpMem.MemMem size space)) = + recordMapping mem (size, space) +lookInLParam _ = return () + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStmRec bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStmRec bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern patctxelems patvalelems) _ e) = do + case patvalelems of + [PatElem mem _] -> case lookForAllocSize e of + Just (size, space) -> + recordMapping mem (size, space) + Nothing -> return () + _ -> return () + mapM_ lookInPatCtxElem patctxelems + +lookInStmRec :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStmRec stm@(Let _ _ e) = do + lookInStm stm + + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + , walkOnFParam = lookInFParam + , walkOnLParam = lookInLParam + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInLambda + , walkOnKernelLParam = lookInLParam + } + +lookInPatCtxElem :: LoreConstraints lore => + PatElem lore -> FindM lore () +lookInPatCtxElem (PatElem mem (ExpMem.MemMem size space)) = + recordMapping mem (size, space) +lookInPatCtxElem _ = return () + +lookInLambda :: LoreConstraints lore => + Lambda lore -> FindM lore () +lookInLambda (Lambda params body _) = do + forM_ params lookInLParam + lookInBody body + +class AllocSizeUtils lore where + lookForAllocSize :: Exp lore -> Maybe (SubExp, Space) + +instance AllocSizeUtils ExplicitMemory where + lookForAllocSize (Op (ExpMem.Alloc size space)) = Just (size, space) + lookForAllocSize _ = Nothing + +instance AllocSizeUtils InKernel where + lookForAllocSize (Op (ExpMem.Alloc size space)) = Just (size, space) + lookForAllocSize _ = Nothing diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/Core.hs b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/Core.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Reuse/Core.hs @@ -0,0 +1,747 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE TupleSections #-} +-- | Find array creations that can be set to use existing memory blocks instead +-- of new allocations. +module Futhark.Optimise.MemoryBlockMerging.Reuse.Core + ( coreReuseFunDef + ) where + +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import qualified Data.List as L +import Data.Maybe (catMaybes, fromMaybe, isJust) +import Control.Monad +import Control.Monad.RWS +import Control.Monad.State +import Control.Monad.Identity + +import Futhark.MonadFreshNames +import Futhark.Binder +import Futhark.Construct +import Futhark.Representation.AST +import Futhark.Analysis.PrimExp +import Futhark.Analysis.PrimExp.Convert +import Futhark.Representation.ExplicitMemory + (ExplicitMemory, ExplicitMemorish) +import Futhark.Pass.ExplicitAllocations() +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.PrimExps (findPrimExpsFunDef) +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types +import Futhark.Optimise.MemoryBlockMerging.MemoryUpdater + +import Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizes +import Futhark.Optimise.MemoryBlockMerging.Reuse.AllocationSizeUses + + +data Context = Context { ctxFirstUses :: FirstUses + -- ^ From the module Liveness.FirstUses + , ctxInterferences :: Interferences + , ctxPotentialKernelInterferences + :: PotentialKernelDataRaceInterferences + -- ^ From the module Liveness.Interferences + , ctxSizes :: Sizes + -- ^ maps a memory block to its size and space + , ctxVarToMem :: VarMemMappings MemorySrc + -- ^ From the module VariableMemory + , ctxActualVars :: M.Map VName Names + -- ^ From the module ActualVariables + , ctxExistentials :: Names + -- ^ From the module Existentials + , ctxVarPrimExps :: M.Map VName (PrimExp VName) + -- ^ From the module PrimExps + , ctxSizeVarsUsesBefore :: M.Map VName Names + -- ^ maps a memory name to the size variables available + -- at that memory block allocation point + } + deriving (Show) + +data Current = Current { curUses :: M.Map MName MNames + -- ^ maps a memory block to the memory blocks that + -- have been merged into it so far + , curEqAsserts :: M.Map VName Names + -- ^ maps a variable name to other semantically equal + -- variable names + + , curVarToMemRes :: VarMemMappings MemoryLoc + -- ^ The result of the core analysis: maps an array + -- name to its memory block. + + , curVarToMaxExpRes :: M.Map MName Names + -- ^ Changes in variable uses where allocation sizes + -- are maxed from its elements. Keyed by statement + -- memory name (alloc stmt). Maps an alloc stmt to the + -- sizes that need to be taken max for. + + , curKernelMaxSizedRes :: M.Map MName (VName, + ((VName, VName), + (VName, VName))) + -- ^ Maps an alloc stmt to + -- (size0, + -- ((array0, size_var0, ixfun0), + -- (array1, size_var1, ixfun1))). + -- + -- Needed for array creations in kernel + -- bodies that can only reuse memory if index functions + -- are changed, and the allocation size is maxed. + -- + -- size_var0 is *not* the size of the entire allocation + -- of the key memory, but *part of* the allocation + -- size. This part will be replaced by the maximum of + -- the two sizes. + } + deriving (Show) + +emptyCurrent :: Current +emptyCurrent = Current { curUses = M.empty + , curEqAsserts = M.empty + , curVarToMemRes = M.empty + , curVarToMaxExpRes = M.empty + , curKernelMaxSizedRes = M.empty + } + +newtype FindM lore a = FindM { unFindM :: RWS Context () Current a } + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadState Current) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- Lookup the memory block statically associated with a variable. +lookupVarMem :: MonadReader Context m => + VName -> m MemorySrc +lookupVarMem var = + -- This should always be called from a place where it is certain that 'var' + -- refers to a statement with an array expression. + fromJust ("lookup memory block from " ++ pretty var) . M.lookup var + <$> asks ctxVarToMem + +lookupActualVars' :: ActualVariables -> VName -> Names +lookupActualVars' actual_vars var = + -- Do this recursively. + let actual_vars' = expandWithAliases actual_vars actual_vars + in fromMaybe (S.singleton var) $ M.lookup var actual_vars' + +lookupActualVars :: MonadReader Context m => + VName -> m Names +lookupActualVars var = asks $ flip lookupActualVars' var . ctxActualVars + +lookupSize :: MonadReader Context m => + VName -> m SubExp +lookupSize var = + fst . fromJust ("lookup size from " ++ pretty var) . M.lookup var + <$> asks ctxSizes + +lookupSpace :: MonadReader Context m => + MName -> m Space +lookupSpace mem = + snd . fromJust ("lookup space from " ++ pretty mem) . M.lookup mem + <$> asks ctxSizes + +-- Record that the existing old_mem now also "is the same as" new_mem. +insertUse :: VName -> VName -> FindM lore () +insertUse old_mem new_mem = + modify $ \cur -> cur { curUses = insertOrUpdate old_mem new_mem $ curUses cur } + +recordMemMapping :: VName -> MemoryLoc -> FindM lore () +recordMemMapping x mem = + modify $ \cur -> cur { curVarToMemRes = M.insert x mem $ curVarToMemRes cur } + +recordMaxMapping :: MName -> VName -> FindM lore () +recordMaxMapping mem y = + modify $ \cur -> cur { curVarToMaxExpRes = insertOrUpdate mem y + $ curVarToMaxExpRes cur } + +recordKernelMaxMapping :: MName -> (VName, ((VName, VName), (VName, VName))) + -> FindM lore () +recordKernelMaxMapping mem info = + modify $ \cur -> cur { curKernelMaxSizedRes = + M.insert mem info $ curKernelMaxSizedRes cur + } + +modifyCurEqAsserts :: (M.Map VName Names -> M.Map VName Names) -> FindM lore () +modifyCurEqAsserts f = modify $ \c -> c { curEqAsserts = f $ curEqAsserts c } + +-- Run a monad with a local copy of the uses. We don't want any new uses in +-- nested bodies to be available for merging into when we are back in the main +-- body, but we do want updates to existing uses to be propagated. +withLocalUses :: FindM lore a -> FindM lore a +withLocalUses m = do + uses_before <- gets curUses + res <- m + uses_after <- gets curUses + -- Only take the results whose memory block keys were also present prior to + -- traversing the sub-body. + let uses_before_updated = M.filterWithKey + (\mem _ -> mem `S.member` M.keysSet uses_before) + uses_after + modify $ \cur -> cur { curUses = uses_before_updated } + return res + +coreReuseFunDef :: MonadFreshNames m => + FunDef ExplicitMemory -> FirstUses -> + Interferences -> PotentialKernelDataRaceInterferences -> + VarMemMappings MemorySrc -> ActualVariables -> Names -> + m (FunDef ExplicitMemory) +coreReuseFunDef fundef first_uses interferences potential_kernel_interferences var_to_mem actual_vars existentials = do + let sizes = memBlockSizesFunDef fundef + size_uses = findSizeUsesFunDef fundef + var_to_pe = findPrimExpsFunDef fundef + context = Context + { ctxFirstUses = first_uses + , ctxInterferences = interferences + , ctxPotentialKernelInterferences = potential_kernel_interferences + , ctxSizes = sizes + , ctxVarToMem = var_to_mem + , ctxActualVars = actual_vars + , ctxExistentials = existentials + , ctxVarPrimExps = var_to_pe + , ctxSizeVarsUsesBefore = size_uses + } + m = unFindM $ do + forM_ (funDefParams fundef) lookInFParam + lookInBody $ funDefBody fundef + (res, ()) = execRWS m context emptyCurrent + var_to_mem_res = curVarToMemRes res + fundef' <- transformFromVarMemMappings var_to_mem_res (M.map memSrcName var_to_mem) (M.map fst sizes) (M.map fst sizes) False fundef + let sizes' = memBlockSizesFunDef fundef' + fundef'' <- transformFromVarMaxExpMappings (curVarToMaxExpRes res) fundef' + transformFromKernelMaxSizedMappings var_to_pe var_to_mem (M.map memLocName var_to_mem_res) sizes' actual_vars (curKernelMaxSizedRes res) fundef'' + +lookInFParam :: LoreConstraints lore => + FParam lore -> FindM lore () +lookInFParam (Param _ membound) = + -- Unique array function parameters also count as "allocations" in which + -- memory can be reused. + case membound of + ExpMem.MemArray _ _ Unique (ExpMem.ArrayIn mem _) -> + insertUse mem mem + _ -> return () + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern _patctxelems patvalelems) _ e) = do + var_to_pe <- asks ctxVarPrimExps + let eqs | BasicOp (Assert (Var v) _ _) <- e + , Just (CmpOpExp (CmpEq _) (LeafExp v0 _) (LeafExp v1 _)) <- M.lookup v var_to_pe = do + modifyCurEqAsserts $ insertOrUpdate v0 v1 + modifyCurEqAsserts $ insertOrUpdate v1 v0 + | otherwise = return () + eqs + + forM_ patvalelems $ \(PatElem var membound) -> do + -- For every declaration with a first memory use, check (through + -- handleNewArray) if it can reuse some earlier memory block. + first_uses_var <- lookupEmptyable var <$> asks ctxFirstUses + actual_vars_var <- lookupActualVars var + existentials <- asks ctxExistentials + case membound of + ExpMem.MemArray _ _ _ (ExpMem.ArrayIn mem _) -> + when (-- We require that it must be a first use, i.e. an array creation. + mem `S.member` first_uses_var + -- If the array is existential or "aliases" something that is + -- existential, we do not try to make it reuse any memory. + && not (var `S.member` existentials) + && not (any (`S.member` existentials) actual_vars_var)) + $ handleNewArray var mem + _ -> return () + + fullWalkExpM walker walker_kernel e + + where walker = identityWalker + { walkOnBody = withLocalUses . lookInBody } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . withLocalUses . lookInBody + , walkOnKernelKernelBody = coerce . withLocalUses . lookInKernelBody + , walkOnKernelLambda = coerce . withLocalUses . lookInBody . lambdaBody + } + +-- Check if a new array declaration x with a first use of the memory xmem can be +-- set to use a previously encountered memory block. +handleNewArray :: VName -> MName -> FindM lore () +handleNewArray x xmem = do + interferences <- asks ctxInterferences + actual_vars <- lookupActualVars x + + let notTheSame :: Monad m => MName -> MNames -> m Bool + notTheSame kmem _used_mems = return (kmem /= xmem) + + let noneInterfere :: Monad m => MName -> MNames -> m Bool + noneInterfere _kmem used_mems = + -- A memory block can have already been reused. We also check for + -- interference with any previously merged blocks. + return $ all (\used_mem -> not $ S.member xmem + $ lookupEmptyable used_mem interferences) + $ S.toList used_mems + + let noneInterfereKernelArray :: MonadReader Context m => MNames -> m Bool + noneInterfereKernelArray used_mems = + not <$> anyM (interferesInKernel xmem) (S.toList used_mems) + + let sameSpace :: MonadReader Context m => + MName -> MNames -> m Bool + sameSpace kmem _used_mems = do + kspace <- lookupSpace kmem + xspace <- lookupSpace xmem + return (kspace == xspace) + + -- Is the size of the new memory block (xmem) equal to any of the memory + -- blocks (used_mems) using an already used memory block? + let sizesMatch :: MNames -> FindM lore Bool + sizesMatch used_mems = do + ok_sizes <- mapM lookupSize $ S.toList used_mems + new_size <- lookupSize xmem + -- Check for size equality by checking for variable name equality. + let eq_simple = new_size `L.elem` ok_sizes + + -- Check for size equality by constructing 'PrimExp's and comparing + -- those. Use the custom VarWithLooseEquality type to compare inner + -- sizes: If an equality assert statement was found earlier, consider + -- its two operands to be the same. + var_to_pe <- asks ctxVarPrimExps + eq_asserts <- gets curEqAsserts + let sePrimExp se = do + v <- subExpVar se + pe <- M.lookup v var_to_pe + let pe_expanded = expandPrimExp var_to_pe pe + traverse (\v_inner -> -- Has custom Eq instance. + pure $ VarWithLooseEquality v_inner + $ lookupEmptyable v_inner eq_asserts + ) pe_expanded + let ok_sizes_pe = map sePrimExp ok_sizes + let new_size_pe = sePrimExp new_size + + -- If new_size_pe actually denotes a PrimExp, check if it is among the + -- constructed 'PrimExp's of the sizes of the memory blocks that have + -- already been set to use the target memory block. + let eq_advanced = isJust new_size_pe && new_size_pe `L.elem` ok_sizes_pe + + return (eq_simple || eq_advanced) + + -- In case sizes do not match: Is it possible to change the size of the target + -- memory block to be a maximum of itself and the new memory block? + let sizesCanBeMaxed :: MName -> FindM lore Bool + sizesCanBeMaxed kmem = do + ksize <- lookupSize kmem + xsize <- lookupSize xmem + uses_before <- asks ctxSizeVarsUsesBefore + let ok = fromMaybe False $ do + ksize' <- subExpVar ksize + xsize' <- subExpVar xsize + return (xsize' `S.member` fromJust ("is recorded for all size variables " + ++ pretty ksize') + (M.lookup ksize' uses_before)) + return ok + + let sizesCanBeMaxedKernelArray :: MName -> MNames -> + FindM lore (Maybe (VName, ((VName, VName), + (VName, VName)))) + sizesCanBeMaxedKernelArray kmem used_mems = do + -- Let a kernel body have two indexed array creations result_0 and + -- result_1 with the index functions + -- + -- result_0: ixfun_start_0[indices_start_0, 0i64:+res_0*1i64] + -- result_1: ixfun_start_1[indices_start_1, 0i64:+res_1*1i64] + -- + -- with the additional requirements that + -- + -- + ixfun_start_0 is equal to ixfun_start_1 except for mentions of + -- res_0 and res_1. + -- + -- + indices_start_0 is equal to indices_start_1. + -- + -- Example: + -- + -- result_0: Direct(num_groups, res_0, group_size)[0, 2, 1][group_id, local_tid, 0i64:+res_0*1i64] + -- result_1: Direct(num_groups, res_1, group_size)[0, 2, 1][group_id, local_tid, 0i64:+res_1*1i64] + -- + -- By default result_0 and result_1 will be set to interfere because + -- each thread can access parts of the memory of another thread if they + -- are merged. We can fix this my making both index functions describe + -- the same access pattern except for the final dimension. We want this + -- to happen for the example above: + -- + -- result_0': Direct(num_groups, res_max, group_size)[0, 2, 1][group_id, local_tid, 0i64:+res_0*1i64] + -- result_1': Direct(num_groups, res_max, group_size)[0, 2, 1][group_id, local_tid, 0i64:+res_1*1i64] + -- + -- Where res_max = max(res_0, res_1). Now they cover the same area in + -- space. The final index slices are kept as they were, since the shape + -- of the created array should stay the same. This means that the + -- smallest array will not be writing to all of its available space. + -- + -- We need to check: + -- + -- + Is res_1 in scope at the allocation? Allocation size hoisting + -- has probably been helpful here. + -- + -- + Does res_0 and res_1 have the same base type size? + -- + -- If true, modify the program as such: + -- + -- + Insert a res_max statement before the allocation. + -- + -- + Change the allocation size to use res_max instead of res_0. + -- + -- + Modify both index functions to use res_max instead of res_0 and + -- res_1, respectively, except for at the final index slice. + -- + -- Extension: If an array reuses an already reused array, remember to + -- update *all* index functions. Currently we avoid these cases for + -- simplicity of implementation. + + potentials <- asks ctxPotentialKernelInterferences + uses_before <- asks ctxSizeVarsUsesBefore + + let first_usess = filter (\p -> + let pot_mems = map (\(m, _, _, _) -> m) p + in kmem `elem` pot_mems && xmem `elem` pot_mems) + potentials + kmem_size <- fromJust "should be a var" . subExpVar <$> lookupSize kmem + + return $ case (S.toList used_mems, first_usess) of + -- We only support the basic case for now. FIXME (or, at the very + -- least, manage to create a program where this will have an effect). + -- + -- A used_mems list of size > 1 means that kmem has already been + -- reused. This is okay, but a bit harder to keep track of. + -- + -- A first_usess list of size > 1 means that xmem and kmem + -- data-race-interfere in multiple kernels. This will never happen in + -- the current implementation, but could *potentially* happen in the + -- future. + ([_], [first_uses]) -> do + (_, kmem_array, kmem_pt, kmem_ixfun) <- + L.find (\(mname, _, _, _) -> mname == kmem) first_uses + (_, xmem_array, xmem_pt, xmem_ixfun) <- + L.find (\(mname, _, _, _) -> mname == xmem) first_uses + + if (kmem, kmem_ixfun) `ixFunsCompatible` (xmem, xmem_ixfun) + then Nothing -- These are not special, and need not special handling. + else do + (kmem_ixfun_start, kmem_indices_start, kmem_final_dim) <- + IxFun.getInfoMaxUnification kmem_ixfun + (xmem_ixfun_start, xmem_indices_start, xmem_final_dim) <- + IxFun.getInfoMaxUnification xmem_ixfun + + let xmem_final_dim_before_kmem_final_dim = + maybe False (xmem_final_dim `S.member`) $ + M.lookup kmem_final_dim uses_before + kmem_ixfun_start' = getIxFun' kmem_ixfun_start + (M.singleton kmem_final_dim xmem_final_dim) + xmem_ixfun_start' = getIxFun' xmem_ixfun_start + (M.singleton xmem_final_dim kmem_final_dim) + + res = if kmem_indices_start == xmem_indices_start && + (kmem, kmem_ixfun_start') `ixFunsCompatible` + (xmem, xmem_ixfun_start') && + (primByteSize kmem_pt :: Int) == primByteSize xmem_pt && + xmem_final_dim_before_kmem_final_dim + then return (kmem_size, + ((kmem_array, kmem_final_dim), + (xmem_array, xmem_final_dim))) + else Nothing + + in res + _ -> Nothing + + where getIxFun' :: ExpMem.IxFun -> M.Map VName VName -> + IxFun.IxFun (PrimExp VarWithLooseEquality) + getIxFun' ixfun others = + let loose_eq_map name_inner = + -- Has custom Eq instance. + pure $ VarWithLooseEquality name_inner + $ maybe S.empty S.singleton $ M.lookup name_inner others + in runIdentity $ traverse (traverse loose_eq_map) ixfun + + let sizesCanBeMaxedKernelArray' :: MName -> MNames -> FindM lore Bool + sizesCanBeMaxedKernelArray' kmem used_mems = + isJust <$> sizesCanBeMaxedKernelArray kmem used_mems + + let noOtherUsesOfMemory :: MName -> MNames -> FindM lore Bool + noOtherUsesOfMemory _kmem _used_mems = + -- If the array in question 'x' is not the only array that uses the + -- memory (ignoring aliasing), then do not perform memory reuse. We + -- only want to reuse memory if it means we can remove an allocation. + -- FIXME: If we can check that all arrays using the memory in question + -- 'xmem' can be set to reuse some other memory, so that 'xmem' does not + -- have to be allocated, then this restriction can go away. It also + -- might be the case that the ActualVariables module does not find all + -- array connections, i.e. it concludes that two arrays are distinct + -- when they are actually not; this can happen with streams. + and . M.elems . M.mapWithKey ( + \v m -> (memSrcName m /= xmem) + || (v `L.elem` actual_vars) + ) <$> asks ctxVarToMem + + let notCurrentlyDisabled :: FindM lore Bool + notCurrentlyDisabled = + -- FIXME: We currently disable reusing memory of constant size. This is + -- a problem in the misc/heston/heston32.fut benchmark (but not the + -- heston64.fut one). It would be nice to not have to disable this + -- feature, as it works well for the most part. Why is this a problem? + -- Or is it maybe something else that causes heston32 to segfault? + isJust . subExpVar <$> lookupSize xmem + + let sizesWorkOut :: MName -> MNames -> FindM lore Bool + sizesWorkOut kmem used_mems = + -- The size of an allocation is okay to reuse if it is the same as the + -- current memory size, or if it can be changed to be the maximum size + -- of the two sizes. + (notCurrentlyDisabled <&&> noneInterfereKernelArray used_mems <&&> + (sizesMatch used_mems <||> sizesCanBeMaxed kmem)) + <||> sizesCanBeMaxedKernelArray' kmem used_mems + + let canBeUsed t = and <$> mapM (($ t) . uncurry) + [notTheSame, noneInterfere, sameSpace, noOtherUsesOfMemory, + sizesWorkOut] + cur_uses <- gets curUses + found_use <- catMaybes <$> mapM (maybeFromBoolM canBeUsed) (M.assocs cur_uses) + + case found_use of + (kmem, used_mems) : _ -> do + -- There is a previous memory block that we can use. Record the mapping. + insertUse kmem xmem + forM_ actual_vars $ \var -> do + ixfun <- memSrcIxFun <$> lookupVarMem var + recordMemMapping var $ MemoryLoc kmem ixfun -- Only change the memory block. + + -- Record any size-maximum change in case of sizesCanBeMaxed returning + -- True. + whenM (sizesCanBeMaxed kmem) $ do + ksize <- lookupSize kmem + xsize <- lookupSize xmem + fromMaybe (return ()) $ do + ksize' <- subExpVar ksize + xsize' <- subExpVar xsize + return $ do + recordMaxMapping kmem ksize' + recordMaxMapping kmem xsize' + + -- If we are inside a kernel body, and the current array can use the + -- memory block of another array if its size gets maximised, record this + -- change. The actual program transformation will happen later. + kernel_maxing <- sizesCanBeMaxedKernelArray kmem used_mems + forM_ kernel_maxing $ \info -> + recordKernelMaxMapping kmem info + + _ -> + -- There is no previous memory block available for use. Record that this + -- memory block is available. + insertUse xmem xmem + +data VarWithLooseEquality = VarWithLooseEquality VName Names + deriving (Show) + +instance Eq VarWithLooseEquality where + VarWithLooseEquality v0 vs0 == VarWithLooseEquality v1 vs1 = + not $ S.null $ S.intersection (S.insert v0 vs0) (S.insert v1 vs1) + +interferesInKernel :: MonadReader Context m => MName -> MName -> m Bool +interferesInKernel mem0 mem1 = do + potentials <- asks ctxPotentialKernelInterferences + + let interferesInGroup :: PotentialKernelDataRaceInterferenceGroup -> Bool + interferesInGroup first_uses = fromMaybe False $ do + (_, _, pt0, ixfun0) <- L.find (\(mname, _, _, _) -> mname == mem0) first_uses + (_, _, pt1, ixfun1) <- L.find (\(mname, _, _, _) -> mname == mem1) first_uses + return $ interferes (pt0, ixfun0) (pt1, ixfun1) + + interferes :: (PrimType, ExpMem.IxFun) -> (PrimType, ExpMem.IxFun) -> Bool + interferes (pt0, ixfun0) (pt1, ixfun1) = + -- Must be different. + mem0 /= mem1 && + ( + -- Do the index functions range over different memory areas? + ((ixFunHasIndex ixfun0 || ixFunHasIndex ixfun1) && + not (ixFunsCompatible (mem0, ixfun0) (mem1, ixfun1))) + || + -- Do the arrays have different base type size? If so, they take + -- up different amounts of space, and will not be compatible. + ((primByteSize pt0 :: Int) /= primByteSize pt1) + ) + + return $ any interferesInGroup potentials + +-- Does an index function contain an Index expression? +-- +-- If the index function of the memory annotation uses an index, it means that +-- the array creation does not refer to the entire array. It is an array +-- creation, but only partially: It creates part of the array, and another part +-- is created in another loop iteration or kernel thread. The danger in +-- declaring this memory a first use lies in how it can then be reused later in +-- the iteration/thread by some memory with a *different* index in its memory +-- annotation index function, which can affect reads in other threads. +ixFunHasIndex :: IxFun.IxFun num -> Bool +ixFunHasIndex = IxFun.ixFunHasIndex + +-- Do the two index functions describe the same range? In other words, does one +-- array take up precisely the same location (offset) and size as another array +-- relative to the beginning of their respective memory blocks? FIXME: This can +-- be less conservative, for example by handling that different reshapes of the +-- same array can describe the same offset and space, but do we have any tests +-- or benchmarks where that occurs? +ixFunsCompatible :: Eq v => + (MName, IxFun.IxFun (PrimExp v)) -> (MName, IxFun.IxFun (PrimExp v)) -> + Bool +ixFunsCompatible (_mem0, ixfun0) (_mem1, ixfun1) = + IxFun.ixFunsCompatibleRaw ixfun0 ixfun1 + +-- Replace certain allocation sizes in a program with new variables describing +-- the maximum of two or more allocation sizes. +transformFromVarMaxExpMappings :: MonadFreshNames m => + M.Map VName Names + -> FunDef ExplicitMemory -> m (FunDef ExplicitMemory) +transformFromVarMaxExpMappings var_to_max fundef = do + var_to_new_var <- + M.fromList <$> mapM (\(k, v) -> (k,) <$> maxsToReplacement (S.toList v)) + (M.assocs var_to_max) + return $ insertAndReplace var_to_new_var fundef + +-- A replacement is a new size variable and any new subexpressions that the new +-- variable depends on. +data Replacement = Replacement + { replName :: VName -- The new variable + , replStms :: [Stm ExplicitMemory] -- The new expressions + } + deriving (Show) + +-- Take a list of size variables. Return a replacement consisting of a size +-- variable denoting the maximum of the input sizes. +maxsToReplacement :: MonadFreshNames m => + [VName] -> m Replacement +maxsToReplacement [] = error "maxsToReplacements: Cannot take max of zero variables" +maxsToReplacement [v] = return $ Replacement v [] +maxsToReplacement vs = do + -- Should be O(lg N) number of new expressions. + let (vs0, vs1) = splitAt (length vs `div` 2) vs + Replacement m0 es0 <- maxsToReplacement vs0 + Replacement m1 es1 <- maxsToReplacement vs1 + vmax <- newVName "max" + let emax = BasicOp $ BinOp (SMax Int64) (Var m0) (Var m1) + new_stm = Let (Pattern [] [PatElem vmax + (ExpMem.MemPrim (IntType Int64))]) (defAux ()) emax + prev_stms = es0 ++ es1 ++ [new_stm] + return $ Replacement vmax prev_stms + +-- Modify a function to use the new replacements. +insertAndReplace :: M.Map MName Replacement -> FunDef ExplicitMemory -> + FunDef ExplicitMemory +insertAndReplace replaces0 fundef = + let body' = evalState (transformBody $ funDefBody fundef) replaces0 + in fundef { funDefBody = body' } + + where transformBody :: Body ExplicitMemory -> + State (M.Map VName Replacement) (Body ExplicitMemory) + transformBody body = do + stms' <- concat <$> mapM transformStm (stmsToList $ bodyStms body) + return $ body { bodyStms = stmsFromList stms' } + + transformStm :: Stm ExplicitMemory -> + State (M.Map VName Replacement) [Stm ExplicitMemory] + transformStm stm@(Let (Pattern [] [PatElem mem_name + (ExpMem.MemMem _ pat_space)]) _ + (Op (ExpMem.Alloc _ space))) = do + replaces <- get + case M.lookup mem_name replaces of + Just repl -> do + let prev = replStms repl + new = Let (Pattern [] [PatElem mem_name + (ExpMem.MemMem (Var (replName repl)) + pat_space)]) (defAux ()) + (Op (ExpMem.Alloc (Var (replName repl)) space)) + -- We should only generate the new statements once. + modify $ M.adjust (\repl0 -> repl0 { replStms = [] }) mem_name + return (prev ++ [new]) + Nothing -> return [stm] + transformStm (Let pat attr e) = do + let mapper = identityMapper { mapOnBody = const transformBody } + e' <- mapExpM mapper e + return [Let pat attr e'] + + +-- Change certain allocation sizes in a program. +transformFromKernelMaxSizedMappings :: MonadFreshNames m => + M.Map VName (PrimExp VName) -> VarMemMappings MemorySrc -> VarMemMappings MName -> + Sizes -> ActualVariables -> M.Map MName (VName, ((VName, VName), + (VName, VName))) -> + FunDef ExplicitMemory -> m (FunDef ExplicitMemory) +transformFromKernelMaxSizedMappings + var_to_pe var_to_mem var_to_mem_res sizes_orig actual_vars mem_to_info fundef = do + (mem_to_size_var, arr_to_mem_ixfun) <- + unzip <$> mapM (uncurry withNewMaxVar) (M.assocs mem_to_info) + let mem_to_size_var' = M.fromList mem_to_size_var + arr_to_memloc = M.fromList $ map (\(arr, destmem, ixfun) -> + (arr, MemoryLoc destmem ixfun)) + $ concat arr_to_mem_ixfun + + fundef' = insertAndReplace mem_to_size_var' fundef + sizes = memBlockSizesFunDef fundef' + transformFromVarMemMappings arr_to_memloc (M.union var_to_mem_res (M.map memSrcName var_to_mem)) (M.map fst sizes) (M.map fst sizes_orig) True fundef' + + where withNewMaxVar :: MonadFreshNames m => + MName -> (VName, + ((VName, VName), + (VName, VName))) -> + m ((MName, Replacement), + [(VName, MName, ExpMem.IxFun)]) + withNewMaxVar mem (kmem_size, + ((kmem_array, kmem_final_dim), + (xmem_array, xmem_final_dim))) = do + final_dim_max_v <- newVName "max_final_dim" + let final_dim_max_e = + BasicOp (BinOp (SMax Int32) + (Var kmem_final_dim) (Var xmem_final_dim)) + + var_to_pe_extension = + M.singleton kmem_final_dim (LeafExp final_dim_max_v (IntType Int32)) + var_to_pe' = M.union var_to_pe_extension var_to_pe + full_size_pe = fromJust "should exist" $ M.lookup kmem_size var_to_pe + full_size_pe_expanded = expandPrimExp var_to_pe' full_size_pe + new_full_size_m = + letExp "max" =<< primExpToExp (return . BasicOp . SubExp . Var) + full_size_pe_expanded + (alloc_size_var, alloc_size_stms) <- + modifyNameSource $ runState $ runBinderT new_full_size_m mempty + let alloc_size_fd_stm = + Let (Pattern [] [PatElem final_dim_max_v + (ExpMem.MemPrim (IntType Int32))]) (defAux ()) final_dim_max_e + alloc_size_stms' = oneStm alloc_size_fd_stm <> alloc_size_stms + + vars_kmem = + S.insert kmem_array $ lookupActualVars' actual_vars kmem_array + vars_xmem = + S.insert xmem_array $ lookupActualVars' actual_vars xmem_array + + arrayToMapping final_dim v = + let ixfun = memSrcIxFun $ fromJust "should exist" + $ M.lookup v var_to_mem + ixfun_new = IxFun.subsInIndexIxFun ixfun final_dim final_dim_max_v --newIxFun ixfun final_dim + in (v, mem, ixfun_new) + arr_to_mem_ixfun_kmem = map (arrayToMapping kmem_final_dim) + $ S.toList vars_kmem + arr_to_mem_ixfun_xmem = map (arrayToMapping xmem_final_dim) + $ S.toList vars_xmem + arr_to_mem_ixfun = arr_to_mem_ixfun_kmem ++ arr_to_mem_ixfun_xmem + + return ((mem, Replacement alloc_size_var $ stmsToList alloc_size_stms'), + arr_to_mem_ixfun) diff --git a/src/Futhark/Optimise/MemoryBlockMerging/Types.hs b/src/Futhark/Optimise/MemoryBlockMerging/Types.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/Types.hs @@ -0,0 +1,91 @@ +module Futhark.Optimise.MemoryBlockMerging.Types + ( MName + , MNames + , MemorySrc(..) + , MemoryLoc(..) + , VarMemMappings + , MemAliases + , VarAliases + , FirstUses + , StmOrRes(..) + , LastUses + , Interferences + , ActualVariables + , PotentialKernelDataRaceInterferences + , PotentialKernelDataRaceInterferenceGroup + , KernelFirstUse + ) +where + +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem + +import Futhark.Representation.AST +import qualified Futhark.Representation.ExplicitMemory as ExpMem + + +-- | Memory block VName. +type MName = VName + +-- | Memory block names. +type MNames = Names + +data MemorySrc = MemorySrc + { memSrcName :: MName -- ^ the memory block name + , memSrcIxFun :: ExpMem.IxFun -- ^ the index function into the memory + , memSrcShape :: Shape -- ^ the shape of the original array + } + deriving (Show, Eq) + +data MemoryLoc = MemoryLoc + { memLocName :: MName -- ^ the memory block name + , memLocIxFun :: ExpMem.IxFun -- ^ the index function into the memory + } + deriving (Show, Eq) + +-- A mapping from variable names to memory blocks (with varying details) +type VarMemMappings t = M.Map VName t + +-- Aliasing of memory blocks, meaning multiple memory blocks refer to the same +-- actualy memory. Aliasing is not commutative. +type MemAliases = M.Map MName MNames + +-- Aliasing of variables, meaning the use the same memory blocks. Aliasing is +-- commutative? +type VarAliases = M.Map VName Names + +-- First uses of memory blocks in statement denoted by variable name. +type FirstUses = M.Map VName MNames + +-- A last use can occur in a statement OR in a body result. +data StmOrRes = FromStm VName + | FromRes VName + deriving (Show, Eq, Ord) +type LastUses = M.Map StmOrRes MNames + +-- Interferences between memory blocks. +type Interferences = M.Map MName MNames + +-- Sets of potential interferences inside kernels because of potential data +-- races. For each set, every memory block *can* interfere with every other +-- memory block, but only in dire edge cases. Usually some of them can be said +-- to not interfere, and sometimes array creation statements can be modified to +-- have fewer interferences. See Reuse/Core.hs. +type PotentialKernelDataRaceInterferences = + [PotentialKernelDataRaceInterferenceGroup] +type PotentialKernelDataRaceInterferenceGroup = [KernelFirstUse] +type KernelFirstUse = (MName, VName, PrimType, ExpMem.IxFun) + +-- "Links" for handling how variables belong together. +type ActualVariables = M.Map VName Names + +-- Log keeping. Statement variable names to a list of topic-content-mappings. +newtype Log = Log (M.Map VName [(String, String)]) + deriving (Show, Eq, Ord) + +instance Sem.Semigroup Log where + Log a <> Log b = Log $ M.unionWith (++) a b + +instance Monoid Log where + mempty = Log M.empty + mappend = (Sem.<>) diff --git a/src/Futhark/Optimise/MemoryBlockMerging/VariableAliases.hs b/src/Futhark/Optimise/MemoryBlockMerging/VariableAliases.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/VariableAliases.hs @@ -0,0 +1,82 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find all variable aliases. Avoids having to use the Aliases representation +-- in other modules. +-- +-- FIXME: This module is silly. It should be able to go away, with the other +-- modules getting variable aliases by using the Aliases representation +-- directly. +module Futhark.Optimise.MemoryBlockMerging.VariableAliases + ( findVarAliases + ) where + +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.Aliases (Aliases, unNames) +import Futhark.Representation.ExplicitMemory + (ExplicitMemorish, ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel +import Futhark.Analysis.Alias (analyseFun) + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types + + +newtype FindM lore a = FindM { unFindM :: Writer [VarAliases] a } + deriving (Monad, Functor, Applicative, + MonadWriter [VarAliases]) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalkAliases lore) + +recordMapping :: VName -> Names -> FindM lore () +recordMapping var names = tell [M.singleton var names] + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- | Find all variable aliases in a function definition. +findVarAliases :: FunDef ExplicitMemory -> VarAliases +findVarAliases fundef = + let fundef' = analyseFun fundef + m = unFindM $ lookInBody $ funDefBody fundef' + var_aliases = M.unionsWith S.union $ execWriter m + var_aliases' = removeEmptyMaps $ expandWithAliases var_aliases var_aliases + in var_aliases' + +lookInBody :: LoreConstraints lore => + Body (Aliases lore) -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody (Aliases lore) -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm (Aliases lore) -> FindM lore () +lookInStm (Let (Pattern _patctxelems patvalelems) _ e) = do + mapM_ lookInPatValElem patvalelems + fullWalkAliasesExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInBody . lambdaBody + } + +lookInPatValElem :: LoreConstraints lore => + PatElem (Aliases lore) -> FindM lore () +lookInPatValElem (PatElem x (names', ExpMem.MemArray{})) = do + let aliases = unNames names' + recordMapping x aliases +lookInPatValElem _ = return () diff --git a/src/Futhark/Optimise/MemoryBlockMerging/VariableMemory.hs b/src/Futhark/Optimise/MemoryBlockMerging/VariableMemory.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MemoryBlockMerging/VariableMemory.hs @@ -0,0 +1,99 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Find all variable-to-memory mappings, so that other modules can lookup the +-- relation. Maps array names to memory blocks. + +module Futhark.Optimise.MemoryBlockMerging.VariableMemory + ( findVarMemMappings + ) where + +import qualified Data.Map.Strict as M +import Control.Monad.Writer + +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory + (ExplicitMemorish, ExplicitMemory) +import qualified Futhark.Representation.ExplicitMemory as ExpMem +import Futhark.Representation.Kernels.Kernel + +import Futhark.Optimise.MemoryBlockMerging.Miscellaneous +import Futhark.Optimise.MemoryBlockMerging.Types + + +newtype FindM lore a = FindM { unFindM :: Writer (VarMemMappings MemorySrc) a } + deriving (Monad, Functor, Applicative, + MonadWriter (VarMemMappings MemorySrc)) + +type LoreConstraints lore = (ExplicitMemorish lore, + FullWalk lore) + +recordMapping :: VName -> MemorySrc -> FindM lore () +recordMapping var memloc = tell $ M.singleton var memloc + +coerce :: FindM flore a -> FindM tlore a +coerce = FindM . unFindM + +-- | Find all variable-memory block mappings in a function definition. +findVarMemMappings :: FunDef ExplicitMemory -> VarMemMappings MemorySrc +findVarMemMappings fundef = + let m = unFindM $ do + mapM_ lookInFParam $ funDefParams fundef + lookInBody $ funDefBody fundef + var_to_mem = execWriter m + in var_to_mem + +lookInFParam :: LoreConstraints lore => + FParam lore -> FindM lore () +lookInFParam (Param x (ExpMem.MemArray _ shape _ (ExpMem.ArrayIn xmem xixfun))) = do + let memloc = MemorySrc xmem xixfun shape + recordMapping x memloc +lookInFParam _ = return () + +lookInLParam :: LoreConstraints lore => + LParam lore -> FindM lore () +lookInLParam (Param x (ExpMem.MemArray _ shape _ (ExpMem.ArrayIn xmem xixfun))) = do + let memloc = MemorySrc xmem xixfun shape + recordMapping x memloc +lookInLParam _ = return () + +lookInBody :: LoreConstraints lore => + Body lore -> FindM lore () +lookInBody (Body _ bnds _res) = + mapM_ lookInStm bnds + +lookInKernelBody :: LoreConstraints lore => + KernelBody lore -> FindM lore () +lookInKernelBody (KernelBody _ bnds _res) = + mapM_ lookInStm bnds + +lookInStm :: LoreConstraints lore => + Stm lore -> FindM lore () +lookInStm (Let (Pattern _patctxelems patvalelems) _ e) = do + mapM_ lookInPatValElem patvalelems + fullWalkExpM walker walker_kernel e + where walker = identityWalker + { walkOnBody = lookInBody + , walkOnFParam = lookInFParam + , walkOnLParam = lookInLParam + } + walker_kernel = identityKernelWalker + { walkOnKernelBody = coerce . lookInBody + , walkOnKernelKernelBody = coerce . lookInKernelBody + , walkOnKernelLambda = coerce . lookInLambda + , walkOnKernelLParam = lookInLParam + } + +lookInPatValElem :: LoreConstraints lore => + PatElem lore -> FindM lore () +lookInPatValElem (PatElem x (ExpMem.MemArray _ shape _ (ExpMem.ArrayIn xmem xixfun))) = do + let memloc = MemorySrc xmem xixfun shape + recordMapping x memloc +lookInPatValElem _ = return () + +lookInLambda :: LoreConstraints lore => + Lambda lore -> FindM lore () +lookInLambda (Lambda params body _) = do + forM_ params lookInLParam + lookInBody body diff --git a/src/Futhark/Optimise/Simplify.hs b/src/Futhark/Optimise/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +module Futhark.Optimise.Simplify + ( simplifyProg + , simplifySomething + , simplifyFun + , simplifyLambda + , simplifyStms + + , Engine.SimpleOps (..) + , Engine.SimpleM + , Engine.SimplifyOp + , Engine.bindableSimpleOps + , Engine.noExtraHoistBlockers + , Engine.SimplifiableLore + , Engine.HoistBlockers + , RuleBook + ) + where + +import Data.Semigroup ((<>)) + +import Futhark.Representation.AST +import Futhark.MonadFreshNames +import qualified Futhark.Optimise.Simplify.Engine as Engine +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Optimise.Simplify.Rule +import Futhark.Optimise.Simplify.Lore +import Futhark.Pass + +-- | Simplify the given program. Even if the output differs from the +-- 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) +simplifyProg simpl rules blockers = + intraproceduralTransformation $ simplifyFun simpl rules blockers + +-- | Run a simplification operation to convergence. +simplifySomething :: (MonadFreshNames m, HasScope lore m, + Engine.SimplifiableLore lore) => + (a -> Engine.SimpleM lore b) + -> (b -> a) + -> Engine.SimpleOps lore + -> RuleBook (Wise lore) + -> Engine.HoistBlockers lore + -> a + -> m a +simplifySomething f g simpl rules blockers x = do + scope <- askScope + let f' x' = Engine.localVtable (ST.fromScope (addScopeWisdom scope)<>) $ f x' + loopUntilConvergence env simpl f' g x + where env = Engine.emptyEnv rules blockers + +-- | Simplify the given function. Even if the output differs from the +-- output, meaningful simplification may not have taken place - the +-- 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 + -> FunDef lore + -> m (FunDef lore) +simplifyFun simpl rules blockers = + loopUntilConvergence env simpl Engine.simplifyFun removeFunDefWisdom + where env = Engine.emptyEnv rules blockers + +-- | Simplify just a single 'Lambda'. +simplifyLambda :: (MonadFreshNames m, HasScope lore m, Engine.SimplifiableLore lore) => + Engine.SimpleOps lore + -> RuleBook (Engine.Wise lore) + -> Engine.HoistBlockers lore + -> Lambda lore -> [Maybe VName] + -> m (Lambda lore) +simplifyLambda simpl rules blockers orig_lam args = + simplifySomething f removeLambdaWisdom simpl rules blockers orig_lam + where f lam' = Engine.simplifyLambdaNoHoisting lam' args + +-- | Simplify a list of 'Stm's. +simplifyStms :: (MonadFreshNames m, HasScope lore m, Engine.SimplifiableLore lore) => + Engine.SimpleOps lore + -> RuleBook (Engine.Wise lore) + -> Engine.HoistBlockers lore + -> Stms lore + -> m (Stms lore) +simplifyStms = simplifySomething f g + where f stms = fmap snd $ Engine.simplifyStms stms $ return ((), mempty) + g = fmap removeStmWisdom + +loopUntilConvergence :: (MonadFreshNames m, Engine.SimplifiableLore lore) => + Engine.Env lore + -> Engine.SimpleOps lore + -> (a -> Engine.SimpleM lore b) + -> (b -> a) + -> a + -> m a +loopUntilConvergence env simpl f g x = do + (x', changed) <- modifyNameSource $ Engine.runSimpleM (f x) simpl env + if changed then loopUntilConvergence env simpl f g (g x') else return $ g x' diff --git a/src/Futhark/Optimise/Simplify/ClosedForm.hs b/src/Futhark/Optimise/Simplify/ClosedForm.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/ClosedForm.hs @@ -0,0 +1,180 @@ +{-# LANGUAGE FlexibleContexts #-} +-- | This module implements facilities for determining whether a +-- reduction or fold can be expressed in a closed form (i.e. not as a +-- SOAC). +-- +-- Right now, the module can detect only trivial cases. In the +-- future, we would like to make it more powerful, as well as possibly +-- also being able to analyse sequential loops. +module Futhark.Optimise.Simplify.ClosedForm + ( foldClosedForm + , loopClosedForm + ) +where + +import Control.Monad +import Data.Maybe +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Semigroup ((<>)) + +import Futhark.Construct +import Futhark.Representation.AST +import Futhark.Transform.Rename +import Futhark.Optimise.Simplify.Rule + +-- | A function that, given a variable name, returns its definition. +type VarLookup lore = VName -> Maybe (Exp lore, Certificates) + +{- +Motivation: + + let {*[int,x_size_27] map_computed_shape_1286} = replicate(x_size_27, + all_equal_shape_1044) in + let {*[bool,x_size_27] map_size_checks_1292} = replicate(x_size_27, x_1291) in + let {bool all_equal_checked_1298, int all_equal_shape_1299} = + reduceT(fn {bool, int} (bool bacc_1293, int nacc_1294, bool belm_1295, + int nelm_1296) => + let {bool tuplit_elems_1297} = bacc_1293 && belm_1295 in + {tuplit_elems_1297, nelm_1296}, + {True, 0}, map_size_checks_1292, map_computed_shape_1286) +-} + +-- | @foldClosedForm look foldfun accargs arrargs@ determines whether +-- each of the results of @foldfun@ can be expressed in a closed form. +foldClosedForm :: (Attributes lore, BinderOps lore) => + VarLookup lore + -> Pattern lore + -> Lambda lore + -> [SubExp] -> [VName] + -> RuleM lore () + +foldClosedForm look pat lam accs arrs = do + inputsize <- arraysSize 0 <$> mapM lookupType arrs + + t <- case patternTypes pat of [Prim t] -> return t + _ -> cannotSimplify + + closedBody <- checkResults (patternNames pat) inputsize mempty knownBnds + (map paramName (lambdaParams lam)) + (lambdaBody lam) accs + isEmpty <- newVName "fold_input_is_empty" + letBindNames_ [isEmpty] $ + BasicOp $ CmpOp (CmpEq int32) inputsize (intConst Int32 0) + letBind_ pat =<< (If (Var isEmpty) + <$> resultBodyM accs + <*> renameBody closedBody + <*> pure (IfAttr [primBodyType t] IfNormal)) + where knownBnds = determineKnownBindings look lam accs arrs + +-- | @loopClosedForm pat respat merge bound bodys@ determines whether +-- the do-loop can be expressed in a closed form. +loopClosedForm :: (Attributes lore, BinderOps lore) => + Pattern lore + -> [(FParam lore,SubExp)] + -> Names -> SubExp -> Body lore + -> RuleM lore () +loopClosedForm pat merge i bound body = do + t <- case patternTypes pat of [Prim t] -> return t + _ -> cannotSimplify + + closedBody <- checkResults mergenames bound i knownBnds + (map identName mergeidents) body mergeexp + isEmpty <- newVName "bound_is_zero" + letBindNames_ [isEmpty] $ + BasicOp $ CmpOp (CmpSlt Int32) bound (intConst Int32 0) + + letBind_ pat =<< (If (Var isEmpty) + <$> resultBodyM mergeexp + <*> renameBody closedBody + <*> pure (IfAttr [primBodyType t] IfNormal)) + where (mergepat, mergeexp) = unzip merge + mergeidents = map paramIdent mergepat + mergenames = map paramName mergepat + knownBnds = M.fromList $ zip mergenames mergeexp + +checkResults :: BinderOps lore => + [VName] + -> SubExp + -> Names + -> M.Map VName SubExp + -> [VName] -- ^ Lambda-bound + -> Body lore + -> [SubExp] + -> RuleM lore (Body lore) +checkResults pat size untouchable knownBnds params body accs = do + ((), bnds) <- collectStms $ + zipWithM_ checkResult (zip pat res) (zip accparams accs) + mkBodyM bnds $ map Var pat + + where bndMap = makeBindMap body + (accparams, _) = splitAt (length accs) params + res = bodyResult body + + nonFree = boundInBody body <> + S.fromList params <> + untouchable + + checkResult (p, Var v) (accparam, acc) + | Just (BasicOp (BinOp bop x y)) <- M.lookup v bndMap = do + -- One of x,y must be *this* accumulator, and the other must + -- be something that is free in the body. + let isThisAccum = (==Var accparam) + (this, el) <- liftMaybe $ + case ((asFreeSubExp x, isThisAccum y), + (asFreeSubExp y, isThisAccum x)) of + ((Just free, True), _) -> Just (acc, free) + (_, (Just free, True)) -> Just (acc, free) + _ -> Nothing + + case bop of + LogAnd -> + letBindNames_ [p] $ BasicOp $ BinOp LogAnd this el + Add t | Just properly_typed_size <- properIntSize t -> do + size' <- properly_typed_size + letBindNames_ [p] =<< + eBinOp (Add t) (eSubExp this) + (pure $ BasicOp $ BinOp (Mul t) el size') + FAdd t | Just properly_typed_size <- properFloatSize t -> do + size' <- properly_typed_size + letBindNames_ [p] =<< + eBinOp (FAdd t) (eSubExp this) + (pure $ BasicOp $ BinOp (FMul t) el size') + _ -> cannotSimplify -- Um... sorry. + + checkResult _ _ = cannotSimplify + + asFreeSubExp :: SubExp -> Maybe SubExp + asFreeSubExp (Var v) + | S.member v nonFree = M.lookup v knownBnds + asFreeSubExp se = Just se + + properIntSize Int32 = Just $ return size + properIntSize t = Just $ letSubExp "converted_size" $ + BasicOp $ ConvOp (SExt Int32 t) size + + properFloatSize t = + Just $ letSubExp "converted_size" $ + BasicOp $ ConvOp (SIToFP Int32 t) size + +determineKnownBindings :: VarLookup lore -> Lambda lore -> [SubExp] -> [VName] + -> M.Map VName SubExp +determineKnownBindings look lam accs arrs = + accBnds <> arrBnds + where (accparams, arrparams) = + splitAt (length accs) $ lambdaParams lam + accBnds = M.fromList $ + zip (map paramName accparams) accs + arrBnds = M.fromList $ mapMaybe isReplicate $ + zip (map paramName arrparams) arrs + + isReplicate (p, v) + | Just (BasicOp (Replicate _ ve), cs) <- look v, + cs == mempty = Just (p, ve) + isReplicate _ = Nothing + +makeBindMap :: Body lore -> M.Map VName (Exp lore) +makeBindMap = M.fromList . mapMaybe isSingletonStm . stmsToList . bodyStms + where isSingletonStm (Let pat _ e) = case patternNames pat of + [v] -> Just (v,e) + _ -> Nothing diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -0,0 +1,878 @@ +{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, TypeFamilies, FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE UndecidableInstances #-} +-- | +-- +-- Perform general rule-based simplification based on data dependency +-- information. This module will: +-- +-- * Perform common-subexpression elimination (CSE). +-- +-- * Hoist expressions out of loops (including lambdas) and +-- branches. This is done as aggressively as possible. +-- +-- * Apply simplification rules (see +-- "Futhark.Optimise.Simplification.Rules"). +-- +-- If you just want to run the simplifier as simply as possible, you +-- may prefer to use the "Futhark.Optimise.Simplify" module. +-- +module Futhark.Optimise.Simplify.Engine + ( -- * Monadic interface + SimpleM + , runSimpleM + , subSimpleM + , SimpleOps (..) + , SimplifyOp + , bindableSimpleOps + + , Env (envHoistBlockers, envRules) + , emptyEnv + , HoistBlockers(..) + , neverBlocks + , noExtraHoistBlockers + , BlockPred + , orIf + , hasFree + , isConsumed + , isFalse + , isOp + , isNotSafe + , asksEngineEnv + , changed + , askVtable + , localVtable + + -- * Building blocks + , SimplifiableLore + , Simplifiable (..) + , simplifyStms + , simplifyFun + , simplifyLambda + , simplifyLambdaSeq + , simplifyLambdaNoHoisting + , simplifyParam + , bindLParams + , bindChunkLParams + , bindLoopVar + , enterLoop + , simplifyBody + , SimplifiedBody + + , blockIf + , constructBody + , protectIf + + , module Futhark.Optimise.Simplify.Lore + ) where + +import Control.Monad.Writer +import Control.Monad.RWS.Strict +import Data.Either +import Data.List +import Data.Maybe +import qualified Data.Set as S + +import Futhark.Representation.AST +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Optimise.Simplify.Rule +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Analysis.Usage +import Futhark.Construct +import Futhark.Optimise.Simplify.Lore +import Futhark.Util (splitFromEnd) + +data HoistBlockers lore = HoistBlockers + { blockHoistPar :: BlockPred (Wise lore) + -- ^ Blocker for hoisting out of parallel loops. + , blockHoistSeq :: BlockPred (Wise lore) + -- ^ Blocker for hoisting out of sequential loops. + , blockHoistBranch :: BlockPred (Wise lore) + -- ^ Blocker for hoisting out of branches. + , getArraySizes :: Stm (Wise lore) -> Names + -- ^ gets the sizes of arrays from a binding. + , isAllocation :: Stm (Wise lore) -> Bool + } + +noExtraHoistBlockers :: HoistBlockers lore +noExtraHoistBlockers = HoistBlockers neverBlocks neverBlocks neverBlocks (const S.empty) (const False) + +data Env lore = Env { envRules :: RuleBook (Wise lore) + , envHoistBlockers :: HoistBlockers lore + , envVtable :: ST.SymbolTable (Wise lore) + } + +emptyEnv :: RuleBook (Wise lore) -> HoistBlockers lore -> Env lore +emptyEnv rules blockers = + Env { envRules = rules + , envHoistBlockers = blockers + , envVtable = mempty + } + +data SimpleOps lore = + SimpleOps { mkExpAttrS :: ST.SymbolTable (Wise lore) + -> Pattern (Wise lore) -> Exp (Wise lore) + -> SimpleM lore (ExpAttr (Wise lore)) + , mkBodyS :: ST.SymbolTable (Wise lore) + -> Stms (Wise lore) -> Result + -> SimpleM lore (Body (Wise lore)) + , mkLetNamesS :: ST.SymbolTable (Wise lore) + -> [VName] -> Exp (Wise lore) + -> SimpleM lore (Stm (Wise lore), Stms (Wise lore)) + , simplifyOpS :: SimplifyOp lore + } + +type SimplifyOp lore = Op lore -> SimpleM lore (OpWithWisdom (Op lore), Stms (Wise lore)) + +bindableSimpleOps :: (SimplifiableLore lore, Bindable lore) => + SimplifyOp lore -> SimpleOps lore +bindableSimpleOps = SimpleOps mkExpAttrS' mkBodyS' mkLetNamesS' + where mkExpAttrS' _ pat e = return $ mkExpAttr pat e + mkBodyS' _ bnds res = return $ mkBody bnds res + mkLetNamesS' _ name e = (,) <$> mkLetNames name e <*> pure mempty + +newtype SimpleM lore a = + SimpleM (RWS (SimpleOps lore, Env lore) Certificates (VNameSource, Bool) a) + deriving (Applicative, Functor, Monad, + MonadReader (SimpleOps lore, Env lore), + MonadState (VNameSource, Bool), + MonadWriter Certificates) + +instance MonadFreshNames (SimpleM lore) where + putNameSource src = modify $ \(_, b) -> (src, b) + getNameSource = gets fst + +instance SimplifiableLore lore => HasScope (Wise lore) (SimpleM lore) where + askScope = ST.toScope <$> askVtable + lookupType name = do + vtable <- askVtable + case ST.lookupType name vtable of + Just t -> return t + Nothing -> fail $ + "SimpleM.lookupType: cannot find variable " ++ + pretty name ++ " in symbol table." + +instance SimplifiableLore lore => + LocalScope (Wise lore) (SimpleM lore) where + localScope types = localVtable (<>ST.fromScope types) + +runSimpleM :: SimpleM lore a + -> SimpleOps lore + -> Env lore + -> VNameSource + -> ((a, Bool), VNameSource) +runSimpleM (SimpleM m) simpl env src = + let (x, (src', b), _) = runRWS m (simpl, env) (src, False) + in ((x, b), src') + +subSimpleM :: (MonadFreshNames m, + SameScope outerlore lore, + ExpAttr outerlore ~ ExpAttr lore, + BodyAttr outerlore ~ BodyAttr lore, + RetType outerlore ~ RetType lore, + BranchType outerlore ~ BranchType lore) => + SimpleOps lore + -> Env lore + -> ST.SymbolTable (Wise outerlore) + -> SimpleM lore a + -> m (a, Bool) +subSimpleM simpl env outer_vtable m = do + let inner_vtable = ST.castSymbolTable outer_vtable + src <- getNameSource + let SimpleM m' = localVtable (<>inner_vtable) m + (x, (src', b), _) = runRWS m' (simpl, env) (src, False) + putNameSource src' + return (x, b) + +askEngineEnv :: SimpleM lore (Env lore) +askEngineEnv = snd <$> ask + +asksEngineEnv :: (Env lore -> a) -> SimpleM lore a +asksEngineEnv f = f <$> askEngineEnv + +askVtable :: SimpleM lore (ST.SymbolTable (Wise lore)) +askVtable = asksEngineEnv envVtable + +localVtable :: (ST.SymbolTable (Wise lore) -> ST.SymbolTable (Wise lore)) + -> SimpleM lore a -> SimpleM lore a +localVtable f = local $ \(ops, env) -> (ops, env { envVtable = f $ envVtable env }) + +collectCerts :: SimpleM lore a -> SimpleM lore (a, Certificates) +collectCerts m = pass $ do (x, cs) <- listen m + return ((x, cs), const mempty) + +-- | Mark that we have changed something and it would be a good idea +-- to re-run the simplifier. +changed :: SimpleM lore () +changed = modify $ \(src, _) -> (src, True) + +usedCerts :: Certificates -> SimpleM lore () +usedCerts = tell + +enterLoop :: SimpleM lore a -> SimpleM lore a +enterLoop = localVtable ST.deepen + +bindFParams :: SimplifiableLore lore => + [FParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a +bindFParams params = + localVtable $ ST.insertFParams params + +bindLParams :: SimplifiableLore lore => + [LParam (Wise lore)] -> SimpleM lore a -> SimpleM lore a +bindLParams params = + localVtable $ \vtable -> + foldr ST.insertLParam vtable params + +bindArrayLParams :: SimplifiableLore lore => + [(LParam (Wise lore),Maybe VName)] -> SimpleM lore a -> SimpleM lore a +bindArrayLParams params = + localVtable $ \vtable -> + foldr (uncurry ST.insertArrayLParam) vtable params + +bindChunkLParams :: SimplifiableLore lore => + VName -> [(LParam (Wise lore),VName)] -> SimpleM lore a -> SimpleM lore a +bindChunkLParams offset params = + localVtable $ \vtable -> + foldr (uncurry $ ST.insertChunkLParam offset) vtable params + +bindLoopVar :: SimplifiableLore lore => + VName -> IntType -> SubExp -> SimpleM lore a -> SimpleM lore a +bindLoopVar var it bound = + localVtable $ clampUpper . clampVar + where clampVar = ST.insertLoopVar var it bound + -- If we enter the loop, then 'bound' is at least one. + clampUpper = case bound of Var v -> ST.isAtLeast v 1 + _ -> id + +-- | We are willing to hoist potentially unsafe statements out of +-- branches, but they most be protected by adding a branch on top of +-- them. (This means such hoisting is not worth it unless they are in +-- turn hoisted out of a loop somewhere.) +protectIfHoisted :: SimplifiableLore lore => + SubExp -- ^ Branch condition. + -> Bool -- ^ Which side of the branch are we + -- protecting here? + -> SimpleM lore (a, Stms (Wise lore)) + -> SimpleM lore (a, Stms (Wise lore)) +protectIfHoisted cond side m = do + (x, stms) <- m + runBinder $ do + if any (not . safeExp . stmExp) stms + then do cond' <- if side then return cond + else letSubExp "cond_neg" $ BasicOp $ UnOp Not cond + mapM_ (protectIf unsafeOrCostly cond') stms + else addStms stms + return x + where unsafeOrCostly e = not (safeExp e) || not (cheapExp e) + +-- | We are willing to hoist potentially unsafe statements out of +-- 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)) +protectLoopHoisted ctx val form m = do + (x, stms) <- m + runBinder $ do + if any (not . safeExp . stmExp) stms + then do is_nonempty <- checkIfNonEmpty + mapM_ (protectIf (not . safeExp) is_nonempty) stms + else addStms stms + return x + where checkIfNonEmpty = + case form of + WhileLoop cond + | Just (_, cond_init) <- + find ((==cond) . paramName . fst) $ ctx ++ val -> + return cond_init + | otherwise -> return $ constant True -- infinite loop + ForLoop _ it bound _ -> + letSubExp "loop_nonempty" $ + BasicOp $ CmpOp (CmpSlt it) (intConst it 0) bound + +protectIf :: MonadBinder m => (Exp (Lore m) -> Bool) -> SubExp -> Stm (Lore m) -> m () +protectIf _ taken (Let pat (StmAux cs _) + (If cond taken_body untaken_body (IfAttr if_ts IfFallback))) = do + cond' <- letSubExp "protect_cond_conj" $ BasicOp $ BinOp LogAnd taken cond + certifying cs $ + letBind_ pat $ If cond' taken_body untaken_body $ + IfAttr if_ts IfFallback +protectIf f taken (Let pat (StmAux cs _) e) + | f e = do + taken_body <- eBody [pure e] + untaken_body <- eBody $ map (emptyOfType $ patternContextNames pat) + (patternValueTypes pat) + if_ts <- expTypesFromPattern pat + certifying cs $ + letBind_ pat $ If taken taken_body untaken_body $ + IfAttr if_ts IfFallback +protectIf _ _ stm = + addStm stm + +emptyOfType :: MonadBinder m => [VName] -> Type -> m (Exp (Lore m)) +emptyOfType _ Mem{} = + fail "emptyOfType: Cannot hoist non-existential memory." +emptyOfType _ (Prim pt) = + return $ BasicOp $ SubExp $ Constant $ blankPrimValue pt +emptyOfType ctx_names (Array pt shape _) = do + let dims = map zeroIfContext $ shapeDims shape + return $ BasicOp $ Scratch pt dims + where zeroIfContext (Var v) | v `elem` ctx_names = intConst Int32 0 + zeroIfContext se = se + +-- | Statements that are not worth hoisting out of loops, because they +-- are unsafe, and added safety (by 'protectLoopHoisted') may inhibit +-- further optimisation.. +notWorthHoisting :: Attributes lore => BlockPred lore +notWorthHoisting _ (Let pat _ e) = + not (safeExp e) && any (>0) (map arrayRank $ patternTypes pat) + +hoistStms :: SimplifiableLore lore => + RuleBook (Wise lore) -> BlockPred (Wise lore) + -> ST.SymbolTable (Wise lore) -> UT.UsageTable + -> Stms (Wise lore) + -> SimpleM lore (Stms (Wise lore), + Stms (Wise lore)) +hoistStms rules block vtable uses orig_stms = do + (blocked, hoisted) <- simplifyStmsBottomUp vtable uses orig_stms + unless (null hoisted) changed + return (stmsFromList blocked, stmsFromList hoisted) + where simplifyStmsBottomUp vtable' uses' stms = do + (_, stms') <- simplifyStmsBottomUp' vtable' uses' stms + -- We need to do a final pass to ensure that nothing is + -- hoisted past something that it depends on. + let (blocked, hoisted) = partitionEithers $ blockUnhoistedDeps stms' + return (blocked, hoisted) + + simplifyStmsBottomUp' vtable' uses' stms = + foldM hoistable (uses',[]) $ reverse $ zip (stmsToList stms) vtables + where vtables = scanl (flip ST.insertStm) vtable' $ stmsToList stms + + hoistable (uses',stms) (stm, vtable') + | not $ any (`UT.isUsedDirectly` uses') $ provides stm = -- Dead statement. + return (uses', stms) + | otherwise = do + res <- localVtable (const vtable') $ + bottomUpSimplifyStm rules (vtable', uses') stm + case res of + Nothing -- Nothing to optimise - see if hoistable. + | block uses' stm -> + return (expandUsage vtable' uses' stm `UT.without` provides stm, + Left stm : stms) + | otherwise -> + return (expandUsage vtable' uses' stm, Right stm : stms) + Just optimstms -> do + changed + (uses'',stms') <- simplifyStmsBottomUp' vtable' uses' optimstms + return (uses'', stms'++stms) + +blockUnhoistedDeps :: Attributes lore => + [Either (Stm lore) (Stm lore)] + -> [Either (Stm lore) (Stm lore)] +blockUnhoistedDeps = snd . mapAccumL block S.empty + where block blocked (Left need) = + (blocked <> S.fromList (provides need), Left need) + block blocked (Right need) + | blocked `intersects` requires need = + (blocked <> S.fromList (provides need), Left need) + | otherwise = + (blocked, Right need) + +provides :: Stm lore -> [VName] +provides = patternNames . stmPattern + +requires :: Attributes lore => Stm lore -> Names +requires = freeInStm + +expandUsage :: (Attributes lore, Aliased lore, UsageInOp (Op lore)) => + ST.SymbolTable lore -> UT.UsageTable -> Stm lore -> UT.UsageTable +expandUsage vtable utable bnd = + UT.expand (`ST.lookupAliases` vtable) (usageInStm bnd <> usageThroughAliases) <> + utable + where pat = stmPattern bnd + usageThroughAliases = + mconcat $ mapMaybe usageThroughBindeeAliases $ + zip (patternNames pat) (patternAliases pat) + usageThroughBindeeAliases (name, aliases) = do + uses <- UT.lookup name utable + return $ mconcat $ map (`UT.usage` uses) $ S.toList aliases + +intersects :: Ord a => S.Set a -> S.Set a -> Bool +intersects a b = not $ S.null $ a `S.intersection` b + +type BlockPred lore = UT.UsageTable -> Stm lore -> Bool + +neverBlocks :: BlockPred lore +neverBlocks _ _ = False + +isFalse :: Bool -> BlockPred lore +isFalse b _ _ = not b + +orIf :: BlockPred lore -> BlockPred lore -> BlockPred lore +orIf p1 p2 body need = p1 body need || p2 body need + +andAlso :: BlockPred lore -> BlockPred lore -> BlockPred lore +andAlso p1 p2 body need = p1 body need && p2 body need + +isConsumed :: BlockPred lore +isConsumed utable = any (`UT.isConsumed` utable) . patternNames . stmPattern + +isOp :: BlockPred lore +isOp _ (Let _ _ Op{}) = True +isOp _ _ = False + +constructBody :: SimplifiableLore lore => Stms (Wise lore) -> Result + -> SimpleM lore (Body (Wise lore)) +constructBody stms res = + fmap fst $ runBinder $ insertStmsM $ do addStms stms + resultBodyM res + +type SimplifiedBody lore a = ((a, UT.UsageTable), Stms (Wise lore)) + +blockIf :: SimplifiableLore lore => + BlockPred (Wise lore) + -> SimpleM lore (SimplifiedBody lore a) + -> SimpleM lore ((Stms (Wise lore), a), Stms (Wise lore)) +blockIf block m = do + ((x, usages), stms) <- m + vtable <- askVtable + rules <- asksEngineEnv envRules + (blocked, hoisted) <- hoistStms rules block vtable usages stms + return ((blocked, x), hoisted) + +insertAllStms :: SimplifiableLore lore => + SimpleM lore (SimplifiedBody lore Result) + -> SimpleM lore (Body (Wise lore)) +insertAllStms = uncurry constructBody . fst <=< blockIf (isFalse False) + +hasFree :: Attributes lore => Names -> BlockPred lore +hasFree ks _ need = ks `intersects` requires need + +isNotSafe :: Attributes lore => BlockPred lore +isNotSafe _ = not . safeExp . stmExp + +isInPlaceBound :: BlockPred m +isInPlaceBound _ = isUpdate . stmExp + where isUpdate (BasicOp Update{}) = True + isUpdate _ = False + +isNotCheap :: Attributes lore => BlockPred lore +isNotCheap _ = not . cheapStm + +cheapStm :: Attributes lore => Stm lore -> Bool +cheapStm = cheapExp . stmExp + +cheapExp :: Attributes lore => Exp lore -> Bool +cheapExp (BasicOp BinOp{}) = True +cheapExp (BasicOp SubExp{}) = True +cheapExp (BasicOp UnOp{}) = True +cheapExp (BasicOp CmpOp{}) = True +cheapExp (BasicOp ConvOp{}) = True +cheapExp (BasicOp Copy{}) = False +cheapExp DoLoop{} = False +cheapExp (If _ tbranch fbranch _) = all cheapStm (bodyStms tbranch) && + all cheapStm (bodyStms fbranch) +cheapExp (Op op) = cheapOp op +cheapExp _ = True -- Used to be False, but + -- let's try it out. + +stmIs :: (Stm lore -> Bool) -> BlockPred lore +stmIs f _ = f + +loopInvariantStm :: Attributes lore => ST.SymbolTable lore -> Stm lore -> Bool +loopInvariantStm vtable = + all (`S.member` ST.availableAtClosestLoop vtable) . freeInStm + +hoistCommon :: SimplifiableLore lore => + SubExp -> IfSort + -> SimplifiedBody lore Result + -> SimplifiedBody lore Result + -> SimpleM lore (Body (Wise lore), Body (Wise lore), Stms (Wise lore)) +hoistCommon cond ifsort ((res1, usages1), stms1) ((res2, usages2), stms2) = do + is_alloc_fun <- asksEngineEnv $ isAllocation . envHoistBlockers + getArrSz_fun <- asksEngineEnv $ getArraySizes . envHoistBlockers + branch_blocker <- asksEngineEnv $ blockHoistBranch . envHoistBlockers + vtable <- askVtable + let -- We are unwilling to hoist things that are unsafe or costly, + -- *except* if they are invariant to the most enclosing loop, + -- because in that case they will also be hoisted past that + -- loop. + -- + -- "isNotHoistableBnd hoistbl_nms" ensures that only the + -- (transitive closure) of the bindings used for allocations, + -- shape computations, and expensive loop-invariant operations + -- are if-hoistable. + cond_loop_invariant = + all (`S.member` ST.availableAtClosestLoop vtable) $ freeIn cond + desirableToHoist stm = + is_alloc_fun stm || + (ST.loopDepth vtable > 0 && + cond_loop_invariant && + ifsort /= IfFallback && + loopInvariantStm vtable stm) + hoistbl_nms = filterBnds desirableToHoist getArrSz_fun $ + stmsToList $ stms1<>stms2 + block = branch_blocker `orIf` + ((isNotSafe `orIf` isNotCheap) `andAlso` stmIs (not . desirableToHoist)) + `orIf` isInPlaceBound `orIf` isNotHoistableBnd hoistbl_nms + rules <- asksEngineEnv envRules + (body1_bnds', safe1) <- protectIfHoisted cond True $ + hoistStms rules block vtable usages1 stms1 + (body2_bnds', safe2) <- protectIfHoisted cond False $ + hoistStms rules block vtable usages2 stms2 + let hoistable = safe1 <> safe2 + body1' <- constructBody body1_bnds' res1 + body2' <- constructBody body2_bnds' res2 + return (body1', body2', hoistable) + where filterBnds interesting getArrSz_fn all_bnds = + let sz_nms = mconcat $ map getArrSz_fn all_bnds + sz_needs = transClosSizes all_bnds sz_nms [] + alloc_bnds = filter interesting all_bnds + sel_nms = S.fromList $ + concatMap (patternNames . stmPattern) + (sz_needs ++ alloc_bnds) + in sel_nms + transClosSizes all_bnds scal_nms hoist_bnds = + let new_bnds = filter (hasPatName scal_nms) all_bnds + new_nms = mconcat $ map (freeInExp . stmExp) new_bnds + in if null new_bnds + then hoist_bnds + else transClosSizes all_bnds new_nms (new_bnds ++ hoist_bnds) + hasPatName nms bnd = intersects nms $ S.fromList $ + patternNames $ stmPattern bnd + isNotHoistableBnd _ _ (Let _ _ (BasicOp ArrayLit{})) = False + isNotHoistableBnd nms _ stm = not (hasPatName nms stm) + +-- | Simplify a single 'Body'. The @[Diet]@ only covers the value +-- elements, because the context cannot be consumed. +simplifyBody :: SimplifiableLore lore => + [Diet] -> Body lore -> SimpleM lore (SimplifiedBody lore Result) +simplifyBody ds (Body _ bnds res) = + simplifyStms bnds $ do res' <- simplifyResult ds res + return (res', mempty) + +-- | Simplify a single 'Result'. The @[Diet]@ only covers the value +-- elements, because the context cannot be consumed. +simplifyResult :: SimplifiableLore lore => + [Diet] -> Result -> SimpleM lore (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 + -- place to put the certificates when copy-propagating a certified + -- statement. However, for results in the *context*, it is OK to + -- just throw away the certificates, because for the program to be + -- type-correct, those statements must anyway be used (or + -- copy-propagated into) the statements producing the value result. + (ctx_res', _ctx_res_cs) <- collectCerts $ mapM simplify ctx_res + val_res' <- mapM simplify' val_res + + let consumption = consumeResult $ zip ds val_res' + res' = ctx_res' <> val_res' + return (res', UT.usages (freeIn res') <> consumption) + + where simplify' (Var name) = do + bnd <- ST.lookupSubExp name <$> askVtable + case bnd of + Just (Constant v, cs) + | cs == mempty -> return $ Constant v + Just (Var id', cs) + | cs == mempty -> return $ Var id' + _ -> return $ Var name + simplify' (Constant v) = + return $ Constant v + +isDoLoopResult :: Result -> UT.UsageTable +isDoLoopResult = mconcat . map checkForVar + where checkForVar (Var ident) = UT.inResultUsage ident + checkForVar _ = mempty + +simplifyStms :: SimplifiableLore lore => + Stms lore -> SimpleM lore (a, Stms (Wise lore)) + -> SimpleM lore (a, Stms (Wise lore)) +simplifyStms stms m = + case stmsHead stms of + Nothing -> inspectStms mempty m + Just (Let pat (StmAux stm_cs attr) e, stms') -> do + stm_cs' <- simplify stm_cs + ((e', e_stms), e_cs) <- collectCerts $ simplifyExp e + (pat', pat_cs) <- collectCerts $ simplifyPattern pat + let cs = stm_cs'<>e_cs<>pat_cs + inspectStms e_stms $ + inspectStm (mkWiseLetStm pat' (StmAux cs attr) e') $ + simplifyStms stms' m + +inspectStm :: SimplifiableLore lore => + Stm (Wise lore) -> SimpleM lore (a, Stms (Wise lore)) + -> SimpleM lore (a, Stms (Wise lore)) +inspectStm = inspectStms . oneStm + +inspectStms :: SimplifiableLore lore => + Stms (Wise lore) + -> SimpleM lore (a, Stms (Wise lore)) + -> SimpleM lore (a, Stms (Wise lore)) +inspectStms stms m = + case stmsHead stms of + Nothing -> m + Just (stm, stms') -> do + vtable <- askVtable + rules <- asksEngineEnv envRules + simplified <- topDownSimplifyStm rules vtable stm + case simplified of + Just newbnds -> changed >> inspectStms (newbnds <> stms') m + Nothing -> do (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 = do f <- asks $ simplifyOpS . fst + f op + +simplifyExp :: SimplifiableLore lore => + Exp lore -> SimpleM lore (Exp (Wise lore), Stms (Wise lore)) + +simplifyExp (If cond tbranch fbranch (IfAttr 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 + -- across branches. + cond' <- simplify cond + ts' <- mapM simplify ts + -- FIXME: we have to be conservative about the diet here, because we + -- lack proper ifnormation. Something is wrong with the order in + -- which the simplifier does things - it should be purely bottom-up + -- (or else, If expressions should indicate explicitly the diet of + -- their return types). + let ds = map (const Consume) ts + tbranch' <- localVtable (ST.updateBounds True cond) $ simplifyBody ds tbranch + fbranch' <- localVtable (ST.updateBounds False cond) $ simplifyBody ds fbranch + (tbranch'',fbranch'', hoisted) <- hoistCommon cond' ifsort tbranch' fbranch' + return (If cond' tbranch'' fbranch'' $ IfAttr ts' ifsort, hoisted) + +simplifyExp (DoLoop ctx val form loopbody) = do + let (ctxparams, ctxinit) = unzip ctx + (valparams, valinit) = unzip val + ctxparams' <- mapM (simplifyParam simplify) ctxparams + ctxinit' <- mapM simplify ctxinit + valparams' <- mapM (simplifyParam simplify) valparams + valinit' <- mapM simplify valinit + let ctx' = zip ctxparams' ctxinit' + val' = zip valparams' valinit' + diets = map (diet . paramDeclType) valparams' + (form', boundnames, wrapbody) <- case form of + ForLoop loopvar it boundexp loopvars -> do + boundexp' <- simplify boundexp + let (loop_params, loop_arrs) = unzip loopvars + loop_params' <- mapM (simplifyParam simplify) loop_params + loop_arrs' <- mapM simplify loop_arrs + let form' = ForLoop loopvar it boundexp' (zip loop_params' loop_arrs') + return (form', + S.fromList (loopvar : map paramName loop_params') <> fparamnames, + bindLoopVar loopvar it boundexp' . + protectLoopHoisted ctx' val' form' . + bindArrayLParams (zip loop_params' (map Just loop_arrs'))) + WhileLoop cond -> do + cond' <- simplify cond + return (WhileLoop cond', + fparamnames, + protectLoopHoisted ctx' val' (WhileLoop cond')) + seq_blocker <- asksEngineEnv $ blockHoistSeq . envHoistBlockers + ((loopstms, loopres), hoisted) <- + enterLoop $ + bindFParams (ctxparams'++valparams') $ wrapbody $ + blockIf + (hasFree boundnames `orIf` isConsumed + `orIf` seq_blocker `orIf` notWorthHoisting) $ do + ((res, uses), stms) <- simplifyBody diets loopbody + return ((res, uses <> isDoLoopResult res), stms) + loopbody' <- constructBody loopstms loopres + return (DoLoop ctx' val' form' loopbody', hoisted) + where fparamnames = S.fromList (map (paramName . fst) $ ctx++val) + +simplifyExp (Op op) = do (op', stms) <- simplifyOp op + return (Op op', stms) + +-- Special case for simplification of commutative BinOps where we +-- arrange the operands in sorted order. This can make expressions +-- more identical, which helps CSE. +simplifyExp (BasicOp (BinOp op x y)) + | commutativeBinOp op = do + x' <- simplify x + y' <- simplify y + return (BasicOp $ BinOp op (min x' y') (max x' y'), mempty) + +simplifyExp e = do e' <- simplifyExpBase e + return (e', mempty) + +simplifyExpBase :: SimplifiableLore lore => + Exp lore -> SimpleM lore (Exp (Wise lore)) +simplifyExpBase = mapExpM hoist + where hoist = Mapper { + -- Bodies are handled explicitly because we need to + -- provide their result diet. + mapOnBody = fail "Unhandled body in simplification engine." + , mapOnSubExp = simplify + -- Lambdas are handled explicitly because we need to + -- bind their parameters. + , mapOnVName = simplify + , mapOnCertificates = simplify + , mapOnRetType = simplify + , mapOnBranchType = simplify + , mapOnFParam = + fail "Unhandled FParam in simplification engine." + , mapOnLParam = + fail "Unhandled LParam in simplification engine." + , mapOnOp = + fail "Unhandled Op in simplification engine." + } + +type SimplifiableLore lore = (Attributes lore, + Simplifiable (LetAttr lore), + Simplifiable (FParamAttr lore), + Simplifiable (LParamAttr lore), + Simplifiable (RetType lore), + Simplifiable (BranchType lore), + CanBeWise (Op lore), + ST.IndexOp (OpWithWisdom (Op lore)), + BinderOps (Wise lore), + IsOp (Op lore)) + +class Simplifiable e where + simplify :: SimplifiableLore lore => e -> SimpleM lore e + +instance (Simplifiable a, Simplifiable b) => Simplifiable (a, b) where + simplify (x,y) = (,) <$> simplify x <*> simplify y + +instance (Simplifiable a, Simplifiable b, Simplifiable c) => Simplifiable (a, b, c) where + simplify (x,y,z) = (,,) <$> simplify x <*> simplify y <*> simplify z + +-- Convenient for Scatter. +instance Simplifiable Int where + simplify = pure + +instance Simplifiable a => Simplifiable (Maybe a) where + simplify Nothing = return Nothing + simplify (Just x) = Just <$> simplify x + +instance Simplifiable a => Simplifiable [a] where + simplify = mapM simplify + +instance Simplifiable SubExp where + simplify (Var name) = do + bnd <- ST.lookupSubExp name <$> askVtable + case bnd of + Just (Constant v, cs) -> do changed + usedCerts cs + return $ Constant v + Just (Var id', cs) -> do changed + usedCerts cs + return $ Var id' + _ -> return $ Var name + simplify (Constant v) = + return $ Constant v + +simplifyPattern :: (SimplifiableLore lore, Simplifiable attr) => + PatternT attr + -> SimpleM lore (PatternT attr) +simplifyPattern pat = + Pattern <$> + mapM inspect (patternContextElements pat) <*> + mapM inspect (patternValueElements pat) + where inspect (PatElem name lore) = PatElem name <$> simplify lore + +simplifyParam :: (attr -> SimpleM lore attr) -> ParamT attr -> SimpleM lore (ParamT attr) +simplifyParam simplifyAttribute (Param name attr) = + Param name <$> simplifyAttribute attr + +instance Simplifiable VName where + simplify v = do + se <- ST.lookupSubExp v <$> askVtable + case se of + Just (Var v', cs) -> do changed + usedCerts cs + return v' + _ -> return v + +instance Simplifiable d => Simplifiable (ShapeBase d) where + simplify = fmap Shape . simplify . shapeDims + +instance Simplifiable ExtSize where + simplify (Free se) = Free <$> simplify se + simplify (Ext x) = return $ Ext x + +instance Simplifiable shape => Simplifiable (TypeBase shape u) where + simplify (Array et shape u) = do + shape' <- simplify shape + return $ Array et shape' u + simplify (Mem size space) = + Mem <$> simplify size <*> pure space + simplify (Prim bt) = + return $ Prim bt + +instance Simplifiable d => Simplifiable (DimIndex d) where + simplify (DimFix i) = DimFix <$> simplify i + simplify (DimSlice i n s) = DimSlice <$> simplify i <*> simplify n <*> simplify s + +simplifyLambda :: SimplifiableLore lore => + Lambda lore + -> [Maybe VName] + -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) +simplifyLambda lam arrs = do + par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers + simplifyLambdaMaybeHoist par_blocker lam arrs + +simplifyLambdaSeq :: SimplifiableLore lore => + Lambda lore + -> [Maybe VName] + -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) +simplifyLambdaSeq = simplifyLambdaMaybeHoist neverBlocks + +simplifyLambdaNoHoisting :: SimplifiableLore lore => + Lambda lore + -> [Maybe VName] + -> SimpleM lore (Lambda (Wise lore)) +simplifyLambdaNoHoisting lam arr = + fst <$> simplifyLambdaMaybeHoist (isFalse False) lam arr + +simplifyLambdaMaybeHoist :: SimplifiableLore lore => + BlockPred (Wise lore) -> Lambda lore + -> [Maybe VName] + -> SimpleM lore (Lambda (Wise lore), Stms (Wise lore)) +simplifyLambdaMaybeHoist blocked lam@(Lambda params body rettype) arrs = do + params' <- mapM (simplifyParam simplify) params + let (nonarrayparams, arrayparams) = + splitAt (length params' - length arrs) params' + paramnames = S.fromList $ boundByLambda lam + ((lamstms, lamres), hoisted) <- + enterLoop $ + bindLParams nonarrayparams $ + bindArrayLParams (zip arrayparams arrs) $ + blockIf (blocked `orIf` hasFree paramnames `orIf` isConsumed) $ + simplifyBody (map (const Observe) rettype) body + body' <- constructBody lamstms lamres + rettype' <- simplify rettype + return (Lambda params' body' rettype', hoisted) + +consumeResult :: [(Diet, SubExp)] -> UT.UsageTable +consumeResult = mconcat . map inspect + where inspect (Consume, se) = + mconcat $ map UT.consumedUsage $ S.toList $ subExpAliases se + inspect _ = mempty + +instance Simplifiable Certificates where + simplify (Certificates ocs) = Certificates . nub . concat <$> mapM check ocs + where check idd = do + vv <- ST.lookupSubExp idd <$> askVtable + case vv of + Just (Constant Checked, Certificates cs) -> return cs + Just (Var idd', _) -> return [idd'] + _ -> return [idd] + +simplifyFun :: SimplifiableLore lore => FunDef lore -> SimpleM lore (FunDef (Wise lore)) +simplifyFun (FunDef entry fname rettype params body) = do + rettype' <- simplify rettype + let ds = map diet (retTypeValues rettype') + body' <- bindFParams params $ insertAllStms $ simplifyBody ds body + return $ FunDef entry fname rettype' params body' diff --git a/src/Futhark/Optimise/Simplify/Lore.hs b/src/Futhark/Optimise/Simplify/Lore.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Lore.hs @@ -0,0 +1,269 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +-- | Definition of the lore used by the simplification engine. +module Futhark.Optimise.Simplify.Lore + ( + Wise + , VarWisdom (..) + , ExpWisdom + , removeStmWisdom + , removeLambdaWisdom + , removeProgWisdom + , removeFunDefWisdom + , removeExpWisdom + , removePatternWisdom + , removePatElemWisdom + , removeBodyWisdom + , removeScopeWisdom + , addScopeWisdom + , addWisdomToPattern + , mkWiseBody + , mkWiseLetStm + , mkWiseExpAttr + + , CanBeWise (..) + ) + where + +import Control.Monad.Identity +import Control.Monad.Reader +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M + +import Futhark.Representation.AST +import Futhark.Representation.AST.Attributes.Ranges +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.Aliases + (unNames, Names' (..), VarAliases, ConsumedInExp) +import qualified Futhark.Representation.Aliases as Aliases +import qualified Futhark.Representation.Ranges as Ranges +import Futhark.Binder +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import Futhark.Analysis.Rephrase +import Futhark.Analysis.Usage (UsageInOp) + +data Wise lore + +-- | The wisdom of the let-bound variable. +data VarWisdom = VarWisdom { varWisdomAliases :: VarAliases + , varWisdomRange :: Range + } + deriving (Eq, Ord, Show) + +instance Rename VarWisdom where + rename = substituteRename + +instance Substitute VarWisdom where + substituteNames substs (VarWisdom als range) = + VarWisdom (substituteNames substs als) (substituteNames substs range) + +instance FreeIn VarWisdom where + freeIn (VarWisdom als range) = freeIn als <> freeIn range + +-- | Wisdom about an expression. +data ExpWisdom = ExpWisdom { _expWisdomConsumed :: ConsumedInExp + , expWisdomFree :: Names' + } + deriving (Eq, Ord, Show) + +instance FreeIn ExpWisdom where + freeIn = mempty + +instance FreeAttr ExpWisdom where + precomputed = const . unNames . 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 + , bodyWisdomRanges :: [Range] + , bodyWisdomFree :: Names' + } + deriving (Eq, Ord, Show) + +instance Rename BodyWisdom where + rename = substituteRename + +instance Substitute BodyWisdom where + substituteNames substs (BodyWisdom als cons rs free) = + BodyWisdom + (substituteNames substs als) + (substituteNames substs cons) + (substituteNames substs rs) + (substituteNames substs free) + +instance FreeIn BodyWisdom where + freeIn (BodyWisdom als cons rs free) = + freeIn als <> freeIn cons <> freeIn rs <> freeIn free + +instance FreeAttr BodyWisdom where + precomputed = const . unNames . bodyWisdomFree + +instance (Annotations lore, + CanBeWise (Op lore)) => Annotations (Wise lore) where + type LetAttr (Wise lore) = (VarWisdom, LetAttr lore) + type ExpAttr (Wise lore) = (ExpWisdom, ExpAttr lore) + type BodyAttr (Wise lore) = (BodyWisdom, BodyAttr lore) + type FParamAttr (Wise lore) = FParamAttr lore + type LParamAttr (Wise lore) = LParamAttr 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 (Attributes lore, CanBeWise (Op lore)) => Attributes (Wise lore) where + expTypesFromPattern = + withoutWisdom . expTypesFromPattern . removePatternWisdom + +instance PrettyAnnot (PatElemT attr) => PrettyAnnot (PatElemT (VarWisdom, attr)) where + ppAnnot = ppAnnot . fmap snd + +instance (PrettyLore lore, CanBeWise (Op lore)) => PrettyLore (Wise lore) where + ppExpLore (_, attr) = ppExpLore attr . removeExpWisdom + +instance AliasesOf (VarWisdom, attr) where + aliasesOf = unNames . varWisdomAliases . fst + +instance RangeOf (VarWisdom, attr) where + rangeOf = varWisdomRange . fst + +instance RangesOf (BodyWisdom, attr) where + rangesOf = bodyWisdomRanges . fst + +instance (Attributes lore, CanBeWise (Op lore)) => Aliased (Wise lore) where + bodyAliases = map unNames . bodyWisdomAliases . fst . bodyAttr + consumedInBody = unNames . bodyWisdomConsumed . fst . bodyAttr + +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 (LetInfo (_, attr)) = LetInfo attr + unAlias (FParamInfo attr) = FParamInfo attr + unAlias (LParamInfo attr) = LParamInfo attr + unAlias (IndexInfo it) = IndexInfo it + +addScopeWisdom :: Scope lore -> Scope (Wise lore) +addScopeWisdom = M.map alias + where alias (LetInfo attr) = LetInfo (VarWisdom mempty unknownRange, attr) + alias (FParamInfo attr) = FParamInfo attr + alias (LParamInfo attr) = LParamInfo attr + alias (IndexInfo it) = IndexInfo it + +removeProgWisdom :: CanBeWise (Op lore) => Prog (Wise lore) -> Prog lore +removeProgWisdom = runIdentity . rephraseProg removeWisdom + +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) + +removePatElemWisdom :: PatElemT (VarWisdom, a) -> PatElemT a +removePatElemWisdom = runIdentity . rephrasePatElem (return . snd) + +addWisdomToPattern :: (Attributes lore, CanBeWise (Op lore)) => + Pattern lore + -> Exp (Wise lore) + -> Pattern (Wise lore) +addWisdomToPattern pat e = + Pattern + (map (`addRanges` unknownRange) ctxals) + (zipWith addRanges valals ranges) + where (ctxals, valals) = Aliases.mkPatternAliases pat e + addRanges patElem range = + let (als, innerlore) = patElemAttr patElem + in patElem `setPatElemLore` (VarWisdom als range, innerlore) + ranges = expRanges e + +mkWiseBody :: (Attributes lore, CanBeWise (Op lore)) => + BodyAttr lore -> Stms (Wise lore) -> Result -> Body (Wise lore) +mkWiseBody innerlore bnds res = + Body (BodyWisdom aliases consumed ranges (Names' $ freeInStmsAndRes bnds res), + innerlore) bnds res + where (aliases, consumed) = Aliases.mkBodyAliases bnds res + ranges = Ranges.mkBodyRanges bnds res + +mkWiseLetStm :: (Attributes lore, CanBeWise (Op lore)) => + Pattern lore + -> StmAux (ExpAttr lore) -> Exp (Wise lore) + -> Stm (Wise lore) +mkWiseLetStm pat (StmAux cs attr) e = + let pat' = addWisdomToPattern pat e + in Let pat' (StmAux cs $ mkWiseExpAttr pat' attr e) e + +mkWiseExpAttr :: (Attributes lore, CanBeWise (Op lore)) => + Pattern (Wise lore) -> ExpAttr lore -> Exp (Wise lore) + -> ExpAttr (Wise lore) +mkWiseExpAttr pat explore e = + (ExpWisdom + (Names' $ consumedInExp e) + (Names' $ freeIn pat <> freeIn explore <> freeInExp e), + explore) + +instance (Bindable lore, + CanBeWise (Op lore)) => Bindable (Wise lore) where + mkExpPat ctx val e = + addWisdomToPattern (mkExpPat ctx val $ removeExpWisdom e) e + + mkExpAttr pat e = + mkWiseExpAttr pat (mkExpAttr (removePatternWisdom pat) $ removeExpWisdom e) e + + mkLetNames names e = do + env <- asksScope removeScopeWisdom + flip runReaderT env $ do + Let pat attr _ <- mkLetNames names $ removeExpWisdom e + return $ mkWiseLetStm pat attr e + + mkBody bnds res = + let Body bodylore _ _ = mkBody (fmap removeStmWisdom bnds) res + in mkWiseBody bodylore bnds res + +class (AliasedOp (OpWithWisdom op), + RangedOp (OpWithWisdom op), + IsOp (OpWithWisdom op), + UsageInOp (OpWithWisdom op)) => CanBeWise op where + type OpWithWisdom op :: * + 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 new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rule.hs @@ -0,0 +1,271 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeFamilies #-} +-- | This module defines the concept of a simplification rule for +-- bindings. The intent is that you pass some context (such as symbol +-- table) and a binding, and is given back a sequence of bindings that +-- compute the same result, but are "better" in some sense. +-- +-- These rewrite rules are "local", in that they do not maintain any +-- state or look at the program as a whole. Compare this to the +-- fusion algorithm in @Futhark.Optimise.Fusion.Fusion@, which must be implemented +-- as its own pass. +module Futhark.Optimise.Simplify.Rule + ( -- * The rule monad + RuleM + , cannotSimplify + , liftMaybe + + -- * Rule definition + , SimplificationRule(..) + , RuleGeneric + , RuleBasicOp + , RuleIf + , RuleDoLoop + + -- * Top-down rules + , TopDown + , TopDownRule + , TopDownRuleGeneric + , TopDownRuleBasicOp + , TopDownRuleIf + , TopDownRuleDoLoop + , TopDownRuleOp + + -- * Bottom-up rules + , BottomUp + , BottomUpRule + , BottomUpRuleGeneric + , BottomUpRuleBasicOp + , BottomUpRuleIf + , BottomUpRuleDoLoop + , BottomUpRuleOp + + -- * Assembling rules + , RuleBook + , ruleBook + + -- * Applying rules + , topDownSimplifyStm + , bottomUpSimplifyStm + ) where + +import Data.Semigroup ((<>)) +import Control.Monad.State +import qualified Data.Semigroup as Sem +import qualified Control.Monad.Fail as Fail +import Control.Monad.Except + +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Representation.AST +import Futhark.Binder + +data RuleError = CannotSimplify + | OtherError String + +-- | The monad in which simplification rules are evaluated. +newtype RuleM lore a = RuleM (BinderT lore (StateT VNameSource (Except RuleError)) a) + deriving (Functor, Applicative, Monad, + MonadFreshNames, HasScope lore, LocalScope lore, + MonadError RuleError) + +instance Fail.MonadFail (RuleM lore) where + fail = throwError . OtherError + +instance (Attributes lore, BinderOps lore) => MonadBinder (RuleM lore) where + type Lore (RuleM lore) = lore + mkExpAttrM pat e = RuleM $ mkExpAttrM pat e + mkBodyM bnds res = RuleM $ mkBodyM bnds res + mkLetNamesM pat e = RuleM $ mkLetNamesM pat e + + addStms = RuleM . addStms + collectStms (RuleM m) = RuleM $ collectStms m + certifying cs (RuleM m) = RuleM $ certifying cs m + +-- | Execute a 'RuleM' action. If succesful, returns the result and a +-- list of new bindings. Even if the action fail, there may still be +-- a monadic effect - particularly, the name source may have been +-- modified. +simplify :: (MonadFreshNames m, HasScope lore m) => + RuleM lore a + -> m (Maybe (a, Stms lore)) +simplify (RuleM m) = do + scope <- askScope + modifyNameSource $ \src -> + case runExcept $ runStateT (runBinderT m scope) src of + Left CannotSimplify -> (Nothing, src) + Left (OtherError err) -> error $ "simplify: " ++ err + Right (x, src') -> (Just x, src') + +cannotSimplify :: RuleM lore a +cannotSimplify = throwError CannotSimplify + +liftMaybe :: Maybe a -> RuleM lore a +liftMaybe Nothing = cannotSimplify +liftMaybe (Just x) = return x + +type RuleGeneric lore a = a -> Stm lore -> RuleM lore () +type RuleBasicOp lore a = (a -> Pattern lore -> StmAux (ExpAttr lore) -> + BasicOp lore -> RuleM lore ()) +type RuleIf lore a = a -> Pattern lore -> StmAux (ExpAttr lore) -> + (SubExp, BodyT lore, BodyT lore, + IfAttr (BranchType lore)) -> + RuleM lore () +type RuleDoLoop lore a = a -> Pattern lore -> StmAux (ExpAttr lore) -> + ([(FParam lore, SubExp)], [(FParam lore, SubExp)], + LoopForm lore, BodyT lore) -> + RuleM lore () +type RuleOp lore a = a -> Pattern lore -> StmAux (ExpAttr lore) -> + Op lore -> RuleM lore () + +-- | 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) + +-- | 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] + } + +instance Sem.Semigroup (Rules lore 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 + mempty = Rules mempty mempty mempty mempty mempty + mappend = (Sem.<>) + +-- | 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 TopDownRuleGeneric lore = RuleGeneric lore (TopDown lore) +type TopDownRuleBasicOp lore = RuleBasicOp lore (TopDown lore) +type TopDownRuleIf lore = RuleIf lore (TopDown lore) +type TopDownRuleDoLoop lore = RuleDoLoop lore (TopDown lore) +type TopDownRuleOp lore = RuleOp lore (TopDown lore) +type TopDownRule lore = SimplificationRule lore (TopDown lore) + +-- | 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 BottomUpRuleGeneric lore = RuleGeneric lore (BottomUp lore) +type BottomUpRuleBasicOp lore = RuleBasicOp lore (BottomUp lore) +type BottomUpRuleIf lore = RuleIf lore (BottomUp lore) +type BottomUpRuleDoLoop lore = RuleDoLoop lore (BottomUp lore) +type BottomUpRuleOp lore = RuleOp lore (BottomUp lore) +type BottomUpRule lore = SimplificationRule lore (BottomUp lore) + +-- | A collection of top-down rules. +type TopDownRules lore = Rules lore (TopDown lore) + +-- | A collection of bottom-up rules. +type BottomUpRules lore = Rules lore (BottomUp lore) + +-- | A collection of both top-down and bottom-up rules. +data RuleBook lore = RuleBook { bookTopDownRules :: TopDownRules lore + , bookBottomUpRules :: BottomUpRules lore + } + +instance Sem.Semigroup (RuleBook lore) where + RuleBook ts1 bs1 <> RuleBook ts2 bs2 = RuleBook (ts1<>ts2) (bs1<>bs2) + +instance Monoid (RuleBook lore) where + mempty = RuleBook mempty mempty + mappend = (Sem.<>) + +-- | Construct a rule book from a collection of rules. +ruleBook :: [TopDownRule m] + -> [BottomUpRule m] + -> RuleBook m +ruleBook topdowns bottomups = + RuleBook (groupRules topdowns) (groupRules bottomups) + where groupRules :: [SimplificationRule m a] -> Rules m a + groupRules rs = Rules rs + (filter forBasicOp rs) + (filter forIf rs) + (filter forDoLoop rs) + (filter forOp rs) + + forBasicOp RuleBasicOp{} = True + forBasicOp RuleGeneric{} = True + forBasicOp _ = False + + forIf RuleIf{} = True + forIf RuleGeneric{} = True + forIf _ = False + + forDoLoop RuleDoLoop{} = True + forDoLoop RuleGeneric{} = True + forDoLoop _ = False + + forOp RuleOp{} = True + forOp RuleGeneric{} = True + forOp _ = False + +-- | @simplifyStm lookup bnd@ performs simplification of the +-- binding @bnd@. If simplification is possible, a replacement list +-- 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, BinderOps lore) => + RuleBook lore + -> ST.SymbolTable lore + -> Stm lore + -> m (Maybe (Stms lore)) +topDownSimplifyStm = applyRules . bookTopDownRules + +-- | @simplifyStm uses bnd@ performs simplification of the binding +-- @bnd@. If simplification is possible, a replacement list of +-- bindings is returned, that bind at least the same names as the +-- 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, BinderOps lore) => + RuleBook lore + -> (ST.SymbolTable lore, UT.UsageTable) + -> Stm lore + -> m (Maybe (Stms lore)) +bottomUpSimplifyStm = applyRules . bookBottomUpRules + +rulesForStm :: Stm lore -> Rules lore a -> [SimplificationRule lore a] +rulesForStm stm = case stmExp stm of BasicOp{} -> rulesBasicOp + DoLoop{} -> rulesDoLoop + Op{} -> rulesOp + If{} -> rulesIf + _ -> rulesAny + +applyRule :: SimplificationRule lore a -> a -> Stm lore -> RuleM lore () +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)) = + f a pat aux (ctx, val, form, body) +applyRule (RuleIf f) a (Let pat aux (If cond tbody fbody ifsort)) = + f a pat aux (cond, tbody, fbody, ifsort) +applyRule (RuleOp f) a (Let pat aux (Op op)) = + f a pat aux op +applyRule _ _ _ = + cannotSimplify + +applyRules :: (MonadFreshNames m, HasScope lore m, BinderOps lore) => + Rules lore a -> a -> Stm lore + -> m (Maybe (Stms lore)) +applyRules rules context stm = applyRules' (rulesForStm stm rules) context stm + +applyRules' :: (MonadFreshNames m, HasScope lore m, BinderOps lore) => + [SimplificationRule lore a] -> a -> Stm lore + -> m (Maybe (Stms lore)) +applyRules' [] _ _ = return Nothing +applyRules' (rule:rules) context bnd = do + res <- simplify $ applyRule rule context bnd + case res of Just ((), bnds) -> return $ Just bnds + Nothing -> applyRules' rules context bnd diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rules.hs @@ -0,0 +1,1239 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeFamilies #-} +-- | This module defines a collection of simplification rules, as per +-- "Futhark.Optimise.Simplify.Rule". They are used in the +-- simplifier. +-- +-- For performance reasons, many sufficiently simple logically +-- separate rules are merged into single "super-rules", like ruleIf +-- and ruleBasicOp. This is because it is relatively expensive to +-- activate a rule just to determine that it does not apply. Thus, it +-- is more efficient to have a few very fat rules than a lot of small +-- rules. This does not affect the compiler result in any way; it is +-- purely an optimisation to speed up compilation. +module Futhark.Optimise.Simplify.Rules + ( standardRules + , removeUnnecessaryCopy + ) +where + +import Control.Monad +import Data.Either +import Data.Foldable (all) +import Data.List hiding (all) +import Data.Maybe +import Data.Semigroup ((<>)) + +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Analysis.DataDependencies +import Futhark.Optimise.Simplify.ClosedForm +import Futhark.Optimise.Simplify.Rule +import Futhark.Analysis.PrimExp.Convert +import Futhark.Representation.AST +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Construct +import Futhark.Transform.Substitute +import Futhark.Util + +topDownRules :: (BinderOps lore, Aliased lore) => [TopDownRule lore] +topDownRules = [ RuleDoLoop hoistLoopInvariantMergeVariables + , RuleDoLoop simplifyClosedFormLoop + , RuleDoLoop simplifKnownIterationLoop + , RuleDoLoop simplifyLoopVariables + , RuleGeneric constantFoldPrimFun + , RuleIf ruleIf + , RuleIf hoistBranchInvariant + , RuleBasicOp ruleBasicOp + ] + +bottomUpRules :: BinderOps lore => [BottomUpRule lore] +bottomUpRules = [ RuleDoLoop removeRedundantMergeVariables + , RuleIf removeDeadBranchResult + , RuleBasicOp simplifyIndex + , RuleBasicOp simplifyConcat + ] + +asInt32PrimExp :: PrimExp v -> PrimExp v +asInt32PrimExp pe + | IntType it <- primExpType pe, it /= Int32 = + ConvOpExp (SExt it Int32) pe + | otherwise = + pe + +-- | 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 = ruleBook topDownRules bottomUpRules + +-- This next one is tricky - it's easy enough to determine that some +-- loop result is not used after the loop, but here, we must also make +-- sure that it does not affect any other values. +-- +-- 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 (_, used) pat _ (ctx, val, form, body) + | not $ all (usedAfterLoop . fst) val, + null ctx = -- FIXME: things get tricky if we can remove all vals + -- but some ctxs are still used. We take the easy way + -- out for now. + let (ctx_es, val_es) = splitAt (length ctx) $ bodyResult body + necessaryForReturned = + findNecessaryForReturned usedAfterLoopOrInForm + (zip (map fst $ ctx++val) $ ctx_es++val_es) (dataDependencies body) + + resIsNecessary ((v,_), _) = + usedAfterLoop v || + paramName v `S.member` necessaryForReturned || + referencedInPat v || + referencedInForm v + + (keep_ctx, discard_ctx) = + partition resIsNecessary $ zip ctx ctx_es + (keep_valpart, discard_valpart) = + partition (resIsNecessary . snd) $ + zip (patternValueElements pat) $ zip val val_es + + (keep_valpatelems, keep_val) = unzip keep_valpart + (_discard_valpatelems, discard_val) = unzip discard_valpart + (ctx', ctx_es') = unzip keep_ctx + (val', val_es') = unzip keep_val + + body' = body { bodyResult = ctx_es' ++ val_es' } + free_in_keeps = freeIn keep_valpatelems + + stillUsedContext pat_elem = + patElemName pat_elem `S.member` + (free_in_keeps <> + freeIn (filter (/=pat_elem) $ patternContextElements pat)) + + pat' = pat { patternValueElements = keep_valpatelems + , patternContextElements = + filter stillUsedContext $ patternContextElements pat } + in if ctx' ++ val' == ctx ++ val + then cannotSimplify + else do + -- We can't just remove the bindings in 'discard', since the loop + -- body may still use their names in (now-dead) expressions. + -- Hence, we add them inside the loop, fully aware that dead-code + -- removal will eventually get rid of them. Some care is + -- necessary to handle unique bindings. + body'' <- insertStmsM $ do + mapM_ (uncurry letBindNames) $ dummyStms discard_ctx + mapM_ (uncurry letBindNames) $ dummyStms discard_val + return body' + letBind_ pat' $ DoLoop ctx' val' form body'' + where pat_used = map (`UT.isUsedDirectly` used) $ patternValueNames pat + used_vals = map fst $ filter snd $ zip (map (paramName . fst) val) pat_used + usedAfterLoop = flip elem used_vals . paramName + usedAfterLoopOrInForm p = + usedAfterLoop p || paramName p `S.member` freeIn form + patAnnotNames = freeIn $ map fst $ ctx++val + referencedInPat = (`S.member` patAnnotNames) . paramName + referencedInForm = (`S.member` freeIn form) . paramName + + dummyStms = map dummyStm + dummyStm ((p,e), _) + | unique (paramDeclType p), + Var v <- e = ([paramName p], BasicOp $ Copy v) + | otherwise = ([paramName p], BasicOp $ SubExp e) +removeRedundantMergeVariables _ _ _ _ = + cannotSimplify + +-- 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 _ pat _ (ctx, val, form, loopbody) = + -- Figure out which of the elements of loopresult are + -- loop-invariant, and hoist them out. + case foldr checkInvariance ([], explpat, [], []) $ + zip merge res of + ([], _, _, _) -> + -- Nothing is invariant. + cannotSimplify + (invariant, explpat', merge', res') -> do + -- We have moved something invariant out of the loop. + let loopbody' = loopbody { bodyResult = res' } + invariantShape :: (a, VName) -> Bool + invariantShape (_, shapemerge) = shapemerge `elem` + map (paramName . fst) merge' + (implpat',implinvariant) = partition invariantShape implpat + implinvariant' = [ (patElemIdent p, Var v) | (p,v) <- implinvariant ] + implpat'' = map fst implpat' + explpat'' = map fst explpat' + (ctx', val') = splitAt (length implpat') merge' + forM_ (invariant ++ implinvariant') $ \(v1,v2) -> + letBindNames_ [identName v1] $ BasicOp $ SubExp v2 + letBind_ (Pattern implpat'' explpat'') $ + DoLoop ctx' val' form loopbody' + where merge = ctx ++ val + res = bodyResult loopbody + + implpat = zip (patternContextElements pat) $ + map paramName $ loopResultContext (map fst ctx) (map fst val) + explpat = zip (patternValueElements pat) $ + map (paramName . fst) val + + namesOfMergeParams = S.fromList $ map (paramName . fst) $ ctx++val + + removeFromResult (mergeParam,mergeInit) explpat' = + case partition ((==paramName mergeParam) . snd) explpat' of + ([(patelem,_)], rest) -> + (Just (patElemIdent patelem, mergeInit), rest) + (_, _) -> + (Nothing, explpat') + + checkInvariance + ((mergeParam,mergeInit), resExp) + (invariant, explpat', merge', resExps) + | not (unique (paramDeclType mergeParam)) || + arrayRank (paramDeclType mergeParam) == 1, + isInvariant resExp, + -- Also do not remove the condition in a while-loop. + not $ paramName mergeParam `S.member` freeIn form = + let (bnd, explpat'') = + removeFromResult (mergeParam,mergeInit) explpat' + in (maybe id (:) bnd $ (paramIdent mergeParam, mergeInit) : invariant, + explpat'', merge', resExps) + where + -- A non-unique merge variable is invariant if the corresponding + -- subexp in the result is EITHER: + -- + -- (0) a variable of the same name as the parameter, where + -- all existential parameters are already known to be + -- invariant + isInvariant (Var v2) + | paramName mergeParam == v2 = + allExistentialInvariant + (S.fromList $ map (identName . fst) invariant) mergeParam + -- (1) or identical to the initial value of the parameter. + isInvariant _ = mergeInit == resExp + + checkInvariance ((mergeParam,mergeInit), resExp) (invariant, explpat', merge', resExps) = + (invariant, explpat', (mergeParam,mergeInit):merge', resExp:resExps) + + allExistentialInvariant namesOfInvariant mergeParam = + all (invariantOrNotMergeParam namesOfInvariant) + (paramName mergeParam `S.delete` freeIn mergeParam) + invariantOrNotMergeParam namesOfInvariant name = + not (name `S.member` namesOfMergeParams) || + name `S.member` namesOfInvariant + +-- | A function that, given a variable name, returns its definition. +type VarLookup lore = VName -> Maybe (Exp lore, 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 lore -> Maybe (BasicOp lore, Certificates) + +simpleRules :: [SimpleRule lore] +simpleRules = [ simplifyBinOp + , simplifyCmpOp + , simplifyUnOp + , simplifyConvOp + , simplifyAssert + , copyScratchToScratch + , simplifyIdentityReshape + , simplifyReshapeReshape + , simplifyReshapeScratch + , simplifyReshapeReplicate + , simplifyReshapeIota + , improveReshape ] + +simplifyClosedFormLoop :: BinderOps lore => TopDownRuleDoLoop lore +simplifyClosedFormLoop _ pat _ ([], val, ForLoop i _ bound [], body) = + loopClosedForm pat val (S.singleton i) bound body +simplifyClosedFormLoop _ _ _ _ = cannotSimplify + +simplifyLoopVariables :: (BinderOps lore, Aliased lore) => TopDownRuleDoLoop lore +simplifyLoopVariables vtable pat _ (ctx, val, form@(ForLoop i it num_iters loop_vars), body) + | simplifiable <- map checkIfSimplifiable loop_vars, + not $ all isNothing simplifiable = do + -- Check if the simplifications throw away more information than + -- we are comfortable with at this stage. + (maybe_loop_vars, body_prefix_stms) <- + localScope (scopeOf form) $ + unzip <$> zipWithM onLoopVar loop_vars simplifiable + if maybe_loop_vars == map Just loop_vars + then cannotSimplify + else do body' <- insertStmsM $ do + addStms $ mconcat body_prefix_stms + resultBodyM =<< bodyBind body + letBind_ pat $ DoLoop ctx val + (ForLoop i it num_iters $ catMaybes maybe_loop_vars) body' + + where seType (Var v) + | v == i = Just $ Prim $ IntType it + | otherwise = ST.lookupType v vtable + seType (Constant v) = Just $ Prim $ primValueType v + consumed_in_body = consumedInBody body + + vtable' = ST.fromScope (scopeOf form) <> vtable + + checkIfSimplifiable (p,arr) = + simplifyIndexing vtable' seType arr + (DimFix (Var i) : fullSlice (paramType p) []) $ + paramName p `S.member` consumed_in_body + + -- We only want this simplification if the result does not refer + -- to 'i' at all, or does not contain accesses. + onLoopVar (p,arr) Nothing = + return (Just (p,arr), mempty) + onLoopVar (p,arr) (Just m) = do + (x,x_stms) <- collectStms m + case x of + IndexResult cs arr' slice + | all (not . (i `S.member`) . freeInStm) x_stms, + DimFix (Var j) : slice' <- slice, + j == i, not $ i `S.member` freeIn slice -> do + addStms x_stms + w <- arraySize 0 <$> lookupType arr' + for_in_partial <- + certifying cs $ letExp "for_in_partial" $ BasicOp $ Index arr' $ + DimSlice (intConst Int32 0) w (intConst Int32 1) : slice' + return (Just (p, for_in_partial), mempty) + + SubExpResult cs se + | all (notIndex . stmExp) x_stms -> do + x_stms' <- collectStms_ $ certifying cs $ do + addStms x_stms + letBindNames_ [paramName p] $ BasicOp $ SubExp se + return (Nothing, x_stms') + + _ -> return (Just (p,arr), mempty) + + notIndex (BasicOp Index{}) = False + notIndex _ = True +simplifyLoopVariables _ _ _ _ = cannotSimplify + +simplifKnownIterationLoop :: BinderOps lore => TopDownRuleDoLoop lore +simplifKnownIterationLoop _ pat _ (ctx, val, ForLoop i it (Constant iters) loop_vars, body) + | zeroIsh iters = do + let bindResult p r = letBindNames [patElemName p] $ BasicOp $ SubExp r + zipWithM_ bindResult (patternContextElements pat) (map snd ctx) + zipWithM_ bindResult (patternValueElements pat) (map snd val) + + | oneIsh iters = do + + forM_ (ctx++val) $ \(mergevar, mergeinit) -> + letBindNames [paramName mergevar] $ BasicOp $ SubExp mergeinit + + letBindNames_ [i] $ BasicOp $ SubExp $ intConst it 0 + + forM_ loop_vars $ \(p,arr) -> + letBindNames_ [paramName p] $ BasicOp $ Index arr $ + DimFix (intConst Int32 0) : fullSlice (paramType p) [] + + (loop_body_ctx, loop_body_val) <- splitAt (length ctx) <$> (mapM asVar =<< bodyBind body) + let subst = M.fromList $ zip (map (paramName . fst) ctx) loop_body_ctx + ctx_params = substituteNames subst $ map fst ctx + val_params = substituteNames subst $ map fst val + res_context = loopResultContext ctx_params val_params + forM_ (zip (patternContextElements pat) res_context) $ \(pat_elem, p) -> + letBind_ (Pattern [] [pat_elem]) $ BasicOp $ SubExp $ Var $ paramName p + forM_ (zip (patternValueElements pat) loop_body_val) $ \(pat_elem, v) -> + letBind_ (Pattern [] [pat_elem]) $ BasicOp $ SubExp $ Var v + where asVar (Var v) = return v + asVar (Constant v) = letExp "named" $ BasicOp $ SubExp $ Constant v +simplifKnownIterationLoop _ _ _ _ = + cannotSimplify + +-- | 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 => BottomUpRuleBasicOp lore +removeUnnecessaryCopy (vtable,used) (Pattern [] [d]) _ (Copy v) + | not (v `UT.used` used), + consumable || not (patElemName d `UT.isConsumed` used) = + letBind_ (Pattern [] [d]) $ BasicOp $ SubExp $ Var v + where -- We need to make sure we can even consume the original. + -- This is currently a hacky check, much too conservative, + -- because we don't have the information conveniently + -- available. + consumable = case M.lookup v $ ST.toScope vtable of + Just (FParamInfo info) -> unique $ declTypeOf info + _ -> False +removeUnnecessaryCopy _ _ _ _ = cannotSimplify + +simplifyCmpOp :: SimpleRule lore +simplifyCmpOp _ _ (CmpOp cmp e1 e2) + | e1 == e2 = constRes $ BoolValue $ + case cmp of CmpEq{} -> True + CmpSlt{} -> False + CmpUlt{} -> False + CmpSle{} -> True + CmpUle{} -> True + FCmpLt{} -> False + FCmpLe{} -> True + CmpLlt -> False + CmpLle -> True +simplifyCmpOp _ _ (CmpOp cmp (Constant v1) (Constant v2)) = + constRes =<< BoolValue <$> doCmpOp cmp v1 v2 +simplifyCmpOp _ _ _ = Nothing + +simplifyBinOp :: SimpleRule lore + +simplifyBinOp _ _ (BinOp op (Constant v1) (Constant v2)) + | Just res <- doBinOp op v1 v2 = + constRes res + +simplifyBinOp _ _ (BinOp Add{} e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + +simplifyBinOp _ _ (BinOp FAdd{} e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + + +simplifyBinOp look _ (BinOp Sub{} e1 e2) + | isCt0 e2 = subExpRes e1 + -- Cases for simplifying (a+b)-b and permutations. + | Var v1 <- e1, + Just (BasicOp (BinOp Add{} e1_a e1_b), cs) <- look v1, + e1_a == e2 = Just (SubExp e1_b, cs) + | Var v1 <- e1, + Just (BasicOp (BinOp Add{} e1_a e1_b), cs) <- look v1, + e1_b == e2 = Just (SubExp e1_a, cs) + | Var v2 <- e2, + Just (BasicOp (BinOp Add{} e2_a e2_b), cs) <- look v2, + e2_a == e1 = Just (SubExp e2_b, cs) + | Var v2 <- e1, + Just (BasicOp (BinOp Add{} e2_a e2_b), cs) <- look v2, + e2_b == e1 = Just (SubExp e2_a, cs) + +simplifyBinOp _ _ (BinOp FSub{} e1 e2) + | isCt0 e2 = subExpRes e1 + +simplifyBinOp _ _ (BinOp Mul{} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt0 e2 = subExpRes e2 + | isCt1 e1 = subExpRes e2 + | isCt1 e2 = subExpRes e1 + +simplifyBinOp _ _ (BinOp FMul{} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt0 e2 = subExpRes e2 + | isCt1 e1 = subExpRes e2 + | isCt1 e2 = subExpRes e1 + +simplifyBinOp look _ (BinOp (SMod t) e1 e2) + | isCt1 e2 = constRes $ IntValue $ intValue t (0 :: Int) + | e1 == e2 = constRes $ IntValue $ intValue t (0 :: Int) + | Var v1 <- e1, + Just (BasicOp (BinOp SMod{} _ e4), v1_cs) <- look v1, + e4 == e2 = Just (SubExp e1, v1_cs) + +simplifyBinOp _ _ (BinOp SDiv{} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt1 e2 = subExpRes e1 + | isCt0 e2 = Nothing + +simplifyBinOp _ _ (BinOp FDiv{} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt1 e2 = subExpRes e1 + | isCt0 e2 = Nothing + +simplifyBinOp _ _ (BinOp (SRem t) e1 e2) + | isCt1 e2 = constRes $ IntValue $ intValue t (0 :: Int) + | e1 == e2 = constRes $ IntValue $ intValue t (1 :: Int) + +simplifyBinOp _ _ (BinOp SQuot{} e1 e2) + | isCt1 e2 = subExpRes e1 + | isCt0 e2 = Nothing + +simplifyBinOp _ _ (BinOp (FPow t) e1 e2) + | isCt0 e2 = subExpRes $ floatConst t 1 + | isCt0 e1 || isCt1 e1 || isCt1 e2 = subExpRes e1 + +simplifyBinOp _ _ (BinOp (Shl t) e1 e2) + | isCt0 e2 = subExpRes e1 + | isCt0 e1 = subExpRes $ intConst t 0 + +simplifyBinOp _ _ (BinOp AShr{} e1 e2) + | isCt0 e2 = subExpRes e1 + +simplifyBinOp _ _ (BinOp (And t) e1 e2) + | isCt0 e1 = subExpRes $ intConst t 0 + | isCt0 e2 = subExpRes $ intConst t 0 + | e1 == e2 = subExpRes e1 + +simplifyBinOp _ _ (BinOp Or{} e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + | e1 == e2 = subExpRes e1 + +simplifyBinOp _ _ (BinOp (Xor t) e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + | e1 == e2 = subExpRes $ intConst t 0 + +simplifyBinOp defOf _ (BinOp LogAnd e1 e2) + | isCt0 e1 = constRes $ BoolValue False + | isCt0 e2 = constRes $ BoolValue False + | isCt1 e1 = subExpRes e2 + | isCt1 e2 = subExpRes e1 + | Var v <- e1, + Just (BasicOp (UnOp Not e1'), v_cs) <- defOf v, + e1' == e2 = Just (SubExp $ Constant $ BoolValue False, v_cs) + | Var v <- e2, + Just (BasicOp (UnOp Not e2'), v_cs) <- defOf v, + e2' == e1 = Just (SubExp $ Constant $ BoolValue False, v_cs) + +simplifyBinOp defOf _ (BinOp LogOr e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + | isCt1 e1 = constRes $ BoolValue True + | isCt1 e2 = constRes $ BoolValue True + | Var v <- e1, + Just (BasicOp (UnOp Not e1'), v_cs) <- defOf v, + e1' == e2 = Just (SubExp $ Constant $ BoolValue True, v_cs) + | Var v <- e2, + Just (BasicOp (UnOp Not e2'), v_cs) <- defOf v, + e2' == e1 = Just (SubExp $ Constant $ BoolValue True, v_cs) + +simplifyBinOp defOf _ (BinOp (SMax it) e1 e2) + | e1 == e2 = + subExpRes e1 + | Var v1 <- e1, + Just (BasicOp (BinOp (SMax _) e1_1 e1_2), v1_cs) <- defOf v1, + e1_1 == e2 = + Just (BinOp (SMax it) e1_2 e2, v1_cs) + | Var v1 <- e1, + Just (BasicOp (BinOp (SMax _) e1_1 e1_2), v1_cs) <- defOf v1, + e1_2 == e2 = + Just (BinOp (SMax it) e1_1 e2, v1_cs) + | Var v2 <- e2, + Just (BasicOp (BinOp (SMax _) e2_1 e2_2), v2_cs) <- defOf v2, + e2_1 == e1 = + Just (BinOp (SMax it) e2_2 e1, v2_cs) + | Var v2 <- e2, + Just (BasicOp (BinOp (SMax _) e2_1 e2_2), v2_cs) <- defOf v2, + e2_2 == e1 = + Just (BinOp (SMax it) e2_1 e1, v2_cs) + +simplifyBinOp _ _ _ = Nothing + +constRes :: PrimValue -> Maybe (BasicOp lore, Certificates) +constRes = Just . (,mempty) . SubExp . Constant + +subExpRes :: SubExp -> Maybe (BasicOp lore, Certificates) +subExpRes = Just . (,mempty) . SubExp + +simplifyUnOp :: SimpleRule lore +simplifyUnOp _ _ (UnOp op (Constant v)) = + constRes =<< doUnOp op v +simplifyUnOp defOf _ (UnOp Not (Var v)) + | Just (BasicOp (UnOp Not v2), v_cs) <- defOf v = + Just (SubExp v2, v_cs) +simplifyUnOp _ _ _ = + Nothing + +simplifyConvOp :: SimpleRule lore +simplifyConvOp _ _ (ConvOp op (Constant v)) = + constRes =<< doConvOp op v +simplifyConvOp _ _ (ConvOp op se) + | (from, to) <- convOpType op, from == to = + subExpRes se +simplifyConvOp lookupVar _ (ConvOp (SExt t2 t1) (Var v)) + | Just (BasicOp (ConvOp (SExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (SExt t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (ZExt t2 t1) (Var v)) + | Just (BasicOp (ConvOp (ZExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (ZExt t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (SIToFP t2 t1) (Var v)) + | Just (BasicOp (ConvOp (SExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (SIToFP t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (UIToFP t2 t1) (Var v)) + | Just (BasicOp (ConvOp (ZExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (UIToFP t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (FPConv t2 t1) (Var v)) + | Just (BasicOp (ConvOp (FPConv t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (FPConv t3 t1) se, v_cs) +simplifyConvOp _ _ _ = + Nothing + +-- If expression is true then just replace assertion. +simplifyAssert :: SimpleRule lore +simplifyAssert _ _ (Assert (Constant (BoolValue True)) _ _) = + constRes Checked +simplifyAssert _ _ _ = + Nothing + +constantFoldPrimFun :: BinderOps lore => TopDownRuleGeneric lore +constantFoldPrimFun _ (Let pat (StmAux cs _) (Apply fname args _ _)) + | Just args' <- mapM (isConst . fst) args, + Just (_, _, fun) <- M.lookup (nameToString fname) primFuns, + Just result <- fun args' = + certifying cs $ letBind_ pat $ BasicOp $ SubExp $ Constant result + where isConst (Constant v) = Just v + isConst _ = Nothing +constantFoldPrimFun _ _ = cannotSimplify + +simplifyIndex :: BinderOps lore => BottomUpRuleBasicOp lore +simplifyIndex (vtable, used) pat@(Pattern [] [pe]) (StmAux cs _) (Index idd inds) + | Just m <- simplifyIndexing vtable seType idd inds consumed = do + res <- m + case res of + SubExpResult cs' se -> + certifying (cs<>cs') $ letBindNames_ (patternNames pat) $ + BasicOp $ SubExp se + IndexResult extra_cs idd' inds' -> + certifying (cs<>extra_cs) $ + letBindNames_ (patternNames pat) $ BasicOp $ Index idd' inds' + where consumed = patElemName pe `UT.isConsumed` used + seType (Var v) = ST.lookupType v vtable + seType (Constant v) = Just $ Prim $ primValueType v + +simplifyIndex _ _ _ _ = cannotSimplify + +data IndexResult = IndexResult Certificates VName (Slice SubExp) + | SubExpResult Certificates SubExp + +simplifyIndexing :: MonadBinder m => + ST.SymbolTable (Lore m) -> TypeLookup + -> VName -> Slice SubExp -> Bool + -> Maybe (m IndexResult) +simplifyIndexing vtable seType idd inds consuming = + case defOf idd of + _ | Just t <- seType (Var idd), + inds == fullSlice t [] -> + Just $ pure $ SubExpResult mempty $ Var idd + + | Just inds' <- sliceIndices inds, + Just (e, cs) <- ST.index idd inds' vtable, + worthInlining e -> + Just $ SubExpResult cs <$> (letSubExp "index_primexp" =<< toExp e) + + Nothing -> Nothing + + Just (SubExp (Var v), cs) -> Just $ pure $ IndexResult cs v inds + + Just (Iota _ x s to_it, cs) + | [DimFix ii] <- inds, + Just (Prim (IntType from_it)) <- seType ii -> + Just $ + fmap (SubExpResult cs) $ letSubExp "index_iota" <=< toExp $ + ConvOpExp (SExt from_it to_it) (primExpFromSubExp (IntType from_it) ii) + * primExpFromSubExp (IntType to_it) s + + primExpFromSubExp (IntType to_it) x + | [DimSlice i_offset i_n i_stride] <- inds -> + Just $ do + i_offset' <- asIntS to_it i_offset + i_stride' <- asIntS to_it i_stride + i_offset'' <- letSubExp "iota_offset" $ + BasicOp $ BinOp (Add Int32) x i_offset' + i_stride'' <- letSubExp "iota_offset" $ + BasicOp $ BinOp (Mul Int32) s i_stride' + fmap (SubExpResult cs) $ letSubExp "slice_iota" $ + BasicOp $ Iota i_n i_offset'' i_stride'' to_it + + Just (Rotate offsets a, cs) -> Just $ do + dims <- arrayDims <$> lookupType a + let adjustI i o d = do + i_p_o <- letSubExp "i_p_o" $ BasicOp $ BinOp (Add Int32) i o + letSubExp "rot_i" (BasicOp $ BinOp (SMod Int32) i_p_o d) + adjust (DimFix i, o, d) = + DimFix <$> adjustI i o d + adjust (DimSlice i n s, o, d) = + DimSlice <$> adjustI i o d <*> pure n <*> pure s + IndexResult cs a <$> mapM adjust (zip3 inds offsets dims) + + Just (Index aa ais, cs) -> + Just $ IndexResult cs aa <$> sliceSlice ais inds + + Just (Replicate (Shape [_]) (Var vv), cs) + | [DimFix{}] <- inds, not consuming -> Just $ pure $ SubExpResult cs $ Var vv + | DimFix{}:is' <- inds, not consuming -> Just $ pure $ IndexResult cs vv is' + + Just (Replicate (Shape [_]) val@(Constant _), cs) + | [DimFix{}] <- inds, not consuming -> Just $ pure $ SubExpResult cs val + + Just (Replicate (Shape ds) v, cs) + | (ds_inds, rest_inds) <- splitAt (length ds) inds, + (ds', ds_inds') <- unzip $ mapMaybe index ds_inds, + ds' /= ds -> + Just $ do + arr <- letExp "smaller_replicate" $ BasicOp $ Replicate (Shape ds') v + return $ IndexResult cs arr $ ds_inds' ++ rest_inds + where index DimFix{} = Nothing + index (DimSlice _ n s) = Just (n, DimSlice (constant (0::Int32)) n s) + + Just (Rearrange perm src, cs) + | rearrangeReach perm <= length (takeWhile isIndex inds) -> + let inds' = rearrangeShape (rearrangeInverse perm) inds + in Just $ pure $ IndexResult cs src inds' + where isIndex DimFix{} = True + isIndex _ = False + + Just (Copy src, cs) + | Just dims <- arrayDims <$> seType (Var src), + length inds == length dims, + not consuming -> + Just $ pure $ IndexResult cs src inds + + Just (Reshape newshape src, cs) + | Just newdims <- shapeCoercion newshape, + Just olddims <- arrayDims <$> seType (Var src), + changed_dims <- zipWith (/=) newdims olddims, + not $ or $ drop (length inds) changed_dims -> + Just $ pure $ IndexResult cs src inds + + | Just newdims <- shapeCoercion newshape, + Just olddims <- arrayDims <$> seType (Var src), + length newshape == length inds, + length olddims == length newdims -> + Just $ pure $ IndexResult cs src inds + + Just (Reshape [_] v2, cs) + | Just [_] <- arrayDims <$> seType (Var v2) -> + Just $ pure $ IndexResult cs v2 inds + + Just (Concat d x xs _, cs) + | Just (ibef, DimFix i, iaft) <- focusNth d inds, + Just (Prim res_t) <- (`setArrayDims` sliceDims inds) <$> + ST.lookupType x vtable -> Just $ do + x_len <- arraySize d <$> lookupType x + xs_lens <- mapM (fmap (arraySize d) . lookupType) xs + + let add n m = do + added <- letSubExp "index_concat_add" $ BasicOp $ BinOp (Add Int32) n m + return (added, n) + (_, starts) <- mapAccumLM add x_len xs_lens + let xs_and_starts = reverse $ zip xs starts + + let mkBranch [] = + letSubExp "index_concat" $ BasicOp $ Index x $ ibef ++ DimFix i : iaft + mkBranch ((x', start):xs_and_starts') = do + cmp <- letSubExp "index_concat_cmp" $ BasicOp $ CmpOp (CmpSle Int32) start i + (thisres, thisbnds) <- collectStms $ do + i' <- letSubExp "index_concat_i" $ BasicOp $ BinOp (Sub Int32) i start + letSubExp "index_concat" $ BasicOp $ Index x' $ ibef ++ DimFix i' : iaft + thisbody <- mkBodyM thisbnds [thisres] + (altres, altbnds) <- collectStms $ mkBranch xs_and_starts' + altbody <- mkBodyM altbnds [altres] + letSubExp "index_concat_branch" $ If cmp thisbody altbody $ + IfAttr [primBodyType res_t] IfNormal + SubExpResult cs <$> mkBranch xs_and_starts + + Just (ArrayLit ses _, cs) + | DimFix (Constant (IntValue (Int32Value i))) : inds' <- inds, + Just se <- maybeNth i ses -> + case inds' of + [] -> Just $ pure $ SubExpResult cs se + _ | Var v2 <- se -> Just $ pure $ IndexResult cs v2 inds' + _ -> Nothing + + -- Indexing single-element arrays. We know the index must be 0. + _ | Just t <- seType $ Var idd, isCt1 $ arraySize 0 t, + DimFix i : inds' <- inds, not $ isCt0 i -> + Just $ pure $ IndexResult mempty idd $ + DimFix (constant (0::Int32)) : inds' + + _ -> Nothing + + where defOf v = do (BasicOp op, def_cs) <- ST.lookupExp v vtable + return (op, def_cs) + + -- | A crude heuristic for determining when a PrimExp is + -- worth inlining over keeping it in an array and reading it + -- from memory. + worthInlining e + | length e > 10 = False -- totally ad-hoc. + worthInlining (BinOpExp Pow{} _ _) = False + worthInlining (BinOpExp FPow{} _ _) = False + worthInlining (BinOpExp _ x y) = worthInlining x && worthInlining y + worthInlining (CmpOpExp _ x y) = worthInlining x && worthInlining y + worthInlining (ConvOpExp _ x) = worthInlining x + worthInlining (UnOpExp _ x) = worthInlining x + worthInlining FunExp{} = False + worthInlining _ = True + +sliceSlice :: MonadBinder m => + [DimIndex SubExp] -> [DimIndex SubExp] -> m [DimIndex SubExp] +sliceSlice (DimFix j:js') is' = (DimFix j:) <$> sliceSlice js' is' +sliceSlice (DimSlice j _ s:js') (DimFix i:is') = do + i_t_s <- letSubExp "j_t_s" $ BasicOp $ BinOp (Mul Int32) i s + j_p_i_t_s <- letSubExp "j_p_i_t_s" $ BasicOp $ BinOp (Add Int32) j i_t_s + (DimFix j_p_i_t_s:) <$> sliceSlice js' is' +sliceSlice (DimSlice j _ s0:js') (DimSlice i n s1:is') = do + s0_t_i <- letSubExp "s0_t_i" $ BasicOp $ BinOp (Mul Int32) s0 i + j_p_s0_t_i <- letSubExp "j_p_s0_t_i" $ BasicOp $ BinOp (Add Int32) j s0_t_i + (DimSlice j_p_s0_t_i n s1:) <$> sliceSlice js' is' +sliceSlice _ _ = return [] + + +simplifyConcat :: BinderOps lore => BottomUpRuleBasicOp lore + +-- 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, + let perm = [i] ++ [0..i-1] ++ [i+1..r-1], + Just (x',x_cs) <- transposedBy perm x, + Just (xs',xs_cs) <- unzip <$> mapM (transposedBy perm) xs = do + concat_rearrange <- + certifying (x_cs<>mconcat xs_cs) $ + letExp "concat_rearrange" $ BasicOp $ Concat 0 x' xs' new_d + letBind_ pat $ BasicOp $ Rearrange perm concat_rearrange + where transposedBy perm1 v = + case ST.lookupExp v vtable of + Just (BasicOp (Rearrange perm2 v'), vcs) + | perm1 == perm2 -> Just (v', vcs) + _ -> Nothing + +-- concat xs (concat ys zs) == concat xs ys zs +simplifyConcat (vtable, _) pat (StmAux cs _) (Concat i x xs new_d) + | x' /= x || concat xs' /= xs = + certifying (cs<>x_cs<>mconcat xs_cs) $ + letBind_ pat $ BasicOp $ Concat i x' (zs++concat xs') new_d + where (x':zs, x_cs) = isConcat x + (xs', xs_cs) = unzip $ map isConcat xs + isConcat v = case ST.lookupBasicOp v vtable of + Just (Concat j y ys _, v_cs) | j == i -> (y : ys, v_cs) + _ -> ([v], mempty) + +-- If concatenating a bunch of array literals (or equivalent +-- replicate), just construct the array literal instead. +simplifyConcat (vtable, _) pat (StmAux cs _) (Concat 0 x xs _) + | Just (vs, vcs) <- unzip <$> mapM isArrayLit (x:xs) = do + rt <- rowType <$> lookupType x + certifying (cs <> mconcat vcs) $ + letBind_ pat $ BasicOp $ ArrayLit vs rt + where isArrayLit v + | Just (Replicate shape se, vcs) <- ST.lookupBasicOp v vtable, + unitShape shape = Just (se, vcs) + | Just (ArrayLit [se] _, vcs) <- ST.lookupBasicOp v vtable = + Just (se, vcs) + | otherwise = + Nothing + + unitShape = (==Shape [Constant $ IntValue $ Int32Value 1]) + +simplifyConcat _ _ _ _ = cannotSimplify + +ruleIf :: BinderOps lore => TopDownRuleIf lore + +ruleIf _ pat _ (e1, tb, fb, IfAttr t ifsort) + | Just branch <- checkBranch, + ifsort /= IfFallback || isCt1 e1 = do + let ses = bodyResult branch + addStms $ bodyStms branch + ctx <- subExpShapeContext (bodyTypeValues t) ses + let ses' = ctx ++ ses + sequence_ [ letBind (Pattern [] [p]) $ BasicOp $ SubExp se + | (p,se) <- zip (patternElements pat) ses'] + + where checkBranch + | isCt1 e1 = Just tb + | isCt0 e1 = Just fb + | otherwise = Nothing + +-- IMPROVE: the following two rules can be generalised to work in more +-- cases, especially when the branches have bindings, or return more +-- than one value. +-- +-- if c then True else v == c || v +ruleIf _ pat _ + (cond, Body _ tstms [Constant (BoolValue True)], + Body _ fstms [se], IfAttr ts _) + | null tstms, null fstms, [Prim Bool] <- bodyTypeValues ts = + letBind_ pat $ BasicOp $ BinOp LogOr cond se + +-- When type(x)==bool, if c then x else y == (c && x) || (!c && y) +ruleIf _ pat _ (cond, tb, fb, IfAttr ts _) + | Body _ tstms [tres] <- tb, + Body _ fstms [fres] <- fb, + all (safeExp . stmExp) $ tstms <> fstms, + all (==Prim Bool) $ bodyTypeValues ts = do + addStms tstms + addStms fstms + e <- eBinOp LogOr (pure $ BasicOp $ BinOp LogAnd cond tres) + (eBinOp LogAnd (pure $ BasicOp $ UnOp Not cond) + (pure $ BasicOp $ SubExp fres)) + letBind_ pat e + +ruleIf _ pat _ (_, tbranch, _, IfAttr _ IfFallback) + | null $ patternContextNames pat, + all (safeExp . stmExp) $ bodyStms tbranch = do + let ses = bodyResult tbranch + addStms $ bodyStms tbranch + sequence_ [ letBind (Pattern [] [p]) $ BasicOp $ SubExp se + | (p,se) <- zip (patternElements pat) ses] + +ruleIf _ _ _ _ = cannotSimplify + +-- | 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 _ pat _ (cond, tb, fb, IfAttr ret ifsort) = do + let tses = bodyResult tb + fses = bodyResult fb + (hoistings, (pes, ts, res)) <- + fmap (fmap unzip3 . partitionEithers) $ mapM branchInvariant $ + zip3 (patternElements pat) + (map Left [0..num_ctx-1] ++ map Right ret) + (zip tses fses) + let ctx_fixes = catMaybes hoistings + (tses', fses') = unzip res + tb' = tb { bodyResult = tses' } + fb' = fb { bodyResult = fses' } + ret' = foldr (uncurry fixExt) (rights ts) ctx_fixes + (ctx_pes, val_pes) = splitFromEnd (length ret') pes + if not $ null hoistings -- Was something hoisted? + then do -- We may have to add some reshapes if we made the type + -- less existential. + tb'' <- reshapeBodyResults tb' $ map extTypeOf ret' + fb'' <- reshapeBodyResults fb' $ map extTypeOf ret' + letBind_ (Pattern ctx_pes val_pes) $ + If cond tb'' fb'' (IfAttr ret' ifsort) + else cannotSimplify + where num_ctx = length $ patternContextElements pat + bound_in_branches = S.fromList $ concatMap (patternNames . stmPattern) $ + bodyStms tb <> bodyStms fb + mem_sizes = freeIn $ filter (isMem . patElemType) $ patternElements pat + invariant Constant{} = True + invariant (Var v) = not $ v `S.member` bound_in_branches + + isMem Mem{} = True + isMem _ = False + sizeOfMem v = v `S.member` mem_sizes + + branchInvariant (pe, t, (tse, fse)) + -- Do both branches return the same value? + | tse == fse = do + letBind_ (Pattern [] [pe]) $ BasicOp $ SubExp tse + hoisted pe t + + -- Do both branches return values that are free in the + -- branch, and are we not the only pattern element? The + -- latter is to avoid infinite application of this rule. + | invariant tse, invariant fse, patternSize pat > 1, + Prim _ <- patElemType pe, not $ sizeOfMem $ patElemName pe = do + bt <- expTypesFromPattern $ Pattern [] [pe] + letBind_ (Pattern [] [pe]) =<< + (If cond <$> resultBodyM [tse] + <*> resultBodyM [fse] + <*> pure (IfAttr bt ifsort)) + hoisted pe t + + | otherwise = + return $ Right (pe, t, (tse,fse)) + + hoisted pe (Left i) = return $ Left $ Just (i, Var $ patElemName pe) + hoisted _ Right{} = return $ Left Nothing + + reshapeBodyResults body rets = insertStmsM $ do + ses <- bodyBind body + let (ctx_ses, val_ses) = splitFromEnd (length rets) ses + resultBodyM . (ctx_ses++) =<< zipWithM reshapeResult val_ses rets + reshapeResult (Var v) t@Array{} = do + v_t <- lookupType v + let newshape = arrayDims $ removeExistentials t v_t + if newshape /= arrayDims v_t + then letSubExp "branch_ctx_reshaped" $ shapeCoerce newshape v + else return $ Var v + reshapeResult se _ = + return se + +simplifyIdentityReshape :: SimpleRule lore +simplifyIdentityReshape _ seType (Reshape newshape v) + | Just t <- seType $ Var v, + newDims newshape == arrayDims t = -- No-op reshape. + subExpRes $ Var v +simplifyIdentityReshape _ _ _ = Nothing + +simplifyReshapeReshape :: SimpleRule lore +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 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 defOf seType (Reshape newshape v) + | Just (BasicOp (Replicate _ se), v_cs) <- defOf v, + Just oldshape <- arrayShape <$> seType se, + shapeDims oldshape `isSuffixOf` newDims newshape = + let new = take (length newshape - shapeRank oldshape) $ + newDims newshape + in Just (Replicate (Shape new) se, v_cs) +simplifyReshapeReplicate _ _ _ = Nothing + +simplifyReshapeIota :: SimpleRule lore +simplifyReshapeIota defOf _ (Reshape newshape v) + | Just (BasicOp (Iota _ offset stride it), v_cs) <- defOf v, + [n] <- newDims newshape = + Just (Iota n offset stride it, v_cs) +simplifyReshapeIota _ _ _ = Nothing + +improveReshape :: SimpleRule lore +improveReshape _ seType (Reshape newshape v) + | Just t <- seType $ Var v, + newshape' <- informReshape (arrayDims t) newshape, + newshape' /= newshape = + Just (Reshape newshape' v, mempty) +improveReshape _ _ _ = Nothing + +-- | If we are copying a scratch array (possibly indirectly), just turn it into a scratch by +-- itself. +copyScratchToScratch :: SimpleRule lore +copyScratchToScratch defOf seType (Copy src) = do + t <- seType $ Var src + if isActuallyScratch src then + Just (Scratch (elemType t) (arrayDims t), mempty) + else Nothing + where isActuallyScratch v = + case asBasicOp . fst =<< defOf v of + Just Scratch{} -> True + Just (Rearrange _ v') -> isActuallyScratch v' + Just (Reshape _ v') -> isActuallyScratch v' + _ -> False +copyScratchToScratch _ _ _ = + Nothing + +ruleBasicOp :: BinderOps lore => TopDownRuleBasicOp lore + +-- Check all the simpleRules. +ruleBasicOp vtable pat aux op + | Just (op', cs) <- msum [ rule defOf seType op | rule <- simpleRules ] = + certifying (cs <> stmAuxCerts aux) $ letBind_ pat $ BasicOp op' + where defOf = (`ST.lookupExp` vtable) + seType (Var v) = ST.lookupType v vtable + seType (Constant v) = Just $ Prim $ primValueType v + +ruleBasicOp vtable pat _ (Update src _ (Var v)) + | Just (BasicOp Scratch{}, _) <- ST.lookupExp v vtable = + letBind_ pat $ BasicOp $ SubExp $ Var src + +ruleBasicOp vtable pat _ (Update dest destis (Var v)) + | Just (e, _) <- ST.lookupExp v vtable, + arrayFrom e = + letBind_ pat $ BasicOp $ SubExp $ Var dest + where arrayFrom (BasicOp (Copy copy_v)) + | Just (e',_) <- ST.lookupExp copy_v vtable = + arrayFrom e' + arrayFrom (BasicOp (Index src srcis)) = + src == dest && destis == srcis + arrayFrom (BasicOp (Replicate v_shape v_se)) + | Just (Replicate dest_shape dest_se, _) <- ST.lookupBasicOp dest vtable, + v_se == dest_se, + shapeDims v_shape `isSuffixOf` shapeDims dest_shape = + True + arrayFrom _ = + False + +-- | Turn in-place updates that replace an entire array into just +-- array literals. +ruleBasicOp vtable pat _ (Update dest is se) + | Just dest_t <- ST.lookupType dest vtable, + isFullSlice (arrayShape dest_t) is = + letBind_ pat $ BasicOp $ + case se of + Var v | not $ null $ sliceDims is -> + Reshape (map DimNew $ arrayDims dest_t) v + _ -> ArrayLit [se] $ rowType dest_t + +-- | Simplify a chain of in-place updates and copies. This chain is +-- often produced by in-place lowering. +ruleBasicOp vtable pat (StmAux cs1 _) (Update dest1 is1 (Var v1)) + | Just (Update dest2 is2 se2, cs2) <- ST.lookupBasicOp v1 vtable, + Just (Copy v3, cs3) <- ST.lookupBasicOp dest2 vtable, + Just (Index v4 is4, cs4) <- ST.lookupBasicOp v3 vtable, + is4 == is1, v4 == dest1 = certifying (cs1 <> cs2 <> cs3 <> cs4) $ do + is5 <- sliceSlice is1 is2 + letBind_ pat $ BasicOp $ Update dest1 is5 se2 + +-- | If we are comparing X against the result of a branch of the form +-- @if P then Y else Z@ then replace comparison with '(P && X == Y) || +-- (!P && X == Z'). This may allow us to get rid of a branch, and the +-- extra comparisons may be constant-folded out. Question: maybe we +-- should have some more checks to ensure that we only do this if that +-- is actually the case, such as if we will obtain at least one +-- constant-to-constant comparison? +ruleBasicOp vtable pat _ (CmpOp (CmpEq t) se1 se2) + | Just m <- simplifyWith se1 se2 = m + | Just m <- simplifyWith se2 se1 = m + where simplifyWith (Var v) x + | Just bnd <- ST.entryStm =<< ST.lookup v vtable, + If p tbranch fbranch _ <- stmExp bnd, + Just (y, z) <- + returns v (stmPattern bnd) tbranch fbranch, + S.null $ freeIn y `S.intersection` boundInBody tbranch, + S.null $ freeIn z `S.intersection` boundInBody fbranch = Just $ do + eq_x_y <- + letSubExp "eq_x_y" $ BasicOp $ CmpOp (CmpEq t) x y + eq_x_z <- + letSubExp "eq_x_z" $ BasicOp $ CmpOp (CmpEq t) x z + p_and_eq_x_y <- + letSubExp "p_and_eq_x_y" $ BasicOp $ BinOp LogAnd p eq_x_y + not_p <- + letSubExp "not_p" $ BasicOp $ UnOp Not p + not_p_and_eq_x_z <- + letSubExp "p_and_eq_x_y" $ BasicOp $ BinOp LogAnd not_p eq_x_z + letBind_ pat $ + BasicOp $ BinOp LogOr p_and_eq_x_y not_p_and_eq_x_z + simplifyWith _ _ = + Nothing + + returns v ifpat tbranch fbranch = + fmap snd $ + find ((==v) . patElemName . fst) $ + zip (patternValueElements ifpat) $ + zip (bodyResult tbranch) (bodyResult fbranch) + +ruleBasicOp _ pat _ (Replicate (Shape []) se@Constant{}) = + letBind_ pat $ BasicOp $ SubExp se +ruleBasicOp _ pat _ (Replicate (Shape []) (Var v)) = do + v_t <- lookupType v + letBind_ pat $ BasicOp $ if primType v_t + then SubExp $ Var v + else Copy v +ruleBasicOp vtable pat _ (Replicate shape (Var v)) + | Just (BasicOp (Replicate shape2 se), cs) <- ST.lookupExp v vtable = + certifying cs $ letBind_ pat $ BasicOp $ Replicate (shape<>shape2) se + +-- | Turn array literals with identical elements into replicates. +ruleBasicOp _ pat _ (ArrayLit (se:ses) _) + | all (==se) ses = + let n = constant (fromIntegral (length ses) + 1 :: Int32) + in letBind_ pat $ BasicOp $ Replicate (Shape [n]) se + +ruleBasicOp vtable pat (StmAux cs _) (Index idd slice) + | Just inds <- sliceIndices slice, + Just (BasicOp (Reshape newshape idd2), idd_cs) <- ST.lookupExp idd vtable, + length newshape == length inds = + case shapeCoercion newshape of + Just _ -> + certifying (cs<>idd_cs) $ + letBind_ pat $ BasicOp $ Index idd2 slice + Nothing -> do + -- Linearise indices and map to old index space. + oldshape <- arrayDims <$> lookupType idd2 + let new_inds = + reshapeIndex (map (primExpFromSubExp int32) oldshape) + (map (primExpFromSubExp int32) $ newDims newshape) + (map (primExpFromSubExp int32) inds) + new_inds' <- + mapM (letSubExp "new_index" <=< toExp . asInt32PrimExp) new_inds + certifying (cs<>idd_cs) $ + letBind_ pat $ BasicOp $ Index idd2 $ map DimFix new_inds' + +ruleBasicOp _ pat _ (BinOp (Pow t) e1 e2) + | e1 == intConst t 2 = + letBind_ pat $ BasicOp $ BinOp (Shl t) (intConst t 1) e2 + +-- Handle identity permutation. +ruleBasicOp _ pat _ (Rearrange perm v) + | sort perm == perm = + letBind_ pat $ BasicOp $ SubExp $ Var v + +ruleBasicOp vtable pat (StmAux cs _) (Rearrange perm v) + | Just (BasicOp (Rearrange perm2 e), v_cs) <- ST.lookupExp v vtable = + -- Rearranging a rearranging: compose the permutations. + certifying (cs<>v_cs) $ + letBind_ pat $ BasicOp $ Rearrange (perm `rearrangeCompose` perm2) e + +ruleBasicOp vtable pat (StmAux cs _) (Rearrange perm v) + | Just (BasicOp (Rotate offsets v2), v_cs) <- ST.lookupExp v vtable, + Just (BasicOp (Rearrange perm3 v3), v2_cs) <- ST.lookupExp v2 vtable = do + let offsets' = rearrangeShape (rearrangeInverse perm3) offsets + rearrange_rotate <- letExp "rearrange_rotate" $ BasicOp $ Rotate offsets' v3 + certifying (cs<>v_cs<>v2_cs) $ + letBind_ pat $ BasicOp $ Rearrange (perm `rearrangeCompose` perm3) rearrange_rotate + +-- Rearranging a replicate where the outer dimension is left untouched. +ruleBasicOp vtable pat (StmAux cs _) (Rearrange perm v1) + | Just (BasicOp (Replicate dims (Var v2)), v1_cs) <- ST.lookupExp v1 vtable, + num_dims <- shapeRank dims, + (rep_perm, rest_perm) <- splitAt num_dims perm, + not $ null rest_perm, + rep_perm == [0..length rep_perm-1] = certifying (cs<>v1_cs) $ do + v <- letSubExp "rearrange_replicate" $ + BasicOp $ Rearrange (map (subtract num_dims) rest_perm) v2 + letBind_ pat $ BasicOp $ Replicate dims v + +-- A zero-rotation is identity. +ruleBasicOp _ pat _ (Rotate offsets v) + | all isCt0 offsets = letBind_ pat $ BasicOp $ SubExp $ Var v + +ruleBasicOp vtable pat (StmAux cs _) (Rotate offsets v) + | Just (BasicOp (Rearrange perm v2), v_cs) <- ST.lookupExp v vtable, + Just (BasicOp (Rotate offsets2 v3), v2_cs) <- ST.lookupExp v2 vtable = do + let offsets2' = rearrangeShape (rearrangeInverse perm) offsets2 + addOffsets x y = letSubExp "summed_offset" $ BasicOp $ BinOp (Add Int32) x y + offsets' <- zipWithM addOffsets offsets offsets2' + rotate_rearrange <- + certifying cs $ letExp "rotate_rearrange" $ BasicOp $ Rearrange perm v3 + certifying (v_cs <> v2_cs) $ + letBind_ pat $ BasicOp $ Rotate offsets' rotate_rearrange + +-- Combining Rotates. +ruleBasicOp vtable pat (StmAux cs _) (Rotate offsets1 v) + | Just (BasicOp (Rotate offsets2 v2), v_cs) <- ST.lookupExp v vtable = do + offsets <- zipWithM add offsets1 offsets2 + certifying (cs<>v_cs) $ + letBind_ pat $ BasicOp $ Rotate offsets v2 + where add x y = letSubExp "offset" $ BasicOp $ BinOp (Add Int32) x y + +ruleBasicOp _ _ _ _ = + cannotSimplify + +-- | Remove the return values of a branch, that are not actually used +-- 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 (_, used) pat _ (e1, tb, fb, IfAttr rettype ifsort) + | -- Only if there is no existential context... + patternSize pat == length rettype, + -- Figure out which of the names in 'pat' are used... + patused <- map (`UT.isUsedDirectly` used) $ patternNames pat, + -- If they are not all used, then this rule applies. + not (and patused) = + -- Remove the parts of the branch-results that correspond to dead + -- return value bindings. Note that this leaves dead code in the + -- branch bodies, but that will be removed later. + let tses = bodyResult tb + fses = bodyResult fb + pick :: [a] -> [a] + pick = map snd . filter fst . zip patused + tb' = tb { bodyResult = pick tses } + fb' = fb { bodyResult = pick fses } + pat' = pick $ patternElements pat + rettype' = pick rettype + in letBind_ (Pattern [] pat') $ If e1 tb' fb' $ IfAttr rettype' ifsort + | otherwise = cannotSimplify + + +-- Some helper functions + +isCt1 :: SubExp -> Bool +isCt1 (Constant v) = oneIsh v +isCt1 _ = False + +isCt0 :: SubExp -> Bool +isCt0 (Constant v) = zeroIsh v +isCt0 _ = False diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/TileLoops.hs @@ -0,0 +1,385 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | Perform a restricted form of loop tiling within kernel streams. +-- We only tile primitive types, to avoid excessive local memory use. +module Futhark.Optimise.TileLoops + ( tileLoops ) + where + +import Control.Applicative +import Control.Monad.State +import Control.Monad.Reader +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import Data.Semigroup ((<>)) +import Data.List +import Data.Maybe + +import Futhark.MonadFreshNames +import Futhark.Representation.Kernels + +import Futhark.Pass +import Futhark.Tools +import Futhark.Util (mapAccumLM) + +tileLoops :: Pass Kernels Kernels +tileLoops = Pass "tile loops" "Tile stream loops inside kernels" $ + intraproceduralTransformation optimiseFunDef + +optimiseFunDef :: MonadFreshNames m => FunDef Kernels -> m (FunDef Kernels) +optimiseFunDef fundec = do + body' <- modifyNameSource $ runState $ + runReaderT m (scopeOfFParams (funDefParams fundec)) + return fundec { funDefBody = body' } + where m = optimiseBody $ funDefBody fundec + +type TileM = ReaderT (Scope Kernels) (State VNameSource) + +optimiseBody :: Body Kernels -> TileM (Body Kernels) +optimiseBody (Body () bnds res) = + Body () <$> (mconcat <$> mapM optimiseStm (stmsToList bnds)) <*> pure res + +optimiseStm :: Stm Kernels -> TileM (Stms Kernels) +optimiseStm (Let pat aux (Op old_kernel@(Kernel desc space ts body))) = do + (extra_bnds, space', body') <- tileInKernelBody mempty initial_variance space body + let new_kernel = Kernel desc space' ts body' + -- XXX: we should not change the type of the kernel (such as by + -- changing the number of groups being used for a kernel that + -- returns a result-per-group). + if kernelType old_kernel == kernelType new_kernel + then return $ extra_bnds <> oneStm (Let pat aux $ Op new_kernel) + else return $ oneStm $ Let pat aux $ Op old_kernel + where initial_variance = M.map mempty $ scopeOfKernelSpace space +optimiseStm (Let pat aux e) = + pure <$> (Let pat aux <$> mapExpM optimise e) + where optimise = identityMapper { mapOnBody = const optimiseBody } + +tileInKernelBody :: Names -> VarianceTable + -> KernelSpace -> KernelBody InKernel + -> TileM (Stms Kernels, KernelSpace, KernelBody InKernel) +tileInKernelBody branch_variant initial_variance initial_kspace (KernelBody () kstms kres) = do + (extra_bnds, kspace', kstms') <- + tileInStms branch_variant initial_variance initial_kspace kstms + return (extra_bnds, kspace', KernelBody () kstms' kres) + +tileInBody :: Names -> VarianceTable + -> KernelSpace -> Body InKernel + -> TileM (Stms Kernels, KernelSpace, Body InKernel) +tileInBody branch_variant initial_variance initial_kspace (Body () stms res) = do + (extra_bnds, kspace', stms') <- + tileInStms branch_variant initial_variance initial_kspace stms + return (extra_bnds, kspace', Body () stms' res) + +tileInStms :: Names -> VarianceTable + -> KernelSpace -> Stms InKernel + -> TileM (Stms Kernels, KernelSpace, Stms InKernel) +tileInStms branch_variant initial_variance initial_kspace kstms = do + ((kspace, extra_bndss), kstms') <- + mapAccumLM tileInKernelStatement (initial_kspace,mempty) $ stmsToList kstms + return (extra_bndss, kspace, stmsFromList kstms') + where variance = varianceInStms initial_variance kstms + + tileInKernelStatement (kspace, extra_bnds) + (Let pat attr (Op (GroupStream w max_chunk lam accs arrs))) + | max_chunk == w, + not $ null arrs, + chunk_size <- Var $ groupStreamChunkSize lam, + arr_chunk_params <- groupStreamArrParams lam, + maybe_1d_tiles <- + zipWith (is1dTileable branch_variant kspace variance chunk_size) arrs arr_chunk_params, + maybe_1_5d_tiles <- + zipWith (is1_5dTileable branch_variant kspace variance chunk_size) arrs arr_chunk_params, + Just mk_tilings <- + zipWithM (<|>) maybe_1d_tiles maybe_1_5d_tiles = do + + (kspaces, arr_chunk_params', tile_kstms) <- unzip3 <$> sequence mk_tilings + + let (kspace', kspace_bnds) = + case kspaces of + [] -> (kspace, mempty) + new_kspace : _ -> new_kspace + Body () lam_kstms lam_res <- syncAtEnd $ groupStreamLambdaBody lam + let lam_kstms' = mconcat tile_kstms <> lam_kstms + group_size = spaceGroupSize kspace + lam' = lam { groupStreamLambdaBody = Body () lam_kstms' lam_res + , groupStreamArrParams = arr_chunk_params' + } + + return ((kspace', extra_bnds <> kspace_bnds), + Let pat attr $ Op $ GroupStream w group_size lam' accs arrs) + + tileInKernelStatement (kspace, extra_bnds) + (Let pat attr (Op (GroupStream w max_chunk lam accs arrs))) + | w == max_chunk, + not $ null arrs, + FlatThreadSpace gspace <- spaceStructure kspace, + chunk_size <- Var $ groupStreamChunkSize lam, + arr_chunk_params <- groupStreamArrParams lam, + + Just mk_tilings <- + zipWithM (is2dTileable branch_variant kspace variance chunk_size) + arrs arr_chunk_params = do + + ((tile_size, tiled_group_size), tile_size_bnds) <- runBinder $ do + tile_size_key <- newVName "tile_size" + tile_size <- letSubExp "tile_size" $ Op $ GetSize tile_size_key SizeTile + tiled_group_size <- letSubExp "tiled_group_size" $ + BasicOp $ BinOp (Mul Int32) tile_size tile_size + return (tile_size, tiled_group_size) + + let (tiled_gspace,untiled_gspace) = splitAt 2 $ reverse gspace + -- Play with reversion to ensure we get increasing IDs for + -- ltids. This affects readability of generated code. + untiled_gspace' <- fmap reverse $ forM (reverse untiled_gspace) $ \(gtid,gdim) -> do + ltid <- newVName "ltid" + return (gtid,gdim, + ltid, constant (1::Int32)) + tiled_gspace' <- fmap reverse $ forM (reverse tiled_gspace) $ \(gtid,gdim) -> do + ltid <- newVName "ltid" + return (gtid,gdim, + ltid, tile_size) + let gspace' = reverse $ tiled_gspace' ++ untiled_gspace' + + -- We have to recalculate number of workgroups and + -- number of threads to fit the new workgroup size. + ((num_threads, num_groups), num_bnds) <- + runBinder $ sufficientGroups gspace' tiled_group_size + + let kspace' = kspace { spaceStructure = NestedThreadSpace gspace' + , spaceGroupSize = tiled_group_size + , spaceNumThreads = num_threads + , spaceNumGroups = num_groups + } + local_ids = map (\(_, _, ltid, _) -> ltid) gspace' + + (arr_chunk_params', tile_kstms) <- + fmap unzip $ forM mk_tilings $ \mk_tiling -> + mk_tiling tile_size local_ids + + Body () lam_kstms lam_res <- syncAtEnd $ groupStreamLambdaBody lam + let lam_kstms' = mconcat tile_kstms <> lam_kstms + lam' = lam { groupStreamLambdaBody = Body () lam_kstms' lam_res + , groupStreamArrParams = arr_chunk_params' + } + + return ((kspace', extra_bnds <> tile_size_bnds <> num_bnds), + Let pat attr $ Op $ GroupStream w tile_size lam' accs arrs) + + tileInKernelStatement (kspace, extra_bnds) + (Let pat attr (Op (GroupStream w maxchunk lam accs arrs))) = do + let branch_variant' = branch_variant <> + fromMaybe mempty (flip M.lookup variance =<< subExpVar w) + (bnds, kspace', lam') <- tileInStreamLambda branch_variant' variance kspace lam + return ((kspace', extra_bnds <> bnds), + Let pat attr $ Op $ GroupStream w maxchunk lam' accs arrs) + + tileInKernelStatement acc stm = + return (acc, stm) + +tileInStreamLambda :: Names -> VarianceTable -> KernelSpace -> GroupStreamLambda InKernel + -> TileM (Stms Kernels, KernelSpace, GroupStreamLambda InKernel) +tileInStreamLambda branch_variant variance kspace lam = do + (bnds, kspace', kbody') <- + tileInBody branch_variant variance' kspace $ groupStreamLambdaBody lam + return (bnds, kspace', lam { groupStreamLambdaBody = kbody' }) + where variance' = varianceInStms variance $ + bodyStms $ groupStreamLambdaBody lam + +is1dTileable :: MonadFreshNames m => + Names -> KernelSpace -> VarianceTable -> SubExp -> VName -> LParam InKernel + -> Maybe (m ((KernelSpace, Stms Kernels), + LParam InKernel, + Stms InKernel)) +is1dTileable branch_variant kspace variance block_size arr block_param = do + guard $ S.null $ M.findWithDefault mempty arr variance + guard $ S.null branch_variant + guard $ primType $ rowType $ paramType block_param + + return $ do + (outer_block_param, kstms) <- tile1d kspace block_size block_param + return ((kspace, mempty), outer_block_param, kstms) + +is1_5dTileable :: (MonadFreshNames m, HasScope Kernels m) => + Names -> KernelSpace -> VarianceTable + -> SubExp -> VName -> LParam InKernel + -> Maybe (m ((KernelSpace, Stms Kernels), + LParam InKernel, + Stms InKernel)) +is1_5dTileable branch_variant kspace variance block_size arr block_param = do + guard $ primType $ rowType $ paramType block_param + + (inner_gtid, inner_gdim) <- invariantToInnermostDimension + mk_structure <- + case spaceStructure kspace of + NestedThreadSpace{} -> Nothing + FlatThreadSpace gtids_and_gdims -> + return $ do + -- Force a functioning group size. XXX: not pretty. + let n_dims = length gtids_and_gdims + outer <- forM (take (n_dims-1) gtids_and_gdims) $ \(gtid, gdim) -> do + ltid <- newVName "ltid" + return (gtid, gdim, ltid, gdim) + + inner_ltid <- newVName "inner_ltid" + inner_ldim <- newVName "inner_ldim" + let compute_tiled_group_size = + mkLet [] [Ident inner_ldim $ Prim int32] $ + BasicOp $ BinOp (SMin Int32) (spaceGroupSize kspace) inner_gdim + structure = NestedThreadSpace $ outer ++ [(inner_gtid, inner_gdim, + inner_ltid, Var inner_ldim)] + ((num_threads, num_groups), num_bnds) <- runBinder $ do + threads_necessary <- + letSubExp "threads_necessary" =<< + foldBinOp (Mul Int32) + (constant (1::Int32)) (map snd gtids_and_gdims) + groups_necessary <- + letSubExp "groups_necessary" =<< + eDivRoundingUp Int32 (eSubExp threads_necessary) (eSubExp $ Var inner_ldim) + num_threads <- + letSubExp "num_threads" $ + BasicOp $ BinOp (Mul Int32) groups_necessary (Var inner_ldim) + return (num_threads, groups_necessary) + + let kspace' = kspace { spaceGroupSize = Var inner_ldim + , spaceNumGroups = num_groups + , spaceNumThreads = num_threads + , spaceStructure = structure + } + return (oneStm compute_tiled_group_size <> num_bnds, + kspace') + return $ do + (outer_block_param, kstms) <- tile1d kspace block_size block_param + (structure_bnds, kspace') <- mk_structure + return ((kspace', structure_bnds), outer_block_param, kstms) + where invariantToInnermostDimension :: Maybe (VName, SubExp) + invariantToInnermostDimension = + case reverse $ spaceDimensions kspace of + (i,d) : _ + | not $ i `S.member` M.findWithDefault mempty arr variance, + not $ i `S.member` branch_variant -> Just (i,d) + _ -> Nothing + +tile1d :: MonadFreshNames m => + KernelSpace + -> SubExp + -> LParam InKernel + -> m (LParam InKernel, Stms InKernel) +tile1d kspace block_size block_param = do + outer_block_param <- do + name <- newVName $ baseString (paramName block_param) ++ "_outer" + return block_param { paramName = name } + + let ltid = spaceLocalId kspace + read_elem_bnd <- do + name <- newVName $ baseString (paramName outer_block_param) ++ "_elem" + return $ + mkLet [] [Ident name $ rowType $ paramType outer_block_param] $ + BasicOp $ Index (paramName outer_block_param) [DimFix $ Var ltid] + + cid <- newVName "cid" + let block_cspace = combineSpace [(cid, block_size)] + block_pe = + PatElem (paramName block_param) $ paramType outer_block_param + write_block_stms = + [ Let (Pattern [] [block_pe]) (defAux ()) $ Op $ + Combine block_cspace [patElemType pe] [] $ + Body () (oneStm read_elem_bnd) [Var $ patElemName pe] + | pe <- patternElements $ stmPattern read_elem_bnd ] + + return (outer_block_param, stmsFromList write_block_stms) + +is2dTileable :: MonadFreshNames m => + Names -> KernelSpace -> VarianceTable -> SubExp -> VName -> LParam InKernel + -> Maybe (SubExp -> [VName] -> m (LParam InKernel, Stms InKernel)) +is2dTileable branch_variant kspace variance block_size arr block_param = do + guard $ primType $ rowType $ paramType block_param + + pt <- case rowType $ paramType block_param of + Prim pt -> return pt + _ -> Nothing + inner_perm <- invariantToOneOfTwoInnerDims + Just $ \tile_size local_is -> do + let num_outer = length local_is - 2 + perm = [0..num_outer-1] ++ map (+num_outer) inner_perm + invariant_i : variant_i : _ = reverse $ rearrangeShape perm local_is + (global_i,global_d):_ = rearrangeShape inner_perm $ drop num_outer $ spaceDimensions kspace + outer_block_param <- do + name <- newVName $ baseString (paramName block_param) ++ "_outer" + return block_param { paramName = name } + + elem_name <- newVName $ baseString (paramName outer_block_param) ++ "_elem" + let read_elem_bnd = mkLet [] [Ident elem_name $ Prim pt] $ + BasicOp $ Index (paramName outer_block_param) $ + fullSlice (paramType outer_block_param) [DimFix $ Var invariant_i] + + cids <- replicateM (length local_is - num_outer) $ newVName "cid" + let block_size_2d = Shape $ rearrangeShape inner_perm [tile_size, block_size] + block_cspace = combineSpace $ zip cids $ + rearrangeShape inner_perm [tile_size,block_size] + + block_name_2d <- newVName $ baseString (paramName block_param) ++ "_2d" + let block_pe = + PatElem block_name_2d $ + rowType (paramType outer_block_param) `arrayOfShape` block_size_2d + write_block_stm = + Let (Pattern [] [block_pe]) (defAux ()) $ + Op $ Combine block_cspace [Prim pt] [(global_i, global_d)] $ + Body () (oneStm read_elem_bnd) [Var elem_name] + + let index_block_kstms = + [mkLet [] [paramIdent block_param] $ + BasicOp $ Index block_name_2d $ + rearrangeShape inner_perm $ + fullSlice (rearrangeType inner_perm $ patElemType block_pe) + [DimFix $ Var variant_i]] + + return (outer_block_param, + oneStm write_block_stm <> stmsFromList index_block_kstms) + + where invariantToOneOfTwoInnerDims :: Maybe [Int] + invariantToOneOfTwoInnerDims = do + (j,_) : (i,_) : _ <- Just $ reverse $ spaceDimensions kspace + let variant_to = M.findWithDefault mempty arr variance + branch_invariant = not $ S.member j branch_variant || S.member i branch_variant + if branch_invariant && i `S.member` variant_to && not (j `S.member` variant_to) then + Just [0,1] + else if branch_invariant && j `S.member` variant_to && not (i `S.member` variant_to) then + Just [1,0] + else + Nothing + +syncAtEnd :: MonadFreshNames m => Body InKernel -> m (Body InKernel) +syncAtEnd (Body () stms res) = do + (res', stms') <- (`runBinderT` mempty) $ do + mapM_ addStm stms + map Var <$> letTupExp "sync" (Op $ Barrier res) + return $ Body () stms' res' + +-- | The variance table keeps a mapping from a variable name +-- (something produced by a 'Stm') to the kernel thread indices +-- that name depends on. If a variable is not present in this table, +-- that means it is bound outside the kernel (and so can be considered +-- invariant to all dimensions). +type VarianceTable = M.Map VName Names + +varianceInStms :: VarianceTable -> Stms InKernel -> VarianceTable +varianceInStms = foldl varianceInStm + +varianceInStm :: VarianceTable -> Stm InKernel -> VarianceTable +varianceInStm variance bnd = + foldl' add variance $ patternNames $ stmPattern bnd + where add variance' v = M.insert v binding_variance variance' + look variance' v = S.insert v $ M.findWithDefault mempty v variance' + binding_variance = mconcat $ map (look variance) $ S.toList (freeInStm bnd) + +sufficientGroups :: MonadBinder m => + [(VName, SubExp, VName, SubExp)] -> SubExp + -> m (SubExp, SubExp) +sufficientGroups gspace group_size = do + groups_in_dims <- forM gspace $ \(_, gd, _, ld) -> + letSubExp "groups_in_dim" =<< eDivRoundingUp Int32 (eSubExp gd) (eSubExp ld) + num_groups <- letSubExp "num_groups" =<< + foldBinOp (Mul Int32) (constant (1::Int32)) groups_in_dims + num_threads <- letSubExp "num_threads" $ + BasicOp $ BinOp (Mul Int32) num_groups group_size + return (num_threads, num_groups) diff --git a/src/Futhark/Optimise/Unstream.hs b/src/Futhark/Optimise/Unstream.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Unstream.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | Turn GroupStreams that operate on entire input or thread-variant +-- sizes into do-loops, thus aiding subsequent optimisation. It is +-- very important that this is run *after* any access-pattern-related +-- optimisation, because this pass will destroy information. +module Futhark.Optimise.Unstream + ( unstream ) + where + +import Control.Monad.State +import Control.Monad.Reader +import qualified Data.Set as S +import qualified Data.Map as M + +import Futhark.Representation.AST.Attributes.Aliases +import qualified Futhark.Analysis.Alias as Alias +import Futhark.MonadFreshNames +import Futhark.Representation.Kernels +import Futhark.Pass +import Futhark.Tools + +unstream :: Pass Kernels Kernels +unstream = Pass "unstream" "Remove whole-array streams in kernels" $ + intraproceduralTransformation optimiseFunDef + +optimiseFunDef :: MonadFreshNames m => FunDef Kernels -> m (FunDef Kernels) +optimiseFunDef fundec = do + body' <- modifyNameSource $ runState $ + runReaderT m (scopeOfFParams (funDefParams fundec)) + return fundec { funDefBody = body' } + where m = optimiseBody $ funDefBody fundec + +type UnstreamM = ReaderT (Scope Kernels) (State VNameSource) + +optimiseBody :: Body Kernels -> UnstreamM (Body Kernels) +optimiseBody (Body () stms res) = + localScope (scopeOf stms) $ + Body () <$> (stmsFromList . concat <$> mapM optimiseStm (stmsToList stms)) <*> pure res + +optimiseStm :: Stm Kernels -> UnstreamM [Stm Kernels] +optimiseStm (Let pat aux (Op (Kernel desc space ts body))) = do + inv <- S.fromList . M.keys <$> askScope + stms' <- localScope (scopeOfKernelSpace space) $ + runBinder_ $ optimiseInKernelStms inv $ kernelBodyStms body + return [Let pat aux $ Op $ Kernel desc space ts $ body { kernelBodyStms = stms' }] +optimiseStm (Let pat aux e) = + pure <$> (Let pat aux <$> mapExpM optimise e) + where optimise = identityMapper { mapOnBody = \scope -> localScope scope . optimiseBody } + +type Invariant = S.Set VName + +type InKernelM = Binder InKernel + +optimiseInKernelStms :: Invariant -> Stms InKernel -> InKernelM () +optimiseInKernelStms inv = mapM_ (optimiseInKernelStm inv) . stmsToList + +optimiseInKernelStm :: Invariant -> Stm InKernel -> InKernelM () +optimiseInKernelStm inv (Let pat aux (Op (GroupStream w max_chunk lam accs arrs))) + | max_chunk == w || maybe False (`S.notMember` inv) (subExpVar w) = do + let GroupStreamLambda chunk_size chunk_offset acc_params arr_params body = lam + letBindNames_ [chunk_size] $ BasicOp $ SubExp $ constant (1::Int32) + + loop_body <- insertStmsM $ do + forM_ (zip arr_params arrs) $ \(p,a) -> + letBindNames_ [paramName p] $ + BasicOp $ Index a $ fullSlice (paramType p) + [DimSlice (Var chunk_offset) (Var chunk_size) (constant (1::Int32))] + localScope (scopeOfLParams acc_params) $ optimiseInBody inv body + + -- Some accumulators may be updated in-place and must hence be unique. + let lam_consumed = consumedInBody $ Alias.analyseBody $ groupStreamLambdaBody lam + uniqueIfConsumed p | paramName p `S.member` lam_consumed = + fmap (`toDecl` Unique) p + | otherwise = fmap (`toDecl` Nonunique) p + merge = zip (map uniqueIfConsumed acc_params) accs + certifying (stmAuxCerts aux) $ + letBind_ pat $ DoLoop [] merge (ForLoop chunk_offset Int32 w []) loop_body +optimiseInKernelStm inv (Let pat aux e) = + addStm =<< (Let pat aux <$> mapExpM optimise e) + where optimise = identityMapper + { mapOnBody = \scope -> localScope scope . optimiseInBody inv } + +optimiseInBody :: Invariant -> Body InKernel -> InKernelM (Body InKernel) +optimiseInBody inv body = do + stms' <- collectStms_ $ optimiseInKernelStms inv $ bodyStms body + return body { bodyStms = stms' } diff --git a/src/Futhark/Pass.hs b/src/Futhark/Pass.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleContexts #-} +-- | Definition of a polymorphic (generic) pass that can work with programs of any +-- lore. +module Futhark.Pass + ( PassM + , runPassM + , liftEither + , liftEitherM + , Pass (..) + , passLongOption + , intraproceduralTransformation + ) where + +import Control.Monad.Writer.Strict +import Control.Monad.Except hiding (liftEither) +import Control.Monad.State.Strict +import Control.Parallel.Strategies +import Data.Char +import Data.Either + +import Prelude hiding (log) + +import Futhark.Error +import Futhark.Representation.AST +import Futhark.Util.Log +import Futhark.MonadFreshNames + +-- | The monad in which passes execute. +newtype PassM a = PassM (ExceptT InternalError (WriterT Log (State VNameSource)) a) + deriving (Functor, Applicative, Monad, + MonadError InternalError) + +instance MonadLogger PassM where + addLog = PassM . tell + +instance MonadFreshNames PassM where + putNameSource = PassM . put + getNameSource = PassM get + +-- | Execute a 'PassM' action, yielding logging information and either +-- an error text or a result. +runPassM :: MonadFreshNames m => + PassM a -> m (Either InternalError a, Log) +runPassM (PassM m) = modifyNameSource $ \src -> + runState (runWriterT $ runExceptT m) src + +-- | Turn an 'Either' computation into a 'PassM'. If the 'Either' is +-- 'Left', the result is a 'CompilerBug'. +liftEither :: Show err => Either err a -> PassM a +liftEither (Left e) = compilerBugS $ show e +liftEither (Right v) = return v + +-- | Turn an 'Either' monadic computation into a 'PassM'. If the 'Either' is +-- 'Left', the result is an exception. +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 { passName :: String + -- ^ Name of the pass. Keep this short and simple. It will + -- be used to automatically generate a command-line option + -- name via 'passLongOption'. + , passDescription :: String + -- ^ A slightly longer description, which will show up in the + -- command-line help text. + , passFunction :: Prog fromlore -> PassM (Prog tolore) + } + +-- | Take the name of the pass, turn spaces into dashes, and make all +-- characters lowercase. +passLongOption :: Pass fromlore tolore -> String +passLongOption = map (spaceToDash . toLower) . passName + where spaceToDash ' ' = '-' + spaceToDash c = c + +intraproceduralTransformation :: (FunDef fromlore -> PassM (FunDef tolore)) + -> Prog fromlore -> PassM (Prog tolore) +intraproceduralTransformation ft prog = + either onError onSuccess <=< modifyNameSource $ \src -> + case partitionEithers $ parMap rpar (onFunction src) (progFunctions prog) of + ([], rs) -> let (funs, logs, srcs) = unzip3 rs + in (Right (Prog funs, mconcat logs), mconcat srcs) + ((err,log,src'):_, _) -> (Left (err, log), src') + where onFunction src f = case runState (runPassM (ft f)) src of + ((Left x, log), src') -> Left (x, log, src') + ((Right x, log), src') -> Right (x, log, src') + + onError (err, log) = addLog log >> throwError err + onSuccess (x, log) = addLog log >> return x diff --git a/src/Futhark/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExpandAllocations.hs @@ -0,0 +1,460 @@ +{-# LANGUAGE TypeFamilies, FlexibleContexts #-} +-- | Expand allocations inside of maps when possible. +module Futhark.Pass.ExpandAllocations + ( expandAllocations ) +where + +import Control.Monad.Identity +import Control.Monad.Except +import Control.Monad.State +import Control.Monad.Reader +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import Data.List +import Data.Semigroup ((<>)) + +import Prelude hiding (quot) + +import Futhark.Analysis.Rephrase +import Futhark.Error +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Pass +import Futhark.Representation.AST +import Futhark.Representation.ExplicitMemory +import qualified Futhark.Representation.ExplicitMemory.Simplify as ExplicitMemory +import qualified Futhark.Representation.Kernels as Kernels +import Futhark.Representation.Kernels.Simplify as Kernels +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Pass.ExtractKernels.BlockedKernel (blockedReduction) +import Futhark.Pass.ExplicitAllocations (explicitAllocationsInStms) +import Futhark.Util.IntegralExp +import Futhark.Util (mapAccumLM) + +expandAllocations :: Pass ExplicitMemory ExplicitMemory +expandAllocations = + Pass "expand allocations" "Expand allocations" $ + fmap Prog . mapM transformFunDef . progFunctions + -- Cannot use intraproceduralTransformation because it might create + -- duplicate size keys (which are not fixed by renamer, and size + -- keys must currently be globally unique). + +type ExpandM = ExceptT InternalError (ReaderT (Scope ExplicitMemory) (State VNameSource)) + +transformFunDef :: FunDef ExplicitMemory -> PassM (FunDef ExplicitMemory) +transformFunDef fundec = do + body' <- either throwError return <=< modifyNameSource $ + runState $ runReaderT (runExceptT m) mempty + return fundec { funDefBody = body' } + where m = inScopeOf fundec $ transformBody $ funDefBody fundec + +transformBody :: Body ExplicitMemory -> ExpandM (Body ExplicitMemory) +transformBody (Body () stms res) = Body () <$> transformStms stms <*> pure res + +transformStms :: Stms ExplicitMemory -> ExpandM (Stms ExplicitMemory) +transformStms stms = + inScopeOf stms $ mconcat <$> mapM transformStm (stmsToList stms) + +transformStm :: Stm ExplicitMemory -> ExpandM (Stms ExplicitMemory) + +transformStm (Let pat aux e) = do + (bnds, e') <- transformExp =<< mapExpM transform e + return $ bnds <> oneStm (Let pat aux e') + where transform = identityMapper { mapOnBody = \scope -> localScope scope . transformBody + } + +transformExp :: Exp ExplicitMemory -> ExpandM (Stms ExplicitMemory, Exp ExplicitMemory) + +transformExp (Op (Inner (Kernel desc kspace ts kbody))) = do + let (kbody', allocs) = extractKernelBodyAllocations kbody + variantAlloc (Var v) = v `S.member` bound_in_kernel + variantAlloc _ = False + (variant_allocs, invariant_allocs) = M.partition (variantAlloc . fst) allocs + + num_threads64 <- newVName "num_threads64" + let num_threads64_pat = Pattern [] [PatElem num_threads64 $ MemPrim int64] + num_threads64_bnd = Let num_threads64_pat (defAux ()) $ BasicOp $ + ConvOp (SExt Int32 Int64) (spaceNumThreads kspace) + + (invariant_alloc_stms, invariant_alloc_offsets) <- + expandedInvariantAllocations + (Var num_threads64, spaceNumGroups kspace, spaceGroupSize kspace) + (spaceGlobalId kspace, spaceGroupId kspace, spaceLocalId kspace) invariant_allocs + + (variant_alloc_stms, variant_alloc_offsets) <- + expandedVariantAllocations kspace kbody variant_allocs + + let alloc_offsets = invariant_alloc_offsets <> variant_alloc_offsets + alloc_stms = invariant_alloc_stms <> variant_alloc_stms + + kbody'' <- either compilerLimitationS pure $ + offsetMemoryInKernelBody alloc_offsets + kbody' { kernelBodyStms = kernelBodyStms kbody' } + + return (oneStm num_threads64_bnd <> alloc_stms, + Op $ Inner $ Kernel desc kspace ts kbody'') + + where bound_in_kernel = + S.fromList $ M.keys $ scopeOfKernelSpace kspace <> + scopeOf (kernelBodyStms kbody) + +transformExp e = + return (mempty, e) + +-- | Extract allocations from 'Thread' statements with +-- 'extractThreadAllocations'. +extractKernelBodyAllocations :: KernelBody InKernel + -> (KernelBody InKernel, + M.Map VName (SubExp, Space)) +extractKernelBodyAllocations kbody = + let (allocs, stms) = mapAccumL extract M.empty $ stmsToList $ kernelBodyStms kbody + in (kbody { kernelBodyStms = mconcat stms }, allocs) + where extract allocs bnd = + let (bnds, body_allocs) = extractThreadAllocations $ oneStm bnd + in (allocs <> body_allocs, bnds) + +extractThreadAllocations :: Stms InKernel + -> (Stms InKernel, M.Map VName (SubExp, Space)) +extractThreadAllocations bnds = + let (allocs, bnds') = mapAccumL isAlloc M.empty $ stmsToList bnds + in (stmsFromList $ catMaybes bnds', allocs) + where isAlloc allocs (Let (Pattern [] [patElem]) _ (Op (Alloc size space))) = + (M.insert (patElemName patElem) (size, space) allocs, + Nothing) + + isAlloc allocs bnd = + (allocs, Just bnd) + +expandedInvariantAllocations :: (SubExp,SubExp, SubExp) + -> (VName, VName, VName) + -> M.Map VName (SubExp, Space) + -> ExpandM (Stms ExplicitMemory, RebaseMap) +expandedInvariantAllocations (num_threads64, num_groups, group_size) + (_thread_index, group_id, local_id) + invariant_allocs = do + -- We expand the invariant allocations by adding an inner dimension + -- equal to the number of kernel threads. + (alloc_bnds, rebases) <- unzip <$> mapM expand (M.toList invariant_allocs) + + return (mconcat alloc_bnds, mconcat rebases) + where expand (mem, (per_thread_size, Space "local")) = do + let allocpat = Pattern [] [PatElem mem $ + MemMem per_thread_size $ Space "local"] + return (oneStm $ Let allocpat (defAux ()) $ + Op $ Alloc per_thread_size $ Space "local", + mempty) + + expand (mem, (per_thread_size, space)) = do + total_size <- newVName "total_size" + let sizepat = Pattern [] [PatElem total_size $ MemPrim int64] + allocpat = Pattern [] [PatElem mem $ + MemMem (Var total_size) space] + return (stmsFromList + [Let sizepat (defAux ()) $ + BasicOp $ BinOp (Mul Int64) num_threads64 per_thread_size, + Let allocpat (defAux ()) $ + Op $ Alloc (Var total_size) space], + M.singleton mem newBase) + + newBase (old_shape, _) = + let num_dims = length old_shape + perm = [0, num_dims+1] ++ [1..num_dims] + root_ixfun = IxFun.iota (primExpFromSubExp int32 num_groups : old_shape + ++ [primExpFromSubExp int32 group_size]) + permuted_ixfun = IxFun.permute root_ixfun perm + untouched d = DimSlice (fromInt32 0) d (fromInt32 1) + offset_ixfun = IxFun.slice permuted_ixfun $ + [DimFix (LeafExp group_id int32), + DimFix (LeafExp local_id int32)] ++ + map untouched old_shape + in offset_ixfun + +expandedVariantAllocations :: KernelSpace -> KernelBody InKernel + -> M.Map VName (SubExp, Space) + -> ExpandM (Stms ExplicitMemory, RebaseMap) +expandedVariantAllocations _ _ variant_allocs + | null variant_allocs = return (mempty, mempty) +expandedVariantAllocations kspace kbody variant_allocs = do + let sizes_to_blocks = removeCommonSizes variant_allocs + variant_sizes = map fst sizes_to_blocks + + (slice_stms, offsets, size_sums) <- + sliceKernelSizes variant_sizes kspace kbody + -- Note the recursive call to expand allocations inside the newly + -- produced kernels. + slice_stms_tmp <- ExplicitMemory.simplifyStms =<< explicitAllocationsInStms slice_stms + slice_stms' <- transformStms slice_stms_tmp + + let variant_allocs' :: [(VName, (SubExp, SubExp, Space))] + variant_allocs' = concat $ zipWith memInfo (map snd sizes_to_blocks) + (zip offsets size_sums) + memInfo blocks (offset, total_size) = + [ (mem, (Var offset, Var total_size, space)) | (mem, space) <- blocks ] + + -- We expand the invariant allocations by adding an inner dimension + -- equal to the sum of the sizes required by different threads. + (alloc_bnds, rebases) <- unzip <$> mapM expand variant_allocs' + + return (slice_stms' <> stmsFromList alloc_bnds, mconcat rebases) + where expand (mem, (offset, total_size, space)) = do + let allocpat = Pattern [] [PatElem mem $ + MemMem total_size space] + return (Let allocpat (defAux ()) $ Op $ Alloc total_size space, + M.singleton mem $ newBase offset) + + num_threads = primExpFromSubExp int32 $ spaceNumThreads kspace + gtid = LeafExp (spaceGlobalId kspace) int32 + + -- For the variant allocations, we add an inner dimension, + -- which is then offset by a thread-specific amount. + newBase size_per_thread (old_shape, pt) = + let pt_size = fromInt32 $ primByteSize pt + elems_per_thread = ConvOpExp (SExt Int64 Int32) + (primExpFromSubExp int64 size_per_thread) + `quot` pt_size + root_ixfun = IxFun.iota [elems_per_thread, num_threads] + offset_ixfun = IxFun.slice root_ixfun + [DimSlice (fromInt32 0) num_threads (fromInt32 1), + DimFix gtid] + shapechange = if length old_shape == 1 + then map DimCoercion old_shape + else map DimNew old_shape + in IxFun.reshape offset_ixfun shapechange + +-- | A map from memory block names to new index function bases. + +type RebaseMap = M.Map VName (([PrimExp VName], PrimType) -> IxFun) + +lookupNewBase :: VName -> ([PrimExp VName], PrimType) -> RebaseMap -> Maybe IxFun +lookupNewBase name x = fmap ($ x) . M.lookup name + +offsetMemoryInKernelBody :: RebaseMap -> KernelBody InKernel + -> Either String (KernelBody InKernel) +offsetMemoryInKernelBody initial_offsets kbody = do + stms' <- snd <$> mapAccumLM offsetMemoryInStm initial_offsets + (stmsToList $ kernelBodyStms kbody) + return kbody { kernelBodyStms = stmsFromList stms' } + +offsetMemoryInBody :: RebaseMap -> Body InKernel -> Either String (Body InKernel) +offsetMemoryInBody offsets (Body attr stms res) = do + stms' <- stmsFromList . snd <$> mapAccumLM offsetMemoryInStm offsets (stmsToList stms) + return $ Body attr stms' res + +offsetMemoryInStm :: RebaseMap -> Stm InKernel + -> Either String (RebaseMap, Stm InKernel) +offsetMemoryInStm offsets (Let pat attr e) = do + (offsets', pat') <- offsetMemoryInPattern offsets pat + e' <- offsetMemoryInExp offsets e + return (offsets', Let pat' attr e') + +offsetMemoryInPattern :: RebaseMap -> Pattern InKernel + -> Either String (RebaseMap, Pattern InKernel) +offsetMemoryInPattern offsets (Pattern ctx vals) = do + offsets' <- foldM inspectCtx offsets ctx + return (offsets', Pattern ctx $ map (inspectVal offsets') vals) + where inspectVal offsets' = fmap $ offsetMemoryInMemBound offsets' + inspectCtx ctx_offsets patElem + | Mem _ space <- patElemType patElem, + space /= Space "local" = + throwError $ unwords ["Cannot deal with existential memory block", + pretty (patElemName patElem), + "when expanding inside kernels."] + | otherwise = + return ctx_offsets + +offsetMemoryInParam :: RebaseMap -> Param (MemBound u) -> Param (MemBound u) +offsetMemoryInParam offsets fparam = + fparam { paramAttr = offsetMemoryInMemBound offsets $ paramAttr fparam } + +offsetMemoryInMemBound :: RebaseMap -> MemBound u -> MemBound u +offsetMemoryInMemBound offsets (MemArray pt shape u (ArrayIn mem ixfun)) + | Just new_base <- lookupNewBase mem (IxFun.base ixfun, pt) offsets = + MemArray pt shape u $ ArrayIn mem $ IxFun.rebase new_base ixfun +offsetMemoryInMemBound _ summary = + summary + +offsetMemoryInBodyReturns :: RebaseMap -> BodyReturns -> BodyReturns +offsetMemoryInBodyReturns offsets (MemArray pt shape u (ReturnsInBlock mem ixfun)) + | Just ixfun' <- isStaticIxFun ixfun, + Just new_base <- lookupNewBase mem (IxFun.base ixfun', pt) offsets = + MemArray pt shape u $ ReturnsInBlock mem $ + IxFun.rebase (fmap (fmap Free) new_base) ixfun +offsetMemoryInBodyReturns _ br = br + +offsetMemoryInExp :: RebaseMap -> Exp InKernel -> Either String (Exp InKernel) +offsetMemoryInExp offsets (DoLoop ctx val form body) = + DoLoop (zip ctxparams' ctxinit) (zip valparams' valinit) form <$> + offsetMemoryInBody offsets body + where (ctxparams, ctxinit) = unzip ctx + (valparams, valinit) = unzip val + ctxparams' = map (offsetMemoryInParam offsets) ctxparams + valparams' = map (offsetMemoryInParam offsets) valparams +offsetMemoryInExp offsets (Op (Inner (GroupStream w max_chunk lam accs arrs))) = do + body <- offsetMemoryInBody offsets $ groupStreamLambdaBody lam + let lam' = lam { groupStreamLambdaBody = body + , groupStreamAccParams = map (offsetMemoryInParam offsets) $ + groupStreamAccParams lam + , groupStreamArrParams = map (offsetMemoryInParam offsets) $ + groupStreamArrParams lam + } + return $ Op $ Inner $ GroupStream w max_chunk lam' accs arrs +offsetMemoryInExp offsets (Op (Inner (GroupReduce w lam input))) = do + body <- offsetMemoryInBody offsets $ lambdaBody lam + let lam' = lam { lambdaBody = body } + return $ Op $ Inner $ GroupReduce w lam' input +offsetMemoryInExp offsets (Op (Inner (GroupGenReduce w dests lam nes vals locks))) = do + body <- offsetMemoryInBody offsets $ lambdaBody lam + let lam' = lam { lambdaBody = body + , lambdaParams = map (offsetMemoryInParam offsets) $ lambdaParams lam + } + return $ Op $ Inner $ GroupGenReduce w dests lam' nes vals locks +offsetMemoryInExp offsets (Op (Inner (Combine cspace ts active body))) = + Op . Inner . Combine cspace ts active <$> offsetMemoryInBody offsets body +offsetMemoryInExp offsets e = mapExpM recurse e + where recurse = identityMapper + { mapOnBody = const $ offsetMemoryInBody offsets + , mapOnBranchType = return . offsetMemoryInBodyReturns offsets + } + +---- Slicing allocation sizes out of a kernel. + +unAllocInKernelBody :: KernelBody InKernel + -> Either String (KernelBody Kernels.InKernel) +unAllocInKernelBody = unAllocKernelBody False + where + unAllocBody (Body attr stms res) = + Body attr <$> unAllocStms True stms <*> pure res + + unAllocKernelBody nested (KernelBody attr stms res) = + KernelBody attr <$> unAllocStms nested stms <*> pure res + + unAllocStms nested = + fmap (stmsFromList . catMaybes) . mapM (unAllocStm nested) . stmsToList + + unAllocStm nested stm@(Let _ _ (Op Alloc{})) + | nested = throwError $ "Cannot handle nested allocation: " ++ pretty stm + | otherwise = return Nothing + unAllocStm _ (Let pat attr e) = + Just <$> (Let <$> unAllocPattern pat <*> pure attr <*> mapExpM unAlloc' e) + + unAllocKernelExp (Barrier se) = + return $ Barrier se + unAllocKernelExp (SplitSpace o w i elems_per_thread) = + return $ SplitSpace o w i elems_per_thread + unAllocKernelExp (Combine cspace ts active body) = + Combine cspace ts active <$> unAllocBody body + unAllocKernelExp (GroupReduce w lam input) = + GroupReduce w <$> unAllocLambda lam <*> pure input + unAllocKernelExp (GroupScan w lam input) = + GroupScan w <$> unAllocLambda lam <*> pure input + unAllocKernelExp (GroupStream w maxchunk lam accs arrs) = + GroupStream w maxchunk <$> unAllocStreamLambda lam <*> pure accs <*> pure arrs + unAllocKernelExp (GroupGenReduce w arrs op bucket vals locks) = + GroupGenReduce w arrs <$> unAllocLambda op <*> + pure bucket <*> pure vals <*> pure locks + + unAllocStreamLambda (GroupStreamLambda chunk_size chunk_offset + acc_params arr_params body) = + GroupStreamLambda chunk_size chunk_offset + (unParams acc_params) (unParams arr_params) <$> + unAllocBody body + + unAllocLambda (Lambda params body ret) = + Lambda (unParams params) <$> unAllocBody body <*> pure ret + + unParams = mapMaybe $ traverse unAttr + + unAllocPattern pat@(Pattern ctx val) = + Pattern <$> maybe bad return (mapM (rephrasePatElem unAttr) ctx) + <*> maybe bad return (mapM (rephrasePatElem unAttr) val) + where bad = Left $ "Cannot handle memory in pattern " ++ pretty pat + + unAllocOp Alloc{} = Left "unhandled Op" + unAllocOp (Inner op) = unAllocKernelExp op + + unParam p = maybe bad return $ traverse unAttr p + where bad = Left $ "Cannot handle memory-typed parameter '" ++ pretty p ++ "'" + + unT t = maybe bad return $ unAttr t + where bad = Left $ "Cannot handle memory type '" ++ pretty t ++ "'" + + unAlloc' :: Mapper InKernel Kernels.InKernel (Either String) + unAlloc' = Mapper { mapOnBody = const unAllocBody + , mapOnRetType = unT + , mapOnBranchType = unT + , mapOnFParam = unParam + , mapOnLParam = unParam + , mapOnOp = unAllocOp + , mapOnSubExp = Right + , mapOnVName = Right + , mapOnCertificates = Right + } + +unAttr :: MemInfo d u ret -> Maybe (TypeBase (ShapeBase d) u) +unAttr (MemPrim pt) = Just $ Prim pt +unAttr (MemArray pt shape u _) = Just $ Array pt shape u +unAttr MemMem{} = Nothing + +unAllocScope :: Scope ExplicitMemory -> Scope Kernels.InKernel +unAllocScope = M.mapMaybe unInfo + where unInfo (LetInfo attr) = LetInfo <$> unAttr attr + unInfo (FParamInfo attr) = FParamInfo <$> unAttr attr + unInfo (LParamInfo attr) = LParamInfo <$> unAttr attr + unInfo (IndexInfo it) = Just $ IndexInfo it + +removeCommonSizes :: M.Map VName (SubExp, Space) + -> [(SubExp, [(VName, Space)])] +removeCommonSizes = M.toList . foldl' comb mempty . M.toList + where comb m (mem, (size, space)) = M.insertWith (++) size [(mem, space)] m + +sliceKernelSizes :: [SubExp] -> KernelSpace -> KernelBody InKernel + -> ExpandM (Stms Kernels.Kernels, [VName], [VName]) +sliceKernelSizes sizes kspace kbody = do + kbody' <- either compilerLimitationS return $ unAllocInKernelBody kbody + let num_sizes = length sizes + i64s = replicate num_sizes $ Prim int64 + inkernels_scope <- asks unAllocScope + + let kernels_scope = castScope inkernels_scope + + (max_lam, _) <- flip runBinderT inkernels_scope $ do + xs <- replicateM num_sizes $ newParam "x" (Prim int64) + ys <- replicateM num_sizes $ newParam "y" (Prim int64) + (zs, stms) <- localScope (scopeOfLParams $ xs ++ ys) $ collectStms $ + forM (zip xs ys) $ \(x,y) -> + letSubExp "z" $ BasicOp $ BinOp (SMax Int64) (Var $ paramName x) (Var $ paramName y) + return $ Lambda (xs ++ ys) (mkBody stms zs) i64s + + (size_lam', _) <- flip runBinderT inkernels_scope $ do + params <- replicateM num_sizes $ newParam "x" (Prim int64) + (zs, stms) <- localScope (scopeOfLParams params <> + scopeOfKernelSpace kspace) $ collectStms $ do + mapM_ addStm $ kernelBodyStms kbody' + return sizes + localScope (scopeOfKernelSpace kspace) $ + Kernels.simplifyLambda kspace -- XXX, is this the right KernelSpace? + (Lambda mempty (Body () stms zs) i64s) [] + + ((maxes_per_thread, size_sums), slice_stms) <- flip runBinderT kernels_scope $ do + space_size <- letSubExp "space_size" =<< + foldBinOp (Mul Int32) (intConst Int32 1) + (map snd $ spaceDimensions kspace) + num_threads_64 <- letSubExp "num_threads" $ + BasicOp $ ConvOp (SExt Int32 Int64) $ spaceNumThreads kspace + + pat <- basicPattern [] <$> replicateM num_sizes + (newIdent "max_per_thread" $ Prim int64) + + addStms =<< + blockedReduction pat space_size Commutative + max_lam size_lam' (spaceDimensions kspace) + (replicate num_sizes $ intConst Int64 0) [] + + size_sums <- forM (patternNames pat) $ \threads_max -> + letExp "size_sum" $ + BasicOp $ BinOp (Mul Int64) (Var threads_max) num_threads_64 + + return (patternNames pat, size_sums) + + return (slice_stms, maxes_per_thread, size_sums) diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -0,0 +1,1014 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving, TypeFamilies, FlexibleContexts, TupleSections, FlexibleInstances, MultiParamTypeClasses #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Pass.ExplicitAllocations + ( explicitAllocations + , explicitAllocationsInStms + , simplifiable + + , arraySizeInBytesExp + ) +where + +import Control.Monad.State +import Control.Monad.Writer +import Control.Monad.Reader +import Control.Monad.RWS.Strict +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Control.Monad.Fail as Fail +import Data.Maybe + +import Futhark.Representation.Kernels +import Futhark.Optimise.Simplify.Lore + (mkWiseBody, + mkWiseLetStm, + removeExpWisdom, + + removeScopeWisdom) +import Futhark.MonadFreshNames +import Futhark.Representation.ExplicitMemory +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Tools +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Optimise.Simplify.Engine (SimpleOps (..)) +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Pass +import Futhark.Util (splitFromEnd, takeLast) + +type InInKernel = Futhark.Representation.Kernels.InKernel +type OutInKernel = Futhark.Representation.ExplicitMemory.InKernel + +data AllocStm = SizeComputation VName (PrimExp VName) + | Allocation VName SubExp Space + | ArrayCopy VName VName + deriving (Eq, Ord, Show) + +bindAllocStm :: (MonadBinder m, Op (Lore m) ~ MemOp inner) => + AllocStm -> m () +bindAllocStm (SizeComputation name pe) = + letBindNames_ [name] =<< toExp (coerceIntPrimExp Int64 pe) +bindAllocStm (Allocation name size space) = + letBindNames_ [name] $ Op $ Alloc size space +bindAllocStm (ArrayCopy name src) = + letBindNames_ [name] $ BasicOp $ Copy src + +class (MonadFreshNames m, HasScope lore m, ExplicitMemorish lore) => + Allocator lore m where + addAllocStm :: AllocStm -> m () + -- | The subexpression giving the number of elements we should + -- allocate space for. See 'ChunkMap' comment. + dimAllocationSize :: SubExp -> m SubExp + + expHints :: Exp lore -> m [ExpHint] + expHints e = return $ replicate (expExtTypeSize e) NoHint + +allocateMemory :: Allocator lore m => + String -> SubExp -> Space -> m VName +allocateMemory desc size space = do + v <- newVName desc + addAllocStm $ Allocation v size space + return v + +computeSize :: Allocator lore m => + String -> PrimExp VName -> m SubExp +computeSize desc se = do + v <- newVName desc + addAllocStm $ SizeComputation v se + return $ Var v + +type Allocable fromlore tolore = + (ExplicitMemorish tolore, + SameScope fromlore Kernels, + RetType fromlore ~ RetType Kernels, + BranchType fromlore ~ BranchType Kernels, + BodyAttr fromlore ~ (), + BodyAttr tolore ~ (), + ExpAttr tolore ~ (), + SizeSubst (Op tolore), + BinderOps tolore) + +-- | A mapping from chunk names to their maximum size. XXX FIXME +-- HACK: This is part of a hack to add loop-invariant allocations to +-- reduce kernels, because memory expansion does not use range +-- analysis yet (it should). +type ChunkMap = M.Map VName SubExp + +data AllocEnv fromlore tolore = + AllocEnv { chunkMap :: ChunkMap + , aggressiveReuse :: Bool + -- ^ Aggressively try to reuse memory in do-loops - + -- should be True inside kernels, False outside. + , allocInOp :: Op fromlore -> AllocM fromlore tolore (Op tolore) + } + +boundDims :: ChunkMap -> AllocEnv fromlore tolore + -> AllocEnv fromlore tolore +boundDims m env = env { chunkMap = m <> chunkMap env } + +boundDim :: VName -> SubExp -> AllocEnv fromlore tolore + -> AllocEnv fromlore tolore +boundDim name se = boundDims $ M.singleton name se + +-- | Monad for adding allocations to an entire program. +newtype AllocM fromlore tolore a = + AllocM (BinderT tolore (ReaderT (AllocEnv fromlore tolore) (State VNameSource)) a) + deriving (Applicative, Functor, Monad, + MonadFreshNames, + HasScope tolore, + LocalScope tolore, + MonadReader (AllocEnv fromlore tolore)) + +instance Fail.MonadFail (AllocM fromlore tolore) where + fail = error . ("AllocM.fail: "++) + +instance (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + MonadBinder (AllocM fromlore tolore) where + type Lore (AllocM fromlore tolore) = tolore + + mkExpAttrM _ _ = return () + + mkLetNamesM names e = do + pat <- patternWithAllocations names e + return $ Let pat (defAux ()) e + + mkBodyM bnds res = return $ Body () bnds res + + addStms binding = AllocM $ addBinderStms binding + collectStms (AllocM m) = AllocM $ collectBinderStms m + certifying cs (AllocM m) = AllocM $ certifyingBinder cs m + +instance Allocable fromlore OutInKernel => + Allocator ExplicitMemory (AllocM fromlore ExplicitMemory) where + addAllocStm (SizeComputation name se) = + letBindNames_ [name] =<< toExp (coerceIntPrimExp Int64 se) + addAllocStm (Allocation name size space) = + letBindNames_ [name] $ Op $ Alloc size space + addAllocStm (ArrayCopy name src) = + letBindNames_ [name] $ BasicOp $ Copy src + + dimAllocationSize (Var v) = + -- It is important to recurse here, as the substitution may itself + -- be a chunk size. + maybe (return $ Var v) dimAllocationSize =<< asks (M.lookup v . chunkMap) + dimAllocationSize size = + return size + + expHints = kernelExpHints + +instance Allocable fromlore OutInKernel => + Allocator OutInKernel (AllocM fromlore OutInKernel) where + addAllocStm (SizeComputation name se) = + letBindNames_ [name] =<< toExp (coerceIntPrimExp Int64 se) + addAllocStm (Allocation name size space) = + letBindNames_ [name] $ Op $ Alloc size space + addAllocStm (ArrayCopy name src) = + letBindNames_ [name] $ BasicOp $ Copy src + + dimAllocationSize (Var v) = + -- It is important to recurse here, as the substitution may itself + -- be a chunk size. + maybe (return $ Var v) dimAllocationSize =<< asks (M.lookup v . chunkMap) + dimAllocationSize size = + return size + + expHints = inKernelExpHints + +runAllocM :: MonadFreshNames m => + (Op fromlore -> AllocM fromlore tolore (Op tolore)) + -> AllocM fromlore tolore a -> m a +runAllocM handleOp (AllocM m) = + fmap fst $ modifyNameSource $ runState $ runReaderT (runBinderT m mempty) env + where env = AllocEnv mempty False handleOp + +subAllocM :: (SameScope tolore1 tolore2, ExplicitMemorish tolore2) => + (Op fromlore1 -> AllocM fromlore1 tolore1 (Op tolore1)) -> Bool + -> AllocM fromlore1 tolore1 a + -> AllocM fromlore2 tolore2 a +subAllocM handleOp b (AllocM m) = do + scope <- castScope <$> askScope + chunks <- asks chunkMap + let env = AllocEnv chunks b handleOp + fmap fst $ modifyNameSource $ runState $ runReaderT (runBinderT m scope) env + +-- | Monad for adding allocations to a single pattern. +newtype PatAllocM lore a = PatAllocM (RWS + (Scope lore) + [AllocStm] + VNameSource + a) + deriving (Applicative, Functor, Monad, + HasScope lore, + MonadWriter [AllocStm], + MonadFreshNames) + +instance Allocator ExplicitMemory (PatAllocM ExplicitMemory) where + addAllocStm = tell . pure + dimAllocationSize = return + +instance Allocator OutInKernel (PatAllocM OutInKernel) where + addAllocStm = tell . pure + dimAllocationSize = return + +runPatAllocM :: MonadFreshNames m => + PatAllocM lore a -> Scope lore + -> m (a, [AllocStm]) +runPatAllocM (PatAllocM m) mems = + modifyNameSource $ frob . runRWS m mems + where frob (a,s,w) = ((a,w),s) + +arraySizeInBytesExp :: Type -> PrimExp VName +arraySizeInBytesExp t = + product + [ toInt64 $ product $ map (primExpFromSubExp int32) (arrayDims t) + , ValueExp $ IntValue $ Int64Value $ primByteSize $ elemType t ] + where toInt64 = ConvOpExp $ SExt Int32 Int64 + +arraySizeInBytesExpM :: Allocator lore m => Type -> m (PrimExp VName) +arraySizeInBytesExpM t = do + dims <- mapM dimAllocationSize (arrayDims t) + let dim_prod_i32 = product $ map (primExpFromSubExp int32) dims + let elm_size_i64 = ValueExp $ IntValue $ Int64Value $ primByteSize $ elemType t + return $ product [ toInt64 dim_prod_i32, elm_size_i64 ] + where toInt64 = ConvOpExp $ SExt Int32 Int64 + +arraySizeInBytes :: Allocator lore m => Type -> m SubExp +arraySizeInBytes = computeSize "bytes" <=< arraySizeInBytesExpM + +allocForArray :: Allocator lore m => + Type -> Space -> m (SubExp, VName) +allocForArray t space = do + size <- arraySizeInBytes t + m <- allocateMemory "mem" size space + return (size, m) + +allocsForStm :: (Allocator lore m, ExpAttr lore ~ ()) => + [Ident] -> [Ident] -> Exp lore + -> m (Stm lore, [AllocStm]) +allocsForStm sizeidents validents e = do + rts <- expReturns e + hints <- expHints e + (ctxElems, valElems, postbnds) <- allocsForPattern sizeidents validents rts hints + return (Let (Pattern ctxElems valElems) (defAux ()) e, + postbnds) + +patternWithAllocations :: (Allocator lore m, ExpAttr lore ~ ()) => + [VName] + -> Exp lore + -> m (Pattern lore) +patternWithAllocations names e = do + (ts',sizes) <- instantiateShapes' =<< expExtType e + let identForBindage name t = + pure $ Ident name t + vals <- sequence [ identForBindage name t | (name, t) <- zip names ts' ] + (Let pat _ _, extrabnds) <- allocsForStm sizes vals e + case extrabnds of + [] -> return pat + _ -> fail $ "Cannot make allocations for pattern of " ++ pretty e + +allocsForPattern :: Allocator lore m => + [Ident] -> [Ident] -> [ExpReturns] -> [ExpHint] + -> m ([PatElem ExplicitMemory], + [PatElem ExplicitMemory], + [AllocStm]) +allocsForPattern sizeidents validents rts hints = do + let sizes' = [ PatElem size $ MemPrim int32 | size <- map identName sizeidents ] + (vals,(mems_and_sizes, postbnds)) <- + runWriterT $ forM (zip3 validents rts hints) $ \(ident, rt, hint) -> do + let shape = arrayShape $ identType ident + case rt of + MemPrim _ -> do + summary <- lift $ summaryForBindage (identType ident) hint + return $ PatElem (identName ident) summary + + MemMem (Free size) space -> + return $ PatElem (identName ident) $ + MemMem size space + + MemMem Ext{} space -> + return $ PatElem (identName ident) $ + MemMem (intConst Int32 0) space + + MemArray bt _ u (Just (ReturnsInBlock mem ixfun)) -> + PatElem (identName ident) . MemArray bt shape u . + ArrayIn mem <$> instantiateIxFun ixfun + + MemArray _ extshape _ Nothing + | Just _ <- knownShape extshape -> do + summary <- lift $ summaryForBindage (identType ident) hint + return $ PatElem (identName ident) summary + + MemArray bt _ u ret -> do + let space = case ret of + Just (ReturnsNewBlock mem_space _ _ _) -> mem_space + _ -> DefaultSpace + (memsize,mem,(ident',ixfun)) <- lift $ memForBindee ident + tell ([PatElem (identName memsize) $ MemPrim int64, + PatElem (identName mem) $ MemMem (Var $ identName memsize) space], + []) + return $ PatElem (identName ident') $ MemArray bt shape u $ + ArrayIn (identName mem) ixfun + + return (sizes' <> mems_and_sizes, + vals, + postbnds) + where knownShape = mapM known . shapeDims + known (Free v) = Just v + known Ext{} = Nothing + +instantiateIxFun :: Monad m => ExtIxFun -> m IxFun +instantiateIxFun = traverse $ traverse inst + where inst Ext{} = fail "instantiateIxFun: not yet" + inst (Free x) = return x + +summaryForBindage :: Allocator lore m => + Type -> ExpHint + -> m (MemBound NoUniqueness) +summaryForBindage (Prim bt) _ = + return $ MemPrim bt +summaryForBindage (Mem size space) _ = + return $ MemMem size space +summaryForBindage t@(Array bt shape u) NoHint = do + (_, m) <- allocForArray t DefaultSpace + return $ directIndexFunction bt shape u m t +summaryForBindage t (Hint ixfun space) = do + let bt = elemType t + bytes <- computeSize "bytes" $ + product [ConvOpExp (SExt Int32 Int64) (product (IxFun.base ixfun)), + fromIntegral (primByteSize (elemType t)::Int64)] + m <- allocateMemory "mem" bytes space + return $ MemArray bt (arrayShape t) NoUniqueness $ ArrayIn m ixfun + +memForBindee :: (MonadFreshNames m) => + Ident + -> m (Ident, + Ident, + (Ident, IxFun)) +memForBindee ident = do + size <- newIdent (memname <> "_size") (Prim int64) + mem <- newIdent memname $ Mem (Var $ identName size) DefaultSpace + return (size, + mem, + (ident, IxFun.iota $ map (primExpFromSubExp int32) $ arrayDims t)) + where memname = baseString (identName ident) <> "_mem" + t = identType ident + +directIndexFunction :: PrimType -> Shape -> u -> VName -> Type -> MemBound u +directIndexFunction bt shape u mem t = + MemArray bt shape u $ ArrayIn mem $ + IxFun.iota $ map (primExpFromSubExp int32) $ arrayDims t + +allocInFParams :: (Allocable fromlore tolore) => + [(FParam fromlore, Space)] -> + ([FParam tolore] -> AllocM fromlore tolore a) + -> AllocM fromlore tolore a +allocInFParams params m = do + (valparams, memparams) <- + runWriterT $ mapM (uncurry allocInFParam) params + let params' = memparams <> valparams + summary = scopeOfFParams params' + localScope summary $ m params' + +allocInFParam :: (Allocable fromlore tolore) => + FParam fromlore + -> Space + -> WriterT [FParam tolore] + (AllocM fromlore tolore) (FParam tolore) +allocInFParam param pspace = + case paramDeclType param of + Array bt shape u -> do + let memname = baseString (paramName param) <> "_mem" + ixfun = IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape + memsize <- lift $ newVName (memname <> "_size") + mem <- lift $ newVName memname + tell [ Param memsize $ MemPrim int64 + , Param mem $ MemMem (Var memsize) pspace] + return param { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun } + Prim bt -> + return param { paramAttr = MemPrim bt } + Mem size space -> + return param { paramAttr = MemMem size space } + +allocInMergeParams :: (Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore)) => + [VName] + -> [(FParam fromlore,SubExp)] + -> ([FParam tolore] + -> [FParam tolore] + -> ([SubExp] -> AllocM fromlore tolore ([SubExp], [SubExp])) + -> AllocM fromlore tolore a) + -> AllocM fromlore tolore a +allocInMergeParams variant merge m = do + ((valparams, handle_loop_subexps), mem_and_size_params) <- + runWriterT $ unzip <$> mapM allocInMergeParam merge + let mergeparams' = mem_and_size_params <> valparams + summary = scopeOfFParams mergeparams' + + mk_loop_res ses = do + (valargs, memargs) <- + runWriterT $ zipWithM ($) handle_loop_subexps ses + return (memargs, valargs) + + localScope summary $ m mem_and_size_params valparams mk_loop_res + where allocInMergeParam (mergeparam, Var v) + | Array bt shape u <- paramDeclType mergeparam = do + (mem, ixfun) <- lift $ lookupArraySummary v + Mem _ space <- lift $ lookupType mem + reuse <- asks aggressiveReuse + if space /= Space "local" && + reuse && + u == Unique && + loopInvariantShape mergeparam && + IxFun.isLinear ixfun + then return (mergeparam { paramAttr = MemArray bt shape Unique $ ArrayIn mem ixfun }, + lift . ensureArrayIn (paramType mergeparam) mem ixfun) + else doDefault mergeparam space + + allocInMergeParam (mergeparam, _) = doDefault mergeparam DefaultSpace + + doDefault mergeparam space = do + mergeparam' <- allocInFParam mergeparam space + return (mergeparam', linearFuncallArg (paramType mergeparam) space) + + variant_names = variant ++ map (paramName . fst) merge + loopInvariantShape = + not . any (`elem` variant_names) . subExpVars . arrayDims . paramType + +ensureArrayIn :: (Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore)) => + Type -> VName -> IxFun -> SubExp + -> AllocM fromlore tolore SubExp +ensureArrayIn _ _ _ (Constant v) = + fail $ "ensureArrayIn: " ++ pretty v ++ " cannot be an array." +ensureArrayIn t mem ixfun (Var v) = do + (src_mem, src_ixfun) <- lookupArraySummary v + if src_mem == mem && src_ixfun == ixfun + then return $ Var v + else do copy <- newIdent (baseString v ++ "_ensure_copy") t + let summary = MemArray (elemType t) (arrayShape t) NoUniqueness $ + ArrayIn mem ixfun + pat = Pattern [] [PatElem (identName copy) summary] + letBind_ pat $ BasicOp $ Copy v + return $ Var $ identName copy + +ensureDirectArray :: (Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore)) => + Maybe Space -> VName -> AllocM fromlore tolore (SubExp, VName, SubExp) +ensureDirectArray space_ok v = do + (mem, ixfun) <- lookupArraySummary v + Mem size mem_space <- lookupType mem + if IxFun.isDirect ixfun && maybe True (==mem_space) space_ok + then return (size, mem, Var v) + else needCopy (fromMaybe DefaultSpace space_ok) + where needCopy space = + -- We need to do a new allocation, copy 'v', and make a new + -- binding for the size of the memory block. + allocLinearArray space (baseString v) v + +allocLinearArray :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Space -> String -> VName + -> AllocM fromlore tolore (SubExp, VName, SubExp) +allocLinearArray space s v = do + t <- lookupType v + (size, mem) <- allocForArray t space + v' <- newIdent (s ++ "_linear") t + let pat = Pattern [] [PatElem (identName v') $ + directIndexFunction (elemType t) (arrayShape t) + NoUniqueness mem t] + addStm $ Let pat (defAux ()) $ BasicOp $ Copy v + return (size, mem, Var $ identName v') + +funcallArgs :: (Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore)) => + [(SubExp,Diet)] -> AllocM fromlore tolore [(SubExp,Diet)] +funcallArgs args = do + (valargs, mem_and_size_args) <- runWriterT $ forM args $ \(arg,d) -> do + t <- lift $ subExpType arg + arg' <- linearFuncallArg t DefaultSpace arg + return (arg', d) + return $ map (,Observe) mem_and_size_args <> valargs + +linearFuncallArg :: (Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore)) => + Type -> Space -> SubExp + -> WriterT [SubExp] (AllocM fromlore tolore) SubExp +linearFuncallArg Array{} space (Var v) = do + (size, mem, arg') <- lift $ ensureDirectArray (Just space) v + tell [size, Var mem] + return arg' +linearFuncallArg _ _ arg = + return arg + +explicitAllocations :: Pass Kernels ExplicitMemory +explicitAllocations = + Pass "explicit allocations" "Transform program to explicit memory representation" $ + intraproceduralTransformation allocInFun + +explicitAllocationsInStms :: (MonadFreshNames m, HasScope ExplicitMemory m) => + Stms Kernels -> m (Stms ExplicitMemory) +explicitAllocationsInStms stms = do + scope <- askScope + runAllocM handleKernel $ localScope scope $ allocInStms stms return + +memoryInRetType :: [RetType Kernels] -> [RetType ExplicitMemory] +memoryInRetType ts = evalState (mapM addAttr ts) $ startOfFreeIDRange ts + where addAttr (Prim t) = return $ MemPrim t + addAttr Mem{} = fail "memoryInRetType: too much memory" + addAttr (Array bt shape u) = do + i <- get <* modify (+2) + return $ MemArray bt shape u $ ReturnsNewBlock DefaultSpace (i+1) (Ext i) $ + IxFun.iota $ map convert $ shapeDims shape + + convert (Ext i) = LeafExp (Ext i) int32 + convert (Free v) = Free <$> primExpFromSubExp int32 v + +startOfFreeIDRange :: [TypeBase ExtShape u] -> Int +startOfFreeIDRange = S.size . shapeContext + +allocInFun :: MonadFreshNames m => FunDef Kernels -> m (FunDef ExplicitMemory) +allocInFun (FunDef entry fname rettype params fbody) = + runAllocM handleKernel $ + allocInFParams (zip params $ repeat DefaultSpace) $ \params' -> do + fbody' <- insertStmsM $ allocInFunBody + (map (const $ Just DefaultSpace) rettype) fbody + return $ FunDef entry fname (memoryInRetType rettype) params' fbody' + +handleKernel :: Kernel InInKernel + -> AllocM fromlore2 ExplicitMemory (MemOp (Kernel OutInKernel)) +handleKernel (GetSize key size_class) = + return $ Inner $ GetSize key size_class +handleKernel (GetSizeMax size_class) = + return $ Inner $ GetSizeMax size_class +handleKernel (CmpSizeLe key size_class x) = + return $ Inner $ CmpSizeLe key size_class x +handleKernel (Kernel desc space kernel_ts kbody) = subAllocM handleKernelExp True $ + Inner . Kernel desc space kernel_ts <$> + localScope (scopeOfKernelSpace space) (allocInKernelBody kbody) + where handleKernelExp (Barrier se) = + return $ Inner $ Barrier se + + handleKernelExp (SplitSpace o w i elems_per_thread) = + return $ Inner $ SplitSpace o w i elems_per_thread + + handleKernelExp (Combine cspace ts active body) = + Inner . Combine cspace ts active <$> allocInBodyNoDirect body + + handleKernelExp (GroupReduce w lam input) = do + summaries <- mapM lookupArraySummary arrs + lam' <- allocInReduceLambda lam summaries + return $ Inner $ GroupReduce w lam' input + where arrs = map snd input + + handleKernelExp (GroupScan w lam input) = do + summaries <- mapM lookupArraySummary arrs + lam' <- allocInReduceLambda lam summaries + return $ Inner $ GroupScan w lam' input + where arrs = map snd input + + handleKernelExp (GroupGenReduce w dests op bucket vs locks) = do + let (x_params, y_params) = splitAt (length vs) $ lambdaParams op + sliceDest dest = do + dest_t <- lookupType dest + sliceInfo dest $ fullSlice dest_t $ map DimFix bucket + x_params' <- zipWith Param (map paramName x_params) <$> + mapM sliceDest dests + y_params' <- zipWith Param (map paramName y_params) <$> + mapM subExpMemInfo vs + + op' <- allocInLambda (x_params'<>y_params') (lambdaBody op) (lambdaReturnType op) + return $ Inner $ GroupGenReduce w dests op' bucket vs locks + + handleKernelExp (GroupStream w maxchunk lam accs arrs) = do + acc_summaries <- mapM accSummary accs + arr_summaries <- mapM lookupArraySummary arrs + lam' <- allocInGroupStreamLambda maxchunk lam acc_summaries arr_summaries + return $ Inner $ GroupStream w maxchunk lam' accs arrs + where accSummary (Constant v) = return $ MemPrim $ primValueType v + accSummary (Var v) = lookupMemInfo v + +allocInBodyNoDirect :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Body fromlore -> AllocM fromlore tolore (Body tolore) +allocInBodyNoDirect (Body _ bnds res) = + allocInStms bnds $ \bnds' -> + return $ Body () bnds' res + +bodyReturnMemCtx :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + SubExp -> AllocM fromlore tolore [SubExp] +bodyReturnMemCtx Constant{} = + return [] +bodyReturnMemCtx (Var v) = do + info <- lookupMemInfo v + case info of + MemPrim{} -> return [] + MemMem{} -> return [] -- should not happen + MemArray _ _ _ (ArrayIn mem _) -> do + size <- lookupMemSize mem + return [size, Var mem] + +allocInFunBody :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + [Maybe Space] -> Body fromlore -> AllocM fromlore tolore (Body tolore) +allocInFunBody space_oks (Body _ bnds res) = + allocInStms bnds $ \bnds' -> do + (res'', allocs) <- collectStms $ do + res' <- zipWithM ensureDirect space_oks' res + let (ctx_res, val_res) = splitFromEnd num_vals res' + mem_ctx_res <- concat <$> mapM bodyReturnMemCtx val_res + return $ ctx_res <> mem_ctx_res <> val_res + return $ Body () (bnds'<>allocs) res'' + where num_vals = length space_oks + space_oks' = replicate (length res - num_vals) Nothing ++ space_oks + ensureDirect _ se@Constant{} = return se + ensureDirect space_ok (Var v) = do + bt <- primType <$> lookupType v + if bt + then return $ Var v + else do (_, _, v') <- ensureDirectArray space_ok v + return v' + +allocInStms :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Stms fromlore -> (Stms tolore -> AllocM fromlore tolore a) + -> AllocM fromlore tolore a +allocInStms origbnds m = allocInStms' (stmsToList origbnds) mempty + where allocInStms' [] bnds' = + m bnds' + allocInStms' (x:xs) bnds' = do + allocbnds <- allocInStm' x + let summaries = scopeOf allocbnds + localScope summaries $ + local (boundDims $ mconcat $ map sizeSubst $ stmsToList allocbnds) $ + allocInStms' xs (bnds'<>allocbnds) + allocInStm' bnd = do + ((),bnds') <- collectStms $ certifying (stmCerts bnd) $ allocInStm bnd + return bnds' + +allocInStm :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Stm fromlore -> AllocM fromlore tolore () +allocInStm (Let (Pattern sizeElems valElems) _ e) = do + e' <- allocInExp e + let sizeidents = map patElemIdent sizeElems + validents = map patElemIdent valElems + (bnd, bnds) <- allocsForStm sizeidents validents e' + addStm bnd + mapM_ addAllocStm bnds + +allocInExp :: (Allocable fromlore tolore, Allocator tolore (AllocM fromlore tolore)) => + Exp fromlore -> AllocM fromlore tolore (Exp tolore) +allocInExp (DoLoop ctx val form (Body () bodybnds bodyres)) = + allocInMergeParams mempty ctx $ \_ ctxparams' _ -> + allocInMergeParams (map paramName ctxparams') val $ + \new_ctx_params valparams' mk_loop_val -> do + form' <- allocInLoopForm form + localScope (scopeOf form') $ do + (valinit_ctx, valinit') <- mk_loop_val valinit + body' <- insertStmsM $ allocInStms bodybnds $ \bodybnds' -> do + ((val_ses,valres'),val_retbnds) <- collectStms $ mk_loop_val valres + return $ Body () (bodybnds'<>val_retbnds) (ctxres++val_ses++valres') + return $ + DoLoop + (zip (ctxparams'++new_ctx_params) (ctxinit++valinit_ctx)) + (zip valparams' valinit') + form' body' + where (_ctxparams, ctxinit) = unzip ctx + (_valparams, valinit) = unzip val + (ctxres, valres) = splitAt (length ctx) bodyres +allocInExp (Apply fname args rettype loc) = do + args' <- funcallArgs args + return $ Apply fname args' (memoryInRetType rettype) loc +allocInExp (If cond tbranch fbranch (IfAttr rets ifsort)) = do + tbranch' <- allocInFunBody (map (const Nothing) rets) tbranch + space_oks <- mkSpaceOks (length rets) tbranch' + fbranch' <- allocInFunBody space_oks fbranch + let rets' = createBodyReturns rets space_oks + return $ If cond tbranch' fbranch' $ IfAttr rets' ifsort +allocInExp e = mapExpM alloc e + where alloc = + identityMapper { mapOnBody = fail "Unhandled Body in ExplicitAllocations" + , mapOnRetType = fail "Unhandled RetType in ExplicitAllocations" + , mapOnBranchType = fail "Unhandled BranchType in ExplicitAllocations" + , mapOnFParam = fail "Unhandled FParam in ExplicitAllocations" + , mapOnLParam = fail "Unhandled LParam in ExplicitAllocations" + , mapOnOp = \op -> do handle <- asks allocInOp + handle op + } + +mkSpaceOks :: (ExplicitMemorish tolore, LocalScope tolore m) => + Int -> Body tolore -> m [Maybe Space] +mkSpaceOks num_vals (Body _ stms res) = + inScopeOf stms $ + mapM mkSpaceOK $ takeLast num_vals res + where mkSpaceOK (Var v) = do + v_info <- lookupMemInfo v + case v_info of MemArray _ _ _ (ArrayIn mem _) -> do + mem_info <- lookupMemInfo mem + case mem_info of MemMem _ space -> return $ Just space + _ -> return Nothing + _ -> return Nothing + mkSpaceOK _ = return Nothing + +createBodyReturns :: [ExtType] -> [Maybe Space] -> [BodyReturns] +createBodyReturns ts spaces = + evalState (zipWithM inspect ts spaces) $ S.size $ shapeContext ts + where inspect (Array pt shape u) space = do + i <- get <* modify (+2) + let space' = fromMaybe DefaultSpace space + return $ MemArray pt shape u $ ReturnsNewBlock space' (i+1) (Ext i) $ + IxFun.iota $ map convert $ shapeDims shape + inspect (Prim pt) _ = + return $ MemPrim pt + inspect (Mem size space) _ = + return $ MemMem (Free size) space + + convert (Ext i) = LeafExp (Ext i) int32 + convert (Free v) = Free <$> primExpFromSubExp int32 v + +allocInLoopForm :: (Allocable fromlore tolore, + Allocator tolore (AllocM fromlore tolore)) => + LoopForm fromlore -> AllocM fromlore tolore (LoopForm tolore) +allocInLoopForm (WhileLoop v) = return $ WhileLoop v +allocInLoopForm (ForLoop i it n loopvars) = + ForLoop i it n <$> mapM allocInLoopVar loopvars + where allocInLoopVar (p,a) = do + (mem, ixfun) <- lookupArraySummary a + case paramType p of + Array bt shape u -> + let ixfun' = IxFun.slice ixfun $ + fullSliceNum (IxFun.shape ixfun) [DimFix $ LeafExp i int32] + in return (p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' }, a) + Prim bt -> + return (p { paramAttr = MemPrim bt }, a) + Mem size space -> + return (p { paramAttr = MemMem size space }, a) + +allocInReduceLambda :: Lambda InInKernel + -> [(VName, IxFun)] + -> AllocM InInKernel OutInKernel (Lambda OutInKernel) +allocInReduceLambda lam input_summaries = do + let (i, j_param, actual_params) = + partitionChunkedKernelLambdaParameters $ lambdaParams lam + (acc_params, arr_params) = + splitAt (length input_summaries) actual_params + this_index = LeafExp i int32 + other_index = LeafExp (paramName j_param) int32 + acc_params' <- + allocInReduceParameters this_index $ + zip acc_params input_summaries + arr_params' <- + allocInReduceParameters other_index $ + zip arr_params input_summaries + + allocInLambda (Param i (MemPrim int32) : + j_param { paramAttr = MemPrim int32 } : + acc_params' ++ arr_params') + (lambdaBody lam) (lambdaReturnType lam) + +allocInReduceParameters :: PrimExp VName + -> [(LParam InInKernel, (VName, IxFun))] + -> AllocM InInKernel OutInKernel [LParam ExplicitMemory] +allocInReduceParameters my_id = mapM allocInReduceParameter + where allocInReduceParameter (p, (mem, ixfun)) = + case paramType p of + (Array bt shape u) -> + let ixfun' = IxFun.slice ixfun $ + fullSliceNum (IxFun.shape ixfun) [DimFix my_id] + in return p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' } + Prim bt -> + return p { paramAttr = MemPrim bt } + Mem size space -> + return p { paramAttr = MemMem size space } + +allocInChunkedParameters :: PrimExp VName + -> [(LParam InInKernel, (VName, IxFun))] + -> AllocM InInKernel OutInKernel [LParam OutInKernel] +allocInChunkedParameters offset = mapM allocInChunkedParameter + where allocInChunkedParameter (p, (mem, ixfun)) = + case paramType p of + Array bt shape u -> + let ixfun' = IxFun.offsetIndex ixfun offset + in return p { paramAttr = MemArray bt shape u $ ArrayIn mem ixfun' } + Prim bt -> + return p { paramAttr = MemPrim bt } + Mem size space -> + return p { paramAttr = MemMem size space } + +allocInLambda :: [LParam OutInKernel] -> Body InInKernel -> [Type] + -> AllocM InInKernel OutInKernel (Lambda OutInKernel) +allocInLambda params body rettype = do + body' <- localScope (scopeOfLParams params) $ + allocInStms (bodyStms body) $ \bnds' -> + return $ Body () bnds' $ bodyResult body + return $ Lambda params body' rettype + +allocInKernelBody :: KernelBody InInKernel + -> AllocM InInKernel OutInKernel (KernelBody OutInKernel) +allocInKernelBody (KernelBody () stms res) = + allocInStms stms $ \stms' -> + return $ KernelBody () stms' res + +class SizeSubst op where + opSizeSubst :: PatternT attr -> op -> ChunkMap + +instance SizeSubst (Kernel lore) where + opSizeSubst _ _ = mempty + +instance SizeSubst op => SizeSubst (MemOp op) where + opSizeSubst pat (Inner op) = opSizeSubst pat op + opSizeSubst _ _ = mempty + +instance SizeSubst (KernelExp lore) where + opSizeSubst (Pattern _ [size]) (SplitSpace _ _ _ elems_per_thread) = + M.singleton (patElemName size) elems_per_thread + opSizeSubst _ _ = mempty + +sizeSubst :: SizeSubst (Op lore) => Stm lore -> ChunkMap +sizeSubst (Let pat _ (Op op)) = opSizeSubst pat op +sizeSubst _ = mempty + +allocInGroupStreamLambda :: SubExp + -> GroupStreamLambda InInKernel + -> [MemBound NoUniqueness] + -> [(VName, IxFun)] + -> AllocM InInKernel OutInKernel (GroupStreamLambda OutInKernel) +allocInGroupStreamLambda maxchunk lam acc_summaries arr_summaries = do + let GroupStreamLambda block_size block_offset acc_params arr_params body = lam + + acc_params' <- + allocInAccParameters acc_params acc_summaries + arr_params' <- + allocInChunkedParameters (LeafExp block_offset int32) $ + zip arr_params arr_summaries + + body' <- localScope (M.insert block_size (IndexInfo Int32) $ + M.insert block_offset (IndexInfo Int32) $ + scopeOfLParams $ acc_params' ++ arr_params') $ + local (boundDim block_size maxchunk) $ do + body' <- allocInBodyNoDirect body + insertStmsM $ do + -- We copy the result of the body to whereever the accumulators are stored. + addStms (bodyStms body') + let maybeCopyResult r p = + case paramAttr p of + MemArray _ _ _ (ArrayIn mem ixfun) -> + ensureArrayIn (paramType p) mem ixfun r + _ -> + return r + resultBodyM =<< + zipWithM maybeCopyResult (bodyResult body') acc_params' + return $ + GroupStreamLambda block_size block_offset acc_params' arr_params' body' + +allocInAccParameters :: [LParam InInKernel] + -> [MemBound NoUniqueness] + -> AllocM InInKernel OutInKernel [LParam OutInKernel] +allocInAccParameters = zipWithM allocInAccParameter + where allocInAccParameter p attr = return p { paramAttr = attr } + + +mkLetNamesB' :: (Op (Lore m) ~ MemOp inner, + MonadBinder m, ExpAttr (Lore m) ~ (), + Allocator (Lore m) (PatAllocM (Lore m))) => + ExpAttr (Lore m) -> [VName] -> Exp (Lore m) -> m (Stm (Lore m)) +mkLetNamesB' attr names e = do + scope <- askScope + pat <- bindPatternWithAllocations scope names e + return $ Let pat (defAux attr) e + +mkLetNamesB'' :: (Op (Lore m) ~ MemOp inner, ExpAttr lore ~ (), + HasScope (Engine.Wise lore) m, Allocator lore (PatAllocM lore), + MonadBinder m, Engine.CanBeWise (Op lore)) => + [VName] -> Exp (Engine.Wise lore) + -> m (Stm (Engine.Wise lore)) +mkLetNamesB'' names e = do + scope <- Engine.removeScopeWisdom <$> askScope + (pat, prestms) <- runPatAllocM (patternWithAllocations names $ Engine.removeExpWisdom e) scope + mapM_ bindAllocStm prestms + let pat' = Engine.addWisdomToPattern pat e + attr = Engine.mkWiseExpAttr pat' () e + return $ Let pat' (defAux attr) e + +instance BinderOps ExplicitMemory where + mkExpAttrB _ _ = return () + mkBodyB stms res = return $ Body () stms res + mkLetNamesB = mkLetNamesB' () + +instance BinderOps OutInKernel where + mkExpAttrB _ _ = return () + mkBodyB stms res = return $ Body () stms res + mkLetNamesB = mkLetNamesB' () + +instance BinderOps (Engine.Wise ExplicitMemory) where + mkExpAttrB pat e = return $ Engine.mkWiseExpAttr pat () e + mkBodyB stms res = return $ Engine.mkWiseBody () stms res + mkLetNamesB = mkLetNamesB'' + +instance BinderOps (Engine.Wise OutInKernel) where + mkExpAttrB pat e = return $ Engine.mkWiseExpAttr pat () e + mkBodyB stms res = return $ Engine.mkWiseBody () stms res + mkLetNamesB = mkLetNamesB'' + +simplifiable :: (Engine.SimplifiableLore lore, + ExpAttr lore ~ (), + BodyAttr lore ~ (), + Op lore ~ MemOp inner, + Allocator lore (PatAllocM lore)) => + (inner -> Engine.SimpleM lore (Engine.OpWithWisdom inner, Stms (Engine.Wise lore))) + -> SimpleOps lore +simplifiable simplifyInnerOp = + SimpleOps mkExpAttrS' mkBodyS' mkLetNamesS' simplifyOp + where mkExpAttrS' _ pat e = + return $ Engine.mkWiseExpAttr pat () e + + mkBodyS' _ bnds res = return $ mkWiseBody () bnds res + + mkLetNamesS' vtable names e = do + (pat', stms) <- runBinder $ bindPatternWithAllocations env names $ + removeExpWisdom e + return (mkWiseLetStm pat' (defAux ()) e, stms) + where env = removeScopeWisdom $ ST.toScope vtable + + simplifyOp (Alloc size space) = + (,) <$> (Alloc <$> Engine.simplify size <*> pure space) <*> pure mempty + simplifyOp (Inner k) = do (k', hoisted) <- simplifyInnerOp k + return (Inner k', hoisted) + +bindPatternWithAllocations :: (MonadBinder m, + ExpAttr lore ~ (), + Op (Lore m) ~ MemOp inner, + Allocator lore (PatAllocM lore)) => + Scope lore -> [VName] -> Exp lore + -> m (Pattern lore) +bindPatternWithAllocations types names e = do + (pat,prebnds) <- runPatAllocM (patternWithAllocations names e) types + mapM_ bindAllocStm prebnds + return pat + +data ExpHint = NoHint + | Hint IxFun Space + +kernelExpHints :: (Allocator lore m, Op lore ~ MemOp (Kernel somelore)) => + Exp lore -> 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 (primExpFromSubExp int32) dims') + perm_inv + return [Hint ixfun DefaultSpace] +kernelExpHints (Op (Inner (Kernel _ space rets kbody))) = + zipWithM hint rets $ kernelBodyResult kbody + where num_threads = spaceNumThreads space + + spacy AllThreads = Just [num_threads] + spacy ThreadsInSpace = Just $ map snd $ spaceDimensions space + spacy _ = Nothing + + -- Heuristic: do not rearrange for returned arrays that are + -- sufficiently small. + coalesceReturnOfShape _ [] = False + coalesceReturnOfShape bs [Constant (IntValue (Int32Value d))] = bs * d > 4 + coalesceReturnOfShape _ _ = True + + 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 (primExpFromSubExp int32) dims_perm + ixfun_rearranged = IxFun.permute ixfun_base perm_inv + in ixfun_rearranged + + hint t (ThreadsReturn threads _) + | coalesceReturnOfShape (primByteSize (elemType t)) $ arrayDims t, + Just space_dims <- spacy threads = do + 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 + + -- TODO: Can we make hint for ConcatRetuns when it has an offset? + hint Prim{} (ConcatReturns SplitContiguous w elems_per_thread Nothing _) = do + let ixfun_base = IxFun.iota $ map (primExpFromSubExp int32) [num_threads,elems_per_thread] + ixfun_tr = IxFun.permute ixfun_base [1,0] + ixfun = IxFun.reshape ixfun_tr $ map (DimNew . primExpFromSubExp int32) [w] + return $ Hint ixfun DefaultSpace + + hint _ _ = return NoHint +kernelExpHints e = + return $ replicate (expExtTypeSize e) NoHint + +inKernelExpHints :: (Allocator lore m, Op lore ~ MemOp (KernelExp somelore)) => + Exp lore -> m [ExpHint] +inKernelExpHints (Op (Inner (Combine (CombineSpace scatter cspace) ts _ _))) = + fmap (replicate (sum ns) NoHint ++) $ forM (drop (sum ns*2) ts) $ \t -> do + alloc_dims <- mapM dimAllocationSize $ dims ++ arrayDims t + let ixfun = IxFun.iota $ map (primExpFromSubExp int32) alloc_dims + return $ Hint ixfun $ Space "local" + where dims = map snd cspace + (_, ns, _) = unzip3 scatter + +inKernelExpHints e = + return $ replicate (expExtTypeSize e) NoHint diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -0,0 +1,1595 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | Kernel extraction. +-- +-- In the following, I will use the term "width" to denote the amount +-- of immediate parallelism in a map - that is, the outer size of the +-- array(s) being used as input. +-- +-- = Basic Idea +-- +-- If we have: +-- +-- @ +-- map +-- map(f) +-- bnds_a... +-- map(g) +-- @ +-- +-- Then we want to distribute to: +-- +-- @ +-- map +-- map(f) +-- map +-- bnds_a +-- map +-- map(g) +-- @ +-- +-- But for now only if +-- +-- (0) it can be done without creating irregular arrays. +-- Specifically, the size of the arrays created by @map(f)@, by +-- @map(g)@ and whatever is created by @bnds_a@ that is also used +-- in @map(g)@, must be invariant to the outermost loop. +-- +-- (1) the maps are _balanced_. That is, the functions @f@ and @g@ +-- must do the same amount of work for every iteration. +-- +-- The advantage is that the map-nests containing @map(f)@ and +-- @map(g)@ can now be trivially flattened at no cost, thus exposing +-- more parallelism. Note that the @bnds_a@ map constitutes array +-- expansion, which requires additional storage. +-- +-- = Distributing Sequential Loops +-- +-- As a starting point, sequential loops are treated like scalar +-- expressions. That is, not distributed. However, sometimes it can +-- be worthwhile to distribute if they contain a map: +-- +-- @ +-- map +-- loop +-- map +-- map +-- @ +-- +-- If we distribute the loop and interchange the outer map into the +-- loop, we get this: +-- +-- @ +-- loop +-- map +-- map +-- map +-- map +-- @ +-- +-- Now more parallelism may be available. +-- +-- = Unbalanced Maps +-- +-- Unbalanced maps will as a rule be sequentialised, but sometimes, +-- there is another way. Assume we find this: +-- +-- @ +-- map +-- map(f) +-- map(g) +-- map +-- @ +-- +-- Presume that @map(f)@ is unbalanced. By the simple rule above, we +-- would then fully sequentialise it, resulting in this: +-- +-- @ +-- map +-- loop +-- map +-- map +-- @ +-- +-- == Balancing by Loop Interchange +-- +-- The above is not ideal, as we cannot flatten the @map-loop@ nest, +-- and we are thus limited in the amount of parallelism available. +-- +-- But assume now that the width of @map(g)@ is invariant to the outer +-- loop. Then if possible, we can interchange @map(f)@ and @map(g)@, +-- sequentialise @map(f)@ and distribute, interchanging the outer +-- parallel loop into the sequential loop: +-- +-- @ +-- loop(f) +-- map +-- map(g) +-- map +-- map +-- @ +-- +-- After flattening the two nests we can obtain more parallelism. +-- +-- When distributing a map, we also need to distribute everything that +-- the map depends on - possibly as its own map. When distributing a +-- set of scalar bindings, we will need to know which of the binding +-- results are used afterwards. Hence, we will need to compute usage +-- information. +-- +-- = Redomap +-- +-- Redomap can be handled much like map. Distributed loops are +-- distributed as maps, with the parameters corresponding to the +-- neutral elements added to their bodies. The remaining loop will +-- remain a redomap. Example: +-- +-- @ +-- redomap(op, +-- fn (v) => +-- map(f) +-- map(g), +-- e,a) +-- @ +-- +-- distributes to +-- +-- @ +-- let b = map(fn v => +-- let acc = e +-- map(f), +-- a) +-- redomap(op, +-- fn (v,dist) => +-- map(g), +-- e,a,b) +-- @ +-- +-- Note that there may be further kernel extraction opportunities +-- inside the @map(f)@. The downside of this approach is that the +-- intermediate array (@b@ above) must be written to main memory. An +-- often better approach is to just turn the entire @redomap@ into a +-- single kernel. +-- +module Futhark.Pass.ExtractKernels + (extractKernels) + where + +import Control.Monad.RWS.Strict +import Control.Monad.Reader +import Control.Monad.Trans.Maybe +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import Data.List +import qualified Data.Semigroup as Sem + +import Futhark.Representation.SOACS +import Futhark.Representation.SOACS.Simplify (simplifyStms, simpleSOACS) +import qualified Futhark.Representation.Kernels as Out +import Futhark.Representation.Kernels.Kernel +import Futhark.MonadFreshNames +import Futhark.Tools +import qualified Futhark.Transform.FirstOrderTransform as FOT +import qualified Futhark.Pass.ExtractKernels.Kernelise as Kernelise +import Futhark.Transform.Rename +import Futhark.Pass +import Futhark.Transform.CopyPropagate +import Futhark.Pass.ExtractKernels.Distribution +import Futhark.Pass.ExtractKernels.ISRWIM +import Futhark.Pass.ExtractKernels.BlockedKernel +import Futhark.Pass.ExtractKernels.Segmented +import Futhark.Pass.ExtractKernels.Interchange +import Futhark.Pass.ExtractKernels.Intragroup +import Futhark.Util +import Futhark.Util.Log + +type KernelsStms = Out.Stms Out.Kernels +type InKernelStms = Out.Stms Out.InKernel +type InKernelLambda = Out.Lambda Out.InKernel + +-- | Transform a program using SOACs to a program using explicit +-- kernels, using the kernel extraction transformation. +extractKernels :: Pass SOACS Out.Kernels +extractKernels = + Pass { passName = "extract kernels" + , passDescription = "Perform kernel extraction" + , passFunction = runDistribM . fmap Prog . mapM transformFunDef . progFunctions + } + +newtype DistribM a = DistribM (RWS (Scope Out.Kernels) Log VNameSource a) + deriving (Functor, Applicative, Monad, + HasScope Out.Kernels, + LocalScope Out.Kernels, + MonadFreshNames, + MonadLogger) + +runDistribM :: (MonadLogger m, MonadFreshNames m) => + DistribM a -> m a +runDistribM (DistribM m) = do + (x, msgs) <- modifyNameSource $ positionNameSource . runRWS m M.empty + addLog msgs + return x + where positionNameSource (x, src, msgs) = ((x, msgs), src) + +runDistribM' :: MonadFreshNames m => DistribM a -> m a +runDistribM' (DistribM m) = + fmap fst $ modifyNameSource $ positionNameSource . runRWS m M.empty + where positionNameSource (x, src, msgs) = ((x, msgs), src) + +transformFunDef :: FunDef -> DistribM (Out.FunDef Out.Kernels) +transformFunDef (FunDef entry name rettype params body) = do + body' <- localScope (scopeOfFParams params) $ + transformBody mempty body + return $ FunDef entry name rettype params body' + +transformBody :: KernelPath -> Body -> DistribM (Out.Body Out.Kernels) +transformBody path body = do bnds <- transformStms path $ stmsToList $ bodyStms body + return $ mkBody bnds $ bodyResult body + +transformStms :: KernelPath -> [Stm] -> DistribM KernelsStms +transformStms _ [] = + return mempty +transformStms path (bnd:bnds) = + sequentialisedUnbalancedStm bnd >>= \case + Nothing -> do + bnd' <- transformStm path bnd + inScopeOf bnd' $ + (bnd'<>) <$> transformStms path bnds + Just bnds' -> + transformStms path $ stmsToList bnds' <> bnds + +sequentialisedUnbalancedStm :: Stm -> DistribM (Maybe (Stms SOACS)) +sequentialisedUnbalancedStm (Let pat _ (Op soac@(Screma _ form _))) + | Just (_, _, _, lam2) <- isRedomapSOAC form, + unbalancedLambda lam2, lambdaContainsParallelism lam2 = do + types <- asksScope scopeForSOACs + Just . snd <$> runBinderT (FOT.transformSOAC pat soac) types +sequentialisedUnbalancedStm _ = + return Nothing + +scopeForSOACs :: Scope Out.Kernels -> Scope SOACS +scopeForSOACs = castScope + +scopeForKernels :: Scope SOACS -> Scope Out.Kernels +scopeForKernels = castScope + +transformStm :: KernelPath -> Stm -> DistribM KernelsStms + +transformStm path (Let pat aux (Op (CmpThreshold what s))) = + runBinder_ $ do + (r, _) <- cmpSizeLe s (Out.SizeThreshold path) what + addStm $ Let pat aux $ BasicOp $ SubExp r + +transformStm path (Let pat aux (If c tb fb rt)) = do + tb' <- transformBody path tb + fb' <- transformBody path fb + return $ oneStm $ Let pat aux $ If c tb' fb' rt + +transformStm path (Let pat aux (DoLoop ctx val form body)) = + localScope (castScope (scopeOf form) <> + scopeOfFParams mergeparams) $ + oneStm . Let pat aux . DoLoop ctx val form' <$> transformBody path body + where mergeparams = map fst $ ctx ++ val + form' = case form of + WhileLoop cond -> + WhileLoop cond + ForLoop i it bound ps -> + ForLoop i it bound ps + +transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) + | Just lam <- isMapSOAC form = + distributeMap path $ MapLoop pat cs w lam arrs + +transformStm path (Let res_pat (StmAux cs _) (Op (Screma w form arrs))) + | Just (scan_lam, nes) <- isScanSOAC form, + Just do_iswim <- iswim res_pat w scan_lam $ zip nes arrs = do + types <- asksScope scopeForSOACs + transformStms path =<< (stmsToList . snd <$> runBinderT (certifying cs do_iswim) types) + + | Just (scan_lam, scan_nes) <- isScanSOAC form, + ScremaForm _ _ map_lam <- form = + doScan (scan_lam, scan_nes) (mempty, nilFn, mempty) map_lam + + | ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam <- form, + not $ null scan_nes, all primType $ lambdaReturnType scan_lam, + not $ lambdaContainsParallelism map_lam = + doScan (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam + + where doScan (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam = do + scan_lam_sequential <- Kernelise.transformLambda scan_lam + red_lam_sequential <- Kernelise.transformLambda red_lam + map_lam_sequential <- Kernelise.transformLambda map_lam + runBinder_ $ certifying cs $ + blockedScan res_pat w + (scan_lam_sequential, scan_nes) + (comm, red_lam_sequential, red_nes) + map_lam_sequential (intConst Int32 16) [] [] arrs + +transformStm path (Let res_pat (StmAux cs _) (Op (Screma w form arrs))) + | Just (comm, red_fun, nes) <- isReduceSOAC form, + let comm' | commutativeLambda red_fun = Commutative + | otherwise = comm, + Just do_irwim <- irwim res_pat w comm' red_fun $ zip nes arrs = do + types <- asksScope scopeForSOACs + bnds <- fst <$> runBinderT (simplifyStms =<< collectStms_ (certifying cs do_irwim)) types + transformStms path $ stmsToList bnds + +transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) + | Just (comm, red_lam, nes, map_lam) <- isRedomapSOAC form = do + + let paralleliseOuter = do + red_lam_sequential <- Kernelise.transformLambda red_lam + map_lam_sequential <- Kernelise.transformLambda map_lam + fmap (certify cs) <$> + blockedReduction pat w comm' red_lam_sequential map_lam_sequential [] nes arrs + + outerParallelBody = + renameBody =<< + (mkBody <$> paralleliseOuter <*> pure (map Var (patternNames pat))) + + paralleliseInner path' = do + (mapbnd, redbnd) <- redomapToMapAndReduce pat (w, comm', red_lam, map_lam, nes, arrs) + transformStms path' [certify cs mapbnd, certify cs redbnd] + + innerParallelBody path' = + renameBody =<< + (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) + + + comm' | commutativeLambda red_lam = Commutative + | otherwise = comm + + if not $ lambdaContainsParallelism map_lam + then paralleliseOuter + else if incrementalFlattening then do + ((outer_suff, outer_suff_key), suff_stms) <- + runBinder $ sufficientParallelism "suff_outer_redomap" w path + + outer_stms <- outerParallelBody + inner_stms <- innerParallelBody ((outer_suff_key, False):path) + + (suff_stms<>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] + else paralleliseOuter + +-- Streams can be handled in two different ways - either we +-- sequentialise the body or we keep it parallel and distribute. +transformStm path (Let pat (StmAux cs _) (Op (Stream w (Parallel _ _ _ []) map_fun arrs))) = do + -- No reduction part. Remove the stream and leave the body + -- parallel. It will be distributed. + types <- asksScope scopeForSOACs + transformStms path =<< + (stmsToList . snd <$> runBinderT (certifying cs $ sequentialStreamWholeArray pat w [] map_fun arrs) types) + +transformStm path (Let pat aux@(StmAux cs _) (Op (Stream w (Parallel o comm red_fun nes) fold_fun arrs))) + | incrementalFlattening = do + ((outer_suff, outer_suff_key), suff_stms) <- + runBinder $ sufficientParallelism "suff_outer_stream" w path + + outer_stms <- outerParallelBody ((outer_suff_key, True) : path) + inner_stms <- innerParallelBody ((outer_suff_key, False) : path) + + (suff_stms<>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)] + + | otherwise = paralleliseOuter path + + where + paralleliseOuter path' + | any (not . primType) $ lambdaReturnType red_fun = do + -- Split into a chunked map and a reduction, with the latter + -- further transformed. + fold_fun_sequential <- Kernelise.transformLambda fold_fun + + let (red_pat_elems, concat_pat_elems) = + splitAt (length nes) $ patternValueElements pat + red_pat = Pattern [] red_pat_elems + concat_pat = Pattern [] concat_pat_elems + + (map_bnd, map_misc_bnds) <- blockedMap concat_pat w InOrder fold_fun_sequential nes arrs + let num_threads = arraysSize 0 $ patternTypes $ stmPattern map_bnd + + reduce_soac <- reduceSOAC comm' red_fun nes + + ((map_misc_bnds<>oneStm map_bnd)<>) <$> + inScopeOf (map_misc_bnds<>oneStm map_bnd) + (transformStm path' $ Let red_pat aux $ + Op (Screma num_threads reduce_soac $ patternNames $ stmPattern map_bnd)) + + | otherwise = do + red_fun_sequential <- Kernelise.transformLambda red_fun + fold_fun_sequential <- Kernelise.transformLambda fold_fun + fmap (certify cs) <$> + blockedReductionStream pat w comm' red_fun_sequential fold_fun_sequential [] nes arrs + + outerParallelBody path' = + renameBody =<< + (mkBody <$> paralleliseOuter path' <*> pure (map Var (patternNames pat))) + + paralleliseInner path' = do + types <- asksScope scopeForSOACs + transformStms path' . fmap (certify cs) =<< + (stmsToList . snd <$> runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types) + + innerParallelBody path' = + renameBody =<< + (mkBody <$> paralleliseInner path' <*> pure (map Var (patternNames pat))) + + comm' | commutativeLambda red_fun, o /= InOrder = Commutative + | otherwise = comm + +transformStm path (Let pat (StmAux cs _) (Op (Screma w form arrs))) = do + -- This with-loop is too complicated for us to immediately do + -- anything, so split it up and try again. + scope <- asksScope scopeForSOACs + transformStms path . map (certify cs) . stmsToList . snd =<< + runBinderT (dissectScrema pat w form arrs) scope + +transformStm path (Let pat _ (Op (Stream w (Sequential nes) fold_fun arrs))) = do + -- Remove the stream and leave the body parallel. It will be + -- distributed. + types <- asksScope scopeForSOACs + 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 + lam' <- Kernelise.transformLambda lam + write_i <- newVName "write_i" + let (as_ws, as_ns, as_vs) = unzip3 as + (i_res, v_res) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam' + kstms = bodyStms $ lambdaBody lam' + krets = do (a_w, a, is_vs) <- zip3 as_ws as_vs $ chunks as_ns $ zip i_res v_res + return $ WriteReturn [a_w] a [ ([i],v) | (i,v) <- is_vs ] + body = KernelBody () kstms krets + inputs = do (p, p_a) <- zip (lambdaParams lam') ivs + return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] + (bnds, kernel) <- + mapKernel w (FlatThreadSpace [(write_i,w)]) inputs (map rowType $ patternTypes pat) body + certifying cs $ do + addStms bnds + letBind_ pat $ Op kernel + +transformStm path (Let orig_pat (StmAux cs _) (Op (GenReduce w ops bucket_fun imgs))) = do + bfun' <- Kernelise.transformLambda bucket_fun + genReduceKernel path [] orig_pat [] [] cs w ops bfun' imgs + +transformStm _ bnd = + runBinder_ $ FOT.transformStmRecursively bnd + +data MapLoop = MapLoop Pattern Certificates SubExp Lambda [VName] + +mapLoopStm :: MapLoop -> Stm +mapLoopStm (MapLoop pat cs w lam arrs) = Let pat (StmAux cs ()) $ Op $ Screma w (mapSOAC lam) arrs + +sufficientParallelism :: (Op (Lore m) ~ Kernel innerlore, MonadBinder m) => + String -> SubExp -> KernelPath -> m (SubExp, VName) +sufficientParallelism desc what path = cmpSizeLe desc (Out.SizeThreshold path) what + +distributeMap :: (HasScope Out.Kernels m, + MonadFreshNames m, MonadLogger m) => + KernelPath -> MapLoop -> m KernelsStms +distributeMap path (MapLoop pat cs w lam arrs) = do + types <- askScope + let loopnest = MapNesting pat cs w $ zip (lambdaParams lam) arrs + env path' = KernelEnv { kernelNest = + singleNesting (Nesting mempty loopnest) + , kernelScope = + scopeForKernels (scopeOf lam) <> types + , kernelPath = + path' + } + exploitInnerParallelism path' = do + (acc', postkernels) <- runKernelM (env path') $ + distribute =<< distributeMapBodyStms acc (bodyStms $ lambdaBody lam) + + -- There may be a few final targets remaining - these correspond to + -- arrays that are identity mapped, and must have statements + -- inserted here. + return $ postKernelsStms postkernels <> + identityStms (outerTarget $ kernelTargets acc') + + if not incrementalFlattening then exploitInnerParallelism path + else do + + let exploitOuterParallelism path' = do + soactypes <- asksScope scopeForSOACs + (seq_lam, _) <- runBinderT (Kernelise.transformLambda lam) soactypes + (acc', postkernels) <- runKernelM (env path') $ distribute $ + addStmsToKernel (bodyStms $ lambdaBody seq_lam) acc + -- As above, we deal with identity mappings. + return $ postKernelsStms postkernels <> + identityStms (outerTarget $ kernelTargets acc') + + distributeMap' (newKernel loopnest) path exploitOuterParallelism exploitInnerParallelism pat w lam + where acc = KernelAcc { kernelTargets = singleTarget (pat, bodyResult $ lambdaBody lam) + , kernelStms = mempty + } + + params_to_arrs = zip (map paramName $ lambdaParams lam) arrs + identityStms (rem_pat, res) = + stmsFromList $ zipWith identityStm (patternValueElements rem_pat) res + identityStm pe (Var v) + | Just arr <- lookup v params_to_arrs = + Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ Copy arr + identityStm pe se = + Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ Replicate (Shape [w]) se + +distributeMap' :: (HasScope Out.Kernels m, MonadFreshNames m) => + KernelNest -> KernelPath + -> (KernelPath -> m (Out.Stms Out.Kernels)) + -> (KernelPath -> m (Out.Stms Out.Kernels)) + -> PatternT Type + -> SubExp + -> LambdaT SOACS + -> m (Out.Stms Out.Kernels) +distributeMap' loopnest path mk_seq_stms mk_par_stms pat nest_w lam = do + let res = map Var $ patternNames pat + + types <- askScope + ((outer_suff, outer_suff_key), outer_suff_stms) <- runBinder $ + sufficientParallelism "suff_outer_par" nest_w path + + intra <- if worthIntraGroup lam then + flip runReaderT types $ intraGroupParallelise loopnest lam + else return Nothing + seq_body <- renameBody =<< mkBody <$> + mk_seq_stms ((outer_suff_key, True) : path) <*> pure res + let seq_alts = [(outer_suff, seq_body) | worthSequentialising lam] + + case intra of + Nothing -> do + par_body <- renameBody =<< mkBody <$> + mk_par_stms ((outer_suff_key, False) : path) <*> pure res + + (outer_suff_stms<>) <$> kernelAlternatives pat par_body seq_alts + + Just ((_intra_min_par, intra_avail_par), group_size, intra_prelude, intra_stms) -> do + -- We must check that all intra-group parallelism fits in a group. + ((intra_ok, intra_suff_key), intra_suff_stms) <- runBinder $ do + addStms intra_prelude + + max_group_size <- + letSubExp "max_group_size" $ Op $ Out.GetSizeMax Out.SizeGroup + fits <- letSubExp "fits" $ BasicOp $ + CmpOp (CmpSle Int32) group_size max_group_size + + (intra_suff, suff_key) <- sufficientParallelism "suff_intra_par" intra_avail_par $ + (outer_suff_key, False) : path + + intra_ok <- letSubExp "intra_suff_and_fits" $ BasicOp $ BinOp LogAnd fits intra_suff + return (intra_ok, suff_key) + + group_par_body <- renameBody $ mkBody intra_stms res + + par_body <- renameBody =<< mkBody <$> + mk_par_stms ([(outer_suff_key, False), + (intra_suff_key, False)] + ++ path) <*> pure res + + ((outer_suff_stms<>intra_suff_stms)<>) <$> + kernelAlternatives pat par_body (seq_alts ++ [(intra_ok, group_par_body)]) + +data KernelEnv = KernelEnv { kernelNest :: Nestings + , kernelScope :: Scope Out.Kernels + , kernelPath :: KernelPath + } + +data KernelAcc = KernelAcc { kernelTargets :: Targets + , kernelStms :: InKernelStms + } + +data KernelRes = KernelRes { accPostKernels :: PostKernels + , accLog :: Log + } + +instance Sem.Semigroup KernelRes where + KernelRes ks1 log1 <> KernelRes ks2 log2 = + KernelRes (ks1 <> ks2) (log1 <> log2) + +instance Monoid KernelRes where + mempty = KernelRes mempty mempty + mappend = (Sem.<>) + +newtype PostKernel = PostKernel { unPostKernel :: KernelsStms } + +newtype PostKernels = PostKernels [PostKernel] + +instance Sem.Semigroup PostKernels where + PostKernels xs <> PostKernels ys = PostKernels $ ys ++ xs + +instance Monoid PostKernels where + mempty = PostKernels mempty + mappend = (Sem.<>) + +postKernelsStms :: PostKernels -> KernelsStms +postKernelsStms (PostKernels kernels) = mconcat $ map unPostKernel kernels + +typeEnvFromKernelAcc :: KernelAcc -> Scope Out.Kernels +typeEnvFromKernelAcc = scopeOfPattern . fst . outerTarget . kernelTargets + +addStmsToKernel :: InKernelStms -> KernelAcc -> KernelAcc +addStmsToKernel stms acc = + acc { kernelStms = stms <> kernelStms acc } + +addStmToKernel :: (LocalScope Out.Kernels m, MonadFreshNames m) => + Stm -> KernelAcc -> m KernelAcc +addStmToKernel bnd acc = do + stms <- runBinder_ $ Kernelise.transformStm bnd + return acc { kernelStms = stms <> kernelStms acc } + +newtype KernelM a = KernelM (RWS KernelEnv KernelRes VNameSource a) + deriving (Functor, Applicative, Monad, + MonadReader KernelEnv, + MonadWriter KernelRes, + MonadFreshNames) + +instance HasScope Out.Kernels KernelM where + askScope = asks kernelScope + +instance LocalScope Out.Kernels KernelM where + localScope types = local $ \env -> + env { kernelScope = types <> kernelScope env } + +instance MonadLogger KernelM where + addLog msgs = tell mempty { accLog = msgs } + +runKernelM :: (MonadFreshNames m, MonadLogger m) => + KernelEnv -> KernelM a -> m (a, PostKernels) +runKernelM env (KernelM m) = do + (x, res) <- modifyNameSource $ getKernels . runRWS m env + addLog $ accLog res + return (x, accPostKernels res) + where getKernels (x,s,a) = ((x, a), s) + +collectKernels :: KernelM a -> KernelM (a, PostKernels) +collectKernels m = pass $ do + (x, res) <- listen m + return ((x, accPostKernels res), + const res { accPostKernels = mempty }) + +collectKernels_ :: KernelM () -> KernelM PostKernels +collectKernels_ = fmap snd . collectKernels + +localPath :: KernelPath -> KernelM a -> KernelM a +localPath path = local $ \env -> env { kernelPath = path } + +addKernels :: PostKernels -> KernelM () +addKernels ks = tell $ mempty { accPostKernels = ks } + +addKernel :: KernelsStms -> KernelM () +addKernel bnds = addKernels $ PostKernels [PostKernel bnds] + +withStm :: Stm -> KernelM a -> KernelM a +withStm bnd = local $ \env -> + env { kernelScope = + scopeForKernels (scopeOf [bnd]) <> kernelScope env + , kernelNest = + letBindInInnerNesting provided $ + kernelNest env + } + where provided = S.fromList $ patternNames $ stmPattern bnd + +mapNesting :: Pattern -> Certificates -> SubExp -> Lambda -> [VName] + -> KernelM a + -> KernelM a +mapNesting pat cs w lam arrs = local $ \env -> + env { kernelNest = pushInnerNesting nest $ kernelNest env + , kernelScope = scopeForKernels (scopeOf lam) <> kernelScope env + } + where nest = Nesting mempty $ + MapNesting pat cs w $ + zip (lambdaParams lam) arrs + +inNesting :: KernelNest -> KernelM a -> KernelM a +inNesting (outer, nests) = local $ \env -> + env { kernelNest = (inner, nests') + , kernelScope = mconcat (map scopeOf $ outer : nests) <> kernelScope env + } + where (inner, nests') = + case reverse nests of + [] -> (asNesting outer, []) + (inner' : ns) -> (asNesting inner', map asNesting $ outer : reverse ns) + asNesting = Nesting mempty + +unbalancedLambda :: Lambda -> Bool +unbalancedLambda lam = + unbalancedBody + (S.fromList $ map paramName $ lambdaParams lam) $ + lambdaBody lam + + where subExpBound (Var i) bound = i `S.member` bound + subExpBound (Constant _) _ = False + + unbalancedBody bound body = + any (unbalancedStm (bound <> boundInBody body) . stmExp) $ + bodyStms body + + -- XXX - our notion of balancing is probably still too naive. + unbalancedStm bound (Op (Stream w _ _ _)) = + w `subExpBound` bound + unbalancedStm bound (Op (Screma w _ _)) = + w `subExpBound` bound + unbalancedStm _ Op{} = + False + unbalancedStm _ DoLoop{} = False + + unbalancedStm bound (If cond tbranch fbranch _) = + cond `subExpBound` bound && + (unbalancedBody bound tbranch || unbalancedBody bound fbranch) + + unbalancedStm _ (BasicOp _) = + False + unbalancedStm _ (Apply fname _ _ _) = + not $ isBuiltInFunction fname + +bodyContainsParallelism :: Body -> Bool +bodyContainsParallelism = any (isMap . stmExp) . bodyStms + where isMap Op{} = True + isMap _ = False + +lambdaContainsParallelism :: Lambda -> Bool +lambdaContainsParallelism = bodyContainsParallelism . lambdaBody + +-- | Returns the sizes of nested parallelism. +nestedParallelism :: Body -> [SubExp] +nestedParallelism = concatMap (parallelism . stmExp) . bodyStms + where parallelism (Op (Scatter w _ _ _)) = [w] + parallelism (Op (Screma w _ _)) = [w] + parallelism (Op (Stream w Sequential{} lam _)) + | chunk_size_param : _ <- lambdaParams lam = + let update (Var v) | v == paramName chunk_size_param = w + update se = se + in map update $ nestedParallelism $ lambdaBody lam + parallelism (DoLoop _ _ _ body) = nestedParallelism body + parallelism _ = [] + +-- | A lambda is worth sequentialising if it contains nested +-- parallelism of an interesting kind. +worthSequentialising :: Lambda -> Bool +worthSequentialising lam = interesting $ lambdaBody lam + where interesting body = any (interesting' . stmExp) $ bodyStms body + interesting' (Op (Screma _ form@(ScremaForm _ _ lam') _)) + | isJust $ isMapSOAC form = worthSequentialising lam' + interesting' (Op Scatter{}) = False -- Basically a map. + interesting' (DoLoop _ _ _ body) = interesting body + interesting' (Op _) = True + interesting' _ = False + +-- | Intra-group parallelism is worthwhile if the lambda contains +-- non-map nested parallelism, or any nested parallelism inside a +-- loop. +worthIntraGroup :: Lambda -> Bool +worthIntraGroup lam = interesting $ lambdaBody lam + where interesting body = not (null $ nestedParallelism body) && + not (onlyMaps $ bodyStms body) + onlyMaps = all $ isMapOrSeq . stmExp + isMapOrSeq (Op (Screma _ form@(ScremaForm _ _ lam') _)) + | isJust $ isMapSOAC form = not $ worthIntraGroup lam' + isMapOrSeq (Op Scatter{}) = True -- Basically a map. + isMapOrSeq (DoLoop _ _ _ body) = + null $ nestedParallelism body + isMapOrSeq (Op _) = False + isMapOrSeq _ = True + +-- Enable if you want the cool new versioned code. Beware: may be +-- slower in practice. Caveat emptor (and you are the emptor). +incrementalFlattening :: Bool +incrementalFlattening = isJust $ lookup "FUTHARK_INCREMENTAL_FLATTENING" unixEnvironment + +distributeInnerMap :: MapLoop -> KernelAcc + -> KernelM KernelAcc +distributeInnerMap maploop@(MapLoop pat cs w lam arrs) acc + | unbalancedLambda lam, lambdaContainsParallelism lam = + addStmToKernel (mapLoopStm maploop) acc + | not incrementalFlattening = + distributeNormally + | otherwise = + distributeSingleStm acc (mapLoopStm maploop) >>= \case + Just (post_kernels, res, nest, acc') + | Just (perm, _pat_unused) <- permutationAndMissing pat res -> do + addKernels post_kernels + multiVersion perm nest acc' + _ -> distributeNormally + where + lam_bnds = bodyStms $ lambdaBody lam + lam_res = bodyResult $ lambdaBody lam + + def_acc = KernelAcc { kernelTargets = pushInnerTarget + (pat, bodyResult $ lambdaBody lam) $ + kernelTargets acc + , kernelStms = mempty + } + + distributeNormally = + distribute =<< + leavingNesting maploop =<< + mapNesting pat cs w lam arrs + (distribute =<< distributeMapBodyStms def_acc lam_bnds) + + multiVersion perm nest acc' = do + -- The kernel can be distributed by itself, so now we can + -- decide whether to just sequentialise, or exploit inner + -- parallelism. + let map_nesting = MapNesting pat cs w $ zip (lambdaParams lam) arrs + lam_res' = rearrangeShape perm lam_res + nest' = pushInnerKernelNesting (pat, lam_res') map_nesting nest + extra_scope = targetsScope $ kernelTargets acc' + + exploitInnerParallelism path' = + fmap postKernelsStms $ collectKernels_ $ localPath path' $ + localScope extra_scope $ inNesting nest' $ void $ + distribute =<< leavingNesting maploop =<< distribute =<< + distributeMapBodyStms def_acc lam_bnds + + -- XXX: we do not construct a new KernelPath when + -- sequentialising. This is only OK as long as further + -- versioning does not take place down that branch (it currently + -- does not). + (nestw_bnds, nestw, sequentialised_kernel) <- localScope extra_scope $ do + sequentialised_map_body <- + localScope (scopeOfLParams (lambdaParams lam)) $ runBinder_ $ + Kernelise.transformStms lam_bnds + let kbody = KernelBody () sequentialised_map_body $ + map (ThreadsReturn ThreadsInSpace) lam_res' + constructKernel nest' kbody + + let outer_pat = loopNestingPattern $ fst nest + path <- asks kernelPath + addKernel =<< (nestw_bnds<>) <$> + localScope extra_scope (distributeMap' nest' path + (const $ return $ oneStm sequentialised_kernel) + exploitInnerParallelism + outer_pat nestw + lam { lambdaBody = (lambdaBody lam) { bodyResult = lam_res' }}) + + return acc' + +leavingNesting :: MapLoop -> KernelAcc -> KernelM KernelAcc +leavingNesting (MapLoop _ cs w lam arrs) acc = + case popInnerTarget $ kernelTargets acc of + Nothing -> + fail "The kernel targets list is unexpectedly small" + Just ((pat,res), newtargets) -> do + let acc' = acc { kernelTargets = newtargets } + if null $ kernelStms acc' + then return acc' + else do let kbody = Body () (kernelStms acc') res + used_in_body = freeInBody kbody + (used_params, used_arrs) = + unzip $ + filter ((`S.member` used_in_body) . paramName . fst) $ + zip (lambdaParams lam) arrs + stms <- runBinder_ $ Kernelise.mapIsh pat cs w used_params kbody used_arrs + return $ addStmsToKernel stms acc' { kernelStms = mempty } + +distributeMapBodyStms :: KernelAcc -> Stms SOACS -> KernelM KernelAcc +distributeMapBodyStms orig_acc = onStms orig_acc . stmsToList + where + onStms acc [] = return acc + + onStms acc (Let pat (StmAux cs _) (Op (Stream w (Sequential accs) lam arrs)):stms) = do + types <- asksScope scopeForSOACs + stream_stms <- + snd <$> runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types + stream_stms' <- + runReaderT (copyPropagateInStms simpleSOACS stream_stms) types + onStms acc $ stmsToList (fmap (certify cs) stream_stms') ++ stms + + onStms acc (stm:stms) = + -- It is important that stm is in scope if 'maybeDistributeStm' + -- wants to distribute, even if this causes the slightly silly + -- situation that stm is in scope of itself. + withStm stm $ maybeDistributeStm stm =<< onStms acc stms + +maybeDistributeStm :: Stm -> KernelAcc -> KernelM KernelAcc + +maybeDistributeStm bnd@(Let pat _ (Op (Screma w form arrs))) acc + | Just lam <- isMapSOAC form = + -- Only distribute inside the map if we can distribute everything + -- following the map. + distributeIfPossible acc >>= \case + Nothing -> addStmToKernel bnd acc + Just acc' -> distribute =<< distributeInnerMap (MapLoop pat (stmCerts bnd) w lam arrs) acc' + +maybeDistributeStm bnd@(Let pat _ (DoLoop [] val form@ForLoop{} body)) acc + | null (patternContextElements pat), bodyContainsParallelism body = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | S.null $ freeIn form `S.intersection` boundInKernelNest nest, + Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromKernelAcc acc') $ do + addKernels kernels + nest' <- expandKernelNest pat_unused nest + types <- asksScope scopeForSOACs + scope <- askScope + bnds <- runReaderT + (interchangeLoops nest' (SeqLoop perm pat val form body)) types + -- runDistribM starts out with an empty scope, so we have to + -- immmediately insert the real one. + path <- asks kernelPath + bnds' <- runDistribM $ localScope scope $ transformStms path $ stmsToList bnds + addKernel bnds' + return acc' + _ -> + addStmToKernel bnd acc + +maybeDistributeStm stm@(Let pat _ (If cond tbranch fbranch ret)) acc + | null (patternContextElements pat), + bodyContainsParallelism tbranch || bodyContainsParallelism fbranch || + any (not . primType) (ifReturns ret) = + distributeSingleStm acc stm >>= \case + Just (kernels, res, nest, acc') + | S.null $ (freeIn cond <> freeIn ret) `S.intersection` + boundInKernelNest nest, + Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromKernelAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + addKernels kernels + types <- asksScope scopeForSOACs + let branch = Branch perm pat cond tbranch fbranch ret + stms <- runReaderT (interchangeBranch nest' branch) types + -- runDistribM starts out with an empty scope, so we have to + -- immmediately insert the real one. + scope <- askScope + path <- asks kernelPath + stms' <- runDistribM $ localScope scope $ transformStms path $ stmsToList stms + addKernel stms' + return acc' + _ -> + addStmToKernel stm acc + +maybeDistributeStm (Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | Just (comm, lam, nes) <- isReduceSOAC form, + Just m <- irwim pat w comm lam $ zip nes arrs = do + types <- asksScope scopeForSOACs + (_, bnds) <- runBinderT (certifying cs m) types + distributeMapBodyStms acc bnds + +-- Parallelise segmented scatters. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Scatter w lam ivs as))) acc = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + localScope (typeEnvFromKernelAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + lam' <- Kernelise.transformLambda lam + addKernels kernels + addKernel =<< segmentedScatterKernel nest' perm pat cs w lam' ivs as + return acc' + _ -> + addStmToKernel bnd acc + +-- Parallelise segmented GenReduce. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (GenReduce w ops lam as))) acc = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + localScope (typeEnvFromKernelAcc acc') $ do + lam' <- Kernelise.transformLambda lam + nest' <- expandKernelNest pat_unused nest + addKernels kernels + addKernel =<< segmentedGenReduceKernel nest' perm cs w ops lam' as + return acc' + _ -> + addStmToKernel bnd acc + +-- If the scan can be distributed by itself, we will turn it into a +-- segmented scan. +-- +-- If the scan cannot be distributed by itself, it will be +-- sequentialised in the default case for this function. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | Just (lam, nes, map_lam) <- isScanomapSOAC form = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromKernelAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + map_lam' <- Kernelise.transformLambda map_lam + lam' <- Kernelise.transformLambda lam + localScope (typeEnvFromKernelAcc acc') $ + segmentedScanomapKernel nest' perm w lam' map_lam' nes arrs >>= + kernelOrNot cs bnd acc kernels acc' + _ -> + addStmToKernel bnd acc + +-- If the reduction can be distributed by itself, we will turn it into a +-- segmented reduce. +-- +-- If the reduction cannot be distributed by itself, it will be +-- sequentialised in the default case for this function. +maybeDistributeStm bnd@(Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | Just (comm, lam, nes, map_lam) <- isRedomapSOAC form, + isIdentityLambda map_lam || incrementalFlattening = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | Just (perm, pat_unused) <- permutationAndMissing pat res -> + -- We need to pretend pat_unused was used anyway, by adding + -- it to the kernel nest. + localScope (typeEnvFromKernelAcc acc') $ do + nest' <- expandKernelNest pat_unused nest + lam' <- Kernelise.transformLambda lam + map_lam' <- Kernelise.transformLambda map_lam + + let comm' | commutativeLambda lam = Commutative + | otherwise = comm + + regularSegmentedRedomapKernel nest' perm w comm' lam' map_lam' nes arrs >>= + kernelOrNot cs bnd acc kernels acc' + _ -> + addStmToKernel bnd acc + +maybeDistributeStm (Let pat (StmAux cs _) (Op (Screma w form arrs))) acc + | incrementalFlattening || isNothing (isRedomapSOAC form) = do + -- This with-loop is too complicated for us to immediately do + -- anything, so split it up and try again. + scope <- asksScope scopeForSOACs + distributeMapBodyStms acc . fmap (certify cs) . snd =<< + runBinderT (dissectScrema pat w form arrs) scope + +maybeDistributeStm (Let pat aux (BasicOp (Replicate (Shape (d:ds)) v))) acc + | [t] <- patternTypes pat = do + -- XXX: We need a temporary dummy binding to prevent an empty + -- map body. The kernel extractor does not like empty map + -- bodies. + tmp <- newVName "tmp" + let rowt = rowType t + newbnd = Let pat aux $ Op $ Screma d (mapSOAC lam) [] + tmpbnd = Let (Pattern [] [PatElem tmp rowt]) aux $ + BasicOp $ Replicate (Shape ds) v + lam = Lambda { lambdaReturnType = [rowt] + , lambdaParams = [] + , lambdaBody = mkBody (oneStm tmpbnd) [Var tmp] + } + maybeDistributeStm newbnd acc + +maybeDistributeStm bnd@(Let _ aux (BasicOp Copy{})) acc = + distributeSingleUnaryStm acc bnd $ \_ outerpat arr -> + return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr + +-- Opaques are applied to the full array, because otherwise they can +-- drastically inhibit parallelisation in some cases. +maybeDistributeStm bnd@(Let (Pattern [] [pe]) aux (BasicOp Opaque{})) acc + | not $ primType $ typeOf pe = + distributeSingleUnaryStm acc bnd $ \_ outerpat arr -> + return $ oneStm $ Let outerpat aux $ BasicOp $ Copy arr + +maybeDistributeStm bnd@(Let _ aux (BasicOp (Rearrange perm _))) acc = + distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do + let r = length (snd nest) + 1 + perm' = [0..r-1] ++ map (+r) perm + -- We need to add a copy, because the original map nest + -- will have produced an array without aliases, and so must we. + arr' <- newVName $ baseString arr + arr_t <- lookupType arr + return $ stmsFromList + [Let (Pattern [] [PatElem arr' arr_t]) aux $ BasicOp $ Copy arr, + Let outerpat aux $ BasicOp $ Rearrange perm' arr'] + +maybeDistributeStm bnd@(Let _ aux (BasicOp (Reshape reshape _))) acc = + distributeSingleUnaryStm acc bnd $ \nest outerpat arr -> do + let reshape' = map DimNew (kernelNestWidths nest) ++ + map DimNew (newDims reshape) + return $ oneStm $ Let outerpat aux $ BasicOp $ Reshape reshape' arr + +maybeDistributeStm stm@(Let _ aux (BasicOp (Rotate rots _))) acc = + distributeSingleUnaryStm acc stm $ \nest outerpat arr -> do + let rots' = map (const $ intConst Int32 0) (kernelNestWidths nest) ++ rots + return $ oneStm $ Let outerpat aux $ BasicOp $ Rotate rots' arr + +-- XXX? This rule is present to avoid the case where an in-place +-- update is distributed as its own kernel, as this would mean thread +-- then writes the entire array that it updated. This is problematic +-- because the in-place updates is O(1), but writing the array is +-- O(n). It is OK if the in-place update is preceded, followed, or +-- nested inside a sequential loop or similar, because that will +-- probably be O(n) by itself. As a hack, we only distribute if there +-- does not appear to be a loop following. The better solution is to +-- depend on memory block merging for this optimisation, but it is not +-- ready yet. +maybeDistributeStm (Let pat aux (BasicOp (Update arr [DimFix i] v))) acc + | [t] <- patternTypes pat, + arrayRank t == 1, + not $ any (amortises . stmExp) $ kernelStms acc = do + let w = arraySize 0 t + et = stripArray 1 t + lam = Lambda { lambdaParams = [] + , lambdaReturnType = [Prim int32, et] + , lambdaBody = mkBody mempty [i, v] } + maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int32 1) lam [] [(w, 1, arr)]) acc + where amortises DoLoop{} = True + amortises Op{} = True + amortises _ = False + +maybeDistributeStm stm@(Let _ aux (BasicOp (Concat d x xs w))) acc = + distributeSingleStm acc stm >>= \case + Just (kernels, _, nest, acc') -> + localScope (typeEnvFromKernelAcc acc') $ + segmentedConcat nest >>= + kernelOrNot (stmAuxCerts aux) stm acc kernels acc' + _ -> + addStmToKernel stm acc + + where segmentedConcat nest = + isSegmentedOp nest [0] w [] mempty mempty [] (x:xs) $ + \pat _ _ _ _ _ _ (x':xs') _ -> + let d' = d + length (snd nest) + 1 + in addStm $ Let pat aux $ BasicOp $ Concat d' x' xs' w + +maybeDistributeStm bnd acc = + addStmToKernel bnd acc + +distributeSingleUnaryStm :: KernelAcc + -> Stm + -> (KernelNest -> Pattern -> VName -> KernelM (Stms Out.Kernels)) + -> KernelM KernelAcc +distributeSingleUnaryStm acc bnd f = + distributeSingleStm acc bnd >>= \case + Just (kernels, res, nest, acc') + | res == map Var (patternNames $ stmPattern bnd), + (outer, inners) <- nest, + [(arr_p, arr)] <- loopNestingParamsAndArrs outer, + boundInKernelNest nest `S.intersection` freeInStm bnd + == S.singleton (paramName arr_p) -> do + addKernels kernels + let outerpat = loopNestingPattern $ fst nest + localScope (typeEnvFromKernelAcc acc') $ do + (arr', pre_stms) <- repeatMissing arr (outer:inners) + f_stms <- inScopeOf pre_stms $ f nest outerpat arr' + addKernel $ pre_stms <> f_stms + return acc' + _ -> addStmToKernel bnd acc + where -- | For an imperfectly mapped array, repeat the missing + -- dimensions to make it look like it was in fact perfectly + -- mapped. + repeatMissing arr inners = do + arr_t <- lookupType arr + let shapes = determineRepeats arr arr_t inners + if all (==Shape []) shapes then return (arr, mempty) + else do + let (outer_shapes, inner_shape) = repeatShapes shapes arr_t + arr_t' = repeatDims outer_shapes inner_shape arr_t + arr' <- newVName $ baseString arr + return (arr', oneStm $ Let (Pattern [] [PatElem arr' arr_t']) (defAux ()) $ + BasicOp $ Repeat outer_shapes inner_shape arr) + + determineRepeats arr arr_t nests + | (skipped, arr_nest:nests') <- break (hasInput arr) nests, + [(arr_p, _)] <- loopNestingParamsAndArrs arr_nest = + Shape (map loopNestingWidth skipped) : + determineRepeats (paramName arr_p) (rowType arr_t) nests' + | otherwise = + Shape (map loopNestingWidth nests) : replicate (arrayRank arr_t) (Shape []) + + hasInput arr nest + | [(_, arr')] <- loopNestingParamsAndArrs nest, arr' == arr = True + | otherwise = False + + +distribute :: KernelAcc -> KernelM KernelAcc +distribute acc = + fromMaybe acc <$> distributeIfPossible acc + +distributeIfPossible :: KernelAcc -> KernelM (Maybe KernelAcc) +distributeIfPossible acc = do + nest <- asks kernelNest + tryDistribute nest (kernelTargets acc) (kernelStms acc) >>= \case + Nothing -> return Nothing + Just (targets, kernel) -> do + addKernel kernel + return $ Just KernelAcc { kernelTargets = targets + , kernelStms = mempty + } + +distributeSingleStm :: KernelAcc -> Stm + -> KernelM (Maybe (PostKernels, Result, KernelNest, KernelAcc)) +distributeSingleStm acc bnd = do + nest <- asks kernelNest + tryDistribute nest (kernelTargets acc) (kernelStms acc) >>= \case + Nothing -> return Nothing + Just (targets, distributed_bnds) -> + tryDistributeStm nest targets bnd >>= \case + Nothing -> return Nothing + Just (res, targets', new_kernel_nest) -> + return $ Just (PostKernels [PostKernel distributed_bnds], + res, + new_kernel_nest, + KernelAcc { kernelTargets = targets' + , kernelStms = mempty + }) + +segmentedScatterKernel :: KernelNest + -> [Int] + -> Pattern + -> Certificates + -> SubExp + -> InKernelLambda + -> [VName] -> [(SubExp,Int,VName)] + -> KernelM KernelsStms +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 + -- scan. + -- + -- First, pretend that the scatter is also part of the nesting. The + -- KernelNest we produce here is technically not sensible, but it's + -- good enough for flatKernel to work. + let nest' = pushInnerKernelNesting (scatter_pat, bodyResult $ lambdaBody lam) + (MapNesting scatter_pat cs scatter_w $ zip (lambdaParams lam) ivs) nest + (nest_bnds, w, ispace, kernel_inps, _rets) <- flatKernel nest' + + let (as_ws, as_ns, as) = unzip3 dests + + -- The input/output arrays ('as') _must_ correspond to some kernel + -- input, or else the original nested scatter would have been + -- ill-typed. Find them. + as_inps <- mapM (findInput kernel_inps) as + + runBinder_ $ do + addStms nest_bnds + + let rts = concatMap (take 1) $ chunks as_ns $ + drop (sum as_ns) $ lambdaReturnType lam + (is,vs) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam + k_body = KernelBody () (bodyStms $ lambdaBody lam) $ + map (inPlaceReturn ispace) $ + zip3 as_ws as_inps $ chunks as_ns $ zip is vs + + (k_bnds, k) <- + mapKernel w (FlatThreadSpace ispace) kernel_inps rts k_body + + addStms k_bnds + + let pat = Pattern [] $ rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + + certifying cs $ letBind_ pat $ Op k + where findInput kernel_inps a = + maybe bad return $ find ((==a) . kernelInputName) kernel_inps + bad = fail "Ill-typed nested scatter encountered." + + inPlaceReturn ispace (aw, inp, is_vs) = + WriteReturn (init ws++[aw]) (kernelInputArray inp) + [ (map Var (init gtids)++[i], v) | (i,v) <- is_vs ] + where (gtids,ws) = unzip ispace + +segmentedGenReduceKernel :: KernelNest + -> [Int] + -> Certificates + -> SubExp + -> [GenReduceOp SOACS] + -> InKernelLambda + -> [VName] + -> KernelM KernelsStms +segmentedGenReduceKernel nest perm cs genred_w ops lam arrs = do + -- We replicate some of the checking done by 'isSegmentedOp', but + -- things are different because a GenReduce is not a reduction or + -- scan. + (nest_stms, _, ispace, inputs, _rets) <- flatKernel nest + let orig_pat = Pattern [] $ rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + path <- asks kernelPath + -- The input/output arrays _must_ correspond to some kernel input, + -- or else the original nested GenReduce would have been ill-typed. + -- Find them. + ops' <- forM ops $ \(GenReduceOp num_bins dests nes op) -> + GenReduceOp num_bins + <$> mapM (fmap kernelInputArray . findInput inputs) dests + <*> pure nes + <*> pure op + -- We should also remove those from the kernel nest, as otherwise + -- the generated code may be ill-typed (referencing a consumed + -- array). They will not be used anywhere else (due to uniqueness + -- constraints), so this is safe. + let all_dests = concatMap genReduceDest ops' + (nest_stms<>) <$> + inScopeOf nest_stms + (genReduceKernel path (kernelNestLoops $ removeArraysFromNest all_dests nest) + orig_pat ispace inputs cs genred_w ops' lam arrs) + where findInput kernel_inps a = + maybe bad return $ find ((==a) . kernelInputName) kernel_inps + bad = fail "Ill-typed nested GenReduce encountered." + +genReduceKernel :: (HasScope Out.Kernels m, MonadFreshNames m) => + KernelPath -> [LoopNesting] + -> Pattern -> [(VName, SubExp)] -> [KernelInput] + -> Certificates -> SubExp -> [GenReduceOp SOACS] + -> InKernelLambda -> [VName] + -> m KernelsStms +genReduceKernel path nests orig_pat ispace inputs cs genred_w ops lam arrs = do + ops' <- forM ops $ \(GenReduceOp num_bins dests nes op) -> + GenReduceOp num_bins dests nes <$> Kernelise.transformLambda op + + let isDest = flip elem $ concatMap genReduceDest ops' + inputs' = filter (not . isDest . kernelInputArray) inputs + + runBinder_ $ do + (histos, k_stms) <- blockedGenReduce genred_w ispace inputs' ops' lam arrs + + addStms $ fmap (certify cs) k_stms + + let histos' = chunks (map (length . genReduceDest) ops') histos + pes = chunks (map (length . genReduceDest) ops') $ patternElements orig_pat + + mapM_ combineIntermediateResults (zip3 pes ops histos') + + where depth = length nests + + combineIntermediateResults (pes, GenReduceOp num_bins _ nes op, histos) = do + num_histos <- arraysSize depth <$> mapM lookupType histos + + -- Avoid the segmented reduction if num_histos is 1. + num_histos_is_one <- + letSubExp "num_histos_is_one" $ + BasicOp $ CmpOp (CmpEq int32) num_histos $ intConst Int32 1 + + body_with_reshape <- runBodyBinder $ + fmap resultBody $ forM histos $ \histo -> do + histo_dims <- arrayDims <$> lookupType histo + -- Drop the num_histos dimension dimension. + let final_dims = take depth histo_dims ++ drop (depth+1) histo_dims + letSubExp "histo_flattened" $ BasicOp $ Reshape (map DimNew final_dims) histo + + -- Move the num_histos dimension innermost wrt. segments and bins. + histos_tr <- forM histos $ \h -> do + h_t <- lookupType h + let histo_perm = [0..depth-1] ++ [depth+1,depth] ++ [depth+2..arrayRank h_t-1] + letExp (baseString h <> "_tr") $ BasicOp $ Rearrange histo_perm h + histos_tr_t <- mapM lookupType histos_tr + + op_renamed <- renameLambda op + map_params <- forM (lambdaReturnType op) $ \t -> + newParam "bin" $ t `arrayOfRow` num_histos + (map_res, map_stms) <- runBinder $ do + form <- reduceSOAC Commutative op_renamed nes + letTupExp "bin_combined" $ Op $ + Screma num_histos form $ map paramName map_params + inner_segred_pat <- fmap (Pattern []) <$> forM pes $ \pe -> + PatElem <$> newVName "inner_segred" <*> + pure (stripArray depth $ patElemType pe) + nests' <- + moreArrays (map paramName map_params) histos_tr_t histos_tr $ + nests ++ [MapNesting inner_segred_pat cs num_bins $ zip (lambdaParams lam) arrs] + let collapse_body = reconstructMapNest nests' (map (rowType . patElemType) pes) $ + mkBody map_stms $ map Var map_res + + scope <- askScope + segmented_reduce_stms <- + runDistribM' $ localScope scope $ transformStms path $ + stmsToList $ bodyStms collapse_body + + let body_with_segred = mkBody segmented_reduce_stms $ + bodyResult collapse_body + letBindNames (map patElemName pes) $ + If num_histos_is_one body_with_reshape body_with_segred $ + IfAttr (staticShapes $ map patElemType pes) IfNormal + +reconstructMapNest :: [LoopNesting] -> [Type] -> BodyT SOACS -> BodyT SOACS +reconstructMapNest [] _ body = body +reconstructMapNest (MapNesting pat cs w ps_and_arrs : nests) ts body = + mkBody (oneStm $ Let pat (StmAux cs ()) $ Op $ Screma w (mapSOAC map_lam) arrs) $ + map Var $ patternNames pat + where (ps, arrs) = unzip ps_and_arrs + map_lam = Lambda { lambdaReturnType = ts + , lambdaParams = ps + , lambdaBody = reconstructMapNest nests (map rowType ts) body + } + +moreArrays :: MonadFreshNames m => + [VName] -> [Type] -> [VName] -> [LoopNesting] + -> m [LoopNesting] +moreArrays _ _ _ [] = return [] +moreArrays ps more_ts more_arrs (MapNesting pat cs w ps_and_arrs : nests) = do + ps' <- case nests of [] -> return $ zipWith Param ps row_ts + _ -> zipWithM newParam (map baseString ps) row_ts + pat' <- renamePattern pat + let outer = MapNesting pat' cs w $ ps_and_arrs ++ zip ps' more_arrs + (outer:) <$> moreArrays ps row_ts (map paramName ps') nests + where row_ts = map rowType more_ts + +segmentedScanomapKernel :: KernelNest + -> [Int] + -> SubExp + -> InKernelLambda -> InKernelLambda + -> [SubExp] -> [VName] + -> KernelM (Maybe KernelsStms) +segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = + isSegmentedOp nest perm segment_size + (lambdaReturnType map_lam) (freeInLambda lam) (freeInLambda map_lam) nes arrs $ + \pat flat_pat _num_segments total_num_elements ispace inps nes' _ arrs' -> do + regularSegmentedScan segment_size flat_pat total_num_elements + lam map_lam ispace inps nes' arrs' + + forM_ (zip (patternValueElements pat) (patternNames flat_pat)) $ + \(dst_pat_elem, flat) -> do + let ident = patElemIdent dst_pat_elem + dims = arrayDims $ identType ident + addStm $ mkLet [] [ident] $ BasicOp $ Reshape (map DimNew dims) flat + +regularSegmentedRedomapKernel :: KernelNest + -> [Int] + -> SubExp -> Commutativity + -> InKernelLambda -> InKernelLambda -> [SubExp] -> [VName] + -> KernelM (Maybe KernelsStms) +regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs = + isSegmentedOp nest perm segment_size + (lambdaReturnType map_lam) (freeInLambda lam) (freeInLambda map_lam) nes arrs $ + \pat flat_pat num_segments total_num_elements ispace inps nes' _ arrs' -> do + fold_lam <- composeLambda nilFn lam map_lam + regularSegmentedRedomap + segment_size num_segments (kernelNestWidths nest) + flat_pat pat total_num_elements comm lam fold_lam ispace inps nes' arrs' + +isSegmentedOp :: KernelNest + -> [Int] + -> SubExp + -> [Type] + -> Names -> Names + -> [SubExp] -> [VName] + -> (Pattern + -> Pattern + -> SubExp + -> SubExp + -> [(VName, SubExp)] + -> [KernelInput] + -> [SubExp] -> [VName] -> [VName] + -> Binder Out.Kernels ()) + -> KernelM (Maybe KernelsStms) +isSegmentedOp nest perm segment_size ret 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 + -- free in the op may be bound by the nest. Furthermore, the + -- neutral elements must be free in the loop nest. + -- + -- We must summarise any names from free_in_op that are bound in the + -- nest, and describe how to obtain them given segment indices. + + let bound_by_nest = boundInKernelNest nest + + (pre_bnds, nesting_size, ispace, kernel_inps, _rets) <- flatKernel nest + + unless (S.null $ free_in_op `S.intersection` bound_by_nest) $ + fail "Non-fold lambda uses nest-bound parameters." + + let indices = map fst ispace + + prepareNe (Var v) | v `S.member` bound_by_nest = + fail "Neutral element bound in nest" + prepareNe ne = return ne + + prepareArr arr = + case find ((==arr) . kernelInputName) kernel_inps of + Just inp + | kernelInputIndices inp == map Var indices -> + return $ return $ kernelInputArray inp + | not (kernelInputArray inp `S.member` bound_by_nest) -> + return $ replicateMissing ispace inp + Nothing | not (arr `S.member` bound_by_nest) -> + -- This input is something that is free inside + -- the loop nesting. We will have to replicate + -- it. + return $ + letExp (baseString arr ++ "_repd") + (BasicOp $ Replicate (Shape [nesting_size]) $ Var arr) + _ -> + fail "Input not free or outermost." + + nes' <- mapM prepareNe nes + + mk_arrs <- mapM prepareArr arrs + + lift $ runBinder_ $ do + addStms pre_bnds + + -- We must make sure all inputs are of size + -- segment_size*nesting_size. + total_num_elements <- + letSubExp "total_num_elements" $ BasicOp $ BinOp (Mul Int32) segment_size nesting_size + + let flatten arr = do + arr_shape <- arrayShape <$> lookupType arr + -- CHECKME: is the length the right thing here? We want to + -- reproduce the parameter type. + let reshape = reshapeOuter [DimNew total_num_elements] + (2+length (snd nest)) arr_shape + letExp (baseString arr ++ "_flat") $ + BasicOp $ Reshape reshape arr + + nested_arrs <- sequence mk_arrs + arrs' <- mapM flatten nested_arrs + + let pat = Pattern [] $ rearrangeShape perm $ + patternValueElements $ loopNestingPattern $ fst nest + flatPatElem pat_elem t = do + let t' = arrayOfRow t total_num_elements + name <- newVName $ baseString (patElemName pat_elem) ++ "_flat" + return $ PatElem name t' + flat_pat <- Pattern [] <$> zipWithM flatPatElem (patternValueElements pat) ret + + m pat flat_pat nesting_size total_num_elements ispace kernel_inps nes' nested_arrs arrs' + + where replicateMissing ispace inp = do + t <- lookupType $ kernelInputArray inp + let inp_is = kernelInputIndices inp + shapes = determineRepeats ispace inp_is + (outer_shapes, inner_shape) = repeatShapes shapes t + letExp "repeated" $ BasicOp $ + Repeat outer_shapes inner_shape $ kernelInputArray inp + + determineRepeats ispace (i:is) + | (skipped_ispace, ispace') <- span ((/=i) . Var . fst) ispace = + Shape (map snd skipped_ispace) : determineRepeats (drop 1 ispace') is + determineRepeats ispace _ = + [Shape $ map snd ispace] + +permutationAndMissing :: Pattern -> [SubExp] -> Maybe ([Int], [PatElem]) +permutationAndMissing pat res = do + let pes = patternValueElements pat + (_used,unused) = + partition ((`S.member` freeIn res) . patElemName) pes + res_expanded = res ++ map (Var . patElemName) unused + perm <- map (Var . patElemName) pes `isPermutationOf` res_expanded + return (perm, unused) + +-- Add extra pattern elements to every kernel nesting level. +expandKernelNest :: MonadFreshNames m => + [PatElem] -> KernelNest -> m KernelNest +expandKernelNest pes (outer_nest, inner_nests) = do + let outer_size = loopNestingWidth outer_nest : + map loopNestingWidth inner_nests + inner_sizes = tails $ map loopNestingWidth inner_nests + outer_nest' <- expandWith outer_nest outer_size + inner_nests' <- zipWithM expandWith inner_nests inner_sizes + return (outer_nest', inner_nests') + where expandWith nest dims = do + pes' <- mapM (expandPatElemWith dims) pes + return nest { loopNestingPattern = + Pattern [] $ + patternElements (loopNestingPattern nest) <> pes' + } + + expandPatElemWith dims pe = do + name <- newVName $ baseString $ patElemName pe + return pe { patElemName = name + , patElemAttr = patElemType pe `arrayOfShape` Shape dims + } + +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) +kernelAlternatives pat default_body [] = runBinder_ $ do + ses <- bodyBind default_body + forM_ (zip (patternNames pat) ses) $ \(name, se) -> + letBindNames_ [name] $ BasicOp $ SubExp se +kernelAlternatives pat default_body ((cond,alt):alts) = runBinder_ $ do + alts_pat <- fmap (Pattern []) $ forM (patternElements pat) $ \pe -> do + name <- newVName $ baseString $ patElemName pe + return pe { patElemName = name } + + alt_stms <- kernelAlternatives alts_pat default_body alts + let alt_body = mkBody alt_stms $ map Var $ patternValueNames alts_pat + + letBind_ pat $ If cond alt alt_body $ ifCommon $ patternTypes pat + +kernelOrNot :: Certificates -> Stm -> KernelAcc + -> PostKernels -> KernelAcc -> Maybe KernelsStms + -> KernelM KernelAcc +kernelOrNot cs bnd acc _ _ Nothing = + addStmToKernel (certify cs bnd) acc +kernelOrNot cs _ _ kernels acc' (Just bnds) = do + addKernels kernels + addKernel $ fmap (certify cs) bnds + return acc' diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs @@ -0,0 +1,1069 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +module Futhark.Pass.ExtractKernels.BlockedKernel + ( blockedReduction + , blockedReductionStream + , blockedGenReduce + , blockedMap + , blockedScan + + , mapKernel + , mapKernelFromBody + , KernelInput(..) + , readKernelInput + + -- Helper functions shared with at least Segmented.hs + , kerneliseLambda + , newKernelSpace + , chunkLambda + , splitArrays + , getSize + , cmpSizeLe + ) + where + +import Control.Monad +import Data.Maybe +import Data.List +import Data.Semigroup ((<>)) +import qualified Data.Set as S + +import Prelude hiding (quot) + +import Futhark.Analysis.PrimExp +import Futhark.Representation.AST +import Futhark.Representation.Kernels + hiding (Prog, Body, Stm, Pattern, PatElem, + BasicOp, Exp, Lambda, FunDef, FParam, LParam, RetType) +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Transform.Rename +import qualified Futhark.Pass.ExtractKernels.Kernelise as Kernelise +import Futhark.Representation.AST.Attributes.Aliases +import qualified Futhark.Analysis.Alias as Alias +import Futhark.Representation.SOACS.SOAC (composeLambda, Scan, Reduce, nilFn, GenReduceOp(..)) +import Futhark.Util +import Futhark.Util.IntegralExp + +getSize :: (MonadBinder m, Op (Lore m) ~ Kernel innerlore) => + String -> SizeClass -> m SubExp +getSize desc size_class = do + size_key <- newVName desc + letSubExp desc $ Op $ GetSize size_key size_class + +cmpSizeLe :: (MonadBinder m, Op (Lore m) ~ Kernel innerlore) => + String -> SizeClass -> SubExp -> m (SubExp, VName) +cmpSizeLe desc size_class to_what = do + size_key <- newVName desc + cmp_res <- letSubExp desc $ Op $ CmpSizeLe size_key size_class to_what + return (cmp_res, size_key) + +blockedReductionStream :: (MonadFreshNames m, HasScope Kernels m) => + Pattern Kernels + -> SubExp + -> Commutativity + -> Lambda InKernel -> Lambda InKernel + -> [(VName, SubExp)] -> [SubExp] -> [VName] + -> m (Stms Kernels) +blockedReductionStream pat w comm reduce_lam fold_lam ispace nes arrs = runBinder_ $ do + (max_step_one_num_groups, step_one_size) <- blockedKernelSize =<< asIntS Int64 w + + let one = constant (1 :: Int32) + num_chunks = kernelWorkgroups step_one_size + + let (acc_idents, arr_idents) = splitAt (length nes) $ patternIdents pat + step_one_pat <- basicPattern [] <$> + ((++) <$> + mapM (mkIntermediateIdent num_chunks) acc_idents <*> + pure arr_idents) + let (_fold_chunk_param, _fold_acc_params, _fold_inp_params) = + partitionChunkedFoldParameters (length nes) $ lambdaParams fold_lam + + fold_lam' <- kerneliseLambda nes fold_lam + + my_index <- newVName "my_index" + other_index <- newVName "other_index" + let my_index_param = Param my_index (Prim int32) + other_index_param = Param other_index (Prim int32) + reduce_lam' = reduce_lam { lambdaParams = my_index_param : + other_index_param : + lambdaParams reduce_lam + } + params_to_arrs = zip (map paramName $ drop 1 $ lambdaParams fold_lam') arrs + consumedArray v = fromMaybe v $ lookup v params_to_arrs + consumed_in_fold = + S.map consumedArray $ consumedByLambda $ Alias.analyseLambda fold_lam + + arrs_copies <- forM arrs $ \arr -> + if arr `S.member` consumed_in_fold then + letExp (baseString arr <> "_copy") $ BasicOp $ Copy arr + else return arr + + step_one <- chunkedReduceKernel w step_one_size comm reduce_lam' fold_lam' + ispace nes arrs_copies + addStm =<< renameStm (Let step_one_pat (defAux ()) $ Op step_one) + + step_two_pat <- basicPattern [] <$> + mapM (mkIntermediateIdent $ constant (1 :: Int32)) acc_idents + + let step_two_size = KernelSize one max_step_one_num_groups one num_chunks max_step_one_num_groups + + step_two <- reduceKernel step_two_size reduce_lam' nes $ take (length nes) $ patternNames step_one_pat + + addStm $ Let step_two_pat (defAux ()) $ Op step_two + + forM_ (zip (patternIdents step_two_pat) (patternIdents pat)) $ \(arr, x) -> + addStm $ mkLet [] [x] $ BasicOp $ Index (identName arr) $ + fullSlice (identType arr) [DimFix $ constant (0 :: Int32)] + where mkIntermediateIdent chunk_size ident = + newIdent (baseString $ identName ident) $ + arrayOfRow (identType ident) chunk_size + +chunkedReduceKernel :: (MonadBinder m, Lore m ~ Kernels) => + SubExp + -> KernelSize + -> Commutativity + -> Lambda InKernel -> Lambda InKernel + -> [(VName, SubExp)] -> [SubExp] -> [VName] + -> m (Kernel InKernel) +chunkedReduceKernel w step_one_size comm reduce_lam' fold_lam' ispace nes arrs = do + let ordering = case comm of Commutative -> Disorder + Noncommutative -> InOrder + group_size = kernelWorkgroupSize step_one_size + num_nonconcat = length nes + + space <- newKernelSpace (kernelWorkgroups step_one_size, group_size, kernelNumThreads step_one_size) $ FlatThreadSpace ispace + ((chunk_red_pes, chunk_map_pes), chunk_and_fold) <- + runBinder $ blockedPerThread (spaceGlobalId space) + w step_one_size ordering fold_lam' num_nonconcat arrs + let red_ts = map patElemType chunk_red_pes + map_ts = map (rowType . patElemType) chunk_map_pes + ts = red_ts ++ map_ts + ordering' = + case ordering of InOrder -> SplitContiguous + Disorder -> SplitStrided $ kernelNumThreads step_one_size + + chunk_red_pes' <- forM red_ts $ \red_t -> do + pe_name <- newVName "chunk_fold_red" + return $ PatElem pe_name $ red_t `arrayOfRow` group_size + combine_reds <- forM (zip chunk_red_pes' chunk_red_pes) $ \(pe', pe) -> do + combine_id <- newVName "combine_id" + return $ Let (Pattern [] [pe']) (defAux ()) $ Op $ + Combine (combineSpace [(combine_id, group_size)]) [patElemType pe] [] $ + Body () mempty [Var $ patElemName pe] + + final_red_pes <- forM (lambdaReturnType reduce_lam') $ \t -> do + pe_name <- newVName "final_result" + return $ PatElem pe_name t + let reduce_chunk = Let (Pattern [] final_red_pes) (defAux ()) $ Op $ + GroupReduce group_size reduce_lam' $ + zip nes $ map patElemName chunk_red_pes' + + red_rets <- forM final_red_pes $ \pe -> + return $ ThreadsReturn OneResultPerGroup $ Var $ patElemName pe + elems_per_thread <- asIntS Int32 $ kernelElementsPerThread step_one_size + map_rets <- forM chunk_map_pes $ \pe -> + return $ ConcatReturns ordering' w elems_per_thread Nothing $ patElemName pe + let rets = red_rets ++ map_rets + + return $ Kernel (KernelDebugHints "chunked_reduce" [("input size", w)]) space ts $ + KernelBody () (chunk_and_fold<>stmsFromList combine_reds<>oneStm reduce_chunk) rets + +reduceKernel :: (MonadBinder m, Lore m ~ Kernels) => + KernelSize + -> Lambda InKernel + -> [SubExp] + -> [VName] + -> m (Kernel InKernel) +reduceKernel step_two_size reduce_lam' nes arrs = do + let group_size = kernelWorkgroupSize step_two_size + red_ts = lambdaReturnType reduce_lam' + space <- newKernelSpace (kernelWorkgroups step_two_size, group_size, kernelNumThreads step_two_size) $ + FlatThreadSpace [] + let thread_id = spaceGlobalId space + + (rets, kstms) <- runBinder $ localScope (scopeOfKernelSpace space) $ do + in_bounds <- letSubExp "in_bounds" $ BasicOp $ CmpOp (CmpSlt Int32) + (Var $ spaceLocalId space) + (kernelTotalElements step_two_size) + + combine_body <- runBodyBinder $ + fmap resultBody $ forM (zip arrs nes) $ \(arr, ne) -> do + arr_t <- lookupType arr + letSubExp "elem" =<< + eIf (eSubExp in_bounds) + (eBody [pure $ BasicOp $ Index arr $ + fullSlice arr_t [DimFix (Var thread_id)]]) + (resultBodyM [ne]) + + combine_pat <- fmap (Pattern []) $ forM (zip arrs red_ts) $ \(arr, red_t) -> do + arr' <- newVName $ baseString arr ++ "_combined" + return $ PatElem arr' $ red_t `arrayOfRow` group_size + + combine_id <- newVName "combine_id" + letBind_ combine_pat $ + Op $ Combine (combineSpace [(combine_id, group_size)]) + (map rowType $ patternTypes combine_pat) [] combine_body + + let arrs' = patternNames combine_pat + + final_res_pes <- forM (lambdaReturnType reduce_lam') $ \t -> do + pe_name <- newVName "final_result" + return $ PatElem pe_name t + letBind_ (Pattern [] final_res_pes) $ + Op $ GroupReduce group_size reduce_lam' $ zip nes arrs' + + forM final_res_pes $ \pe -> + return $ ThreadsReturn OneResultPerGroup $ Var $ patElemName pe + + return $ Kernel (KernelDebugHints "reduce" []) space (lambdaReturnType reduce_lam') $ + KernelBody () kstms rets + +-- | Requires a fold lambda that includes accumulator parameters. +chunkLambda :: (MonadFreshNames m, HasScope Kernels m) => + Pattern Kernels -> [SubExp] -> Lambda InKernel -> m (Lambda InKernel) +chunkLambda pat nes fold_lam = do + chunk_size <- newVName "chunk_size" + + let arr_idents = drop (length nes) $ patternIdents pat + (fold_acc_params, fold_arr_params) = + splitAt (length nes) $ lambdaParams fold_lam + chunk_size_param = Param chunk_size (Prim int32) + arr_chunk_params <- mapM (mkArrChunkParam $ Var chunk_size) fold_arr_params + + map_arr_params <- forM arr_idents $ \arr -> + newParam (baseString (identName arr) <> "_in") $ + setOuterSize (identType arr) (Var chunk_size) + + fold_acc_params' <- forM fold_acc_params $ \p -> + newParam (baseString $ paramName p) $ paramType p + + let seq_rt = + let (acc_ts, arr_ts) = + splitAt (length nes) $ lambdaReturnType fold_lam + in acc_ts ++ map (`arrayOfRow` Var chunk_size) arr_ts + + res_idents = zipWith Ident (patternValueNames pat) seq_rt + + param_scope = + scopeOfLParams $ fold_acc_params' ++ arr_chunk_params ++ map_arr_params + + seq_loop_stms <- + runBinder_ $ localScope param_scope $ + Kernelise.groupStreamMapAccumL + (patternElements (basicPattern [] res_idents)) + (Var chunk_size) fold_lam (map (Var . paramName) fold_acc_params') + (map paramName arr_chunk_params) + + let seq_body = mkBody seq_loop_stms $ map (Var . identName) res_idents + + return Lambda { lambdaParams = chunk_size_param : + fold_acc_params' ++ + arr_chunk_params ++ + map_arr_params + , lambdaReturnType = seq_rt + , lambdaBody = seq_body + } + where mkArrChunkParam chunk_size arr_param = + newParam (baseString (paramName arr_param) <> "_chunk") $ + arrayOfRow (paramType arr_param) chunk_size + +-- | Given a chunked fold lambda that takes its initial accumulator +-- value as parameters, bind those parameters to the neutral element +-- instead. +kerneliseLambda :: MonadFreshNames m => + [SubExp] -> Lambda InKernel -> m (Lambda InKernel) +kerneliseLambda nes lam = do + thread_index <- newVName "thread_index" + let thread_index_param = Param thread_index $ Prim int32 + (fold_chunk_param, fold_acc_params, fold_inp_params) = + partitionChunkedFoldParameters (length nes) $ lambdaParams lam + + mkAccInit p (Var v) + | not $ primType $ paramType p = + mkLet [] [paramIdent p] $ BasicOp $ Copy v + mkAccInit p x = mkLet [] [paramIdent p] $ BasicOp $ SubExp x + acc_init_bnds = stmsFromList $ zipWith mkAccInit fold_acc_params nes + return lam { lambdaBody = insertStms acc_init_bnds $ + lambdaBody lam + , lambdaParams = thread_index_param : + fold_chunk_param : + fold_inp_params + } + +blockedReduction :: (MonadFreshNames m, HasScope Kernels m) => + Pattern Kernels + -> SubExp + -> Commutativity + -> Lambda InKernel -> Lambda InKernel + -> [(VName, SubExp)] -> [SubExp] -> [VName] + -> m (Stms Kernels) +blockedReduction pat w comm reduce_lam map_lam ispace nes arrs = runBinder_ $ do + fold_lam <- composeLambda nilFn reduce_lam map_lam + fold_lam' <- chunkLambda pat nes fold_lam + + let arr_idents = drop (length nes) $ patternIdents pat + map_out_arrs <- forM arr_idents $ \(Ident name t) -> + letExp (baseString name <> "_out_in") $ + BasicOp $ Scratch (elemType t) (arrayDims t) + + addStms =<< + blockedReductionStream pat w comm reduce_lam fold_lam' + ispace nes (arrs ++ map_out_arrs) + +blockedGenReduce :: (MonadFreshNames m, HasScope Kernels m) => + SubExp + -> [(VName,SubExp)] -- ^ Segment indexes and sizes. + -> [KernelInput] + -> [GenReduceOp InKernel] + -> Lambda InKernel -> [VName] + -> m ([VName], Stms Kernels) +blockedGenReduce arr_w segments inputs ops lam arrs = runBinder $ do + let (segment_is, segment_sizes) = unzip segments + depth = length segments + arr_w_64 <- letSubExp "arr_w_64" =<< eConvOp (SExt Int32 Int64) (toExp arr_w) + segment_sizes_64 <- mapM (letSubExp "segment_size_64" <=< eConvOp (SExt Int32 Int64) . toExp) segment_sizes + total_w <- letSubExp "genreduce_elems" =<< foldBinOp (Mul Int64) arr_w_64 segment_sizes_64 + (_, KernelSize num_groups group_size elems_per_thread_64 _ num_threads) <- + blockedKernelSize total_w + + kspace <- newKernelSpace (num_groups, group_size, num_threads) $ FlatThreadSpace [] + let ltid = spaceLocalId kspace + gtid = spaceGlobalId kspace + nthreads = spaceNumThreads kspace + + -- Determining the degree of cooperation (heuristic): + -- coop_lvl := size of histogram (Cooperation level) + -- num_histos := (threads / coop_lvl) (Number of histograms) + -- threads := min(physical_threads, segment_size) + num_histos <- forM ops $ \(GenReduceOp w _ _ _) -> + letSubExp "num_histos" =<< eDivRoundingUp Int32 (eSubExp nthreads) + (foldBinOp (Mul Int32) w segment_sizes) + + -- Initialize sub-histograms. + sub_histos <- forM (zip ops num_histos) $ \(GenReduceOp w dests nes _, num_histos') -> do + -- If num_histos' 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. This + -- only works if neither the Reshape nor the If results in a copy. + let num_histos_is_one = BasicOp $ CmpOp (CmpEq int32) num_histos' $ intConst Int32 1 + + reuse_dest = + fmap resultBody $ forM dests $ \dest -> do + (segment_dims, hist_dims) <- splitAt depth . arrayDims <$> lookupType dest + letSubExp "sub_histo" $ BasicOp $ + Reshape (map DimNew $ segment_dims ++ num_histos' : hist_dims) dest + + make_subhistograms = + -- To incorporate the original values of the genreduce target, we + -- copy those values to the first subhistogram here. + fmap resultBody $ forM (zip nes dests) $ \(ne, dest) -> do + blank <- letExp "sub_histo_blank" $ + BasicOp $ Replicate (Shape $ segment_sizes ++ [num_histos', w]) ne + let (zero, one) = (intConst Int32 0, intConst Int32 1) + slice <- fullSlice <$> lookupType blank <*> + pure (map (flip (DimSlice zero) one) segment_sizes ++ [DimFix zero]) + letSubExp "sub_histo" $ BasicOp $ Update blank slice $ Var dest + + letTupExp "histo_dests" =<< + eIf (pure num_histos_is_one) reuse_dest make_subhistograms + + let sub_histos' = concat sub_histos + dest_ts <- mapM lookupType sub_histos' + + lock_arrs <- forM (zip ops num_histos) $ \(GenReduceOp w _ _ _, num_histos') -> + letExp "locks_arr" $ BasicOp $ + Replicate (Shape $ segment_sizes ++ [num_histos', w]) (intConst Int32 0) + + (kres, kstms) <- runBinder $ localScope (scopeOfKernelSpace kspace) $ do + let toInt64 = eConvOp (SExt Int32 Int64) + i <- newVName "i" + -- The merge parameters are the histogram we are constructing. + merge_params <- zipWithM newParam (map baseString sub_histos') + (map (`toDecl` Unique) dest_ts) + group_size_64 <- letSubExp "group_size_64" =<< + toInt64 (toExp group_size) + let merge = zip merge_params $ map Var sub_histos' + form = ForLoop i Int64 elems_per_thread_64 [] + + loop_body <- runBodyBinder $ localScope (scopeOfFParams (map fst merge) <> + scopeOf form) $ do + -- Compute the offset into the input and output. To this a + -- thread can add its local ID to figure out which element it is + -- responsible for. The calculation is done with 64-bit + -- integers to avoid overflow, but the final segment indexes are + -- 32 bit. + offset <- letSubExp "offset" =<< + eBinOp (Add Int64) + (eBinOp (Mul Int64) + (toInt64 $ toExp $ spaceGroupId kspace) + (eBinOp (Mul Int64) (toExp elems_per_thread_64) (toExp group_size_64))) + (eBinOp (Mul Int64) (toExp i) (toExp group_size_64)) + + -- Construct segment indices. + j <- letSubExp "j" =<< eBinOp (Add Int64) (toExp offset) (toInt64 $ toExp ltid) + l <- newVName "l" + let bindIndex v = letBindNames_ [v] <=< toExp + zipWithM_ bindIndex (segment_is++[l]) $ + map (ConvOpExp (SExt Int64 Int32)) . + unflattenIndex (map (ConvOpExp (SExt Int32 Int64) . + primExpFromSubExp int32) $ segment_sizes ++ [arr_w]) $ + primExpFromSubExp int64 j + + -- 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 in_bounds = pure $ BasicOp $ CmpOp (CmpSlt Int64) j total_w + + in_bounds_branch = do + -- Read segment inputs. + mapM_ (addStm <=< readKernelInput) inputs + + -- Read array input. + arr_elems <- forM arrs $ \a -> do + a_t <- lookupType a + let slice = fullSlice a_t [DimFix $ Var l] + letSubExp (baseString a ++ "_elem") $ BasicOp $ Index a slice + + -- Apply bucket function. + resultBody <$> eLambda lam (map eSubExp arr_elems) + + not_in_bounds_branch = + return $ resultBody $ replicate (length ops) (intConst Int32 (-1)) ++ + concatMap genReduceNeutral ops + + lam_res <- letTupExp "bucket_fun_res" =<< + eIf in_bounds in_bounds_branch not_in_bounds_branch + + let (buckets, vs) = splitAt (length ops) $ map Var lam_res + perOp :: [a] -> [[a]] + perOp = chunks $ map (length . genReduceDest) ops + + ops_res <- forM (zip6 ops (perOp $ map paramName merge_params) buckets (perOp vs) lock_arrs num_histos) $ + \(GenReduceOp dest_w _ _ comb_op, subhistos, bucket, vs', lock_arrs', num_histos') -> do + -- Compute subhistogram index for each thread. + subhisto_ind <- letSubExp "subhisto_ind" =<< + eBinOp (SDiv Int32) + (toExp gtid) + (eDivRoundingUp Int32 (toExp nthreads) (eSubExp num_histos')) + fmap (map Var) $ letTupExp "genreduce_res" $ Op $ + GroupGenReduce (segment_sizes ++ [num_histos', dest_w]) + subhistos comb_op (map Var segment_is ++ [subhisto_ind, bucket]) vs' lock_arrs' + + return $ resultBody $ concat ops_res + + result <- letTupExp "result" $ DoLoop [] merge form loop_body + return $ map KernelInPlaceReturn result + + let kbody = KernelBody () kstms kres + letTupExp "histograms" $ Op $ Kernel (KernelDebugHints "gen_reduce" []) kspace dest_ts kbody + +blockedMap :: (MonadFreshNames m, HasScope Kernels m) => + Pattern Kernels -> SubExp + -> StreamOrd -> Lambda InKernel -> [SubExp] -> [VName] + -> m (Stm Kernels, Stms Kernels) +blockedMap concat_pat w ordering lam nes arrs = runBinder $ do + (_, kernel_size) <- blockedKernelSize =<< asIntS Int64 w + let num_nonconcat = length (lambdaReturnType lam) - patternSize concat_pat + num_groups = kernelWorkgroups kernel_size + group_size = kernelWorkgroupSize kernel_size + num_threads = kernelNumThreads kernel_size + ordering' = + case ordering of InOrder -> SplitContiguous + Disorder -> SplitStrided $ kernelNumThreads kernel_size + + space <- newKernelSpace (num_groups, group_size, num_threads) (FlatThreadSpace []) + lam' <- kerneliseLambda nes lam + ((chunk_red_pes, chunk_map_pes), chunk_and_fold) <- runBinder $ + blockedPerThread (spaceGlobalId space) w kernel_size ordering lam' num_nonconcat arrs + + nonconcat_pat <- + fmap (Pattern []) $ forM (take num_nonconcat $ lambdaReturnType lam) $ \t -> do + name <- newVName "nonconcat" + return $ PatElem name $ t `arrayOfRow` num_threads + + let pat = nonconcat_pat <> concat_pat + ts = map patElemType chunk_red_pes ++ + map (rowType . patElemType) chunk_map_pes + + nonconcat_rets <- forM chunk_red_pes $ \pe -> + return $ ThreadsReturn AllThreads $ Var $ patElemName pe + elems_per_thread <- asIntS Int32 $ kernelElementsPerThread kernel_size + concat_rets <- forM chunk_map_pes $ \pe -> + return $ ConcatReturns ordering' w elems_per_thread Nothing $ patElemName pe + + return $ Let pat (defAux ()) $ Op $ Kernel (KernelDebugHints "chunked_map" []) space ts $ + KernelBody () chunk_and_fold $ nonconcat_rets ++ concat_rets + +blockedPerThread :: (MonadBinder m, Lore m ~ InKernel) => + VName -> SubExp -> KernelSize -> StreamOrd -> Lambda InKernel + -> Int -> [VName] + -> m ([PatElem InKernel], [PatElem InKernel]) +blockedPerThread thread_gtid w kernel_size ordering lam num_nonconcat arrs = do + let (_, chunk_size, [], arr_params) = + partitionChunkedKernelFoldParameters 0 $ lambdaParams lam + + ordering' = + case ordering of InOrder -> SplitContiguous + Disorder -> SplitStrided $ kernelNumThreads kernel_size + red_ts = take num_nonconcat $ lambdaReturnType lam + map_ts = map rowType $ drop num_nonconcat $ lambdaReturnType lam + + per_thread <- asIntS Int32 $ kernelElementsPerThread kernel_size + splitArrays (paramName chunk_size) (map paramName arr_params) ordering' w + (Var thread_gtid) per_thread arrs + + chunk_red_pes <- forM red_ts $ \red_t -> do + pe_name <- newVName "chunk_fold_red" + return $ PatElem pe_name red_t + chunk_map_pes <- forM map_ts $ \map_t -> do + pe_name <- newVName "chunk_fold_map" + return $ PatElem pe_name $ map_t `arrayOfRow` Var (paramName chunk_size) + + let (chunk_red_ses, chunk_map_ses) = + splitAt num_nonconcat $ bodyResult $ lambdaBody lam + + addStms $ + bodyStms (lambdaBody lam) <> + stmsFromList + [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se + | (pe,se) <- zip chunk_red_pes chunk_red_ses ] <> + stmsFromList + [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se + | (pe,se) <- zip chunk_map_pes chunk_map_ses ] + + return (chunk_red_pes, chunk_map_pes) + +splitArrays :: (MonadBinder m, Lore m ~ InKernel) => + VName -> [VName] + -> SplitOrdering -> SubExp -> SubExp -> SubExp -> [VName] + -> m () +splitArrays chunk_size split_bound ordering w i elems_per_i arrs = do + letBindNames_ [chunk_size] $ Op $ SplitSpace ordering w i elems_per_i + case ordering of + SplitContiguous -> do + offset <- letSubExp "slice_offset" $ BasicOp $ BinOp (Mul Int32) i elems_per_i + zipWithM_ (contiguousSlice offset) split_bound arrs + SplitStrided stride -> zipWithM_ (stridedSlice stride) split_bound arrs + where contiguousSlice offset slice_name arr = do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimSlice offset (Var chunk_size) (constant (1::Int32))] + letBindNames_ [slice_name] $ BasicOp $ Index arr slice + + stridedSlice stride slice_name arr = do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimSlice i (Var chunk_size) stride] + letBindNames_ [slice_name] $ BasicOp $ Index arr slice + +data KernelSize = KernelSize { kernelWorkgroups :: SubExp + -- ^ Int32 + , kernelWorkgroupSize :: SubExp + -- ^ Int32 + , kernelElementsPerThread :: SubExp + -- ^ Int64 + , kernelTotalElements :: SubExp + -- ^ Int64 + , kernelNumThreads :: SubExp + -- ^ Int32 + } + deriving (Eq, Ord, Show) + +numberOfGroups :: MonadBinder m => SubExp -> SubExp -> SubExp -> m (SubExp, SubExp) +numberOfGroups w group_size max_num_groups = do + -- If 'w' is small, we launch fewer groups than we normally would. + -- We don't want any idle groups. + w_div_group_size <- letSubExp "w_div_group_size" =<< + eDivRoundingUp Int64 (eSubExp w) (eSubExp group_size) + -- We also don't want zero groups. + num_groups_maybe_zero <- letSubExp "num_groups_maybe_zero" $ BasicOp $ BinOp (SMin Int64) + w_div_group_size max_num_groups + num_groups <- letSubExp "num_groups" $ + BasicOp $ BinOp (SMax Int64) (intConst Int64 1) + num_groups_maybe_zero + num_threads <- + letSubExp "num_threads" $ BasicOp $ BinOp (Mul Int64) num_groups group_size + return (num_groups, num_threads) + +blockedKernelSize :: (MonadBinder m, Lore m ~ Kernels) => + SubExp -> m (SubExp, KernelSize) +blockedKernelSize w = do + group_size <- getSize "group_size" SizeGroup + max_num_groups <- getSize "max_num_groups" SizeNumGroups + + group_size' <- asIntS Int64 group_size + max_num_groups' <- asIntS Int64 max_num_groups + (num_groups, num_threads) <- numberOfGroups w group_size' max_num_groups' + num_groups' <- asIntS Int32 num_groups + num_threads' <- asIntS Int32 num_threads + + per_thread_elements <- + letSubExp "per_thread_elements" =<< + eDivRoundingUp Int64 (toExp =<< asIntS Int64 w) (toExp =<< asIntS Int64 num_threads) + + return (max_num_groups, + KernelSize num_groups' group_size per_thread_elements w num_threads') + +-- First stage scan kernel. +scanKernel1 :: (MonadBinder m, Lore m ~ Kernels) => + SubExp -> KernelSize + -> Scan InKernel + -> Reduce InKernel + -> Lambda InKernel -> [VName] + -> m (Kernel InKernel) +scanKernel1 w scan_sizes (scan_lam, scan_nes) (_comm, red_lam, red_nes) foldlam arrs = do + num_elements <- asIntS Int32 $ kernelTotalElements scan_sizes + + let (scan_ts, red_ts, map_ts) = + splitAt3 (length scan_nes) (length red_nes) $ lambdaReturnType foldlam + (_, foldlam_acc_params, _) = + partitionChunkedFoldParameters (length scan_nes + length red_nes) $ lambdaParams foldlam + + -- Scratch arrays for scanout and mapout parts. + (scanout_arrs, scanout_arr_params, scanout_arr_ts) <- + unzip3 <$> mapM (mkOutArray "scanout") scan_ts + (mapout_arrs, mapout_arr_params, mapout_arr_ts) <- + unzip3 <$> mapM (mkOutArray "scanout") map_ts + + last_thread <- letSubExp "last_thread" $ BasicOp $ + BinOp (Sub Int32) group_size (constant (1::Int32)) + kspace <- newKernelSpace (num_groups, group_size, num_threads) $ FlatThreadSpace [] + let lid = spaceLocalId kspace + + (res, stms) <- runBinder $ localScope (scopeOfKernelSpace kspace) $ do + -- We create a loop that moves in group_size chunks over the input. + num_iterations <- letSubExp "num_iterations" =<< + eDivRoundingUp Int32 (eSubExp w) (eSubExp num_threads) + + -- The merge parameters are the scanout arrays, the reduction + -- results, the mapout arrays, and the (renamed) scan accumulator + -- parameters of foldlam (which function as carries). We do not + -- need to keep accumulator parameters/carries for the reduction, + -- because the reduction result suffices. + (acc_params, nes') <- unzip <$> zipWithM mkAccMergeParam foldlam_acc_params + (scan_nes ++ red_nes) + let (scan_acc_params, red_acc_params) = + splitAt (length scan_nes) acc_params + (scan_nes', red_nes') = + splitAt (length scan_nes) nes' + let merge = zip scanout_arr_params (map Var scanout_arrs) ++ + zip red_acc_params red_nes' ++ + zip mapout_arr_params (map Var mapout_arrs) ++ + zip scan_acc_params scan_nes' + i <- newVName "i" + let form = ForLoop i Int32 num_iterations [] + + loop_body <- runBodyBinder $ localScope (scopeOfFParams (map fst merge) <> + scopeOf form) $ do + -- Compute the offset into the input and output. To this a + -- thread can add its local ID to figure out which element it is + -- responsible for. + offset <- letSubExp "offset" =<< + eBinOp (Add Int32) + (eBinOp (Mul Int32) + (eSubExp $ Var $ spaceGroupId kspace) + (pure $ BasicOp $ BinOp (Mul Int32) num_iterations group_size)) + (pure $ BasicOp $ BinOp (Mul Int32) (Var i) group_size) + + -- Now we apply the fold function if j=offset+lid is less than + -- num_elements. This also involves writing to the mapout + -- arrays. + j <- letSubExp "j" $ BasicOp $ BinOp (Add Int32) offset (Var lid) + let in_bounds = pure $ BasicOp $ CmpOp (CmpSlt Int32) j num_elements + + in_bounds_fold_branch = do + -- Read array input. + arr_elems <- forM arrs $ \arr -> do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimFix j] + letSubExp (baseString arr ++ "_elem") $ BasicOp $ Index arr slice + + -- Apply the body of the fold function. + fold_res <- + eLambda foldlam $ map eSubExp $ j : map (Var . paramName) acc_params ++ arr_elems + + -- Scatter the to_map parts to the mapout arrays using + -- in-place updates, and return the to_scan parts. + let (to_scan, to_red, to_map) = splitAt3 (length scan_nes) (length red_nes) fold_res + mapout_arrs' <- forM (zip to_map mapout_arr_params) $ \(se,arr) -> do + let slice = fullSlice (paramType arr) [DimFix j] + letInPlace "mapout" (paramName arr) slice $ BasicOp $ SubExp se + return $ resultBody $ to_scan ++ to_red ++ map Var mapout_arrs' + + not_in_bounds_fold_branch = return $ resultBody $ map (Var . paramName) $ + scan_acc_params ++ red_acc_params ++ mapout_arr_params + + (to_scan_res, to_red_res, mapout_arrs') <- + fmap (splitAt3 (length scan_nes) (length red_nes)) . letTupExp "foldres" =<< + eIf in_bounds in_bounds_fold_branch not_in_bounds_fold_branch + + (scanned_arrs, scanout_arrs') <- + doScan j kspace in_bounds scanout_arr_params to_scan_res + + new_scan_carries <- + resetCarries "scan" lid scan_acc_params scan_nes' $ runBodyBinder $ do + carries <- forM scanned_arrs $ \arr -> do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimFix last_thread] + letSubExp "carry" $ BasicOp $ Index arr slice + return $ resultBody carries + + red_res <- doReduce to_red_res + + new_red_carries <- resetCarries "red" lid red_acc_params red_nes' $ + return $ resultBody $ map Var red_res + + -- HACK + new_scan_carries' <- letTupExp "new_carry_sync" $ Op $ Barrier $ map Var new_scan_carries + return $ resultBody $ map Var $ + scanout_arrs' ++ new_red_carries ++ mapout_arrs' ++ new_scan_carries' + + result <- letTupExp "result" $ DoLoop [] merge form loop_body + let (scanout_result, red_result, mapout_result, scan_carry_result) = + splitAt4 (length scan_ts) (length red_ts) (length mapout_arrs) result + return (map KernelInPlaceReturn scanout_result ++ + map (ThreadsReturn OneResultPerGroup . Var) scan_carry_result ++ + map (ThreadsReturn OneResultPerGroup . Var) red_result ++ + map KernelInPlaceReturn mapout_result) + + let kts = scanout_arr_ts ++ scan_ts ++ red_ts ++ mapout_arr_ts + kbody = KernelBody () stms res + + return $ Kernel (KernelDebugHints "scan1" []) kspace kts kbody + where num_groups = kernelWorkgroups scan_sizes + group_size = kernelWorkgroupSize scan_sizes + num_threads = kernelNumThreads scan_sizes + consumed_in_foldlam = consumedInBody $ lambdaBody $ Alias.analyseLambda foldlam + + mkOutArray desc t = do + let arr_t = t `arrayOfRow` w + arr <- letExp desc $ BasicOp $ Scratch (elemType arr_t) (arrayDims arr_t) + pname <- newVName $ desc++"param" + return (arr, Param pname $ toDecl arr_t Unique, arr_t) + + mkAccMergeParam (Param pname ptype) se = do + pname' <- newVName $ baseString pname ++ "_merge" + -- We have to copy the initial merge parameter (the neutral + -- element) if it is consumed inside the lambda. + case se of + Var v | pname `S.member` consumed_in_foldlam -> do + se' <- letSubExp "scan_ne_copy" $ BasicOp $ Copy v + return (Param pname' $ toDecl ptype Unique, + se') + _ -> return (Param pname' $ toDecl ptype Nonunique, + se) + + doScan j kspace in_bounds scanout_arr_params to_scan_res = do + let lid = spaceLocalId kspace + scan_ts = map (rowType . paramType) scanout_arr_params + -- Create an array of per-thread fold results and scan it. + combine_id <- newVName "combine_id" + to_scan_arrs <- letTupExp "combined" $ + Op $ Combine (combineSpace [(combine_id, group_size)]) scan_ts [] $ + Body () mempty $ map Var to_scan_res + scanned_arrs <- letTupExp "scanned" $ + Op $ GroupScan group_size scan_lam $ zip scan_nes to_scan_arrs + + -- If we are in bounds, we write scanned_arrs[lid] to scanout[j]. + let in_bounds_scan_branch = do + -- Read scanned_arrs[j]. + arr_elems <- forM scanned_arrs $ \arr -> do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimFix $ Var lid] + letSubExp (baseString arr ++ "_elem") $ BasicOp $ Index arr slice + + -- Scatter the to_map parts to the scanout arrays using + -- in-place updates. + scanout_arrs' <- forM (zip arr_elems scanout_arr_params) $ \(se,p) -> do + let slice = fullSlice (paramType p) [DimFix j] + letInPlace "mapout" (paramName p) slice $ BasicOp $ SubExp se + return $ resultBody $ map Var scanout_arrs' + + not_in_bounds_scan_branch = + return $ resultBody $ map (Var . paramName) scanout_arr_params + + scanres <- letTupExp "scanres" =<< + eIf in_bounds in_bounds_scan_branch not_in_bounds_scan_branch + return (scanned_arrs, scanres) + + doReduce to_red_res = do + red_ts <- mapM lookupType to_red_res + + -- Create an array of per-thread fold results and reduce it. + combine_id <- newVName "combine_id" + to_red_arrs <- letTupExp "combined" $ + Op $ Combine (combineSpace [(combine_id, group_size)]) red_ts [] $ + Body () mempty $ map Var to_red_res + letTupExp "reduced" $ + Op $ GroupReduce group_size red_lam $ zip red_nes to_red_arrs + + resetCarries what lid acc_params nes mk_read_res = do + -- All threads but the first in the group reset the accumulator + -- to the neutral element. The first resets it to the carry-out + -- of the scan or reduction. + is_first_thread <- letSubExp "is_first_thread" $ BasicOp $ + CmpOp (CmpEq int32) (Var lid) (constant (0::Int32)) + + read_res <- mk_read_res + + reset_carry_outs <- runBodyBinder $ do + carries <- forM (zip acc_params nes) $ \(p, se) -> + case se of + Var v | unique $ declTypeOf p -> + letSubExp "reset_acc_copy" $ BasicOp $ Copy v + _ -> return se + return $ resultBody carries + + letTupExp ("new_" ++ what ++ "_carry") $ + If is_first_thread read_res reset_carry_outs $ + ifCommon $ map paramType acc_params + +-- Second stage scan kernel with no fold part. +scanKernel2 :: (MonadBinder m, Lore m ~ Kernels) => + KernelSize + -> Lambda InKernel + -> [(SubExp,VName)] + -> m (Kernel InKernel) +scanKernel2 scan_sizes lam input = do + let (nes, arrs) = unzip input + scan_ts = lambdaReturnType lam + + kspace <- newKernelSpace (kernelWorkgroups scan_sizes, + group_size, + kernelNumThreads scan_sizes) (FlatThreadSpace []) + (res, stms) <- runBinder $ localScope (scopeOfKernelSpace kspace) $ do + -- Create an array of the elements we are to scan. + let indexMine cid arr = do + arr_t <- lookupType arr + let slice = fullSlice arr_t [DimFix $ Var cid] + letSubExp (baseString arr <> "_elem") $ BasicOp $ Index arr slice + combine_id <- newVName "combine_id" + read_elements <- runBodyBinder $ resultBody <$> mapM (indexMine combine_id) arrs + to_scan_arrs <- letTupExp "combined" $ + Op $ Combine (combineSpace [(combine_id, group_size)]) scan_ts [] read_elements + scanned_arrs <- letTupExp "scanned" $ + Op $ GroupScan group_size lam $ zip nes to_scan_arrs + + -- Each thread returns scanned_arrs[i]. + res_elems <- mapM (indexMine $ spaceLocalId kspace) scanned_arrs + return $ map (ThreadsReturn AllThreads) res_elems + + return $ Kernel (KernelDebugHints "scan2" []) kspace (lambdaReturnType lam) $ KernelBody () stms res + where group_size = kernelWorkgroupSize scan_sizes + +-- | The 'VName's returned are the names of variables bound to the +-- carry-out of the last thread. You can ignore them if you don't +-- need them. +blockedScan :: (MonadBinder m, Lore m ~ Kernels) => + Pattern Kernels -> SubExp + -> Scan InKernel + -> Reduce InKernel + -> Lambda InKernel -> SubExp -> [(VName, SubExp)] -> [KernelInput] + -> [VName] + -> m [VName] +blockedScan pat w (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam segment_size ispace inps arrs = do + foldlam <- composeLambda scan_lam red_lam map_lam + + (_, first_scan_size) <- blockedKernelSize =<< asIntS Int64 w + my_index <- newVName "my_index" + other_index <- newVName "other_index" + let num_groups = kernelWorkgroups first_scan_size + group_size = kernelWorkgroupSize first_scan_size + num_threads = kernelNumThreads first_scan_size + my_index_param = Param my_index (Prim int32) + other_index_param = Param other_index (Prim int32) + + let foldlam_scope = scopeOfLParams $ my_index_param : lambdaParams foldlam + bindIndex i v = letBindNames_ [i] =<< toExp v + compute_segments <- runBinder_ $ localScope foldlam_scope $ + zipWithM_ bindIndex (map fst ispace) $ + unflattenIndex (map (primExpFromSubExp int32 . snd) ispace) + (LeafExp (paramName my_index_param) int32 `quot` + primExpFromSubExp int32 segment_size) + read_inps <- stmsFromList <$> mapM readKernelInput inps + first_scan_foldlam <- renameLambda + foldlam { lambdaParams = my_index_param : + lambdaParams foldlam + , lambdaBody = insertStms (compute_segments<>read_inps) $ + lambdaBody foldlam + } + first_scan_lam <- renameLambda + scan_lam { lambdaParams = my_index_param : + other_index_param : + lambdaParams scan_lam + } + first_scan_red_lam <- renameLambda + red_lam { lambdaParams = my_index_param : + other_index_param : + lambdaParams red_lam + } + + let (scan_idents, red_idents, arr_idents) = + splitAt3 (length scan_nes) (length red_nes) $ patternIdents pat + final_res_pat = Pattern [] $ take (length scan_nes) $ patternValueElements pat + first_scan_pat <- basicPattern [] . concat <$> + sequence [mapM (mkIntermediateIdent "seq_scanned" [w]) scan_idents, + mapM (mkIntermediateIdent "scan_carry_out" [num_groups]) scan_idents, + mapM (mkIntermediateIdent "red_carry_out" [num_groups]) red_idents, + pure arr_idents] + + addStm . Let first_scan_pat (defAux ()) . Op =<< scanKernel1 w first_scan_size + (first_scan_lam, scan_nes) + (comm, first_scan_red_lam, red_nes) + first_scan_foldlam arrs + + let (sequentially_scanned, group_carry_out, group_red_res, _) = + splitAt4 (length scan_nes) (length scan_nes) (length red_nes) $ patternNames first_scan_pat + + let second_scan_size = KernelSize one num_groups one num_groups num_groups + unless (null group_red_res) $ do + second_stage_red_lam <- renameLambda first_scan_red_lam + red_res <- letTupExp "red_res" . Op =<< + reduceKernel second_scan_size second_stage_red_lam red_nes group_red_res + forM_ (zip red_idents red_res) $ \(dest, arr) -> do + arr_t <- lookupType arr + addStm $ mkLet [] [dest] $ BasicOp $ Index arr $ + fullSlice arr_t [DimFix $ constant (0 :: Int32)] + + second_scan_lam <- renameLambda first_scan_lam + + group_carry_out_scanned <- + letTupExp "group_carry_out_scanned" . Op =<< + scanKernel2 second_scan_size + second_scan_lam (zip scan_nes group_carry_out) + + last_group <- letSubExp "last_group" $ BasicOp $ BinOp (Sub Int32) num_groups one + carries <- forM group_carry_out_scanned $ \carry_outs -> do + arr_t <- lookupType carry_outs + letExp "carry_out" $ BasicOp $ Index carry_outs $ fullSlice arr_t [DimFix last_group] + + scan_lam''' <- renameLambda scan_lam + j <- newVName "j" + let (acc_params, arr_params) = + splitAt (length scan_nes) $ lambdaParams scan_lam''' + result_map_input = + zipWith (mkKernelInput [Var j]) arr_params sequentially_scanned + + chunks_per_group <- letSubExp "chunks_per_group" =<< + eDivRoundingUp Int32 (eSubExp w) (eSubExp num_threads) + elems_per_group <- letSubExp "elements_per_group" $ + BasicOp $ BinOp (Mul Int32) chunks_per_group group_size + + result_map_body <- runBodyBinder $ localScope (scopeOfLParams $ map kernelInputParam result_map_input) $ do + group_id <- + letSubExp "group_id" $ + BasicOp $ BinOp (SQuot Int32) (Var j) elems_per_group + let do_nothing = + pure $ resultBody $ map (Var . paramName) arr_params + add_carry_in = runBodyBinder $ do + forM_ (zip acc_params group_carry_out_scanned) $ \(p, arr) -> do + carry_in_index <- + letSubExp "carry_in_index" $ + BasicOp $ BinOp (Sub Int32) group_id one + arr_t <- lookupType arr + letBindNames_ [paramName p] $ + BasicOp $ Index arr $ fullSlice arr_t [DimFix carry_in_index] + return $ lambdaBody scan_lam''' + group_lasts <- + letTupExp "final_result" =<< + eIf (eCmpOp (CmpEq int32) (eSubExp zero) (eSubExp group_id)) + do_nothing + add_carry_in + return $ resultBody $ map Var group_lasts + + (mapk_bnds, mapk) <- mapKernelFromBody w (FlatThreadSpace [(j, w)]) result_map_input + (lambdaReturnType scan_lam) result_map_body + addStms mapk_bnds + letBind_ final_res_pat $ Op mapk + + return carries + where one = constant (1 :: Int32) + zero = constant (0 :: Int32) + + mkIntermediateIdent desc shape ident = + newIdent (baseString (identName ident) ++ "_" ++ desc) $ + arrayOf (rowType $ identType ident) (Shape shape) NoUniqueness + + mkKernelInput indices p arr = KernelInput { kernelInputName = paramName p + , kernelInputType = paramType p + , kernelInputArray = arr + , kernelInputIndices = indices + } + +mapKernelSkeleton :: (HasScope Kernels m, MonadFreshNames m) => + SubExp -> SpaceStructure -> [KernelInput] + -> m (KernelSpace, + Stms Kernels, + Stms InKernel) +mapKernelSkeleton w ispace inputs = do + ((group_size, num_threads, num_groups), ksize_bnds) <- + runBinder $ numThreadsAndGroups w + + read_input_bnds <- stmsFromList <$> mapM readKernelInput inputs + + let ksize = (num_groups, group_size, num_threads) + + space <- newKernelSpace ksize ispace + return (space, ksize_bnds, read_input_bnds) + +-- Given the desired minium number of threads, compute the group size, +-- number of groups and total number of threads. +numThreadsAndGroups :: (MonadBinder m, Op (Lore m) ~ Kernel innerlore) => + SubExp -> m (SubExp, SubExp, SubExp) +numThreadsAndGroups w = do + group_size <- getSize "group_size" SizeGroup + num_groups <- letSubExp "num_groups" =<< eDivRoundingUp Int32 + (eSubExp w) (eSubExp group_size) + num_threads <- letSubExp "num_threads" $ + BasicOp $ BinOp (Mul Int32) num_groups group_size + return (group_size, num_threads, num_groups) + +mapKernel :: (HasScope Kernels m, MonadFreshNames m) => + SubExp -> SpaceStructure -> [KernelInput] + -> [Type] -> KernelBody InKernel + -> m (Stms Kernels, Kernel InKernel) +mapKernel w ispace inputs rts (KernelBody () kstms krets) = do + (space, ksize_bnds, read_input_bnds) <- mapKernelSkeleton w ispace inputs + + let kbody' = KernelBody () (read_input_bnds <> kstms) krets + return (ksize_bnds, Kernel (KernelDebugHints "map" []) space rts kbody') + +mapKernelFromBody :: (HasScope Kernels m, MonadFreshNames m) => + SubExp -> SpaceStructure -> [KernelInput] + -> [Type] -> Body InKernel + -> m (Stms Kernels, Kernel InKernel) +mapKernelFromBody w ispace inputs rts body = + mapKernel w ispace inputs rts kbody + where kbody = KernelBody () (bodyStms body) krets + krets = map (ThreadsReturn ThreadsInSpace) $ bodyResult body + +data KernelInput = KernelInput { kernelInputName :: VName + , kernelInputType :: Type + , kernelInputArray :: VName + , kernelInputIndices :: [SubExp] + } + deriving (Show) + +kernelInputParam :: KernelInput -> Param Type +kernelInputParam p = Param (kernelInputName p) (kernelInputType p) + +readKernelInput :: (HasScope scope m, Monad m) => + KernelInput -> m (Stm InKernel) +readKernelInput inp = do + let pe = PatElem (kernelInputName inp) $ kernelInputType inp + arr_t <- lookupType $ kernelInputArray inp + return $ Let (Pattern [] [pe]) (defAux ()) $ + BasicOp $ Index (kernelInputArray inp) $ + fullSlice arr_t $ map DimFix $ kernelInputIndices inp + +newKernelSpace :: MonadFreshNames m => + (SubExp,SubExp,SubExp) -> SpaceStructure -> m KernelSpace +newKernelSpace (num_groups, group_size, num_threads) ispace = + KernelSpace + <$> newVName "global_tid" + <*> newVName "local_tid" + <*> newVName "group_id" + <*> pure num_threads + <*> pure num_groups + <*> pure group_size + <*> pure ispace diff --git a/src/Futhark/Pass/ExtractKernels/Distribution.hs b/src/Futhark/Pass/ExtractKernels/Distribution.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/Distribution.hs @@ -0,0 +1,539 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Futhark.Pass.ExtractKernels.Distribution + ( + Target + , Targets + , ppTargets + , singleTarget + , outerTarget + , innerTarget + , pushInnerTarget + , popInnerTarget + , targetsScope + + , LoopNesting (..) + , ppLoopNesting + + , Nesting (..) + , Nestings + , ppNestings + , letBindInInnerNesting + , singleNesting + , pushInnerNesting + + , KernelNest + , ppKernelNest + , newKernel + , pushKernelNesting + , pushInnerKernelNesting + , removeArraysFromNest + , kernelNestLoops + , kernelNestWidths + , boundInKernelNest + , boundInKernelNests + , flatKernel + , constructKernel + + , tryDistribute + , tryDistributeStm + ) + where + +import Control.Monad.RWS.Strict +import Control.Monad.Trans.Maybe +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Foldable +import Data.Maybe +import Data.List + +import Futhark.Representation.Kernels +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Util +import Futhark.Transform.Rename +import Futhark.Util.Log +import Futhark.Pass.ExtractKernels.BlockedKernel (mapKernel, KernelInput(..)) + +type Target = (Pattern Kernels, Result) + +-- | First pair element is the very innermost ("current") target. In +-- the list, the outermost target comes first. Invariant: Every +-- element of a pattern must be present as the result of the +-- immediately enclosing target. This is ensured by 'pushInnerTarget' +-- by removing unused pattern elements. +data Targets = Targets { _innerTarget :: Target + , _outerTargets :: [Target] + } + +ppTargets :: Targets -> String +ppTargets (Targets target targets) = + unlines $ map ppTarget $ targets ++ [target] + where ppTarget (pat, res) = + pretty pat ++ " <- " ++ pretty res + +singleTarget :: Target -> Targets +singleTarget = flip Targets [] + +outerTarget :: Targets -> Target +outerTarget (Targets inner_target []) = inner_target +outerTarget (Targets _ (outer_target : _)) = outer_target + +innerTarget :: Targets -> Target +innerTarget (Targets inner_target _) = inner_target + +pushOuterTarget :: Target -> Targets -> Targets +pushOuterTarget target (Targets inner_target targets) = + Targets inner_target (target : targets) + +pushInnerTarget :: Target -> Targets -> Targets +pushInnerTarget (pat, res) (Targets inner_target targets) = + Targets (pat', res') (targets ++ [inner_target]) + where (pes', res') = unzip $ filter (used . fst) $ zip (patternElements pat) res + pat' = Pattern [] pes' + inner_used = freeIn $ snd inner_target + used pe = patElemName pe `S.member` inner_used + +popInnerTarget :: Targets -> Maybe (Target, Targets) +popInnerTarget (Targets t ts) = + case reverse ts of + x:xs -> Just (t, Targets x $ reverse xs) + [] -> Nothing + +targetScope :: Target -> Scope Kernels +targetScope = scopeOfPattern . fst + +targetsScope :: Targets -> Scope Kernels +targetsScope (Targets t ts) = mconcat $ map targetScope $ t : ts + +data LoopNesting = MapNesting { loopNestingPattern :: Pattern Kernels + , loopNestingCertificates :: Certificates + , loopNestingWidth :: SubExp + , loopNestingParamsAndArrs :: [(Param Type, VName)] + } + deriving (Show) + +instance Scoped Kernels LoopNesting where + scopeOf = scopeOfLParams . map fst . loopNestingParamsAndArrs + +ppLoopNesting :: LoopNesting -> String +ppLoopNesting (MapNesting _ _ _ params_and_arrs) = + pretty (map fst params_and_arrs) ++ + " <- " ++ + pretty (map snd params_and_arrs) + +loopNestingParams :: LoopNesting -> [LParam Kernels] +loopNestingParams = map fst . loopNestingParamsAndArrs + +instance FreeIn LoopNesting where + freeIn (MapNesting pat cs w params_and_arrs) = + freeIn pat <> + freeIn cs <> + freeIn w <> + freeIn params_and_arrs + +data Nesting = Nesting { nestingLetBound :: Names + , nestingLoop :: LoopNesting + } + deriving (Show) + +letBindInNesting :: Names -> Nesting -> Nesting +letBindInNesting newnames (Nesting oldnames loop) = + Nesting (oldnames <> newnames) loop + +-- ^ First pair element is the very innermost ("current") nest. In +-- the list, the outermost nest comes first. +type Nestings = (Nesting, [Nesting]) + +ppNestings :: Nestings -> String +ppNestings (nesting, nestings) = + unlines $ map ppNesting $ nestings ++ [nesting] + where ppNesting (Nesting _ loop) = + ppLoopNesting loop + +singleNesting :: Nesting -> Nestings +singleNesting = (,[]) + +pushInnerNesting :: Nesting -> Nestings -> Nestings +pushInnerNesting nesting (inner_nesting, nestings) = + (nesting, nestings ++ [inner_nesting]) + +-- | Both parameters and let-bound. +boundInNesting :: Nesting -> Names +boundInNesting nesting = + S.fromList (map paramName (loopNestingParams loop)) <> + nestingLetBound nesting + where loop = nestingLoop nesting + +letBindInInnerNesting :: Names -> Nestings -> Nestings +letBindInInnerNesting names (nest, nestings) = + (letBindInNesting names nest, nestings) + + +-- | Note: first element is *outermost* nesting. This is different +-- from the similar types elsewhere! +type KernelNest = (LoopNesting, [LoopNesting]) + +ppKernelNest :: KernelNest -> String +ppKernelNest (nesting, nestings) = + unlines $ map ppLoopNesting $ nesting : nestings + +-- | Add new outermost nesting, pushing the current outermost to the +-- list, also taking care to swap patterns if necessary. +pushKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest +pushKernelNesting target newnest (nest, nests) = + (fixNestingPatternOrder newnest target (loopNestingPattern nest), + nest : nests) + +-- | Add new innermost nesting, pushing the current outermost to the +-- list. It is important that the 'Target' has the right order +-- (non-permuted compared to what is expected by the outer nests). +pushInnerKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest +pushInnerKernelNesting target newnest (nest, nests) = + (nest, nests ++ [fixNestingPatternOrder newnest target (loopNestingPattern innermost)]) + where innermost = case reverse nests of + [] -> nest + n:_ -> n + +fixNestingPatternOrder :: LoopNesting -> Target -> Pattern Kernels -> LoopNesting +fixNestingPatternOrder nest (_,res) inner_pat = + nest { loopNestingPattern = basicPattern [] pat' } + where pat = loopNestingPattern nest + pat' = map fst fixed_target + fixed_target = sortOn posInInnerPat $ zip (patternValueIdents pat) res + posInInnerPat (_, Var v) = fromMaybe 0 $ elemIndex v $ patternNames inner_pat + posInInnerPat _ = 0 + +-- | Remove these arrays from the outermost nesting, and all +-- uses of corresponding parameters from innermost nesting. +removeArraysFromNest :: [VName] -> KernelNest -> KernelNest +removeArraysFromNest orig_arrs (outer, inners) = + let (arrs, outer') = remove (S.fromList orig_arrs) outer + (_, inners') = mapAccumL remove arrs inners + in (outer', inners') + where remove arrs nest = + let (discard, keep) = partition ((`S.member` arrs) . snd) $ loopNestingParamsAndArrs nest + in (S.fromList (map (paramName . fst) discard) <> arrs, + nest { loopNestingParamsAndArrs = keep }) + +newKernel :: LoopNesting -> KernelNest +newKernel nest = (nest, []) + +kernelNestLoops :: KernelNest -> [LoopNesting] +kernelNestLoops (loop, loops) = loop : loops + +boundInKernelNest :: KernelNest -> Names +boundInKernelNest = mconcat . boundInKernelNests + +boundInKernelNests :: KernelNest -> [Names] +boundInKernelNests = map (S.fromList . + map (paramName . fst) . + loopNestingParamsAndArrs) . + kernelNestLoops + +kernelNestWidths :: KernelNest -> [SubExp] +kernelNestWidths = map loopNestingWidth . kernelNestLoops + +constructKernel :: (MonadFreshNames m, LocalScope Kernels m) => + KernelNest -> KernelBody InKernel + -> m (Stms Kernels, SubExp, Stm Kernels) +constructKernel kernel_nest inner_body = do + (w_bnds, w, ispace, inps, rts) <- flatKernel kernel_nest + let used_inps = filter inputIsUsed inps + cs = loopNestingCertificates first_nest + + (ksize_bnds, k) <- inScopeOf w_bnds $ + mapKernel w (FlatThreadSpace ispace) used_inps rts inner_body + + let kbnds = w_bnds <> ksize_bnds + return (kbnds, + w, + Let (loopNestingPattern first_nest) (StmAux cs ()) $ Op k) + where + first_nest = fst kernel_nest + inputIsUsed input = kernelInputName input `S.member` + freeIn inner_body + +-- | Flatten a kernel nesting to: +-- +-- (0) Ancillary prologue bindings. +-- +-- (1) The total number of threads, equal to the product of all +-- nesting widths, and equal to the product of the index space. +-- +-- (2) The index space. +-- +-- (3) The kernel inputs - not that some of these may be unused. +-- +-- (4) The per-thread return type. +flatKernel :: MonadFreshNames m => + KernelNest + -> m (Stms Kernels, + SubExp, + [(VName, SubExp)], + [KernelInput], + [Type]) +flatKernel (MapNesting pat _ nesting_w params_and_arrs, []) = do + i <- newVName "gtid" + let inps = [ KernelInput pname ptype arr [Var i] | + (Param pname ptype, arr) <- params_and_arrs ] + return (mempty, nesting_w, [(i,nesting_w)], inps, + map rowType $ patternTypes pat) + +flatKernel (MapNesting _ _ nesting_w params_and_arrs, nest : nests) = do + i <- newVName "gtid" + (w_bnds, w, ispace, inps, returns) <- flatKernel (nest, nests) + + w' <- newVName "nesting_size" + let w_bnd = mkLet [] [Ident w' $ Prim int32] $ + BasicOp $ BinOp (Mul Int32) w nesting_w + + let inps' = map fixupInput inps + isParam inp = + snd <$> find ((==kernelInputArray inp) . paramName . fst) params_and_arrs + fixupInput inp + | Just arr <- isParam inp = + inp { kernelInputArray = arr + , kernelInputIndices = Var i : kernelInputIndices inp } + | otherwise = + inp + + return (w_bnds <> oneStm w_bnd, Var w', (i, nesting_w) : ispace, + extra_inps i <> inps', returns) + where extra_inps i = + [ KernelInput pname ptype arr [Var i] | + (Param pname ptype, arr) <- params_and_arrs ] + +-- | Description of distribution to do. +data DistributionBody = DistributionBody { + distributionTarget :: Targets + , distributionFreeInBody :: Names + , distributionIdentityMap :: M.Map VName Ident + , distributionExpandTarget :: Target -> Target + -- ^ Also related to avoiding identity mapping. + } + +distributionInnerPattern :: DistributionBody -> Pattern Kernels +distributionInnerPattern = fst . innerTarget . distributionTarget + +distributionBodyFromStms :: Attributes lore => + Targets -> Stms lore -> (DistributionBody, Result) +distributionBodyFromStms (Targets (inner_pat, inner_res) targets) stms = + let bound_by_stms = S.fromList $ M.keys $ scopeOf stms + (inner_pat', inner_res', inner_identity_map, inner_expand_target) = + removeIdentityMappingGeneral bound_by_stms inner_pat inner_res + in (DistributionBody + { distributionTarget = Targets (inner_pat', inner_res') targets + , distributionFreeInBody = fold (fmap freeInStm stms) `S.difference` bound_by_stms + , distributionIdentityMap = inner_identity_map + , distributionExpandTarget = inner_expand_target + }, + inner_res') + +distributionBodyFromStm :: Attributes lore => + Targets -> Stm lore -> (DistributionBody, Result) +distributionBodyFromStm targets bnd = + distributionBodyFromStms targets $ oneStm bnd + +createKernelNest :: (MonadFreshNames m, HasScope t m) => + Nestings + -> DistributionBody + -> m (Maybe (Targets, KernelNest)) +createKernelNest (inner_nest, nests) distrib_body = do + let Targets target targets = distributionTarget distrib_body + unless (length nests == length targets) $ + fail $ "Nests and targets do not match!\n" ++ + "nests: " ++ ppNestings (inner_nest, nests) ++ + "\ntargets:" ++ ppTargets (Targets target targets) + runMaybeT $ fmap prepare $ recurse $ zip nests targets + + where prepare (x, _, z) = (z, x) + bound_in_nest = + mconcat $ map boundInNesting $ inner_nest : nests + -- | Can something of this type be taken outside the nest? + -- I.e. are none of its dimensions bound inside the nest. + distributableType = + S.null . S.intersection bound_in_nest . freeIn . arrayDims + + distributeAtNesting :: (HasScope t m, MonadFreshNames m) => + Nesting + -> Pattern Kernels + -> (LoopNesting -> KernelNest, Names) + -> M.Map VName Ident + -> [Ident] + -> (Target -> Targets) + -> MaybeT m (KernelNest, Names, Targets) + distributeAtNesting + (Nesting nest_let_bound nest) + pat + (add_to_kernel, free_in_kernel) + identity_map + inner_returned_arrs + addTarget = do + let nest'@(MapNesting _ cs w params_and_arrs) = + removeUnusedNestingParts free_in_kernel nest + (params,arrs) = unzip params_and_arrs + param_names = S.fromList $ map paramName params + free_in_kernel' = + (freeIn nest' <> free_in_kernel) `S.difference` param_names + required_from_nest = + free_in_kernel' `S.intersection` nest_let_bound + + required_from_nest_idents <- + forM (S.toList required_from_nest) $ \name -> do + t <- lift $ lookupType name + return $ Ident name t + + (free_params, free_arrs, bind_in_target) <- + fmap unzip3 $ + forM (inner_returned_arrs++required_from_nest_idents) $ + \(Ident pname ptype) -> + case M.lookup pname identity_map of + Nothing -> do + arr <- newIdent (baseString pname ++ "_r") $ + arrayOfRow ptype w + return (Param pname ptype, + arr, + True) + Just arr -> + return (Param pname ptype, + arr, + False) + + let free_arrs_pat = + basicPattern [] $ map snd $ + filter fst $ zip bind_in_target free_arrs + free_params_pat = + map snd $ filter fst $ zip bind_in_target free_params + + (actual_params, actual_arrs) = + (params++free_params, + arrs++map identName free_arrs) + actual_param_names = + S.fromList $ map paramName actual_params + + nest'' = + removeUnusedNestingParts free_in_kernel $ + MapNesting pat cs w $ zip actual_params actual_arrs + + free_in_kernel'' = + (freeIn nest'' <> free_in_kernel) `S.difference` actual_param_names + + unless (all (distributableType . paramType) $ + loopNestingParams nest'') $ + fail "Would induce irregular array" + return (add_to_kernel nest'', + + free_in_kernel'', + + addTarget (free_arrs_pat, map (Var . paramName) free_params_pat)) + + recurse :: (HasScope t m, MonadFreshNames m) => + [(Nesting,Target)] + -> MaybeT m (KernelNest, Names, Targets) + recurse [] = + distributeAtNesting + inner_nest + (distributionInnerPattern distrib_body) + (newKernel, + distributionFreeInBody distrib_body `S.intersection` bound_in_nest) + (distributionIdentityMap distrib_body) + [] $ + singleTarget . distributionExpandTarget distrib_body + + recurse ((nest, (pat,res)) : nests') = do + (kernel@(outer, _), kernel_free, kernel_targets) <- recurse nests' + + let (pat', res', identity_map, expand_target) = + removeIdentityMappingFromNesting + (S.fromList $ patternNames $ loopNestingPattern outer) pat res + + distributeAtNesting + nest + pat' + (\k -> pushKernelNesting (pat',res') k kernel, + kernel_free) + identity_map + (patternIdents $ fst $ outerTarget kernel_targets) + ((`pushOuterTarget` kernel_targets) . expand_target) + +removeUnusedNestingParts :: Names -> LoopNesting -> LoopNesting +removeUnusedNestingParts used (MapNesting pat cs w params_and_arrs) = + MapNesting pat cs w $ zip used_params used_arrs + where (params,arrs) = unzip params_and_arrs + (used_params, used_arrs) = + unzip $ + filter ((`S.member` used) . paramName . fst) $ + zip params arrs + +removeIdentityMappingGeneral :: Names -> Pattern Kernels -> Result + -> (Pattern Kernels, + Result, + M.Map VName Ident, + Target -> Target) +removeIdentityMappingGeneral bound pat res = + let (identities, not_identities) = + mapEither isIdentity $ zip (patternElements pat) res + (not_identity_patElems, not_identity_res) = unzip not_identities + (identity_patElems, identity_res) = unzip identities + expandTarget (tpat, tres) = + (Pattern [] $ patternElements tpat ++ identity_patElems, + tres ++ map Var identity_res) + identity_map = M.fromList $ zip identity_res $ + map patElemIdent identity_patElems + in (Pattern [] not_identity_patElems, + not_identity_res, + identity_map, + expandTarget) + where isIdentity (patElem, Var v) + | not (v `S.member` bound) = Left (patElem, v) + isIdentity x = Right x + +removeIdentityMappingFromNesting :: Names -> Pattern Kernels -> Result + -> (Pattern Kernels, + Result, + M.Map VName Ident, + Target -> Target) +removeIdentityMappingFromNesting bound_in_nesting pat res = + let (pat', res', identity_map, expand_target) = + removeIdentityMappingGeneral bound_in_nesting pat res + in (pat', res', identity_map, expand_target) + +tryDistribute :: (MonadFreshNames m, LocalScope Kernels m, MonadLogger m) => + Nestings -> Targets -> Stms InKernel + -> m (Maybe (Targets, Stms Kernels)) +tryDistribute _ targets stms | null stms = + -- No point in distributing an empty kernel. + return $ Just (targets, mempty) +tryDistribute nest targets stms = + createKernelNest nest dist_body >>= + \case + Just (targets', distributed) -> do + (w_bnds, _, kernel_bnd) <- localScope (targetsScope targets') $ + constructKernel distributed inner_body + distributed' <- renameStm kernel_bnd + logMsg $ "distributing\n" ++ + unlines (map pretty $ stmsToList stms) ++ + pretty (snd $ innerTarget targets) ++ + "\nas\n" ++ pretty distributed' ++ + "\ndue to targets\n" ++ ppTargets targets ++ + "\nand with new targets\n" ++ ppTargets targets' + return $ Just (targets', w_bnds <> oneStm distributed') + Nothing -> + return Nothing + where (dist_body, inner_body_res) = distributionBodyFromStms targets stms + inner_body = KernelBody () stms $ + map (ThreadsReturn ThreadsInSpace) inner_body_res + +tryDistributeStm :: (MonadFreshNames m, HasScope t m, Attributes lore) => + Nestings -> Targets -> Stm lore + -> m (Maybe (Result, Targets, KernelNest)) +tryDistributeStm nest targets bnd = + fmap addRes <$> createKernelNest nest dist_body + where (dist_body, res) = distributionBodyFromStm targets bnd + addRes (targets', kernel_nest) = (res, targets', kernel_nest) diff --git a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs @@ -0,0 +1,170 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +module Futhark.Pass.ExtractKernels.ISRWIM + ( iswim + , irwim + , rwimPossible + ) + where + +import Control.Arrow (first) +import Control.Monad.State +import Data.Semigroup ((<>)) + +import Futhark.MonadFreshNames +import Futhark.Representation.SOACS +import Futhark.Tools + +-- | Interchange Scan With Inner Map. Tries to turn a @scan(map)@ into a +-- @map(scan) +iswim :: (MonadBinder m, Lore m ~ SOACS) => + Pattern + -> SubExp + -> Lambda + -> [(SubExp, VName)] + -> Maybe (m ()) +iswim res_pat w scan_fun scan_input + | Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible scan_fun = Just $ do + let (accs, arrs) = unzip scan_input + arrs' <- transposedArrays arrs + accs' <- mapM (letExp "acc" . BasicOp . SubExp) accs + + let map_arrs' = accs' ++ arrs' + (scan_acc_params, scan_elem_params) = + splitAt (length arrs) $ lambdaParams scan_fun + map_params = map removeParamOuterDim scan_acc_params ++ + map (setParamOuterDimTo w) scan_elem_params + map_rettype = map (setOuterDimTo w) $ lambdaReturnType scan_fun + + scan_params = lambdaParams map_fun + scan_body = lambdaBody map_fun + scan_rettype = lambdaReturnType map_fun + scan_fun' = Lambda scan_params scan_body scan_rettype + scan_input' = map (first Var) $ + uncurry zip $ splitAt (length arrs') $ map paramName map_params + (nes', scan_arrs) = unzip scan_input' + + scan_soac <- scanSOAC scan_fun' nes' + let map_body = mkBody (oneStm $ Let (setPatternOuterDimTo w map_pat) (defAux ()) $ + Op $ Screma w scan_soac scan_arrs) $ + map Var $ patternNames map_pat + map_fun' = Lambda map_params map_body map_rettype + + res_pat' <- fmap (basicPattern []) $ + mapM (newIdent' (<>"_transposed") . transposeIdentType) $ + patternValueIdents res_pat + + addStm $ Let res_pat' (StmAux map_cs ()) $ Op $ Screma map_w + (ScremaForm (nilFn, mempty) (mempty, nilFn, mempty) map_fun') map_arrs' + + forM_ (zip (patternValueIdents res_pat) + (patternValueIdents res_pat')) $ \(to, from) -> do + let perm = [1,0] ++ [2..arrayRank (identType from)-1] + addStm $ Let (basicPattern [] [to]) (defAux ()) $ + BasicOp $ Rearrange perm $ identName from + | otherwise = Nothing + +-- | Interchange Reduce With Inner Map. Tries to turn a @reduce(map)@ into a +-- @map(reduce) +irwim :: (MonadBinder m, Lore m ~ SOACS) => + Pattern + -> SubExp + -> Commutativity -> Lambda + -> [(SubExp, VName)] + -> Maybe (m ()) +irwim res_pat w comm red_fun red_input + | Just (map_pat, map_cs, map_w, map_fun) <- rwimPossible red_fun = Just $ do + let (accs, arrs) = unzip red_input + arrs' <- transposedArrays arrs + -- FIXME? Can we reasonably assume that the accumulator is a + -- replicate? We also assume that it is non-empty. + let indexAcc (Var v) = do + v_t <- lookupType v + letSubExp "acc" $ BasicOp $ Index v $ + fullSlice v_t [DimFix $ intConst Int32 0] + indexAcc Constant{} = + fail "irwim: array accumulator is a constant." + accs' <- mapM indexAcc accs + + let (_red_acc_params, red_elem_params) = + splitAt (length arrs) $ lambdaParams red_fun + map_rettype = map rowType $ lambdaReturnType red_fun + map_params = map (setParamOuterDimTo w) red_elem_params + + red_params = lambdaParams map_fun + red_body = lambdaBody map_fun + red_rettype = lambdaReturnType map_fun + red_fun' = Lambda red_params red_body red_rettype + red_input' = zip accs' $ map paramName map_params + red_pat = stripPatternOuterDim map_pat + + map_body <- + case irwim red_pat w comm red_fun' red_input' of + Nothing -> do + reduce_soac <- reduceSOAC comm red_fun' $ map fst red_input' + return $ mkBody (oneStm $ Let red_pat (defAux ()) $ + Op $ Screma w reduce_soac $ map snd red_input') $ + map Var $ patternNames map_pat + Just m -> localScope (scopeOfLParams map_params) $ do + map_body_bnds <- collectStms_ m + return $ mkBody map_body_bnds $ map Var $ patternNames map_pat + + let map_fun' = Lambda map_params map_body map_rettype + + addStm $ Let res_pat (StmAux map_cs ()) $ Op $ Screma map_w (mapSOAC map_fun') arrs' + | otherwise = Nothing + +rwimPossible :: Lambda + -> Maybe (Pattern, Certificates, SubExp, Lambda) +rwimPossible fun + | Body _ stms res <- lambdaBody fun, + [bnd] <- stmsToList stms, -- Body has a single binding + map_pat <- stmPattern bnd, + map Var (patternNames map_pat) == res, -- Returned verbatim + Op (Screma map_w form map_arrs) <- stmExp bnd, + Just map_fun <- isMapSOAC form, + map paramName (lambdaParams fun) == map_arrs = + Just (map_pat, stmCerts bnd, map_w, map_fun) + | otherwise = + Nothing + +transposedArrays :: MonadBinder m => [VName] -> m [VName] +transposedArrays arrs = forM arrs $ \arr -> do + t <- lookupType arr + let perm = [1,0] ++ [2..arrayRank t-1] + letExp (baseString arr) $ BasicOp $ Rearrange perm arr + +removeParamOuterDim :: LParam -> LParam +removeParamOuterDim param = + let t = rowType $ paramType param + in param { paramAttr = t } + +setParamOuterDimTo :: SubExp -> LParam -> LParam +setParamOuterDimTo w param = + let t = setOuterDimTo w $ paramType param + in param { paramAttr = t } + +setIdentOuterDimTo :: SubExp -> Ident -> Ident +setIdentOuterDimTo w ident = + let t = setOuterDimTo w $ identType ident + in ident { identType = t } + +setOuterDimTo :: SubExp -> Type -> Type +setOuterDimTo w t = + arrayOfRow (rowType t) w + +setPatternOuterDimTo :: SubExp -> Pattern -> Pattern +setPatternOuterDimTo w pat = + basicPattern [] $ map (setIdentOuterDimTo w) $ patternValueIdents pat + +transposeIdentType :: Ident -> Ident +transposeIdentType ident = + ident { identType = transposeType $ identType ident } + +stripIdentOuterDim :: Ident -> Ident +stripIdentOuterDim ident = + ident { identType = rowType $ identType ident } + +stripPatternOuterDim :: Pattern -> Pattern +stripPatternOuterDim pat = + basicPattern [] $ map stripIdentOuterDim $ patternValueIdents pat diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs @@ -0,0 +1,177 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | It is well known that fully parallel loops can always be +-- interchanged inwards with a sequential loop. This module +-- implements that transformation. +-- +-- This is also where we implement loop-switching (for branches), +-- which is semantically similar to interchange. +module Futhark.Pass.ExtractKernels.Interchange + ( + SeqLoop (..) + , interchangeLoops + , Branch (..) + , interchangeBranch + ) where + +import Control.Monad.RWS.Strict +import qualified Data.Set as S +import Data.Maybe +import Data.List + +import Futhark.Pass.ExtractKernels.Distribution + (LoopNesting(..), KernelNest, kernelNestLoops) +import Futhark.Representation.SOACS +import Futhark.MonadFreshNames +import Futhark.Transform.Rename +import Futhark.Tools + +-- | An encoding of a sequential do-loop with no existential context, +-- alongside its result pattern. +data SeqLoop = SeqLoop [Int] Pattern [(FParam, SubExp)] (LoopForm SOACS) Body + +seqLoopStm :: SeqLoop -> Stm +seqLoopStm (SeqLoop _ pat merge form body) = + Let pat (defAux ()) $ DoLoop [] merge form body + +interchangeLoop :: (MonadBinder m, LocalScope SOACS m) => + SeqLoop -> LoopNesting + -> m SeqLoop +interchangeLoop + (SeqLoop perm loop_pat merge form body) + (MapNesting pat cs w params_and_arrs) = do + merge_expanded <- + localScope (scopeOfLParams $ map fst params_and_arrs) $ + mapM expand merge + + let loop_pat_expanded = + Pattern [] $ map expandPatElem $ patternElements loop_pat + new_params = [ Param pname $ fromDecl ptype + | (Param pname ptype, _) <- merge ] + new_arrs = map (paramName . fst) merge_expanded + rettype = map rowType $ patternTypes loop_pat_expanded + + -- If the map consumes something that is bound outside the loop + -- (i.e. is not a merge parameter), we have to copy() it. As a + -- small simplification, we just remove the parameter outright if + -- it is not used anymore. This might happen if the parameter was + -- used just as the inital value of a merge parameter. + ((params', arrs'), pre_copy_bnds) <- + runBinder $ localScope (scopeOfLParams new_params) $ + unzip . catMaybes <$> mapM copyOrRemoveParam params_and_arrs + + body' <- mkDummyStms (params'<>new_params) body + + let lam = Lambda (params'<>new_params) body' rettype + map_bnd = Let loop_pat_expanded (StmAux cs ()) $ + Op $ Screma w (mapSOAC lam) $ arrs' <> new_arrs + res = map Var $ patternNames loop_pat_expanded + pat' = Pattern [] $ rearrangeShape perm $ patternValueElements pat + + return $ + SeqLoop [0..patternSize pat-1] pat' merge_expanded form $ + mkBody (pre_copy_bnds<>oneStm map_bnd) res + where free_in_body = freeInBody body + + copyOrRemoveParam (param, arr) + | not (paramName param `S.member` free_in_body) = + return Nothing + | otherwise = + return $ Just (param, arr) + + expandedInit _ (Var v) + | Just (_, arr) <- + find ((==v).paramName.fst) params_and_arrs = + return $ Var arr + expandedInit param_name se = + letSubExp (param_name <> "_expanded_init") $ + BasicOp $ Replicate (Shape [w]) se + + expand (merge_param, merge_init) = do + expanded_param <- + newParam (param_name <> "_expanded") $ + arrayOf (paramDeclType merge_param) (Shape [w]) $ + uniqueness $ declTypeOf merge_param + expanded_init <- expandedInit param_name merge_init + return (expanded_param, expanded_init) + where param_name = baseString $ paramName merge_param + + expandPatElem (PatElem name t) = + PatElem name $ arrayOfRow t w + + -- | The kernel extractor cannot handle identity mappings, so + -- insert dummy statements for body results that are just a + -- lambda parameter. + mkDummyStms params (Body () stms res) = do + (res', extra_stms) <- unzip <$> mapM dummyStm res + return $ Body () (stms<>mconcat extra_stms) res' + where dummyStm (Var v) + | Just p <- find ((==v) . paramName) params = do + dummy <- newVName (baseString v ++ "_dummy") + return (Var dummy, + oneStm $ + Let (Pattern [] [PatElem dummy $ paramType p]) + (defAux ()) $ + BasicOp $ SubExp $ Var $ paramName p) + dummyStm se = return (se, mempty) + +-- | Given a (parallel) map nesting and an inner sequential loop, move +-- the maps inside the sequential loop. The result is several +-- statements - one of these will be the loop, which will then contain +-- statements with 'Map' expressions. +interchangeLoops :: (MonadFreshNames m, HasScope SOACS m) => + KernelNest -> SeqLoop + -> m (Stms SOACS) +interchangeLoops nest loop = do + (loop', bnds) <- + runBinder $ foldM interchangeLoop loop $ reverse $ kernelNestLoops nest + return $ bnds <> oneStm (seqLoopStm loop') + +data Branch = Branch [Int] Pattern SubExp Body Body (IfAttr (BranchType SOACS)) + +branchStm :: Branch -> Stm +branchStm (Branch _ pat cond tbranch fbranch ret) = + Let pat (defAux ()) $ If cond tbranch fbranch ret + +interchangeBranch1 :: (MonadBinder m, LocalScope SOACS m) => + Branch -> LoopNesting -> m Branch +interchangeBranch1 + (Branch perm branch_pat cond tbranch fbranch (IfAttr ret if_sort)) + (MapNesting pat cs w params_and_arrs) = do + let ret' = map (`arrayOfRow` Free w) ret + pat' = Pattern [] $ rearrangeShape perm $ patternValueElements pat + + (params, arrs) = unzip params_and_arrs + lam_ret = map rowType $ patternTypes pat + + branch_pat' = + Pattern [] $ map (fmap (`arrayOfRow` w)) $ patternElements branch_pat + + mkBranch branch = (renameBody=<<) $ do + branch' <- if null $ bodyStms branch + then runBodyBinder $ + -- XXX: We need a temporary dummy binding to + -- prevent an empty map body. The kernel + -- extractor does not like empty map bodies. + resultBody <$> mapM dummyBind (bodyResult branch) + else return branch + let lam = Lambda params branch' lam_ret + res = map Var $ patternNames branch_pat' + map_bnd = Let branch_pat' (StmAux cs ()) $ Op $ Screma w (mapSOAC lam) arrs + return $ mkBody (oneStm map_bnd) res + + tbranch' <- mkBranch tbranch + fbranch' <- mkBranch fbranch + return $ Branch [0..patternSize pat-1] pat' cond tbranch' fbranch' $ + IfAttr ret' if_sort + where dummyBind se = do + dummy <- newVName "dummy" + letBindNames_ [dummy] (BasicOp $ SubExp se) + return $ Var dummy + +interchangeBranch :: (MonadFreshNames m, HasScope SOACS m) => + KernelNest -> Branch -> m (Stms SOACS) +interchangeBranch nest loop = do + (loop', bnds) <- + runBinder $ foldM interchangeBranch1 loop $ reverse $ kernelNestLoops nest + return $ bnds <> oneStm (branchStm loop') diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs @@ -0,0 +1,324 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | Extract limited nested parallelism for execution inside +-- individual kernel workgroups. +module Futhark.Pass.ExtractKernels.Intragroup + (intraGroupParallelise) +where + +import Control.Monad.RWS +import Control.Monad.Trans.Maybe +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Futhark.Analysis.PrimExp.Convert +import Futhark.Representation.SOACS +import qualified Futhark.Representation.Kernels as Out +import Futhark.Representation.Kernels.Kernel +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Analysis.DataDependencies +import qualified Futhark.Pass.ExtractKernels.Kernelise as Kernelise +import Futhark.Pass.ExtractKernels.Distribution +import Futhark.Pass.ExtractKernels.BlockedKernel + +-- | Convert the statements inside a map nest to kernel statements, +-- attempting to parallelise any remaining (top-level) parallel +-- statements. Anything that is not a map, scan or reduction will +-- simply be sequentialised. This includes sequential loops that +-- contain maps, scans or reduction. In the future, we could probably +-- do something more clever. Make sure that the amount of parallelism +-- to be exploited does not exceed the group size. Further, as a hack +-- we also consider the size of all intermediate arrays as +-- "parallelism to be exploited" to avoid exploding local memory. +-- +-- We distinguish between "minimum group size" and "maximum +-- exploitable parallelism". +intraGroupParallelise :: (MonadFreshNames m, LocalScope Out.Kernels m) => + KernelNest -> Lambda + -> m (Maybe ((SubExp, SubExp), SubExp, + Out.Stms Out.Kernels, Out.Stms Out.Kernels)) +intraGroupParallelise knest lam = runMaybeT $ do + (w_stms, w, ispace, inps, rts) <- lift $ flatKernel knest + let num_groups = w + body = lambdaBody lam + + ltid <- newVName "ltid" + let group_variant = S.fromList [ltid] + (wss_min, wss_avail, kbody) <- + lift $ localScope (scopeOfLParams $ lambdaParams lam) $ + intraGroupParalleliseBody (dataDependencies body) group_variant ltid body + + known_outside <- lift $ M.keys <$> askScope + unless (all (`elem` known_outside) $ freeIn $ wss_min ++ wss_avail) $ + fail "Irregular parallelism" + + ((intra_avail_par, kspace, read_input_stms), prelude_stms) <- lift $ runBinder $ do + let foldBinOp' _ [] = eSubExp $ intConst Int32 0 + foldBinOp' bop (x:xs) = foldBinOp bop x xs + ws_min <- mapM (letSubExp "one_intra_par_min" <=< foldBinOp' (Mul Int32)) $ + filter (not . null) wss_min + ws_avail <- mapM (letSubExp "one_intra_par_avail" <=< foldBinOp' (Mul Int32)) $ + filter (not . null) wss_avail + + -- The amount of parallelism available *in the worst case* is + -- equal to the smallest parallel loop. + intra_avail_par <- letSubExp "intra_avail_par" =<< foldBinOp' (SMin Int32) ws_avail + + -- The group size is either the maximum of the minimum parallelism + -- exploited, or the desired parallelism (bounded by the max group + -- size) in case there is no minimum. + group_size <- letSubExp "computed_group_size" =<< + if null ws_min + then eBinOp (SMin Int32) + (eSubExp =<< letSubExp "max_group_size" (Op $ Out.GetSizeMax Out.SizeGroup)) + (eSubExp intra_avail_par) + else foldBinOp' (SMax Int32) ws_min + + let inputIsUsed input = kernelInputName input `S.member` freeInBody body + used_inps = filter inputIsUsed inps + + addStms w_stms + + num_threads <- letSubExp "num_threads" $ + BasicOp $ BinOp (Mul Int32) num_groups group_size + + let ksize = (num_groups, group_size, num_threads) + + kspace <- newKernelSpace ksize $ FlatThreadSpace $ ispace ++ [(ltid,group_size)] + + read_input_stms <- mapM readKernelInput used_inps + + return (intra_avail_par, kspace, read_input_stms) + + let kbody' = kbody { kernelBodyStms = stmsFromList read_input_stms <> kernelBodyStms kbody } + + -- The kernel itself is producing a "flat" result of shape + -- [num_groups]. We must explicitly reshape it to match the shapes + -- of our enclosing map-nests. + let nested_pat = loopNestingPattern first_nest + flatPatElem pat_elem = do + let t' = arrayOfRow (length ispace `stripArray` patElemType pat_elem) num_groups + name <- newVName $ baseString (patElemName pat_elem) ++ "_flat" + return $ PatElem name t' + flat_pat <- lift $ Pattern [] <$> mapM flatPatElem (patternValueElements nested_pat) + + let kstm = Let flat_pat (StmAux cs ()) $ Op $ + Kernel (KernelDebugHints "map_intra_group" []) kspace rts kbody' + reshapeStm nested_pe flat_pe = + Let (Pattern [] [nested_pe]) (StmAux cs ()) $ + BasicOp $ Reshape (map DimNew $ arrayDims $ patElemType nested_pe) $ + patElemName flat_pe + reshape_stms = zipWith reshapeStm (patternElements nested_pat) + (patternElements flat_pat) + + let intra_min_par = intra_avail_par + return ((intra_min_par, intra_avail_par), spaceGroupSize kspace, + prelude_stms, oneStm kstm <> stmsFromList reshape_stms) + where first_nest = fst knest + cs = loopNestingCertificates first_nest + +data Env = Env { _localTID :: VName + , _dataDeps :: Dependencies + , _groupVariant :: Names + } + +type IntraGroupM = BinderT Out.InKernel (RWS Env (S.Set [SubExp], S.Set [SubExp]) VNameSource) + +runIntraGroupM :: (MonadFreshNames m, HasScope Out.Kernels m) => + Env -> IntraGroupM () -> m ([[SubExp]], [[SubExp]], Out.Stms Out.InKernel) +runIntraGroupM env m = do + scope <- castScope <$> askScope + modifyNameSource $ \src -> + let (((), kstms), src', (ws_min, ws_avail)) = runRWS (runBinderT m scope) env src + in ((S.toList ws_min, S.toList ws_avail, kstms), src') + +parallelMin :: [SubExp] -> IntraGroupM () +parallelMin ws = tell (S.singleton ws, S.singleton ws) + +parallelAvail :: [SubExp] -> IntraGroupM () +parallelAvail ws = tell (mempty, S.singleton ws) + +intraGroupBody :: Body -> IntraGroupM (Out.Body Out.InKernel) +intraGroupBody body = do + stms <- collectStms_ $ mapM_ intraGroupStm $ bodyStms body + return $ mkBody stms $ bodyResult body + +intraGroupStm :: Stm -> IntraGroupM () +intraGroupStm stm@(Let pat _ e) = do + Env ltid deps group_variant <- ask + let groupInvariant (Var v) = + S.null . S.intersection group_variant . + flip (M.findWithDefault mempty) deps $ v + groupInvariant Constant{} = True + + case e of + DoLoop ctx val (ForLoop i it bound inps) loopbody + | groupInvariant bound -> + localScope (scopeOf form) $ + localScope (scopeOfFParams $ map fst $ ctx ++ val) $ do + loopbody' <- intraGroupBody loopbody + letBind_ pat $ DoLoop ctx val form loopbody' + where form = ForLoop i it bound inps + + If cond tbody fbody ifattr + | groupInvariant cond -> do + tbody' <- intraGroupBody tbody + fbody' <- intraGroupBody fbody + letBind_ pat $ If cond tbody' fbody' ifattr + + Op (Screma w form arrs) | Just fun <- isMapSOAC form -> do + body_stms <- collectStms_ $ do + forM_ (zip (lambdaParams fun) arrs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames [paramName p] $ BasicOp $ Index arr $ + fullSlice arr_t [DimFix $ Var ltid] + Kernelise.transformStms $ bodyStms $ lambdaBody fun + let comb_body = mkBody body_stms $ bodyResult $ lambdaBody fun + ctid <- newVName "ctid" + letBind_ pat $ Op $ + Out.Combine (Out.combineSpace [(ctid, w)]) (lambdaReturnType fun) [] comb_body + mapM_ (parallelMin . arrayDims) $ patternTypes pat + parallelMin [w] + + Op (Screma w form arrs) + | Just (scanfun, nes, foldfun) <- isScanomapSOAC form -> do + let (scan_pes, map_pes) = + splitAt (length nes) $ patternElements pat + scan_input <- procInput ltid (Pattern [] map_pes) w foldfun nes arrs + + scanfun' <- Kernelise.transformLambda scanfun + + -- A GroupScan lambda needs two more parameters. + my_index <- newVName "my_index" + other_index <- newVName "other_index" + let my_index_param = Param my_index (Prim int32) + other_index_param = Param other_index (Prim int32) + scanfun'' = scanfun' { lambdaParams = my_index_param : + other_index_param : + lambdaParams scanfun' + } + letBind_ (Pattern [] scan_pes) $ + Op $ Out.GroupScan w scanfun'' $ zip nes scan_input + parallelMin [w] + + Op (Screma w form arrs) + | Just (_, redfun, nes, foldfun) <- isRedomapSOAC form -> do + let (red_pes, map_pes) = + splitAt (length nes) $ patternElements pat + red_input <- procInput ltid (Pattern [] map_pes) w foldfun nes arrs + + redfun' <- Kernelise.transformLambda redfun + + -- A GroupReduce lambda needs two more parameters. + my_index <- newVName "my_index" + other_index <- newVName "other_index" + let my_index_param = Param my_index (Prim int32) + other_index_param = Param other_index (Prim int32) + redfun'' = redfun' { lambdaParams = my_index_param : + other_index_param : + lambdaParams redfun' + } + letBind_ (Pattern [] red_pes) $ + Op $ Out.GroupReduce w redfun'' $ zip nes red_input + parallelMin [w] + + Op (Stream w (Sequential accs) lam arrs) + | chunk_size_param : _ <- lambdaParams lam -> do + types <- asksScope castScope + ((), stream_bnds) <- + runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types + let replace (Var v) | v == paramName chunk_size_param = w + replace se = se + replaceSets (x, y) = (S.map (map replace) x, S.map (map replace) y) + censor replaceSets $ mapM_ intraGroupStm stream_bnds + + Op (Scatter w lam ivs dests) -> do + parallelMin [w] + ctid <- newVName "ctid" + let cspace = Out.CombineSpace dests [(ctid, w)] + body_stms <- collectStms_ $ do + forM_ (zip (lambdaParams lam) ivs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames [paramName p] $ BasicOp $ Index arr $ + fullSlice arr_t [DimFix $ Var ltid] -- ltid on purpose to enable hoisting. + Kernelise.transformStms $ bodyStms $ lambdaBody lam + let body = mkBody body_stms $ bodyResult $ lambdaBody lam + letBind_ pat $ Op $ Out.Combine cspace (lambdaReturnType lam) mempty body + + BasicOp (Update dest slice (Var v)) -> do + let ws = sliceDims slice + activeForDim w i = BasicOp $ CmpOp (CmpSlt Int32) i w + parallelMin ws + dest' <- letExp "update_inp" $ Op $ Out.Barrier [Var dest] + let new_inds = unflattenIndex (map (primExpFromSubExp int32) ws) + (primExpFromSubExp int32 $ Var ltid) + new_inds' <- mapM (letSubExp "i" <=< toExp) new_inds + active <- letSubExp "active" =<< + foldBinOp LogAnd (constant True) =<< + mapM (letSubExp "active") (zipWith activeForDim ws new_inds') + (active_res, active_stms) <- collectStms $ do + slice' <- + mapM (letSubExp "j" <=< toExp) $ + fixSlice (map (fmap $ primExpFromSubExp int32) slice) new_inds + letInPlace "update_res" dest' (map DimFix slice') $ + BasicOp $ Index v $ map DimFix new_inds' + sync <- letSubExp "update_res" =<< eIf (eSubExp active) + (pure $ mkBody active_stms [Var active_res]) + (pure $ mkBody mempty [Var dest']) + letBind_ pat $ Op $ Out.Barrier [sync] + + BasicOp (Copy arr) -> do + arr_t <- lookupType arr + let w = arraySize 0 arr_t + ctid <- newVName "copy_ctid" + letBind_ pat . Op . Out.Combine (Out.combineSpace [(ctid, w)]) [rowType arr_t] [] <=< + localScope (M.singleton ctid $ IndexInfo Int32) $ + insertStmsM $ resultBodyM . pure <=< letSubExp "v" $ + BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var ctid] + + BasicOp (Replicate (Shape outer_ws) se) + | [inner_ws] <- map (drop (length outer_ws) . arrayDims) $ patternTypes pat -> do + let ws = outer_ws ++ inner_ws + new_inds' <- replicateM (length ws) $ newVName "new_local_index" + let inner_inds' = drop (length outer_ws) new_inds' + space = Out.combineSpace $ zip new_inds' ws + index = case se of Var v -> BasicOp $ Index v $ + map (DimFix . Var) inner_inds' + Constant{} -> BasicOp $ SubExp se + body <- runBodyBinder $ eBody [pure index] + letBind_ pat $ Op $ + Out.Combine space (map (Prim . elemType) $ patternTypes pat) [] body + mapM_ (parallelAvail . arrayDims) $ patternTypes pat + + _ -> + Kernelise.transformStm stm + + where procInput :: VName + -> Out.Pattern Out.InKernel + -> SubExp -> Lambda -> [SubExp] -> [VName] + -> IntraGroupM [VName] + procInput ltid map_pat w map_fun nes arrs = do + fold_stms <- collectStms_ $ do + forM_ (zip (lambdaParams map_fun) arrs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames_ [paramName p] $ BasicOp $ Index arr $ + fullSlice arr_t [DimFix $ Var ltid] + + Kernelise.transformStms $ bodyStms $ lambdaBody map_fun + let fold_body = mkBody fold_stms $ bodyResult $ lambdaBody map_fun + + op_inps <- replicateM (length nes) (newVName "op_input") + ctid <- newVName "ctid" + letBindNames_ (op_inps ++ patternNames map_pat) $ Op $ + Out.Combine (Out.combineSpace [(ctid, w)]) (lambdaReturnType map_fun) [] fold_body + return op_inps + +intraGroupParalleliseBody :: (MonadFreshNames m, HasScope Out.Kernels m) => + Dependencies -> Names -> VName -> Body + -> m ([[SubExp]], [[SubExp]], Out.KernelBody Out.InKernel) +intraGroupParalleliseBody deps group_variant ltid body = do + (min_ws, avail_ws, kstms) <- runIntraGroupM (Env ltid deps group_variant) $ + mapM_ intraGroupStm $ bodyStms body + return (min_ws, avail_ws, + KernelBody () kstms $ map (ThreadsReturn OneResultPerGroup) $ bodyResult body) diff --git a/src/Futhark/Pass/ExtractKernels/Kernelise.hs b/src/Futhark/Pass/ExtractKernels/Kernelise.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/Kernelise.hs @@ -0,0 +1,283 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | Sequentialise to kernel statements. +module Futhark.Pass.ExtractKernels.Kernelise + ( transformStm + , transformStms + , transformBody + , transformLambda + , mapIsh + + , groupStreamMapAccumL + ) + where + +import Control.Monad +import Data.Semigroup ((<>)) +import qualified Data.Set as S + +import qualified Futhark.Analysis.Alias as Alias +import qualified Futhark.Transform.FirstOrderTransform as FOT +import Futhark.Representation.SOACS +import qualified Futhark.Representation.Kernels as Out +import Futhark.MonadFreshNames +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Tools + +type Transformer m = (MonadBinder m, + Lore m ~ Out.InKernel, + LocalScope (Lore m) m) + +transformStms :: Transformer m => Stms SOACS -> m () +transformStms = mapM_ transformStm . stmsToList + +transformStm :: Transformer m => Stm -> m () + +transformStm (Let pat aux (Op (Screma w form arrs))) + -- No map-out part + | Just (_, red_lam, nes, map_lam) <- isRedomapSOAC form, + patternSize pat == length nes = do + + fold_lam <- composeLambda nilFn red_lam map_lam + + chunk_size <- newVName "chunk_size" + chunk_offset <- newVName "chunk_offset" + let arr_idents = drop (length nes) $ patternIdents pat + (fold_acc_params, fold_elem_params) = + splitAt (length nes) $ lambdaParams fold_lam + arr_chunk_params <- mapM (mkArrChunkParam $ Var chunk_size) fold_elem_params + + map_arr_params <- forM arr_idents $ \arr -> + newParam (baseString (identName arr) <> "_in") $ + setOuterSize (identType arr) (Var chunk_size) + + fold_acc_params' <- forM fold_acc_params $ \p -> + newParam (baseString $ paramName p) $ paramType p + + let param_scope = + scopeOfLParams $ fold_acc_params' ++ arr_chunk_params ++ map_arr_params + + redomap_pes <- forM (patternValueElements pat) $ \pe -> + PatElem <$> newVName (baseString $ patElemName pe) <*> pure (patElemType pe) + + redomap_kstms <- collectStms_ $ localScope param_scope $ do + fold_lam' <- transformLambda fold_lam + groupStreamMapAccumL redomap_pes (Var chunk_size) fold_lam' + (map (Var . paramName) fold_acc_params') (map paramName arr_chunk_params) + + let stream_kbody = Out.Body () redomap_kstms $ + map (Var . patElemName) redomap_pes + stream_lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = chunk_size + , Out.groupStreamChunkOffset = chunk_offset + , Out.groupStreamAccParams = fold_acc_params' + , Out.groupStreamArrParams = arr_chunk_params + , Out.groupStreamLambdaBody = stream_kbody + } + + -- Tricky reverse logic: we have to copy all the initial + -- accumulators that were *not* consumed in the original lambda, as + -- a GroupStream will write to its accumulators. + let consumed = consumedByLambda $ Alias.analyseLambda fold_lam + nes' <- forM (zip fold_acc_params nes) $ \(p,e) -> + case e of + Var v | not $ paramName p `S.member` consumed, + not $ primType $ paramType p -> + letSubExp "groupstream_mapaccum_copy" $ BasicOp $ Copy v + _ -> return e + + addStm $ Let pat aux $ Op $ Out.GroupStream w w stream_lam nes' arrs + + where mkArrChunkParam chunk_size arr_param = + newParam (baseString (paramName arr_param) <> "_chunk") $ + arrayOfRow (paramType arr_param) chunk_size + +transformStm (Let pat aux (Op (Stream w (Sequential accs) fold_lam arrs))) = do + let ret = lambdaReturnType fold_lam + -- Sequential streams can be transformed easily to a GroupStream. + -- But we have to create accumulator parameters for mapout. + + chunk_offset <- newVName "streamseq_chunk_offset" + + let (chunk_size_param, fold_acc_params, arr_chunk_params) = + partitionChunkedFoldParameters (length accs) $ lambdaParams fold_lam + chunk_size = paramName chunk_size_param + map_arr_tps = map (`setOuterSize` w) $ drop (length accs) ret + + mapout_arrs <- resultArray map_arr_tps + outarr_params <- forM map_arr_tps $ \map_arr_t -> + Param <$> newVName "redomap_outarr" <*> pure map_arr_t + + lam_body <- localScope (castScope (scopeOf fold_lam) <> + scopeOfLParams outarr_params) $ insertStmsM $ do + res <- bodyBind =<< transformBody (lambdaBody fold_lam) + -- Some results are to be returned; others to be copied into the + -- map-out arrays. + let (acc_res, mapout_res) = splitAt (length accs) res + + mapout_res' <- forM (zip outarr_params mapout_res) $ \(p, r) -> + let slice = fullSlice (paramType p) + [DimSlice (Var chunk_offset) (Var chunk_size) (constant (1::Int32))] + in fmap Var $ letInPlace "mapout_res" (paramName p) slice $ BasicOp $ SubExp r + + return $ resultBody $ acc_res++mapout_res' + + let stream_lam = Out.GroupStreamLambda + { Out.groupStreamChunkSize = chunk_size + , Out.groupStreamChunkOffset = chunk_offset + , Out.groupStreamAccParams = fold_acc_params ++ outarr_params + , Out.groupStreamArrParams = arr_chunk_params + , Out.groupStreamLambdaBody = lam_body + } + + -- Only copy the accs that were not consumed in the original stream. + let consumed = consumedByLambda $ Alias.analyseLambda fold_lam + accs' <- forM (zip fold_acc_params accs) $ \(p, acc) -> + case acc of + Var v | not $ paramName p `S.member` consumed, + not $ primType $ paramType p -> + letSubExp "streamseq_acc_copy" $ BasicOp $ Copy v + _ -> return acc + + addStm $ Let pat aux $ Op $ + Out.GroupStream w w stream_lam (accs'++map Var mapout_arrs) arrs + +transformStm (Let pat aux (DoLoop [] val (ForLoop i Int32 bound []) body)) = do + dummy_chunk_size <- newVName "dummy_chunk_size" + body' <- localScope (scopeOfFParams (map fst val)) $ transformBody body + let lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = dummy_chunk_size + , Out.groupStreamChunkOffset = i + , Out.groupStreamAccParams = map (fmap fromDecl . fst) val + , Out.groupStreamArrParams = [] + , Out.groupStreamLambdaBody = body' } + + -- Copy the initial merge parameters that were not unique in the + -- original stream. + accs' <- forM val $ \(p, initial) -> + case initial of + Var v | not $ unique $ paramDeclType p, + not $ primType $ paramDeclType p -> + letSubExp "streamseq_merge_copy" $ BasicOp $ Copy v + _ -> return initial + + addStm $ Let pat aux $ Op $ Out.GroupStream + bound (constant (1::Int32)) lam accs' [] + +transformStm (Let pat aux (If cond tb fb ts)) = do + tb' <- transformBody tb + fb' <- transformBody fb + addStm $ Let pat aux $ If cond tb' fb' ts + +transformStm bnd = + FOT.transformStmRecursively bnd + +transformBody :: Transformer m => Body -> m (Out.Body Out.InKernel) +transformBody (Body attr bnds res) = do + stms <- collectStms_ $ transformStms bnds + return $ Out.Body attr stms res + +transformLambda :: (MonadFreshNames m, + HasScope lore m, + SameScope lore Out.InKernel) => + Lambda -> m (Out.Lambda Out.InKernel) +transformLambda (Lambda params body rettype) = do + body' <- runBodyBinder $ + localScope (scopeOfLParams params) $ + transformBody body + return $ Lambda params body' rettype + +groupStreamMapAccumL :: Transformer m => + [Out.PatElem Out.InKernel] + -> SubExp + -> Out.Lambda Out.InKernel + -> [SubExp] + -> [VName] + -> m () +groupStreamMapAccumL pes w fold_lam accexps arrexps = do + let acc_num = length accexps + res_tps = lambdaReturnType fold_lam + map_arr_tps = drop acc_num res_tps + + let fold_lam' = fold_lam { lambdaParams = take acc_num $ lambdaParams fold_lam } + fold_lam_aliases = Alias.analyseLambda fold_lam' + + mapout_arrs <- resultArray [ arrayOf t (Shape [w]) NoUniqueness + | t <- map_arr_tps ] + + (merge, i, redomap_loop) <- + FOT.doLoopMapAccumL' w fold_lam_aliases accexps [] mapout_arrs + + -- HACK: we manually inject the indexing here. + dummy_chunk_size <- newVName "groupstream_mapaccum_dummy_chunk_size" + let arr_params = drop acc_num $ lambdaParams fold_lam + arr_params_chunked <- forM arr_params $ \arr_param -> + newParam (baseString (paramName arr_param) <> "_chunked") $ + paramType arr_param `arrayOfRow` Var dummy_chunk_size + let index_bnds = do + (p, arr, arr_t) <- zip3 arr_params (map paramName arr_params_chunked) + (map paramType arr_params_chunked) + return $ mkLet [] [paramIdent p] $ + BasicOp $ Index arr $ fullSlice arr_t [DimFix $ constant (0::Int32)] + + let redomap_kbody = stmsFromList index_bnds `insertStms` redomap_loop + acc_params = map (fmap fromDecl . fst) merge + stream_lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = dummy_chunk_size + , Out.groupStreamChunkOffset = i + , Out.groupStreamAccParams = acc_params + , Out.groupStreamArrParams = arr_params_chunked + , Out.groupStreamLambdaBody = redomap_kbody + } + + letBind_ (Pattern [] pes) $ Op $ + Out.GroupStream w (constant (1::Int32)) stream_lam (accexps++map Var mapout_arrs) arrexps + +resultArray :: MonadBinder m => [Type] -> m [VName] +resultArray = mapM oneArray + where oneArray t = letExp "result" $ BasicOp $ Scratch (elemType t) (arrayDims t) + +mapIsh :: Transformer m => + Pattern + -> Certificates + -> SubExp + -> [LParam] + -> Out.Body Out.InKernel + -> [VName] + -> m () +mapIsh pat cs w params (Out.Body () kstms kres) arrs = do + i <- newVName "i" + + outarrs <- resultArray $ patternTypes pat + + outarr_params <- forM (patternElements pat) $ \pe -> + newParam (baseString (patElemName pe) <> "_out") $ + patElemType pe + + dummy_chunk_size <- newVName "dummy_chunk_size" + params_chunked <- forM params $ \param -> + newParam (baseString (paramName param) <> "_chunked") $ + paramType param `arrayOfRow` Var dummy_chunk_size + + (outarr_params_new, write_elems) <- + fmap unzip $ forM (zip outarr_params kres) $ \(outarr_param, se) -> do + outarr_param_new <- newParam' (<>"_new") outarr_param + return (outarr_param_new, + mkLet [] [paramIdent outarr_param_new] $ BasicOp $ + Update (paramName outarr_param) + (fullSlice (paramType outarr_param) [DimFix $ Var i]) se) + + let index_stms = do + (p, arr, arr_t) <- zip3 params (map paramName params_chunked) $ + map paramType params_chunked + return $ mkLet [] [paramIdent p] $ + BasicOp $ Index arr $ fullSlice arr_t [DimFix $ constant (0::Int32)] + kbody' = Out.Body () (stmsFromList index_stms <> kstms <> stmsFromList write_elems) $ + map (Var . paramName) outarr_params_new + + let stream_lam = Out.GroupStreamLambda { Out.groupStreamChunkSize = dummy_chunk_size + , Out.groupStreamChunkOffset = i + , Out.groupStreamAccParams = outarr_params + , Out.groupStreamArrParams = params_chunked + , Out.groupStreamLambdaBody = kbody' + } + certifying cs $ addStm $ Let pat (StmAux cs ()) $ + Op $ Out.GroupStream w (constant (1::Int32)) stream_lam (map Var outarrs) arrs diff --git a/src/Futhark/Pass/ExtractKernels/Segmented.hs b/src/Futhark/Pass/ExtractKernels/Segmented.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ExtractKernels/Segmented.hs @@ -0,0 +1,899 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | Multiversion segmented reduction. +module Futhark.Pass.ExtractKernels.Segmented + ( regularSegmentedRedomap + , regularSegmentedScan + ) + where + +import Control.Monad +import qualified Data.Map.Strict as M +import Data.Semigroup ((<>)) + +import Futhark.Transform.Rename +import Futhark.Representation.Kernels +import Futhark.Representation.SOACS.SOAC (nilFn) +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Pass.ExtractKernels.BlockedKernel + +data SegmentedVersion = OneGroupOneSegment + | ManyGroupsOneSegment + deriving (Eq, Ord, Show) + +-- | @regularSegmentedRedomap@ will generate code for a segmented redomap using +-- 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 the input array +-- will be transposed (by the KernelBabysitter) to enable memory coalesced +-- accesses. Currently we will always make each thread read as many elements +-- as it can, but this /could/ be unfavorable because of the transpose: in +-- the case where each thread can only read 2 elements, the cost of the +-- transpose might not be worth the performance gained by letting each thread +-- read multiple elements. This could be investigated more in depth in the +-- future (TODO) +-- +-- * 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 normally 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. +regularSegmentedRedomap :: (HasScope Kernels m, MonadBinder m, Lore m ~ Kernels) => + SubExp -- segment_size + -> SubExp -- num_segments + -> [SubExp] -- nest_sizes = the sizes of the maps on "top" of this redomap + -> Pattern Kernels -- flat_pat ... pat where each type is array with dim [w] + -> Pattern Kernels -- pat + -> SubExp -- w = total_num_elements + -> Commutativity -- comm + -> Lambda InKernel -- reduce_lam + -> Lambda InKernel -- fold_lam = this lambda performs both the map-part and + -- reduce-part of a redomap (described in redomap paper) + -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this redomap + -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace + -> [SubExp] -- nes + -> [VName] -- arrs_flat + -> m () +regularSegmentedRedomap segment_size num_segments nest_sizes flat_pat + pat w comm reduce_lam fold_lam ispace inps nes arrs_flat = do + unless (null $ patternContextElements pat) $ fail "regularSegmentedRedomap result pattern contains context elements, and Rasmus did not think this would ever happen." + + -- the result of the "map" part of a redomap has to be stored somewhere within + -- the chunking loop of a kernel. The current way to do this is to make some + -- scratch space initially, and each thread will get a part of this by + -- splitting it. Finally it is returned as a result of the kernel (to not + -- break functional semantics). + map_out_arrs <- forM (drop num_redres $ patternIdents pat) $ \(Ident name t) -> do + tmp <- letExp (baseString name <> "_out_in") $ + BasicOp $ Scratch (elemType t) (arrayDims t) + -- This reshape will not always work. + letExp (baseString name ++ "_out_in") $ + BasicOp $ Reshape (reshapeOuter [DimNew w] (length nest_sizes+1) $ arrayShape t) tmp + + -- Check that we're only dealing with arrays with dimension [w] + forM_ arrs_flat $ \arr -> do + tp <- lookupType arr + case tp of + -- TODO: this won't work if the reduction operator works on lists... but + -- they seem to be handled in some other way (which makes sense) + Array _primtp (Shape (flatsize:_)) _uniqness -> + when (flatsize /= w) $ + fail$ "regularSegmentedRedomap: first dimension of array has incorrect size " ++ pretty arr ++ ":" ++ pretty tp + _ -> + fail $ "regularSegmentedRedomap: non array encountered " ++ pretty arr ++ ":" ++ pretty tp + + -- The pattern passed to chunkLambda must have exactly *one* array dimension, + -- to get the correct size of [chunk_size]type. + -- + -- TODO: not sure if this will work when result of map is multidimensional, + -- or if reduction operator uses lists... must check + chunk_pat <- fmap (Pattern []) $ forM (patternValueElements pat) $ \pat_e -> + case patElemType pat_e of + Array ty (Shape (dim0:_)) u -> do + vn' <- newName $ patElemName pat_e + return $ PatElem vn' $ Array ty (Shape [dim0]) u + _ -> fail $ "segmentedRedomap: result pattern is not array " ++ pretty pat_e + + chunk_fold_lam <- chunkLambda chunk_pat nes fold_lam + + kern_chunk_fold_lam <- kerneliseLambda nes chunk_fold_lam + + let chunk_red_pat = Pattern [] $ take num_redres $ patternValueElements chunk_pat + kern_chunk_reduce_lam <- kerneliseLambda nes =<< chunkLambda chunk_red_pat nes reduce_lam + + -- the lambda for a GroupReduce needs these two extra parameters + my_index <- newVName "my_index" + other_offset <- newVName "other_offset" + let my_index_param = Param my_index (Prim int32) + let other_offset_param = Param other_offset (Prim int32) + let reduce_lam' = reduce_lam { lambdaParams = my_index_param : + other_offset_param : + lambdaParams reduce_lam + } + flag_reduce_lam <- addFlagToLambda nes reduce_lam + let flag_reduce_lam' = flag_reduce_lam { lambdaParams = my_index_param : + other_offset_param : + lambdaParams flag_reduce_lam + } + + + -- TODO: 'blockedReductionStream' in BlockedKernel.hs which is very similar + -- performs a copy here... however, I have not seen a need for it yet. + + group_size <- getSize "group_size" SizeGroup + num_groups_hint <- getSize "num_groups_hint" SizeNumGroups + + -- Here we make a small optimization: if we will use the large kernel, and + -- only one group per segment, we can simplify the calcualtions within the + -- kernel for the indexes of which segment is it working on; therefore we + -- create two different kernels (this will increase the final code size a bit + -- though). TODO: test how much we win by doing this. + + (num_groups_per_segment, _) <- + calcGroupsPerSegmentAndElementsPerThread + segment_size num_segments num_groups_hint group_size ManyGroupsOneSegment + + let all_arrs = arrs_flat ++ map_out_arrs + (large_1_ses, large_1_stms) <- runBinder $ + useLargeOnePerSeg group_size all_arrs reduce_lam' kern_chunk_fold_lam + (large_m_ses, large_m_stms) <- runBinder $ + useLargeMultiRecursiveReduce group_size all_arrs reduce_lam' kern_chunk_fold_lam + kern_chunk_reduce_lam flag_reduce_lam' + + let e_large_seg = eIf (eCmpOp (CmpEq $ IntType Int32) (eSubExp num_groups_per_segment) + (eSubExp one)) + (mkBodyM large_1_stms large_1_ses) + (mkBodyM large_m_stms large_m_ses) + + + (small_ses, small_stms) <- runBinder $ useSmallKernel group_size map_out_arrs flag_reduce_lam' + + -- if (group_size/2) < segment_size, means that we will not be able to fit two + -- segments into one group, and therefore we should not use the kernel that + -- relies on this. + e <- eIf (eCmpOp (CmpSlt Int32) (eBinOp (SQuot Int32) (eSubExp group_size) (eSubExp two)) + (eSubExp segment_size)) + (eBody [e_large_seg]) + (mkBodyM small_stms small_ses) + + redres_pes <- forM (take num_redres (patternValueElements pat)) $ \pe -> do + vn' <- newName $ patElemName pe + return $ PatElem vn' $ replaceSegmentDims num_segments $ patElemType pe + let mapres_pes = drop num_redres $ patternValueElements flat_pat + let unreshaped_pat = Pattern [] $ redres_pes ++ mapres_pes + + letBind_ unreshaped_pat e + + forM_ (zip (patternValueElements unreshaped_pat) + (patternValueElements pat)) $ \(kpe, pe) -> + letBind_ (Pattern [] [pe]) $ + BasicOp $ Reshape [DimNew se | se <- arrayDims $ patElemAttr pe] + (patElemName kpe) + + where + replaceSegmentDims d t = + t `setArrayDims` (d : drop (length nest_sizes) (arrayDims t)) + + one = constant (1 :: Int32) + two = constant (2 :: Int32) + + -- number of reduction results (tuple size for reduction operator) + num_redres = length nes + + ---------------------------------------------------------------------------- + -- The functions below generate all the needed code for the two different + -- version of segmented-redomap (one group per segment, and many groups per + -- segment). + -- + -- We rename statements before adding them because the same lambdas + -- (reduce/fold) are used multiple times, and we do not want to bind the + -- same VName twice (as this is a type error) + ---------------------------------------------------------------------------- + useLargeOnePerSeg group_size all_arrs reduce_lam' kern_chunk_fold_lam = do + mapres_pes <- forM (drop num_redres $ patternValueElements flat_pat) $ \pe -> do + vn' <- newName $ patElemName pe + return $ PatElem vn' $ patElemType pe + + (kernel, _, _) <- + largeKernel group_size segment_size num_segments nest_sizes + all_arrs comm reduce_lam' kern_chunk_fold_lam + nes w OneGroupOneSegment + ispace inps + + kernel_redres_pes <- forM (take num_redres (patternValueElements pat)) $ \pe -> do + vn' <- newName $ patElemName pe + return $ PatElem vn' $ replaceSegmentDims num_segments $ patElemType pe + + let kernel_pat = Pattern [] $ kernel_redres_pes ++ mapres_pes + + addStm =<< renameStm (Let kernel_pat (defAux ()) $ Op kernel) + return $ map (Var . patElemName) $ patternValueElements kernel_pat + + ---------------------------------------------------------------------------- + useLargeMultiRecursiveReduce group_size all_arrs reduce_lam' kern_chunk_fold_lam kern_chunk_reduce_lam flag_reduce_lam' = do + mapres_pes <- forM (drop num_redres $ patternValueElements flat_pat) $ \pe -> do + vn' <- newName $ patElemName pe + return $ PatElem vn' $ patElemType pe + + (firstkernel, num_groups_used, num_groups_per_segment) <- + largeKernel group_size segment_size num_segments nest_sizes + all_arrs comm reduce_lam' kern_chunk_fold_lam + nes w ManyGroupsOneSegment + ispace inps + + firstkernel_redres_pes <- forM (take num_redres (patternValueElements pat)) $ \pe -> do + vn' <- newName $ patElemName pe + return $ PatElem vn' $ replaceSegmentDims num_groups_used $ patElemType pe + + let first_pat = Pattern [] $ firstkernel_redres_pes ++ mapres_pes + addStm =<< renameStm (Let first_pat (defAux ()) $ Op firstkernel) + + let new_segment_size = num_groups_per_segment + let new_total_elems = num_groups_used + let tmp_redres = map patElemName firstkernel_redres_pes + + (finalredres, part_two_stms) <- runBinder $ performFinalReduction + new_segment_size new_total_elems tmp_redres + reduce_lam' kern_chunk_reduce_lam flag_reduce_lam' + + mapM_ (addStm <=< renameStm) part_two_stms + + return $ finalredres ++ map (Var . patElemName) mapres_pes + + ---------------------------------------------------------------------------- + -- The "recursive" reduction step. However, will always do this using + -- exactly one extra step. Either by using the small kernel, or by using the + -- large kernel with one group per segment. + performFinalReduction new_segment_size new_total_elems tmp_redres + reduce_lam' kern_chunk_reduce_lam flag_reduce_lam' = do + group_size <- getSize "group_size" SizeGroup + + -- Large kernel, using one group per segment (ogps) + (large_ses, large_stms) <- runBinder $ do + (large_kernel, _, _) <- largeKernel group_size new_segment_size num_segments nest_sizes + tmp_redres comm reduce_lam' kern_chunk_reduce_lam + nes new_total_elems OneGroupOneSegment + ispace inps + letTupExp' "kernel_result" $ Op large_kernel + + -- Small kernel, using one group many segments (ogms) + (small_ses, small_stms) <- runBinder $ do + red_scratch_arrs <- forM (take num_redres $ patternIdents pat) $ \(Ident name t) -> do + -- We construct a scratch array for writing the result, but + -- we have to flatten the dimensions corresponding to the + -- map nest, because multi-dimensional WriteReturns are/were + -- not supported. + tmp <- letExp (baseString name <> "_redres_scratch") $ + BasicOp $ Scratch (elemType t) (arrayDims t) + let reshape = reshapeOuter [DimNew num_segments] (length nest_sizes) $ arrayShape t + letExp (baseString name ++ "_redres_scratch") $ + BasicOp $ Reshape reshape tmp + kernel <- smallKernel group_size new_segment_size num_segments + tmp_redres red_scratch_arrs + comm flag_reduce_lam' reduce_lam + nes new_total_elems ispace inps + letTupExp' "kernel_result" $ Op kernel + + e <- eIf (eCmpOp (CmpSlt Int32) + (eBinOp (SQuot Int32) (eSubExp group_size) (eSubExp two)) + (eSubExp new_segment_size)) + (mkBodyM large_stms large_ses) + (mkBodyM small_stms small_ses) + + letTupExp' "step_two_kernel_result" e + + ---------------------------------------------------------------------------- + useSmallKernel group_size map_out_arrs flag_reduce_lam' = do + red_scratch_arrs <- + forM (take num_redres $ patternIdents pat) $ \(Ident name t) -> do + tmp <- letExp (baseString name <> "_redres_scratch") $ + BasicOp $ Scratch (elemType t) (arrayDims t) + let shape_change = reshapeOuter [DimNew num_segments] + (length nest_sizes) (arrayShape t) + letExp (baseString name ++ "_redres_scratch") $ + BasicOp $ Reshape shape_change tmp + + let scratch_arrays = red_scratch_arrs ++ map_out_arrs + + kernel <- smallKernel group_size segment_size num_segments + arrs_flat scratch_arrays + comm flag_reduce_lam' fold_lam + nes w ispace inps + letTupExp' "kernel_result" $ Op kernel + +largeKernel :: (MonadBinder m, Lore m ~ Kernels) => + SubExp -- group_size + -> SubExp -- segment_size + -> SubExp -- num_segments + -> [SubExp] -- nest sizes + -> [VName] -- all_arrs: flat arrays (also the "map_out" ones) + -> Commutativity -- comm + -> Lambda InKernel -- reduce_lam + -> Lambda InKernel -- kern_chunk_fold_lam + -> [SubExp] -- nes + -> SubExp -- w = total_num_elements + -> SegmentedVersion -- segver + -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this redomap + -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace + -> m (Kernel InKernel, SubExp, SubExp) +largeKernel group_size segment_size num_segments nest_sizes all_arrs comm + reduce_lam' kern_chunk_fold_lam + nes w segver ispace inps = do + let num_redres = length nes -- number of reduction results (tuple size for + -- reduction operator) + + num_groups_hint <- getSize "num_groups_hint" SizeNumGroups + + (num_groups_per_segment, elements_per_thread) <- + calcGroupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size segver + + num_groups <- letSubExp "num_groups" $ + case segver of + OneGroupOneSegment -> BasicOp $ SubExp num_segments + ManyGroupsOneSegment -> BasicOp $ BinOp (Mul Int32) num_segments num_groups_per_segment + + num_threads <- letSubExp "num_threads" $ + BasicOp $ BinOp (Mul Int32) num_groups group_size + + threads_within_segment <- letSubExp "threads_within_segment" $ + BasicOp $ BinOp (Mul Int32) group_size num_groups_per_segment + + gtid_vn <- newVName "gtid" + gtid_ln <- newVName "gtid" + + -- the array passed here is the structure for how to layout the kernel space + space <- newKernelSpace (num_groups, group_size, num_threads) $ + FlatThreadSpace $ ispace ++ [(gtid_vn, num_groups_per_segment),(gtid_ln,group_size)] + + let red_ts = take num_redres $ lambdaReturnType kern_chunk_fold_lam + let map_ts = map rowType $ drop num_redres $ lambdaReturnType kern_chunk_fold_lam + let kernel_return_types = red_ts ++ map_ts + + let ordering = case comm of Commutative -> SplitStrided threads_within_segment + Noncommutative -> SplitContiguous + + let stride = case ordering of SplitStrided s -> s + SplitContiguous -> one + + let each_thread = do + segment_index <- letSubExp "segment_index" $ + BasicOp $ BinOp (SQuot Int32) (Var $ spaceGroupId space) num_groups_per_segment + + -- localId + (group_size * (groupId % num_groups_per_segment)) + index_within_segment <- letSubExp "index_within_segment" =<< + eBinOp (Add Int32) + (eSubExp $ Var gtid_ln) + (eBinOp (Mul Int32) + (eSubExp group_size) + (eBinOp (SRem Int32) (eSubExp $ Var $ spaceGroupId space) (eSubExp num_groups_per_segment)) + ) + + (in_segment_offset,offset) <- + makeOffsetExp ordering index_within_segment elements_per_thread segment_index + + let (_, chunksize, [], arr_params) = + partitionChunkedKernelFoldParameters 0 $ lambdaParams kern_chunk_fold_lam + let chunksize_se = Var $ paramName chunksize + + patelems_res_of_split <- forM arr_params $ \arr_param -> do + let chunk_t = paramType arr_param `setOuterSize` Var (paramName chunksize) + return $ PatElem (paramName arr_param) chunk_t + + letBind_ (Pattern [] [PatElem (paramName chunksize) $ paramType chunksize]) $ + Op $ SplitSpace ordering segment_size index_within_segment elements_per_thread + + addKernelInputStms inps + + forM_ (zip all_arrs patelems_res_of_split) $ \(arr, pe) -> do + let pe_t = patElemType pe + segment_dims = nest_sizes ++ arrayDims (pe_t `setOuterSize` segment_size) + arr_nested <- letExp (baseString arr ++ "_nested") $ + BasicOp $ Reshape (map DimNew segment_dims) arr + arr_nested_t <- lookupType arr_nested + let slice = fullSlice arr_nested_t $ map (DimFix . Var . fst) ispace ++ + [DimSlice in_segment_offset chunksize_se stride] + letBind_ (Pattern [] [pe]) $ BasicOp $ Index arr_nested slice + + red_pes <- forM red_ts $ \red_t -> do + pe_name <- newVName "chunk_fold_red" + return $ PatElem pe_name red_t + map_pes <- forM map_ts $ \map_t -> do + pe_name <- newVName "chunk_fold_map" + return $ PatElem pe_name $ map_t `arrayOfRow` chunksize_se + + -- we add the lets here, as we practially don't know if the resulting subexp + -- is a Constant or a Var, so better be safe (?) + addStms $ bodyStms (lambdaBody kern_chunk_fold_lam) + addStms $ stmsFromList + [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se + | (pe,se) <- zip (red_pes ++ map_pes) + (bodyResult $ lambdaBody kern_chunk_fold_lam) ] + + -- Combine the reduction results from each thread. This will put results in + -- local memory, so a GroupReduce can be performed on them + combine_red_pes <- forM red_ts $ \red_t -> do + pe_name <- newVName "chunk_fold_red" + return $ PatElem pe_name $ red_t `arrayOfRow` group_size + cids <- replicateM (length red_pes) $ newVName "cid" + addStms $ stmsFromList + [ Let (Pattern [] [pe']) (defAux ()) $ + Op $ Combine (combineSpace [(cid, group_size)]) [patElemType pe] [] $ + Body () mempty [Var $ patElemName pe] + | (cid, pe', pe) <- zip3 cids combine_red_pes red_pes ] + + final_red_pes <- forM (lambdaReturnType reduce_lam') $ \t -> do + pe_name <- newVName "final_result" + return $ PatElem pe_name t + letBind_ (Pattern [] final_red_pes) $ + Op $ GroupReduce group_size reduce_lam' $ + zip nes (map patElemName combine_red_pes) + + return (final_red_pes, map_pes, offset) + + + ((final_red_pes, map_pes, offset), stms) <- runBinder each_thread + + red_returns <- forM final_red_pes $ \pe -> + return $ ThreadsReturn OneResultPerGroup $ Var $ patElemName pe + map_returns <- forM map_pes $ \pe -> + return $ ConcatReturns ordering w elements_per_thread + (Just offset) $ + patElemName pe + let kernel_returns = red_returns ++ map_returns + + let kerneldebughints = KernelDebugHints kernelname + [ ("num_segment", num_segments) + , ("segment_size", segment_size) + , ("num_groups", num_groups) + , ("group_size", group_size) + , ("elements_per_thread", elements_per_thread) + , ("num_groups_per_segment", num_groups_per_segment) + ] + + let kernel = Kernel kerneldebughints space kernel_return_types $ + KernelBody () stms kernel_returns + + return (kernel, num_groups, num_groups_per_segment) + + where + one = constant (1 :: Int32) + + commname = case comm of Commutative -> "comm" + Noncommutative -> "nocomm" + + kernelname = case segver of + OneGroupOneSegment -> "segmented_redomap__large_" ++ commname ++ "_one" + ManyGroupsOneSegment -> "segmented_redomap__large_" ++ commname ++ "_many" + + makeOffsetExp SplitContiguous index_within_segment elements_per_thread segment_index = do + in_segment_offset <- letSubExp "in_segment_offset" $ + BasicOp $ BinOp (Mul Int32) elements_per_thread index_within_segment + offset <- letSubExp "offset" =<< eBinOp (Add Int32) (eSubExp in_segment_offset) + (eBinOp (Mul Int32) (eSubExp segment_size) (eSubExp segment_index)) + return (in_segment_offset, offset) + makeOffsetExp (SplitStrided _) index_within_segment _elements_per_thread segment_index = do + offset <- letSubExp "offset" =<< eBinOp (Add Int32) (eSubExp index_within_segment) + (eBinOp (Mul Int32) (eSubExp segment_size) (eSubExp segment_index)) + return (index_within_segment, offset) + +calcGroupsPerSegmentAndElementsPerThread :: (MonadBinder m, Lore m ~ Kernels) => + SubExp + -> SubExp + -> SubExp + -> SubExp + -> SegmentedVersion + -> m (SubExp, SubExp) +calcGroupsPerSegmentAndElementsPerThread segment_size num_segments + num_groups_hint group_size segver = do + num_groups_per_segment_hint <- + letSubExp "num_groups_per_segment_hint" =<< + case segver of + OneGroupOneSegment -> eSubExp one + ManyGroupsOneSegment -> eDivRoundingUp Int32 (eSubExp num_groups_hint) + (eSubExp num_segments) + elements_per_thread <- + letSubExp "elements_per_thread" =<< + eDivRoundingUp Int32 (eSubExp segment_size) + (eBinOp (Mul Int32) (eSubExp group_size) + (eSubExp num_groups_per_segment_hint)) + + -- if we are using 1 element per thread, we might be launching too many + -- groups. This expression will remedy this. + -- + -- For example, if there are 3 segments of size 512, we are using group size + -- 128, and @num_groups_hint@ is 256; then we would use 1 element per thread, + -- and launch 256 groups. However, we only need 4 groups per segment to + -- process all elements. + num_groups_per_segment <- + letSubExp "num_groups_per_segment" =<< + case segver of + OneGroupOneSegment -> eSubExp one + ManyGroupsOneSegment -> + eIf (eCmpOp (CmpEq $ IntType Int32) (eSubExp elements_per_thread) (eSubExp one)) + (eBody [eDivRoundingUp Int32 (eSubExp segment_size) (eSubExp group_size)]) + (mkBodyM mempty [num_groups_per_segment_hint]) + + return (num_groups_per_segment, elements_per_thread) + + where + one = constant (1 :: Int32) + +smallKernel :: (MonadBinder m, Lore m ~ Kernels) => + SubExp -- group_size + -> SubExp -- segment_size + -> SubExp -- num_segments + -> [VName] -- in_arrs: flat arrays (containing input to fold_lam) + -> [VName] -- scratch_arrs: Preallocated space that we can write into + -> Commutativity -- comm + -> Lambda InKernel -- flag_reduce_lam' + -> Lambda InKernel -- fold_lam + -> [SubExp] -- nes + -> SubExp -- w = total_num_elements + -> [(VName, SubExp)] -- ispace = pair of (gtid, size) for the maps on "top" of this redomap + -> [KernelInput] -- inps = inputs that can be looked up by using the gtids from ispace + -> m (Kernel InKernel) +smallKernel group_size segment_size num_segments in_arrs scratch_arrs + comm flag_reduce_lam' fold_lam_unrenamed + nes w ispace inps = do + let num_redres = length nes -- number of reduction results (tuple size for + -- reduction operator) + + fold_lam <- renameLambda fold_lam_unrenamed + + num_segments_per_group <- letSubExp "num_segments_per_group" $ + BasicOp $ BinOp (SQuot Int32) group_size segment_size + + num_groups <- letSubExp "num_groups" =<< + eDivRoundingUp Int32 (eSubExp num_segments) (eSubExp num_segments_per_group) + + num_threads <- letSubExp "num_threads" $ + BasicOp $ BinOp (Mul Int32) num_groups group_size + + active_threads_per_group <- letSubExp "active_threads_per_group" $ + BasicOp $ BinOp (Mul Int32) segment_size num_segments_per_group + + let remainder_last_group = eBinOp (SRem Int32) (eSubExp num_segments) (eSubExp num_segments_per_group) + + segments_in_last_group <- letSubExp "seg_in_last_group" =<< + eIf (eCmpOp (CmpEq $ IntType Int32) remainder_last_group + (eSubExp zero)) + (eBody [eSubExp num_segments_per_group]) + (eBody [remainder_last_group]) + + active_threads_in_last_group <- letSubExp "active_threads_last_group" $ + BasicOp $ BinOp (Mul Int32) segment_size segments_in_last_group + + -- the array passed here is the structure for how to layout the kernel space + space <- newKernelSpace (num_groups, group_size, num_threads) $ + FlatThreadSpace [] + + ------------------------------------------------------------------------------ + -- What follows is the statements used in the kernel + ------------------------------------------------------------------------------ + + let lid = Var $ spaceLocalId space + + let (red_ts, map_ts) = splitAt num_redres $ lambdaReturnType fold_lam + let kernel_return_types = red_ts ++ map_ts + + let wasted_thread_part1 = do + let create_dummy_val (Prim ty) = return $ Constant $ blankPrimValue ty + create_dummy_val (Array ty sh _) = letSubExp "dummy" $ BasicOp $ Scratch ty (shapeDims sh) + create_dummy_val Mem{} = fail "segredomap, 'Mem' used as result type" + dummy_vals <- mapM create_dummy_val kernel_return_types + return (negone : dummy_vals) + + let normal_thread_part1 = do + segment_index <- letSubExp "segment_index" =<< + eBinOp (Add Int32) + (eBinOp (SQuot Int32) (eSubExp $ Var $ spaceLocalId space) (eSubExp segment_size)) + (eBinOp (Mul Int32) (eSubExp $ Var $ spaceGroupId space) (eSubExp num_segments_per_group)) + + index_within_segment <- letSubExp "index_within_segment" =<< + eBinOp (SRem Int32) (eSubExp $ Var $ spaceLocalId space) (eSubExp segment_size) + + offset <- makeOffsetExp index_within_segment segment_index + + red_pes <- forM red_ts $ \red_t -> do + pe_name <- newVName "fold_red" + return $ PatElem pe_name red_t + map_pes <- forM map_ts $ \map_t -> do + pe_name <- newVName "fold_map" + return $ PatElem pe_name map_t + + addManualIspaceCalcStms segment_index ispace + + addKernelInputStms inps + + -- Index input array to get arguments to fold_lam + let arr_params = drop num_redres $ lambdaParams fold_lam + let nonred_lamparam_pes = map + (\p -> PatElem (paramName p) (paramType p)) arr_params + forM_ (zip in_arrs nonred_lamparam_pes) $ \(arr, pe) -> do + tp <- lookupType arr + let slice = fullSlice tp [DimFix offset] + letBind_ (Pattern [] [pe]) $ BasicOp $ Index arr slice + + -- Bind neutral element (serves as the reduction arguments to fold_lam) + forM_ (zip nes (take num_redres $ lambdaParams fold_lam)) $ \(ne,param) -> do + let pe = PatElem (paramName param) (paramType param) + letBind_ (Pattern [] [pe]) $ BasicOp $ SubExp ne + + addStms $ bodyStms $ lambdaBody fold_lam + + -- we add the lets here, as we practially don't know if the resulting subexp + -- is a Constant or a Var, so better be safe (?) + addStms $ stmsFromList + [ Let (Pattern [] [pe]) (defAux ()) $ BasicOp $ SubExp se + | (pe,se) <- zip (red_pes ++ map_pes) (bodyResult $ lambdaBody fold_lam) ] + + let mapoffset = offset + let mapret_elems = map (Var . patElemName) map_pes + let redres_elems = map (Var . patElemName) red_pes + return (mapoffset : redres_elems ++ mapret_elems) + + let all_threads red_pes = do + isfirstinsegment <- letExp "isfirstinsegment" =<< + eCmpOp (CmpEq $ IntType Int32) + (eBinOp (SRem Int32) (eSubExp lid) (eSubExp segment_size)) + (eSubExp zero) + + -- We will perform a segmented-scan, so all the prime variables here + -- include the flag, which is the first argument to flag_reduce_lam + let red_pes_wflag = PatElem isfirstinsegment (Prim Bool) : red_pes + let red_ts_wflag = Prim Bool : red_ts + + -- Combine the reduction results from each thread. This will put results in + -- local memory, so a GroupReduce/GroupScan can be performed on them + combine_red_pes' <- forM red_ts_wflag $ \red_t -> do + pe_name <- newVName "chunk_fold_red" + return $ PatElem pe_name $ red_t `arrayOfRow` group_size + cids <- replicateM (length red_pes_wflag) $ newVName "cid" + addStms $ stmsFromList [ Let (Pattern [] [pe']) (defAux ()) $ Op $ + Combine (combineSpace [(cid, group_size)]) [patElemType pe] [] $ + Body () mempty [Var $ patElemName pe] + | (cid, pe', pe) <- zip3 cids combine_red_pes' red_pes_wflag ] + + scan_red_pes_wflag <- forM red_ts_wflag $ \red_t -> do + pe_name <- newVName "scanned" + return $ PatElem pe_name $ red_t `arrayOfRow` group_size + let scan_red_pes = drop 1 scan_red_pes_wflag + letBind_ (Pattern [] scan_red_pes_wflag) $ Op $ + GroupScan group_size flag_reduce_lam' $ + zip (false:nes) (map patElemName combine_red_pes') + + return scan_red_pes + + let normal_thread_part2 scan_red_pes = do + segment_index <- letSubExp "segment_index" =<< + eBinOp (Add Int32) + (eBinOp (SQuot Int32) (eSubExp $ Var $ spaceLocalId space) (eSubExp segment_size)) + (eBinOp (Mul Int32) (eSubExp $ Var $ spaceGroupId space) (eSubExp num_segments_per_group)) + + islastinsegment <- letExp "islastinseg" =<< eCmpOp (CmpEq $ IntType Int32) + (eBinOp (SRem Int32) (eSubExp lid) (eSubExp segment_size)) + (eBinOp (Sub Int32) (eSubExp segment_size) (eSubExp one)) + + redoffset <- letSubExp "redoffset" =<< + eIf (eSubExp $ Var islastinsegment) + (eBody [eSubExp segment_index]) + (mkBodyM mempty [negone]) + + redret_elems <- fmap (map Var) $ letTupExp "red_return_elem" =<< + eIf (eSubExp $ Var islastinsegment) + (eBody [return $ BasicOp $ Index (patElemName pe) (fullSlice (patElemType pe) [DimFix lid]) + | pe <- scan_red_pes]) + (mkBodyM mempty nes) + + return (redoffset : redret_elems) + + + let picknchoose = do + is_last_group <- letSubExp "islastgroup" =<< + eCmpOp (CmpEq $ IntType Int32) + (eSubExp $ Var $ spaceGroupId space) + (eBinOp (Sub Int32) (eSubExp num_groups) (eSubExp one)) + + active_threads_this_group <- letSubExp "active_thread_this_group" =<< + eIf (eSubExp is_last_group) + (eBody [eSubExp active_threads_in_last_group]) + (eBody [eSubExp active_threads_per_group]) + + isactive <- letSubExp "isactive" =<< + eCmpOp (CmpSlt Int32) (eSubExp lid) (eSubExp active_threads_this_group) + + -- Part 1: All active threads reads element from input array and applies + -- folding function. "wasted" threads will just create dummy values + (normal_res1, normal_stms1) <- runBinder normal_thread_part1 + (wasted_res1, wasted_stms1) <- runBinder wasted_thread_part1 + + -- we could just have used letTupExp, but this would not give as nice + -- names in the generated code + mapoffset_pe <- (`PatElem` i32) <$> newVName "mapoffset" + redtmp_pes <- forM red_ts $ \red_t -> do + pe_name <- newVName "redtmp_res" + return $ PatElem pe_name red_t + map_pes <- forM map_ts $ \map_t -> do + pe_name <- newVName "map_res" + return $ PatElem pe_name map_t + + e1 <- eIf (eSubExp isactive) + (mkBodyM normal_stms1 normal_res1) + (mkBodyM wasted_stms1 wasted_res1) + letBind_ (Pattern [] (mapoffset_pe:redtmp_pes++map_pes)) e1 + + -- Part 2: All threads participate in Comine & GroupScan + scan_red_pes <- all_threads redtmp_pes + + -- Part 3: Active thread that are the last element in segment, should + -- write the element from local memory to the output array + (normal_res2, normal_stms2) <- runBinder $ normal_thread_part2 scan_red_pes + + redoffset_pe <- (`PatElem` i32) <$> newVName "redoffset" + red_pes <- forM red_ts $ \red_t -> do + pe_name <- newVName "red_res" + return $ PatElem pe_name red_t + + e2 <- eIf (eSubExp isactive) + (mkBodyM normal_stms2 normal_res2) + (mkBodyM mempty (negone : nes)) + letBind_ (Pattern [] (redoffset_pe:red_pes)) e2 + + return $ map (Var . patElemName) $ redoffset_pe:mapoffset_pe:red_pes++map_pes + + (redoffset:mapoffset:redmapres, stms) <- runBinder picknchoose + let (finalredvals, finalmapvals) = splitAt num_redres redmapres + + -- To be able to only return elements from some threads, we exploit the fact + -- that WriteReturn with offset=-1, won't do anything. + red_returns <- forM (zip finalredvals $ take num_redres scratch_arrs) $ \(se, scarr) -> + return $ WriteReturn [num_segments] scarr [([redoffset], se)] + map_returns <- forM (zip finalmapvals $ drop num_redres scratch_arrs) $ \(se, scarr) -> + return $ WriteReturn [w] scarr [([mapoffset], se)] + let kernel_returns = red_returns ++ map_returns + + let kerneldebughints = KernelDebugHints kernelname + [ ("num_segment", num_segments) + , ("segment_size", segment_size) + , ("num_groups", num_groups) + , ("group_size", group_size) + , ("num_segments_per_group", num_segments_per_group) + , ("active_threads_per_group", active_threads_per_group) + ] + + let kernel = Kernel kerneldebughints space kernel_return_types $ + KernelBody () stms kernel_returns + + return kernel + + where + i32 = Prim $ IntType Int32 + zero = constant (0 :: Int32) + one = constant (1 :: Int32) + negone = constant (-1 :: Int32) + false = constant False + + + commname = case comm of Commutative -> "comm" + Noncommutative -> "nocomm" + kernelname = "segmented_redomap__small_" ++ commname + + makeOffsetExp index_within_segment segment_index = do + e <- eBinOp (Add Int32) + (eSubExp index_within_segment) + (eBinOp (Mul Int32) (eSubExp segment_size) (eSubExp segment_index)) + letSubExp "offset" e + +addKernelInputStms :: (MonadBinder m, Lore m ~ InKernel) => + [KernelInput] + -> m () +addKernelInputStms = mapM_ $ \kin -> do + let pe = PatElem (kernelInputName kin) (kernelInputType kin) + let arr = kernelInputArray kin + arrtp <- lookupType arr + let slice = fullSlice arrtp [DimFix se | se <- kernelInputIndices kin] + letBind (Pattern [] [pe]) $ BasicOp $ Index arr slice + +-- | Manually calculate the values for the ispace identifiers, when the +-- 'SpaceStructure' won't do. ispace is the dimensions of the overlaying maps. +-- +-- If the input is @i [(a_vn, a), (b_vn, b), (c_vn, c)]@ then @i@ should hit all +-- the values [0,a*b*c). We can calculate the indexes for the other dimensions: +-- +-- > c_vn = i % c +-- > b_vn = (i/c) % b +-- > a_vn = ((i/c)/b) % a +addManualIspaceCalcStms :: (MonadBinder m, Lore m ~ InKernel) => + SubExp + -> [(VName, SubExp)] + -> m () +addManualIspaceCalcStms outer_index ispace = do + -- TODO: The ispace index is calculated in a bit different way than it + -- would have been done if the ThreadSpace was used. However, this + -- works. Maybe ask Troels if doing it the other way has some benefit? + let calc_ispace_index prev_val (vn,size) = do + let pe = PatElem vn (Prim $ IntType Int32) + letBind_ (Pattern [] [pe]) $ BasicOp $ BinOp (SRem Int32) prev_val size + letSubExp "tmp_val" $ BasicOp $ BinOp (SQuot Int32) prev_val size + foldM_ calc_ispace_index outer_index (reverse ispace) + +addFlagToLambda :: (MonadBinder m, Lore m ~ Kernels) => + [SubExp] -> Lambda InKernel -> m (Lambda InKernel) +addFlagToLambda nes lam = do + let num_accs = length nes + x_flag <- newVName "x_flag" + y_flag <- newVName "y_flag" + let x_flag_param = Param x_flag $ Prim Bool + y_flag_param = Param y_flag $ Prim Bool + (x_params, y_params) = splitAt num_accs $ lambdaParams lam + params = [x_flag_param] ++ x_params ++ [y_flag_param] ++ y_params + + body <- runBodyBinder $ localScope (scopeOfLParams params) $ do + new_flag <- letSubExp "new_flag" $ + BasicOp $ BinOp LogOr (Var x_flag) (Var y_flag) + lhs <- fmap (map Var) $ letTupExp "seg_lhs" $ If (Var y_flag) + (resultBody nes) + (resultBody $ map (Var . paramName) x_params) $ + ifCommon $ map paramType x_params + let rhs = map (Var . paramName) y_params + + lam' <- renameLambda lam -- avoid shadowing + res <- eLambda lam' $ map eSubExp $ lhs ++ rhs + + return $ resultBody $ new_flag : res + + return Lambda { lambdaParams = params + , lambdaBody = body + , lambdaReturnType = Prim Bool : lambdaReturnType lam + } + +regularSegmentedScan :: (MonadBinder m, Lore m ~ Kernels) => + SubExp + -> Pattern Kernels + -> SubExp + -> Lambda InKernel + -> Lambda InKernel + -> [(VName, SubExp)] -> [KernelInput] + -> [SubExp] -> [VName] + -> m () +regularSegmentedScan segment_size pat w lam map_lam ispace inps nes arrs = do + flags_i <- newVName "flags_i" + + unused_flag_array <- newVName "unused_flag_array" + flags_body <- + runBodyBinder $ localScope (M.singleton flags_i $ IndexInfo Int32) $ do + segment_index <- letSubExp "segment_index" $ + BasicOp $ BinOp (SRem Int32) (Var flags_i) segment_size + start_of_segment <- letSubExp "start_of_segment" $ + BasicOp $ CmpOp (CmpEq int32) segment_index zero + let flag = start_of_segment + return $ resultBody [flag] + (mapk_bnds, mapk) <- mapKernelFromBody w (FlatThreadSpace [(flags_i, w)]) [] [Prim Bool] flags_body + addStms mapk_bnds + flags <- letExp "flags" $ Op mapk + + lam' <- addFlagToLambda nes lam + + flag_p <- newParam "flag" $ Prim Bool + let map_lam' = map_lam { lambdaParams = flag_p : lambdaParams map_lam + , lambdaBody = (lambdaBody map_lam) + { bodyResult = Var (paramName flag_p) : bodyResult (lambdaBody map_lam) } + , lambdaReturnType = Prim Bool : lambdaReturnType map_lam + } + + let pat' = pat { patternValueElements = PatElem unused_flag_array + (arrayOf (Prim Bool) (Shape [w]) NoUniqueness) : + patternValueElements pat + } + void $ blockedScan pat' w (lam', false:nes) (Commutative, nilFn, mempty) map_lam' segment_size ispace inps (flags:arrs) + where zero = constant (0 :: Int32) + false = constant False diff --git a/src/Futhark/Pass/FirstOrderTransform.hs b/src/Futhark/Pass/FirstOrderTransform.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/FirstOrderTransform.hs @@ -0,0 +1,15 @@ +module Futhark.Pass.FirstOrderTransform + ( firstOrderTransform + ) + where + +import Futhark.Transform.FirstOrderTransform (transformFunDef) +import Futhark.Representation.SOACS (SOACS) +import Futhark.Representation.Kernels (Kernels) +import Futhark.Pass + +firstOrderTransform :: Pass SOACS Kernels +firstOrderTransform = Pass + "first order transform" + "Transform all second-order array combinators to for-loops." $ + intraproceduralTransformation transformFunDef diff --git a/src/Futhark/Pass/KernelBabysitting.hs b/src/Futhark/Pass/KernelBabysitting.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/KernelBabysitting.hs @@ -0,0 +1,429 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | Do various kernel optimisations - mostly related to coalescing. +module Futhark.Pass.KernelBabysitting + ( babysitKernels + , nonlinearInMemory + ) + where + +import Control.Arrow (first) +import Control.Monad.State.Strict +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Foldable +import Data.List +import Data.Maybe +import Data.Semigroup ((<>)) + +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import Futhark.Representation.Kernels + hiding (Prog, Body, Stm, Pattern, PatElem, + BasicOp, Exp, Lambda, FunDef, FParam, LParam, RetType) +import Futhark.Tools +import Futhark.Pass +import Futhark.Util + +babysitKernels :: Pass Kernels Kernels +babysitKernels = Pass "babysit kernels" + "Transpose kernel input arrays for better performance." $ + intraproceduralTransformation transformFunDef + +transformFunDef :: MonadFreshNames m => FunDef Kernels -> m (FunDef Kernels) +transformFunDef fundec = do + (body', _) <- modifyNameSource $ runState (runBinderT m M.empty) + return fundec { funDefBody = body' } + where m = inScopeOf fundec $ + transformBody mempty $ funDefBody fundec + +type BabysitM = Binder Kernels + +transformBody :: ExpMap -> Body Kernels -> BabysitM (Body Kernels) +transformBody expmap (Body () bnds res) = insertStmsM $ do + foldM_ transformStm expmap bnds + return $ resultBody res + +-- | Map from variable names to defining expression. We use this to +-- hackily determine whether something is transposed or otherwise +-- 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) + +nonlinearInMemory :: VName -> ExpMap -> Maybe (Maybe [Int]) +nonlinearInMemory name m = + case M.lookup name m of + Just (Let _ _ (BasicOp (Rearrange perm _))) -> Just $ Just $ rearrangeInverse perm + Just (Let _ _ (BasicOp (Reshape _ arr))) -> nonlinearInMemory arr m + Just (Let _ _ (BasicOp (Manifest perm _))) -> Just $ Just perm + Just (Let pat _ (Op (Kernel _ _ ts _))) -> + nonlinear =<< find ((==name) . patElemName . fst) + (zip (patternElements pat) ts) + _ -> Nothing + where nonlinear (pe, t) + | inner_r <- arrayRank t, inner_r > 0 = do + let outer_r = arrayRank (patElemType pe) - inner_r + return $ Just $ rearrangeInverse $ [inner_r..inner_r+outer_r-1] ++ [0..inner_r-1] + | otherwise = Nothing + +transformStm :: ExpMap -> Stm Kernels -> BabysitM ExpMap + +transformStm expmap (Let pat aux (Op (Kernel desc space ts kbody))) = do + -- Go spelunking for accesses to arrays that are defined outside the + -- kernel body and where the indices are kernel thread indices. + scope <- askScope + let thread_gids = map fst $ spaceDimensions space + thread_local = S.fromList $ spaceGlobalId space : spaceLocalId space : thread_gids + + kbody'' <- evalStateT (traverseKernelBodyArrayIndexes + thread_local + (castScope scope <> scopeOfKernelSpace space) + (ensureCoalescedAccess expmap (spaceDimensions space) num_threads) + kbody) + mempty + + let bnd' = Let pat aux $ Op $ Kernel desc space ts kbody'' + addStm bnd' + return $ M.fromList [ (name, bnd') | name <- patternNames pat ] <> expmap + where num_threads = spaceNumThreads space + +transformStm expmap (Let pat aux e) = do + e' <- mapExpM (transform expmap) e + let bnd' = Let pat aux e' + addStm bnd' + return $ M.fromList [ (name, bnd') | name <- patternNames pat ] <> expmap + +transform :: ExpMap -> Mapper Kernels Kernels BabysitM +transform expmap = + identityMapper { mapOnBody = \scope -> localScope scope . transformBody expmap } + +type ArrayIndexTransform m = + (VName -> Bool) -> -- thread local? + (SubExp -> Maybe SubExp) -> -- split substitution? + Scope InKernel -> -- type environment + VName -> Slice SubExp -> m (Maybe (VName, Slice SubExp)) + +traverseKernelBodyArrayIndexes :: (Applicative f, Monad f) => + Names + -> Scope InKernel + -> ArrayIndexTransform f + -> KernelBody InKernel + -> f (KernelBody InKernel) +traverseKernelBodyArrayIndexes thread_variant outer_scope f (KernelBody () kstms kres) = + KernelBody () . stmsFromList <$> + mapM (onStm (varianceInStms mempty kstms, + mkSizeSubsts kstms, + outer_scope)) (stmsToList kstms) <*> + pure kres + where onLambda (variance, szsubst, scope) lam = + (\body' -> lam { lambdaBody = body' }) <$> + onBody (variance, szsubst, scope') (lambdaBody lam) + where scope' = scope <> scopeOfLParams (lambdaParams lam) + + onStreamLambda (variance, szsubst, scope) lam = + (\body' -> lam { groupStreamLambdaBody = body' }) <$> + onBody (variance, szsubst, scope') (groupStreamLambdaBody lam) + where scope' = scope <> scopeOf lam + + onBody (variance, szsubst, scope) (Body battr stms bres) = do + stms' <- stmsFromList <$> mapM (onStm (variance', szsubst', scope')) (stmsToList stms) + Body battr stms' <$> pure bres + where variance' = varianceInStms variance stms + szsubst' = mkSizeSubsts stms <> szsubst + scope' = scope <> scopeOf stms + + onStm (variance, szsubst, _) (Let pat attr (BasicOp (Index arr is))) = + Let pat attr . oldOrNew <$> f isThreadLocal sizeSubst outer_scope arr is + where oldOrNew Nothing = + BasicOp $ Index arr is + oldOrNew (Just (arr', is')) = + BasicOp $ Index arr' is' + + isThreadLocal v = + not $ S.null $ + thread_variant `S.intersection` + M.findWithDefault (S.singleton v) v variance + + sizeSubst (Constant v) = Just $ Constant v + sizeSubst (Var v) + | v `M.member` outer_scope = Just $ Var v + | Just v' <- M.lookup v szsubst = sizeSubst v' + | otherwise = Nothing + + onStm (variance, szsubst, scope) (Let pat attr e) = + Let pat attr <$> mapExpM (mapper (variance, szsubst, scope)) e + + mapper ctx = identityMapper { mapOnBody = const (onBody ctx) + , mapOnOp = onOp ctx + } + + onOp ctx (GroupReduce w lam input) = + GroupReduce w <$> onLambda ctx lam <*> pure input + onOp ctx (GroupStream w maxchunk lam accs arrs) = + GroupStream w maxchunk <$> onStreamLambda ctx lam <*> pure accs <*> pure arrs + onOp _ stm = pure stm + + mkSizeSubsts = fold . fmap mkStmSizeSubst + where mkStmSizeSubst (Let (Pattern [] [pe]) _ (Op (SplitSpace _ _ _ elems_per_i))) = + M.singleton (patElemName pe) elems_per_i + mkStmSizeSubst _ = mempty + +-- Not a hashmap, as SubExp is not hashable. +type Replacements = M.Map (VName, Slice SubExp) VName + +ensureCoalescedAccess :: MonadBinder m => + ExpMap + -> [(VName,SubExp)] + -> SubExp + -> ArrayIndexTransform (StateT Replacements m) +ensureCoalescedAccess expmap thread_space num_threads isThreadLocal sizeSubst outer_scope arr slice = do + seen <- gets $ M.lookup (arr, slice) + + case (seen, isThreadLocal arr, typeOf <$> M.lookup arr outer_scope) of + -- Already took care of this case elsewhere. + (Just arr', _, _) -> + pure $ Just (arr', slice) + + (Nothing, False, Just t) + -- We are fully indexing the array with thread IDs, but the + -- indices are in a permuted order. + | Just is <- sliceIndices slice, + length is == arrayRank t, + Just is' <- coalescedIndexes (map Var thread_gids) is, + Just perm <- is' `isPermutationOf` is -> + replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr) + + -- Check whether the access is already coalesced because of a + -- previous rearrange being applied to the current array: + -- 1. get the permutation of the source-array rearrange + -- 2. apply it to the slice + -- 3. check that the innermost index is actually the gid + -- of the innermost kernel dimension. + -- If so, the access is already coalesced, nothing to do! + -- (Cosmin's Heuristic.) + | Just (Let _ _ (BasicOp (Rearrange perm _))) <- M.lookup arr expmap, + ---- Just (Just perm) <- nonlinearInMemory arr expmap, + not $ null perm, + length slice >= length perm, + slice' <- map (\i -> slice !! i) perm, + DimFix inner_ind <- last slice', + not $ null thread_gids, + inner_ind == (Var $ last thread_gids) -> + return Nothing + + -- We are not fully indexing an array, but the remaining slice + -- is invariant to the innermost-kernel dimension. We assume + -- the remaining slice will be sequentially streamed, hence + -- tiling will be applied later and will solve coalescing. + -- Hence nothing to do at this point. (Cosmin's Heuristic.) + | (is, rem_slice) <- splitSlice slice, + not $ null rem_slice, + allDimAreSlice rem_slice, + Nothing <- M.lookup arr expmap, + not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, + is /= map Var (take (length is) thread_gids) || length is == length thread_gids, + not (null thread_gids || null is), + not ( S.member (last thread_gids) (S.union (freeIn is) (freeIn rem_slice)) ) -> + return Nothing + + -- We are not fully indexing the array, and the indices are not + -- a proper prefix of the thread indices, and some indices are + -- thread local, so we assume (HEURISTIC!) that the remaining + -- dimensions will be traversed sequentially. + | (is, rem_slice) <- splitSlice slice, + not $ null rem_slice, + not $ tooSmallSlice (primByteSize (elemType t)) rem_slice, + is /= map Var (take (length is) thread_gids) || length is == length thread_gids, + any isThreadLocal (S.toList $ freeIn is) -> do + let perm = coalescingPermutation (length is) $ arrayRank t + replace =<< lift (rearrangeInput (nonlinearInMemory arr expmap) perm arr) + + -- We are taking a slice of the array with a unit stride. We + -- assume that the slice will be traversed sequentially. + -- + -- We will really want to treat the sliced dimension like two + -- dimensions so we can transpose them. This may require + -- padding. + | (is, rem_slice) <- splitSlice slice, + and $ zipWith (==) is $ map Var thread_gids, + DimSlice offset len (Constant stride):_ <- rem_slice, + isThreadLocalSubExp offset, + Just {} <- sizeSubst len, + oneIsh stride -> do + let num_chunks = if null is + then primExpFromSubExp int32 num_threads + else coerceIntPrimExp Int32 $ + product $ map (primExpFromSubExp int32) $ + drop (length is) thread_gdims + replace =<< lift (rearrangeSlice (length is) (arraySize (length is) t) num_chunks arr) + + -- Everything is fine... assuming that the array is in row-major + -- order! Make sure that is the case. + | Just{} <- nonlinearInMemory arr expmap -> + case sliceIndices slice of + Just is | Just _ <- coalescedIndexes (map Var thread_gids) is -> + replace =<< lift (rowMajorArray arr) + | otherwise -> + return Nothing + _ -> replace =<< lift (rowMajorArray arr) + + _ -> return Nothing + + where (thread_gids, thread_gdims) = unzip thread_space + + replace arr' = do + modify $ M.insert (arr, slice) arr' + return $ Just (arr', slice) + + isThreadLocalSubExp (Var v) = isThreadLocal v + isThreadLocalSubExp Constant{} = False + +-- Heuristic for avoiding rearranging too small arrays. +tooSmallSlice :: Int32 -> Slice SubExp -> Bool +tooSmallSlice bs = fst . foldl comb (True,bs) . sliceDims + where comb (True, x) (Constant (IntValue (Int32Value d))) = (d*x < 4, d*x) + comb (_, x) _ = (False, x) + +splitSlice :: Slice SubExp -> ([SubExp], Slice SubExp) +splitSlice [] = ([], []) +splitSlice (DimFix i:is) = first (i:) $ splitSlice is +splitSlice is = ([], is) + +allDimAreSlice :: Slice SubExp -> Bool +allDimAreSlice [] = True +allDimAreSlice (DimFix _:_) = False +allDimAreSlice (_:is) = allDimAreSlice is + +-- Try to move thread indexes into their proper position. +coalescedIndexes :: [SubExp] -> [SubExp] -> Maybe [SubExp] +coalescedIndexes tgids is + -- Do Nothing if: + -- 1. the innermost index is the innermost thread id + -- (because access is already coalesced) + -- 2. any of the indices is a constant, i.e., kernel free variable + -- (because it would transpose a bigger array then needed -- big overhead). + | any isCt is = + Nothing + | num_is > 0 && not (null tgids) && last is == last tgids = + Just is + -- Otherwise try fix coalescing + | otherwise = + Just $ reverse $ foldl move (reverse is) $ zip [0..] (reverse tgids) + where num_is = length is + + move is_rev (i, tgid) + -- If tgid is in is_rev anywhere but at position i, and + -- position i exists, we move it to position i instead. + | Just j <- elemIndex tgid is_rev, i /= j, i < num_is = + swap i j is_rev + | otherwise = + is_rev + + swap i j l + | Just ix <- maybeNth i l, + Just jx <- maybeNth j l = + update i jx $ update j ix l + | otherwise = + error $ "coalescedIndexes swap: invalid indices" ++ show (i, j, l) + + update 0 x (_:ys) = x : ys + update i x (y:ys) = y : update (i-1) x ys + update _ _ [] = error "coalescedIndexes: update" + + isCt :: SubExp -> Bool + isCt (Constant _) = True + isCt (Var _) = False + +coalescingPermutation :: Int -> Int -> [Int] +coalescingPermutation num_is rank = + [num_is..rank-1] ++ [0..num_is-1] + +rearrangeInput :: MonadBinder m => + Maybe (Maybe [Int]) -> [Int] -> VName -> m VName +rearrangeInput (Just (Just current_perm)) perm arr + | current_perm == perm = return arr -- Already has desired representation. + +rearrangeInput Nothing perm arr + | sort perm == perm = return arr -- We don't know the current + -- representation, but the indexing + -- is linear, so let's hope the + -- array is too. +rearrangeInput (Just Just{}) perm arr + | sort perm == perm = rowMajorArray arr -- We just want a row-major array, no tricks. +rearrangeInput manifest perm arr = do + -- We may first manifest the array to ensure that it is flat in + -- memory. This is sometimes unnecessary, in which case the copy + -- will hopefully be removed by the simplifier. + manifested <- if isJust manifest then rowMajorArray arr else return arr + letExp (baseString arr ++ "_coalesced") $ + BasicOp $ Manifest perm manifested + +rowMajorArray :: MonadBinder m => + VName -> m VName +rowMajorArray arr = do + rank <- arrayRank <$> lookupType arr + letExp (baseString arr ++ "_rowmajor") $ BasicOp $ Manifest [0..rank-1] arr + +rearrangeSlice :: MonadBinder m => + Int -> SubExp -> PrimExp VName -> VName + -> m VName +rearrangeSlice d w num_chunks arr = do + num_chunks' <- letSubExp "num_chunks" =<< toExp num_chunks + + (w_padded, padding) <- paddedScanReduceInput w num_chunks' + + per_chunk <- letSubExp "per_chunk" $ BasicOp $ BinOp (SQuot Int32) w_padded num_chunks' + arr_t <- lookupType arr + arr_padded <- padArray w_padded padding arr_t + rearrange num_chunks' w_padded per_chunk (baseString arr) arr_padded arr_t + + where padArray w_padded padding arr_t = do + let arr_shape = arrayShape arr_t + padding_shape = setDim d arr_shape padding + arr_padding <- + letExp (baseString arr <> "_padding") $ + BasicOp $ Scratch (elemType arr_t) (shapeDims padding_shape) + letExp (baseString arr <> "_padded") $ + BasicOp $ Concat d arr [arr_padding] w_padded + + rearrange num_chunks' w_padded per_chunk arr_name arr_padded arr_t = do + let arr_dims = arrayDims arr_t + pre_dims = take d arr_dims + post_dims = drop (d+1) arr_dims + extradim_shape = Shape $ pre_dims ++ [num_chunks', per_chunk] ++ post_dims + tr_perm = [0..d-1] ++ map (+d) ([1] ++ [2..shapeRank extradim_shape-1-d] ++ [0]) + arr_extradim <- + letExp (arr_name <> "_extradim") $ + BasicOp $ Reshape (map DimNew $ shapeDims extradim_shape) arr_padded + arr_extradim_tr <- + letExp (arr_name <> "_extradim_tr") $ + BasicOp $ Manifest tr_perm arr_extradim + arr_inv_tr <- letExp (arr_name <> "_inv_tr") $ + BasicOp $ Reshape (map DimCoercion pre_dims ++ map DimNew (w_padded : post_dims)) + arr_extradim_tr + letExp (arr_name <> "_inv_tr_init") =<< + eSliceArray d arr_inv_tr (eSubExp $ constant (0::Int32)) (eSubExp w) + +paddedScanReduceInput :: MonadBinder m => + SubExp -> SubExp + -> m (SubExp, SubExp) +paddedScanReduceInput w stride = do + w_padded <- letSubExp "padded_size" =<< + eRoundToMultipleOf Int32 (eSubExp w) (eSubExp stride) + padding <- letSubExp "padding" $ BasicOp $ BinOp (Sub Int32) w_padded w + return (w_padded, padding) + +--- Computing variance. + +type VarianceTable = M.Map VName Names + +varianceInStms :: VarianceTable -> Stms InKernel -> VarianceTable +varianceInStms t = foldl varianceInStm t . stmsToList + +varianceInStm :: VarianceTable -> Stm InKernel -> VarianceTable +varianceInStm variance bnd = + foldl' add variance $ patternNames $ stmPattern bnd + where add variance' v = M.insert v binding_variance variance' + look variance' v = S.insert v $ M.findWithDefault mempty v variance' + binding_variance = mconcat $ map (look variance) $ S.toList (freeInStm bnd) diff --git a/src/Futhark/Pass/ResolveAssertions.hs b/src/Futhark/Pass/ResolveAssertions.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/ResolveAssertions.hs @@ -0,0 +1,55 @@ +-- | Go through the program and use algebraic simplification and range +-- analysis to try to figure out which assertions are statically true. +-- +-- Currently implemented by running the simplifier with a special rule +-- that is too expensive to run all the time. + +module Futhark.Pass.ResolveAssertions + ( resolveAssertions + ) + where + +import Data.Maybe +import Data.Monoid + +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Optimise.Simplify.Rule +import qualified Futhark.Analysis.AlgSimplify as AS +import qualified Futhark.Analysis.ScalExp as SE +import Futhark.Analysis.PrimExp.Convert +import Futhark.Representation.AST.Syntax +import Futhark.Construct +import Futhark.Pass +import Futhark.Representation.SOACS (SOACS) +import qualified Futhark.Representation.SOACS.Simplify as Simplify +import qualified Futhark.Optimise.Simplify as Simplify +import Futhark.Optimise.Simplify.Rules + +import Prelude + +-- | The assertion-resolver pass. +resolveAssertions :: Pass SOACS SOACS +resolveAssertions = Pass + "resolve assertions" + "Try to statically resolve bounds checks and similar." $ + Simplify.simplifyProg Simplify.simpleSOACS rulebook Simplify.noExtraHoistBlockers + where rulebook = standardRules <> ruleBook [ RuleBasicOp simplifyScalExp ] [] + +simplifyScalExp :: BinderOps lore => TopDownRuleBasicOp lore +simplifyScalExp vtable pat _ e = do + res <- SE.toScalExp (`ST.lookupScalExp` vtable) $ BasicOp e + case res of + -- If the sufficient condition is 'True', then it statically succeeds. + Just se + | SE.scalExpType se == Bool, + isNothing $ valOrVar se, + SE.scalExpSize se < size_bound, + Just se' <- valOrVar $ AS.simplify se ranges -> + letBind_ pat $ BasicOp $ SubExp se' + _ -> cannotSimplify + where ranges = ST.rangesRep vtable + size_bound = 10 -- don't touch scalexps bigger than this. + + valOrVar (SE.Val v) = Just $ Constant v + valOrVar (SE.Id v _) = Just $ Var v + valOrVar _ = Nothing diff --git a/src/Futhark/Pass/Simplify.hs b/src/Futhark/Pass/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Simplify.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE FlexibleContexts #-} +module Futhark.Pass.Simplify + ( simplify + , simplifySOACS + , simplifyKernels + , simplifyExplicitMemory + ) + where + +import qualified Futhark.Representation.SOACS as R +import qualified Futhark.Representation.SOACS.Simplify as R +import qualified Futhark.Representation.Kernels as R +import qualified Futhark.Representation.Kernels.Simplify as R +import qualified Futhark.Representation.ExplicitMemory as R +import qualified Futhark.Representation.ExplicitMemory.Simplify as R + +import Futhark.Pass +import Futhark.Representation.AST.Syntax + +simplify :: (Prog lore -> PassM (Prog lore)) + -> Pass lore lore +simplify = Pass "simplify" "Perform simple enabling optimisations." + +simplifySOACS :: Pass R.SOACS R.SOACS +simplifySOACS = simplify R.simplifySOACS + +simplifyKernels :: Pass R.Kernels R.Kernels +simplifyKernels = simplify R.simplifyKernels + +simplifyExplicitMemory :: Pass R.ExplicitMemory R.ExplicitMemory +simplifyExplicitMemory = simplify R.simplifyExplicitMemory diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Passes.hs @@ -0,0 +1,157 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +-- | Optimisation pipelines. +module Futhark.Passes + ( standardPipeline + , sequentialPipeline + , kernelsPipeline + , sequentialCpuPipeline + , gpuPipeline + ) +where + +import Control.Category ((>>>)) + +import Futhark.Optimise.CSE +import Futhark.Optimise.Fusion +import Futhark.Optimise.InPlaceLowering +import Futhark.Optimise.InliningDeadFun +import Futhark.Optimise.TileLoops +import Futhark.Optimise.DoubleBuffer +import Futhark.Optimise.Unstream +import Futhark.Optimise.MemoryBlockMerging +import Futhark.Pass.ExpandAllocations +import Futhark.Pass.ExplicitAllocations +import Futhark.Pass.ExtractKernels +import Futhark.Pass.FirstOrderTransform +import Futhark.Pass.KernelBabysitting +import Futhark.Pass.ResolveAssertions +import Futhark.Pass.Simplify +import Futhark.Pipeline +import Futhark.Representation.ExplicitMemory (ExplicitMemory) +import Futhark.Representation.Kernels (Kernels) +import Futhark.Representation.SOACS (SOACS) +import Futhark.Util + +standardPipeline :: Pipeline SOACS SOACS +standardPipeline = + passes [ simplifySOACS + , inlineAndRemoveDeadFunctions + , performCSE True + , simplifySOACS + -- We run fusion twice + , fuseSOACs + , performCSE True + , simplifySOACS + , fuseSOACs + , performCSE True + , simplifySOACS + , resolveAssertions + , removeDeadFunctions + ] + +-- Do we use in-place lowering? Currently enabled by default. Disable by +-- setting the environment variable IN_PLACE_LOWERING=0. +usesInPlaceLowering :: Bool +usesInPlaceLowering = + isEnvVarSet "IN_PLACE_LOWERING" True + +inPlaceLoweringMaybe :: Pipeline Kernels Kernels +inPlaceLoweringMaybe = + if usesInPlaceLowering + then onePass inPlaceLowering + else passes [] + +-- Do we use the coalescing part of memory block merging? Currently disabled by +-- default. Enable by setting the environment variable +-- MEMORY_BLOCK_MERGING_COALESCING=1. +usesMemoryBlockMergingCoalescing :: Bool +usesMemoryBlockMergingCoalescing = + isEnvVarSet "MEMORY_BLOCK_MERGING_COALESCING" False + +memoryBlockMergingCoalescingMaybe :: Pipeline ExplicitMemory ExplicitMemory +memoryBlockMergingCoalescingMaybe = + passes $ if usesMemoryBlockMergingCoalescing + then [ memoryBlockMergingCoalescing + , simplifyExplicitMemory + ] + else [] + +memoryBlockMergingCoalescingMaybeCPU :: Pipeline ExplicitMemory ExplicitMemory +memoryBlockMergingCoalescingMaybeCPU = memoryBlockMergingCoalescingMaybe + +memoryBlockMergingCoalescingMaybeGPU :: Pipeline ExplicitMemory ExplicitMemory +memoryBlockMergingCoalescingMaybeGPU = memoryBlockMergingCoalescingMaybe + +-- Do we use the reuse part of memory block merging? Currently disabled by +-- default. Enable by setting the environment variable +-- MEMORY_BLOCK_MERGING_REUSE=1. +usesMemoryBlockMergingReuse :: Bool +usesMemoryBlockMergingReuse = + isEnvVarSet "MEMORY_BLOCK_MERGING_REUSE" False + +memoryBlockMergingReuseMaybe :: Pipeline ExplicitMemory ExplicitMemory +memoryBlockMergingReuseMaybe = + passes $ if usesMemoryBlockMergingReuse + then [ memoryBlockMergingReuse + , simplifyExplicitMemory + ] + else [] + +memoryBlockMergingReuseMaybeCPU :: Pipeline ExplicitMemory ExplicitMemory +memoryBlockMergingReuseMaybeCPU = memoryBlockMergingReuseMaybe + +memoryBlockMergingReuseMaybeGPU :: Pipeline ExplicitMemory ExplicitMemory +memoryBlockMergingReuseMaybeGPU = memoryBlockMergingReuseMaybe + + +kernelsPipeline :: Pipeline SOACS Kernels +kernelsPipeline = + standardPipeline >>> + onePass extractKernels >>> + passes [ simplifyKernels + , babysitKernels + , simplifyKernels + , tileLoops + , unstream + , simplifyKernels + , performCSE True + , simplifyKernels + ] >>> + inPlaceLoweringMaybe + +sequentialPipeline :: Pipeline SOACS Kernels +sequentialPipeline = + standardPipeline >>> + onePass firstOrderTransform >>> + passes [ simplifyKernels + ] >>> + inPlaceLoweringMaybe + +sequentialCpuPipeline :: Pipeline SOACS ExplicitMemory +sequentialCpuPipeline = + sequentialPipeline >>> + onePass explicitAllocations >>> + passes [ simplifyExplicitMemory + , performCSE False + , simplifyExplicitMemory + , doubleBuffer + , simplifyExplicitMemory + ] >>> + memoryBlockMergingCoalescingMaybeCPU >>> + memoryBlockMergingReuseMaybeCPU + +gpuPipeline :: Pipeline SOACS ExplicitMemory +gpuPipeline = + kernelsPipeline >>> + onePass explicitAllocations >>> + passes [ simplifyExplicitMemory + , performCSE False + , simplifyExplicitMemory + , doubleBuffer + , simplifyExplicitMemory + , expandAllocations + , simplifyExplicitMemory + ] >>> + memoryBlockMergingCoalescingMaybeGPU >>> + memoryBlockMergingReuseMaybeGPU diff --git a/src/Futhark/Pipeline.hs b/src/Futhark/Pipeline.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pipeline.hs @@ -0,0 +1,154 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, OverloadedStrings #-} +module Futhark.Pipeline + ( Pipeline + , PipelineConfig (..) + , Action (..) + + , FutharkM + , runFutharkM + , Verbosity(..) + + , internalErrorS + + , module Futhark.Error + + , onePass + , passes + , runPasses + , runPipeline + ) + where + +import Control.Category +import Control.Monad +import Control.Monad.Writer.Strict hiding (pass) +import Control.Monad.Except +import Control.Monad.State +import Control.Monad.Reader +import qualified Data.Text as T +import qualified Data.Text.IO as T +import Data.Time.Clock +import System.IO +import Text.Printf + +import Prelude hiding (id, (.)) + +import Futhark.Error +import Futhark.Representation.AST (Prog, PrettyLore) +import Futhark.TypeCheck +import Futhark.Pass +import Futhark.Util.Log +import Futhark.Util.Pretty (Pretty, prettyText) +import Futhark.MonadFreshNames + +-- | If Verbose, print log messages to standard error. If +-- VeryVerbose, also print logs from individual passes. +data Verbosity = NotVerbose | Verbose | VeryVerbose deriving (Eq, Ord) + +newtype FutharkEnv = FutharkEnv { futharkVerbose :: Verbosity } + +data FutharkState = FutharkState { futharkPrevLog :: UTCTime + , futharkNameSource :: VNameSource } + +newtype FutharkM a = FutharkM (ExceptT CompilerError (StateT FutharkState (ReaderT FutharkEnv IO)) a) + deriving (Applicative, Functor, Monad, + MonadError CompilerError, + MonadState FutharkState, + MonadReader FutharkEnv, + MonadIO) + +instance MonadFreshNames FutharkM where + getNameSource = gets futharkNameSource + putNameSource src = modify $ \s -> s { futharkNameSource = src } + +instance MonadLogger FutharkM where + addLog = mapM_ perLine . T.lines . toText + where perLine msg = do + verb <- asks $ (>=Verbose) . futharkVerbose + prev <- gets futharkPrevLog + now <- liftIO getCurrentTime + let delta :: Double + delta = fromRational $ toRational (now `diffUTCTime` prev) + prefix = printf "[ +%.6f] " delta + modify $ \s -> s { futharkPrevLog = now } + when verb $ liftIO $ T.hPutStrLn stderr $ T.pack prefix <> msg + +runFutharkM :: FutharkM a -> Verbosity -> IO (Either CompilerError a) +runFutharkM (FutharkM m) verbose = do + s <- FutharkState <$> getCurrentTime <*> pure blankNameSource + runReaderT (evalStateT (runExceptT m) s) newEnv + where newEnv = FutharkEnv verbose + +internalErrorS :: Pretty t => String -> t -> FutharkM a +internalErrorS s p = throwError $ InternalError (T.pack s) (prettyText p) CompilerBug + +data Action lore = + Action { actionName :: String + , actionDescription :: String + , actionProcedure :: Prog lore -> FutharkM () + } + +data PipelineConfig = + PipelineConfig { pipelineVerbose :: Bool + , pipelineValidate :: Bool + } + +newtype Pipeline fromlore tolore = + Pipeline { unPipeline :: PipelineConfig -> Prog fromlore -> FutharkM (Prog tolore) } + +instance Category Pipeline where + id = Pipeline $ const return + p2 . p1 = Pipeline perform + where perform cfg prog = + runPasses p2 cfg =<< runPasses p1 cfg prog + +runPasses :: Pipeline fromlore tolore + -> PipelineConfig + -> Prog fromlore + -> FutharkM (Prog tolore) +runPasses = unPipeline + +runPipeline :: Pipeline fromlore tolore + -> PipelineConfig + -> Prog fromlore + -> Action tolore + -> FutharkM () +runPipeline p cfg prog a = do + prog' <- runPasses p cfg prog + when (pipelineVerbose cfg) $ logMsg $ + "Running action " <> T.pack (actionName a) + actionProcedure a prog' + +onePass :: (Checkable fromlore, Checkable tolore) => + Pass fromlore tolore -> Pipeline fromlore tolore +onePass pass = Pipeline perform + where perform cfg prog = do + when (pipelineVerbose cfg) $ logMsg $ + "Running pass " <> T.pack (passName pass) + prog' <- runPass pass prog + when (pipelineValidate cfg) $ + case checkProg prog' of + Left err -> validationError pass prog' $ show err + Right () -> return () + return prog' + +passes :: Checkable lore => + [Pass lore lore] -> Pipeline lore lore +passes = foldl (>>>) id . map onePass + +validationError :: PrettyLore tolore => + Pass fromlore tolore -> Prog tolore -> String -> FutharkM a +validationError pass prog err = + throwError $ InternalError msg (prettyText prog) CompilerBug + where msg = "Type error after pass '" <> T.pack (passName pass) <> "':\n" <> T.pack err + +runPass :: PrettyLore fromlore => + Pass fromlore tolore + -> Prog fromlore + -> FutharkM (Prog tolore) +runPass pass prog = do + (res, logged) <- runPassM (passFunction pass prog) + verb <- asks $ (>=VeryVerbose) . futharkVerbose + when verb $ addLog logged + case res of Left err -> internalError err $ prettyText prog + Right x -> return x diff --git a/src/Futhark/Pkg/Info.hs b/src/Futhark/Pkg/Info.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pkg/Info.hs @@ -0,0 +1,338 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | Obtaining information about packages over THE INTERNET! +module Futhark.Pkg.Info + ( -- * Package info + PkgInfo(..) + , lookupPkgRev + , pkgInfo + , PkgRevInfo (..) + , GetManifest (getManifest) + , downloadZipball + + -- * Package registry + , PkgRegistry + , MonadPkgRegistry(..) + , lookupPackage + , lookupPackageRev + , lookupNewestRev + ) + where + +import Control.Monad.IO.Class +import Data.Maybe +import Data.IORef +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.ByteString as BS +import qualified Data.Text.Encoding as T +import qualified Data.Semigroup as Sem +import Data.List +import Data.Monoid ((<>)) +import qualified System.FilePath.Posix as Posix +import System.Exit +import System.IO + +import qualified Codec.Archive.Zip as Zip +import Data.Time (UTCTime, UTCTime, defaultTimeLocale, formatTime, getCurrentTime) +import Data.Versions (SemVer(..), semver, prettySemVer) +import System.Process.ByteString (readProcessWithExitCode) +import Network.HTTP.Client hiding (path) +import Network.HTTP.Simple + +import Futhark.Pkg.Types +import Futhark.Util.Log +import Futhark.Util (maybeHead) + +-- | The manifest is stored as a monadic action, because we want to +-- fetch them on-demand. It would be a waste to fetch it information +-- for every version of every package if we only actually need a small +-- subset of them. +newtype GetManifest m = GetManifest { getManifest :: m PkgManifest } + +instance Show (GetManifest m) where + show _ = "#" + +instance Eq (GetManifest m) where + _ == _ = True + +-- | Information about a version of a single package. The version +-- number is stored separately. +data PkgRevInfo m = PkgRevInfo { pkgRevZipballUrl :: T.Text + , pkgRevZipballDir :: FilePath + -- ^ The directory inside the zipball + -- containing the 'lib' directory, in + -- which the package files themselves + -- are stored (Based on the package + -- path). + , pkgRevCommit :: T.Text + -- ^ The commit ID can be used for + -- verification ("freezing"), by + -- storing what it was at the time this + -- version was last selected. + , pkgRevGetManifest :: GetManifest m + , pkgRevTime :: UTCTime + -- ^ Timestamp for when the revision + -- was made (rarely used). + } + deriving (Eq, Show) + +-- | Create memoisation around a 'GetManifest' action to ensure that +-- multiple inspections of the same revisions will not result in +-- potentially expensive network round trips. +memoiseGetManifest :: MonadIO m => GetManifest m -> m (GetManifest m) +memoiseGetManifest (GetManifest m) = do + ref <- liftIO $ newIORef Nothing + return $ GetManifest $ do + v <- liftIO $ readIORef ref + case v of Just v' -> return v' + Nothing -> do + v' <- m + liftIO $ writeIORef ref $ Just v' + return v' + +downloadZipball :: (MonadLogger m, MonadIO m) => + T.Text -> m Zip.Archive +downloadZipball url = do + logMsg $ "Downloading " <> T.unpack url + r <- liftIO $ parseRequest $ T.unpack url + + r' <- liftIO $ httpLBS r + let bad = fail . (("When downloading " <> T.unpack url <> ": ")<>) + case getResponseStatusCode r' of + 200 -> + case Zip.toArchiveOrFail $ getResponseBody r' of + Left e -> bad $ show e + Right a -> return a + x -> bad $ "got HTTP status " ++ show x + +-- | Information about a package. The name of the package is stored +-- separately. +data PkgInfo m = PkgInfo { pkgVersions :: M.Map SemVer (PkgRevInfo m) + , pkgLookupCommit :: Maybe T.Text -> m (PkgRevInfo m) + -- ^ Look up information about a specific + -- commit, or HEAD in case of Nothing. + } + +lookupPkgRev :: SemVer -> PkgInfo m -> Maybe (PkgRevInfo m) +lookupPkgRev v = M.lookup v . pkgVersions + +majorRevOfPkg :: PkgPath -> (PkgPath, [Word]) +majorRevOfPkg p = + case T.splitOn "@" p of + [p', v] | [(v', "")] <- reads $ T.unpack v -> (p', [v']) + _ -> (p, [0, 1]) + +-- | Retrieve information about a package based on its package path. +-- This uses Semantic Import Versioning when interacting with +-- repositories. For example, a package @github.com/user/repo@ will +-- match version 0.* or 1.* tags only, a package +-- @github.com/user/repo/v2@ will match 2.* tags, and so forth.. +pkgInfo :: (MonadIO m, MonadLogger m) => + PkgPath -> m (Either T.Text (PkgInfo m)) +pkgInfo path + | ["github.com", owner, repo] <- T.splitOn "/" path = + let (repo', vs) = majorRevOfPkg repo + in ghPkgInfo owner repo' vs + | "github.com": owner : repo : _ <- T.splitOn "/" path = + return $ Left $ T.intercalate "\n" + [nope, "Do you perhaps mean 'github.com/" <> owner <> "/" <> repo <> "'?"] + | ["gitlab.com", owner, repo] <- T.splitOn "/" path = + let (repo', vs) = majorRevOfPkg repo + in glPkgInfo owner repo' vs + | "gitlab.com": owner : repo : _ <- T.splitOn "/" path = + return $ Left $ T.intercalate "\n" + [nope, "Do you perhaps mean 'gitlab.com/" <> owner <> "/" <> repo <> "'?"] + | otherwise = + return $ Left nope + where nope = "Unable to handle package paths of the form '" <> path <> "'" + +-- For GitHub, we unfortunately cannot use the (otherwise very nice) +-- GitHub web API, because it is rate-limited to 60 requests per hour +-- for non-authenticated users. Instead we fall back to a combination +-- of calling 'git' directly and retrieving things from the GitHub +-- webserver, which is not rate-limited. This approach is also used +-- by other systems (Go most notably), so we should not be stepping on +-- any toes. + +gitCmd :: MonadIO m => [String] -> m BS.ByteString +gitCmd opts = do + (code, out, err) <- liftIO $ readProcessWithExitCode "git" opts mempty + liftIO $ BS.hPutStr stderr err + case code of + ExitFailure 127 -> fail $ "'" <> unwords ("git" : opts) <> "' failed (program not found?)." + ExitFailure _ -> fail $ "'" <> unwords ("git" : opts) <> "' failed." + ExitSuccess -> return out + +-- The GitLab and GitHub interactions are very similar, so we define a +-- couple of generic functions that are used to implement support for +-- both. + +ghglRevGetManifest :: (MonadIO m, MonadLogger m) => + T.Text -> T.Text -> T.Text -> T.Text -> GetManifest m +ghglRevGetManifest url owner repo tag = GetManifest $ do + logMsg $ "Downloading package manifest from " <> url + r <- liftIO $ parseRequest $ T.unpack url + + r' <- liftIO $ httpBS r + let path = T.unpack $ owner <> "/" <> repo <> "@" <> + tag <> "/" <> T.pack futharkPkg + msg = (("When reading " <> path <> ": ")<>) + case getResponseStatusCode r' of + 200 -> + case T.decodeUtf8' $ getResponseBody r' of + Left e -> fail $ msg $ show e + Right s -> + case parsePkgManifest path s of + Left e -> fail $ msg $ errorBundlePretty e + Right pm -> return pm + x -> fail $ msg $ "got HTTP status " ++ show x + +ghglLookupCommit :: (MonadIO m, MonadLogger m) => + T.Text -> T.Text + -> T.Text -> T.Text -> T.Text -> T.Text -> T.Text -> m (PkgRevInfo m) +ghglLookupCommit archive_url manifest_url owner repo d ref hash = do + gd <- memoiseGetManifest $ ghglRevGetManifest manifest_url owner repo ref + let dir = Posix.addTrailingPathSeparator $ T.unpack repo <> "-" <> T.unpack d + time <- liftIO getCurrentTime -- FIXME + return $ PkgRevInfo archive_url dir hash gd time + +ghglPkgInfo :: (MonadIO m, MonadLogger m) => + T.Text -> (T.Text -> T.Text) -> (T.Text -> T.Text) + -> T.Text -> T.Text -> [Word] -> m (Either T.Text (PkgInfo m)) +ghglPkgInfo repo_url mk_archive_url mk_manifest_url owner repo versions = do + logMsg $ "Retrieving list of tags from " <> repo_url + remote_lines <- T.lines . T.decodeUtf8 <$> gitCmd ["ls-remote", T.unpack repo_url] + + head_ref <- maybe (fail $ "Cannot find HEAD ref for " <> T.unpack repo_url) return $ + maybeHead $ mapMaybe isHeadRef remote_lines + let def = fromMaybe head_ref + + rev_info <- M.fromList . catMaybes <$> mapM revInfo remote_lines + + return $ Right $ PkgInfo rev_info $ \r -> + ghglLookupCommit (mk_archive_url (def r)) (mk_manifest_url (def r)) + owner repo (def r) (def r) (def r) + where isHeadRef l + | [hash, "HEAD"] <- T.words l = Just hash + | otherwise = Nothing + + revInfo l + | [hash, ref] <- T.words l, + ["refs", "tags", t] <- T.splitOn "/" ref, + "v" `T.isPrefixOf` t, + Right v <- semver $ T.drop 1 t, + _svMajor v `elem` versions = do + pinfo <- ghglLookupCommit (mk_archive_url t) (mk_manifest_url t) + owner repo (prettySemVer v) t hash + return $ Just (v, pinfo) + | otherwise = return Nothing + +ghPkgInfo :: (MonadIO m, MonadLogger m) => + T.Text -> T.Text -> [Word] -> m (Either T.Text (PkgInfo m)) +ghPkgInfo owner repo versions = + ghglPkgInfo repo_url mk_archive_url mk_manifest_url owner repo versions + where repo_url = "https://github.com/" <> owner <> "/" <> repo + mk_archive_url r = repo_url <> "/archive/" <> r <> ".zip" + mk_manifest_url r = "https://raw.githubusercontent.com/" <> + owner <> "/" <> repo <> "/" <> + r <> "/" <> T.pack futharkPkg + +glPkgInfo :: (MonadIO m, MonadLogger m) => + T.Text -> T.Text -> [Word] -> m (Either T.Text (PkgInfo m)) +glPkgInfo owner repo versions = + ghglPkgInfo repo_url mk_archive_url mk_manifest_url owner repo versions + where base_url = "https://gitlab.com/" <> owner <> "/" <> repo + repo_url = base_url <> ".git" + mk_archive_url r = base_url <> "/-/archive/" <> r <> + "/" <> repo <> "-" <> r <> ".zip" + mk_manifest_url r = base_url <> "/raw/" <> + r <> "/" <> T.pack futharkPkg + +-- | A package registry is a mapping from package paths to information +-- about the package. It is unlikely that any given registry is +-- global; rather small registries are constructed on-demand based on +-- the package paths referenced by the user, and may also be combined +-- monoidically. In essence, the PkgRegistry is just a cache. +newtype PkgRegistry m = PkgRegistry (M.Map PkgPath (PkgInfo m)) + +instance Sem.Semigroup (PkgRegistry m) where + PkgRegistry x <> PkgRegistry y = PkgRegistry $ x <> y + +instance Monoid (PkgRegistry m) where + mempty = PkgRegistry mempty + mappend = (Sem.<>) + +lookupKnownPackage :: PkgPath -> PkgRegistry m -> Maybe (PkgInfo m) +lookupKnownPackage p (PkgRegistry m) = M.lookup p m + +-- | Monads that support a stateful package registry. These are also +-- required to be instances of 'MonadIO' because most package registry +-- operations involve network operations. +class (MonadIO m, MonadLogger m) => MonadPkgRegistry m where + getPkgRegistry :: m (PkgRegistry m) + putPkgRegistry :: PkgRegistry m -> m () + modifyPkgRegistry :: (PkgRegistry m -> PkgRegistry m) -> m () + modifyPkgRegistry f = putPkgRegistry . f =<< getPkgRegistry + +lookupPackage :: MonadPkgRegistry m => + PkgPath -> m (PkgInfo m) +lookupPackage p = do + r@(PkgRegistry m) <- getPkgRegistry + case lookupKnownPackage p r of + Just info -> + return info + Nothing -> do + e <- pkgInfo p + case e of + Left e' -> fail $ T.unpack e' + Right pinfo -> do + putPkgRegistry $ PkgRegistry $ M.insert p pinfo m + return pinfo + +lookupPackageCommit :: MonadPkgRegistry m => + PkgPath -> Maybe T.Text -> m (SemVer, PkgRevInfo m) +lookupPackageCommit p ref = do + pinfo <- lookupPackage p + rev_info <- pkgLookupCommit pinfo ref + let timestamp = T.pack $ formatTime defaultTimeLocale "%Y%m%d%H%M%S" $ + pkgRevTime rev_info + v = commitVersion timestamp $ pkgRevCommit rev_info + pinfo' = pinfo { pkgVersions = M.insert v rev_info $ pkgVersions pinfo } + modifyPkgRegistry $ \(PkgRegistry m) -> + PkgRegistry $ M.insert p pinfo' m + return (v, rev_info) + +-- | Look up information about a specific version of a package. +lookupPackageRev :: MonadPkgRegistry m => + PkgPath -> SemVer -> m (PkgRevInfo m) +lookupPackageRev p v + | Just commit <- isCommitVersion v = + snd <$> lookupPackageCommit p (Just commit) + | otherwise = do + pinfo <- lookupPackage p + case lookupPkgRev v pinfo of + Nothing -> + let versions = case M.keys $ pkgVersions pinfo of + [] -> "Package " <> p <> " has no versions. Invalid package path?" + ks -> "Known versions: " <> + T.concat (intersperse ", " $ map prettySemVer ks) + major | (_, vs) <- majorRevOfPkg p, + _svMajor v `notElem` vs = + "\nFor major version " <> T.pack (show (_svMajor v)) <> + ", use package path " <> p <> "@" <> T.pack (show (_svMajor v)) + | otherwise = mempty + in fail $ T.unpack $ + "package " <> p <> " does not have a version " <> prettySemVer v <> ".\n" <> + versions <> major + Just v' -> return v' + +-- | Find the newest version of a package. +lookupNewestRev :: MonadPkgRegistry m => + PkgPath -> m SemVer +lookupNewestRev p = do + pinfo <- lookupPackage p + case M.keys $ pkgVersions pinfo of + [] -> do + logMsg $ "Package " <> p <> " has no released versions. Using HEAD." + fst <$> lookupPackageCommit p Nothing + v:vs -> return $ foldl' max v vs diff --git a/src/Futhark/Pkg/Solve.hs b/src/Futhark/Pkg/Solve.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pkg/Solve.hs @@ -0,0 +1,114 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts #-} +-- | Dependency solver +-- +-- This is a relatively simple problem due to the choice of the +-- Minimum Package Version algorithm. In fact, the only failure mode +-- is referencing an unknown package or revision. +module Futhark.Pkg.Solve + ( solveDeps + , solveDepsPure + , PkgRevDepInfo + ) where + +import Control.Monad.State +import qualified Data.Set as S +import qualified Data.Map as M +import qualified Data.Text as T +import Data.Monoid ((<>)) + +import Control.Monad.Free.Church + +import Futhark.Pkg.Info +import Futhark.Pkg.Types + +import Prelude + +data PkgOp a = OpGetDeps PkgPath SemVer (Maybe T.Text) (PkgRevDeps -> a) + +instance Functor PkgOp where + fmap f (OpGetDeps p v h c) = OpGetDeps p v h (f . c) + +-- | A rough build list is like a build list, but may contain packages +-- that are not reachable from the root. Also contains the +-- dependencies of each package. +newtype RoughBuildList = RoughBuildList (M.Map PkgPath (SemVer, [PkgPath])) + deriving (Show) + +emptyRoughBuildList :: RoughBuildList +emptyRoughBuildList = RoughBuildList mempty + +depRoots :: PkgRevDeps -> S.Set PkgPath +depRoots (PkgRevDeps m) = S.fromList $ M.keys m + +-- | Construct a 'BuildList' from a 'RoughBuildList'. This involves +-- pruning all packages that cannot be reached from the root. +buildList :: S.Set PkgPath -> RoughBuildList -> BuildList +buildList roots (RoughBuildList pkgs) = + BuildList $ execState (mapM_ addPkg roots) mempty + where addPkg p = case M.lookup p pkgs of + Nothing -> return () + Just (v, deps) -> do + listed <- gets $ M.member p + modify $ M.insert p v + unless listed $ mapM_ addPkg deps + +type SolveM = StateT RoughBuildList (F PkgOp) + +getDeps :: PkgPath -> SemVer -> Maybe T.Text -> SolveM PkgRevDeps +getDeps p v h = lift $ liftF $ OpGetDeps p v h id + +-- | Given a list of immediate dependency minimum version constraints, +-- find dependency versions that fit, including transitive +-- dependencies. +doSolveDeps :: PkgRevDeps -> SolveM () +doSolveDeps (PkgRevDeps deps) = mapM_ add $ M.toList deps + where add (p, (v, maybe_h)) = do + RoughBuildList l <- get + case M.lookup p l of + -- Already satisfied? + Just (cur_v, _) | v <= cur_v -> return () + -- No; add 'p' and its dependencies. + _ -> do + PkgRevDeps p_deps <- getDeps p v maybe_h + put $ RoughBuildList $ M.insert p (v, M.keys p_deps) l + mapM_ add $ M.toList p_deps + +-- | Run the solver, producing both a package registry containing +-- a cache of the lookups performed, as well as a build list. +solveDeps :: MonadPkgRegistry m => + PkgRevDeps -> m BuildList +solveDeps deps = buildList (depRoots deps) <$> runF + (execStateT (doSolveDeps deps) emptyRoughBuildList) + return step + where step (OpGetDeps p v h c) = do + pinfo <- lookupPackageRev p v + + checkHash p v pinfo h + + d <- fmap pkgRevDeps . getManifest $ pkgRevGetManifest pinfo + c d + + checkHash _ _ _ Nothing = return () + checkHash p v pinfo (Just h) + | h == pkgRevCommit pinfo = return () + | otherwise = fail $ T.unpack $ "Package " <> p <> " " <> prettySemVer v <> + " has commit hash " <> pkgRevCommit pinfo <> + ", but expected " <> h <> " from package manifest." + +-- | A mapping of package revisions to the dependencies of that +-- package. Can be considered a 'PkgRegistry' without the option of +-- obtaining more information from the Internet. Probably useful only +-- for testing the solver. +type PkgRevDepInfo = M.Map (PkgPath, SemVer) PkgRevDeps + +-- | Perform package resolution with only pre-known information. This +-- is useful for testing. +solveDepsPure :: PkgRevDepInfo -> PkgRevDeps -> Either T.Text BuildList +solveDepsPure r deps = buildList (depRoots deps) <$> runF + (execStateT (doSolveDeps deps) emptyRoughBuildList) + Right step + where step (OpGetDeps p v _ c) = do + let errmsg = "Unknown package/version: " <> p <> "-" <> prettySemVer v + d <- maybe (Left errmsg) Right $ M.lookup (p,v) r + c d diff --git a/src/Futhark/Pkg/Types.hs b/src/Futhark/Pkg/Types.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pkg/Types.hs @@ -0,0 +1,296 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | Types (and a few other simple definitions) for futhark-pkg. +module Futhark.Pkg.Types + ( PkgPath + , pkgPathFilePath + , PkgRevDeps(..) + , module Data.Versions + + -- * Versions + , commitVersion + , isCommitVersion + , parseVersion + + -- * Package manifests + , PkgManifest(..) + , newPkgManifest + , pkgRevDeps + , pkgDir + , addRequiredToManifest + , removeRequiredFromManifest + , prettyPkgManifest + , Comment + , Commented(..) + , Required(..) + , futharkPkg + + -- * Parsing package manifests + , parsePkgManifest + , parsePkgManifestFromFile + , errorBundlePretty + + -- * Build list + , BuildList(..) + , prettyBuildList + ) where + +import Control.Applicative +import Control.Monad +import Data.Either +import Data.Foldable +import Data.List +import Data.Maybe +import Data.Traversable +import Data.Void +import Data.Semigroup ((<>)) +import qualified Data.Semigroup as Sem +import qualified Data.Text as T +import qualified Data.Text.IO as T +import qualified Data.Map as M +import System.FilePath +import qualified System.FilePath.Posix as Posix + +import Data.Versions (SemVer(..), VUnit(..), prettySemVer) +import Text.Megaparsec hiding (many, some) +import Text.Megaparsec.Char + +import Prelude + +-- | A package path is a unique identifier for a package, for example +-- @github.com/user/foo@. +type PkgPath = T.Text + +-- | Turn a package path (which always uses forward slashes) into a +-- file path in the local file system (which might use different +-- slashes). +pkgPathFilePath :: PkgPath -> FilePath +pkgPathFilePath = joinPath . Posix.splitPath . T.unpack + +-- | Versions of the form (0,0,0)-timestamp+hash are treated +-- specially, as a reference to the commit identified uniquely with +-- 'hash' (typically the Git commit ID). This function detects such +-- versions. +isCommitVersion :: SemVer -> Maybe T.Text +isCommitVersion (SemVer 0 0 0 [_] [[Str s]]) = Just s +isCommitVersion _ = Nothing + +-- | @commitVersion timestamp commit@ constructs a commit version. +commitVersion :: T.Text -> T.Text -> SemVer +commitVersion time commit = + SemVer 0 0 0 [[Str time]] [[Str commit]] + +-- | Unfortunately, Data.Versions has a buggy semver parser that +-- collapses consecutive zeroes in the metadata field. So, we define +-- our own parser here. It's a little simpler too, since we don't +-- need full semver. +parseVersion :: T.Text -> Either (ParseErrorBundle T.Text Void) SemVer +parseVersion = parse (semver' <* eof) "Semantic Version" + +semver' :: Parsec Void T.Text SemVer +semver' = SemVer <$> majorP <*> minorP <*> patchP <*> preRel <*> metaData + where majorP = digitsP <* char '.' + minorP = majorP + patchP = digitsP + digitsP = read <$> ((T.unpack <$> string "0") <|> some digitChar) + preRel = maybe [] pure <$> optional preRel' + preRel' = char '-' *> (pure . Str . T.pack <$> some digitChar) + metaData = maybe [] pure <$> optional metaData' + metaData' = char '+' *> (pure . Str . T.pack <$> some alphaNumChar) + +-- | The dependencies of a (revision of a) package is a mapping from +-- package paths to minimum versions (and an optional hash pinning). +newtype PkgRevDeps = PkgRevDeps (M.Map PkgPath (SemVer, Maybe T.Text)) + deriving (Show) + +instance Sem.Semigroup PkgRevDeps where + PkgRevDeps x <> PkgRevDeps y = PkgRevDeps $ x <> y + +instance Monoid PkgRevDeps where + mempty = PkgRevDeps mempty + mappend = (Sem.<>) + +--- Package manifest + +-- | A line comment. +type Comment = T.Text + +-- | Wraps a value with an annotation of preceding line comments. +-- This is important to our goal of being able to programmatically +-- modify the @futhark.pkg@ file while keeping comments intact. +data Commented a = Commented { comments :: [Comment] + , commented :: a + } + deriving (Show, Eq) + +instance Functor Commented where + fmap = fmapDefault + +instance Foldable Commented where + foldMap = foldMapDefault + +instance Traversable Commented where + traverse f (Commented cs x) = Commented cs <$> f x + +-- | An entry in the @required@ section of a @futhark.pkg@ file. +data Required = Required + { requiredPkg :: PkgPath + -- ^ Name of the required package. + , requiredPkgRev :: SemVer + -- ^ The minimum revision. + , requiredHash :: Maybe T.Text + -- ^ An optional hash indicating what + -- this revision looked like the last + -- time we saw it. Used for integrity + -- checking. + } + deriving (Show, Eq) + +-- | The name of the file containing the futhark-pkg manifest. +futharkPkg :: FilePath +futharkPkg = "futhark.pkg" + +-- | A structure corresponding to a @futhark.pkg@ file, including +-- comments. It is an invariant that duplicate required packages do +-- not occcur (the parser will verify this). +data PkgManifest = PkgManifest { manifestPkgPath :: Commented (Maybe PkgPath) + -- ^ The name of the package. + , manifestRequire :: Commented [Either Comment Required] + , manifestEndComments :: [Comment] + } + deriving (Show, Eq) + +-- | Possibly given a package path, construct an otherwise-empty manifest file. +newPkgManifest :: Maybe PkgPath -> PkgManifest +newPkgManifest p = + PkgManifest (Commented mempty p) (Commented mempty mempty) mempty + +-- | Prettyprint a package manifest such that it can be written to a +-- @futhark.pkg@ file. +prettyPkgManifest :: PkgManifest -> T.Text +prettyPkgManifest (PkgManifest name required endcs) = + T.unlines $ concat [ prettyComments name + , maybe [] (pure . ("package "<>) . (<>"\n")) $ commented name + , prettyComments required + , ["require {"] + , map ((" "<>) . prettyRequired) $ commented required + , ["}"] + , map prettyComment endcs + ] + where prettyComments = map prettyComment . comments + prettyComment = ("--"<>) + prettyRequired (Left c) = prettyComment c + prettyRequired (Right (Required p r h)) = + T.unwords $ catMaybes [Just p, + Just $ prettySemVer r, + ("#"<>) <$> h] + +-- | The required packages listed in a package manifest. +pkgRevDeps :: PkgManifest -> PkgRevDeps +pkgRevDeps = PkgRevDeps . M.fromList . mapMaybe onR . + commented . manifestRequire + where onR (Right r) = Just (requiredPkg r, (requiredPkgRev r, requiredHash r)) + onR (Left _) = Nothing + +-- | Where in the corresponding repository archive we can expect to +-- find the package files. +pkgDir :: PkgManifest -> Maybe Posix.FilePath +pkgDir = fmap (Posix.addTrailingPathSeparator . ("lib" Posix.) . + T.unpack) . commented . manifestPkgPath + +-- | Add new required package to the package manifest. If the package +-- was already present, return the old version. +addRequiredToManifest :: Required -> PkgManifest -> (PkgManifest, Maybe Required) +addRequiredToManifest new_r pm = + let (old, requires') = mapAccumL add Nothing $ commented $ manifestRequire pm + in (if isJust old + then pm { manifestRequire = const requires' <$> manifestRequire pm } + else pm { manifestRequire = (++[Right new_r]) <$> manifestRequire pm }, + old) + where add acc (Left c) = (acc, Left c) + add acc (Right r) + | requiredPkg r == requiredPkg new_r = (Just r, Right new_r) + | otherwise = (acc, Right r) + +-- | Check if the manifest specifies a required package with the given +-- package path. +requiredInManifest :: PkgPath -> PkgManifest -> Maybe Required +requiredInManifest p = + find ((==p) . requiredPkg) . rights . commented . manifestRequire + +-- | Remove a required package from the manifest. Returns 'Nothing' +-- if the package was not found in the manifest, and otherwise the new +-- manifest and the 'Required' that was present. +removeRequiredFromManifest :: PkgPath -> PkgManifest -> Maybe (PkgManifest, Required) +removeRequiredFromManifest p pm = do + r <- requiredInManifest p pm + return (pm { manifestRequire = filter (not . matches) <$> manifestRequire pm }, + r) + where matches = either (const False) ((==p) . requiredPkg) + +--- Parsing futhark.pkg. + +type Parser = Parsec Void T.Text + +pPkgManifest :: Parser PkgManifest +pPkgManifest = do + c1 <- pComments + p <- optional $ lexstr "package" *> pPkgPath + space + c2 <- pComments + required <- (lexstr "require" *> + braces (many $ (Left <$> pComment) <|> (Right <$> pRequired))) + <|> pure [] + c3 <- pComments + eof + return $ PkgManifest (Commented c1 p) (Commented c2 required) c3 + where lexeme :: Parser a -> Parser a + lexeme p = p <* space + + lexeme' p = p <* spaceNoEol + + lexstr :: T.Text -> Parser () + lexstr = void . try . lexeme . string + + braces :: Parser a -> Parser a + braces p = lexstr "{" *> p <* lexstr "}" + + spaceNoEol = many $ oneOf (" \t" :: String) + + pPkgPath = T.pack <$> some (alphaNumChar <|> oneOf ("@-/.:" :: String)) + "package path" + + pRequired = space *> (Required <$> lexeme' pPkgPath + <*> lexeme' semver' + <*> optional (lexeme' pHash)) <* space + "package requirement" + + pHash = char '#' *> (T.pack <$> some alphaNumChar) + + pComment = lexeme $ T.pack <$> (string "--" >> anySingle `manyTill` (void eol <|> eof)) + + pComments :: Parser [Comment] + pComments = catMaybes <$> many (comment <|> blankLine) + where comment = Just <$> pComment + blankLine = some spaceChar >> pure Nothing + + +parsePkgManifest :: FilePath -> T.Text -> Either (ParseErrorBundle T.Text Void) PkgManifest +parsePkgManifest = parse pPkgManifest + +parsePkgManifestFromFile :: FilePath -> IO PkgManifest +parsePkgManifestFromFile f = do + s <- T.readFile f + case parsePkgManifest f s of + Left err -> fail $ errorBundlePretty err + Right m -> return m + +-- | A mapping from package paths to their chosen revisions. This is +-- the result of the version solver. +newtype BuildList = BuildList { unBuildList :: M.Map PkgPath SemVer } + deriving (Eq, Show) + +-- | Prettyprint a build list; one package per line and +-- newline-terminated. +prettyBuildList :: BuildList -> T.Text +prettyBuildList (BuildList m) = T.unlines $ map f $ sortOn fst $ M.toList m + where f (p, v) = T.unwords [p, "=>", prettySemVer v] diff --git a/src/Futhark/Representation/AST.hs b/src/Futhark/Representation/AST.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST.hs @@ -0,0 +1,16 @@ +-- | A convenient re-export of basic AST modules. Note that +-- "Futhark.Representation.AST.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. +module Futhark.Representation.AST + ( module Futhark.Representation.AST.Attributes + , module Futhark.Representation.AST.Traversals + , module Futhark.Representation.AST.Pretty + , module Futhark.Representation.AST.Syntax + ) +where + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Pretty diff --git a/src/Futhark/Representation/AST/Annotations.hs b/src/Futhark/Representation/AST/Annotations.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Annotations.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE TypeFamilies, FlexibleContexts #-} +module Futhark.Representation.AST.Annotations + ( Annotations (..) + , module Futhark.Representation.AST.RetType + ) + where + +import Futhark.Representation.AST.Syntax.Core +import Futhark.Representation.AST.RetType +import Futhark.Representation.AST.Attributes.Types + +class (Show (LetAttr l), Show (ExpAttr l), Show (BodyAttr l), Show (FParamAttr l), Show (LParamAttr l), Show (RetType l), Show (BranchType l), Show (Op l), + Eq (LetAttr l), Eq (ExpAttr l), Eq (BodyAttr l), Eq (FParamAttr l), Eq (LParamAttr l), Eq (RetType l), Eq (BranchType l), Eq (Op l), + Ord (LetAttr l), Ord (ExpAttr l), Ord (BodyAttr l), Ord (FParamAttr l), Ord (LParamAttr l), Ord (RetType l), Ord (BranchType l), Ord (Op l), + IsRetType (RetType l), IsBodyType (BranchType l), + Typed (FParamAttr l), Typed (LParamAttr l), Typed (LetAttr l), + DeclTyped (FParamAttr l)) + => Annotations l where + -- | Annotation for every let-pattern element. + type LetAttr l :: * + type LetAttr l = Type + -- | Annotation for every expression. + type ExpAttr l :: * + type ExpAttr l = () + -- | Annotation for every body. + type BodyAttr l :: * + type BodyAttr l = () + -- | Annotation for every (non-lambda) function parameter. + type FParamAttr l :: * + type FParamAttr l = DeclType + -- | Annotation for every lambda function parameter. + type LParamAttr l :: * + type LParamAttr l = Type + + -- | The return type annotation of function calls. + type RetType l :: * + type RetType l = DeclExtType + + -- | The return type annotation of branches. + type BranchType l :: * + type BranchType l = ExtType + + -- | Extensible operation. + type Op l :: * + type Op l = () diff --git a/src/Futhark/Representation/AST/Attributes.hs b/src/Futhark/Representation/AST/Attributes.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes.hs @@ -0,0 +1,224 @@ +{-# LANGUAGE TypeFamilies, FlexibleContexts, FlexibleInstances, ConstraintKinds #-} +-- | This module provides various simple ways to query and manipulate +-- fundamental Futhark terms, such as types and values. The intent is to +-- keep "Futhark.Reprsentation.AST.Syntax" simple, and put whatever +-- embellishments we need here. This is an internal, desugared +-- representation. +module Futhark.Representation.AST.Attributes + ( module Futhark.Representation.AST.Attributes.Reshape + , module Futhark.Representation.AST.Attributes.Rearrange + , module Futhark.Representation.AST.Attributes.Types + , module Futhark.Representation.AST.Attributes.Constants + , module Futhark.Representation.AST.Attributes.TypeOf + , module Futhark.Representation.AST.Attributes.Patterns + , module Futhark.Representation.AST.Attributes.Names + , module Futhark.Representation.AST.RetType + + -- * Built-in functions + , isBuiltInFunction + , builtInFunctions + + -- * Extra tools + , funDefByName + , asBasicOp + , safeExp + , subExpVars + , subExpVar + , shapeVars + , commutativeLambda + , entryPointSize + , defAux + , stmCerts + , certify + , expExtTypesFromPattern + + , IsOp (..) + , Attributes (..) + ) + where + +import Data.List +import Data.Maybe (mapMaybe, isJust) +import Data.Monoid ((<>)) +import qualified Data.Map.Strict as M + +import Futhark.Representation.AST.Attributes.Reshape +import Futhark.Representation.AST.Attributes.Rearrange +import Futhark.Representation.AST.Attributes.Types +import Futhark.Representation.AST.Attributes.Constants +import Futhark.Representation.AST.Attributes.Patterns +import Futhark.Representation.AST.Attributes.Names +import Futhark.Representation.AST.Attributes.TypeOf +import Futhark.Representation.AST.RetType +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Pretty +import Futhark.Transform.Rename (Rename, Renameable) +import Futhark.Transform.Substitute (Substitute, Substitutable) +import Futhark.Util.Pretty + +-- | @isBuiltInFunction k@ is 'True' if @k@ is an element of 'builtInFunctions'. +isBuiltInFunction :: Name -> Bool +isBuiltInFunction fnm = fnm `M.member` builtInFunctions + +-- | A map of all built-in functions and their types. +builtInFunctions :: M.Map Name (PrimType,[PrimType]) +builtInFunctions = M.fromList $ map namify $ M.toList primFuns + where namify (k,(paramts,ret,_)) = (nameFromString k, (ret, paramts)) + +-- | Find the function of the given name in the Futhark program. +funDefByName :: Name -> Prog lore -> Maybe (FunDef lore) +funDefByName fname = find ((fname ==) . funDefName) . progFunctions + +-- | If the expression is a 'BasicOp', return that 'BasicOp', otherwise 'Nothing'. +asBasicOp :: Exp lore -> Maybe (BasicOp lore) +asBasicOp (BasicOp op) = Just op +asBasicOp _ = Nothing + +-- | An expression is safe if it is always well-defined (assuming that +-- 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 (BasicOp op) = safeBasicOp op + where safeBasicOp (BinOp SDiv{} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SDiv{} _ _) = False + safeBasicOp (BinOp UDiv{} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp UDiv{} _ _) = False + safeBasicOp (BinOp SMod{} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SMod{} _ _) = False + safeBasicOp (BinOp UMod{} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp UMod{} _ _) = False + + safeBasicOp (BinOp SQuot{} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SQuot{} _ _) = False + safeBasicOp (BinOp SRem{} _ (Constant y)) = not $ zeroIsh y + safeBasicOp (BinOp SRem{} _ _) = False + + safeBasicOp (BinOp Pow{} _ (Constant y)) = not $ negativeIsh y + safeBasicOp (BinOp Pow{} _ _) = False + safeBasicOp ArrayLit{} = True + safeBasicOp BinOp{} = True + safeBasicOp SubExp{} = True + safeBasicOp UnOp{} = True + safeBasicOp CmpOp{} = True + safeBasicOp ConvOp{} = True + safeBasicOp Scratch{} = True + safeBasicOp Concat{} = True + safeBasicOp Reshape{} = True + safeBasicOp Manifest{} = True + safeBasicOp Iota{} = True + safeBasicOp Replicate{} = True + safeBasicOp Copy{} = True + safeBasicOp _ = False + +safeExp (DoLoop _ _ _ body) = safeBody body +safeExp (Apply fname _ _ _) = isBuiltInFunction fname +safeExp (If _ tbranch fbranch _) = + all (safeExp . stmExp) (bodyStms tbranch) && + all (safeExp . stmExp) (bodyStms fbranch) +safeExp (Op op) = safeOp op + +safeBody :: IsOp (Op lore) => Body lore -> Bool +safeBody = all (safeExp . stmExp) . bodyStms + +-- | Return the variable names used in 'Var' subexpressions. May contain +-- duplicates. +subExpVars :: [SubExp] -> [VName] +subExpVars = mapMaybe subExpVar + +-- | If the 'SubExp' is a 'Var' return the variable name. +subExpVar :: SubExp -> Maybe VName +subExpVar (Var v) = Just v +subExpVar Constant{} = Nothing + +-- | Return the variable dimension sizes. May contain +-- duplicates. +shapeVars :: Shape -> [VName] +shapeVars = subExpVars . shapeDims + +-- | Does the given lambda represent a known commutative function? +-- 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 lam = + let body = lambdaBody lam + n2 = length (lambdaParams lam) `div` 2 + (xps,yps) = splitAt n2 (lambdaParams lam) + + okComponent c = isJust $ find (okBinOp c) $ bodyStms body + okBinOp (xp,yp,Var r) (Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y)))) = + patElemName pe == r && + commutativeBinOp op && + ((x == paramName xp && y == paramName yp) || + (y == paramName xp && x == paramName yp)) + okBinOp _ _ = False + + in n2 * 2 == length (lambdaParams lam) && + n2 == length (bodyResult body) && + all okComponent (zip3 xps yps $ bodyResult body) + +-- | How many value parameters are accepted by this entry point? This +-- is used to determine which of the function parameters correspond to +-- 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 + +-- | A 'StmAux' with empty 'Certificates'. +defAux :: attr -> StmAux attr +defAux = StmAux mempty + +-- | The certificates associated with a statement. +stmCerts :: Stm lore -> Certificates +stmCerts = stmAuxCerts . stmAux + +-- | Add certificates to a statement. +certify :: Certificates -> Stm lore -> Stm lore +certify cs1 (Let pat (StmAux cs2 attr) e) = Let pat (StmAux (cs2<>cs1) attr) e + +-- | A type class for operations. +class (Eq op, Ord op, Show op, + TypedOp op, + Rename op, + Substitute op, + FreeIn op, + Pretty op) => IsOp op where + -- | Like 'safeExp', but for arbitrary ops. + safeOp :: op -> Bool + -- | Should we try to hoist this out of branches? + cheapOp :: op -> Bool + +instance IsOp () where + safeOp () = True + cheapOp () = True + +-- | Lore-specific attributes; also means the lore supports some basic +-- facilities. +class (Annotations lore, + + PrettyLore lore, + + Renameable lore, Substitutable lore, + FreeAttr (ExpAttr lore), + FreeIn (LetAttr lore), + FreeAttr (BodyAttr lore), + FreeIn (FParamAttr lore), + FreeIn (LParamAttr lore), + FreeIn (RetType lore), + FreeIn (BranchType lore), + + IsOp (Op lore)) => Attributes lore where + -- | Given a pattern, construct the type of a body that would match + -- it. An implementation for many lores would be + -- 'expExtTypesFromPattern'. + expTypesFromPattern :: (HasScope lore m, Monad m) => + Pattern lore -> m [BranchType lore] + +-- | Construct the type of an expression that would match the pattern. +expExtTypesFromPattern :: Typed attr => PatternT attr -> [ExtType] +expExtTypesFromPattern pat = + existentialiseExtTypes (patternContextNames pat) $ + staticShapes $ map patElemType $ patternValueElements pat diff --git a/src/Futhark/Representation/AST/Attributes/Aliases.hs b/src/Futhark/Representation/AST/Attributes/Aliases.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Aliases.hs @@ -0,0 +1,172 @@ +{-# LANGUAGE TypeFamilies #-} +{-# Language FlexibleInstances, FlexibleContexts #-} +module Futhark.Representation.AST.Attributes.Aliases + ( vnameAliases + , subExpAliases + , primOpAliases + , expAliases + , patternAliases + , Aliased (..) + , AliasesOf (..) + -- * Consumption + , consumedInStm + , consumedInExp + , consumedByLambda + -- * Extensibility + , AliasedOp (..) + , CanBeAliased (..) + ) + where + +import Control.Arrow (first) +import Data.Monoid ((<>)) +import qualified Data.Set as S + +import Futhark.Representation.AST.Attributes (IsOp) +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes.Patterns +import Futhark.Representation.AST.Attributes.Types + +class (Annotations lore, AliasedOp (Op lore), + AliasesOf (LetAttr lore)) => Aliased lore where + bodyAliases :: Body lore -> [Names] + consumedInBody :: Body lore -> Names + +vnameAliases :: VName -> Names +vnameAliases = S.singleton + +subExpAliases :: SubExp -> Names +subExpAliases Constant{} = mempty +subExpAliases (Var v) = vnameAliases v + +primOpAliases :: BasicOp lore -> [Names] +primOpAliases (SubExp se) = [subExpAliases se] +primOpAliases (Opaque se) = [subExpAliases se] +primOpAliases (ArrayLit _ _) = [mempty] +primOpAliases BinOp{} = [mempty] +primOpAliases ConvOp{} = [mempty] +primOpAliases CmpOp{} = [mempty] +primOpAliases UnOp{} = [mempty] + +primOpAliases (Index ident _) = + [vnameAliases ident] +primOpAliases Update{} = + [mempty] +primOpAliases Iota{} = + [mempty] +primOpAliases Replicate{} = + [mempty] +primOpAliases (Repeat _ _ v) = + [vnameAliases v] +primOpAliases Scratch{} = + [mempty] +primOpAliases (Reshape _ e) = + [vnameAliases e] +primOpAliases (Rearrange _ e) = + [vnameAliases e] +primOpAliases (Rotate _ e) = + [vnameAliases e] +primOpAliases Concat{} = + [mempty] +primOpAliases Copy{} = + [mempty] +primOpAliases Manifest{} = + [mempty] +primOpAliases Assert{} = + [mempty] +primOpAliases (Partition n _ arr) = + replicate n mempty ++ map vnameAliases arr + +ifAliases :: ([Names], Names) -> ([Names], Names) -> [Names] +ifAliases (als1,cons1) (als2,cons2) = + map (S.filter notConsumed) $ zipWith mappend als1 als2 + where notConsumed = not . (`S.member` cons) + cons = cons1 <> cons2 + +funcallAliases :: [(SubExp, Diet)] -> [TypeBase shape Uniqueness] -> [Names] +funcallAliases args t = + returnAliases t [(subExpAliases se, d) | (se,d) <- args ] + +expAliases :: (Aliased lore) => Exp lore -> [Names] +expAliases (If _ tb fb attr) = + drop (length all_aliases - length ts) all_aliases + where ts = ifReturns attr + all_aliases = ifAliases + (bodyAliases tb, consumedInBody tb) + (bodyAliases fb, consumedInBody fb) +expAliases (BasicOp op) = primOpAliases op +expAliases (DoLoop ctxmerge valmerge _ loopbody) = + map (`S.difference` merge_names) val_aliases + where (_ctx_aliases, val_aliases) = + splitAt (length ctxmerge) $ bodyAliases loopbody + merge_names = S.fromList $ map (paramName . fst) $ ctxmerge ++ valmerge +expAliases (Apply _ args t _) = + funcallAliases args $ retTypeValues t +expAliases (Op op) = opAliases op + +returnAliases :: [TypeBase shaper Uniqueness] -> [(Names, Diet)] -> [Names] +returnAliases rts args = map returnType' rts + where returnType' (Array _ _ Nonunique) = + mconcat $ map (uncurry maskAliases) args + returnType' (Array _ _ Unique) = + mempty + returnType' (Prim _) = + mempty + returnType' Mem{} = + error "returnAliases Mem" + +maskAliases :: Names -> Diet -> Names +maskAliases _ Consume = mempty +maskAliases als Observe = als + +consumedInStm :: Aliased lore => Stm lore -> Names +consumedInStm = consumedInExp . stmExp + +consumedInExp :: (Aliased lore) => Exp lore -> Names +consumedInExp (Apply _ args _ _) = + mconcat (map (consumeArg . first subExpAliases) args) + where consumeArg (als, Consume) = als + consumeArg (_, Observe) = mempty +consumedInExp (If _ tb fb _) = + consumedInBody tb <> consumedInBody fb +consumedInExp (DoLoop _ merge _ _) = + mconcat (map (subExpAliases . snd) $ + filter (unique . paramDeclType . fst) merge) +consumedInExp (BasicOp (Update src _ _)) = S.singleton src +consumedInExp (Op op) = consumedInOp op +consumedInExp _ = mempty + +consumedByLambda :: Aliased lore => Lambda lore -> Names +consumedByLambda = consumedInBody . lambdaBody + +patternAliases :: AliasesOf attr => PatternT attr -> [Names] +patternAliases = map (aliasesOf . patElemAttr) . patternElements + +-- | Something that contains alias information. +class AliasesOf a where + -- | The alias of the argument element. + aliasesOf :: a -> Names + +instance AliasesOf Names where + aliasesOf = id + +instance AliasesOf attr => AliasesOf (PatElemT attr) where + aliasesOf = aliasesOf . patElemAttr + +class IsOp op => AliasedOp op where + opAliases :: op -> [Names] + consumedInOp :: op -> Names + +instance AliasedOp () where + opAliases () = [] + consumedInOp () = mempty + +class AliasedOp (OpWithAliases op) => CanBeAliased op where + type OpWithAliases op :: * + removeOpAliases :: OpWithAliases op -> op + addOpAliases :: op -> OpWithAliases op + +instance CanBeAliased () where + type OpWithAliases () = () + removeOpAliases = id + addOpAliases = id diff --git a/src/Futhark/Representation/AST/Attributes/Constants.hs b/src/Futhark/Representation/AST/Attributes/Constants.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Constants.hs @@ -0,0 +1,76 @@ +-- | Possibly convenient facilities for constructing constants. +module Futhark.Representation.AST.Attributes.Constants + ( + IsValue (..) + , constant + , intConst + , floatConst + ) + where + +import Futhark.Representation.AST.Syntax.Core + +-- | If a Haskell type is an instance of 'IsValue', it means that a +-- value of that type can be converted to a Futhark 'PrimValue'. +-- This is intended to cut down on boilerplate when writing compiler +-- code - for example, you'll quickly grow tired of writing @Constant +-- (LogVal True) loc@. +class IsValue a where + value :: a -> PrimValue + +instance IsValue Int where + value = IntValue . Int32Value . fromIntegral + +instance IsValue Int8 where + value = IntValue . Int8Value + +instance IsValue Int16 where + value = IntValue . Int16Value + +instance IsValue Int32 where + value = IntValue . Int32Value + +instance IsValue Int64 where + value = IntValue . Int64Value + +instance IsValue Word8 where + value = IntValue . Int8Value . fromIntegral + +instance IsValue Word16 where + value = IntValue . Int16Value . fromIntegral + +instance IsValue Word32 where + value = IntValue . Int32Value . fromIntegral + +instance IsValue Word64 where + value = IntValue . Int64Value . fromIntegral + +instance IsValue Double where + value = FloatValue . Float64Value + +instance IsValue Float where + value = FloatValue . Float32Value + +instance IsValue Bool where + value = BoolValue + +instance IsValue PrimValue where + value = id + +instance IsValue IntValue where + value = IntValue + +instance IsValue FloatValue where + value = FloatValue + +-- | Create a 'Constant' 'SubExp' containing the given value. +constant :: IsValue v => v -> SubExp +constant = Constant . value + +-- | Utility definition for reasons of type ambiguity. +intConst :: IntType -> Integer -> SubExp +intConst t v = constant $ intValue t v + +-- | Utility definition for reasons of type ambiguity. +floatConst :: FloatType -> Double -> SubExp +floatConst t v = constant $ floatValue t v diff --git a/src/Futhark/Representation/AST/Attributes/Names.hs b/src/Futhark/Representation/AST/Attributes/Names.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Names.hs @@ -0,0 +1,243 @@ +{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, FlexibleContexts, UndecidableInstances #-} +-- | Facilities for determining which names are used in some syntactic +-- construct. The most important interface is the 'FreeIn' class and +-- its instances, but for reasons related to the Haskell type system, +-- some constructs have specialised functions. +module Futhark.Representation.AST.Attributes.Names + ( + -- * Class + FreeIn (..) + , Names + -- * Specialised Functions + , freeInStmsAndRes + , freeInBody + , freeInExp + , freeInStm + , freeInLambda + -- * Bound Names + , boundInBody + , boundByStm + , boundByStms + , boundByLambda + + , FreeAttr(..) + ) + where + +import Control.Monad.Writer +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Foldable + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Attributes.Patterns +import Futhark.Representation.AST.Attributes.Scope + +freeWalker :: (FreeAttr (ExpAttr lore), + FreeAttr (BodyAttr lore), + FreeIn (FParamAttr lore), + FreeIn (LParamAttr lore), + FreeIn (LetAttr lore), + FreeIn (Op lore)) => + Walker lore (Writer Names) +freeWalker = identityWalker { + walkOnSubExp = tell . freeIn + , walkOnBody = tell . freeInBody + , walkOnVName = tell . S.singleton + , walkOnCertificates = tell . freeIn + , walkOnOp = tell . freeIn + } + +-- | Return the set of variable names that are free in the given +-- statements and result. Filters away the names that are bound by +-- the statements. +freeInStmsAndRes :: (FreeIn (Op lore), + FreeIn (LetAttr lore), + FreeIn (LParamAttr lore), + FreeIn (FParamAttr lore), + FreeAttr (BodyAttr lore), + FreeAttr (ExpAttr lore)) => + Stms lore -> Result -> Names +freeInStmsAndRes stms res = + (freeIn res `mappend` fold (fmap freeInStm stms)) + `S.difference` boundByStms stms + +-- | Return the set of variable names that are free in the given body. +freeInBody :: (FreeAttr (ExpAttr lore), + FreeAttr (BodyAttr lore), + FreeIn (FParamAttr lore), + FreeIn (LParamAttr lore), + FreeIn (LetAttr lore), + FreeIn (Op lore)) => + Body lore -> Names +freeInBody (Body attr stms res) = + precomputed attr $ freeIn attr <> freeInStmsAndRes stms res + +-- | Return the set of variable names that are free in the given +-- expression. +freeInExp :: (FreeAttr (ExpAttr lore), + FreeAttr (BodyAttr lore), + FreeIn (FParamAttr lore), + FreeIn (LParamAttr lore), + FreeIn (LetAttr lore), + FreeIn (Op lore)) => + Exp lore -> Names +freeInExp (DoLoop ctxmerge valmerge form loopbody) = + let (ctxparams, ctxinits) = unzip ctxmerge + (valparams, valinits) = unzip valmerge + bound_here = S.fromList $ M.keys $ + scopeOf form <> + scopeOfFParams (ctxparams ++ valparams) + in (freeIn (ctxinits ++ valinits) <> freeIn form <> + freeIn (ctxparams ++ valparams) <> freeInBody loopbody) + `S.difference` bound_here +freeInExp e = execWriter $ walkExpM freeWalker e + +-- | Return the set of variable names that are free in the given +-- binding. +freeInStm :: (FreeAttr (ExpAttr lore), + FreeAttr (BodyAttr lore), + FreeIn (FParamAttr lore), + FreeIn (LParamAttr lore), + FreeIn (LetAttr lore), + FreeIn (Op lore)) => + Stm lore -> Names +freeInStm (Let pat (StmAux cs attr) e) = + freeIn cs <> precomputed attr (freeIn attr <> freeInExp e <> freeIn pat) + +-- | Return the set of variable names that are free in the given +-- lambda, including shape annotations in the parameters. +freeInLambda :: (FreeAttr (ExpAttr lore), + FreeAttr (BodyAttr lore), + FreeIn (FParamAttr lore), + FreeIn (LParamAttr lore), + FreeIn (LetAttr lore), + FreeIn (Op lore)) => + Lambda lore -> Names +freeInLambda (Lambda params body rettype) = + S.filter (`notElem` paramnames) $ inRet <> inParams <> inBody + where inRet = mconcat $ map freeIn rettype + inParams = mconcat $ map freeIn params + inBody = freeInBody body + paramnames = map paramName params + +-- | A class indicating that we can obtain free variable information +-- from values of this type. +class FreeIn a where + freeIn :: a -> Names + +instance FreeIn () where + freeIn () = mempty + +instance FreeIn Int where + freeIn = const mempty + +instance (FreeIn a, FreeIn b) => FreeIn (a,b) where + freeIn (a,b) = freeIn a <> freeIn b + +instance (FreeIn a, FreeIn b, FreeIn c) => FreeIn (a,b,c) where + freeIn (a,b,c) = freeIn a <> freeIn b <> freeIn c + +instance FreeIn a => FreeIn [a] where + freeIn = fold . fmap freeIn + +instance FreeIn (Stm lore) => FreeIn (Stms lore) where + freeIn = fold . fmap freeIn + +instance FreeIn Names where + freeIn = id + +instance FreeIn Bool where + freeIn _ = mempty + +instance FreeIn a => FreeIn (Maybe a) where + freeIn = maybe mempty freeIn + +instance FreeIn VName where + freeIn = S.singleton + +instance FreeIn Ident where + freeIn = freeIn . identType + +instance FreeIn SubExp where + freeIn (Var v) = freeIn v + freeIn Constant{} = mempty + +instance FreeIn d => FreeIn (ShapeBase d) where + freeIn = mconcat . map freeIn . shapeDims + +instance FreeIn d => FreeIn (Ext d) where + freeIn (Free x) = freeIn x + freeIn (Ext _) = mempty + +instance FreeIn shape => FreeIn (TypeBase shape u) where + freeIn (Array _ shape _) = freeIn shape + freeIn (Mem size _) = freeIn size + freeIn (Prim _) = mempty + +instance FreeIn attr => FreeIn (ParamT attr) where + freeIn (Param _ attr) = freeIn attr + +instance FreeIn attr => FreeIn (PatElemT attr) where + freeIn (PatElem _ attr) = freeIn attr + +instance FreeIn (LParamAttr lore) => FreeIn (LoopForm lore) where + freeIn (ForLoop _ _ bound loop_vars) = freeIn bound <> freeIn loop_vars + freeIn (WhileLoop cond) = freeIn cond + +instance FreeIn d => FreeIn (DimChange d) where + freeIn = Data.Foldable.foldMap freeIn + +instance FreeIn d => FreeIn (DimIndex d) where + freeIn = Data.Foldable.foldMap freeIn + +instance FreeIn attr => FreeIn (PatternT attr) where + freeIn (Pattern context values) = + mconcat (map freeIn $ context ++ values) `S.difference` bound_here + where bound_here = S.fromList $ map patElemName $ context ++ values + +instance FreeIn Certificates where + freeIn (Certificates cs) = freeIn cs + +instance FreeIn attr => FreeIn (StmAux attr) where + freeIn (StmAux cs attr) = freeIn cs <> freeIn attr + +instance FreeIn a => FreeIn (IfAttr a) where + freeIn (IfAttr r _) = freeIn r + +-- | Either return precomputed free names stored in the attribute, or +-- the freshly computed names. Relies on lazy evaluation to avoid the +-- work. +class FreeIn attr => FreeAttr attr where + precomputed :: attr -> Names -> Names + precomputed _ = id + +instance FreeAttr () where + +instance (FreeAttr a, FreeIn b) => FreeAttr (a,b) where + precomputed (a,_) = precomputed a + +instance FreeAttr a => FreeAttr [a] where + precomputed [] = id + precomputed (a:_) = precomputed a + +instance FreeAttr a => FreeAttr (Maybe a) where + precomputed Nothing = id + precomputed (Just a) = precomputed a + +-- | The names bound by the bindings immediately in a 'Body'. +boundInBody :: Body lore -> Names +boundInBody = boundByStms . bodyStms + +-- | The names bound by a binding. +boundByStm :: Stm lore -> Names +boundByStm = S.fromList . patternNames . stmPattern + +-- | The names bound by the bindings. +boundByStms :: Stms lore -> Names +boundByStms = fold . fmap boundByStm + +-- | The names of the lambda parameters plus the index parameter. +boundByLambda :: Lambda lore -> [VName] +boundByLambda lam = map paramName (lambdaParams lam) diff --git a/src/Futhark/Representation/AST/Attributes/Patterns.hs b/src/Futhark/Representation/AST/Attributes/Patterns.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Patterns.hs @@ -0,0 +1,111 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | Inspecing and modifying 'Pattern's, function parameters and +-- pattern elements. +module Futhark.Representation.AST.Attributes.Patterns + ( + -- * Function parameters + paramIdent + , paramType + , paramDeclType + -- * Pattern elements + , patElemIdent + , patElemType + , setPatElemLore + , patternElements + , patternIdents + , patternContextIdents + , patternValueIdents + , patternNames + , patternValueNames + , patternContextNames + , patternTypes + , patternValueTypes + , patternExtTypes + , patternSize + -- * Pattern construction + , basicPattern + ) + where + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes.Types + (existentialiseExtTypes, staticShapes, Typed(..), DeclTyped(..)) + +-- | The 'Type' of a parameter. +paramType :: Typed attr => ParamT attr -> Type +paramType = typeOf + +-- | The 'DeclType' of a parameter. +paramDeclType :: DeclTyped attr => ParamT attr -> DeclType +paramDeclType = declTypeOf + +-- | An 'Ident' corresponding to a parameter. +paramIdent :: Typed attr => ParamT attr -> Ident +paramIdent param = Ident (paramName param) (typeOf param) + +-- | An 'Ident' corresponding to a pattern element. +patElemIdent :: Typed attr => PatElemT attr -> Ident +patElemIdent pelem = Ident (patElemName pelem) (typeOf pelem) + +-- | The type of a name bound by a 'PatElem'. +patElemType :: Typed attr => PatElemT attr -> Type +patElemType = typeOf + +-- | Set the lore of a 'PatElem'. +setPatElemLore :: PatElemT oldattr -> newattr -> PatElemT newattr +setPatElemLore pe x = fmap (const x) pe + +-- | All pattern elements in the pattern - context first, then values. +patternElements :: PatternT attr -> [PatElemT attr] +patternElements pat = patternContextElements pat ++ patternValueElements pat + +-- | Return a list of the 'Ident's bound by the 'Pattern'. +patternIdents :: Typed attr => PatternT attr -> [Ident] +patternIdents pat = patternContextIdents pat ++ patternValueIdents pat + +-- | Return a list of the context 'Ident's bound by the 'Pattern'. +patternContextIdents :: Typed attr => PatternT attr -> [Ident] +patternContextIdents = map patElemIdent . patternContextElements + +-- | Return a list of the value 'Ident's bound by the 'Pattern'. +patternValueIdents :: Typed attr => PatternT attr -> [Ident] +patternValueIdents = map patElemIdent . patternValueElements + +-- | Return a list of the 'Name's bound by the 'Pattern'. +patternNames :: PatternT attr -> [VName] +patternNames = map patElemName . patternElements + +-- | Return a list of the 'Name's bound by the context part of the 'Pattern'. +patternContextNames :: PatternT attr -> [VName] +patternContextNames = map patElemName . patternContextElements + +-- | Return a list of the 'Name's bound by the value part of the 'Pattern'. +patternValueNames :: PatternT attr -> [VName] +patternValueNames = map patElemName . patternValueElements + +-- | Return a list of the 'types's bound by the 'Pattern'. +patternTypes :: Typed attr => PatternT attr -> [Type] +patternTypes = map identType . patternIdents + +-- | Return a list of the 'Types's bound by the value part of the 'Pattern'. +patternValueTypes :: Typed attr => PatternT attr -> [Type] +patternValueTypes = map identType . patternValueIdents + +-- | Return a list of the 'ExtTypes's bound by the value part of the +-- 'Pattern', with existentials where the sizes are part of the +-- context part of the 'Pattern'. +patternExtTypes :: Typed attr => PatternT attr -> [ExtType] +patternExtTypes pat = + existentialiseExtTypes (patternContextNames pat) + (staticShapes (patternValueTypes pat)) + +-- | Return the number of names bound by the 'Pattern'. +patternSize :: PatternT attr -> Int +patternSize (Pattern context values) = length context + length values + +-- | Create a pattern using 'Type' as the attribute. +basicPattern :: [Ident] -> [Ident] -> PatternT Type +basicPattern context values = + Pattern (map patElem context) (map patElem values) + where patElem (Ident name t) = PatElem name t diff --git a/src/Futhark/Representation/AST/Attributes/Ranges.hs b/src/Futhark/Representation/AST/Attributes/Ranges.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Ranges.hs @@ -0,0 +1,274 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE FlexibleInstances #-} +-- | Utility declarations for performing range analysis. +module Futhark.Representation.AST.Attributes.Ranges + ( Bound + , KnownBound (..) + , boundToScalExp + , minimumBound + , maximumBound + , Range + , unknownRange + , ScalExpRange + , Ranged + , RangeOf (..) + , RangesOf (..) + , expRanges + , RangedOp (..) + , CanBeRanged (..) + ) + where + +import Data.Monoid ((<>)) +import qualified Data.Set as S +import qualified Data.Map.Strict as M + +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Syntax +import qualified Futhark.Analysis.ScalExp as SE +import qualified Futhark.Analysis.AlgSimplify as AS +import Futhark.Transform.Substitute +import Futhark.Transform.Rename +import qualified Futhark.Util.Pretty as PP + +-- | A known bound on a value. +data KnownBound = VarBound VName + -- ^ Has the same bounds as this variable. VERY + -- IMPORTANT: this variable may be an array, so it + -- cannot be immediately translated to a 'ScalExp'. + | MinimumBound KnownBound KnownBound + -- ^ Bounded by the minimum of these two bounds. + | MaximumBound KnownBound KnownBound + -- ^ Bounded by the maximum of these two bounds. + | ScalarBound SE.ScalExp + -- ^ Bounded by this scalar expression. + deriving (Eq, Ord, Show) + +instance Substitute KnownBound where + substituteNames substs (VarBound name) = + VarBound $ substituteNames substs name + substituteNames substs (MinimumBound b1 b2) = + MinimumBound (substituteNames substs b1) (substituteNames substs b2) + substituteNames substs (MaximumBound b1 b2) = + MaximumBound (substituteNames substs b1) (substituteNames substs b2) + substituteNames substs (ScalarBound se) = + ScalarBound $ substituteNames substs se + +instance Rename KnownBound where + rename = substituteRename + +instance FreeIn KnownBound where + freeIn (VarBound v) = freeIn v + freeIn (MinimumBound b1 b2) = freeIn b1 <> freeIn b2 + freeIn (MaximumBound b1 b2) = freeIn b1 <> freeIn b2 + freeIn (ScalarBound e) = freeIn e + +instance FreeAttr KnownBound where + precomputed _ = id + +instance PP.Pretty KnownBound where + ppr (VarBound v) = + PP.text "variable " <> PP.ppr v + ppr (MinimumBound b1 b2) = + PP.text "min" <> PP.parens (PP.ppr b1 <> PP.comma PP.<+> PP.ppr b2) + ppr (MaximumBound b1 b2) = + PP.text "max" <> PP.parens (PP.ppr b1 <> PP.comma PP.<+> PP.ppr b2) + ppr (ScalarBound e) = + PP.ppr e + +-- | Convert the bound to a scalar expression if possible. This is +-- possible for all bounds that do not contain 'VarBound's. +boundToScalExp :: KnownBound -> Maybe SE.ScalExp +boundToScalExp (VarBound _) = Nothing +boundToScalExp (ScalarBound se) = Just se +boundToScalExp (MinimumBound b1 b2) = do + b1' <- boundToScalExp b1 + b2' <- boundToScalExp b2 + return $ SE.MaxMin True [b1', b2'] +boundToScalExp (MaximumBound b1 b2) = do + b1' <- boundToScalExp b1 + b2' <- boundToScalExp b2 + return $ SE.MaxMin False [b1', b2'] + +-- | A possibly undefined bound on a value. +type Bound = Maybe KnownBound + +-- | Construct a 'MinimumBound' from two possibly known bounds. The +-- resulting bound will be unknown unless both of the given 'Bound's +-- are known. This may seem counterintuitive, but it actually makes +-- sense when you consider the task of combining the lower bounds for +-- two different flows of execution (like an @if@ expression). If we +-- only have knowledge about one of the branches, this means that we +-- have no useful information about the combined lower bound, as the +-- other branch may take any value. +minimumBound :: Bound -> Bound -> Bound +minimumBound (Just x) (Just y) = Just $ MinimumBound x y +minimumBound _ _ = Nothing + +-- | Like 'minimumBound', but constructs a 'MaximumBound'. +maximumBound :: Bound -> Bound -> Bound +maximumBound (Just x) (Just y) = Just $ MaximumBound x y +maximumBound _ _ = Nothing + +-- | Upper and lower bound, both inclusive. +type Range = (Bound, Bound) + +-- | A range in which both upper and lower bounds are 'Nothing. +unknownRange :: Range +unknownRange = (Nothing, Nothing) + +-- | The range as a pair of scalar expressions. +type ScalExpRange = (Maybe SE.ScalExp, Maybe SE.ScalExp) + +-- | The lore has embedded range information. Note that it may not be +-- up to date, unless whatever maintains the syntax tree is careful. +type Ranged lore = (Attributes lore, + RangedOp (Op lore), + RangeOf (LetAttr lore), + RangesOf (BodyAttr lore)) + +-- | Something that contains range information. +class RangeOf a where + -- | The range of the argument element. + rangeOf :: a -> Range + +instance RangeOf Range where + rangeOf = id + +instance RangeOf attr => RangeOf (PatElemT attr) where + rangeOf = rangeOf . patElemAttr + +instance RangeOf SubExp where + rangeOf se = (Just lower, Just upper) + where (lower, upper) = subExpKnownRange se + +-- | Something that contains range information for several things, +-- most notably 'Body' or 'Pattern'. +class RangesOf a where + -- | The ranges of the argument. + rangesOf :: a -> [Range] + +instance RangeOf a => RangesOf [a] where + rangesOf = map rangeOf + +instance RangeOf attr => RangesOf (PatternT attr) where + rangesOf = map rangeOf . patternElements + +instance Ranged lore => RangesOf (Body lore) where + rangesOf = rangesOf . bodyAttr + +subExpKnownRange :: SubExp -> (KnownBound, KnownBound) +subExpKnownRange (Var v) = + (VarBound v, + VarBound v) +subExpKnownRange (Constant val) = + (ScalarBound $ SE.Val val, + ScalarBound $ SE.Val val) + +-- | The range of a scalar expression. +scalExpRange :: SE.ScalExp -> Range +scalExpRange se = + (Just $ ScalarBound se, Just $ ScalarBound se) + +primOpRanges :: BasicOp lore -> [Range] +primOpRanges (SubExp se) = + [rangeOf se] + +primOpRanges (BinOp (Add t) x y) = + [scalExpRange $ SE.SPlus (SE.subExpToScalExp x $ IntType t) (SE.subExpToScalExp y $ IntType t)] +primOpRanges (BinOp (Sub t) x y) = + [scalExpRange $ SE.SMinus (SE.subExpToScalExp x $ IntType t) (SE.subExpToScalExp y $ IntType t)] +primOpRanges (BinOp (Mul t) x y) = + [scalExpRange $ SE.STimes (SE.subExpToScalExp x $ IntType t) (SE.subExpToScalExp y $ IntType t)] +primOpRanges (BinOp (SDiv t) x y) = + [scalExpRange $ SE.SDiv (SE.subExpToScalExp x $ IntType t) (SE.subExpToScalExp y $ IntType t)] + +primOpRanges (ConvOp (SExt from to) x) + | from < to = [rangeOf x] + +primOpRanges (Iota n x s Int32) = + [(Just $ ScalarBound x', + Just $ ScalarBound $ x' + (n' - 1) * s')] + where n' = case n of + Var v -> SE.Id v $ IntType Int32 + Constant val -> SE.Val val + x' = case x of + Var v -> SE.Id v $ IntType Int32 + Constant val -> SE.Val val + s' = case s of + Var v -> SE.Id v $ IntType Int32 + Constant val -> SE.Val val +primOpRanges (Replicate _ v) = + [rangeOf v] +primOpRanges (Rearrange _ v) = + [rangeOf $ Var v] +primOpRanges (Copy se) = + [rangeOf $ Var se] +primOpRanges (Index v _) = + [rangeOf $ Var v] +primOpRanges (Partition n _ arr) = + replicate n unknownRange ++ map (rangeOf . Var) arr +primOpRanges (ArrayLit (e:es) _) = + [(Just lower, Just upper)] + where (e_lower, e_upper) = subExpKnownRange e + (es_lower, es_upper) = unzip $ map subExpKnownRange es + lower = foldl MinimumBound e_lower es_lower + upper = foldl MaximumBound e_upper es_upper +primOpRanges _ = + [unknownRange] + +-- | Ranges of the value parts of the expression. +expRanges :: Ranged lore => + Exp lore -> [Range] +expRanges (BasicOp op) = + primOpRanges op +expRanges (If _ tbranch fbranch _) = + zip + (zipWith minimumBound t_lower f_lower) + (zipWith maximumBound t_upper f_upper) + where (t_lower, t_upper) = unzip $ rangesOf tbranch + (f_lower, f_upper) = unzip $ rangesOf fbranch +expRanges (DoLoop ctxmerge valmerge (ForLoop i Int32 iterations _) body) = + zipWith returnedRange valmerge $ rangesOf body + where bound_in_loop = + S.fromList $ i : map (paramName . fst) (ctxmerge++valmerge) ++ + concatMap (patternNames . stmPattern) (bodyStms body) + + returnedRange mergeparam (lower, upper) = + (returnedBound mergeparam lower, + returnedBound mergeparam upper) + + returnedBound (param, mergeinit) (Just bound) + | paramType param == Prim (IntType Int32), + Just bound' <- boundToScalExp bound, + let se_diff = + AS.simplify (SE.SMinus (SE.Id (paramName param) $ IntType Int32) bound') M.empty, + S.null $ S.intersection bound_in_loop $ freeIn se_diff = + Just $ ScalarBound $ SE.SPlus (SE.subExpToScalExp mergeinit $ IntType Int32) $ + SE.STimes se_diff $ SE.MaxMin False + [SE.subExpToScalExp iterations $ IntType Int32, 0] + returnedBound _ _ = Nothing +expRanges (Op ranges) = opRanges ranges +expRanges e = + replicate (expExtTypeSize e) unknownRange + +class IsOp op => RangedOp op where + opRanges :: op -> [Range] + +instance RangedOp () where + opRanges () = [] + +class RangedOp (OpWithRanges op) => + CanBeRanged op where + type OpWithRanges op :: * + removeOpRanges :: OpWithRanges op -> op + addOpRanges :: op -> OpWithRanges op + +instance CanBeRanged () where + type OpWithRanges () = () + removeOpRanges = id + addOpRanges = id diff --git a/src/Futhark/Representation/AST/Attributes/Rearrange.hs b/src/Futhark/Representation/AST/Attributes/Rearrange.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Rearrange.hs @@ -0,0 +1,106 @@ +module Futhark.Representation.AST.Attributes.Rearrange + ( rearrangeShape + , rearrangeInverse + , rearrangeReach + , rearrangeCompose + , isPermutationOf + , transposeIndex + , isMapTranspose + ) where + +import Data.List + +import Futhark.Util + +-- | Calculate the given permutation of the list. It is an error if +-- the permutation goes out of bounds. +rearrangeShape :: [Int] -> [a] -> [a] +rearrangeShape perm l = map pick perm + where pick i + | 0 <= i, i < n = l!!i + | otherwise = + error $ show perm ++ " is not a valid permutation for input." + n = length l + +-- | Produce the inverse permutation. +rearrangeInverse :: [Int] -> [Int] +rearrangeInverse perm = map snd $ sortOn fst $ zip perm [0..] + +-- | Return the first dimension not affected by the permutation. For +-- example, the permutation @[1,0,2]@ would return @2@. +rearrangeReach :: [Int] -> Int +rearrangeReach perm = case dropWhile (uncurry (/=)) $ zip (tails perm) (tails [0..n-1]) of + [] -> n + 1 + (perm',_):_ -> n - length perm' + where n = length perm + +-- | Compose two permutations, with the second given permutation being +-- applied first. +rearrangeCompose :: [Int] -> [Int] -> [Int] +rearrangeCompose = rearrangeShape + +-- | Check whether the first list is a permutation of the second, and +-- if so, return the permutation. This will also find identity +-- permutations (i.e. the lists are the same) The implementation is +-- naive and slow. +isPermutationOf :: Eq a => [a] -> [a] -> Maybe [Int] +isPermutationOf l1 l2 = + case mapAccumLM (pick 0) (map Just l2) l1 of + Just (l2', perm) + | all (==Nothing) l2' -> Just perm + _ -> Nothing + where pick :: Eq a => Int -> [Maybe a] -> a -> Maybe ([Maybe a], Int) + pick _ [] _ = Nothing + pick i (x:xs) y + | Just y == x = Just (Nothing : xs, i) + | otherwise = do + (xs', v) <- pick (i+1) xs y + return (x : xs', v) + +-- | If @l@ is an index into the array @a@, then @transposeIndex k n +-- l@ is an index to the same element in the array @transposeArray k n +-- a@. +transposeIndex :: Int -> Int -> [a] -> [a] +transposeIndex k n l + | k + n >= length l = + let n' = ((k + n) `mod` length l)-k + in transposeIndex k n' l + | n < 0, + (pre,needle:end) <- splitAt k l, + (beg,mid) <- splitAt (length pre+n) pre = + beg ++ [needle] ++ mid ++ end + | (beg,needle:post) <- splitAt k l, + (mid,end) <- splitAt n post = + beg ++ mid ++ [needle] ++ end + | otherwise = l + +-- | If @perm@ is conceptually a map of a transposition, +-- @isMapTranspose perm@ returns the number of dimensions being mapped +-- and the number dimension being transposed. For example, we can +-- consider the permutation @[0,1,4,5,2,3]@ as a map of a transpose, +-- by considering dimensions @[0,1]@, @[4,5]@, and @[2,3]@ as single +-- dimensions each. +-- +-- If the input is not a valid permutation, then the result is +-- undefined. +isMapTranspose :: [Int] -> Maybe (Int, Int, Int) +isMapTranspose perm + | posttrans == [length mapped..length mapped+length posttrans-1], + not $ null pretrans, not $ null posttrans = + Just (length mapped, length pretrans, length posttrans) + | otherwise = + Nothing + where (mapped, notmapped) = findIncreasingFrom 0 perm + (pretrans, posttrans) = findTransposed notmapped + + findIncreasingFrom x (i:is) + | i == x = + let (js, ps) = findIncreasingFrom (x+1) is + in (i : js, ps) + findIncreasingFrom _ is = + ([], is) + + findTransposed [] = + ([], []) + findTransposed (i:is) = + findIncreasingFrom i (i:is) diff --git a/src/Futhark/Representation/AST/Attributes/Reshape.hs b/src/Futhark/Representation/AST/Attributes/Reshape.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Reshape.hs @@ -0,0 +1,181 @@ +module Futhark.Representation.AST.Attributes.Reshape + ( + -- * Basic tools + newDim + , newDims + , newShape + + -- * Construction + , shapeCoerce + , repeatShapes + + -- * Execution + , reshapeOuter + , reshapeInner + , repeatDims + + -- * Inspection + , shapeCoercion + + -- * Simplification + , fuseReshape + , fuseReshapes + , informReshape + + -- * Shape calculations + , reshapeIndex + , flattenIndex + , unflattenIndex + , sliceSizes + ) + where + +import Data.Foldable + +import Prelude hiding (sum, product, quot) + +import Futhark.Representation.AST.Attributes.Types +import Futhark.Representation.AST.Syntax +import Futhark.Util.IntegralExp + +-- | The new dimension. +newDim :: DimChange d -> d +newDim (DimCoercion se) = se +newDim (DimNew se) = se + +-- | The new dimensions resulting from a reshape operation. +newDims :: ShapeChange d -> [d] +newDims = map newDim + +-- | The new shape resulting from a reshape operation. +newShape :: ShapeChange SubExp -> Shape +newShape = Shape . newDims + +-- ^ Construct a 'Reshape' where all dimension changes are +-- 'DimCoercion's. +shapeCoerce :: [SubExp] -> VName -> Exp lore +shapeCoerce newdims arr = + BasicOp $ Reshape (map DimCoercion newdims) arr + +-- | Construct a pair suitable for a 'Repeat'. +repeatShapes :: [Shape] -> Type -> ([Shape], Shape) +repeatShapes shapes t = + case splitAt t_rank shapes of + (outer_shapes, [inner_shape]) -> + (outer_shapes, inner_shape) + _ -> + (shapes ++ replicate (length shapes - t_rank) (Shape []), Shape []) + where t_rank = arrayRank t + +-- | Modify the shape of an array type as 'Repeat' would do +repeatDims :: [Shape] -> Shape -> Type -> Type +repeatDims shape innershape = modifyArrayShape repeatDims' + where repeatDims' (Shape ds) = + Shape $ concat (zipWith (++) (map shapeDims shape) (map pure ds)) ++ + shapeDims innershape + +-- | @reshapeOuter newshape n oldshape@ returns a 'Reshape' expression +-- that replaces the outer @n@ dimensions of @oldshape@ with @newshape@. +reshapeOuter :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp +reshapeOuter newshape n oldshape = + newshape ++ map coercion_or_new (drop n (shapeDims oldshape)) + where coercion_or_new + | length newshape == n = DimCoercion + | otherwise = DimNew + +-- | @reshapeInner newshape n oldshape@ returns a 'Reshape' expression +-- that replaces the inner @m-n@ dimensions (where @m@ is the rank of +-- @oldshape@) of @src@ with @newshape@. +reshapeInner :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp +reshapeInner newshape n oldshape = + map coercion_or_new (take n (shapeDims oldshape)) ++ newshape + where coercion_or_new + | length newshape == m-n = DimCoercion + | otherwise = DimNew + m = shapeRank oldshape + +-- | If the shape change is nothing but shape coercions, return the new dimensions. Otherwise, return +-- 'Nothing'. +shapeCoercion :: ShapeChange d -> Maybe [d] +shapeCoercion = mapM dimCoercion + where dimCoercion (DimCoercion d) = Just d + dimCoercion (DimNew _) = Nothing + +-- | @fuseReshape s1 s2@ creates a new 'ShapeChange' that is +-- semantically the same as first applying @s1@ and then @s2@. This +-- may take advantage of properties of 'DimCoercion' versus 'DimNew' +-- to preserve information. +fuseReshape :: Eq d => ShapeChange d -> ShapeChange d -> ShapeChange d +fuseReshape s1 s2 + | length s1 == length s2 = + zipWith comb s1 s2 + where comb (DimNew _) (DimCoercion d2) = + DimNew d2 + comb (DimCoercion d1) (DimNew d2) + | d1 == d2 = DimCoercion d2 + | otherwise = DimNew d2 + comb _ d2 = + d2 +-- TODO: intelligently handle case where s1 is a prefix of s2. +fuseReshape _ s2 = s2 + +-- | @fuseReshapes s ss@ creates a fused 'ShapeChange' that is +-- logically the same as first applying @s@ and then the changes in +-- @ss@ from left to right. +fuseReshapes :: (Eq d, Data.Foldable.Foldable t) => + ShapeChange d -> t (ShapeChange d) -> ShapeChange d +fuseReshapes = Data.Foldable.foldl fuseReshape + +-- | Given concrete information about the shape of the source array, +-- convert some 'DimNew's into 'DimCoercion's. +informReshape :: Eq d => [d] -> ShapeChange d -> ShapeChange d +informReshape shape sc + | length shape == length sc = + zipWith inform shape sc + where inform d1 (DimNew d2) + | d1 == d2 = DimCoercion d2 + inform _ dc = + dc +informReshape _ sc = sc + +-- | @reshapeIndex to_dims from_dims is@ transforms the index list +-- @is@ (which is into an array of shape @from_dims@) into an index +-- list @is'@, which is into an array of shape @to_dims@. @is@ must +-- have the same length as @from_dims@, and @is'@ will have the same +-- length as @to_dims@. +reshapeIndex :: IntegralExp num => + [num] -> [num] -> [num] -> [num] +reshapeIndex to_dims from_dims is = + unflattenIndex to_dims $ flattenIndex from_dims is + +-- | @unflattenIndex dims i@ computes a list of indices into an array +-- with dimension @dims@ given the flat index @i@. The resulting list +-- will have the same size as @dims@. +unflattenIndex :: IntegralExp num => + [num] -> num -> [num] +unflattenIndex = unflattenIndexFromSlices . drop 1 . sliceSizes + +unflattenIndexFromSlices :: IntegralExp num => + [num] -> num -> [num] +unflattenIndexFromSlices [] _ = [] +unflattenIndexFromSlices (size : slices) i = + (i `quot` size) : unflattenIndexFromSlices slices (i - (i `quot` size) * size) + +-- | @flattenIndex dims is@ computes the flat index of @is@ into an +-- array with dimensions @dims@. The length of @dims@ and @is@ must +-- be the same. +flattenIndex :: IntegralExp num => + [num] -> [num] -> num +flattenIndex dims is = + sum $ zipWith (*) is slicesizes + where slicesizes = drop 1 $ sliceSizes dims + +-- | Given a length @n@ list of dimensions @dims@, @sizeSizes dims@ +-- will compute a length @n+1@ list of the size of each possible array +-- slice. The first element of this list will be the product of +-- @dims@, and the last element will be 1. +sliceSizes :: IntegralExp num => + [num] -> [num] +sliceSizes [] = [1] +sliceSizes (n:ns) = + product (n : ns) : sliceSizes ns diff --git a/src/Futhark/Representation/AST/Attributes/Scope.hs b/src/Futhark/Representation/AST/Attributes/Scope.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Scope.hs @@ -0,0 +1,219 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE StandaloneDeriving #-} +-- | This module defines the concept of a type environment as a +-- mapping from variable names to 'Type's. Convenience facilities are +-- also provided to communicate that some monad or applicative functor +-- maintains type information. +module Futhark.Representation.AST.Attributes.Scope + ( HasScope (..) + , NameInfo (..) + , LocalScope (..) + , Scope + , Scoped(..) + , inScopeOf + , scopeOfLParams + , scopeOfFParams + , scopeOfPattern + , scopeOfPatElem + + , SameScope + , castScope + , castNameInfo + + -- * Extended type environment + , ExtendedScope + , extendedScope + ) where + +import Control.Monad.Except +import Control.Monad.Reader +import qualified Control.Monad.RWS.Strict +import qualified Control.Monad.RWS.Lazy +import Data.Foldable +import qualified Data.Map.Strict as M + +import Futhark.Representation.AST.Annotations +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes.Types +import Futhark.Representation.AST.Attributes.Patterns +import Futhark.Representation.AST.Pretty () + +-- | How some name in scope was bound. +data NameInfo lore = LetInfo (LetAttr lore) + | FParamInfo (FParamAttr lore) + | LParamInfo (LParamAttr lore) + | IndexInfo IntType + +deriving instance Annotations lore => Show (NameInfo lore) + +instance Annotations lore => Typed (NameInfo lore) where + typeOf (LetInfo attr) = typeOf attr + typeOf (FParamInfo attr) = typeOf attr + typeOf (LParamInfo attr) = typeOf attr + typeOf (IndexInfo it) = Prim $ IntType it + +-- | A scope is a mapping from variable names to information about +-- that name. +type Scope lore = M.Map VName (NameInfo lore) + +-- | 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, Annotations lore) => HasScope lore m | m -> lore where + -- | Return the type of the given variable, or fail if it is not in + -- the type environment. + lookupType :: VName -> m Type + lookupType = fmap typeOf . lookupInfo + + -- | Return the info of the given variable, or fail if it is not in + -- the type environment. + lookupInfo :: VName -> m (NameInfo lore) + lookupInfo name = + asksScope (M.findWithDefault notFound name) + where notFound = + error $ "Scope.lookupInfo: Name " ++ pretty name ++ + " not found in type environment." + + -- | Return the type environment contained in the applicative + -- functor. + askScope :: m (Scope lore) + + -- | Return the result of applying some function to the type + -- environment. + asksScope :: (Scope lore -> a) -> m a + asksScope f = f <$> askScope + +instance (Applicative m, Monad m, Annotations lore) => + HasScope lore (ReaderT (Scope lore) m) where + askScope = ask + +instance (Monad m, HasScope lore m) => HasScope lore (ExceptT e m) where + askScope = lift askScope + +instance (Applicative m, Monad m, Monoid w, Annotations lore) => + HasScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) where + askScope = ask + +instance (Applicative m, Monad m, Monoid w, Annotations lore) => + HasScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) 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 + -- | 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 + +instance (Monad m, LocalScope lore m) => LocalScope lore (ExceptT e m) where + localScope = mapExceptT . localScope + +instance (Applicative m, Monad m, Annotations lore) => + LocalScope lore (ReaderT (Scope lore) m) where + localScope = local . M.union + +instance (Applicative m, Monad m, Monoid w, Annotations lore) => + LocalScope lore (Control.Monad.RWS.Strict.RWST (Scope lore) w s m) where + localScope = local . M.union + +instance (Applicative m, Monad m, Monoid w, Annotations lore) => + LocalScope lore (Control.Monad.RWS.Lazy.RWST (Scope lore) 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 'Lambda', is is the parameters +-- (including index). +class Scoped lore a | a -> lore where + scopeOf :: a -> Scope lore + +inScopeOf :: (Scoped lore a, LocalScope lore m) => a -> m b -> m b +inScopeOf = localScope . scopeOf + +instance Scoped lore a => Scoped lore [a] where + scopeOf = mconcat . map scopeOf + +instance Scoped lore (Stms lore) where + scopeOf = fold . fmap scopeOf + +instance Scoped lore (Stm lore) where + scopeOf = scopeOfPattern . stmPattern + +instance Scoped lore (FunDef lore) where + scopeOf = scopeOfFParams . funDefParams + +instance Scoped lore (VName, NameInfo lore) where + scopeOf = uncurry M.singleton + +instance Scoped lore (LoopForm lore) where + scopeOf (WhileLoop _) = mempty + scopeOf (ForLoop i it _ xs) = + M.insert i (IndexInfo it) $ scopeOfLParams (map fst xs) + +scopeOfPattern :: LetAttr lore ~ attr => PatternT attr -> Scope lore +scopeOfPattern = + mconcat . map scopeOfPatElem . patternElements + +scopeOfPatElem :: LetAttr lore ~ attr => PatElemT attr -> Scope lore +scopeOfPatElem (PatElem name attr) = M.singleton name $ LetInfo attr + +scopeOfLParams :: LParamAttr lore ~ attr => + [ParamT attr] -> Scope lore +scopeOfLParams = M.fromList . map f + where f param = (paramName param, LParamInfo $ paramAttr param) + +scopeOfFParams :: FParamAttr lore ~ attr => + [ParamT attr] -> Scope lore +scopeOfFParams = M.fromList . map f + where f param = (paramName param, FParamInfo $ paramAttr param) + +instance Scoped lore (Lambda lore) where + scopeOf lam = scopeOfLParams $ lambdaParams lam + +type SameScope lore1 lore2 = (LetAttr lore1 ~ LetAttr lore2, + FParamAttr lore1 ~ FParamAttr lore2, + LParamAttr lore1 ~ LParamAttr lore2) + +-- | If two scopes are really the same, then you can convert one to +-- the other. +castScope :: SameScope fromlore tolore => + Scope fromlore -> Scope tolore +castScope = M.map castNameInfo + +castNameInfo :: SameScope fromlore tolore => + NameInfo fromlore -> NameInfo tolore +castNameInfo (LetInfo attr) = LetInfo attr +castNameInfo (FParamInfo attr) = FParamInfo attr +castNameInfo (LParamInfo attr) = LParamInfo attr +castNameInfo (IndexInfo it) = IndexInfo it + +-- | 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) + deriving (Functor, Applicative, Monad, + MonadReader (Scope lore)) + +instance (HasScope lore m, Monad m) => + HasScope lore (ExtendedScope lore m) where + lookupType name = do + res <- asks $ fmap typeOf . M.lookup name + maybe (ExtendedScope $ lift $ lookupType name) return res + askScope = asks M.union <*> ExtendedScope (lift askScope) + +-- | Run a computation in the extended type environment. +extendedScope :: ExtendedScope lore m a + -> Scope lore + -> m a +extendedScope (ExtendedScope m) = runReaderT m diff --git a/src/Futhark/Representation/AST/Attributes/TypeOf.hs b/src/Futhark/Representation/AST/Attributes/TypeOf.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/TypeOf.hs @@ -0,0 +1,196 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +-- | This module provides facilities for obtaining the types of +-- various Futhark constructs. Typically, you will need to execute +-- these in a context where type information is available as a +-- 'Scope'; usually by using a monad that is an instance of +-- 'HasScope'. The information is returned as a list of 'ExtType' +-- values - one for each of the values the Futhark construct returns. +-- Some constructs (such as subexpressions) can produce only a single +-- value, and their typing functions hence do not return a list. +-- +-- Some representations may have more specialised facilities enabling +-- even more information - for example, +-- "Futhark.Representation.ExplicitMemory" exposes functionality for +-- also obtaining information about the storage location of results. +module Futhark.Representation.AST.Attributes.TypeOf + ( + expExtType + , expExtTypeSize + , subExpType + , bodyExtType + , primOpType + , mapType + , subExpShapeContext + , loopResultContext + , loopExtType + + -- * Return type + , module Futhark.Representation.AST.RetType + -- * Type environment + , module Futhark.Representation.AST.Attributes.Scope + + -- * Extensibility + , TypedOp(..) + ) + where + +import Data.Maybe +import Data.Semigroup ((<>)) +import Data.Foldable +import qualified Data.Set as S + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes.Reshape +import Futhark.Representation.AST.Attributes.Types +import Futhark.Representation.AST.Attributes.Patterns +import Futhark.Representation.AST.Attributes.Constants +import Futhark.Representation.AST.Attributes.Names +import Futhark.Representation.AST.RetType +import Futhark.Representation.AST.Attributes.Scope + +-- | The type of a subexpression. +subExpType :: HasScope t m => SubExp -> m Type +subExpType (Constant val) = pure $ Prim $ primValueType val +subExpType (Var name) = lookupType name + +-- | @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 outersize f = [ arrayOf t (Shape [outersize]) NoUniqueness + | t <- lambdaReturnType f ] + +-- | The type of a primitive operation. +primOpType :: HasScope t m => + BasicOp lore -> m [Type] +primOpType (SubExp se) = + pure <$> subExpType se +primOpType (Opaque se) = + pure <$> subExpType se +primOpType (ArrayLit es rt) = + pure [arrayOf rt (Shape [n]) NoUniqueness] + where n = Constant (value (length es)) +primOpType (BinOp bop _ _) = + pure [Prim $ binOpType bop] +primOpType (UnOp _ x) = + pure <$> subExpType x +primOpType CmpOp{} = + pure [Prim Bool] +primOpType (ConvOp conv _) = + pure [Prim $ snd $ convOpType conv] +primOpType (Index ident slice) = + result <$> lookupType ident + where result t = [Prim (elemType t) `arrayOfShape` shape] + shape = Shape $ mapMaybe dimSize slice + dimSize (DimSlice _ d _) = Just d + dimSize DimFix{} = Nothing +primOpType (Update src _ _) = + pure <$> lookupType src +primOpType (Iota n _ _ et) = + pure [arrayOf (Prim (IntType et)) (Shape [n]) NoUniqueness] +primOpType (Replicate (Shape []) e) = + pure <$> subExpType e +primOpType (Repeat shape innershape v) = + pure . repeatDims shape innershape <$> lookupType v +primOpType (Replicate shape e) = + pure . flip arrayOfShape shape <$> subExpType e +primOpType (Scratch t shape) = + pure [arrayOf (Prim t) (Shape shape) NoUniqueness] +primOpType (Reshape [] e) = + result <$> lookupType e + where result t = [Prim $ elemType t] +primOpType (Reshape shape e) = + result <$> lookupType e + where result t = [t `setArrayShape` newShape shape] +primOpType (Rearrange perm e) = + result <$> lookupType e + where result t = [rearrangeType perm t] +primOpType (Rotate _ e) = + pure <$> lookupType e +primOpType (Concat i x _ ressize) = + result <$> lookupType x + where result xt = [setDimSize i xt ressize] +primOpType (Copy v) = + pure <$> lookupType v +primOpType (Manifest _ v) = + pure <$> lookupType v +primOpType Assert{} = + pure [Prim Cert] +primOpType (Partition n _ arrays) = + result <$> traverse lookupType arrays + where result ts = replicate n (Prim $ IntType Int32) ++ ts + + +-- | The type of an expression. +expExtType :: (HasScope lore m, TypedOp (Op lore)) => + Exp lore -> m [ExtType] +expExtType (Apply _ _ rt _) = pure $ map fromDecl $ retTypeValues rt +expExtType (If _ _ _ rt) = pure $ bodyTypeValues $ ifReturns rt +expExtType (DoLoop ctxmerge valmerge _ _) = + pure $ loopExtType (map (paramIdent . fst) ctxmerge) (map (paramIdent . fst) valmerge) +expExtType (BasicOp op) = staticShapes <$> primOpType op +expExtType (Op op) = opType op + +-- | The number of values returned by an expression. +expExtTypeSize :: (Annotations lore, TypedOp (Op lore)) => + Exp lore -> Int +expExtTypeSize = length . feelBad . expExtType + +-- FIXME, this is a horrible quick hack. +newtype FeelBad lore a = FeelBad { feelBad :: a } + +instance Functor (FeelBad lore) where + fmap f = FeelBad . f . feelBad + +instance Applicative (FeelBad lore) where + pure = FeelBad + f <*> x = FeelBad $ feelBad f $ feelBad x + +instance Annotations lore => HasScope lore (FeelBad lore) where + lookupType = const $ pure $ Prim $ IntType Int32 + askScope = pure mempty + +-- | 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 (Body _ stms res) = + existentialiseExtTypes bound . staticShapes <$> + extendedScope (traverse subExpType res) bndscope + where bndscope = scopeOf stms + boundInLet (Let pat _ _) = S.fromList $ patternNames pat + bound = S.toList $ fold $ fmap boundInLet stms + +-- | Given the return type of a function and the subexpressions +-- returned by that function, return the size context. +subExpShapeContext :: HasScope t m => + [TypeBase ExtShape u] -> [SubExp] -> m [SubExp] +subExpShapeContext rettype ses = + extractShapeContext rettype <$> traverse (fmap arrayDims . subExpType) ses + +-- | A loop returns not only its value merge parameters, but may also +-- have an existential context. Thus, @loopResult ctxmergeparams +-- valmergeparams@ returns those paramters in @ctxmergeparams@ that +-- constitute the returned context. +loopResultContext :: FreeIn attr => [Param attr] -> [Param attr] -> [Param attr] +loopResultContext ctx val = filter usedInValue ctx + where usedInValue = (`S.member` used) . paramName + used = freeIn val <> freeIn ctx + +-- | Given the context and value merge parameters of a Futhark @loop@, +-- produce the return type. +loopExtType :: [Ident] -> [Ident] -> [ExtType] +loopExtType ctx val = + existentialiseExtTypes inaccessible $ staticShapes $ map identType val + where inaccessible = map identName ctx + +-- | Any operation must define an instance of this class, which +-- describes the type of the operation (at the value level). +class TypedOp op where + opType :: HasScope t m => op -> m [ExtType] + +instance TypedOp () where + opType () = pure [] diff --git a/src/Futhark/Representation/AST/Attributes/Types.hs b/src/Futhark/Representation/AST/Attributes/Types.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Attributes/Types.hs @@ -0,0 +1,562 @@ +{-# LANGUAGE FlexibleContexts, FlexibleInstances, TypeSynonymInstances #-} +-- | Functions for inspecting and constructing various types. +module Futhark.Representation.AST.Attributes.Types + ( + rankShaped + , arrayRank + , arrayShape + , modifyArrayShape + , setArrayShape + , existential + , uniqueness + , setUniqueness + , unique + , staticShapes + , staticShapes1 + , primType + + , arrayOf + , arrayOfRow + , arrayOfShape + , setOuterSize + , setDimSize + , setOuterDim + , setDim + , setArrayDims + , setArrayExtDims + , peelArray + , stripArray + , arrayDims + , arrayExtDims + , shapeSize + , arraySize + , arraysSize + , rowType + , elemType + + , transposeType + , rearrangeType + + , diet + + , subtypeOf + , subtypesOf + + , toDecl + , fromDecl + + , extractShapeContext + , shapeContext + , shapeContextSize + , hasStaticShape + , hasStaticShapes + , generaliseExtTypes + , existentialiseExtTypes + , shapeMapping + , shapeMapping' + , shapeExtMapping + + -- * Abbreviations + , int8, int16, int32, int64 + , float32, float64 + + -- * The Typed typeclass + , Typed (..) + , DeclTyped (..) + , ExtTyped (..) + , DeclExtTyped (..) + , SetType (..) + , FixExt (..) + ) + where + +import Control.Monad.State +import Data.Maybe +import Data.Monoid ((<>)) +import Data.List (elemIndex) +import qualified Data.Set as S +import qualified Data.Map.Strict as M + +import Futhark.Representation.AST.Syntax.Core +import Futhark.Representation.AST.Attributes.Constants +import Futhark.Representation.AST.Attributes.Rearrange + +-- | Remove shape information from a type. +rankShaped :: ArrayShape shape => TypeBase shape u -> TypeBase Rank u +rankShaped (Array et sz u) = Array et (Rank $ shapeRank sz) u +rankShaped (Prim et) = Prim et +rankShaped (Mem size space) = Mem size space + +-- | Return the dimensionality of a type. For non-arrays, this is +-- zero. For a one-dimensional array it is one, for a two-dimensional +-- it is two, and so forth. +arrayRank :: ArrayShape shape => TypeBase shape u -> Int +arrayRank = shapeRank . arrayShape + +-- | Return the shape of a type - for non-arrays, this is the +-- 'mempty'. +arrayShape :: ArrayShape shape => TypeBase shape u -> shape +arrayShape (Array _ ds _) = ds +arrayShape _ = mempty + +-- | Modify the shape of an array - for non-arrays, this does nothing. +modifyArrayShape :: ArrayShape newshape => + (oldshape -> newshape) + -> TypeBase oldshape u + -> TypeBase newshape u +modifyArrayShape f (Array t ds u) + | shapeRank ds' == 0 = Prim t + | otherwise = Array t (f ds) u + where ds' = f ds +modifyArrayShape _ (Prim t) = Prim t +modifyArrayShape _ (Mem size space) = Mem size space + +-- | Set the shape of an array. If the given type is not an +-- array, return the type unchanged. +setArrayShape :: ArrayShape newshape => + TypeBase oldshape u + -> newshape + -> TypeBase newshape u +setArrayShape t ds = modifyArrayShape (const ds) t + +-- | True if the given type has a dimension that is existentially sized. +existential :: ExtType -> Bool +existential = any ext . shapeDims . arrayShape + where ext (Ext _) = True + ext (Free _) = False + +-- | Return the uniqueness of a type. +uniqueness :: TypeBase shape Uniqueness -> Uniqueness +uniqueness (Array _ _ u) = u +uniqueness _ = Nonunique + +-- | @unique t@ is 'True' if the type of the argument is unique. +unique :: TypeBase shape Uniqueness -> Bool +unique = (==Unique) . uniqueness + +-- | Set the uniqueness attribute of a type. +setUniqueness :: TypeBase shape Uniqueness + -> Uniqueness + -> TypeBase shape Uniqueness +setUniqueness (Array et dims _) u = Array et dims u +setUniqueness t _ = t + +-- | Convert types with non-existential shapes to types with +-- non-existential shapes. Only the representation is changed, so all +-- the shapes will be 'Free'. +staticShapes :: [TypeBase Shape u] -> [TypeBase ExtShape u] +staticShapes = map staticShapes1 + +-- | As 'staticShapes', but on a single type. +staticShapes1 :: TypeBase Shape u -> TypeBase ExtShape u +staticShapes1 (Prim bt) = + Prim bt +staticShapes1 (Array bt (Shape shape) u) = + Array bt (Shape $ map Free shape) u +staticShapes1 (Mem size space) = + Mem size space + +-- | @arrayOf t s u@ constructs an array type. The convenience +-- compared to using the 'Array' constructor directly is that @t@ can +-- itself be an array. If @t@ is an @n@-dimensional array, and @s@ is +-- a list of length @n@, the resulting type is of an @n+m@ dimensions. +-- The uniqueness of the new array will be @u@, no matter the +-- uniqueness of @t@. If the shape @s@ has rank 0, then the @t@ will +-- be returned, although if it is an array, with the uniqueness +-- changed to @u@. +arrayOf :: ArrayShape shape => + TypeBase shape u_unused -> shape -> u -> TypeBase shape u +arrayOf (Array et size1 _) size2 u = + Array et (size2 <> size1) u +arrayOf (Prim et) s _ + | 0 <- shapeRank s = Prim et +arrayOf (Prim et) size u = + Array et size u +arrayOf Mem{} _ _ = + error "arrayOf Mem" + +-- | Construct an array whose rows are the given type, and the outer +-- size is the given dimension. This is just a convenient wrapper +-- around 'arrayOf'. +arrayOfRow :: ArrayShape (ShapeBase d) => + TypeBase (ShapeBase d) NoUniqueness + -> d + -> TypeBase (ShapeBase d) NoUniqueness +arrayOfRow t size = arrayOf t (Shape [size]) NoUniqueness + +-- | Construct an array whose rows are the given type, and the outer +-- size is the given 'Shape'. This is just a convenient wrapper +-- around 'arrayOf'. +arrayOfShape :: Type -> Shape -> Type +arrayOfShape t shape = arrayOf t shape NoUniqueness + +-- | Set the dimensions of an array. If the given type is not an +-- array, return the type unchanged. +setArrayDims :: TypeBase oldshape u -> [SubExp] -> TypeBase Shape u +setArrayDims t dims = t `setArrayShape` Shape dims + +-- | Set the existential dimensions of an array. If the given type is +-- not an array, return the type unchanged. +setArrayExtDims :: TypeBase oldshape u -> [ExtSize] -> TypeBase ExtShape u +setArrayExtDims t dims = t `setArrayShape` Shape dims + +-- | Replace the size of the outermost dimension of an array. If the +-- given type is not an array, it is returned unchanged. +setOuterSize :: ArrayShape (ShapeBase d) => + TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u +setOuterSize = setDimSize 0 + +-- | Replace the size of the given dimension of an array. If the +-- given type is not an array, it is returned unchanged. +setDimSize :: ArrayShape (ShapeBase d) => + Int -> TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u +setDimSize i t e = t `setArrayShape` setDim i (arrayShape t) e + +-- | Replace the outermost dimension of an array shape. +setOuterDim :: ShapeBase d -> d -> ShapeBase d +setOuterDim = setDim 0 + +-- | Replace the specified dimension of an array shape. +setDim :: Int -> ShapeBase d -> d -> ShapeBase d +setDim i (Shape ds) e = Shape $ take i ds ++ e : drop (i+1) ds + +-- | @peelArray n t@ returns the type resulting from peeling the first +-- @n@ array dimensions from @t@. Returns @Nothing@ if @t@ has less +-- than @n@ dimensions. +peelArray :: ArrayShape shape => + Int -> TypeBase shape u -> Maybe (TypeBase shape u) +peelArray 0 t = Just t +peelArray n (Array et shape u) + | shapeRank shape == n = Just $ Prim et + | shapeRank shape > n = Just $ Array et (stripDims n shape) u +peelArray _ _ = Nothing + +-- | @stripArray n t@ removes the @n@ outermost layers of the array. +-- Essentially, it is the type of indexing an array of type @t@ with +-- @n@ indexes. +stripArray :: ArrayShape shape => Int -> TypeBase shape u -> TypeBase shape u +stripArray n (Array et shape u) + | n < shapeRank shape = Array et (stripDims n shape) u + | otherwise = Prim et +stripArray _ t = t + +-- | Return the size of the given dimension. If the dimension does +-- not exist, the zero constant is returned. +shapeSize :: Int -> Shape -> SubExp +shapeSize i shape = case drop i $ shapeDims shape of + e : _ -> e + [] -> constant (0 :: Int32) + +-- | Return the dimensions of a type - for non-arrays, this is the +-- empty list. +arrayDims :: TypeBase Shape u -> [SubExp] +arrayDims = shapeDims . arrayShape + +-- | Return the existential dimensions of a type - for non-arrays, +-- this is the empty list. +arrayExtDims :: TypeBase ExtShape u -> [ExtSize] +arrayExtDims = shapeDims . arrayShape + +-- | Return the size of the given dimension. If the dimension does +-- not exist, the zero constant is returned. +arraySize :: Int -> TypeBase Shape u -> SubExp +arraySize i = shapeSize i . arrayShape + +-- | Return the size of the given dimension in the first element of +-- the given type list. If the dimension does not exist, or no types +-- are given, the zero constant is returned. +arraysSize :: Int -> [TypeBase Shape u] -> SubExp +arraysSize _ [] = constant (0 :: Int32) +arraysSize i (t:_) = arraySize i t + +-- | Return the immediate row-type of an array. For @[[int]]@, this +-- would be @[int]@. +rowType :: ArrayShape shape => TypeBase shape u -> TypeBase shape u +rowType = stripArray 1 + +-- | A type is a primitive type if it is not an array or memory block. +primType :: TypeBase shape u -> Bool +primType Array{} = False +primType Mem{} = False +primType _ = True + +-- | Returns the bottommost type of an array. For @[[int]]@, this +-- would be @int@. If the given type is not an array, it is returned. +elemType :: TypeBase shape u -> PrimType +elemType (Array t _ _) = t +elemType (Prim t) = t +elemType Mem{} = error "elemType Mem" + +-- | Swap the two outer dimensions of the type. +transposeType :: Type -> Type +transposeType = rearrangeType [1,0] + +-- | Rearrange the dimensions of the type. If the length of the +-- permutation does not match the rank of the type, the permutation +-- will be extended with identity. +rearrangeType :: [Int] -> Type -> Type +rearrangeType perm t = + t `setArrayShape` Shape (rearrangeShape perm' $ arrayDims t) + where perm' = perm ++ [length perm .. arrayRank t - 1] + +-- | @diet t@ returns a description of how a function parameter of +-- type @t@ might consume its argument. +diet :: TypeBase shape Uniqueness -> Diet +diet (Prim _) = Observe +diet (Array _ _ Unique) = Consume +diet (Array _ _ Nonunique) = Observe +diet Mem{} = Observe + +-- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal to +-- @y@), meaning @x@ is valid whenever @y@ is. +subtypeOf :: (Ord u, ArrayShape shape) => + TypeBase shape u + -> TypeBase shape u + -> Bool +subtypeOf (Array t1 shape1 u1) (Array t2 shape2 u2) = + u2 <= u1 && + t1 == t2 && + shape1 `subShapeOf` shape2 +subtypeOf (Prim t1) (Prim t2) = t1 == t2 +subtypeOf (Mem _ space1) (Mem _ space2) = space1 == space2 +subtypeOf _ _ = False + +-- | @xs \`subtypesOf\` ys@ is true if @xs@ is the same size as @ys@, +-- and each element in @xs@ is a subtype of the corresponding element +-- in @ys@.. +subtypesOf :: (Ord u, ArrayShape shape) => + [TypeBase shape u] + -> [TypeBase shape u] + -> Bool +subtypesOf xs ys = length xs == length ys && + and (zipWith subtypeOf xs ys) + +toDecl :: TypeBase shape NoUniqueness + -> Uniqueness + -> TypeBase shape Uniqueness +toDecl (Prim bt) _ = Prim bt +toDecl (Array et shape _) u = Array et shape u +toDecl (Mem size space) _ = Mem size space + +fromDecl :: TypeBase shape Uniqueness + -> TypeBase shape NoUniqueness +fromDecl (Prim bt) = Prim bt +fromDecl (Array et shape _) = Array et shape NoUniqueness +fromDecl (Mem size space) = Mem size space + +-- | Given the existential return type of a function, and the shapes +-- of the values returned by the function, return the existential +-- shape context. That is, those sizes that are existential in the +-- return type. +extractShapeContext :: [TypeBase ExtShape u] -> [[a]] -> [a] +extractShapeContext ts shapes = + evalState (concat <$> zipWithM extract ts shapes) S.empty + where extract t shape = + catMaybes <$> zipWithM extract' (shapeDims $ arrayShape t) shape + extract' (Ext x) v = do + seen <- gets $ S.member x + if seen then return Nothing + else do modify $ S.insert x + return $ Just v + extract' (Free _) _ = return Nothing + +-- | The set of identifiers used for the shape context in the given +-- 'ExtType's. +shapeContext :: [TypeBase ExtShape u] -> S.Set Int +shapeContext = S.fromList + . concatMap (mapMaybe ext . shapeDims . arrayShape) + where ext (Ext x) = Just x + ext (Free _) = Nothing + +-- | The size of the set that would be returned by 'shapeContext'. +shapeContextSize :: [ExtType] -> Int +shapeContextSize = S.size . shapeContext + +-- | If all dimensions of the given 'RetType' are statically known, +-- return the corresponding list of 'Type'. +hasStaticShape :: ExtType -> Maybe Type +hasStaticShape (Prim bt) = + Just $ Prim bt +hasStaticShape (Mem size space) = + Just $ Mem size space +hasStaticShape (Array bt (Shape shape) u) = + Array bt <$> (Shape <$> mapM isFree shape) <*> pure u + where isFree (Free s) = Just s + isFree (Ext _) = Nothing + +hasStaticShapes :: [ExtType] -> Maybe [Type] +hasStaticShapes = mapM hasStaticShape + +-- | Given two lists of 'ExtType's of the same length, return a list +-- of 'ExtType's that is a subtype (as per 'isSubtypeOf') of the two +-- operands. +generaliseExtTypes :: [TypeBase ExtShape u] + -> [TypeBase ExtShape u] + -> [TypeBase ExtShape u] +generaliseExtTypes rt1 rt2 = + evalState (zipWithM unifyExtShapes rt1 rt2) (0, M.empty) + where unifyExtShapes t1 t2 = + setArrayShape t1 . Shape <$> + zipWithM unifyExtDims + (shapeDims $ arrayShape t1) + (shapeDims $ arrayShape t2) + unifyExtDims (Free se1) (Free se2) + | se1 == se2 = return $ Free se1 -- Arbitrary + | otherwise = do (n,m) <- get + put (n + 1, m) + return $ Ext n + unifyExtDims (Ext x) (Ext y) + | x == y = Ext <$> (maybe (new x) return =<< + gets (M.lookup x . snd)) + unifyExtDims (Ext x) _ = Ext <$> new x + unifyExtDims _ (Ext x) = Ext <$> new x + new x = do (n,m) <- get + put (n + 1, M.insert x n m) + return n + +-- | Given a list of 'ExtType's and a list of "forbidden" names, +-- modify the dimensions of the 'ExtType's such that they are 'Ext' +-- where they were previously 'Free' with a variable in the set of +-- forbidden names. +existentialiseExtTypes :: [VName] -> [ExtType] -> [ExtType] +existentialiseExtTypes inaccessible = map makeBoundShapesFree + where makeBoundShapesFree = + modifyArrayShape $ fmap checkDim + checkDim (Free (Var v)) + | Just i <- v `elemIndex` inaccessible = + Ext i + checkDim d = d + +-- | In the call @shapeMapping ts1 ts2@, the lists @ts1@ and @ts@ must +-- be of equal length and their corresponding elements have the same +-- types modulo exact dimensions (but matching array rank is +-- important). The result is a mapping from named dimensions of @ts1@ +-- to the corresponding dimension in @ts2@. +-- +-- This function is useful when @ts1@ are the value parameters of some +-- function and @ts2@ are the value arguments, and we need to figure +-- out which shape context to pass. +shapeMapping :: [TypeBase Shape u0] -> [TypeBase Shape u1] -> M.Map VName SubExp +shapeMapping ts = shapeMapping' ts . map arrayDims + +-- | Like @shapeMapping@, but works with explicit dimensions. +shapeMapping' :: [TypeBase Shape u] -> [[a]] -> M.Map VName a +shapeMapping' = dimMapping arrayDims id match + where match Constant{} _ = M.empty + match (Var v) dim = M.singleton v dim + +-- | Like 'shapeMapping', but produces a mapping for the dimensions context. +shapeExtMapping :: [TypeBase ExtShape u] -> [TypeBase Shape u1] -> M.Map Int SubExp +shapeExtMapping = dimMapping arrayExtDims arrayDims match + where match Free{} _ = mempty + match (Ext i) dim = M.singleton i dim + +dimMapping :: Monoid res => + (t1 -> [dim1]) -> (t2 -> [dim2]) -> (dim1 -> dim2 -> res) + -> [t1] -> [t2] + -> res +dimMapping getDims1 getDims2 f ts1 ts2 = + mconcat $ concat $ zipWith (zipWith f) (map getDims1 ts1) (map getDims2 ts2) + +int8 :: PrimType +int8 = IntType Int8 + +int16 :: PrimType +int16 = IntType Int16 + +int32 :: PrimType +int32 = IntType Int32 + +int64 :: PrimType +int64 = IntType Int64 + +float32 :: PrimType +float32 = FloatType Float32 + +float64 :: PrimType +float64 = FloatType Float64 + +-- | Typeclass for things that contain 'Type's. +class Typed t where + typeOf :: t -> Type + +instance Typed Type where + typeOf = id + +instance Typed DeclType where + typeOf = fromDecl + +instance Typed Ident where + typeOf = identType + +instance Typed attr => Typed (Param attr) where + typeOf = typeOf . paramAttr + +instance Typed attr => Typed (PatElemT attr) where + typeOf = typeOf . patElemAttr + +instance Typed b => Typed (a,b) where + typeOf = typeOf . snd + +-- | Typeclass for things that contain 'DeclType's. +class DeclTyped t where + declTypeOf :: t -> DeclType + +instance DeclTyped DeclType where + declTypeOf = id + +instance DeclTyped attr => DeclTyped (Param attr) where + declTypeOf = declTypeOf . paramAttr + +-- | Typeclass for things that contain 'ExtType's. +class FixExt t => ExtTyped t where + extTypeOf :: t -> ExtType + +instance ExtTyped ExtType where + extTypeOf = id + +-- | Typeclass for things that contain 'DeclExtType's. +class FixExt t => DeclExtTyped t where + declExtTypeOf :: t -> DeclExtType + +instance DeclExtTyped DeclExtType where + declExtTypeOf = id + +-- | Typeclass for things whose type can be changed. +class Typed a => SetType a where + setType :: a -> Type -> a + +instance SetType Type where + setType _ t = t + +instance SetType b => SetType (a, b) where + setType (a, b) t = (a, setType b t) + +instance SetType attr => SetType (PatElemT attr) where + setType (PatElem name attr) t = + PatElem name $ setType attr t + +-- | Something with an existential context that can be (partially) +-- fixed. +class FixExt t where + -- | Fix the given existentional variable to the indicated free + -- value. + fixExt :: Int -> SubExp -> t -> t + +instance (FixExt shape, ArrayShape shape) => FixExt (TypeBase shape u) where + fixExt i se = modifyArrayShape $ fixExt i se + +instance FixExt d => FixExt (ShapeBase d) where + fixExt i se = fmap $ fixExt i se + +instance FixExt a => FixExt [a] where + fixExt i se = fmap $ fixExt i se + +instance FixExt ExtSize where + fixExt i se (Ext j) | j > i = Ext $ j - 1 + | j == i = Free se + | otherwise = Ext j + fixExt _ _ (Free x) = Free x + +instance FixExt () where + fixExt _ _ () = () diff --git a/src/Futhark/Representation/AST/Pretty.hs b/src/Futhark/Representation/AST/Pretty.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Pretty.hs @@ -0,0 +1,289 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +-- | Futhark prettyprinter. This module defines 'Pretty' instances +-- for the AST defined in "Futhark.Representation.AST.Syntax", +-- but also a number of convenience functions if you don't want to use +-- the interface from 'Pretty'. +module Futhark.Representation.AST.Pretty + ( prettyTuple + , pretty + , PrettyAnnot (..) + , PrettyLore (..) + , ppTuple' + , bindingAnnotation + ) + where + +import Data.Maybe +import Data.Monoid ((<>)) + +import Futhark.Util.Pretty + +import Futhark.Representation.AST.Attributes.Patterns +import Futhark.Representation.AST.Syntax + +-- | Class for values that may have some prettyprinted annotation. +class PrettyAnnot a where + ppAnnot :: a -> Maybe Doc + +instance PrettyAnnot (PatElemT (TypeBase shape u)) where + ppAnnot = const Nothing + +instance PrettyAnnot (ParamT (TypeBase shape u)) where + ppAnnot = const Nothing + +instance PrettyAnnot () where + ppAnnot = const Nothing + +-- | The class of lores whose annotations can be prettyprinted. +class (Annotations lore, + Pretty (RetType lore), + Pretty (BranchType lore), + Pretty (ParamT (FParamAttr lore)), + Pretty (ParamT (LParamAttr lore)), + Pretty (PatElemT (LetAttr lore)), + PrettyAnnot (PatElem lore), + PrettyAnnot (FParam lore), + PrettyAnnot (LParam lore), + Pretty (Op lore)) => PrettyLore lore where + ppExpLore :: ExpAttr lore -> Exp lore -> Maybe Doc + ppExpLore _ (If _ _ _ (IfAttr ts _)) = + Just $ stack $ map (text . ("-- "++)) $ lines $ pretty $ + text "Branch returns:" <+> ppTuple' ts + ppExpLore _ _ = Nothing + +commastack :: [Doc] -> Doc +commastack = align . stack . punctuate comma + +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 Shape where + ppr = brackets . commasep . map ppr . shapeDims + +instance Pretty a => Pretty (Ext a) where + ppr (Free e) = ppr e + ppr (Ext x) = text "?" <> text (show x) + +instance Pretty ExtShape where + ppr = brackets . commasep . map ppr . shapeDims + +instance Pretty Space where + ppr DefaultSpace = mempty + ppr (Space s) = text "@" <> text s + +instance Pretty u => Pretty (TypeBase Shape u) where + ppr (Prim et) = ppr et + ppr (Array et (Shape ds) u) = + ppr u <> mconcat (map (brackets . ppr) ds) <> ppr et + ppr (Mem s DefaultSpace) = text "mem" <> parens (ppr s) + ppr (Mem s (Space sp)) = text "mem" <> parens (ppr s) <> text "@" <> text sp + +instance Pretty u => Pretty (TypeBase ExtShape u) where + ppr (Prim et) = ppr et + ppr (Array et (Shape ds) u) = + ppr u <> mconcat (map (brackets . ppr) ds) <> ppr et + ppr (Mem s DefaultSpace) = text "mem" <> parens (ppr s) + ppr (Mem s (Space sp)) = text "mem" <> parens (ppr s) <> text "@" <> text sp + +instance Pretty u => Pretty (TypeBase Rank u) where + ppr (Prim et) = ppr et + ppr (Array et (Rank n) u) = + ppr u <> mconcat (replicate n $ brackets mempty) <> ppr et + ppr (Mem s DefaultSpace) = text "mem" <> parens (ppr s) + ppr (Mem s (Space sp)) = text "mem" <> parens (ppr s) <> text "@" <> text sp + +instance Pretty Ident where + ppr ident = ppr (identType ident) <+> ppr (identName ident) + +instance Pretty SubExp where + ppr (Var v) = ppr v + ppr (Constant v) = ppr v + +instance Pretty Certificates where + ppr (Certificates []) = empty + ppr (Certificates cs) = text "<" <> commasep (map ppr cs) <> text ">" + +instance PrettyLore lore => Pretty (Stms lore) where + ppr = stack . map ppr . stmsToList + +instance PrettyLore lore => Pretty (Body lore) where + ppr (Body _ stms res) + | null stms = braces (commasep $ map ppr res) + | otherwise = stack (map ppr $ stmsToList stms) + text "in" <+> braces (commasep $ map ppr res) + +bindingAnnotation :: PrettyLore lore => Stm lore -> Doc -> Doc +bindingAnnotation bnd = + case mapMaybe ppAnnot $ patternElements $ stmPattern bnd of + [] -> id + annots -> (stack annots ) + +instance Pretty (PatElemT attr) => Pretty (PatternT attr) where + ppr pat = ppPattern (patternContextElements pat) (patternValueElements pat) + +instance Pretty (PatElemT b) => Pretty (PatElemT (a,b)) where + ppr = ppr . fmap snd + +instance Pretty (PatElemT Type) where + ppr (PatElem name t) = ppr t <+> ppr name + +instance Pretty (ParamT b) => Pretty (ParamT (a,b)) where + ppr = ppr . fmap snd + +instance Pretty (ParamT DeclType) where + ppr (Param name t) = + ppr t <+> + ppr name + +instance Pretty (ParamT Type) where + ppr (Param name t) = + ppr t <+> + ppr name + +instance PrettyLore lore => Pretty (Stm lore) where + ppr bnd@(Let pat (StmAux cs attr) e) = + bindingAnnotation bnd $ align $ hang 2 $ + text "let" <+> align (ppr pat) <+> + case (linebreak, ppExpLore attr e) of + (True, Nothing) -> equals e' + (_, Just ann) -> equals (ann e') + (False, Nothing) -> equals <+/> e' + where e' = ppr cs <> ppr e + linebreak = case e of + DoLoop{} -> True + Op{} -> True + If{} -> True + BasicOp ArrayLit{} -> False + _ -> False + +instance Pretty (BasicOp lore) where + ppr (SubExp se) = ppr se + ppr (Opaque e) = text "opaque" <> apply [ppr e] + ppr (ArrayLit [] rt) = + text "empty" <> parens (ppr rt) + ppr (ArrayLit es rt) = + case rt of + Array {} -> brackets $ commastack $ map ppr es + _ -> brackets $ commasep $ map ppr es + ppr (BinOp bop x y) = ppr bop <> parens (ppr x <> comma <+> ppr y) + ppr (CmpOp op x y) = ppr op <> parens (ppr x <> comma <+> ppr y) + ppr (ConvOp conv x) = + text (convOpFun conv) <+> ppr fromtype <+> ppr x <+> text "to" <+> ppr totype + where (fromtype, totype) = convOpType conv + ppr (UnOp op e) = ppr op <+> pprPrec 9 e + ppr (Index v idxs) = + ppr v <> brackets (commasep (map ppr idxs)) + ppr (Update src idxs se) = + ppr src <+> text "with" <+> brackets (commasep (map ppr idxs)) <+> + text "<-" <+> ppr se + ppr (Iota e x s et) = text "iota" <> et' <> apply [ppr e, ppr x, ppr s] + where et' = text $ show $ primBitSize $ IntType et + ppr (Replicate ne ve) = + text "replicate" <> apply [ppr ne, align (ppr ve)] + ppr (Repeat shapes innershape v) = + text "repeat" <> apply [apply $ map ppr $ shapes ++ [innershape], ppr v] + ppr (Scratch t shape) = + text "scratch" <> apply (ppr t : map ppr shape) + ppr (Reshape shape e) = + text "reshape" <> apply [apply (map ppr shape), ppr e] + ppr (Rearrange perm e) = + text "rearrange" <> apply [apply (map ppr perm), ppr e] + ppr (Rotate es e) = + text "rotate" <> apply [apply (map ppr es), ppr e] + ppr (Concat i x ys _) = + text "concat" <> text "@" <> ppr i <> apply (ppr x : map ppr ys) + ppr (Copy e) = text "copy" <> parens (ppr e) + ppr (Manifest perm e) = text "manifest" <> apply [apply (map ppr perm), ppr e] + ppr (Assert e msg (loc, _)) = + text "assert" <> apply [ppr e, ppr msg, text $ show $ locStr loc] + ppr (Partition n flags arrs) = + text "partition" <> + parens (commasep $ [ ppr n, ppr flags ] ++ map ppr arrs) + +instance Pretty a => Pretty (ErrorMsg a) where + ppr (ErrorMsg parts) = commasep $ map p parts + where p (ErrorString s) = text $ show s + p (ErrorInt32 x) = ppr x + +instance PrettyLore lore => Pretty (Exp lore) where + ppr (If c t f (IfAttr _ ifsort)) = + text "if" <+> info' <+> ppr c + text "then" <+> maybeNest t <+> + text "else" <+> maybeNest f + where info' = case ifsort of IfNormal -> mempty + IfFallback -> text "" + maybeNest b | null $ bodyStms b = ppr b + | otherwise = nestedBlock "{" "}" $ ppr b + ppr (BasicOp op) = ppr op + ppr (Apply fname args _ (safety, _, _)) = + text (nameToString fname) <> safety' <> apply (map (align . pprArg) args) + where pprArg (arg, Consume) = text "*" <> ppr arg + pprArg (arg, Observe) = ppr arg + safety' = case safety of Unsafe -> text "" + Safe -> mempty + ppr (Op op) = ppr op + ppr (DoLoop ctx val form loopbody) = + annot (mapMaybe ppAnnot (ctxparams++valparams)) $ + text "loop" <+> ppPattern ctxparams valparams <+> + equals <+> ppTuple' (ctxinit++valinit) + (case form of + ForLoop i it bound [] -> + text "for" <+> align (ppr i <> text ":" <> ppr it <+> + text "<" <+> align (ppr bound)) + ForLoop i it bound loop_vars -> + annot (mapMaybe (ppAnnot . fst) loop_vars) $ + text "for" <+> align (ppr i <> text ":" <> ppr it <+> + text "<" <+> align (ppr bound) + stack (map pprLoopVar loop_vars)) + WhileLoop cond -> + text "while" <+> ppr cond + ) <+> text "do" <+> nestedBlock "{" "}" (ppr loopbody) + where (ctxparams, ctxinit) = unzip ctx + (valparams, valinit) = unzip val + pprLoopVar (p,a) = ppr p <+> text "in" <+> ppr a + +instance PrettyLore lore => Pretty (Lambda lore) where + ppr (Lambda [] _ []) = text "nilFn" + ppr (Lambda params body rettype) = + annot (mapMaybe ppAnnot params) $ + text "fn" <+> ppTuple' rettype <+> + parens (commasep (map ppr params)) <+> + text "=>" indent 2 (ppr body) + +instance PrettyLore lore => Pretty (FunDef lore) where + ppr (FunDef entry name rettype fparams body) = + annot (mapMaybe ppAnnot fparams) $ + text fun <+> ppTuple' rettype <+> + text (nameToString name) + apply (map ppr fparams) <+> + equals <+> nestedBlock "{" "}" (ppr body) + where fun | isJust entry = "entry" + | otherwise = "fun" + +instance PrettyLore lore => Pretty (Prog lore) where + ppr = stack . punctuate line . map ppr . progFunctions + +instance Pretty d => Pretty (DimChange d) where + ppr (DimCoercion se) = text "~" <> ppr se + ppr (DimNew se) = ppr se + +instance Pretty d => Pretty (DimIndex d) where + ppr (DimFix i) = ppr i + ppr (DimSlice i n s) = ppr i <> text ":+" <> ppr n <> text "*" <> ppr s + +ppPattern :: (Pretty a, Pretty b) => [a] -> [b] -> Doc +ppPattern [] bs = braces $ commasep $ map ppr bs +ppPattern as bs = braces $ commasep (map ppr as) <> semi <+> commasep (map ppr bs) + +ppTuple' :: Pretty a => [a] -> Doc +ppTuple' ets = braces $ commasep $ map ppr ets diff --git a/src/Futhark/Representation/AST/RetType.hs b/src/Futhark/Representation/AST/RetType.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/RetType.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE FlexibleInstances, TypeFamilies #-} +-- | This module exports a type class covering representations of +-- function return types. +module Futhark.Representation.AST.RetType + ( + IsBodyType (..) + , bodyTypeValues + , IsRetType (..) + , retTypeValues + , expectedTypes + ) + where + +import qualified Data.Map.Strict as M + +import Futhark.Representation.AST.Syntax.Core +import Futhark.Representation.AST.Attributes.Types + +-- | A type representing the return type of a body. It should contain +-- at least the information contained in a list of 'ExtType's, but may +-- have more, notably an existential context. +class (Show rt, Eq rt, Ord rt, ExtTyped rt) => IsBodyType rt where + -- | Construct a body type from a primitive type. + primBodyType :: PrimType -> rt + +bodyTypeValues :: IsBodyType rt => [rt] -> [ExtType] +bodyTypeValues = map extTypeOf + +instance IsBodyType ExtType where + primBodyType = Prim + +-- | A type representing the return type of a function. In practice, +-- a list of these will be used. It should contain at least the +-- information contained in an 'ExtType', but may have more, notably +-- an existential context. +class (Show rt, Eq rt, Ord rt, DeclExtTyped rt) => IsRetType rt where + -- | Contruct a return type from a primitive type. + primRetType :: PrimType -> rt + + -- | Given a function return type, the parameters of the function, + -- and the arguments for a concrete call, return the instantiated + -- return type for the concrete call, if valid. + applyRetType :: Typed attr => + [rt] + -> [Param attr] + -> [(SubExp, Type)] + -> Maybe [rt] + +retTypeValues :: IsRetType rt => [rt] -> [DeclExtType] +retTypeValues = map declExtTypeOf + +-- | Given shape parameter names and value parameter types, produce the +-- types of arguments accepted. +expectedTypes :: Typed t => [VName] -> [t] -> [SubExp] -> [Type] +expectedTypes shapes value_ts args = map (correctDims . typeOf) value_ts + where parammap :: M.Map VName SubExp + parammap = M.fromList $ zip shapes args + + correctDims t = + t `setArrayShape` + Shape (map correctDim $ shapeDims $ arrayShape t) + + correctDim (Constant v) = Constant v + correctDim (Var v) + | Just se <- M.lookup v parammap = se + | otherwise = Var v + +instance IsRetType DeclExtType where + primRetType = Prim + + applyRetType extret params args = + if length args == length params && + and (zipWith subtypeOf argtypes $ + expectedTypes (map paramName params) params $ map fst args) + then Just $ map correctExtDims extret + else Nothing + where argtypes = map snd args + + parammap :: M.Map VName SubExp + parammap = M.fromList $ zip (map paramName params) (map fst args) + + correctExtDims t = + t `setArrayShape` + Shape (map correctExtDim $ shapeDims $ arrayShape t) + + correctExtDim (Ext i) = Ext i + correctExtDim (Free d) = Free $ correctDim d + + correctDim (Constant v) = Constant v + correctDim (Var v) + | Just se <- M.lookup v parammap = se + | otherwise = Var v diff --git a/src/Futhark/Representation/AST/Syntax.hs b/src/Futhark/Representation/AST/Syntax.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Syntax.hs @@ -0,0 +1,385 @@ +{-# LANGUAGE TypeFamilies, FlexibleContexts, FlexibleInstances, StandaloneDeriving #-} +-- | Futhark core language skeleton. Concrete representations further +-- extend this skeleton by defining a "lore", which specifies concrete +-- annotations ("Futhark.Representation.AST.Annotations") and +-- semantics. +module Futhark.Representation.AST.Syntax + ( + module Language.Futhark.Core + , module Futhark.Representation.AST.Annotations + , module Futhark.Representation.AST.Syntax.Core + + -- * Types + , Uniqueness(..) + , NoUniqueness(..) + , Rank(..) + , ArrayShape(..) + , Space (..) + , TypeBase(..) + , Diet(..) + + -- * Abstract syntax tree + , Ident (..) + , SubExp(..) + , PatElem + , PatElemT (..) + , PatternT (..) + , Pattern + , StmAux(..) + , Stm(..) + , Stms + , Result + , BodyT(..) + , Body + , BasicOp (..) + , UnOp (..) + , BinOp (..) + , CmpOp (..) + , ConvOp (..) + , DimChange (..) + , ShapeChange + , ExpT(..) + , Exp + , LoopForm (..) + , IfAttr (..) + , IfSort (..) + , Safety (..) + , LambdaT(..) + , Lambda + + -- * Definitions + , ParamT (..) + , FParam + , LParam + , FunDefT (..) + , FunDef + , EntryPoint + , EntryPointType(..) + , ProgT(..) + , Prog + + -- * Utils + , oneStm + , stmsFromList + , stmsToList + , stmsHead + ) + where + +import Data.Foldable +import Data.Loc +import qualified Data.Sequence as Seq +import qualified Data.Semigroup as Sem + +import Language.Futhark.Core +import Futhark.Representation.AST.Annotations +import Futhark.Representation.AST.Syntax.Core + +-- | A type alias for namespace control. +type PatElem lore = PatElemT (LetAttr lore) + +-- | A pattern is conceptually just a list of names and their types. +data PatternT attr = + Pattern { patternContextElements :: [PatElemT attr] + -- ^ existential context (sizes and memory blocks) + , patternValueElements :: [PatElemT attr] + -- ^ "real" values + } + deriving (Ord, Show, Eq) + +instance Functor PatternT where + fmap f (Pattern ctx val) = Pattern (map (fmap f) ctx) (map (fmap f) val) + +instance Sem.Semigroup (PatternT attr) where + Pattern cs1 vs1 <> Pattern cs2 vs2 = Pattern (cs1++cs2) (vs1++vs2) + +instance Monoid (PatternT attr) where + mempty = Pattern [] [] + mappend = (Sem.<>) + +-- | A type alias for namespace control. +type Pattern lore = PatternT (LetAttr lore) + +-- | Auxilliary Information associated with a statement. +data StmAux attr = StmAux { stmAuxCerts :: !Certificates + , stmAuxAttr :: attr + } + deriving (Ord, Show, Eq) + +-- | A local variable binding. +data Stm lore = Let { stmPattern :: Pattern lore + , stmAux :: StmAux (ExpAttr lore) + , stmExp :: Exp lore + } + +deriving instance Annotations lore => Ord (Stm lore) +deriving instance Annotations lore => Show (Stm lore) +deriving instance Annotations lore => Eq (Stm lore) + +-- | A sequence of statements. +type Stms lore = Seq.Seq (Stm lore) + +oneStm :: Stm lore -> Stms lore +oneStm = Seq.singleton + +stmsFromList :: [Stm lore] -> Stms lore +stmsFromList = Seq.fromList + +stmsToList :: Stms lore -> [Stm lore] +stmsToList = toList + +stmsHead :: Stms lore -> Maybe (Stm lore, Stms lore) +stmsHead stms = case Seq.viewl stms of stm Seq.:< stms' -> Just (stm, stms') + Seq.EmptyL -> Nothing + +-- | The result of a body is a sequence of subexpressions. +type Result = [SubExp] + +-- | A body consists of a number of bindings, terminating in a result +-- (essentially a tuple literal). +data BodyT lore = Body { bodyAttr :: BodyAttr lore + , bodyStms :: Stms lore + , bodyResult :: Result + } + +deriving instance Annotations lore => Ord (BodyT lore) +deriving instance Annotations lore => Show (BodyT lore) +deriving instance Annotations lore => Eq (BodyT lore) + +-- | Type alias for namespace reasons. +type Body = BodyT + +-- | The new dimension in a 'Reshape'-like operation. This allows us to +-- disambiguate "real" reshapes, that change the actual shape of the +-- array, from type coercions that are just present to make the types +-- work out. +data DimChange d = DimCoercion d + -- ^ The new dimension is guaranteed to be numerically + -- equal to the old one. + | DimNew d + -- ^ The new dimension is not necessarily numerically + -- equal to the old one. + deriving (Eq, Ord, Show) + +instance Functor DimChange where + fmap f (DimCoercion d) = DimCoercion $ f d + fmap f (DimNew d) = DimNew $ f d + +instance Foldable DimChange where + foldMap f (DimCoercion d) = f d + foldMap f (DimNew d) = f d + +instance Traversable DimChange where + traverse f (DimCoercion d) = DimCoercion <$> f d + traverse f (DimNew d) = DimNew <$> f d + +-- | A list of 'DimChange's, indicating the new dimensions of an array. +type ShapeChange d = [DimChange d] + +-- | A primitive operation that returns something of known size and +-- does not itself contain any bindings. +data BasicOp lore + = SubExp SubExp + -- ^ A variable or constant. + + | Opaque SubExp + -- ^ Semantically and operationally just identity, but is + -- invisible/impenetrable to optimisations (hopefully). This is + -- just a hack to avoid optimisation (so, to work around compiler + -- limitations). + + | ArrayLit [SubExp] Type + -- ^ Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@. + -- Second arg is the element type of the rows of the array. + -- Scalar operations + + | UnOp UnOp SubExp + -- ^ Unary operation. + + | BinOp BinOp SubExp SubExp + -- ^ Binary operation. + + | CmpOp CmpOp SubExp SubExp + -- ^ Comparison - result type is always boolean. + + | ConvOp ConvOp SubExp + -- ^ Conversion "casting". + + | Assert SubExp (ErrorMsg SubExp) (SrcLoc, [SrcLoc]) + -- ^ Turn a boolean into a certificate, halting the program with the + -- given error message if the boolean is false. + + -- Primitive array operations + + | Index VName (Slice SubExp) + -- ^ The certificates for bounds-checking are part of the 'Stm'. + + | Update VName (Slice SubExp) SubExp + -- ^ An in-place update of the given array at the given position. + -- Consumes the array. + + | Concat Int VName [VName] SubExp + -- ^ @concat@0([1],[2, 3, 4]) = [1, 2, 3, 4]@. + + | Copy VName + -- ^ Copy the given array. The result will not alias anything. + + | Manifest [Int] VName + -- ^ Manifest an array with dimensions represented in the given + -- order. The result will not alias anything. + + -- Array construction. + | Iota SubExp SubExp SubExp IntType + -- ^ @iota(n, x, s) = [x,x+s,..,x+(n-1)*s]@. + -- + -- The 'IntType' indicates the type of the array returned and the + -- offset/stride arguments, but not the length argument. + + | Replicate Shape SubExp + -- ^ @replicate([3][2],1) = [[1,1], [1,1], [1,1]]@ + + | Repeat [Shape] Shape VName + -- ^ Repeat each dimension of the input array some number of times, + -- given by the corresponding shape. For an array of rank @k@, the + -- list must contain @k@ shapes. A shape may be empty (in which + -- case the dimension is not repeated, but it is still present). + -- The last shape indicates the amount of extra innermost + -- dimensions. All other extra dimensions are added *before* the original dimension. + + | Scratch PrimType [SubExp] + -- ^ Create array of given type and shape, with undefined elements. + + -- Array index space transformation. + | Reshape (ShapeChange SubExp) VName + -- ^ 1st arg is the new shape, 2nd arg is the input array *) + + | Rearrange [Int] VName + -- ^ Permute the dimensions of the input array. The list + -- of integers is a list of dimensions (0-indexed), which + -- must be a permutation of @[0,n-1]@, where @n@ is the + -- number of dimensions in the input array. + + | Rotate [SubExp] VName + -- ^ Rotate the dimensions of the input array. The list of + -- subexpressions specify how much each dimension is rotated. The + -- length of this list must be equal to the rank of the array. + + | Partition Int VName [VName] + -- ^ First variable is the flag array, second is the element + -- arrays. If no arrays are given, the returned offsets are zero, + -- and no arrays are returned. + deriving (Eq, Ord, Show) + +-- | The root Futhark expression type. The 'Op' constructor contains +-- a lore-specific operation. Do-loops, branches and function calls +-- are special. Everything else is a simple 'BasicOp'. +data ExpT lore + = BasicOp (BasicOp lore) + -- ^ A simple (non-recursive) operation. + + | Apply Name [(SubExp, Diet)] [RetType lore] (Safety, SrcLoc, [SrcLoc]) + + | If SubExp (BodyT lore) (BodyT lore) (IfAttr (BranchType lore)) + + | DoLoop [(FParam lore, SubExp)] [(FParam lore, SubExp)] (LoopForm lore) (BodyT lore) + -- ^ @loop {a} = {v} (for i < n|while b) do b@. The merge + -- parameters are divided into context and value part. + + | Op (Op lore) + +deriving instance Annotations lore => Eq (ExpT lore) +deriving instance Annotations lore => Show (ExpT lore) +deriving instance Annotations lore => Ord (ExpT lore) + +-- | Whether something is safe or unsafe (mostly function calls, and +-- in the context of whether operations are dynamically checked). +-- When we inline an 'Unsafe' function, we remove all safety checks in +-- its body. The 'Ord' instance picks 'Unsafe' as being less than +-- 'Safe'. +data Safety = Unsafe | Safe deriving (Eq, Ord, Show) + +-- | For-loop or while-loop? +data LoopForm lore = ForLoop VName IntType SubExp [(LParam lore,VName)] + | WhileLoop VName + +deriving instance Annotations lore => Eq (LoopForm lore) +deriving instance Annotations lore => Show (LoopForm lore) +deriving instance Annotations lore => Ord (LoopForm lore) + +-- | Data associated with a branch. +data IfAttr rt = IfAttr { ifReturns :: [rt] + , ifSort :: IfSort + } + deriving (Eq, Show, Ord) + +data IfSort = IfNormal -- ^ An ordinary branch. + | IfFallback -- ^ A branch where the "true" case is what + -- we are actually interested in, and the + -- "false" case is only present as a fallback + -- for when the true case cannot be safely + -- evaluated. the compiler is permitted to + -- optimise away the branch if the true case + -- contains only safe statements. + deriving (Eq, Show, Ord) + +-- | A type alias for namespace control. +type Exp = ExpT + +-- | Anonymous function for use in a SOAC. +data LambdaT lore = Lambda { lambdaParams :: [LParam lore] + , lambdaBody :: BodyT lore + , lambdaReturnType :: [Type] + } + +deriving instance Annotations lore => Eq (LambdaT lore) +deriving instance Annotations lore => Show (LambdaT lore) +deriving instance Annotations lore => Ord (LambdaT lore) + +-- | Type alias for namespacing reasons. +type Lambda = LambdaT + +type FParam lore = ParamT (FParamAttr lore) + +type LParam lore = ParamT (LParamAttr lore) + +-- | Function Declarations +data FunDefT lore = FunDef { funDefEntryPoint :: Maybe EntryPoint + -- ^ Contains a value if this function is + -- an entry point. + , funDefName :: Name + , funDefRetType :: [RetType lore] + , funDefParams :: [FParam lore] + , funDefBody :: BodyT lore + } + +deriving instance Annotations lore => Eq (FunDefT lore) +deriving instance Annotations lore => Show (FunDefT lore) +deriving instance Annotations lore => Ord (FunDefT lore) + +-- | Information about the parameters and return value of an entry +-- point. The first element is for parameters, the second for return +-- value. +type EntryPoint = ([EntryPointType], [EntryPointType]) + +-- | Every entry point argument and return value has an annotation +-- indicating how it maps to the original source program type. +data EntryPointType = TypeUnsigned + -- ^ Is an unsigned integer or array of unsigned + -- integers. + | TypeOpaque String Int + -- ^ A black box type comprising this many core + -- values. The string is a human-readable + -- description with no other semantics. + | TypeDirect + -- ^ Maps directly. + deriving (Eq, Show, Ord) + +-- | Type alias for namespace reasons. +type FunDef = FunDefT + +-- | An entire Futhark program. +newtype ProgT lore = Prog { progFunctions :: [FunDef lore] } + deriving (Eq, Ord, Show) + +-- | Type alias for namespace reasons. +type Prog = ProgT diff --git a/src/Futhark/Representation/AST/Syntax/Core.hs b/src/Futhark/Representation/AST/Syntax/Core.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Syntax/Core.hs @@ -0,0 +1,357 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +-- | The most primitive ("core") aspects of the AST. Split out of +-- "Futhark.Representation.AST.Syntax" in order for +-- "Futhark.Representation.AST.Annotations" to use these definitions. This +-- module is re-exported from "Futhark.Representation.AST.Syntax" and +-- there should be no reason to include it explicitly. +module Futhark.Representation.AST.Syntax.Core + ( + module Language.Futhark.Core + , module Futhark.Representation.Primitive + + -- * Types + , Uniqueness(..) + , NoUniqueness(..) + , ShapeBase(..) + , Shape + , Ext(..) + , ExtSize + , ExtShape + , Rank(..) + , ArrayShape(..) + , Space (..) + , SpaceId + , TypeBase(..) + , Type + , ExtType + , DeclType + , DeclExtType + , Diet(..) + , ErrorMsg (..) + , ErrorMsgPart (..) + + -- * Values + , PrimValue(..) + + -- * Abstract syntax tree + , Ident (..) + , Certificates(..) + , SubExp(..) + , ParamT (..) + , Param + , DimIndex (..) + , Slice + , dimFix + , sliceIndices + , sliceDims + , unitSlice + , fixSlice + , PatElemT (..) + + -- * Miscellaneous + , Names + ) where + +import Control.Monad.State +import Data.Maybe +import Data.Monoid ((<>)) +import Data.String +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem +import Data.Traversable + +import Language.Futhark.Core +import Futhark.Representation.Primitive + +-- | The size of an array type as a list of its dimension sizes, with +-- the type of sizes being parametric. +newtype ShapeBase d = Shape { shapeDims :: [d] } + deriving (Eq, Ord, Show) + +-- | The size of an array as a list of subexpressions. If a variable, +-- that variable must be in scope where this array is used. +type Shape = ShapeBase SubExp + +-- | Something that may be existential. +data Ext a = Ext Int + | Free a + deriving (Eq, Ord, Show) + +-- | The size of this dimension. +type ExtSize = Ext SubExp + +-- | Like 'Shape' but some of its elements may be bound in a local +-- environment instead. These are denoted with integral indices. +type ExtShape = ShapeBase ExtSize + +-- | The size of an array type as merely the number of dimensions, +-- with no further information. +newtype Rank = Rank Int + deriving (Show, Eq, Ord) + +-- | A class encompassing types containing array shape information. +class (Monoid a, Eq a, Ord a) => ArrayShape a where + -- | Return the rank of an array with the given size. + shapeRank :: a -> Int + -- | @stripDims n shape@ strips the outer @n@ dimensions from + -- @shape@. + stripDims :: Int -> a -> a + -- | Check whether one shape if a subset of another shape. + subShapeOf :: a -> a -> Bool + +instance Sem.Semigroup (ShapeBase d) where + Shape l1 <> Shape l2 = Shape $ l1 `mappend` l2 + +instance Monoid (ShapeBase d) where + mempty = Shape mempty + mappend = (Sem.<>) + +instance Functor ShapeBase where + fmap f = Shape . map f . shapeDims + +instance ArrayShape (ShapeBase SubExp) where + shapeRank (Shape l) = length l + stripDims n (Shape dims) = Shape $ drop n dims + subShapeOf = (==) + +instance ArrayShape (ShapeBase ExtSize) where + shapeRank (Shape l) = length l + stripDims n (Shape dims) = Shape $ drop n dims + subShapeOf (Shape ds1) (Shape ds2) = + -- Must agree on Free dimensions, and ds1 may not be existential + -- where ds2 is Free. Existentials must also be congruent. + length ds1 == length ds2 && + evalState (and <$> zipWithM subDimOf ds1 ds2) M.empty + where subDimOf (Free se1) (Free se2) = return $ se1 == se2 + subDimOf (Ext _) (Free _) = return False + subDimOf (Free _) (Ext _) = return True + subDimOf (Ext x) (Ext y) = do + extmap <- get + case M.lookup y extmap of + Just ywas | ywas == x -> return True + | otherwise -> return False + Nothing -> do put $ M.insert y x extmap + return True + +instance Sem.Semigroup Rank where + Rank x <> Rank y = Rank $ x + y + +instance Monoid Rank where + mempty = Rank 0 + mappend = (Sem.<>) + +instance ArrayShape Rank where + shapeRank (Rank x) = x + stripDims n (Rank x) = Rank $ x - n + subShapeOf = (==) + +-- | The memory space of a block. If 'DefaultSpace', this is the "default" +-- space, whatever that is. The exact meaning of the 'SpaceID' +-- depends on the backend used. In GPU kernels, for example, this is +-- used to distinguish between constant, global and shared memory +-- spaces. In GPU-enabled host code, it is used to distinguish +-- between host memory ('DefaultSpace') and GPU space. +data Space = DefaultSpace + | Space SpaceId + deriving (Show, Eq, Ord) + +-- | A string representing a specific non-default memory space. +type SpaceId = String + +-- | A fancier name for '()' - encodes no uniqueness information. +data NoUniqueness = NoUniqueness + deriving (Eq, Ord, Show) + +-- | An Futhark type is either an array or an element type. When +-- comparing types for equality with '==', shapes must match. +data TypeBase shape u = Prim PrimType + | Array PrimType shape u + | Mem SubExp Space + deriving (Show, Eq, Ord) + +-- | A type with shape information, used for describing the type of +-- variables. +type Type = TypeBase Shape NoUniqueness + +-- | A type with existentially quantified shapes - used as part of +-- function (and function-like) return types. Generally only makes +-- sense when used in a list. +type ExtType = TypeBase ExtShape NoUniqueness + +-- | A type with shape and uniqueness information, used declaring +-- return- and parameters types. +type DeclType = TypeBase Shape Uniqueness + +-- | An 'ExtType' with uniqueness information, used for function +-- return types. +type DeclExtType = TypeBase ExtShape Uniqueness + +-- | Information about which parts of a value/type are consumed. For +-- example, we might say that a function taking three arguments of +-- types @([int], *[int], [int])@ has diet @[Observe, Consume, +-- Observe]@. +data Diet = Consume -- ^ Consumes this value. + | Observe -- ^ Only observes value in this position, does + -- not consume. + deriving (Eq, Ord, Show) + +-- | An identifier consists of its name and the type of the value +-- bound to the identifier. +data Ident = Ident { identName :: VName + , identType :: Type + } + deriving (Show) + +instance Eq Ident where + x == y = identName x == identName y + +instance Ord Ident where + x `compare` y = identName x `compare` identName y + +-- | A list of names used for certificates in some expressions. +newtype Certificates = Certificates { unCertificates :: [VName] } + deriving (Eq, Ord, Show) + +instance Sem.Semigroup Certificates where + Certificates x <> Certificates y = Certificates (x <> y) + +instance Monoid Certificates where + mempty = Certificates mempty + mappend = (Sem.<>) + +-- | A subexpression is either a scalar constant or a variable. One +-- important property is that evaluation of a subexpression is +-- guaranteed to complete in constant time. +data SubExp = Constant PrimValue + | Var VName + deriving (Show, Eq, Ord) + +-- | A function parameter. +data ParamT attr = Param + { paramName :: VName + -- ^ Name of the parameter. + , paramAttr :: attr + -- ^ Function parameter attribute. + } + deriving (Ord, Show, Eq) + +-- | A type alias for namespace control. +type Param = ParamT + +instance Foldable ParamT where + foldMap = foldMapDefault + +instance Functor ParamT where + fmap = fmapDefault + +instance Traversable ParamT where + traverse f (Param name attr) = Param name <$> f attr + +-- | How to index a single dimension of an array. +data DimIndex d = DimFix + d -- ^ Fix index in this dimension. + | DimSlice d d d + -- ^ @DimSlice start_offset num_elems stride@. + deriving (Eq, Ord, Show) + +instance Functor DimIndex where + fmap f (DimFix i) = DimFix $ f i + fmap f (DimSlice i j s) = DimSlice (f i) (f j) (f s) + +instance Foldable DimIndex where + foldMap f (DimFix d) = f d + foldMap f (DimSlice i j s) = f i <> f j <> f s + +instance Traversable DimIndex where + traverse f (DimFix d) = DimFix <$> f d + traverse f (DimSlice i j s) = DimSlice <$> f i <*> f j <*> f s + +-- | A list of 'DimFix's, indicating how an array should be sliced. +-- Whenever a function accepts a 'Slice', that slice should be total, +-- i.e, cover all dimensions of the array. Deviators should be +-- indicated by taking a list of 'DimIndex'es instead. +type Slice d = [DimIndex d] + +-- | If the argument is a 'DimFix', return its component. +dimFix :: DimIndex d -> Maybe d +dimFix (DimFix d) = Just d +dimFix _ = Nothing + +-- | If the slice is all 'DimFix's, return the components. +sliceIndices :: Slice d -> Maybe [d] +sliceIndices = mapM dimFix + +-- | The dimensions of the array produced by this slice. +sliceDims :: Slice d -> [d] +sliceDims = mapMaybe dimSlice + where dimSlice (DimSlice _ d _) = Just d + dimSlice DimFix{} = Nothing + +-- | A slice with a stride of one. +unitSlice :: Num d => d -> d -> DimIndex d +unitSlice offset n = DimSlice offset n 1 + +-- | Fix the 'DimSlice's of a slice. The number of indexes must equal +-- the length of 'sliceDims' for the slice. +fixSlice :: Num d => Slice d -> [d] -> [d] +fixSlice (DimFix j:mis') is' = + j : fixSlice mis' is' +fixSlice (DimSlice orig_k _ orig_s:mis') (i:is') = + (orig_k+i*orig_s) : fixSlice mis' is' +fixSlice _ _ = [] + +-- | An element of a pattern - consisting of an name (essentially a +-- pair of the name andtype), a 'Bindage', and an addditional +-- parametric attribute. This attribute is what is expected to +-- contain the type of the resulting variable. +data PatElemT attr = PatElem { patElemName :: VName + -- ^ The name being bound. + , patElemAttr :: attr + -- ^ Pattern element attribute. + } + deriving (Ord, Show, Eq) + +instance Functor PatElemT where + fmap f (PatElem name attr) = PatElem name (f attr) + +-- | A set of names. +type Names = S.Set VName + +-- | An error message is a list of error parts, which are concatenated +-- to form the final message. +newtype ErrorMsg a = ErrorMsg [ErrorMsgPart a] + deriving (Eq, Ord, Show) + +instance IsString (ErrorMsg a) where + fromString = ErrorMsg . pure . fromString + +-- | A part of an error message. +data ErrorMsgPart a = ErrorString String -- ^ A literal string. + | ErrorInt32 a -- ^ A run-time integer value. + deriving (Eq, Ord, Show) + +instance IsString (ErrorMsgPart a) where + fromString = ErrorString + +instance Functor ErrorMsg where + fmap f (ErrorMsg parts) = ErrorMsg $ map (fmap f) parts + +instance Foldable ErrorMsg where + foldMap f (ErrorMsg parts) = foldMap (foldMap f) parts + +instance Traversable ErrorMsg where + traverse f (ErrorMsg parts) = ErrorMsg <$> traverse (traverse f) parts + +instance Functor ErrorMsgPart where + fmap _ (ErrorString s) = ErrorString s + fmap f (ErrorInt32 a) = ErrorInt32 $ f a + +instance Foldable ErrorMsgPart where + foldMap _ ErrorString{} = mempty + foldMap f (ErrorInt32 a) = f a + +instance Traversable ErrorMsgPart where + traverse _ (ErrorString s) = pure $ ErrorString s + traverse f (ErrorInt32 a) = ErrorInt32 <$> f a diff --git a/src/Futhark/Representation/AST/Traversals.hs b/src/Futhark/Representation/AST/Traversals.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/AST/Traversals.hs @@ -0,0 +1,252 @@ +----------------------------------------------------------------------------- +-- | +-- +-- Functions for generic traversals across Futhark syntax trees. The +-- motivation for this module came from dissatisfaction with rewriting +-- the same trivial tree recursions for every module. A possible +-- alternative would be to use normal \"Scrap your +-- boilerplate\"-techniques, but these are rejected for two reasons: +-- +-- * They are too slow. +-- +-- * More importantly, they do not tell you whether you have missed +-- some cases. +-- +-- Instead, this module defines various traversals of the Futhark syntax +-- tree. The implementation is rather tedious, but the interface is +-- easy to use. +-- +-- A traversal of the Futhark syntax tree is expressed as a tuple of +-- functions expressing the operations to be performed on the various +-- types of nodes. +-- +-- The "Futhark.Transform.Rename" is a simple example of how to use +-- this facility. +-- +----------------------------------------------------------------------------- +module Futhark.Representation.AST.Traversals + ( + -- * Mapping + Mapper(..) + , identityMapper + , mapBody + , mapExpM + , mapExp + , mapOnType + , mapOnLoopForm + , mapOnExtType + + -- * Walking + , Walker(..) + , identityWalker + , walkExpM + , walkExp + -- * Simple wrappers + ) + where + +import Control.Monad +import Control.Monad.Identity +import qualified Data.Traversable +import Data.Monoid ((<>)) + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes.Scope + +-- | 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 { + mapOnSubExp :: SubExp -> m SubExp + , mapOnBody :: Scope tlore -> Body flore -> m (Body tlore) + -- ^ Most bodies are enclosed in a scope, which is passed along + -- for convenience. + , mapOnVName :: VName -> m VName + , mapOnCertificates :: Certificates -> m Certificates + , 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) + } + +-- | A mapper that simply returns the tree verbatim. +identityMapper :: Monad m => Mapper lore lore m +identityMapper = Mapper { + mapOnSubExp = return + , mapOnBody = const return + , mapOnVName = return + , mapOnCertificates = return + , mapOnRetType = return + , mapOnBranchType = return + , mapOnFParam = return + , mapOnLParam = return + , mapOnOp = return + } + +-- | Map across the bindings of a 'Body'. +mapBody :: (Stm lore -> Stm lore) -> Body lore -> Body lore +mapBody f (Body attr stms res) = Body attr (fmap f stms) res + +-- | Map a monadic action across the immediate children of an +-- expression. Importantly, the 'mapOnExp' action is not invoked for +-- the expression itself, and the mapping does not descend recursively +-- into subexpressions. The mapping is done left-to-right. +mapExpM :: (Applicative m, Monad m) => + Mapper flore tlore m -> Exp flore -> m (Exp tlore) +mapExpM tv (BasicOp (SubExp se)) = + BasicOp <$> (SubExp <$> mapOnSubExp tv se) +mapExpM tv (BasicOp (ArrayLit els rowt)) = + BasicOp <$> (pure ArrayLit <*> mapM (mapOnSubExp tv) els <*> + mapOnType (mapOnSubExp tv) rowt) +mapExpM tv (BasicOp (BinOp bop x y)) = + BasicOp <$> (BinOp bop <$> mapOnSubExp tv x <*> mapOnSubExp tv y) +mapExpM tv (BasicOp (CmpOp op x y)) = + BasicOp <$> (CmpOp op <$> mapOnSubExp tv x <*> mapOnSubExp tv y) +mapExpM tv (BasicOp (ConvOp conv x)) = + BasicOp <$> (ConvOp conv <$> mapOnSubExp tv x) +mapExpM tv (BasicOp (UnOp op x)) = + BasicOp <$> (UnOp op <$> mapOnSubExp tv x) +mapExpM tv (If c texp fexp (IfAttr ts s)) = + If <$> mapOnSubExp tv c <*> mapOnBody tv mempty texp <*> mapOnBody tv mempty fexp <*> + (IfAttr <$> mapM (mapOnBranchType tv) ts <*> pure s) +mapExpM tv (Apply fname args ret loc) = do + args' <- forM args $ \(arg, d) -> + (,) <$> mapOnSubExp tv arg <*> pure d + Apply fname <$> pure args' <*> mapM (mapOnRetType tv) ret <*> pure loc +mapExpM tv (BasicOp (Index arr slice)) = + BasicOp <$> (Index <$> mapOnVName tv arr <*> + mapM (traverse (mapOnSubExp tv)) slice) +mapExpM tv (BasicOp (Update arr slice se)) = + BasicOp <$> (Update <$> mapOnVName tv arr <*> + mapM (traverse (mapOnSubExp tv)) slice <*> mapOnSubExp tv se) +mapExpM tv (BasicOp (Iota n x s et)) = + BasicOp <$> (pure Iota <*> mapOnSubExp tv n <*> mapOnSubExp tv x <*> mapOnSubExp tv s <*> pure et) +mapExpM tv (BasicOp (Replicate shape vexp)) = + BasicOp <$> (Replicate <$> mapOnShape tv shape <*> mapOnSubExp tv vexp) +mapExpM tv (BasicOp (Repeat shapes innershape v)) = + BasicOp <$> (Repeat <$> mapM (mapOnShape tv) shapes <*> + mapOnShape tv innershape <*> mapOnVName tv v) +mapExpM tv (BasicOp (Scratch t shape)) = + BasicOp <$> (Scratch t <$> mapM (mapOnSubExp tv) shape) +mapExpM tv (BasicOp (Reshape shape arrexp)) = + BasicOp <$> (Reshape <$> + mapM (Data.Traversable.traverse (mapOnSubExp tv)) shape <*> + mapOnVName tv arrexp) +mapExpM tv (BasicOp (Rearrange perm e)) = + BasicOp <$> (Rearrange <$> pure perm <*> mapOnVName tv e) +mapExpM tv (BasicOp (Rotate es e)) = + BasicOp <$> (Rotate <$> mapM (mapOnSubExp tv) es <*> mapOnVName tv e) +mapExpM tv (BasicOp (Concat i x ys size)) = + BasicOp <$> (Concat <$> pure i <*> + mapOnVName tv x <*> mapM (mapOnVName tv) ys <*> + mapOnSubExp tv size) +mapExpM tv (BasicOp (Copy e)) = + BasicOp <$> (pure Copy <*> mapOnVName tv e) +mapExpM tv (BasicOp (Manifest perm e)) = + BasicOp <$> (Manifest perm <$> mapOnVName tv e) +mapExpM tv (BasicOp (Assert e msg loc)) = + BasicOp <$> (Assert <$> mapOnSubExp tv e <*> traverse (mapOnSubExp tv) msg <*> pure loc) +mapExpM tv (BasicOp (Opaque e)) = + BasicOp <$> (Opaque <$> mapOnSubExp tv e) +mapExpM tv (BasicOp (Partition n flags arr)) = + BasicOp <$> (Partition <$> + pure n <*> mapOnVName tv flags <*> mapM (mapOnVName tv) arr) +mapExpM tv (DoLoop ctxmerge valmerge form loopbody) = do + ctxparams' <- mapM (mapOnFParam tv) ctxparams + valparams' <- mapM (mapOnFParam tv) valparams + form' <- mapOnLoopForm tv form + let scope = scopeOf form' <> scopeOfFParams (ctxparams'++valparams') + DoLoop <$> + (zip ctxparams' <$> mapM (mapOnSubExp tv) ctxinits) <*> + (zip valparams' <$> mapM (mapOnSubExp tv) valinits) <*> + pure form' <*> mapOnBody tv scope loopbody + where (ctxparams,ctxinits) = unzip ctxmerge + (valparams,valinits) = unzip valmerge +mapExpM tv (Op op) = + Op <$> mapOnOp tv op + +mapOnShape :: Monad m => Mapper flore tlore m -> Shape -> m Shape +mapOnShape tv (Shape ds) = Shape <$> mapM (mapOnSubExp tv) ds + +mapOnExtType :: Monad m => + Mapper flore tlore m -> TypeBase ExtShape u -> m (TypeBase ExtShape u) +mapOnExtType tv (Array bt (Shape shape) u) = + Array bt <$> (Shape <$> mapM mapOnExtSize shape) <*> + return u + where mapOnExtSize (Ext x) = return $ Ext x + mapOnExtSize (Free se) = Free <$> mapOnSubExp tv se +mapOnExtType _ (Prim bt) = return $ Prim bt +mapOnExtType tv (Mem size space) = Mem <$> mapOnSubExp tv size <*> pure space + +mapOnLoopForm :: Monad m => + Mapper flore tlore m -> LoopForm flore -> m (LoopForm tlore) +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) + where (loop_lparams,loop_arrs) = unzip loop_vars +mapOnLoopForm tv (WhileLoop cond) = + WhileLoop <$> mapOnVName tv cond + +-- | Like 'mapExp', but in the 'Identity' monad. +mapExp :: Mapper flore tlore Identity -> Exp flore -> Exp tlore +mapExp m = runIdentity . mapExpM m + +mapOnType :: Monad m => + (SubExp -> m SubExp) -> Type -> m Type +mapOnType _ (Prim bt) = return $ Prim bt +mapOnType f (Mem size space) = Mem <$> f size <*> pure space +mapOnType f (Array bt shape u) = + Array bt <$> (Shape <$> mapM f (shapeDims shape)) <*> pure u + +-- | 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 { + walkOnSubExp :: SubExp -> m () + , walkOnBody :: Body lore -> m () + , walkOnVName :: VName -> m () + , walkOnCertificates :: Certificates -> m () + , walkOnRetType :: RetType lore -> m () + , walkOnBranchType :: BranchType lore -> m () + , walkOnFParam :: FParam lore -> m () + , walkOnLParam :: LParam lore -> m () + , walkOnOp :: Op lore -> m () + } + +-- | A no-op traversal. +identityWalker :: Monad m => Walker lore m +identityWalker = Walker { + walkOnSubExp = const $ return () + , walkOnBody = const $ return () + , walkOnVName = const $ return () + , walkOnCertificates = const $ return () + , walkOnRetType = const $ return () + , walkOnBranchType = const $ return () + , walkOnFParam = const $ return () + , walkOnLParam = const $ return () + , walkOnOp = const $ return () + } + +walkMapper :: Monad m => Walker lore m -> Mapper lore lore m +walkMapper f = Mapper { + mapOnSubExp = wrap walkOnSubExp + , mapOnBody = const $ wrap walkOnBody + , mapOnVName = wrap walkOnVName + , mapOnCertificates = wrap walkOnCertificates + , mapOnRetType = wrap walkOnRetType + , mapOnBranchType = wrap walkOnBranchType + , mapOnFParam = wrap walkOnFParam + , mapOnLParam = wrap walkOnLParam + , mapOnOp = wrap walkOnOp + } + where wrap op k = op f k >> return k + +-- | As 'walkBodyM', but for expressions. +walkExpM :: Monad m => Walker lore m -> Exp lore -> m () +walkExpM f = void . mapExpM m + where m = walkMapper f + +-- | As 'walkExp', but runs in the 'Identity' monad.. +walkExp :: Walker lore Identity -> Exp lore -> () +walkExp f = runIdentity . walkExpM f diff --git a/src/Futhark/Representation/Aliases.hs b/src/Futhark/Representation/Aliases.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Aliases.hs @@ -0,0 +1,376 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +-- | A representation where all bindings are annotated with aliasing +-- information. +module Futhark.Representation.Aliases + ( -- * The Lore definition + Aliases + , Names' (..) + , VarAliases + , ConsumedInExp + , BodyAliasing + , module Futhark.Representation.AST.Attributes.Aliases + -- * Module re-exports + , module Futhark.Representation.AST.Attributes + , module Futhark.Representation.AST.Traversals + , module Futhark.Representation.AST.Pretty + , module Futhark.Representation.AST.Syntax + -- * Adding aliases + , addAliasesToPattern + , mkAliasedLetStm + , mkAliasedBody + , mkPatternAliases + , mkBodyAliases + -- * Removing aliases + , removeProgAliases + , removeFunDefAliases + , removeExpAliases + , removeBodyAliases + , removeStmAliases + , removeLambdaAliases + , removePatternAliases + , removeScopeAliases + -- * Tracking aliases + , AliasesAndConsumed + , trackAliases + , consumedInStms + ) +where + +import Control.Monad.Identity +import Control.Monad.Reader +import Data.Foldable +import Data.Maybe +import Data.Monoid ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import qualified Data.Semigroup as Sem + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Pretty +import Futhark.Transform.Rename +import Futhark.Binder +import Futhark.Transform.Substitute +import Futhark.Analysis.Rephrase +import Futhark.Representation.AST.Attributes.Ranges() +import qualified Futhark.Util.Pretty as PP + +-- | The lore for the basic representation. +data Aliases lore + +-- | A wrapper around 'Names' to get around the fact that we need an +-- 'Ord' instance, which 'Names' does not have. +newtype Names' = Names' { unNames :: Names } + deriving (Show) + +instance Sem.Semigroup Names' where + x <> y = Names' $ unNames x <> unNames y + +instance Monoid Names' where + mempty = Names' mempty + mappend = (Sem.<>) + +instance Eq Names' where + _ == _ = True + +instance Ord Names' where + _ `compare` _ = EQ + +instance Rename Names' where + rename (Names' names) = Names' <$> rename names + +instance Substitute Names' where + substituteNames substs (Names' names) = Names' $ substituteNames substs names + +instance FreeIn Names' where + freeIn = const mempty + +instance PP.Pretty Names' where + ppr = PP.commasep . map PP.ppr . S.toList . unNames + +-- | The aliases of the let-bound variable. +type VarAliases = Names' + +-- | Everything consumed in the expression. +type ConsumedInExp = Names' + +-- | The aliases of what is returned by the 'Body', and what is +-- consumed inside of it. +type BodyAliasing = ([VarAliases], ConsumedInExp) + +instance (Annotations lore, CanBeAliased (Op lore)) => + Annotations (Aliases lore) where + type LetAttr (Aliases lore) = (VarAliases, LetAttr lore) + type ExpAttr (Aliases lore) = (ConsumedInExp, ExpAttr lore) + type BodyAttr (Aliases lore) = (BodyAliasing, BodyAttr lore) + type FParamAttr (Aliases lore) = FParamAttr lore + type LParamAttr (Aliases lore) = LParamAttr lore + type RetType (Aliases lore) = RetType lore + type BranchType (Aliases lore) = BranchType lore + type Op (Aliases lore) = OpWithAliases (Op lore) + +instance AliasesOf (VarAliases, attr) where + aliasesOf = unNames . fst + +instance FreeAttr Names' where + +withoutAliases :: (HasScope (Aliases lore) m, Monad m) => + ReaderT (Scope lore) m a -> m a +withoutAliases m = do + scope <- asksScope removeScopeAliases + runReaderT m scope + +instance (Attributes lore, CanBeAliased (Op lore)) => Attributes (Aliases lore) where + expTypesFromPattern = + withoutAliases . expTypesFromPattern . removePatternAliases + +instance (Attributes lore, CanBeAliased (Op lore)) => Aliased (Aliases lore) where + bodyAliases = map unNames . fst . fst . bodyAttr + consumedInBody = unNames . snd . fst . bodyAttr + +instance PrettyAnnot (PatElemT attr) => + PrettyAnnot (PatElemT (VarAliases, attr)) where + + ppAnnot (PatElem name (Names' als, attr)) = + let alias_comment = PP.oneLine <$> aliasComment name als + in case (alias_comment, ppAnnot (PatElem name attr)) of + (_, Nothing) -> + alias_comment + (Just alias_comment', Just inner_comment) -> + Just $ alias_comment' PP. inner_comment + (Nothing, Just inner_comment) -> + Just inner_comment + + +instance (Attributes lore, CanBeAliased (Op lore)) => PrettyLore (Aliases lore) where + ppExpLore (consumed, inner) e = + maybeComment $ catMaybes [expAttr, + mergeAttr, + ppExpLore inner $ removeExpAliases e] + where mergeAttr = + case e of + DoLoop _ merge _ body -> + let mergeParamAliases fparam als + | primType (paramType fparam) = + Nothing + | otherwise = + resultAliasComment (paramName fparam) als + in maybeComment $ catMaybes $ + zipWith mergeParamAliases (map fst merge) $ + bodyAliases body + _ -> Nothing + + expAttr = case S.toList $ unNames consumed of + [] -> Nothing + als -> Just $ PP.oneLine $ + PP.text "-- Consumes " <> PP.commasep (map PP.ppr als) + +maybeComment :: [PP.Doc] -> Maybe PP.Doc +maybeComment [] = Nothing +maybeComment cs = Just $ PP.folddoc (PP.) cs + +aliasComment :: (PP.Pretty a, PP.Pretty b) => + a -> S.Set b -> Maybe PP.Doc +aliasComment name als = + case S.toList als of + [] -> Nothing + als' -> Just $ PP.oneLine $ + PP.text "-- " <> PP.ppr name <> PP.text " aliases " <> + PP.commasep (map PP.ppr als') + +resultAliasComment :: (PP.Pretty a, PP.Pretty b) => + a -> S.Set b -> Maybe PP.Doc +resultAliasComment name als = + case S.toList als of + [] -> Nothing + als' -> Just $ PP.oneLine $ + 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 = Rephraser { rephraseExpLore = return . snd + , rephraseLetBoundLore = return . snd + , rephraseBodyLore = return . snd + , rephraseFParamLore = return + , rephraseLParamLore = return + , rephraseRetType = return + , rephraseBranchType = return + , rephraseOp = return . removeOpAliases + } + +removeScopeAliases :: Scope (Aliases lore) -> Scope lore +removeScopeAliases = M.map unAlias + where unAlias (LetInfo (_, attr)) = LetInfo attr + unAlias (FParamInfo attr) = FParamInfo attr + unAlias (LParamInfo attr) = LParamInfo attr + unAlias (IndexInfo it) = IndexInfo it + +removeProgAliases :: CanBeAliased (Op lore) => + Prog (Aliases lore) -> Prog lore +removeProgAliases = runIdentity . rephraseProg removeAliases + +removeFunDefAliases :: CanBeAliased (Op lore) => + FunDef (Aliases lore) -> FunDef lore +removeFunDefAliases = runIdentity . rephraseFunDef removeAliases + +removeExpAliases :: CanBeAliased (Op lore) => + Exp (Aliases lore) -> Exp lore +removeExpAliases = runIdentity . rephraseExp removeAliases + +removeBodyAliases :: CanBeAliased (Op lore) => + Body (Aliases lore) -> Body lore +removeBodyAliases = runIdentity . rephraseBody removeAliases + +removeStmAliases :: CanBeAliased (Op lore) => + Stm (Aliases lore) -> Stm lore +removeStmAliases = runIdentity . rephraseStm removeAliases + +removeLambdaAliases :: CanBeAliased (Op lore) => + Lambda (Aliases lore) -> Lambda lore +removeLambdaAliases = runIdentity . rephraseLambda removeAliases + +removePatternAliases :: PatternT (Names', a) + -> PatternT a +removePatternAliases = runIdentity . rephrasePattern (return . snd) + +addAliasesToPattern :: (Attributes lore, CanBeAliased (Op lore), Typed attr) => + PatternT attr -> Exp (Aliases lore) + -> PatternT (VarAliases, attr) +addAliasesToPattern pat e = + uncurry Pattern $ mkPatternAliases pat e + +mkAliasedBody :: (Attributes lore, CanBeAliased (Op lore)) => + BodyAttr lore -> Stms (Aliases lore) -> Result -> Body (Aliases lore) +mkAliasedBody innerlore bnds res = + Body (mkBodyAliases bnds res, innerlore) bnds res + +mkPatternAliases :: (Attributes lore, Aliased lore, Typed attr) => + PatternT attr -> Exp lore + -> ([PatElemT (VarAliases, attr)], + [PatElemT (VarAliases, attr)]) +mkPatternAliases pat e = + -- Some part of the pattern may be the context. This does not have + -- aliases from expAliases, so we use a hack to compute aliases of + -- the context. + let als = expAliases e ++ repeat mempty -- In case the pattern has + -- more elements (this + -- implies a type error). + context_als = mkContextAliases pat e + in (zipWith annotateBindee (patternContextElements pat) context_als, + zipWith annotateBindee (patternValueElements pat) als) + where annotateBindee bindee names = + bindee `setPatElemLore` (Names' names', patElemAttr bindee) + where names' = + case patElemType bindee of + Array {} -> names + Mem _ _ -> names + _ -> mempty + +mkContextAliases :: (Attributes lore, Aliased lore) => + PatternT attr -> Exp lore + -> [Names] +mkContextAliases pat (DoLoop ctxmerge valmerge _ body) = + let ctx = loopResultContext (map fst ctxmerge) (map fst valmerge) + init_als = zip mergenames $ map (subExpAliases . snd) $ ctxmerge ++ valmerge + expand als = als <> S.unions (mapMaybe (`lookup` init_als) (S.toList als)) + merge_als = zip mergenames $ + map ((`S.difference` mergenames_set) . expand) $ + bodyAliases body + in if length ctx == length (patternContextElements pat) + then map (fromMaybe mempty . flip lookup merge_als . paramName) ctx + else map (const mempty) $ patternContextElements pat + where mergenames = map (paramName . fst) $ ctxmerge ++ valmerge + mergenames_set = S.fromList mergenames +mkContextAliases pat (If _ tbranch fbranch _) = + take (length $ patternContextNames pat) $ + zipWith (<>) (bodyAliases tbranch) (bodyAliases fbranch) +mkContextAliases pat _ = + replicate (length $ patternContextElements pat) mempty + +mkBodyAliases :: Aliased lore => + Stms lore + -> Result + -> BodyAliasing +mkBodyAliases bnds res = + -- We need to remove the names that are bound in bnds from the alias + -- and consumption sets. We do this by computing the transitive + -- closure of the alias map (within bnds), then removing anything + -- bound in bnds. + let (aliases, consumed) = mkStmsAliases bnds res + boundNames = + fold $ fmap (S.fromList . patternNames . stmPattern) bnds + bound = (`S.member` boundNames) + aliases' = map (S.filter (not . bound)) aliases + consumed' = S.filter (not . bound) consumed + in (map Names' aliases', Names' consumed') + +mkStmsAliases :: Aliased lore => + Stms lore -> [SubExp] + -> ([Names], Names) +mkStmsAliases bnds res = delve mempty $ stmsToList bnds + where delve (aliasmap, consumed) [] = + (map (aliasClosure aliasmap . subExpAliases) res, + consumed) + delve (aliasmap, consumed) (bnd:bnds') = + delve (trackAliases (aliasmap, consumed) bnd) bnds' + aliasClosure aliasmap names = + names `S.union` mconcat (map look $ S.toList names) + where look k = M.findWithDefault mempty k aliasmap + +-- | Everything consumed in the given bindings and result (even transitively). +consumedInStms :: Aliased lore => Stms lore -> [SubExp] -> Names +consumedInStms bnds res = snd $ mkStmsAliases bnds res + +type AliasesAndConsumed = (M.Map VName Names, + Names) + +trackAliases :: Aliased lore => + AliasesAndConsumed -> Stm lore + -> AliasesAndConsumed +trackAliases (aliasmap, consumed) bnd = + let pat = stmPattern bnd + als = M.fromList $ + zip (patternNames pat) (map addAliasesOfAliases $ patternAliases pat) + aliasmap' = als <> aliasmap + consumed' = consumed <> addAliasesOfAliases (consumedInStm bnd) + in (aliasmap', consumed') + where addAliasesOfAliases names = names <> aliasesOfAliases names + aliasesOfAliases = mconcat . map look . S.toList + look k = M.findWithDefault mempty k aliasmap + +mkAliasedLetStm :: (Attributes lore, CanBeAliased (Op lore)) => + Pattern lore + -> StmAux (ExpAttr lore) -> Exp (Aliases lore) + -> Stm (Aliases lore) +mkAliasedLetStm pat (StmAux cs attr) e = + Let (addAliasesToPattern pat e) + (StmAux cs (Names' $ consumedInExp e, attr)) + e + +instance (Bindable lore, CanBeAliased (Op lore)) => Bindable (Aliases lore) where + mkExpAttr pat e = + let attr = mkExpAttr (removePatternAliases pat) $ removeExpAliases e + in (Names' $ consumedInExp e, attr) + + mkExpPat ctx val e = + addAliasesToPattern (mkExpPat ctx val $ removeExpAliases e) e + + mkLetNames names e = do + env <- asksScope removeScopeAliases + flip runReaderT env $ do + Let pat attr _ <- mkLetNames names $ removeExpAliases e + return $ mkAliasedLetStm pat attr e + + mkBody bnds res = + let Body bodylore _ _ = mkBody (fmap removeStmAliases bnds) res + in mkAliasedBody bodylore bnds res + +instance (Attributes (Aliases lore), Bindable (Aliases lore)) => BinderOps (Aliases lore) where + mkBodyB = bindableMkBodyB + mkExpAttrB = bindableMkExpAttrB + mkLetNamesB = bindableMkLetNamesB diff --git a/src/Futhark/Representation/ExplicitMemory.hs b/src/Futhark/Representation/ExplicitMemory.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/ExplicitMemory.hs @@ -0,0 +1,1112 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeFamilies, FlexibleInstances, FlexibleContexts, MultiParamTypeClasses #-} +{-# LANGUAGE ConstraintKinds #-} +-- | This representation requires that every array is given +-- information about which memory block is it based in, and how array +-- elements map to memory block offsets. The representation is based +-- on the kernels representation, so nested parallelism does not +-- occur. +-- +-- There are two primary concepts you will need to understand: +-- +-- 1. Memory blocks, which are Futhark values of type 'Mem' +-- (parametrized with their size). These correspond to arbitrary +-- blocks of memory, and are created using the 'Alloc' operation. +-- +-- 2. Index functions, which describe a mapping from the index space +-- of an array (eg. a two-dimensional space for an array of type +-- @[[int]]@) to a one-dimensional offset into a memory block. +-- Thus, index functions describe how arbitrary-dimensional arrays +-- are mapped to the single-dimensional world of memory. +-- +-- At a conceptual level, imagine that we have a two-dimensional array +-- @a@ of 32-bit integers, consisting of @n@ rows of @m@ elements +-- each. This array could be represented in classic row-major format +-- with an index function like the following: +-- +-- @ +-- f(i,j) = i * m + j +-- @ +-- +-- When we want to know the location of element @a[2,3]@, we simply +-- call the index function as @f(2,3)@ and obtain @2*m+3@. We could +-- also have chosen another index function, one that represents the +-- array in column-major (or "transposed") format: +-- +-- @ +-- f(i,j) = j * n + i +-- @ +-- +-- Index functions are not Futhark-level functions, but a special +-- construct that the final code generator will eventually use to +-- generate concrete access code. By modifying the index functions we +-- can change how an array is represented in memory, which can permit +-- memory access pattern optimisations. +-- +-- Every time we bind an array, whether in a @let@-binding, @loop@ +-- merge parameter, or @lambda@ parameter, we have an annotation +-- specifying a memory block and an index function. In some cases, +-- such as @let@-bindings for many expressions, we are free to specify +-- an arbitrary index function and memory block - for example, we get +-- to decide where 'Copy' stores its result - but in other cases the +-- type rules of the expression chooses for us. For example, 'Index' +-- always produces an array in the same memory block as its input, and +-- with the same index function, except with some indices fixed. +module Futhark.Representation.ExplicitMemory + ( -- * The Lore definition + ExplicitMemory + , InKernel + , MemOp (..) + , MemInfo (..) + , MemBound + , MemBind (..) + , MemReturn (..) + , IxFun + , ExtIxFun + , isStaticIxFun + , ExpReturns + , BodyReturns + , FunReturns + , noUniquenessReturns + , bodyReturnsToExpReturns + , ExplicitMemorish + , expReturns + , extReturns + , sliceInfo + , lookupMemInfo + , subExpMemInfo + , lookupMemSize + , lookupArraySummary + , fullyLinear + , ixFunMatchesInnerShape + , existentialiseIxFun + + -- * Module re-exports + , module Futhark.Representation.AST.Attributes + , module Futhark.Representation.AST.Traversals + , module Futhark.Representation.AST.Pretty + , module Futhark.Representation.AST.Syntax + , module Futhark.Representation.Kernels.Kernel + , module Futhark.Representation.Kernels.KernelExp + , module Futhark.Analysis.PrimExp.Convert + ) +where + +import Data.Maybe +import Control.Monad.State +import Control.Monad.Reader +import Control.Monad.Except +import qualified Data.Map.Strict as M +import Data.Foldable (traverse_) +import Data.List +import Data.Monoid ((<>)) + +import Futhark.Analysis.Metrics +import Futhark.Representation.AST.Syntax +import Futhark.Representation.Kernels.Kernel +import Futhark.Representation.Kernels.KernelExp +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Pretty +import Futhark.Transform.Rename +import Futhark.Transform.Substitute +import qualified Futhark.TypeCheck as TC +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Analysis.PrimExp.Convert +import Futhark.Analysis.PrimExp.Simplify +import Futhark.Util +import Futhark.Util.IntegralExp +import qualified Futhark.Util.Pretty as PP +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Lore +import Futhark.Representation.Aliases + (Aliases, removeScopeAliases, removeExpAliases, removePatternAliases) +import Futhark.Representation.AST.Attributes.Ranges +import Futhark.Analysis.Usage +import qualified Futhark.Analysis.SymbolTable as ST + +-- | A lore containing explicit memory information. +data ExplicitMemory +data InKernel + +type ExplicitMemorish lore = (SameScope lore ExplicitMemory, + RetType lore ~ FunReturns, + BranchType lore ~ BodyReturns, + CanBeAliased (Op lore), + Attributes lore, Annotations lore, + TC.Checkable lore, + OpReturns lore) + +instance IsRetType FunReturns where + primRetType = MemPrim + applyRetType = applyFunReturns + +instance IsBodyType BodyReturns where + primBodyType = MemPrim + +data MemOp inner = Alloc SubExp Space + -- ^ Allocate a memory block. This really should not be an + -- expression, but what are you gonna do... + | Inner inner + deriving (Eq, Ord, Show) + +instance FreeIn inner => FreeIn (MemOp inner) where + freeIn (Alloc size _) = freeIn size + freeIn (Inner k) = freeIn k + +instance TypedOp inner => TypedOp (MemOp inner) where + opType (Alloc size space) = pure [Mem size space] + opType (Inner k) = opType k + +instance AliasedOp inner => AliasedOp (MemOp inner) where + opAliases Alloc{} = [mempty] + opAliases (Inner k) = opAliases k + + consumedInOp Alloc{} = mempty + consumedInOp (Inner k) = consumedInOp k + +instance CanBeAliased inner => CanBeAliased (MemOp inner) where + type OpWithAliases (MemOp inner) = MemOp (OpWithAliases inner) + removeOpAliases (Alloc se space) = Alloc se space + removeOpAliases (Inner k) = Inner $ removeOpAliases k + + addOpAliases (Alloc se space) = Alloc se space + addOpAliases (Inner k) = Inner $ addOpAliases k + +instance RangedOp inner => RangedOp (MemOp inner) where + opRanges (Alloc _ _) = + [unknownRange] + opRanges (Inner k) = + opRanges k + +instance CanBeRanged inner => CanBeRanged (MemOp inner) where + type OpWithRanges (MemOp inner) = MemOp (OpWithRanges inner) + removeOpRanges (Alloc size space) = Alloc size space + removeOpRanges (Inner k) = Inner $ removeOpRanges k + + addOpRanges (Alloc size space) = Alloc size space + addOpRanges (Inner k) = Inner $ addOpRanges k + +instance Rename inner => Rename (MemOp inner) where + rename (Alloc size space) = Alloc <$> rename size <*> pure space + rename (Inner k) = Inner <$> rename k + +instance Substitute inner => Substitute (MemOp inner) where + substituteNames subst (Alloc size space) = Alloc (substituteNames subst size) space + substituteNames subst (Inner k) = Inner $ substituteNames subst k + +instance PP.Pretty inner => PP.Pretty (MemOp inner) where + ppr (Alloc e DefaultSpace) = PP.text "alloc" <> PP.apply [PP.ppr e] + ppr (Alloc e (Space sp)) = PP.text "alloc" <> PP.apply [PP.ppr e, PP.text sp] + ppr (Inner k) = PP.ppr k + +instance OpMetrics inner => OpMetrics (MemOp inner) where + opMetrics Alloc{} = seen "Alloc" + opMetrics (Inner k) = opMetrics k + +instance IsOp inner => IsOp (MemOp inner) where + safeOp Alloc{} = True + safeOp (Inner k) = safeOp k + cheapOp (Inner k) = cheapOp k + cheapOp Alloc{} = True + +instance UsageInOp inner => UsageInOp (MemOp inner) where + usageInOp Alloc {} = mempty + usageInOp (Inner k) = usageInOp k + +instance CanBeWise inner => CanBeWise (MemOp inner) where + type OpWithWisdom (MemOp inner) = MemOp (OpWithWisdom inner) + removeOpWisdom (Alloc size space) = Alloc size space + removeOpWisdom (Inner k) = Inner $ removeOpWisdom k + +instance ST.IndexOp inner => ST.IndexOp (MemOp inner) where + indexOp vtable k (Inner op) is = ST.indexOp vtable k op is + indexOp _ _ _ _ = Nothing + +instance Annotations ExplicitMemory where + type LetAttr ExplicitMemory = MemInfo SubExp NoUniqueness MemBind + type FParamAttr ExplicitMemory = MemInfo SubExp Uniqueness MemBind + type LParamAttr ExplicitMemory = MemInfo SubExp NoUniqueness MemBind + type RetType ExplicitMemory = FunReturns + type BranchType ExplicitMemory = BodyReturns + type Op ExplicitMemory = MemOp (Kernel InKernel) + +instance Annotations InKernel where + type LetAttr InKernel = MemInfo SubExp NoUniqueness MemBind + type FParamAttr InKernel = MemInfo SubExp Uniqueness MemBind + type LParamAttr InKernel = MemInfo SubExp NoUniqueness MemBind + type RetType InKernel = FunReturns + type BranchType InKernel = BodyReturns + type Op InKernel = MemOp (KernelExp InKernel) + +-- | The index function representation used for memory annotations. +type IxFun = IxFun.IxFun (PrimExp VName) + +-- | An index function that may contain existential variables. +type ExtIxFun = IxFun.IxFun (PrimExp (Ext VName)) + +-- | A summary of the memory information for every let-bound +-- identifier, function parameter, and return value. Parameterisered +-- over uniqueness, dimension, and auxiliary array information. +data MemInfo d u ret = MemPrim PrimType + -- ^ A primitive value. + | MemMem d Space + -- ^ A memory block. + | MemArray PrimType (ShapeBase d) u ret + -- ^ The array is stored in the named memory block, + -- and with the given index function. The index + -- function maps indices in the array to /element/ + -- offset, /not/ byte offsets! To translate to byte + -- offsets, multiply the offset with the size of the + -- array element type. + deriving (Eq, Show, Ord) --- XXX Ord? + +type MemBound u = MemInfo SubExp u MemBind + +instance FixExt ret => DeclExtTyped (MemInfo ExtSize Uniqueness ret) where + declExtTypeOf (MemPrim pt) = Prim pt + declExtTypeOf (MemMem (Free size) space) = Mem size space + declExtTypeOf (MemMem Ext{} space) = Mem (intConst Int32 0) space -- XXX + declExtTypeOf (MemArray pt shape u _) = Array pt shape u + +instance FixExt ret => ExtTyped (MemInfo ExtSize NoUniqueness ret) where + extTypeOf (MemPrim pt) = Prim pt + extTypeOf (MemMem (Free size) space) = Mem size space + extTypeOf (MemMem Ext{} space) = Mem (intConst Int32 0) space -- XXX + extTypeOf (MemArray pt shape u _) = Array pt shape u + +instance FixExt ret => FixExt (MemInfo ExtSize u ret) where + fixExt _ _ (MemPrim pt) = MemPrim pt + fixExt i se (MemMem size space) = MemMem (fixExt i se size) space + fixExt i se (MemArray pt shape u ret) = + MemArray pt (fixExt i se shape) u (fixExt i se ret) + +instance Typed (MemInfo SubExp Uniqueness ret) where + typeOf = fromDecl . declTypeOf + +instance Typed (MemInfo SubExp NoUniqueness ret) where + typeOf (MemPrim pt) = Prim pt + typeOf (MemMem size space) = Mem size space + typeOf (MemArray bt shape u _) = Array bt shape u + +instance DeclTyped (MemInfo SubExp Uniqueness ret) where + declTypeOf (MemPrim bt) = Prim bt + declTypeOf (MemMem size space) = Mem size space + declTypeOf (MemArray bt shape u _) = Array bt shape u + +instance (FreeIn d, FreeIn ret) => FreeIn (MemInfo d u ret) where + freeIn (MemArray _ shape _ ret) = freeIn shape <> freeIn ret + freeIn (MemMem size _) = freeIn size + freeIn (MemPrim _) = mempty + +instance (Substitute d, Substitute ret) => Substitute (MemInfo d u ret) where + substituteNames subst (MemArray bt shape u ret) = + MemArray bt + (substituteNames subst shape) u + (substituteNames subst ret) + substituteNames substs (MemMem size space) = + MemMem (substituteNames substs size) space + substituteNames _ (MemPrim bt) = + MemPrim bt + +instance (Substitute d, Substitute ret) => Rename (MemInfo d u ret) where + rename = substituteRename + +simplifyIxFun :: Engine.SimplifiableLore lore => + IxFun -> Engine.SimpleM lore IxFun +simplifyIxFun = traverse simplifyPrimExp + +simplifyExtIxFun :: Engine.SimplifiableLore lore => + ExtIxFun -> Engine.SimpleM lore ExtIxFun +simplifyExtIxFun = traverse simplifyExtPrimExp + +isStaticIxFun :: ExtIxFun -> Maybe IxFun +isStaticIxFun = traverse $ traverse inst + where inst Ext{} = Nothing + inst (Free x) = Just x + +instance (Engine.Simplifiable d, Engine.Simplifiable ret) => + Engine.Simplifiable (MemInfo d u ret) where + simplify (MemPrim bt) = + return $ MemPrim bt + simplify (MemMem size space) = + MemMem <$> Engine.simplify size <*> pure space + simplify (MemArray bt shape u ret) = + MemArray bt <$> Engine.simplify shape <*> pure u <*> Engine.simplify ret + +instance (PP.Pretty (TypeBase (ShapeBase d) u), + PP.Pretty d, PP.Pretty u, PP.Pretty ret) => PP.Pretty (MemInfo d u ret) where + ppr (MemPrim bt) = PP.ppr bt + ppr (MemMem s DefaultSpace) = + PP.text "mem" <> PP.parens (PP.ppr s) + ppr (MemMem s (Space sp)) = + PP.text "mem" <> PP.parens (PP.ppr s) <> PP.text "@" <> PP.text sp + ppr (MemArray bt shape u ret) = + PP.ppr (Array bt shape u) <> PP.text "@" <> PP.ppr ret + +instance PP.Pretty (Param (MemInfo SubExp Uniqueness ret)) where + ppr = PP.ppr . fmap declTypeOf + +instance PP.Pretty (Param (MemInfo SubExp NoUniqueness ret)) where + ppr = PP.ppr . fmap typeOf + +instance PP.Pretty (PatElemT (MemInfo SubExp NoUniqueness ret)) where + ppr = PP.ppr . fmap typeOf + +-- | Memory information for an array bound somewhere in the program. +data MemBind = ArrayIn VName IxFun + -- ^ Located in this memory block with this index + -- function. + deriving (Show) + +instance Eq MemBind where + _ == _ = True + +instance Ord MemBind where + _ `compare` _ = EQ + +instance Rename MemBind where + rename = substituteRename + +instance Substitute MemBind where + substituteNames substs (ArrayIn ident ixfun) = + ArrayIn (substituteNames substs ident) (substituteNames substs ixfun) + +instance PP.Pretty MemBind where + ppr (ArrayIn mem ixfun) = + PP.text "@" <> PP.ppr mem <> PP.text "->" <> PP.ppr ixfun + +instance FreeIn MemBind where + freeIn (ArrayIn mem ixfun) = freeIn mem <> freeIn ixfun + +-- | A description of the memory properties of an array being returned +-- by an operation. +data MemReturn = ReturnsInBlock VName ExtIxFun + -- ^ The array is located in a memory block that is + -- already in scope. + | ReturnsNewBlock Space Int ExtSize ExtIxFun + -- ^ The operation returns a new (existential) block, + -- with an existential or known size. + deriving (Show) + +instance Eq MemReturn where + _ == _ = True + +instance Ord MemReturn where + _ `compare` _ = EQ + +instance Rename MemReturn where + rename = substituteRename + +instance Substitute MemReturn where + substituteNames substs (ReturnsInBlock ident ixfun) = + ReturnsInBlock (substituteNames substs ident) (substituteNames substs ixfun) + substituteNames substs (ReturnsNewBlock space i size ixfun) = + ReturnsNewBlock space i (substituteNames substs size) (substituteNames substs ixfun) + +instance FixExt MemReturn where + fixExt i (Var v) (ReturnsNewBlock _ j _ ixfun) + | j == i = ReturnsInBlock v $ fixExtIxFun i + (primExpFromSubExp int32 (Var v)) ixfun + fixExt i se (ReturnsNewBlock space j size ixfun) = + ReturnsNewBlock space j' (fixExt i se size) + (fixExtIxFun i (primExpFromSubExp int32 se) ixfun) + where j' | i < j = j-1 + | otherwise = j + fixExt i se (ReturnsInBlock mem ixfun) = + ReturnsInBlock mem (fixExtIxFun i (primExpFromSubExp int32 se) ixfun) + +fixExtIxFun :: Int -> PrimExp VName -> ExtIxFun -> ExtIxFun +fixExtIxFun i e = fmap $ replaceInPrimExp update + where update (Ext j) t | j > i = LeafExp (Ext $ j - 1) t + | j == i = fmap Free e + | otherwise = LeafExp (Ext j) t + update (Free x) t = LeafExp (Free x) t + +leafExp :: Int -> PrimExp (Ext a) +leafExp i = LeafExp (Ext i) int32 + +existentialiseIxFun :: [VName] -> IxFun -> ExtIxFun +existentialiseIxFun ctx = IxFun.substituteInIxFun ctx' . fmap (fmap Free) + where ctx' = M.map leafExp $ M.fromList $ zip (map Free ctx) [0..] + +instance PP.Pretty MemReturn where + ppr (ReturnsInBlock v ixfun) = + PP.parens $ PP.text (pretty v) <> PP.text "->" <> PP.ppr ixfun + ppr (ReturnsNewBlock space i size ixfun) = + PP.text ("?" ++ show i) <> space' <> PP.parens (PP.ppr size) + <> PP.text "->" <> PP.ppr ixfun + where space' = case space of DefaultSpace -> mempty + Space s -> PP.text $ "@" ++ s + +instance FreeIn MemReturn where + freeIn (ReturnsInBlock v ixfun) = freeIn v <> freeIn ixfun + freeIn _ = mempty + +instance Engine.Simplifiable MemReturn where + simplify (ReturnsNewBlock space i size ixfun) = + ReturnsNewBlock space i <$> Engine.simplify size <*> simplifyExtIxFun ixfun + simplify (ReturnsInBlock v ixfun) = + ReturnsInBlock <$> Engine.simplify v <*> simplifyExtIxFun ixfun + + +instance Engine.Simplifiable MemBind where + simplify (ArrayIn mem ixfun) = + ArrayIn <$> Engine.simplify mem <*> simplifyIxFun ixfun + +instance Engine.Simplifiable [FunReturns] where + simplify = mapM Engine.simplify + +-- | The memory return of an expression. An array is annotated with +-- @Maybe MemReturn@, which can be interpreted as the expression +-- either dictating exactly where the array is located when it is +-- returned (if 'Just'), or able to put it whereever the binding +-- prefers (if 'Nothing'). +-- +-- This is necessary to capture the difference between an expression +-- that is just an array-typed variable, in which the array being +-- "returned" is located where it already is, and a @copy@ expression, +-- whose entire purpose is to store an existing array in some +-- arbitrary location. This is a consequence of the design decision +-- never to have implicit memory copies. +type ExpReturns = MemInfo ExtSize NoUniqueness (Maybe MemReturn) + +-- | The return of a body, which must always indicate where +-- returned arrays are located. +type BodyReturns = MemInfo ExtSize NoUniqueness MemReturn + +-- | The memory return of a function, which must always indicate where +-- returned arrays are located. +type FunReturns = MemInfo ExtSize Uniqueness MemReturn + +maybeReturns :: MemInfo d u r -> MemInfo d u (Maybe r) +maybeReturns (MemArray bt shape u ret) = + MemArray bt shape u $ Just ret +maybeReturns (MemPrim bt) = + MemPrim bt +maybeReturns (MemMem size space) = + MemMem size space + +noUniquenessReturns :: MemInfo d u r -> MemInfo d NoUniqueness r +noUniquenessReturns (MemArray bt shape _ r) = + MemArray bt shape NoUniqueness r +noUniquenessReturns (MemPrim bt) = + MemPrim bt +noUniquenessReturns (MemMem size space) = + MemMem size space + +funReturnsToExpReturns :: FunReturns -> ExpReturns +funReturnsToExpReturns = noUniquenessReturns . maybeReturns + +bodyReturnsToExpReturns :: BodyReturns -> ExpReturns +bodyReturnsToExpReturns = noUniquenessReturns . maybeReturns + +instance TC.Checkable ExplicitMemory where + checkExpLore = return + checkBodyLore = return + checkFParamLore = checkMemInfo + checkLParamLore = checkMemInfo + checkLetBoundLore = checkMemInfo + checkRetType = mapM_ TC.checkExtType . retTypeValues + checkOp (Alloc size _) = TC.require [Prim int64] size + checkOp (Inner k) = TC.subCheck $ typeCheckKernel k + primFParam name t = return $ Param name (MemPrim t) + primLParam name t = return $ Param name (MemPrim t) + matchPattern = matchPatternToExp + matchReturnType = matchFunctionReturnType + matchBranchType = matchBranchReturnType + +instance TC.Checkable InKernel where + checkExpLore = return + checkBodyLore = return + checkFParamLore = checkMemInfo + checkLParamLore = checkMemInfo + checkLetBoundLore = checkMemInfo + checkRetType = mapM_ TC.checkExtType . retTypeValues + checkOp (Alloc size _) = TC.require [Prim int64] size + checkOp (Inner k) = typeCheckKernelExp k + primFParam name t = return $ Param name (MemPrim t) + primLParam name t = return $ Param name (MemPrim t) + matchPattern = matchPatternToExp + matchReturnType = matchFunctionReturnType + matchBranchType = matchBranchReturnType + +matchFunctionReturnType :: ExplicitMemorish lore => + [FunReturns] -> Result -> TC.TypeM lore () +matchFunctionReturnType rettype result = do + TC.matchExtReturnType (fromDecl <$> ts) result + scope <- askScope + result_ts <- runReaderT (mapM subExpMemInfo result) $ removeScopeAliases scope + matchReturnType rettype result result_ts + mapM_ checkResultSubExp result + where ts = map declExtTypeOf rettype + checkResultSubExp Constant{} = + return () + checkResultSubExp (Var v) = do + attr <- varMemInfo v + case attr of + MemPrim _ -> return () + MemMem{} -> return () + MemArray _ _ _ (ArrayIn _ ixfun) + | IxFun.isLinear ixfun -> + return () + | otherwise -> + TC.bad $ TC.TypeError $ + "Array " ++ pretty v ++ + " returned by function, but has nontrivial index function " ++ + pretty ixfun ++ " " ++ show ixfun + +matchBranchReturnType :: ExplicitMemorish lore => + [BodyReturns] + -> Body (Aliases lore) + -> TC.TypeM lore () +matchBranchReturnType rettype (Body _ stms res) = do + scope <- askScope + ts <- runReaderT (mapM subExpMemInfo res) $ removeScopeAliases (scope <> scopeOf stms) + matchReturnType rettype res ts + +-- | Helper function for index function unification. +-- +-- The first return value maps a VName (wrapped in 'Free') to its Int +-- (wrapped in 'Ext'). In case of duplicates, it is mapped to the +-- *first* Int that occurs. +-- +-- The second return value maps each Int (wrapped in an 'Ext') to a +-- 'LeafExp' 'Ext' with the Int at which its associated VName first +-- occurs. +getExtMaps :: [(VName,Int)] -> (M.Map (Ext VName) (PrimExp (Ext VName)), + M.Map (Ext VName) (PrimExp (Ext VName))) +getExtMaps ctx_lst_ids = + (M.map leafExp $ M.mapKeys Free $ M.fromListWith (flip const) ctx_lst_ids, + M.fromList $ + mapMaybe (traverse (fmap (\i -> LeafExp (Ext i) int32) . + (`lookup` ctx_lst_ids)) . + uncurry (flip (,)) . fmap Ext) ctx_lst_ids) + +matchReturnType :: PP.Pretty u => + [MemInfo ExtSize u MemReturn] + -> [SubExp] + -> [MemInfo SubExp NoUniqueness MemBind] + -> TC.TypeM lore () +matchReturnType rettype res ts = do + let (ctx_ts, val_ts) = splitFromEnd (length rettype) ts + (ctx_res, _val_res) = splitFromEnd (length rettype) res + + getId :: (SubExp,Int) -> Maybe (VName,Int) + getId (Var ii, i) = Just (ii,i) + getId (Constant _, _) = Nothing + + (ctx_map_ids, ctx_map_exts) = + getExtMaps $ mapMaybe getId $ zip ctx_res [0..length ctx_res - 1] + + existentialiseIxFun0 :: IxFun -> ExtIxFun + existentialiseIxFun0 = IxFun.substituteInIxFun ctx_map_ids . fmap (fmap Free) + + getCt :: (Int,SubExp) -> Maybe (Ext VName, PrimExp (Ext VName)) + getCt (_, Var _) = Nothing + getCt (i, Constant c) = Just (Ext i, ValueExp c) + + ctx_map_cts = M.fromList $ mapMaybe getCt $ + zip [0..length ctx_res - 1] ctx_res + + substConstsInExtIndFun :: ExtIxFun -> ExtIxFun + substConstsInExtIndFun = IxFun.substituteInIxFun (ctx_map_cts<>ctx_map_exts) + + fetchCtx i = case maybeNth i $ zip ctx_res ctx_ts of + Nothing -> throwError $ "Cannot find context variable " ++ + show i ++ " in context results: " ++ pretty ctx_res + Just (se, t) -> return (se, t) + + checkReturn (MemPrim x) (MemPrim y) + | x == y = return () + checkReturn (MemMem x _) (MemMem y _) = + checkDim x y + checkReturn (MemArray x_pt x_shape _ x_ret) + (MemArray y_pt y_shape _ y_ret) + | x_pt == y_pt, shapeRank x_shape == shapeRank y_shape = do + zipWithM_ checkDim (shapeDims x_shape) (shapeDims y_shape) + checkMemReturn x_ret y_ret + checkReturn x y = + throwError $ unwords ["Expected ", pretty x, " but got ", pretty y] + + checkDim (Free x) y + | x == y = return () + | otherwise = throwError $ unwords ["Expected dim", pretty x, + "but got", pretty y] + checkDim (Ext i) y = do + (x, _) <- fetchCtx i + unless (x == y) $ + throwError $ unwords ["Expected ext dim", pretty i, "=>", pretty x, + "but got", pretty y] + + checkMemReturn (ReturnsInBlock x_mem x_ixfun) (ArrayIn y_mem y_ixfun) + | x_mem == y_mem = do + let x_ixfun' = substConstsInExtIndFun x_ixfun + y_ixfun' = existentialiseIxFun0 y_ixfun + unless (x_ixfun' == y_ixfun') $ + throwError $ unwords ["Index function unification fails1!", + "\nixfun of body result: ", pretty y_ixfun', + "\nixfun of return type: ", pretty x_ixfun', + "\nand context elements: ", pretty ctx_res] + checkMemReturn (ReturnsNewBlock x_space x_ext x_mem_size x_ixfun) + (ArrayIn y_mem y_ixfun) = do + (x_mem, x_mem_type) <- fetchCtx x_ext + let x_ixfun' = substConstsInExtIndFun x_ixfun + y_ixfun' = existentialiseIxFun0 y_ixfun + unless (x_ixfun' == y_ixfun') $ + throwError $ unwords ["Index function unification fails2!", + "\nixfun of body result: ", pretty y_ixfun', + "\nixfun of return type: ", pretty x_ixfun', + "\nand context elements: ", pretty ctx_res] + case x_mem_type of + MemMem y_mem_size y_space -> do + unless (x_mem == Var y_mem) $ + throwError $ unwords ["Expected memory", pretty x_ext, "=>", pretty x_mem, + "but got", pretty y_mem] + unless (x_space == y_space) $ + throwError $ unwords ["Expected memory", pretty y_mem, "in space", pretty x_space, + "but actually in space", pretty y_space] + checkDim x_mem_size y_mem_size + t -> + throwError $ unwords ["Expected memory", pretty x_ext, "=>", pretty x_mem, + "but but has type", pretty t] + checkMemReturn x y = + throwError $ unwords ["Expected array in", pretty x, + "but array returned in", pretty y] + + bad :: String -> TC.TypeM lore a + bad s = TC.bad $ TC.TypeError $ + unlines [ "Return type" + , " " ++ prettyTuple rettype + , "cannot match returns of results" + , " " ++ prettyTuple ts + , s + ] + + either bad return =<< runExceptT (zipWithM_ checkReturn rettype val_ts) + +matchPatternToExp :: (ExplicitMemorish lore) => + Pattern (Aliases lore) + -> Exp (Aliases lore) + -> TC.TypeM lore () +matchPatternToExp pat e = do + scope <- asksScope removeScopeAliases + rt <- runReaderT (expReturns $ removeExpAliases e) scope + + let (ctxs, vals) = bodyReturnsFromPattern $ removePatternAliases pat + (ctx_ids, _ctx_ts) = unzip ctxs + (_val_ids, val_ts) = unzip vals + (ctx_map_ids, ctx_map_exts) = + getExtMaps $ zip ctx_ids [0..length ctx_ids - 1] + + unless (length val_ts == length rt && + and (zipWith (matches ctx_map_ids ctx_map_exts) val_ts rt)) $ + TC.bad $ TC.TypeError $ "Expression type:\n " ++ prettyTuple rt ++ + "\ncannot match pattern type:\n " ++ prettyTuple val_ts ++ + "\nwith context elements: " ++ pretty ctx_ids + where matches _ _ (MemPrim x) (MemPrim y) = x == y + matches _ _ (MemMem x_size x_space) (MemMem y_size y_space) = + x_size == y_size && x_space == y_space + matches ctxids ctxexts (MemArray x_pt x_shape _ x_ret) (MemArray y_pt y_shape _ y_ret) = + x_pt == y_pt && x_shape == y_shape && + case (x_ret, y_ret) of + (ReturnsInBlock x_mem x_ixfun, Just (ReturnsInBlock y_mem y_ixfun)) -> + let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun + y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun + in x_mem == y_mem && x_ixfun' == y_ixfun' + (ReturnsInBlock _ x_ixfun, + Just (ReturnsNewBlock _ _ _ y_ixfun)) -> + let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun + y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun + in x_ixfun' == y_ixfun' + (ReturnsNewBlock x_space x_i x_size x_ixfun, + Just (ReturnsNewBlock y_space y_i y_size y_ixfun)) -> + let x_ixfun' = IxFun.substituteInIxFun ctxids x_ixfun + y_ixfun' = IxFun.substituteInIxFun ctxexts y_ixfun + in x_space == y_space && x_i == y_i && + x_size == y_size && x_ixfun' == y_ixfun' + (_, Nothing) -> True + _ -> False + matches _ _ _ _ = False + +varMemInfo :: ExplicitMemorish lore => + VName -> TC.TypeM lore (MemInfo SubExp NoUniqueness MemBind) +varMemInfo name = do + attr <- TC.lookupVar name + + case attr of + LetInfo (_, summary) -> return summary + FParamInfo summary -> return $ noUniquenessReturns summary + LParamInfo summary -> return summary + IndexInfo it -> return $ MemPrim $ IntType it + +nameInfoToMemInfo :: ExplicitMemorish lore => NameInfo lore -> MemBound NoUniqueness +nameInfoToMemInfo info = + case info of + FParamInfo summary -> noUniquenessReturns summary + LParamInfo summary -> summary + LetInfo summary -> summary + IndexInfo it -> MemPrim $ IntType it + +lookupMemInfo :: (HasScope lore m, ExplicitMemorish lore) => + VName -> m (MemInfo SubExp NoUniqueness MemBind) +lookupMemInfo = fmap nameInfoToMemInfo . lookupInfo + +subExpMemInfo :: (HasScope lore m, Monad m, ExplicitMemorish lore) => + SubExp -> m (MemInfo SubExp NoUniqueness MemBind) +subExpMemInfo (Var v) = lookupMemInfo v +subExpMemInfo (Constant v) = return $ MemPrim $ primValueType v + +lookupArraySummary :: (ExplicitMemorish lore, HasScope lore m, Monad m) => + VName -> m (VName, IxFun.IxFun (PrimExp VName)) +lookupArraySummary name = do + summary <- lookupMemInfo name + case summary of + MemArray _ _ _ (ArrayIn mem ixfun) -> + return (mem, ixfun) + _ -> + fail $ "Variable " ++ pretty name ++ " does not look like an array." + +lookupMemSize :: (HasScope lore m, Monad m) => + VName -> m SubExp +lookupMemSize v = do + t <- lookupType v + case t of Mem size _ -> return size + _ -> fail $ "lookupMemSize: " ++ pretty v ++ " is not a memory block." + +checkMemInfo :: TC.Checkable lore => + VName -> MemInfo SubExp u MemBind + -> TC.TypeM lore () +checkMemInfo _ (MemPrim _) = return () +checkMemInfo _ (MemMem size _) = + TC.require [Prim int64] size +checkMemInfo name (MemArray _ shape _ (ArrayIn v ixfun)) = do + t <- lookupType v + case t of + Mem{} -> + return () + _ -> + TC.bad $ TC.TypeError $ + "Variable " ++ pretty v ++ + " used as memory block, but is of type " ++ + pretty t ++ "." + + TC.context ("in index function " ++ pretty ixfun) $ do + traverse_ (TC.requirePrimExp int32) ixfun + let ixfun_rank = IxFun.rank ixfun + ident_rank = shapeRank shape + unless (ixfun_rank == ident_rank) $ + TC.bad $ TC.TypeError $ + "Arity of index function (" ++ pretty ixfun_rank ++ + ") does not match rank of array " ++ pretty name ++ + " (" ++ show ident_rank ++ ")" + +instance Attributes ExplicitMemory where + expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern + +instance Attributes InKernel where + expTypesFromPattern = return . map snd . snd . bodyReturnsFromPattern + +bodyReturnsFromPattern :: PatternT (MemBound NoUniqueness) + -> ([(VName,BodyReturns)], [(VName,BodyReturns)]) +bodyReturnsFromPattern pat = + (map asReturns $ patternContextElements pat, + map asReturns $ patternValueElements pat) + where ctx = patternContextElements pat + + ext (Var v) + | Just (i, _) <- find ((==v) . patElemName . snd) $ zip [0..] ctx = + Ext i + ext se = Free se + + asReturns pe = + (patElemName pe, + case patElemAttr pe of + MemPrim pt -> MemPrim pt + MemMem size space -> MemMem (ext size) space + MemArray pt shape u (ArrayIn mem ixfun) -> + MemArray pt (Shape $ map ext $ shapeDims shape) u $ + case find ((==mem) . patElemName . snd) $ zip [0..] ctx of + Just (i, PatElem _ (MemMem size space)) -> + ReturnsNewBlock space i (ext size) $ + existentialiseIxFun (map patElemName ctx) ixfun + _ -> ReturnsInBlock mem $ existentialiseIxFun [] ixfun + ) + +instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (PatElemT (MemInfo SubExp u r)) where + ppAnnot = bindeeAnnot patElemName patElemAttr + +instance (PP.Pretty u, PP.Pretty r) => PrettyAnnot (ParamT (MemInfo SubExp u r)) where + ppAnnot = bindeeAnnot paramName paramAttr + +instance PrettyLore ExplicitMemory where +instance PrettyLore InKernel where + +bindeeAnnot :: (PP.Pretty u, PP.Pretty r) => + (a -> VName) -> (a -> MemInfo SubExp u r) + -> a -> Maybe PP.Doc +bindeeAnnot bindeeName bindeeLore bindee = + case bindeeLore bindee of + attr@MemArray{} -> + Just $ + PP.text "-- " <> + PP.oneLine (PP.ppr (bindeeName bindee) <> + PP.text " : " <> + PP.ppr attr) + MemMem {} -> + Nothing + MemPrim _ -> + Nothing + +extReturns :: [ExtType] -> [ExpReturns] +extReturns ts = + evalState (mapM addAttr ts) 0 + where addAttr (Prim bt) = + return $ MemPrim bt + addAttr (Mem size space) = + return $ MemMem (Free size) space + addAttr t@(Array bt shape u) + | existential t = do + i <- get <* modify (+2) + return $ MemArray bt shape u $ Just $ + ReturnsNewBlock DefaultSpace (i+1) (Ext i) $ + IxFun.iota $ map convert $ shapeDims shape + | otherwise = + return $ MemArray bt shape u Nothing + convert (Ext i) = LeafExp (Ext i) int32 + convert (Free v) = Free <$> primExpFromSubExp int32 v + +arrayVarReturns :: (HasScope lore m, Monad m, ExplicitMemorish lore) => + VName + -> m (PrimType, Shape, VName, IxFun.IxFun (PrimExp VName)) +arrayVarReturns v = do + summary <- lookupMemInfo v + case summary of + MemArray et shape _ (ArrayIn mem ixfun) -> + return (et, Shape $ shapeDims shape, mem, ixfun) + _ -> + fail $ "arrayVarReturns: " ++ pretty v ++ " is not an array." + +varReturns :: (HasScope lore m, Monad m, ExplicitMemorish lore) => + VName -> m ExpReturns +varReturns v = do + summary <- lookupMemInfo v + case summary of + MemPrim bt -> + return $ MemPrim bt + MemArray et shape _ (ArrayIn mem ixfun) -> + return $ MemArray et (fmap Free shape) NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun + MemMem size space -> + return $ MemMem (Free size) space + +-- | The return information of an expression. This can be seen as the +-- "return type with memory annotations" of the expression. +expReturns :: (Monad m, HasScope lore m, + ExplicitMemorish lore) => + Exp lore -> m [ExpReturns] + +expReturns (BasicOp (SubExp (Var v))) = + pure <$> varReturns v + +expReturns (BasicOp (Opaque (Var v))) = + pure <$> varReturns v + +expReturns (BasicOp (Repeat outer_shapes inner_shape v)) = do + t <- repeatDims outer_shapes inner_shape <$> lookupType v + (et, _, mem, ixfun) <- arrayVarReturns v + let outer_shapes' = map (map (primExpFromSubExp int32) . shapeDims) outer_shapes + inner_shape' = map (primExpFromSubExp int32) $ shapeDims inner_shape + return [MemArray et (Shape $ map Free $ arrayDims t) NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] $ + IxFun.repeat ixfun outer_shapes' inner_shape'] + +expReturns (BasicOp (Reshape newshape v)) = do + (et, _, mem, ixfun) <- arrayVarReturns v + return [MemArray et (Shape $ map (Free . newDim) newshape) NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] $ + IxFun.reshape ixfun $ map (fmap $ primExpFromSubExp int32) newshape] + +expReturns (BasicOp (Rearrange perm v)) = do + (et, Shape dims, mem, ixfun) <- arrayVarReturns v + let ixfun' = IxFun.permute ixfun perm + dims' = rearrangeShape perm dims + return [MemArray et (Shape $ map Free dims') NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun'] + +expReturns (BasicOp (Rotate offsets v)) = do + (et, Shape dims, mem, ixfun) <- arrayVarReturns v + let offsets' = map (primExpFromSubExp int32) offsets + ixfun' = IxFun.rotate ixfun offsets' + return [MemArray et (Shape $ map Free dims) NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun'] + +expReturns (BasicOp (Index v slice)) = do + info <- sliceInfo v slice + case info of + MemArray et shape u (ArrayIn mem ixfun) -> + return [MemArray et (fmap Free shape) u $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun] + MemPrim pt -> return [MemPrim pt] + MemMem d space -> return [MemMem (Free d) space] + +expReturns (BasicOp (Update v _ _)) = + pure <$> varReturns v + +expReturns (BasicOp op) = + extReturns . staticShapes <$> primOpType op + +expReturns (DoLoop ctx val _ _) = + zipWithM typeWithAttr + (loopExtType (map (paramIdent . fst) ctx) (map (paramIdent . fst) val)) $ map fst val + where typeWithAttr t p = + case (t, paramAttr p) of + (Array bt shape u, MemArray _ _ _ (ArrayIn mem ixfun)) + | Just (i, mem_p) <- isMergeVar mem, + Mem mem_size space <- paramType mem_p -> + let ext_size + | Just (j, _) <- isMergeVar =<< subExpVar mem_size = Ext j + | otherwise = Free mem_size + in return $ MemArray bt shape u $ Just $ ReturnsNewBlock space i ext_size ixfun' + | otherwise -> + return (MemArray bt shape u $ + Just $ ReturnsInBlock mem ixfun') + where ixfun' = existentialiseIxFun (map paramName mergevars) ixfun + (Array{}, _) -> + fail "expReturns: Array return type but not array merge variable." + (Prim bt, _) -> + return $ MemPrim bt + (Mem{}, _) -> + fail "expReturns: loop returns memory block explicitly." + isMergeVar v = find ((==v) . paramName . snd) $ zip [0..] mergevars + mergevars = map fst $ ctx ++ val + +expReturns (Apply _ _ ret _) = + return $ map funReturnsToExpReturns ret + +expReturns (If _ _ _ (IfAttr ret _)) = + return $ map bodyReturnsToExpReturns ret + +expReturns (Op op) = + opReturns op + +sliceInfo :: (Monad m, HasScope lore m, ExplicitMemorish lore) => + VName + -> Slice SubExp -> m (MemInfo SubExp NoUniqueness MemBind) +sliceInfo v slice = do + (et, _, mem, ixfun) <- arrayVarReturns v + case sliceDims slice of + [] -> return $ MemPrim et + dims -> + return $ MemArray et (Shape dims) NoUniqueness $ + ArrayIn mem $ IxFun.slice ixfun + (map (fmap (primExpFromSubExp int32)) slice) + +class TypedOp (Op lore) => OpReturns lore where + opReturns :: (Monad m, HasScope lore m) => + Op lore -> m [ExpReturns] + opReturns op = extReturns <$> opType op + +instance OpReturns ExplicitMemory where + opReturns (Alloc size space) = + return [MemMem (Free size) space] + opReturns (Inner k@(Kernel _ _ _ body)) = + zipWithM correct (kernelBodyResult body) =<< (extReturns <$> opType k) + where correct (WriteReturn _ arr _) _ = varReturns arr + correct (KernelInPlaceReturn arr) _ = + extendedScope (varReturns arr) + (castScope $ scopeOf $ kernelBodyStms body) + correct _ ret = return ret + opReturns k = + extReturns <$> opType k + +instance OpReturns InKernel where + opReturns (Alloc size space) = + return [MemMem (Free size) space] + + opReturns (Inner (GroupStream _ _ lam _ _)) = + forM (groupStreamAccParams lam) $ \param -> + case paramAttr param of + MemPrim bt -> + return $ MemPrim bt + MemArray et shape _ (ArrayIn mem ixfun) -> + return $ MemArray et (Shape $ map Free $ shapeDims shape) NoUniqueness $ + Just $ ReturnsInBlock mem $ existentialiseIxFun [] ixfun + MemMem size space -> + return $ MemMem (Free size) space + + opReturns (Inner (GroupScan _ _ input)) = + mapM varReturns arrs + where arrs = map snd input + + opReturns (Inner (GroupGenReduce _ dests _ _ _ _)) = + mapM varReturns dests + + opReturns (Inner (Barrier res)) = mapM f res + where f (Var v) = varReturns v + f (Constant v) = return $ MemPrim $ primValueType v + + opReturns (Inner (Combine (CombineSpace scatter cspace) ts _ _)) = + (++) <$> mapM varReturns as <*> + pure (extReturns $ staticShapes $ map (`arrayOfShape` shape) $ drop (sum ns*2) ts) + where (_, ns, as) = unzip3 scatter + shape = Shape $ map snd cspace + + opReturns k = + extReturns <$> opType k + +applyFunReturns :: Typed attr => + [FunReturns] + -> [Param attr] + -> [(SubExp,Type)] + -> Maybe [FunReturns] +applyFunReturns rets params args + | Just _ <- applyRetType rettype params args = + Just $ map correctDims rets + | otherwise = + Nothing + where rettype = map declExtTypeOf rets + parammap :: M.Map VName (SubExp, Type) + parammap = M.fromList $ + zip (map paramName params) args + + substSubExp (Var v) + | Just (se,_) <- M.lookup v parammap = se + substSubExp se = se + + correctDims (MemPrim t) = + MemPrim t + correctDims (MemMem (Free se) space) = + MemMem (Free $ substSubExp se) space + correctDims (MemMem (Ext d) space) = + MemMem (Ext d) space + correctDims (MemArray et shape u memsummary) = + MemArray et (correctShape shape) u $ + correctSummary memsummary + + correctShape = Shape . map correctDim . shapeDims + correctDim (Ext i) = Ext i + correctDim (Free se) = Free $ substSubExp se + + correctSummary (ReturnsNewBlock space i size ixfun) = + ReturnsNewBlock space i size ixfun + correctSummary (ReturnsInBlock mem ixfun) = + -- FIXME: we should also do a replacement in ixfun here. + ReturnsInBlock mem' ixfun + where mem' = case M.lookup mem parammap of + Just (Var v, _) -> v + _ -> mem + +-- | Is an array of the given shape stored fully flat row-major with +-- the given index function? +fullyLinear :: (Eq num, IntegralExp num) => + ShapeBase num -> IxFun.IxFun num -> Bool +fullyLinear shape ixfun = + IxFun.isLinear ixfun && ixFunMatchesInnerShape shape ixfun + +ixFunMatchesInnerShape :: (Eq num, IntegralExp num) => + ShapeBase num -> IxFun.IxFun num -> Bool +ixFunMatchesInnerShape shape ixfun = + drop 1 (IxFun.shape ixfun) == drop 1 (shapeDims shape) diff --git a/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs @@ -0,0 +1,447 @@ +-- | An index function represents a mapping from an array index space +-- to a flat byte offset. +module Futhark.Representation.ExplicitMemory.IndexFunction + ( +-- IxFun(..) + IxFun + , index + , iota + , offsetIndex + , strideIndex + , permute + , rotate + , reshape + , slice + , base + , rebase + , repeat + , shape + , rank + , linearWithOffset + , rearrangeWithOffset + , isLinear + , isDirect + , substituteInIxFun + , getInfoMaxUnification + , subsInIndexIxFun + , ixFunsCompatibleRaw + , ixFunHasIndex + , offsetIndexDWIM + ) + where + +import Control.Arrow (first) +import Data.Maybe +import Data.Monoid ((<>)) +import Data.List hiding (repeat) +import Control.Monad.Identity +import Control.Monad.Writer + +import Prelude hiding (mod, repeat) + +import qualified Data.List as L +import qualified Data.Map.Strict as M + +import Futhark.Transform.Substitute +import Futhark.Transform.Rename + +import Futhark.Representation.AST.Syntax + (ShapeChange, DimChange(..), DimIndex(..), Slice, sliceDims, unitSlice, VName(..)) +import Futhark.Representation.AST.Attributes.Names +import Futhark.Representation.AST.Attributes.Reshape +import Futhark.Representation.AST.Attributes.Rearrange +import Futhark.Representation.AST.Pretty () +import Futhark.Util.IntegralExp +import Futhark.Util.Pretty +import Futhark.Util +import Futhark.Analysis.PrimExp.Convert + +type Shape num = [num] +type Indices num = [num] +type Permutation = [Int] + +data IxFun num = Direct (Shape num) + | Permute (IxFun num) Permutation + | Rotate (IxFun num) (Indices num) + | Index (IxFun num) (Slice num) + | Reshape (IxFun num) (ShapeChange num) + | Repeat (IxFun num) [Shape num] (Shape num) + deriving (Eq,Show) + +instance Pretty num => Pretty (IxFun num) where + ppr (Direct dims) = + text "Direct" <> parens (commasep $ map ppr dims) + ppr (Permute fun perm) = ppr fun <> ppr perm + ppr (Rotate fun offsets) = ppr fun <> brackets (commasep $ map ((text "+" <>) . ppr) offsets) + ppr (Index fun is) = ppr fun <> brackets (commasep $ map ppr is) + ppr (Reshape fun oldshape) = + ppr fun <> text "->reshape" <> + parens (commasep (map ppr oldshape)) + ppr (Repeat fun outer_shapes inner_shape) = + ppr fun <> text "->repeat" <> parens (commasep (map ppr $ outer_shapes++ [inner_shape])) + +instance Substitute num => Substitute (IxFun num) where + substituteNames substs = fmap $ substituteNames substs + +instance FreeIn num => FreeIn (IxFun num) where + freeIn = foldMap freeIn + +instance Functor IxFun where + fmap f = runIdentity . traverse (return . f) + +instance Foldable IxFun where + foldMap f = execWriter . traverse (tell . f) + +instance Traversable IxFun where + traverse f (Direct dims) = + Direct <$> traverse f dims + traverse f (Permute ixfun perm) = + Permute <$> traverse f ixfun <*> pure perm + traverse f (Rotate ixfun offsets) = + Rotate <$> traverse f ixfun <*> traverse f offsets + traverse f (Index ixfun is) = + Index <$> traverse f ixfun <*> traverse (traverse f) is + traverse f (Reshape ixfun dims) = + Reshape <$> traverse f ixfun <*> traverse (traverse f) dims + traverse f (Repeat ixfun outer_shapes inner_shape) = + Repeat <$> traverse f ixfun <*> + traverse (traverse f) outer_shapes <*> + traverse f inner_shape + +instance Substitute num => Rename (IxFun num) where + rename = substituteRename + +index :: (Pretty num, IntegralExp num) => + IxFun num -> Indices num -> num -> num + +index (Direct dims) is element_size = + sum (zipWith (*) is slicesizes) * element_size + where slicesizes = drop 1 $ sliceSizes dims + +index (Permute fun perm) is_new element_size = + index fun is_old element_size + where is_old = rearrangeShape (rearrangeInverse perm) is_new + +index (Rotate fun offsets) is element_size = + index fun (zipWith mod (zipWith (+) is offsets) dims) element_size + where dims = shape fun + +index (Index fun js) is element_size = + index fun (adjust js is) element_size + where adjust (DimFix j:js') is' = j : adjust js' is' + adjust (DimSlice j _ s:js') (i:is') = j + i * s : adjust js' is' + adjust _ _ = [] + +index (Reshape fun newshape) is element_size = + let new_indices = reshapeIndex (shape fun) (newDims newshape) is + in index fun new_indices element_size + +index (Repeat fun outer_shapes _) is element_size = + -- Discard those indices that are just repeats. It is intentional + -- that we cut off those indices that correspond to the innermost + -- repeated dimensions. + index fun is' element_size + where flags dims = replicate (length dims) True ++ [False] + is' = map snd $ filter (not . fst) $ zip (concatMap flags outer_shapes) is + +iota :: Shape num -> IxFun num +iota = Direct + +offsetIndex :: (Eq num, IntegralExp num) => + IxFun num -> num -> IxFun num +offsetIndex ixfun i | i == 0 = ixfun +offsetIndex ixfun i = + case shape ixfun of + d:ds -> slice ixfun (DimSlice i (d-i) 1 : map (unitSlice 0) ds) + [] -> error "offsetIndex: underlying index function has rank zero" + +strideIndex :: (Eq num, IntegralExp num) => + IxFun num -> num -> IxFun num +strideIndex ixfun s = + case shape ixfun of + d:ds -> slice ixfun (DimSlice (fromInt32 0) d s : map (unitSlice (fromInt32 0)) ds) + [] -> error "offsetIndex: underlying index function has rank zero" + +permute :: IntegralExp num => + IxFun num -> Permutation -> IxFun num +permute (Permute ixfun oldperm) perm + | rearrangeInverse oldperm == perm = ixfun + | otherwise = permute ixfun (rearrangeCompose perm oldperm) +permute ixfun perm + | perm == sort perm = ixfun + | otherwise = Permute ixfun perm + +rotate :: IntegralExp num => + IxFun num -> Indices num -> IxFun num +rotate (Rotate ixfun old_offsets) offsets = + Rotate ixfun $ zipWith (+) old_offsets offsets +rotate ixfun offsets = Rotate ixfun offsets + +repeat :: IxFun num -> [Shape num] -> Shape num -> IxFun num +repeat = Repeat + +reshape :: (Eq num, IntegralExp num) => + IxFun num -> ShapeChange num -> IxFun num + +reshape Direct{} newshape = + Direct $ map newDim newshape + +reshape (Reshape ixfun _) newshape = + reshape ixfun newshape + +reshape (Permute ixfun perm) newshape + | Just (head_coercions, reshapes, tail_coercions) <- + splitCoercions newshape, + num_coercions <- length (head_coercions ++ tail_coercions), + (head_perms, mid_perms, end_perms) <- + splitAt3 (length head_coercions) (length perm - num_coercions) perm, + sequential mid_perms, + first_reshaped <- foldl min (rank ixfun) mid_perms, + extra_dims <- length newshape - length (shape ixfun), + perm' <- map (shiftDim first_reshaped extra_dims) head_perms ++ + take (length reshapes) [first_reshaped..] ++ + map (shiftDim first_reshaped extra_dims) end_perms, + newshape' <- rearrangeShape (rearrangeInverse perm') newshape = + Permute (reshape ixfun newshape') perm' + where splitCoercions newshape' = do + let (head_coercions, newshape'') = span isCoercion newshape' + let (reshapes, tail_coercions) = break isCoercion newshape'' + guard (all isCoercion tail_coercions) + return (head_coercions, reshapes, tail_coercions) + + isCoercion DimCoercion{} = True + isCoercion _ = False + + shiftDim last_reshaped extra_dims x + | x > last_reshaped = x + extra_dims + | otherwise = x + + sequential [] = True + sequential (x:xs) = and $ zipWith (==) xs [x+1, x+2..] + +reshape (Index ixfun slicing) newshape + | [newdim] <- newDims newshape, + Just slicing' <- findSlice slicing (Just newdim) = + Index ixfun slicing' + | (is, rem_slicing) <- splitSlice slicing, + (fixed_ds, sliced_ds) <- splitAt (length is) $ shape ixfun, + and $ zipWith isSliceOf rem_slicing sliced_ds = + -- Move the reshape beneath the slicing. + let newshape' = map DimCoercion fixed_ds ++ newshape + in Index (reshape ixfun newshape') $ + map DimFix is ++ map (unitSlice (fromInt32 0)) (newDims newshape) + where isSliceOf (DimSlice _ d1 1) d2 = d1 == d2 + isSliceOf _ _ = False + + findSlice (DimFix i:is) d = (DimFix i:) <$> findSlice is d + findSlice (DimSlice j _ stride:is) d = do + d' <- d + (DimSlice j d' stride:) <$> findSlice is Nothing + findSlice [] Just{} = Nothing + findSlice [] Nothing = Just [] + +reshape ixfun newshape + | shape ixfun == map newDim newshape = + ixfun + | rank ixfun == length newshape, + Just _ <- shapeCoercion newshape = + ixfun + | otherwise = + Reshape ixfun newshape + +splitSlice :: Slice num -> ([num], Slice num) +splitSlice [] = ([], []) +splitSlice (DimFix i:is) = first (i:) $ splitSlice is +splitSlice is = ([], is) + +slice :: (Eq num, IntegralExp num) => + IxFun num -> Slice num -> IxFun num +slice ixfun is + -- Avoid identity slicing. + | is == map (unitSlice 0) (shape ixfun) = ixfun +slice (Index ixfun mis) is = + Index ixfun $ reslice mis is + where reslice mis' [] = mis' + reslice (DimFix j:mis') is' = + DimFix j : reslice mis' is' + reslice (DimSlice orig_k _ orig_s:mis') (DimSlice new_k n new_s:is') = + DimSlice (orig_k + new_k * orig_s) n (orig_s*new_s) : reslice mis' is' + reslice (DimSlice orig_k _ orig_s:mis') (DimFix i:is') = + DimFix (orig_k+i*orig_s) : reslice mis' is' + reslice _ _ = error "IndexFunction slice: invalid arguments" +slice ixfun [] = ixfun +slice ixfun is = Index ixfun is + +rank :: IntegralExp num => + IxFun num -> Int +rank = length . shape + +shape :: IntegralExp num => + IxFun num -> Shape num +shape (Direct dims) = + dims +shape (Permute ixfun perm) = + rearrangeShape perm $ shape ixfun +shape (Rotate ixfun _) = + shape ixfun +shape (Index _ how) = + sliceDims how +shape (Reshape _ dims) = + map newDim dims +shape (Repeat ixfun outer_shapes inner_shape) = + concat (zipWith repeated outer_shapes (shape ixfun)) ++ inner_shape + where repeated outer_ds d = outer_ds ++ [d] + +base :: IxFun num -> Shape num +base (Direct dims) = + dims +base (Permute ixfun _) = + base ixfun +base (Rotate ixfun _) = + base ixfun +base (Index ixfun _) = + base ixfun +base (Reshape ixfun _) = + base ixfun +base (Repeat ixfun _ _) = + base ixfun + +rebase :: (Eq num, IntegralExp num) => + IxFun num + -> IxFun num + -> IxFun num +rebase new_base (Direct old_shape) + | old_shape == shape new_base = new_base + | otherwise = reshape new_base $ map DimCoercion old_shape +rebase new_base (Permute ixfun perm) = + permute (rebase new_base ixfun) perm +rebase new_base (Rotate ixfun offsets) = + rotate (rebase new_base ixfun) offsets +rebase new_base (Index ixfun is) = + slice (rebase new_base ixfun) is +rebase new_base (Reshape ixfun new_shape) = + reshape (rebase new_base ixfun) new_shape +rebase new_base (Repeat ixfun outer_shapes inner_shape) = + Repeat (rebase new_base ixfun) outer_shapes inner_shape + +-- This function does not cover all possible cases. It's a "best +-- effort" kind of thing. +linearWithOffset :: (Eq num, IntegralExp num) => + IxFun num -> num -> Maybe num +linearWithOffset (Direct _) _ = + Just 0 +linearWithOffset (Reshape ixfun _) element_size = + linearWithOffset ixfun element_size +linearWithOffset (Index ixfun is) element_size = do + is' <- fixingOuter is inner_shape + inner_offset <- linearWithOffset ixfun element_size + let slices = take m $ drop 1 $ sliceSizes $ shape ixfun + return $ inner_offset + sum (zipWith (*) slices is') * element_size + where m = length is + inner_shape = shape ixfun + fixingOuter (DimFix i:is') (_:ds) = (i:) <$> fixingOuter is' ds + fixingOuter (DimSlice off _ 1:is') (_:ds) + | is' == map (unitSlice 0) ds = Just [off] + fixingOuter is' ds + | is' == map (unitSlice 0) ds = Just [] + fixingOuter _ _ = Nothing +linearWithOffset _ _ = Nothing + +rearrangeWithOffset :: (Eq num, IntegralExp num) => + IxFun num -> num -> Maybe (num, [(Int,num)]) +rearrangeWithOffset (Reshape ixfun _) element_size = + rearrangeWithOffset ixfun element_size +rearrangeWithOffset (Permute ixfun perm) element_size = do + offset <- linearWithOffset ixfun element_size + return (offset, zip perm $ rearrangeShape perm $ shape ixfun) +rearrangeWithOffset _ _ = + Nothing + +isLinear :: (Eq num, IntegralExp num) => IxFun num -> Bool +isLinear = + (==Just 0) . flip linearWithOffset 1 + +isDirect :: IxFun num -> Bool +isDirect Direct{} = True +isDirect _ = False + +-- | Substituting a name with a PrimExp in an index function. +substituteInIxFun :: (Ord a) => M.Map a (PrimExp a) -> IxFun (PrimExp a) + -> IxFun (PrimExp a) +substituteInIxFun tab (Direct pes) = + Direct $ map (substituteInPrimExp tab) pes +substituteInIxFun tab (Permute ixfun p) = + Permute (substituteInIxFun tab ixfun) p +substituteInIxFun tab (Rotate ixfun pes) = + Rotate (substituteInIxFun tab ixfun) $ map (substituteInPrimExp tab) pes +substituteInIxFun tab (Index ixfun sl) = + Index (substituteInIxFun tab ixfun) $ map (fmap $ substituteInPrimExp tab) sl +substituteInIxFun tab (Reshape ixfun newshape) = + Reshape (substituteInIxFun tab ixfun) $ map (fmap $ substituteInPrimExp tab) newshape +substituteInIxFun tab (Repeat ixfun outer_shapes inner_shape) = + Repeat (substituteInIxFun tab ixfun) outer_shapes inner_shape + +----------------------------------------------------------- +--- Niels' functions for memory management: --- +--- these are prime candidates to be removed/re-written --- +----------------------------------------------------------- + +type IxFn = IxFun (PrimExp VName) + +getInfoMaxUnification :: IxFn -> Maybe (IxFn, Slice (PrimExp VName), VName) +getInfoMaxUnification (Index ixfun_start slc) = + case L.span isDimFix slc of + (indices_start, [DimSlice _start_offset + (LeafExp final_dim@VName{} (IntType Int32)) + _stride]) -> + Just (ixfun_start, indices_start, final_dim) + _ -> Nothing + where isDimFix DimFix{} = True + isDimFix _ = False +getInfoMaxUnification _ = Nothing + +-- Are two index functions *identical*? (Silly approach, but the Eq +-- instance is used for something else.) +ixFunsCompatibleRaw :: Eq num => IxFun num -> IxFun num -> Bool +ixFunsCompatibleRaw ixfun0 ixfun1 = ixfun0 `primEq` ixfun1 + where primEq a b = case (a, b) of + (Direct sa, Direct sb) -> + sa == sb + (Permute a1 pa, Permute b1 pb) -> + a1 `primEq` b1 && pa == pb + (Rotate a1 ia, Rotate b1 ib) -> + a1 `primEq` b1 && ia == ib + (Index a1 sa, Index b1 sb) -> + a1 `primEq` b1 && sa == sb + (Reshape a1 sa, Reshape b1 sb) -> + a1 `primEq` b1 && sa == sb + (Repeat a1 ssa sa, Repeat b1 ssb sb) -> + a1 `primEq` b1 && ssa == ssb && sa == sb + _ -> False + +ixFunHasIndex :: IxFun num -> Bool +ixFunHasIndex ixfun = case ixfun of + Direct _ -> False + Permute ixfun' _ -> ixFunHasIndex ixfun' + Rotate ixfun' _ -> ixFunHasIndex ixfun' + Index{} -> True + Reshape ixfun' _ -> ixFunHasIndex ixfun' + Repeat ixfun' _ _ -> ixFunHasIndex ixfun' + +subsInIndexIxFun :: IxFn -> VName -> VName -> IxFn +subsInIndexIxFun (Index ixfun_start slc) final_dim final_dim_max_v = + let tab = M.singleton final_dim (LeafExp final_dim_max_v (IntType Int32)) + ixfun_start' = substituteInIxFun tab ixfun_start + in Index ixfun_start' slc +subsInIndexIxFun _ _ _ = error "In IxFun.subsInIndexIxFun: should-not-happen case reached!" + +offsetIndexDWIM :: Int -> IxFn -> PrimExp VName -> IxFn +offsetIndexDWIM n_ignore_initial ixfun offset = + fromMaybe (offsetIndex ixfun offset) $ case ixfun of + Index ixfun1 dimindices -> + let (dim_first, dim_rest) = L.splitAt n_ignore_initial dimindices + in case dim_rest of + (DimFix i : dim_rest') -> + Just $ Index ixfun1 (dim_first ++ DimFix (i + offset) : dim_rest') + _ -> Nothing + _ -> Nothing diff --git a/src/Futhark/Representation/ExplicitMemory/Lmad.hs b/src/Futhark/Representation/ExplicitMemory/Lmad.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/ExplicitMemory/Lmad.hs @@ -0,0 +1,761 @@ +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} +-- | An index function represents a mapping from an array index space +-- to a flat byte offset. This implements a representation for the +-- index function based on linear-memory accessor descriptors, see +-- Zhu, Hoeflinger and David work. Our specific representation is: +-- LMAD = \overline{s,r,n}^k + o, where `o` is the offset, and `s_j`, +-- `r_j`, and `n_j` are the stride, the rotate factor and the number +-- of elements on dimension j. Dimensions are ordered in row major fashion. +-- By definition, the LMAD above denotes the set of points: +-- \{ o + \Sigma_{j=0}^{k} ((i_j+r_j) `mod` n_j)*s_j, +-- \forall i_j such that 0<=i_j Pretty (Lmad num) where + ppr (Lmad tau srnps) = + let (ss, rs, ns, ps, fs) = unzip5 srnps + in text " | " <> ppr tau <> + text " + " <> brackets (commasep $ map ppr ss) <> + text "v" <> brackets (commasep $ map ppr rs) <> + text "v" <> brackets (commasep $ map ppr ns) <> + text "v" <> brackets (commasep $ map ppr ps) <> + text "v" <> brackets (commasep $ map ppr fs) <> + text " | " + +instance Pretty num => Pretty (IxFun num) where + ppr (IxFun lmads orgshp cg) = + text "Shape: " <> braces (commasep $ map ppr orgshp) <> + text " LMADS: " <> braces (stack $ map ppr lmads) <> + text " CONTIG: "<> text (show cg) + +instance Substitute num => Substitute (Lmad num) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitute num => Substitute (IxFun num) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitute num => Rename (Lmad num) where + rename = substituteRename + +instance Substitute num => Rename (IxFun num) where + rename = substituteRename + + +instance FreeIn num => FreeIn (Lmad num) where + freeIn = foldMap freeIn + +instance FreeIn num => FreeIn (IxFun num) where + freeIn = foldMap freeIn + +instance Functor Lmad where + fmap f = runIdentity . traverse (return . f) + +instance Functor IxFun where + fmap f = runIdentity . traverse (return . f) + +instance Foldable Lmad where + foldMap f = execWriter . traverse (tell . f) + +instance Foldable IxFun where + foldMap f = execWriter . traverse (tell . f) + +instance Traversable Lmad where + traverse f (Lmad x l) = + Lmad <$> f x <*> traverse f' l + where f' (a, b, c, k, info) = + (,,,,) <$> f a <*> f b <*> f c <*> pure k <*> pure info + +instance Traversable IxFun where + traverse f (IxFun lmads shp cg) = + IxFun <$> traverse (traverse f) lmads <*> traverse f shp <*> pure cg + +-- | Substituting a name with a PrimExp in an Lmad. +substituteInLmad :: M.Map VName (PrimExp VName) -> Lmad (PrimExp VName) + -> Lmad (PrimExp VName) +substituteInLmad tab (Lmad off srnpds) = + let off' = substituteInPrimExp tab off + srnpds' = map (\(s,r,n,p,d) -> + ( substituteInPrimExp tab s + , substituteInPrimExp tab r + , substituteInPrimExp tab n + , p, d + ) + ) srnpds + in Lmad off' srnpds' + +-- | Substituting a name with a PrimExp in an index function. +substituteInIxFun :: M.Map VName (PrimExp VName) -> IxFun (PrimExp VName) + -> IxFun (PrimExp VName) +substituteInIxFun tab (IxFun lmads shp b) = + IxFun (map (substituteInLmad tab) lmads) + (map (substituteInPrimExp tab) shp) + b + +------------------------------------------ +--- Index Function/LMAD Implementation --- +------------------------------------------ + +-- | whether this is a row-major array +isDirect :: (Eq num, IntegralExp num) => IxFun num -> Bool +isDirect (IxFun [Lmad off info] shp True) + | length shp == length info, + all (\((s,r,n,p,_),i,d) -> s==1 && r==0 && n==d && p==i) + (zip3 info [0..length info - 1] shp), + off == 0 = True + | otherwise = False +isDirect _ = False + +-- | whether an index function has contiguous memory support +isContiguous :: (Eq num, IntegralExp num) => IxFun num -> Bool +isContiguous (IxFun _ _ cg) = cg + +-- | Shape of an Lmad +shape0 :: (Eq num, IntegralExp num) => Lmad num -> Shape num +shape0 lmad@(Lmad _ srns) = + map (\(_,_,z,_,_)->z) $ permuteInv (getPermutation lmad) srns + +-- | Shape of an index function +shape :: (Eq num, IntegralExp num) => IxFun num -> Shape num +shape (IxFun [] _ _) = error "shape: empty index function" +shape (IxFun (lmad:_) _ _) = shape0 lmad + +-- | Computing the flat memory index for a complete set `inds` +-- of array indices and a certain element size `elem_size`. +index :: (IntegralExp num, Eq num) => + IxFun num -> Indices num -> num -> num +index (IxFun [] _ _) _ _ = error "index: empty index function" +index (IxFun [lmad] _ _) iis elm_size = index0 lmad iis elm_size +index (IxFun (lmad1:lmad2:lmads) oshp c) iis elm_size = + let i_flat = index0 lmad1 iis 1 + new_inds = unflattenIndex (shape0 lmad2) i_flat + in index (IxFun (lmad2:lmads) oshp c) new_inds elm_size + +-- | Helper for index: computing the flat index of an Lmad. +index0 :: (Eq num, IntegralExp num) => + Lmad num -> Indices num -> num -> num +index0 lmad@(Lmad tau srnps) inds elm_size = + let prod = sum $ zipWith flatOneDim + (map (\(s,r,n,_,_) -> (s,r,n)) srnps) + (permuteInv (getPermutation lmad) inds) + ind = tau + prod + in if elm_size == 1 then ind else ind * elm_size + +-- | iota +iota :: (IntegralExp num) => Shape num -> IxFun num +iota ns = IxFun [makeRotIota Inc 0 $ zip rs ns] ns True + where rs = replicate (length ns) 0 + +-- | permute dimensions +permute :: IntegralExp num => + IxFun num -> Permutation -> IxFun num +permute (IxFun [] _ _) _ = error "permute: empty index function" +permute (IxFun (lmad:lmads) oshp cg) ps = + let perm = map (\p -> ps !! p) $ getPermutation lmad + in IxFun (setPermutation perm lmad : lmads) oshp cg + +-- | repeating dimensions +repeat :: (Eq num, IntegralExp num) => + IxFun num -> [Shape num] -> Shape num -> IxFun num +repeat (IxFun [] _ _) _ _ = error "repeat: empty index function" +repeat (IxFun (lmad@(Lmad tau srnps) : lmads) oshp cg) shps shp = + let perm = getPermutation lmad + -- inverse permute the shapes and update the permutation! + lens = map (\s -> 1 + length s) shps + (shps', lens') = unzip $ permuteInv perm $ zip shps lens + scn = drop 1 $ scanl (+) 0 lens' + perm' = concatMap (\(p,l) -> map (\i-> (scn!!p)-l+i) [0..l-1]) + $ zip perm lens + tmp = length perm' + perm'' = perm' ++ [tmp..tmp-1+length shp] + + srnps' = concatMap (\(shp_k, srnp)-> + map fakeDim shp_k ++ [srnp] + ) $ zip shps' srnps + lmad' = setPermutation perm'' $ Lmad tau (srnps' ++ map fakeDim shp) + in IxFun (lmad' : lmads) oshp cg + where fakeDim x = (0,0,x,0,Unknown) + +-- | Rotating an index function: +rotate :: (Eq num, IntegralExp num) => + IxFun num -> Indices num -> IxFun num +rotate (IxFun [] _ _) _ = error "rotate: empty index function" +rotate (IxFun (lmad@(Lmad off srnps) : lmads) oshp cg) offs = + let srnps' = zipWith (\(s,r,n,p,f) o -> + if s == 0 then (0,0,n,p,Unknown) + else (s,r+o,n,p,f) + ) srnps $ permuteInv (getPermutation lmad) offs + in IxFun (Lmad off srnps':lmads) oshp cg + + +-- | Slicing an index function. +slice :: (Eq num, IntegralExp num) => + IxFun num -> Slice num -> IxFun num +slice (IxFun [] _ _) _ = error "slice: empty index function" +slice _ [] = error "slice: empty slice ???" +slice ixfn dim_slices + -- Avoid identity slicing. + | dim_slices == map (unitSlice 0) (shape ixfn) = ixfn +slice (IxFun (lmad@(Lmad _ srnpfs):lmads) oshp cg) is = + let perm= getPermutation lmad + is' = permuteInv perm is + contig = cg && preservesContiguous lmad is' + in if harmlessRotation lmad is' + then let lmad' = foldl sliceOne (Lmad (getOffset lmad) []) + $ zip is' srnpfs + -- need to remove the fixed dims from the permutation + perm' = updatePerm perm $ map fst $ filter isFixedDim $ + zip [0..length is' - 1] is' + in IxFun (setPermutation perm' lmad':lmads) oshp contig + else -- falls outside LMAD formula, hence append a new LMAD + case slice (iota (shape0 lmad)) is of + IxFun [lmad'] _ _ -> IxFun (lmad':lmad:lmads) oshp contig + _ -> error "slice: reached impossible case!" + where isFixedDim (_,DimFix{}) = True + isFixedDim _ = False + + updatePerm ps inds = foldl (\acc p -> acc ++ decrease p) [] ps + where decrease p = + let d = foldl (\n i -> if i == p then (-1) + else if i > p + then n + else if n /= (-1) then n+1 + else n + ) 0 inds + in if d == (-1) then [] else [p-d] + + harmlessRotation0 :: (Eq num, IntegralExp num) => + (num,num,num,Int,DimInfo) -> DimIndex num -> Bool + harmlessRotation0 _ (DimFix _) = True + harmlessRotation0 (0,_,_,_,_) _ = True + harmlessRotation0 (_,0,_,_,_) _ = True + harmlessRotation0 (_,_,n,_,_) dslc + | dslc == DimSlice (n-1) n (-1) || + dslc == unitSlice 0 n = True + harmlessRotation0 _ _ = False + + harmlessRotation :: (Eq num, IntegralExp num) => + Lmad num -> Slice num -> Bool + harmlessRotation (Lmad _ srnps) iss = + and $ zipWith harmlessRotation0 srnps iss + + -- | TODO: what happens to r on a negative-stride slice; is there a such case? + sliceOne :: (Eq num, IntegralExp num) => + Lmad num -> (DimIndex num, (num,num,num,Int,DimInfo)) -> Lmad num + sliceOne (Lmad tau srns) (DimFix i, (s,r,n,_,_)) = + Lmad (tau + flatOneDim (s,r,n) i) srns + sliceOne (Lmad tau srns) (DimSlice _ ne _, (0,_,_,p,_)) = + Lmad tau (srns ++ [(0,0,ne,p,Unknown)]) + sliceOne (Lmad tau srns) (dmind, srn@(_,_,n,_,_)) + | dmind == unitSlice 0 n = Lmad tau (srns ++ [srn]) + sliceOne (Lmad tau srns) (dmind, (s,r,n,p,f)) + | dmind == DimSlice (n-1) n (-1) = + let r' = if r == 0 then 0 else n-r + in Lmad tau' (srns ++ [(s*(-1),r',n,p, invertInfo f)]) + where tau' = tau + flatOneDim (s,0,n) (n-1) + sliceOne (Lmad tau srns) (DimSlice b ne 0, (s,r,n,p,_)) = + Lmad (tau + flatOneDim (s,r,n) b) (srns ++ [(0,0,ne,p,Unknown)]) + sliceOne (Lmad tau srns) (DimSlice bs ns ss, (s,0,_,p,f)) = + let f' = case sgn ss of + Just 1 -> f + Just (-1) -> invertInfo f + _ -> Unknown + in Lmad (tau + s*bs) (srns ++ [(ss*s,0,ns,p,f')]) + sliceOne _ _ = error "slice: reached impossible case!" + + normIndex :: (Eq num, IntegralExp num) => + DimIndex num -> DimIndex num + normIndex (DimSlice b 1 _) = DimFix b + normIndex (DimSlice b _ 0) = DimFix b + normIndex d = d + + preservesContiguous :: (Eq num, IntegralExp num) => + Lmad num -> Slice num -> Bool + preservesContiguous (Lmad _ srnps) slc = + -- remove from the slice the Lmad dimensions who have stride 0. + -- If the Lmad was contiguous in mem, then these dims will not + -- influence the contiguousness of the result. + -- Also normalize the input slice, i.e., 0-stride and size-1 + -- slices are rewritten as DimFixed. + let (srnps', slc') = unzip $ + filter (\((s,_,_,_,_),_) -> s /= 0) $ + zip srnps $ map normIndex slc + -- Check that: + -- 1. a clean split point exists between Fixed and Sliced dims + -- 2. the outermost sliced dim has +/- 1 stride AND is unrottated or full. + -- 3. the rest of inner sliced dims are full. + (_, success) = + foldl (\(found,res) (slcdim, (_,r,n,_,_)) -> + case (slcdim, found) of + (DimFix{}, True ) -> (found, False) + (DimFix{}, False) -> (found, res) + (DimSlice _ ne ds, False) -> -- outermost sliced dim: +/-1 stride + let res' = (r == 0 || n == ne) && (ds == 1 || ds == (-1)) + in (True, res && res') + (DimSlice _ ne ds, True) -> -- inner sliced dim: needs to be full + let res' = (n == ne) && (ds == 1 || ds == (-1)) + in (found, res && res') + ) (False,True) $ zip slc' srnps' + in success + +-- | Reshaping an index function. +-- There are four conditions that all must hold for the result +-- of a reshape operation to remain into the one-Lmad domain: +-- (1) the permutation of the underlying Lmad must leave unchanged +-- the Lmad dimensions that were *not* reshape coercions. +-- (2) the repetition of dimensions of the underlying Lmad must +-- refer only to the coerced-dimensions of the reshape operation. +-- (3) similarly, the rotated dimensions must refer only to +-- dimensions that are coerced by the reshape operation. +-- (4) finally, the underlying memory is contiguous (and monotonous) +-- +-- If any of this conditions does not hold then the reshape operation +-- will conservatively add a new Lmad to the list, leading to a +-- representation that provides less opportunities for further analysis. +-- +-- Actually there are some special cases that need to be treated, +-- for example if everything is a coercion, then it should succeed +-- no matter what. +reshape :: (Eq num, IntegralExp num) => + IxFun num -> ShapeChange num -> IxFun num +reshape (IxFun [] _ _) _ = + error "reshape: empty index function" + +reshape ixfn@(IxFun (lmad@(Lmad tau srnps):lmads) oshp cg) newshape + | -- first take care of the case when this is all a coercion! + perm <- getPermutation lmad, + Just (head_coercions, reshapes, tail_coercions) <- + splitCoercions newshape, + hd_len <- length head_coercions, + num_coercions <- hd_len + length tail_coercions, + srnps' <- permuteFwd perm srnps, + mid_srnps <- take (length srnps - num_coercions) $ + drop hd_len srnps', + num_rshps <- length reshapes, + num_rshps == 0 || (num_rshps == 1 && length mid_srnps == 1), + srnps'' <- map snd $ L.sortBy sortGT $ + zipWith (\(s,r,_,p,f) n -> (p,(s,r,n,p,f))) + srnps' $ newDims newshape + = IxFun (Lmad tau srnps'':lmads) oshp cg + + | perm <- getPermutation lmad, + Just (head_coercions, reshapes, tail_coercions) <- + splitCoercions newshape, + hd_len <- length head_coercions, + num_coercions <- hd_len + length tail_coercions, + srnps_perm <- permuteFwd perm srnps, + mid_srnps <- take (length srnps - num_coercions) $ + drop hd_len srnps_perm, + -- checking conditions (2) and (3) + all (\ (s,r,_,_,_) -> s /= 0 && r == 0) mid_srnps, + -- checking condition (1) + consecutive hd_len $ map (\(_,_,_,p,_)->p) mid_srnps, + -- checking condition (4) + info <- getMonotonicityRots True ixfn, + cg && (info == Inc || info == Dec), + -- make new permutation + rsh_len <- length reshapes, + diff <- length newshape - length srnps, + iota_shape <- [0..length newshape-1], + perm' <- map (\i -> let ind = if i < hd_len + then i else i - diff + in if (i>=hd_len) && (i < hd_len+rsh_len) + then i -- already checked mid_srnps not affected + else let (_,_,_,p,_) = srnps !! ind + in if p < hd_len + then p else p + diff + ) iota_shape, + -- split the dimensions + (suport_inds, repeat_inds) <- + foldl (\(sup,rpt) (i,shpdim,ip) -> + case (i < hd_len, i >= hd_len+rsh_len, shpdim) of + (True, _, DimCoercion n) -> + case srnps_perm !! i of + (0,_,_,_,_) -> ( sup, (ip,n) : rpt ) + (_,r,_,_,_) -> ( (ip,(r,n)) : sup, rpt ) + (_, True, DimCoercion n) -> + case srnps_perm !! (i-diff) of + (0,_,_,_,_) -> ( sup, (ip,n) : rpt ) + (_,r,_,_,_) -> ( (ip,(r,n)) : sup, rpt ) + (False, False, _) -> + ( (ip, (0, newDim shpdim)) : sup, rpt ) + -- already checked that the reshaped + -- dims cannot be repeats or rotates + _ -> error "reshape: reached impossible case!" + ) ([],[]) $ reverse $ zip3 iota_shape newshape perm', + + (sup_inds, support) <- unzip $ L.sortBy sortGT suport_inds, + (rpt_inds, repeats) <- unzip repeat_inds, + Lmad tau' srnps_sup <- makeRotIota info tau support, + repeats' <- map (\n -> (0,0,n,0,Unknown)) repeats, + srnps' <- map snd $ L.sortBy sortGT $ + zip sup_inds srnps_sup ++ zip rpt_inds repeats' + = IxFun (setPermutation perm' (Lmad tau' srnps') : lmads) oshp cg + where splitCoercions newshape' = do + let (head_coercions, newshape'') = span isCoercion newshape' + let (reshapes, tail_coercions) = break isCoercion newshape'' + guard (all isCoercion tail_coercions) + return (head_coercions, reshapes, tail_coercions) + + isCoercion DimCoercion{} = True + isCoercion _ = False + + consecutive _ [] = True + consecutive i [p]= i == p + consecutive i ps = and $ zipWith (==) ps [i, i+1..] + +reshape (IxFun lmads oshp cg) newshape = + let new_dims = newDims newshape + in case iota new_dims of + IxFun [lmad] _ _ -> IxFun (lmad : lmads) oshp cg + _ -> error "reshape: impossible case reached" + + +rank :: IntegralExp num => + IxFun num -> Int +rank (IxFun [] _ _) = error "rank: empty index function" +rank (IxFun (Lmad _ sss : _) _ _) = length sss + +base :: IxFun num -> Shape num +base (IxFun [] _ _) = error "base: empty index function" +base (IxFun _ osh _) = osh + +-- | Correctness assumption: the shape of the new base is +-- equal to the base of the index function (to be rebased). +rebase :: (Eq num, IntegralExp num) => + IxFun num + -> IxFun num + -> IxFun num +rebase (IxFun [] _ _) _ = error "base: empty index function 1" +rebase _ (IxFun [] _ _) = error "base: empty index function 2" + +-- | Special Case: `x[i, (k1,m,s1), (k2,n,s2)] = orig` +-- The new base would be the slice of x. +-- If orig is full (contiguous) and monotonicity is known +-- for all orig's dimensions (i.e., either Inc or Dec) +-- Then we can compose the two into one lmad, the result +-- mainly adapts the index function of the new base. +-- How to handle repeated dimensions in the original? +-- (a) Shave them off of the last lmad of original +-- (b) Compose the result from (a) with the first +-- lmad of the new base +-- (c) apply a repeat operation on the result of (b). +-- However, I strongly suspect that for in-place update +-- what we need is actually the INVERSE of the rebase function, +-- i.e., given an index function new-base and another one orig, +-- compute the index function ixfn0 such that: +-- new-base == rebase ixfn0 ixfn, or equivalently: +-- new-base == ixfn o ixfn0 +-- because then I can go bottom up and compose with ixfn0 +-- all the index functions corresponding to the memory +-- block associated with ixfn. +rebase newbase@(IxFun (lmad_base:lmads_base) shp_base cg_base) + ixfn@(IxFun lmads shp cg) + | lmad_full <- last lmads, + (repeats, lmad) <- shaveoffRepeats lmad_full, + perm <- getPermutation lmad, + srnps<- getLmadDims lmad, + -- sanity condition + base ixfn == shape newbase, + -- TODO: handle repetitions in both lmads. + -- 1) orig is full and monotonicity is known for all dims + cg && length shp == length srnps, + and $ zipWith (\n2 (_,_,n1,_,i1) -> n1 == n2 && i1 /= Unknown) + shp srnps, + -- Building the result srnps: compose permutations, + -- reverse strides and adjust offset if necessary. + perm_base <- getPermutation lmad_base, + perm' <- map (\p -> perm !! p) perm_base, + lmad_base' <- setPermutation perm' lmad_base, + (srnps_base, taus_contrib) <- unzip $ + zipWith (\ (s1,r1,n1,p1,_) (_,r2,_,_,i2) -> + -- assumes the monotonicity of all dimensions is known + let (s', tau') = if i2 == Inc then (s1,0) + else (s1*(-1),s1*(n1-1)) + r' | i2 == Inc = if r2 == 0 then r1 else r1+r2 + | r1 == 0 = r2 + | r2 == 0 = n1-r1 + | otherwise = n1-r1+r2 + in ((s',r',n1,p1,Inc),tau') + ) (getLmadDims lmad_base') $ + permuteInv perm_base srnps, + -- Make resulting lmads: + tau_base' <- getOffset lmad_base' + sum taus_contrib, + lmad_base'' <- Lmad tau_base' srnps_base, + -- Put the repeat back on top of the result + newbase' <- IxFun (lmad_base'':lmads_base) shp_base cg_base, + (reps, rep) <- repeats, + IxFun lmads_base'' _ _ <- repeat newbase' reps rep, + lmads' <- take (length lmads - 1) lmads ++ lmads_base'' + = IxFun lmads' shp_base (cg && cg_base) + +-- General case: just concatenate Lmads since this +-- refers to index-function composition -- always safe! + | base ixfn == shape newbase = + IxFun (lmads ++ lmad_base:lmads_base) shp_base (cg && cg_base) + + | otherwise = + let IxFun lmads' shp_base' _ = reshape newbase $ map DimCoercion shp + in IxFun (lmads ++ lmads') shp_base' (cg && cg_base) + +getMonotonicity :: (Eq num, IntegralExp num) => IxFun num -> DimInfo +getMonotonicity = getMonotonicityRots False + +-- | results in the index function corresponding to indexing +-- with `i` on the outermost dimension. +offsetIndex :: (Eq num, IntegralExp num) => + IxFun num -> num -> IxFun num +offsetIndex ixfun i | i == 0 = ixfun +offsetIndex ixfun i = + case shape ixfun of + d:ds -> slice ixfun (DimSlice i (d-i) 1 : map (unitSlice 0) ds) + [] -> error "offsetIndex: underlying index function has rank zero" + +-- | results in the index function corresponding to making +-- the outermost dimension strided by `s`. +strideIndex :: (Eq num, IntegralExp num) => + IxFun num -> num -> IxFun num +strideIndex ixfun s = + case shape ixfun of + d:ds -> slice ixfun (DimSlice 0 d s : map (unitSlice 0) ds) + [] -> error "offsetIndex: underlying index function has rank zero" + + +-- | If the memory support of the index function is contiguous +-- and row-major (i.e., no transpositions, repetitions, +-- rotates, etc.), then this should return the offset from +-- which the memory-support of this index function starts. +linearWithOffset :: (Eq num, IntegralExp num) => + IxFun num -> num -> Maybe num +linearWithOffset (IxFun [] _ _) _ = + error "linearWithOffset: empty index function" +linearWithOffset ixfn@(IxFun [lmad] _ cg) elem_size + | mon <- getMonotonicity ixfn, + perm <- getPermutation lmad, + cg && mon == Inc, + all (\(s,_,_,_,_) -> s /= 0) (getLmadDims lmad), + perm == [0..length perm - 1], + off <- getOffset lmad = return $ off * elem_size + | otherwise = Nothing +linearWithOffset _ _ = Nothing + +-- | Similar restrictions to `linearWithOffset` except +-- for transpositions, which are returned together +-- with the offset. +rearrangeWithOffset :: (Eq num, IntegralExp num) => + IxFun num -> num -> Maybe (num, [(Int,num)]) +rearrangeWithOffset (IxFun [] _ _) _ = + error "rearrangeWithOffset: empty index function" +rearrangeWithOffset ixfn@(IxFun [lmad] _ cg) elem_size + | perm <- getPermutation lmad, + mon <- getMonotonicity ixfn, + cg && mon == Inc, + all (\(s,_,_,_,_) -> s /= 0) (getLmadDims lmad), + perm /= [0..length perm - 1], + offset <- getOffset lmad * elem_size = + return (offset, zip perm $ rearrangeShape perm $ shape ixfn) + | otherwise = Nothing +rearrangeWithOffset _ _ = Nothing + +isLinear :: (Eq num, IntegralExp num) => IxFun num -> Bool +isLinear = + (==Just 0) . flip linearWithOffset 1 + +------------------------ +--- Helper functions --- +------------------------ + +invertInfo :: DimInfo -> DimInfo +invertInfo Inc = Dec +invertInfo Dec = Inc +invertInfo Unknown = Unknown + +getOffset :: Lmad num -> num +getOffset (Lmad tau _) = tau + +getPermutation :: Lmad num -> Permutation +getPermutation (Lmad _ srns) = map (\(_,_,_,p,_) -> p) srns + +getLmadDims :: Lmad num -> [(num,num,num,Int,DimInfo)] +getLmadDims (Lmad _ srnps) = srnps + +setPermutation :: Permutation -> Lmad num -> Lmad num +setPermutation perm (Lmad tau srnps) = + Lmad tau $ zipWith (\(s,r,n,_,i) p -> (s,r,n,p,i)) srnps perm + +--setOffset :: num -> Lmad num -> Lmad num +--setOffset tau (Lmad _ srnps) = Lmad tau srnps + +-- | Given an input lmad, this function computes a repetition `r` +-- and a new lmad `res`, such that `repeat r res` is identical +-- to the input lmad`. +shaveoffRepeats :: (Eq num, IntegralExp num) => Lmad num -> + (([Shape num], Shape num), Lmad num) +shaveoffRepeats lmad = + let perm = getPermutation lmad + srnps = getLmadDims lmad + -- compute the Repeat: + resacc= foldl (\acc (s,_,n,_,_) -> + case acc of + rpt:acc0 -> + if s == 0 then (n:rpt) : acc0 + else [] : (rpt:acc0) + _ -> error "shaveoffRepeats: empty accum!" + ) [[]] $ L.reverse $ permuteFwd perm srnps + last_shape = last resacc + shapes = take (length resacc - 1) resacc + -- update permutation and lmad: + howManyRepLT k = + foldl (\i (s,_,_,p,_) -> + if s == 0 && p < k then i + 1 else i + ) 0 srnps + srnps' = foldl (\acc (s,r,n,p,info) -> + if s == 0 then acc + else let p' = p - howManyRepLT p + in (s,r,n,p',info):acc + ) [] $ L.reverse srnps + lmad' = Lmad (getOffset lmad) srnps' + in ((shapes,last_shape), lmad') + +permuteFwd :: Permutation -> [a] -> [a] +permuteFwd [] _ = [] +permuteFwd (p:ps) ds = (ds !! p) : permuteFwd ps ds + +permuteInv :: Permutation -> [a] -> [a] +permuteInv ps elems = map snd $ L.sortBy sortGT $ zip ps elems + +sortGT :: Ord a => (a, b1) -> (a, b2) -> Ordering +sortGT (a1, _) (a2, _) + | a1 > a2 = GT + | a1 < a2 = LT + | otherwise = GT + +flatOneDim :: (Eq num, IntegralExp num) => + (num, num, num) -> num -> num +flatOneDim (s,r,n) i + | s == 0 = 0 + | r == 0 = i*s + | otherwise = ((i+r) `mod` n) * s + +makeRotIota :: (IntegralExp num) => + DimInfo -> num -> [(num,num)] -> Lmad num +makeRotIota info tau support + | info == Inc || info == Dec = + let rk = length support + (rs,ns) = unzip support + ss0= L.reverse $ take rk $ scanl (*) 1 $ L.reverse ns + ss = if info == Inc then ss0 + else map (*(-1)) ss0 + ps = map fromIntegral [0..rk-1] + fi = replicate rk info + in Lmad tau $ zip5 ss rs ns ps fi + | otherwise = error "makeRotIota requires Inc or Dec!" + +getMonotonicityRots :: (Eq num, IntegralExp num) => Bool -> IxFun num -> DimInfo +getMonotonicityRots _ (IxFun [] _ _) = + error "getMonotonicityRots: empty index function" +getMonotonicityRots ignore_rots (IxFun (lmad:lmads) _ _) = + let mon1 = getLmadMonotonicity ignore_rots lmad + in if all (==mon1) $ map (getLmadMonotonicity ignore_rots) lmads + then mon1 else Unknown + +getLmadMonotonicity :: (Eq num, IntegralExp num) => Bool -> Lmad num -> DimInfo +getLmadMonotonicity ignore_rots (Lmad _ dims) + | all (isMonDim ignore_rots Inc) dims = Inc + | all (isMonDim ignore_rots Dec) dims = Dec + | otherwise = Unknown + +isMonDim :: (Eq num, IntegralExp num) => Bool -> DimInfo -> + (num, num, num, Int, DimInfo) -> Bool +isMonDim ignore_rots mon (s,r,_,_,info) = + s == 0 || ((ignore_rots || r == 0) && mon == info) diff --git a/src/Futhark/Representation/ExplicitMemory/Simplify.hs b/src/Futhark/Representation/ExplicitMemory/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/ExplicitMemory/Simplify.hs @@ -0,0 +1,209 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE LambdaCase #-} +module Futhark.Representation.ExplicitMemory.Simplify + ( simplifyExplicitMemory + , simplifyStms + ) +where + +import Control.Monad +import qualified Data.Set as S +import Data.Semigroup ((<>)) +import Data.Maybe +import Data.List + +import qualified Futhark.Representation.AST.Syntax as AST +import Futhark.Representation.AST.Syntax + hiding (Prog, BasicOp, Exp, Body, Stm, + Pattern, PatElem, Lambda, FunDef, FParam, LParam, RetType) +import Futhark.Representation.ExplicitMemory +import Futhark.Representation.Kernels.Simplify + (simplifyKernelOp, simplifyKernelExp) +import Futhark.Pass.ExplicitAllocations + (simplifiable, arraySizeInBytesExp) +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import qualified Futhark.Optimise.Simplify.Engine as Engine +import qualified Futhark.Optimise.Simplify as Simplify +import Futhark.Construct +import Futhark.Pass +import Futhark.Optimise.Simplify.Rules +import Futhark.Optimise.Simplify.Rule +import Futhark.Optimise.Simplify.Lore +import Futhark.Util + +simpleExplicitMemory :: Simplify.SimpleOps ExplicitMemory +simpleExplicitMemory = simplifiable (simplifyKernelOp simpleInKernel inKernelEnv) + +simpleInKernel :: KernelSpace -> Simplify.SimpleOps InKernel +simpleInKernel = simplifiable . simplifyKernelExp + +simplifyExplicitMemory :: Prog ExplicitMemory -> PassM (Prog ExplicitMemory) +simplifyExplicitMemory = + Simplify.simplifyProg simpleExplicitMemory callKernelRules + blockers { Engine.blockHoistBranch = isAlloc } + +simplifyStms :: (HasScope ExplicitMemory m, MonadFreshNames m) => + Stms ExplicitMemory -> m (Stms ExplicitMemory) +simplifyStms = + Simplify.simplifyStms simpleExplicitMemory callKernelRules blockers + +isAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore +isAlloc _ (Let _ _ (Op Alloc{})) = True +isAlloc _ _ = False + +isResultAlloc :: Op lore ~ MemOp op => Engine.BlockPred lore +isResultAlloc usage (Let (AST.Pattern [] [bindee]) _ (Op Alloc{})) = + UT.isInResult (patElemName bindee) usage +isResultAlloc _ _ = False + +-- | Getting the roots of what to hoist, for now only variable +-- names that represent array and memory-block sizes. +getShapeNames :: ExplicitMemorish lore => + Stm (Wise lore) -> S.Set VName +getShapeNames bnd = + let tps = map patElemType $ patternElements $ stmPattern bnd + ats = map (snd . patElemAttr) $ patternElements $ stmPattern bnd + nms = mapMaybe (\case + MemMem (Var nm) _ -> Just nm + MemArray _ _ _ (ArrayIn nm _) -> Just nm + _ -> Nothing + ) ats + in S.fromList $ nms ++ subExpVars (concatMap arrayDims tps) + +isAlloc0 :: Op lore ~ MemOp op => AST.Stm lore -> Bool +isAlloc0 (Let _ _ (Op Alloc{})) = True +isAlloc0 _ = False + +inKernelEnv :: Engine.Env InKernel +inKernelEnv = Engine.emptyEnv inKernelRules blockers + +blockers :: (ExplicitMemorish lore, Op lore ~ MemOp op) => + Simplify.HoistBlockers lore +blockers = Engine.noExtraHoistBlockers { + Engine.blockHoistPar = isAlloc + , Engine.blockHoistSeq = isResultAlloc + , Engine.getArraySizes = getShapeNames + , Engine.isAllocation = isAlloc0 + } + +callKernelRules :: RuleBook (Wise ExplicitMemory) +callKernelRules = standardRules <> + ruleBook [RuleBasicOp copyCopyToCopy, + RuleBasicOp removeIdentityCopy] [] + +inKernelRules :: RuleBook (Wise InKernel) +inKernelRules = standardRules <> + ruleBook [RuleBasicOp copyCopyToCopy, + RuleBasicOp removeIdentityCopy, + RuleIf unExistentialiseMemory] [] + +-- | If a branch is returning some existential memory, but the size of +-- the array is existential, then we can create a block of the proper +-- size and always return there. +unExistentialiseMemory :: TopDownRuleIf (Wise InKernel) +unExistentialiseMemory _ pat _ (cond, tbranch, fbranch, ifattr) + | fixable <- foldl hasConcretisableMemory mempty $ patternElements pat, + not $ null fixable = do + + -- Create non-existential memory blocks big enough to hold the + -- arrays. + (arr_to_mem, oldmem_to_mem, oldsize_to_size) <- + fmap unzip3 $ forM fixable $ \(arr_pe, oldmem, oldsize, space) -> do + size <- letSubExp "size" =<< + toExp (arraySizeInBytesExp $ patElemType arr_pe) + mem <- letExp "mem" $ Op $ Alloc size space + return ((patElemName arr_pe, mem), (oldmem, mem), (oldsize, size)) + + -- Update the branches to contain Copy expressions putting the + -- arrays where they are expected. + let updateBody body = insertStmsM $ do + res <- bodyBind body + resultBodyM =<< + zipWithM updateResult (patternElements pat) res + updateResult pat_elem (Var v) + | Just mem <- lookup (patElemName pat_elem) arr_to_mem, + (_, MemArray pt shape u (ArrayIn _ ixfun)) <- patElemAttr pat_elem = do + v_copy <- newVName $ baseString v <> "_nonext_copy" + let v_pat = Pattern [] [PatElem v_copy $ + MemArray pt shape u $ ArrayIn mem ixfun] + addStm $ mkWiseLetStm v_pat (defAux ()) $ BasicOp (Copy v) + return $ Var v_copy + | Just mem <- lookup (patElemName pat_elem) oldmem_to_mem = + return $ Var mem + | Just size <- lookup (Var (patElemName pat_elem)) oldsize_to_size = + return size + updateResult _ se = + return se + tbranch' <- updateBody tbranch + fbranch' <- updateBody fbranch + letBind_ pat $ If cond tbranch' fbranch' ifattr + where onlyUsedIn name here = not $ any ((name `S.member`) . freeIn) $ + filter ((/=here) . patElemName) $ + patternValueElements pat + knownSize Constant{} = True + knownSize (Var v) = not $ inContext v + inContext = (`elem` patternContextNames pat) + + hasConcretisableMemory fixable pat_elem + | (_, MemArray _ shape _ (ArrayIn mem _)) <- patElemAttr pat_elem, + Just (j, Mem old_size space) <- + fmap patElemType <$> find ((mem==) . patElemName . snd) + (zip [(0::Int)..] $ patternElements pat), + Just tse <- maybeNth j $ bodyResult tbranch, + Just fse <- maybeNth j $ bodyResult fbranch, + mem `onlyUsedIn` patElemName pat_elem, + all knownSize (shapeDims shape), + fse /= tse = + (pat_elem, mem, old_size, space) : fixable + | otherwise = + fixable +unExistentialiseMemory _ _ _ _ = cannotSimplify + +-- | If we are copying something that is itself a copy, just copy the +-- original one instead. +copyCopyToCopy :: (BinderOps lore, + LetAttr lore ~ (VarWisdom, MemBound u)) => + TopDownRuleBasicOp lore +copyCopyToCopy vtable pat@(Pattern [] [pat_elem]) _ (Copy v1) + | Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable, + + Just (_, MemArray _ _ _ (ArrayIn srcmem src_ixfun)) <- + ST.entryLetBoundAttr =<< ST.lookup v1 vtable, + + Just (Mem _ src_space) <- ST.lookupType srcmem vtable, + + (_, MemArray _ _ _ (ArrayIn destmem dest_ixfun)) <- patElemAttr pat_elem, + + Just (Mem _ dest_space) <- ST.lookupType destmem vtable, + + src_space == dest_space, dest_ixfun == src_ixfun = + + certifying v1_cs $ letBind_ pat $ BasicOp $ Copy v2 + +copyCopyToCopy vtable pat _ (Copy v0) + | Just (BasicOp (Rearrange perm v1), v0_cs) <- ST.lookupExp v0 vtable, + Just (BasicOp (Copy v2), v1_cs) <- ST.lookupExp v1 vtable = do + v0' <- certifying (v0_cs<>v1_cs) $ + letExp "rearrange_v0" $ BasicOp $ Rearrange perm v2 + letBind_ pat $ BasicOp $ Copy v0' + +copyCopyToCopy _ _ _ _ = cannotSimplify + +-- | If the destination of a copy is the same as the source, just +-- remove it. +removeIdentityCopy :: (BinderOps lore, + LetAttr lore ~ (VarWisdom, MemBound u)) => + TopDownRuleBasicOp lore +removeIdentityCopy vtable pat@(Pattern [] [pe]) _ (Copy v) + | (_, MemArray _ _ _ (ArrayIn dest_mem dest_ixfun)) <- patElemAttr pe, + Just (_, MemArray _ _ _ (ArrayIn src_mem src_ixfun)) <- + ST.entryLetBoundAttr =<< ST.lookup v vtable, + dest_mem == src_mem, dest_ixfun == src_ixfun = + letBind_ pat $ BasicOp $ SubExp $ Var v + +removeIdentityCopy _ _ _ _ = cannotSimplify diff --git a/src/Futhark/Representation/Kernels.hs b/src/Futhark/Representation/Kernels.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Kernels.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +-- | A representation with flat parallelism via GPU-oriented kernels. +module Futhark.Representation.Kernels + ( -- * The Lore definition + Kernels + , InKernel + -- * Module re-exports + , module Futhark.Representation.AST.Attributes + , module Futhark.Representation.AST.Traversals + , module Futhark.Representation.AST.Pretty + , module Futhark.Representation.AST.Syntax + , module Futhark.Representation.Kernels.Kernel + , module Futhark.Representation.Kernels.KernelExp + , module Futhark.Representation.Kernels.Sizes + ) +where + +import Control.Monad + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.Kernels.Kernel +import Futhark.Representation.Kernels.KernelExp +import Futhark.Representation.Kernels.Sizes +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Pretty +import Futhark.Binder +import Futhark.Construct +import qualified Futhark.TypeCheck as TypeCheck + +-- This module could be written much nicer if Haskell had functors +-- like Standard ML. Instead, we have to abuse the namespace/module +-- system. + +data Kernels + +instance Annotations Kernels where + type Op Kernels = Kernel InKernel +instance Attributes Kernels where + expTypesFromPattern = return . expExtTypesFromPattern + +data InKernel +instance Annotations InKernel where + type Op InKernel = KernelExp InKernel +instance Attributes InKernel where + expTypesFromPattern = return . expExtTypesFromPattern +instance PrettyLore InKernel where + +instance TypeCheck.Checkable Kernels where + checkExpLore = return + checkBodyLore = return + checkFParamLore _ = TypeCheck.checkType + checkLParamLore _ = TypeCheck.checkType + checkLetBoundLore _ = TypeCheck.checkType + checkRetType = mapM_ TypeCheck.checkExtType . retTypeValues + checkOp = TypeCheck.subCheck . typeCheckKernel + matchPattern pat = TypeCheck.matchExtPattern pat <=< expExtType + primFParam name t = + return $ Param name (Prim t) + primLParam name t = + return $ Param name (Prim t) + matchReturnType = TypeCheck.matchExtReturnType . map fromDecl + matchBranchType = TypeCheck.matchExtBranchType + +instance TypeCheck.Checkable InKernel where + checkExpLore = return + checkBodyLore = return + checkFParamLore _ = TypeCheck.checkType + checkLParamLore _ = TypeCheck.checkType + checkLetBoundLore _ = TypeCheck.checkType + checkRetType = mapM_ TypeCheck.checkExtType . retTypeValues + checkOp = typeCheckKernelExp + matchPattern pat = TypeCheck.matchExtPattern pat <=< expExtType + primFParam name t = + return $ Param name (Prim t) + primLParam name t = + return $ Param name (Prim t) + matchReturnType = TypeCheck.matchExtReturnType . map fromDecl + matchBranchType = TypeCheck.matchExtBranchType + +instance Bindable Kernels where + mkBody = Body () + mkExpPat ctx val _ = basicPattern ctx val + mkExpAttr _ _ = () + mkLetNames = simpleMkLetNames + +instance BinderOps Kernels where + mkExpAttrB = bindableMkExpAttrB + mkBodyB = bindableMkBodyB + mkLetNamesB = bindableMkLetNamesB + +instance Bindable InKernel where + mkBody = Body () + mkExpPat ctx val _ = basicPattern ctx val + mkExpAttr _ _ = () + mkLetNames = simpleMkLetNames + +instance BinderOps InKernel where + mkExpAttrB = bindableMkExpAttrB + mkBodyB = bindableMkBodyB + mkLetNamesB = bindableMkLetNamesB + +instance PrettyLore Kernels where diff --git a/src/Futhark/Representation/Kernels/Kernel.hs b/src/Futhark/Representation/Kernels/Kernel.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Kernels/Kernel.hs @@ -0,0 +1,692 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Futhark.Representation.Kernels.Kernel + ( Kernel(..) + , kernelType + , KernelDebugHints(..) + , KernelBody(..) + , KernelSpace(..) + , spaceDimensions + , SpaceStructure(..) + , scopeOfKernelSpace + , WhichThreads(..) + , KernelResult(..) + , kernelResultSubExp + , KernelPath + + , chunkedKernelNonconcatOutputs + + , typeCheckKernel + + -- * Generic traversal + , KernelMapper(..) + , identityKernelMapper + , mapKernelM + , KernelWalker(..) + , identityKernelWalker + , walkKernelM + ) + where + +import Control.Monad.Writer hiding (mapM_) +import Control.Monad.Identity hiding (mapM_) +import qualified Data.Set as S +import qualified Data.Map.Strict as M +import Data.Foldable +import Data.List + +import Futhark.Representation.AST +import qualified Futhark.Analysis.Alias as Alias +import qualified Futhark.Analysis.UsageTable as UT +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Util.Pretty as PP +import Futhark.Util.Pretty + ((), (<+>), ppr, commasep, Pretty, parens, text) +import Futhark.Transform.Substitute +import Futhark.Transform.Rename +import Futhark.Optimise.Simplify.Lore +import Futhark.Representation.Ranges + (Ranges, removeLambdaRanges, removeBodyRanges, mkBodyRanges) +import Futhark.Representation.AST.Attributes.Ranges +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Representation.Aliases + (Aliases, removeLambdaAliases, removeBodyAliases, removeStmAliases) +import Futhark.Representation.Kernels.KernelExp (SplitOrdering(..)) +import Futhark.Representation.Kernels.Sizes +import Futhark.Analysis.Usage +import qualified Futhark.TypeCheck as TC +import Futhark.Analysis.Metrics +import Futhark.Tools (partitionChunkedKernelLambdaParameters) +import qualified Futhark.Analysis.Range as Range +import Futhark.Util (maybeNth) + +-- | Some information about what goes into a kernel, and where it came +-- from. Has no semantic meaning; only used for debugging generated +-- code. +data KernelDebugHints = + KernelDebugHints { kernelName :: String + , kernelHints :: [(String, SubExp)] + -- ^ A mapping from a description to some + -- PrimType value. + } + deriving (Eq, Show, Ord) + +data Kernel lore = + GetSize VName SizeClass -- ^ Produce some runtime-configurable size. + | GetSizeMax SizeClass -- ^ The maximum size of some class. + | CmpSizeLe VName SizeClass SubExp + -- ^ Compare size (likely a threshold) with some Int32 value. + | Kernel KernelDebugHints KernelSpace [Type] (KernelBody lore) + deriving (Eq, Show, Ord) + +data KernelSpace = KernelSpace { spaceGlobalId :: VName + , spaceLocalId :: VName + , spaceGroupId :: VName + , spaceNumThreads :: SubExp + , spaceNumGroups :: SubExp + , spaceGroupSize :: SubExp -- flat group size + , spaceStructure :: SpaceStructure + -- TODO: document what this spaceStructure is + -- used for + } + deriving (Eq, Show, Ord) +-- ^ first three bound in the kernel, the rest are params to kernel + +-- | Indices computed for each thread (or group) inside the kernel. +-- This is an arbitrary-dimensional space that is generated from the +-- flat GPU thread space. +data SpaceStructure = FlatThreadSpace + [(VName, SubExp)] -- gtids and dim sizes + | NestedThreadSpace + [(VName, -- gtid + SubExp, -- global dim size + VName, -- ltid + SubExp -- local dim sizes + )] + deriving (Eq, Show, Ord) + +-- | Global thread IDs and their upper bound. +spaceDimensions :: KernelSpace -> [(VName, SubExp)] +spaceDimensions = structureDimensions . spaceStructure + where structureDimensions (FlatThreadSpace dims) = dims + structureDimensions (NestedThreadSpace dims) = + let (gtids, gdim_sizes, _, _) = unzip4 dims + in zip gtids gdim_sizes + +-- | The body of a 'Kernel'. +data KernelBody lore = KernelBody { kernelBodyLore :: BodyAttr lore + , kernelBodyStms :: Stms lore + , kernelBodyResult :: [KernelResult] + } + +deriving instance Annotations lore => Ord (KernelBody lore) +deriving instance Annotations lore => Show (KernelBody lore) +deriving instance Annotations lore => Eq (KernelBody lore) + +data KernelResult = ThreadsReturn WhichThreads SubExp + | WriteReturn + [SubExp] -- Size of array. Must match number of dims. + VName -- Which array + [([SubExp], SubExp)] + -- Arbitrary number of index/value pairs. + | ConcatReturns + SplitOrdering -- Permuted? + SubExp -- The final size. + SubExp -- Per-thread (max) chunk size. + (Maybe SubExp) -- Optional precalculated offset. + VName -- Chunk by this thread. + | KernelInPlaceReturn VName -- HACK! + deriving (Eq, Show, Ord) + +kernelResultSubExp :: KernelResult -> SubExp +kernelResultSubExp (ThreadsReturn _ se) = se +kernelResultSubExp (WriteReturn _ arr _) = Var arr +kernelResultSubExp (ConcatReturns _ _ _ _ v) = Var v +kernelResultSubExp (KernelInPlaceReturn v) = Var v + +data WhichThreads = AllThreads + | OneResultPerGroup + | ThreadsPerGroup [(VName,SubExp)] -- All threads before this one. + | ThreadsInSpace + deriving (Eq, Show, Ord) + +-- | Like 'Mapper', but just for 'Kernel's. +data KernelMapper flore tlore m = KernelMapper { + mapOnKernelSubExp :: SubExp -> m SubExp + , mapOnKernelLambda :: Lambda flore -> m (Lambda tlore) + , mapOnKernelBody :: Body flore -> m (Body tlore) + , mapOnKernelVName :: VName -> m VName + , mapOnKernelLParam :: LParam flore -> m (LParam tlore) + , mapOnKernelKernelBody :: KernelBody flore -> m (KernelBody tlore) + } + +-- | A mapper that simply returns the 'Kernel' verbatim. +identityKernelMapper :: Monad m => KernelMapper lore lore m +identityKernelMapper = KernelMapper { mapOnKernelSubExp = return + , mapOnKernelLambda = return + , mapOnKernelBody = return + , mapOnKernelVName = return + , mapOnKernelLParam = return + , mapOnKernelKernelBody = return + } + +-- | Map a monadic action across the immediate children of a +-- Kernel. The mapping does not descend recursively into subexpressions +-- and is done left-to-right. +mapKernelM :: (Applicative m, Monad m) => + KernelMapper flore tlore m -> Kernel flore -> m (Kernel tlore) +mapKernelM _ (GetSize name size_class) = + pure $ GetSize name size_class +mapKernelM _ (GetSizeMax size_class) = + pure $ GetSizeMax size_class +mapKernelM tv (CmpSizeLe name size_class x) = + CmpSizeLe name size_class <$> mapOnKernelSubExp tv x +mapKernelM tv (Kernel desc space ts kernel_body) = + Kernel <$> mapOnKernelDebugHints desc <*> + mapOnKernelSpace space <*> + mapM (mapOnKernelType tv) ts <*> + mapOnKernelKernelBody tv kernel_body + where mapOnKernelDebugHints (KernelDebugHints name kvs) = + KernelDebugHints name <$> + (zip (map fst kvs) <$> mapM (mapOnKernelSubExp tv . snd) kvs) + mapOnKernelSpace (KernelSpace gtid ltid gid num_threads num_groups group_size structure) = + KernelSpace gtid ltid gid -- all in binding position + <$> mapOnKernelSubExp tv num_threads + <*> mapOnKernelSubExp tv num_groups + <*> mapOnKernelSubExp tv group_size + <*> mapOnKernelStructure structure + mapOnKernelStructure (FlatThreadSpace dims) = + FlatThreadSpace <$> (zip gtids <$> mapM (mapOnKernelSubExp tv) gdim_sizes) + where (gtids, gdim_sizes) = unzip dims + mapOnKernelStructure (NestedThreadSpace dims) = + NestedThreadSpace <$> (zip4 gtids + <$> mapM (mapOnKernelSubExp tv) gdim_sizes + <*> pure ltids + <*> mapM (mapOnKernelSubExp tv) ldim_sizes) + where (gtids, gdim_sizes, ltids, ldim_sizes) = unzip4 dims + +mapOnKernelType :: Monad m => + KernelMapper flore tlore m -> Type -> m Type +mapOnKernelType _tv (Prim pt) = pure $ Prim pt +mapOnKernelType tv (Array pt shape u) = Array pt <$> f shape <*> pure u + where f (Shape dims) = Shape <$> mapM (mapOnKernelSubExp tv) dims +mapOnKernelType _tv (Mem se s) = pure $ Mem se s + +instance (Attributes lore, FreeIn (LParamAttr lore)) => + FreeIn (Kernel lore) where + freeIn e = execWriter $ mapKernelM free e + where walk f x = tell (f x) >> return x + free = KernelMapper { mapOnKernelSubExp = walk freeIn + , mapOnKernelLambda = walk freeInLambda + , mapOnKernelBody = walk freeInBody + , mapOnKernelVName = walk freeIn + , mapOnKernelLParam = walk freeIn + , mapOnKernelKernelBody = walk freeIn + } + +-- | Like 'Walker', but just for 'Kernel's. +data KernelWalker lore m = KernelWalker { + walkOnKernelSubExp :: SubExp -> m () + , walkOnKernelLambda :: Lambda lore -> m () + , walkOnKernelBody :: Body lore -> m () + , walkOnKernelVName :: VName -> m () + , walkOnKernelLParam :: LParam lore -> m () + , walkOnKernelKernelBody :: KernelBody lore -> m () + } + +-- | A no-op traversal. +identityKernelWalker :: Monad m => KernelWalker lore m +identityKernelWalker = KernelWalker { + walkOnKernelSubExp = const $ return () + , walkOnKernelLambda = const $ return () + , walkOnKernelBody = const $ return () + , walkOnKernelVName = const $ return () + , walkOnKernelLParam = const $ return () + , walkOnKernelKernelBody = const $ return () + } + +walkKernelMapper :: forall lore m. Monad m => + KernelWalker lore m -> KernelMapper lore lore m +walkKernelMapper f = KernelMapper { + mapOnKernelSubExp = wrap walkOnKernelSubExp + , mapOnKernelLambda = wrap walkOnKernelLambda + , mapOnKernelBody = wrap walkOnKernelBody + , mapOnKernelVName = wrap walkOnKernelVName + , mapOnKernelLParam = wrap walkOnKernelLParam + , mapOnKernelKernelBody = wrap walkOnKernelKernelBody + } + where wrap :: (KernelWalker lore m -> a -> m ()) -> a -> m a + wrap op k = op f k >> return k + +-- | As 'mapKernelM', but ignoring the results. +walkKernelM :: Monad m => KernelWalker lore m -> Kernel lore -> m () +walkKernelM f = void . mapKernelM m + where m = walkKernelMapper f + +instance FreeIn KernelResult where + freeIn (ThreadsReturn which what) = freeIn which <> freeIn what + freeIn (WriteReturn rws arr res) = freeIn rws <> freeIn arr <> freeIn res + freeIn (ConcatReturns o w per_thread_elems moffset v) = + freeIn o <> freeIn w <> freeIn per_thread_elems <> freeIn moffset <> freeIn v + freeIn (KernelInPlaceReturn what) = freeIn what + +instance FreeIn WhichThreads where + freeIn AllThreads = mempty + freeIn OneResultPerGroup = mempty + freeIn (ThreadsPerGroup limit) = freeIn limit + freeIn ThreadsInSpace = mempty + +instance Attributes lore => FreeIn (KernelBody lore) where + freeIn (KernelBody attr stms res) = + (freeIn attr <> free_in_stms <> free_in_res) `S.difference` bound_in_stms + where free_in_stms = fold $ fmap freeInStm stms + free_in_res = freeIn res + bound_in_stms = fold $ fmap boundByStm stms + +instance Attributes lore => Substitute (KernelBody lore) where + substituteNames subst (KernelBody attr stms res) = + KernelBody + (substituteNames subst attr) + (substituteNames subst stms) + (substituteNames subst res) + +instance Substitute KernelResult where + substituteNames subst (ThreadsReturn who se) = + ThreadsReturn (substituteNames subst who) (substituteNames subst se) + substituteNames subst (WriteReturn rws arr res) = + WriteReturn + (substituteNames subst rws) (substituteNames subst arr) + (substituteNames subst res) + substituteNames subst (ConcatReturns o w per_thread_elems moffset v) = + ConcatReturns + (substituteNames subst o) + (substituteNames subst w) + (substituteNames subst per_thread_elems) + (substituteNames subst moffset) + (substituteNames subst v) + substituteNames subst (KernelInPlaceReturn what) = + KernelInPlaceReturn (substituteNames subst what) + +instance Substitute WhichThreads where + substituteNames _ AllThreads = AllThreads + substituteNames _ OneResultPerGroup = OneResultPerGroup + substituteNames _ ThreadsInSpace = ThreadsInSpace + substituteNames subst (ThreadsPerGroup limit) = + ThreadsPerGroup $ substituteNames subst limit + +instance Substitute KernelSpace where + substituteNames subst (KernelSpace gtid ltid gid num_threads num_groups group_size structure) = + KernelSpace (substituteNames subst gtid) + (substituteNames subst ltid) + (substituteNames subst gid) + (substituteNames subst num_threads) + (substituteNames subst num_groups) + (substituteNames subst group_size) + (substituteNames subst structure) + +instance Substitute SpaceStructure where + substituteNames subst (FlatThreadSpace dims) = + FlatThreadSpace (map (substituteNames subst) dims) + substituteNames subst (NestedThreadSpace dims) = + NestedThreadSpace (map (substituteNames subst) dims) + +instance Attributes lore => Substitute (Kernel lore) where + substituteNames subst (Kernel desc space ts kbody) = + Kernel desc + (substituteNames subst space) + (substituteNames subst ts) + (substituteNames subst kbody) + substituteNames subst k = runIdentity $ mapKernelM substitute k + where substitute = + KernelMapper { mapOnKernelSubExp = return . substituteNames subst + , mapOnKernelLambda = return . substituteNames subst + , mapOnKernelBody = return . substituteNames subst + , mapOnKernelVName = return . substituteNames subst + , mapOnKernelLParam = return . substituteNames subst + , mapOnKernelKernelBody = return . substituteNames subst + } + +instance Attributes lore => Rename (KernelBody lore) where + rename (KernelBody attr stms res) = do + attr' <- rename attr + renamingStms stms $ \stms' -> + KernelBody attr' stms' <$> rename res + +instance Rename KernelResult where + rename = substituteRename + +instance Rename WhichThreads where + rename = substituteRename + +scopeOfKernelSpace :: KernelSpace -> Scope lore +scopeOfKernelSpace (KernelSpace gtid ltid gid _ _ _ structure) = + M.fromList $ zip ([gtid, ltid, gid] ++ structure') $ repeat $ IndexInfo Int32 + where structure' = case structure of + FlatThreadSpace dims -> map fst dims + NestedThreadSpace dims -> + let (gtids, _, ltids, _) = unzip4 dims + in gtids ++ ltids + +instance Attributes lore => Rename (Kernel lore) where + rename = mapKernelM renamer + where renamer = KernelMapper rename rename rename rename rename rename + +kernelType :: Kernel lore -> [Type] +kernelType (Kernel _ space ts body) = + zipWith resultShape ts $ kernelBodyResult body + where dims = map snd $ spaceDimensions space + num_groups = spaceNumGroups space + num_threads = spaceNumThreads space + resultShape t (WriteReturn rws _ _) = + t `arrayOfShape` Shape rws + resultShape t (ThreadsReturn AllThreads _) = + t `arrayOfRow` num_threads + resultShape t (ThreadsReturn OneResultPerGroup _) = + t `arrayOfRow` num_groups + resultShape t (ThreadsReturn (ThreadsPerGroup limit) _) = + t `arrayOfShape` Shape (map snd limit) `arrayOfRow` num_groups + resultShape t (ThreadsReturn ThreadsInSpace _) = + foldr (flip arrayOfRow) t dims + resultShape t (ConcatReturns _ w _ _ _) = + t `arrayOfRow` w + resultShape t KernelInPlaceReturn{} = + t + +kernelType GetSize{} = [Prim int32] +kernelType GetSizeMax{} = [Prim int32] +kernelType CmpSizeLe{} = [Prim Bool] + +chunkedKernelNonconcatOutputs :: Lambda lore -> Int +chunkedKernelNonconcatOutputs fun = + length $ takeWhile (not . outerSizeIsChunk) $ lambdaReturnType fun + where outerSizeIsChunk = (==Var (paramName chunk)) . arraySize 0 + (_, chunk, _) = partitionChunkedKernelLambdaParameters $ lambdaParams fun + +instance TypedOp (Kernel lore) where + opType = pure . staticShapes . kernelType + +instance (Attributes lore, Aliased lore) => AliasedOp (Kernel lore) where + opAliases = map (const mempty) . kernelType + + consumedInOp (Kernel _ _ _ kbody) = + consumedInKernelBody kbody <> + mconcat (map consumedByReturn (kernelBodyResult kbody)) + where consumedByReturn (WriteReturn _ a _) = S.singleton a + consumedByReturn _ = mempty + consumedInOp _ = mempty + +aliasAnalyseKernelBody :: (Attributes lore, + CanBeAliased (Op lore)) => + KernelBody lore + -> KernelBody (Aliases lore) +aliasAnalyseKernelBody (KernelBody attr stms res) = + let Body attr' stms' _ = Alias.analyseBody $ Body attr stms [] + in KernelBody attr' stms' $ map aliasAnalyseKernelResult res + where aliasAnalyseKernelResult (ThreadsReturn which what) = + ThreadsReturn which what + aliasAnalyseKernelResult (WriteReturn rws arr res') = + WriteReturn rws arr res' + aliasAnalyseKernelResult (ConcatReturns o w per_thread_elems moffset v) = + ConcatReturns o w per_thread_elems moffset v + aliasAnalyseKernelResult (KernelInPlaceReturn what) = + KernelInPlaceReturn what + +instance (Attributes lore, + Attributes (Aliases lore), + CanBeAliased (Op lore)) => CanBeAliased (Kernel lore) where + type OpWithAliases (Kernel lore) = Kernel (Aliases lore) + + addOpAliases = runIdentity . mapKernelM alias + where alias = KernelMapper return (return . Alias.analyseLambda) + (return . Alias.analyseBody) return return + (return . aliasAnalyseKernelBody) + + removeOpAliases = runIdentity . mapKernelM remove + where remove = KernelMapper return (return . removeLambdaAliases) + (return . removeBodyAliases) return return + (return . removeKernelBodyAliases) + removeKernelBodyAliases :: KernelBody (Aliases lore) + -> KernelBody lore + removeKernelBodyAliases (KernelBody (_, attr) stms res) = + KernelBody attr (fmap removeStmAliases stms) res + +instance Attributes lore => IsOp (Kernel lore) where + safeOp _ = True + cheapOp Kernel{} = False + cheapOp _ = True + +instance Ranged inner => RangedOp (Kernel inner) where + opRanges op = replicate (length $ kernelType op) unknownRange + +instance (Attributes lore, CanBeRanged (Op lore)) => CanBeRanged (Kernel lore) where + type OpWithRanges (Kernel lore) = Kernel (Ranges lore) + + removeOpRanges = runIdentity . mapKernelM remove + where remove = KernelMapper return (return . removeLambdaRanges) + (return . removeBodyRanges) return return + (return . removeKernelBodyRanges) + removeKernelBodyRanges = error "removeKernelBodyRanges" + addOpRanges = Range.runRangeM . mapKernelM add + where add = KernelMapper return Range.analyseLambda + Range.analyseBody return return addKernelBodyRanges + addKernelBodyRanges (KernelBody attr stms res) = + Range.analyseStms stms $ \stms' -> do + let attr' = (mkBodyRanges stms $ map kernelResultSubExp res, attr) + res' <- mapM addKernelResultRanges res + return $ KernelBody attr' stms' res' + + addKernelResultRanges (ThreadsReturn which what) = + return $ ThreadsReturn which what + addKernelResultRanges (WriteReturn rws arr res) = + return $ WriteReturn rws arr res + addKernelResultRanges (ConcatReturns o w per_thread_elems moffset v) = + return $ ConcatReturns o w per_thread_elems moffset v + addKernelResultRanges (KernelInPlaceReturn what) = + return $ KernelInPlaceReturn what + +instance (Attributes lore, CanBeWise (Op lore)) => CanBeWise (Kernel lore) where + type OpWithWisdom (Kernel lore) = Kernel (Wise lore) + + removeOpWisdom = runIdentity . mapKernelM remove + where remove = KernelMapper return + (return . removeLambdaWisdom) + (return . removeBodyWisdom) + return return + (return . removeKernelBodyWisdom) + removeKernelBodyWisdom :: KernelBody (Wise lore) + -> KernelBody lore + removeKernelBodyWisdom (KernelBody attr stms res) = + let Body attr' stms' _ = removeBodyWisdom $ Body attr stms [] + in KernelBody attr' stms' res + +instance Attributes lore => ST.IndexOp (Kernel lore) where + indexOp vtable k (Kernel _ space _ kbody) is = do + ThreadsReturn which se <- maybeNth k $ kernelBodyResult kbody + + prim_table <- case (which, is) of + (AllThreads, [i]) -> + Just $ M.singleton (spaceGlobalId space) (i,mempty) + (ThreadsInSpace, _) + | (gtids, _) <- unzip $ spaceDimensions space, + length gtids == length is -> + Just $ M.fromList $ zip gtids $ zip is $ repeat mempty + _ -> + Nothing + + let prim_table' = foldl expandPrimExpTable prim_table $ kernelBodyStms kbody + case se of + Var v -> M.lookup v prim_table' + _ -> Nothing + where expandPrimExpTable table stm + | [v] <- patternNames $ stmPattern stm, + Just (pe,cs) <- + runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm = + M.insert v (pe, stmCerts stm <> cs) table + | otherwise = + table + + asPrimExp table v + | Just (e,cs) <- M.lookup v table = tell cs >> return e + | Just (Prim pt) <- ST.lookupType v vtable = + return $ LeafExp v pt + | otherwise = lift Nothing + + indexOp _ _ _ _ = Nothing + +instance Aliased lore => UsageInOp (Kernel lore) where + usageInOp (Kernel _ _ _ kbody) = + mconcat $ map UT.consumedUsage $ S.toList $ consumedInKernelBody kbody + usageInOp GetSize{} = mempty + usageInOp GetSizeMax{} = mempty + usageInOp CmpSizeLe{} = mempty + +consumedInKernelBody :: Aliased lore => + KernelBody lore -> Names +consumedInKernelBody (KernelBody attr stms _) = + consumedInBody $ Body attr stms [] + +typeCheckKernel :: TC.Checkable lore => Kernel (Aliases lore) -> TC.TypeM lore () + +typeCheckKernel GetSize{} = return () +typeCheckKernel GetSizeMax{} = return () +typeCheckKernel (CmpSizeLe _ _ x) = TC.require [Prim int32] x + +typeCheckKernel (Kernel _ space kts kbody) = do + checkSpace space + mapM_ TC.checkType kts + mapM_ (TC.require [Prim int32] . snd) $ spaceDimensions space + + TC.binding (scopeOfKernelSpace space) $ + checkKernelBody kts kbody + where checkSpace (KernelSpace _ _ _ num_threads num_groups group_size structure) = do + mapM_ (TC.require [Prim int32]) [num_threads,num_groups,group_size] + case structure of + FlatThreadSpace dims -> + mapM_ (TC.require [Prim int32] . snd) dims + NestedThreadSpace dims -> + let (_, gdim_sizes, _, ldim_sizes) = unzip4 dims + in mapM_ (TC.require [Prim int32]) $ gdim_sizes ++ ldim_sizes + + checkKernelBody ts (KernelBody (_, attr) stms res) = do + TC.checkBodyLore attr + TC.checkStms stms $ do + unless (length ts == length res) $ + TC.bad $ TC.TypeError $ "Kernel return type is " ++ prettyTuple ts ++ + ", but body returns " ++ show (length res) ++ " values." + zipWithM_ checkKernelResult res ts + + checkKernelResult (ThreadsReturn which what) t = do + checkWhich which + TC.require [t] what + checkKernelResult (WriteReturn rws arr res) t = do + mapM_ (TC.require [Prim int32]) rws + arr_t <- lookupType arr + forM_ res $ \(is, e) -> do + mapM_ (TC.require [Prim int32]) is + TC.require [t] e + unless (arr_t == t `arrayOfShape` Shape rws) $ + TC.bad $ TC.TypeError $ "WriteReturn returning " ++ + pretty e ++ " of type " ++ pretty t ++ ", shape=" ++ pretty rws ++ + ", but destination array has type " ++ pretty arr_t + TC.consume =<< TC.lookupAliases arr + checkKernelResult (ConcatReturns o w per_thread_elems moffset v) t = do + case o of + SplitContiguous -> return () + SplitStrided stride -> TC.require [Prim int32] stride + TC.require [Prim int32] w + TC.require [Prim int32] per_thread_elems + mapM_ (TC.require [Prim int32]) moffset + vt <- lookupType v + unless (vt == t `arrayOfRow` arraySize 0 vt) $ + TC.bad $ TC.TypeError $ "Invalid type for ConcatReturns " ++ pretty v + checkKernelResult (KernelInPlaceReturn what) t = + TC.requireI [t] what + + checkWhich AllThreads = return () + checkWhich OneResultPerGroup = return () + checkWhich ThreadsInSpace = return () + checkWhich (ThreadsPerGroup limit) = do + mapM_ (TC.requireI [Prim int32] . fst) limit + mapM_ (TC.require [Prim int32] . snd) limit + +instance OpMetrics (Op lore) => OpMetrics (Kernel lore) where + opMetrics (Kernel _ _ _ kbody) = + inside "Kernel" $ kernelBodyMetrics kbody + where kernelBodyMetrics :: KernelBody lore -> MetricsM () + kernelBodyMetrics = mapM_ bindingMetrics . kernelBodyStms + opMetrics GetSize{} = seen "GetSize" + opMetrics GetSizeMax{} = seen "GetSizeMax" + opMetrics CmpSizeLe{} = seen "CmpSizeLe" + +instance PrettyLore lore => PP.Pretty (Kernel lore) where + ppr (GetSize name size_class) = + text "get_size" <> parens (commasep [ppr name, ppr size_class]) + + ppr (GetSizeMax size_class) = + text "get_size_max" <> parens (ppr size_class) + + ppr (CmpSizeLe name size_class x) = + text "get_size" <> parens (commasep [ppr name, ppr size_class]) <+> + text "<" <+> ppr x + + ppr (Kernel desc space ts body) = + text "kernel" <+> text (kernelName desc) <> + PP.align (ppr space) <+> + PP.colon <+> ppTuple' ts <+> PP.nestedBlock "{" "}" (ppr body) + +instance Pretty KernelSpace where + ppr (KernelSpace f_gtid f_ltid gid num_threads num_groups group_size structure) = + parens (commasep [text "num groups:" <+> ppr num_groups, + text "group size:" <+> ppr group_size, + text "num threads:" <+> ppr num_threads, + text "global TID ->" <+> ppr f_gtid, + text "local TID ->" <+> ppr f_ltid, + text "group ID ->" <+> ppr gid]) structure' + where structure' = + case structure of + FlatThreadSpace dims -> flat dims + NestedThreadSpace space -> + parens (commasep $ do + (gtid,gd,ltid,ld) <- space + return $ ppr (gtid,ltid) <+> "<" <+> ppr (gd,ld)) + flat dims = parens $ commasep $ do + (i,d) <- dims + return $ ppr i <+> "<" <+> ppr d + +instance PrettyLore lore => Pretty (KernelBody lore) where + ppr (KernelBody _ stms res) = + PP.stack (map ppr (stmsToList stms)) + text "return" <+> PP.braces (PP.commasep $ map ppr res) + +instance Pretty KernelResult where + ppr (ThreadsReturn AllThreads what) = + ppr what + ppr (ThreadsReturn OneResultPerGroup what) = + text "group" <+> "returns" <+> ppr what + ppr (ThreadsReturn (ThreadsPerGroup limit) what) = + text "thread <" <+> ppr limit <+> text "returns" <+> ppr what + ppr (ThreadsReturn ThreadsInSpace what) = + text "thread in space returns" <+> ppr what + ppr (WriteReturn rws arr res) = + ppr arr <+> text "with" <+> PP.apply (map ppRes res) + where ppRes (is, e) = + PP.brackets (PP.commasep $ zipWith f is rws) <+> text "<-" <+> ppr e + f i rw = ppr i <+> text "<" <+> ppr rw + ppr (ConcatReturns o w per_thread_elems offset v) = + text "concat" <> suff <> + parens (commasep [ppr w, ppr per_thread_elems] <> offset_text) <+> + ppr v + where suff = case o of SplitContiguous -> mempty + SplitStrided stride -> text "Strided" <> parens (ppr stride) + offset_text = case offset of Nothing -> "" + Just se -> "," <+> "offset=" <> ppr se + ppr (KernelInPlaceReturn what) = + text "kernel returns" <+> ppr what diff --git a/src/Futhark/Representation/Kernels/KernelExp.hs b/src/Futhark/Representation/Kernels/KernelExp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Kernels/KernelExp.hs @@ -0,0 +1,616 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ConstraintKinds #-} +-- | A representation of nested-parallel in-kernel per-workgroup +-- expressions. +module Futhark.Representation.Kernels.KernelExp + ( KernelExp(..) + , GroupStreamLambda(..) + , SplitOrdering(..) + , CombineSpace(..) + , combineSpace + , scopeOfCombineSpace + , typeCheckKernelExp + ) + where + +import Control.Monad +import Data.Monoid ((<>)) +import Data.Maybe +import qualified Data.Set as S +import qualified Data.Map.Strict as M + +import qualified Futhark.Analysis.Alias as Alias +import qualified Futhark.Analysis.Range as Range +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Representation.Aliases +import Futhark.Representation.Ranges +import Futhark.Transform.Substitute +import Futhark.Transform.Rename +import Futhark.Optimise.Simplify.Lore +import Futhark.Analysis.Usage +import Futhark.Analysis.Metrics +import qualified Futhark.Analysis.ScalExp as SE +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Util.Pretty + ((<+>), (), ppr, comma, commasep, Pretty, parens, text, apply, braces, annot, indent) +import qualified Futhark.TypeCheck as TC +import Futhark.Util (chunks) + +-- | How an array is split into chunks. +data SplitOrdering = SplitContiguous + | SplitStrided SubExp + deriving (Eq, Ord, Show) + +-- | A combine can be fully or partially in-place. The initial arrays +-- here work like the ones from the Scatter SOAC. +data CombineSpace = CombineSpace { cspaceScatter :: [(SubExp, Int, VName)] + , cspaceDims :: [(VName,SubExp)] } + deriving (Eq, Ord, Show) + +combineSpace :: [(VName,SubExp)] -> CombineSpace +combineSpace = CombineSpace [] + +scopeOfCombineSpace :: CombineSpace -> Scope lore +scopeOfCombineSpace (CombineSpace _ dims) = + M.fromList $ zip (map fst dims) $ repeat $ IndexInfo Int32 + +data KernelExp lore = SplitSpace SplitOrdering SubExp SubExp SubExp + -- ^ @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@. + | Combine CombineSpace [Type] [(VName,SubExp)] (Body lore) + -- ^ @Combine cspace ts aspace body@ will + -- combine values from threads to a single + -- (multidimensional) array. If we define @(is, + -- ws) = unzip cspace@, then @ws@ is defined the + -- same accross all threads. The @cspace@ + -- defines the shape of the resulting array, and + -- the identifiers used to identify each + -- individual element. Only threads for which + -- @all (\(i,w) -> i < w) aspace@ is true will + -- provide a value (of type @ts@), which is + -- generated by @body@. + -- + -- The result of a combine is always stored in local + -- memory (OpenCL terminology) + -- + -- The same thread may be assigned to multiple + -- elements of 'Combine', if the size of the + -- 'CombineSpace' exceeds the group size. + | GroupReduce SubExp + (Lambda lore) [(SubExp,VName)] + -- ^ @GroupReduce w lam input@ (with @(nes, arrs) = unzip input@), + -- will perform a reduction of the arrays @arrs@ using the + -- associative reduction operator @lam@ and the neutral + -- elements @nes@. + -- + -- The arrays @arrs@ must all have outer + -- dimension @w@, which must not be larger than + -- the group size. + -- + -- Currently a GroupReduce consumes the input arrays, as + -- it uses them for scratch space to store temporary + -- results + -- + -- All threads in a group must participate in a + -- GroupReduce (due to barriers) + -- + -- The length of the arrays @w@ can be smaller than the + -- number of elements in a group (neutral element will be + -- filled in), but @w@ can never be larger than the group + -- size. + | GroupScan SubExp + (Lambda lore) [(SubExp,VName)] + -- ^ Same restrictions as with 'GroupReduce'. + | GroupStream SubExp SubExp + (GroupStreamLambda lore) [SubExp] [VName] + -- Morally a StreamSeq + -- First SubExp is the outersize of the array + -- Second SubExp is the maximal chunk size + -- [SubExp] is the accumulator, [VName] are the input arrays + | GroupGenReduce [SubExp] [VName] (LambdaT lore) [SubExp] [SubExp] VName + -- ^ GroupGenReduce + | Barrier [SubExp] + -- ^ HACK: Semantically identity, but inserts a + -- barrier afterwards. This reflects a weakness + -- in our kernel representation. + deriving (Eq, Ord, Show) + +data GroupStreamLambda lore = GroupStreamLambda + { groupStreamChunkSize :: VName + , groupStreamChunkOffset :: VName + , groupStreamAccParams :: [LParam lore] + , groupStreamArrParams :: [LParam lore] + , groupStreamLambdaBody :: Body lore + } + +deriving instance Annotations lore => Eq (GroupStreamLambda lore) +deriving instance Annotations lore => Show (GroupStreamLambda lore) +deriving instance Annotations lore => Ord (GroupStreamLambda lore) + +instance Attributes lore => IsOp (KernelExp lore) where + safeOp _ = False + cheapOp _ = True + +instance Attributes lore => TypedOp (KernelExp lore) where + opType SplitSpace{} = + pure $ staticShapes [Prim int32] + opType (Combine (CombineSpace scatter cspace) ts _ _) = + pure $ staticShapes $ + zipWith arrayOfRow val_ts ws ++ + map (`arrayOfShape` shape) (drop (sum ns*2) ts) + where shape = Shape $ map snd cspace + val_ts = concatMap (take 1) $ chunks ns $ + take (sum ns) $ drop (sum ns) ts + (ws, ns, _) = unzip3 scatter + opType (GroupReduce _ lam _) = + pure $ staticShapes $ lambdaReturnType lam + opType (GroupScan w lam _) = + pure $ staticShapes $ map (`arrayOfRow` w) (lambdaReturnType lam) + opType (GroupStream _ _ lam _ _) = + pure $ staticShapes $ map paramType $ groupStreamAccParams lam + opType (GroupGenReduce _ dests _ _ _ _) = + staticShapes <$> traverse lookupType dests + opType (Barrier ses) = staticShapes <$> traverse subExpType ses + +instance FreeIn SplitOrdering where + freeIn SplitContiguous = mempty + freeIn (SplitStrided stride) = freeIn stride + +instance Attributes lore => FreeIn (KernelExp lore) where + freeIn (SplitSpace o w i elems_per_thread) = + freeIn o <> freeIn [w, i, elems_per_thread] + freeIn (Combine (CombineSpace scatter cspace) ts active body) = + freeIn scatter <> freeIn (map snd cspace) <> freeIn ts <> freeIn active <> freeInBody body + freeIn (GroupReduce w lam input) = + freeIn w <> freeInLambda lam <> freeIn input + freeIn (GroupScan w lam input) = + freeIn w <> freeInLambda lam <> freeIn input + freeIn (GroupStream w maxchunk lam accs arrs) = + freeIn w <> freeIn maxchunk <> freeIn lam <> freeIn accs <> freeIn arrs + freeIn (GroupGenReduce w dests op bucket values locks) = + freeIn w <> freeIn dests <> freeInLambda op <> freeIn bucket <> freeIn values <> freeIn locks + freeIn (Barrier ses) = freeIn ses + +instance Attributes lore => FreeIn (GroupStreamLambda lore) where + freeIn (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + freeInBody body `S.difference` bound_here + where bound_here = S.fromList $ + chunk_offset : chunk_size : + map paramName (acc_params ++ arr_params) + +instance Ranged inner => RangedOp (KernelExp inner) where + opRanges (SplitSpace _ _ _ elems_per_thread) = + [(Just (ScalarBound 0), + Just (ScalarBound (SE.subExpToScalExp elems_per_thread int32)))] + opRanges _ = repeat unknownRange + +instance (Attributes lore, Aliased lore) => AliasedOp (KernelExp lore) where + opAliases SplitSpace{} = + [mempty] + opAliases Combine{} = + [mempty] + opAliases (GroupReduce _ lam _) = + replicate (length (lambdaReturnType lam)) mempty + opAliases (GroupScan _ lam _) = + replicate (length (lambdaReturnType lam)) mempty + opAliases (GroupStream _ _ lam _ _) = + map (const mempty) $ groupStreamAccParams lam + opAliases (GroupGenReduce _ dests _ _ _ _) = + map S.singleton dests + opAliases (Barrier ses) = map subExpAliases ses + + consumedInOp (GroupReduce _ _ input) = + S.fromList $ map snd input + consumedInOp (GroupScan _ _ input) = + S.fromList $ map snd input + consumedInOp (GroupStream _ _ lam accs arrs) = + -- GroupStream always consumes array-typed accumulators. This + -- guarantees that we can use their storage for the result of the + -- lambda. + S.map consumedArray $ + S.fromList (map paramName acc_params) <> consumedInBody body + where GroupStreamLambda _ _ acc_params arr_params body = lam + consumedArray v = fromMaybe v $ subExpVar =<< lookup v params_to_arrs + params_to_arrs = zip (map paramName $ acc_params ++ arr_params) $ + accs ++ map Var arrs + consumedInOp (GroupGenReduce _ dests _ _ _ _) = + S.fromList dests + + consumedInOp SplitSpace{} = mempty + consumedInOp Barrier{} = mempty + consumedInOp (Combine _ _ _ body) = consumedInBody body + +instance Substitute SplitOrdering where + substituteNames _ SplitContiguous = + SplitContiguous + substituteNames subst (SplitStrided stride) = + SplitStrided $ substituteNames subst stride + +instance Substitute CombineSpace where + substituteNames substs (CombineSpace scatter dims) = + CombineSpace (map sub scatter) (substituteNames substs dims) + where sub (w, n, a) = + (substituteNames substs w, n, substituteNames substs a) + +instance Attributes lore => Substitute (KernelExp lore) 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 subst (Combine cspace ts active v) = + Combine (substituteNames subst cspace) ts + (substituteNames subst active) (substituteNames subst v) + substituteNames subst (GroupReduce w lam input) = + GroupReduce (substituteNames subst w) + (substituteNames subst lam) (substituteNames subst input) + substituteNames subst (GroupScan w lam input) = + GroupScan (substituteNames subst w) + (substituteNames subst lam) (substituteNames subst input) + substituteNames subst (GroupStream w maxchunk lam accs arrs) = + GroupStream + (substituteNames subst w) (substituteNames subst maxchunk) + (substituteNames subst lam) + (substituteNames subst accs) (substituteNames subst arrs) + substituteNames subst (GroupGenReduce w dests op bucket vs locks) = + GroupGenReduce (substituteNames subst w) (substituteNames subst dests) + (substituteNames subst op) (substituteNames subst bucket) (substituteNames subst vs) + (substituteNames subst locks) + substituteNames substs (Barrier ses) = Barrier $ substituteNames substs ses + +instance Attributes lore => Substitute (GroupStreamLambda lore) where + substituteNames + subst (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + GroupStreamLambda + (substituteNames subst chunk_size) + (substituteNames subst chunk_offset) + (substituteNames subst acc_params) + (substituteNames subst arr_params) + (substituteNames subst body) + +instance Rename SplitOrdering where + rename SplitContiguous = + pure SplitContiguous + rename (SplitStrided stride) = + SplitStrided <$> rename stride + +instance Rename CombineSpace where + rename = substituteRename + +instance Renameable lore => Rename (KernelExp lore) where + rename (SplitSpace o w i elems_per_thread) = + SplitSpace + <$> rename o + <*> rename w + <*> rename i + <*> rename elems_per_thread + rename (Combine cspace ts active v) = + Combine <$> rename cspace <*> rename ts <*> rename active <*> rename v + rename (GroupReduce w lam input) = + GroupReduce <$> rename w <*> rename lam <*> rename input + rename (GroupScan w lam input) = + GroupScan <$> rename w <*> rename lam <*> rename input + rename (GroupStream w maxchunk lam accs arrs) = + GroupStream <$> rename w <*> rename maxchunk <*> + rename lam <*> rename accs <*> rename arrs + rename (GroupGenReduce w dests op bucket vs locks) = + GroupGenReduce <$> rename w <*> rename dests <*> rename op <*> + rename bucket <*> rename vs <*> rename locks + rename (Barrier ses) = Barrier <$> mapM rename ses + +instance Renameable lore => Rename (GroupStreamLambda lore) where + rename (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + bindingForRename (chunk_size : chunk_offset : map paramName (acc_params++arr_params)) $ + GroupStreamLambda <$> + rename chunk_size <*> + rename chunk_offset <*> + rename acc_params <*> + rename arr_params <*> + rename body + +instance (Attributes lore, + Attributes (Aliases lore), + CanBeAliased (Op lore)) => CanBeAliased (KernelExp lore) where + type OpWithAliases (KernelExp lore) = KernelExp (Aliases lore) + + addOpAliases (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + addOpAliases (GroupReduce w lam input) = + GroupReduce w (Alias.analyseLambda lam) input + addOpAliases (GroupScan w lam input) = + GroupScan w (Alias.analyseLambda lam) input + addOpAliases (GroupStream w maxchunk lam accs arrs) = + GroupStream w maxchunk lam' accs arrs + where lam' = analyseGroupStreamLambda lam + analyseGroupStreamLambda (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + GroupStreamLambda chunk_size chunk_offset acc_params arr_params $ + Alias.analyseBody body + addOpAliases (GroupGenReduce w dests op bucket vs locks) = + GroupGenReduce w dests (Alias.analyseLambda op) bucket vs locks + addOpAliases (Combine cspace ts active body) = + Combine cspace ts active $ Alias.analyseBody body + addOpAliases (Barrier ses) = Barrier ses + + removeOpAliases (GroupReduce w lam input) = + GroupReduce w (removeLambdaAliases lam) input + removeOpAliases (GroupScan w lam input) = + GroupScan w (removeLambdaAliases lam) input + removeOpAliases (GroupStream w maxchunk lam accs arrs) = + GroupStream w maxchunk (removeGroupStreamLambdaAliases lam) accs arrs + where removeGroupStreamLambdaAliases (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + GroupStreamLambda chunk_size chunk_offset acc_params arr_params $ + removeBodyAliases body + removeOpAliases (GroupGenReduce w dests op bucket vs locks) = + GroupGenReduce w dests (removeLambdaAliases op) bucket vs locks + removeOpAliases (Combine cspace ts active body) = + Combine cspace ts active $ removeBodyAliases body + removeOpAliases (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + removeOpAliases (Barrier ses) = Barrier ses + +instance (Attributes lore, + Attributes (Ranges lore), + CanBeRanged (Op lore)) => CanBeRanged (KernelExp lore) where + type OpWithRanges (KernelExp lore) = KernelExp (Ranges lore) + + addOpRanges (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + addOpRanges (GroupReduce w lam input) = + GroupReduce w (Range.runRangeM $ Range.analyseLambda lam) input + addOpRanges (GroupScan w lam input) = + GroupScan w (Range.runRangeM $ Range.analyseLambda lam) input + addOpRanges (GroupGenReduce w dests op bucket vs locks) = + GroupGenReduce w dests (Range.runRangeM $ Range.analyseLambda op) bucket vs locks + addOpRanges (Combine cspace ts active body) = + Combine cspace ts active $ Range.runRangeM $ Range.analyseBody body + addOpRanges (GroupStream w maxchunk lam accs arrs) = + GroupStream w maxchunk lam' accs arrs + where lam' = analyseGroupStreamLambda lam + analyseGroupStreamLambda (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + GroupStreamLambda chunk_size chunk_offset acc_params arr_params $ + Range.runRangeM $ Range.analyseBody body + addOpRanges (Barrier ses) = Barrier ses + + removeOpRanges (GroupReduce w lam input) = + GroupReduce w (removeLambdaRanges lam) input + removeOpRanges (GroupScan w lam input) = + GroupScan w (removeLambdaRanges lam) input + removeOpRanges (GroupStream w maxchunk lam accs arrs) = + GroupStream w maxchunk (removeGroupStreamLambdaRanges lam) accs arrs + where removeGroupStreamLambdaRanges (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + GroupStreamLambda chunk_size chunk_offset acc_params arr_params $ + removeBodyRanges body + removeOpRanges (GroupGenReduce w dests op bucket vs locks) = + GroupGenReduce w dests (removeLambdaRanges op) bucket vs locks + removeOpRanges (Combine cspace ts active body) = + Combine cspace ts active $ removeBodyRanges body + removeOpRanges (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + removeOpRanges (Barrier ses) = Barrier ses + +instance (Attributes lore, CanBeWise (Op lore)) => CanBeWise (KernelExp lore) where + type OpWithWisdom (KernelExp lore) = KernelExp (Wise lore) + + removeOpWisdom (GroupReduce w lam input) = + GroupReduce w (removeLambdaWisdom lam) input + removeOpWisdom (GroupScan w lam input) = + GroupScan w (removeLambdaWisdom lam) input + removeOpWisdom (GroupStream w maxchunk lam accs arrs) = + GroupStream w maxchunk (removeGroupStreamLambdaWisdom lam) accs arrs + where removeGroupStreamLambdaWisdom + (GroupStreamLambda chunk_size chunk_offset acc_params arr_params body) = + GroupStreamLambda chunk_size chunk_offset acc_params arr_params $ + removeBodyWisdom body + removeOpWisdom (GroupGenReduce w dests op bucket vs locks) = + GroupGenReduce w dests (removeLambdaWisdom op) bucket vs locks + removeOpWisdom (Combine cspace ts active body) = + Combine cspace ts active $ removeBodyWisdom body + removeOpWisdom (SplitSpace o w i elems_per_thread) = + SplitSpace o w i elems_per_thread + removeOpWisdom (Barrier ses) = Barrier ses + +instance ST.IndexOp (KernelExp lore) where + +instance Aliased lore => UsageInOp (KernelExp lore) where + usageInOp (Combine _ _ _ body) = + mconcat $ map UT.consumedUsage $ S.toList $ consumedInBody body + usageInOp _ = mempty + +instance OpMetrics (Op lore) => OpMetrics (KernelExp lore) where + opMetrics SplitSpace{} = seen "SplitSpace" + opMetrics Combine{} = seen "Combine" + opMetrics (GroupReduce _ lam _) = inside "GroupReduce" $ lambdaMetrics lam + opMetrics (GroupScan _ lam _) = inside "GroupScan" $ lambdaMetrics lam + opMetrics (GroupGenReduce _ _ op _ _ _) = inside "GroupGenReduce" $ lambdaMetrics op + opMetrics (GroupStream _ _ lam _ _) = + inside "GroupStream" $ groupStreamLambdaMetrics lam + where groupStreamLambdaMetrics = + bodyMetrics . groupStreamLambdaBody + opMetrics Barrier{} = seen "Barrier" + +typeCheckKernelExp :: TC.Checkable lore => KernelExp (Aliases lore) -> TC.TypeM lore () + +typeCheckKernelExp Barrier{} = return () + +typeCheckKernelExp (SplitSpace o w i elems_per_thread) = do + case o of + SplitContiguous -> return () + SplitStrided stride -> TC.require [Prim int32] stride + mapM_ (TC.require [Prim int32]) [w, i, elems_per_thread] + +typeCheckKernelExp (Combine cspace@(CombineSpace scatter dims) ts aspace body) = do + mapM_ (TC.require [Prim int32]) ws + TC.binding (scopeOfCombineSpace cspace) $ do + let (_as_ws, as_ns, _as_vs) = unzip3 scatter + num_scatters = sum as_ns + ts_is = take num_scatters ts + ts_vs = take num_scatters $ drop num_scatters ts + + unless (length ts_is == num_scatters && length ts_vs == num_scatters) $ + TC.bad $ TC.TypeError "Combine: inconsistent return type annotation." + + forM_ ts_is $ \ts_i -> unless (Prim int32 == ts_i) $ + TC.bad $ TC.TypeError "Combine: index return type must be i32." + + forM_ (zip (chunks as_ns ts_vs) scatter) $ \(ts_vs', (aw, _, a)) -> do + TC.require [Prim int32] aw + forM_ ts_vs' $ \ts_v -> TC.requireI [ts_v `arrayOfRow` aw] a + TC.consume =<< TC.lookupAliases a + + mapM_ TC.checkType ts + mapM_ (TC.requireI [Prim int32]) a_is + mapM_ (TC.require [Prim int32]) a_ws + TC.checkLambdaBody ts body + where ws = map snd dims + (a_is, a_ws) = unzip aspace + +typeCheckKernelExp (GroupReduce w lam input) = + checkScanOrReduce w lam input + +typeCheckKernelExp (GroupScan w lam input) = + checkScanOrReduce w lam input + +typeCheckKernelExp (GroupGenReduce ws dests op bucket vs locks) = do + mapM_ (TC.require [Prim int32]) ws + + mapM_ (TC.require [Prim int32]) bucket + + dest_row_ts <- mapM (fmap (stripArray (length bucket)) . lookupType) dests + + vs_ts <- mapM subExpType vs + unless (vs_ts == dest_row_ts) $ + TC.bad $ TC.TypeError $ "Destination arrays have type " ++ + pretty dest_row_ts ++ ", but values to write have type " ++ pretty vs_ts + + TC.requireI [Prim int32 `arrayOfShape` Shape ws] locks + + let asArg t = (t, mempty) + TC.checkLambda op $ map asArg $ dest_row_ts ++ vs_ts + +typeCheckKernelExp (GroupStream w maxchunk lam accs arrs) = do + TC.require [Prim int32] w + TC.require [Prim int32] maxchunk + + acc_args <- mapM TC.checkArg accs + arr_args <- TC.checkSOACArrayArgs w arrs + + checkGroupStreamLambda acc_args arr_args + where GroupStreamLambda block_size _ acc_params arr_params body = lam + checkGroupStreamLambda acc_args arr_args = do + unless (map TC.argType acc_args == map paramType acc_params) $ + TC.bad $ TC.TypeError + "checkGroupStreamLambda: wrong accumulator arguments." + + let arr_block_ts = + map ((`arrayOfRow` Var block_size) . TC.argType) arr_args + unless (map paramType arr_params == arr_block_ts) $ + TC.bad $ TC.TypeError + "checkGroupStreamLambda: wrong array arguments." + + let acc_consumable = + zip (map paramName acc_params) (map TC.argAliases acc_args) + arr_consumable = + zip (map paramName arr_params) (map TC.argAliases arr_args) + consumable = acc_consumable ++ arr_consumable + TC.binding (scopeOf lam) $ TC.consumeOnlyParams consumable $ do + TC.checkLambdaParams acc_params + TC.checkLambdaParams arr_params + TC.checkLambdaBody (map TC.argType acc_args) body + +checkScanOrReduce :: TC.Checkable lore => + SubExp -> Lambda (Aliases lore) -> [(SubExp, VName)] + -> TC.TypeM lore () +checkScanOrReduce w lam input = do + TC.require [Prim int32] w + let (nes, arrs) = unzip input + asArg t = (t, mempty) + neargs <- mapM TC.checkArg nes + arrargs <- TC.checkSOACArrayArgs w arrs + TC.checkLambda lam $ + map asArg [Prim int32, Prim int32] ++ + map TC.noArgAliases (neargs ++ arrargs) + +instance Scoped lore (GroupStreamLambda lore) where + scopeOf (GroupStreamLambda chunk_size chunk_offset acc_params arr_params _) = + M.insert chunk_size (IndexInfo Int32) $ + M.insert chunk_offset (IndexInfo Int32) $ + scopeOfLParams (acc_params ++ arr_params) + +instance PrettyLore lore => Pretty (KernelExp lore) where + ppr (SplitSpace o w i elems_per_thread) = + text "splitSpace" <> suff <> + parens (commasep [ppr w, ppr i, ppr elems_per_thread]) + where suff = case o of SplitContiguous -> mempty + SplitStrided stride -> text "Strided" <> parens (ppr stride) + ppr (Combine (CombineSpace scatter cspace) ts active body) = + text "combine" <> + apply (map (\(_,n,a) -> text "@" <> ppr (n,a)) scatter ++ + map (\(i,w) -> ppr i <+> text "<" <+> ppr w) cspace ++ + [apply (map ppr ts), ppr active]) <+> text "{" + indent 2 (ppr body) + text "}" + ppr (GroupReduce w lam input) = + text "reduce" <> parens (commasep [ppr w, + ppr lam, + braces (commasep $ map ppr nes), + commasep $ map ppr els]) + where (nes,els) = unzip input + ppr (GroupScan w lam input) = + text "scan" <> parens (commasep [ppr w, + ppr lam, + braces (commasep $ map ppr nes), + commasep $ map ppr els]) + where (nes,els) = unzip input + ppr (GroupStream w maxchunk lam accs arrs) = + text "stream" <> + parens (ppr w <> comma <+> ppr maxchunk <> comma + ppr lam <> comma + braces (commasep $ map ppr accs) <> comma + commasep (map ppr arrs)) + + ppr (GroupGenReduce w dests op bucket vs locks) = + text "gen_reduce" <> + parens (ppr w <> comma + braces (commasep $ map ppr dests) <> comma + ppr op <> comma + braces (commasep $ map ppr bucket) <> comma + braces (commasep $ map ppr vs) <> comma + ppr locks) + + ppr (Barrier ses) = text "barrier" <> parens (commasep $ map ppr ses) + +instance PrettyLore lore => Pretty (GroupStreamLambda lore) where + ppr (GroupStreamLambda block_size block_offset acc_params arr_params body) = + annot (mapMaybe ppAnnot params) $ + text "fn" <+> + parens (commasep (block_size' : block_offset' : map ppr params)) <+> + text "=>" indent 2 (ppr body) + where params = acc_params ++ arr_params + block_size' = text "int" <+> ppr block_size + block_offset' = text "int" <+> ppr block_offset diff --git a/src/Futhark/Representation/Kernels/Simplify.hs b/src/Futhark/Representation/Kernels/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Kernels/Simplify.hs @@ -0,0 +1,463 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ConstraintKinds #-} +module Futhark.Representation.Kernels.Simplify + ( simplifyKernels + , simplifyLambda + + -- * Building blocks + , simplifyKernelOp + , simplifyKernelExp + ) +where + +import Control.Monad +import Data.Either +import Data.Foldable +import Data.List +import Data.Maybe +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Futhark.Representation.Kernels +import qualified Futhark.Optimise.Simplify.Engine as Engine +import Futhark.Optimise.Simplify.Rules +import Futhark.Optimise.Simplify.Lore +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Pass +import qualified Futhark.Optimise.Simplify as Simplify +import Futhark.Optimise.Simplify.Rule +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Analysis.Rephrase (castStm) + +simpleKernels :: Simplify.SimpleOps Kernels +simpleKernels = Simplify.bindableSimpleOps (simplifyKernelOp simpleInKernel inKernelEnv) + +simpleInKernel :: KernelSpace -> Simplify.SimpleOps InKernel +simpleInKernel = Simplify.bindableSimpleOps . simplifyKernelExp + +simplifyKernels :: Prog Kernels -> PassM (Prog Kernels) +simplifyKernels = + Simplify.simplifyProg simpleKernels kernelRules Simplify.noExtraHoistBlockers + +simplifyLambda :: (HasScope InKernel m, MonadFreshNames m) => + KernelSpace -> Lambda InKernel -> [Maybe VName] -> m (Lambda InKernel) +simplifyLambda kspace = + Simplify.simplifyLambda (simpleInKernel kspace) + inKernelRules Engine.noExtraHoistBlockers + +simplifyKernelOp :: (Engine.SimplifiableLore lore, + Engine.SimplifiableLore outerlore, + BodyAttr outerlore ~ (), BodyAttr lore ~ (), + ExpAttr lore ~ ExpAttr outerlore, + SameScope lore outerlore, + RetType lore ~ RetType outerlore, + BranchType lore ~ BranchType outerlore) => + (KernelSpace -> Engine.SimpleOps lore) -> Engine.Env lore + -> Kernel lore -> Engine.SimpleM outerlore (Kernel (Wise lore), Stms (Wise outerlore)) +simplifyKernelOp mk_ops env (Kernel desc space ts kbody) = do + space' <- Engine.simplify space + ts' <- mapM Engine.simplify ts + outer_vtable <- Engine.askVtable + (((kbody_stms, kbody_res), kbody_hoisted), again) <- + Engine.subSimpleM (mk_ops space) env outer_vtable $ do + par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers + Engine.localVtable (<>scope_vtable) $ + Engine.blockIf (Engine.hasFree bound_here + `Engine.orIf` Engine.isOp + `Engine.orIf` par_blocker + `Engine.orIf` Engine.isConsumed) $ + simplifyKernelBodyM kbody + when again Engine.changed + kbody_hoisted' <- mapM processHoistedStm kbody_hoisted + return (Kernel desc space' ts' $ mkWiseKernelBody () kbody_stms kbody_res, + kbody_hoisted') + where scope_vtable = ST.fromScope scope + scope = scopeOfKernelSpace space + bound_here = S.fromList $ M.keys scope + +simplifyKernelOp _ _ (GetSize key size_class) = return (GetSize key size_class, mempty) +simplifyKernelOp _ _ (GetSizeMax size_class) = return (GetSizeMax size_class, mempty) +simplifyKernelOp _ _ (CmpSizeLe key size_class x) = do + x' <- Engine.simplify x + return (CmpSizeLe key size_class x', mempty) + +processHoistedStm :: (Monad m, + PrettyLore from, + ExpAttr from ~ ExpAttr to, + BodyAttr from ~ BodyAttr to, + RetType from ~ RetType to, + BranchType from ~ BranchType to, + LetAttr from ~ LetAttr to, + FParamAttr from ~ FParamAttr to, + LParamAttr from ~ LParamAttr to) => + Stm from -> m (Stm to) +processHoistedStm bnd + | Just bnd' <- castStm bnd = return bnd' + | otherwise = fail $ "Cannot hoist binding: " ++ pretty bnd + +mkWiseKernelBody :: (Attributes lore, CanBeWise (Op lore)) => + BodyAttr lore -> Stms (Wise lore) -> [KernelResult] -> KernelBody (Wise lore) +mkWiseKernelBody attr bnds res = + let Body attr' _ _ = mkWiseBody attr bnds res_vs + in KernelBody attr' bnds res + where res_vs = map resValue res + resValue (ThreadsReturn _ se) = se + resValue (WriteReturn _ arr _) = Var arr + resValue (ConcatReturns _ _ _ _ v) = Var v + resValue (KernelInPlaceReturn v) = Var v + +inKernelEnv :: Engine.Env InKernel +inKernelEnv = Engine.emptyEnv inKernelRules Simplify.noExtraHoistBlockers + +instance Engine.Simplifiable SplitOrdering where + simplify SplitContiguous = + return SplitContiguous + simplify (SplitStrided stride) = + SplitStrided <$> Engine.simplify stride + +instance Engine.Simplifiable CombineSpace where + simplify (CombineSpace scatter cspace) = + CombineSpace <$> mapM Engine.simplify scatter + <*> mapM (traverse Engine.simplify) cspace + +simplifyKernelExp :: Engine.SimplifiableLore lore => + KernelSpace -> KernelExp lore + -> Engine.SimpleM lore (KernelExp (Wise lore), Stms (Wise lore)) + +simplifyKernelExp _ (Barrier se) = + (,) <$> (Barrier <$> Engine.simplify se) <*> pure mempty + +simplifyKernelExp _ (SplitSpace o w i elems_per_thread) = + (,) <$> (SplitSpace <$> Engine.simplify o <*> Engine.simplify w + <*> Engine.simplify i <*> Engine.simplify elems_per_thread) + <*> pure mempty + +simplifyKernelExp kspace (Combine cspace ts active body) = do + ((body_stms', body_res'), hoisted) <- + wrapbody $ Engine.blockIf (Engine.hasFree bound_here `Engine.orIf` + maybeBlockUnsafe) $ + localScope (scopeOfCombineSpace cspace) $ + Engine.simplifyBody (map (const Observe) ts) body + body' <- Engine.constructBody body_stms' body_res' + (,) <$> (Combine <$> Engine.simplify cspace + <*> mapM Engine.simplify ts + <*> mapM Engine.simplify active + <*> pure body') <*> pure hoisted + where bound_here = S.fromList $ M.keys $ scopeOfCombineSpace cspace + + protectCombineHoisted checkIfActive m = do + (x, stms) <- m + runBinder $ do + if any (not . safeExp . stmExp) stms + then do is_active <- checkIfActive + mapM_ (Engine.protectIf (not . safeExp) is_active) stms + else addStms stms + return x + + (maybeBlockUnsafe, wrapbody) + | [d] <- map snd $ cspaceDims cspace, + d == spaceGroupSize kspace = + (Engine.isFalse True, + protectCombineHoisted $ + letSubExp "active" =<< + foldBinOp LogAnd (constant True) =<< + mapM (uncurry check) active) + | otherwise = + (Engine.isNotSafe, id) + + check v se = + letSubExp "is_active" $ BasicOp $ CmpOp (CmpSlt Int32) (Var v) se + +simplifyKernelExp _ (GroupReduce w lam input) = do + arrs' <- mapM Engine.simplify arrs + nes' <- mapM Engine.simplify nes + w' <- Engine.simplify w + (lam', hoisted) <- Engine.simplifyLambdaSeq lam (map (const Nothing) arrs') + return (GroupReduce w' lam' $ zip nes' arrs', hoisted) + where (nes,arrs) = unzip input + +simplifyKernelExp _ (GroupScan w lam input) = do + w' <- Engine.simplify w + nes' <- mapM Engine.simplify nes + arrs' <- mapM Engine.simplify arrs + (lam', hoisted) <- Engine.simplifyLambdaSeq lam (map (const Nothing) arrs') + return (GroupScan w' lam' $ zip nes' arrs', hoisted) + where (nes,arrs) = unzip input + +simplifyKernelExp _ (GroupGenReduce w dests op bucket vs locks) = do + w' <- Engine.simplify w + dests' <- mapM Engine.simplify dests + (op', hoisted) <- Engine.simplifyLambdaSeq op (map (const Nothing) vs) + bucket' <- Engine.simplify bucket + vs' <- mapM Engine.simplify vs + locks' <- Engine.simplify locks + return (GroupGenReduce w' dests' op' bucket' vs' locks', hoisted) + +simplifyKernelExp _ (GroupStream w maxchunk lam accs arrs) = do + w' <- Engine.simplify w + maxchunk' <- Engine.simplify maxchunk + accs' <- mapM Engine.simplify accs + arrs' <- mapM Engine.simplify arrs + (lam', hoisted) <- simplifyGroupStreamLambda lam w' maxchunk' arrs' + return (GroupStream w' maxchunk' lam' accs' arrs', hoisted) + +simplifyKernelBodyM :: Engine.SimplifiableLore lore => + KernelBody lore + -> Engine.SimpleM lore (Engine.SimplifiedBody lore [KernelResult]) +simplifyKernelBodyM (KernelBody _ stms res) = + Engine.simplifyStms stms $ do res' <- mapM Engine.simplify res + return ((res', UT.usages $ freeIn res'), mempty) + +simplifyGroupStreamLambda :: Engine.SimplifiableLore lore => + GroupStreamLambda lore + -> SubExp -> SubExp -> [VName] + -> Engine.SimpleM lore (GroupStreamLambda (Wise lore), Stms (Wise lore)) +simplifyGroupStreamLambda lam w max_chunk arrs = do + let GroupStreamLambda block_size block_offset acc_params arr_params body = lam + bound_here = S.fromList $ block_size : block_offset : + map paramName (acc_params ++ arr_params) + ((body_stms', body_res'), hoisted) <- + Engine.enterLoop $ + Engine.bindLoopVar block_size Int32 max_chunk $ + Engine.bindLoopVar block_offset Int32 w $ + Engine.bindLParams acc_params $ + Engine.bindChunkLParams block_offset (zip arr_params arrs) $ + Engine.blockIf (Engine.hasFree bound_here `Engine.orIf` Engine.isConsumed) $ + Engine.simplifyBody (replicate (length (bodyResult body)) Observe) body + acc_params' <- mapM (Engine.simplifyParam Engine.simplify) acc_params + arr_params' <- mapM (Engine.simplifyParam Engine.simplify) arr_params + body' <- Engine.constructBody body_stms' body_res' + return (GroupStreamLambda block_size block_offset acc_params' arr_params' body', hoisted) + +instance Engine.Simplifiable KernelSpace where + simplify (KernelSpace gtid ltid gid num_threads num_groups group_size structure) = + KernelSpace gtid ltid gid + <$> Engine.simplify num_threads + <*> Engine.simplify num_groups + <*> Engine.simplify group_size + <*> Engine.simplify structure + +instance Engine.Simplifiable SpaceStructure where + simplify (FlatThreadSpace dims) = + FlatThreadSpace <$> (zip gtids <$> mapM Engine.simplify gdims) + where (gtids, gdims) = unzip dims + simplify (NestedThreadSpace dims) = + NestedThreadSpace + <$> (zip4 gtids + <$> mapM Engine.simplify gdims + <*> pure ltids + <*> mapM Engine.simplify ldims) + where (gtids, gdims, ltids, ldims) = unzip4 dims + +instance Engine.Simplifiable KernelResult where + simplify (ThreadsReturn threads what) = + ThreadsReturn <$> Engine.simplify threads <*> Engine.simplify what + simplify (WriteReturn ws a res) = + WriteReturn <$> Engine.simplify ws <*> Engine.simplify a <*> Engine.simplify res + simplify (ConcatReturns o w pte moffset what) = + ConcatReturns + <$> Engine.simplify o + <*> Engine.simplify w + <*> Engine.simplify pte + <*> Engine.simplify moffset + <*> Engine.simplify what + simplify (KernelInPlaceReturn what) = + KernelInPlaceReturn <$> Engine.simplify what + +instance Engine.Simplifiable WhichThreads where + simplify AllThreads = pure AllThreads + simplify OneResultPerGroup = pure OneResultPerGroup + simplify ThreadsInSpace = pure ThreadsInSpace + simplify (ThreadsPerGroup limit) = + ThreadsPerGroup <$> mapM Engine.simplify limit + +instance BinderOps (Wise Kernels) where + mkExpAttrB = bindableMkExpAttrB + mkBodyB = bindableMkBodyB + mkLetNamesB = bindableMkLetNamesB + +instance BinderOps (Wise InKernel) where + mkExpAttrB = bindableMkExpAttrB + mkBodyB = bindableMkBodyB + mkLetNamesB = bindableMkLetNamesB + +kernelRules :: RuleBook (Wise Kernels) +kernelRules = standardRules <> + ruleBook [RuleOp removeInvariantKernelResults] + [RuleOp distributeKernelResults, + RuleBasicOp removeUnnecessaryCopy] + +fuseStreamIota :: TopDownRuleOp (Wise InKernel) +fuseStreamIota vtable pat _ (GroupStream w max_chunk lam accs arrs) + | ([(iota_cs, iota_param, iota_start, iota_stride, iota_t)], params_and_arrs) <- + partitionEithers $ zipWith (isIota vtable) (groupStreamArrParams lam) arrs = do + + let (arr_params', arrs') = unzip params_and_arrs + chunk_size = groupStreamChunkSize lam + offset = groupStreamChunkOffset lam + + body' <- insertStmsM $ inScopeOf lam $ certifying iota_cs $ do + -- Convert index to appropriate type. + offset' <- asIntS iota_t $ Var offset + offset'' <- letSubExp "offset_by_stride" $ + BasicOp $ BinOp (Mul iota_t) offset' iota_stride + start <- letSubExp "iota_start" $ + BasicOp $ BinOp (Add iota_t) offset'' iota_start + letBindNames_ [paramName iota_param] $ + BasicOp $ Iota (Var chunk_size) start iota_stride iota_t + return $ groupStreamLambdaBody lam + let lam' = lam { groupStreamArrParams = arr_params', + groupStreamLambdaBody = body' + } + letBind_ pat $ Op $ GroupStream w max_chunk lam' accs arrs' +fuseStreamIota _ _ _ _ = cannotSimplify + +isIota :: ST.SymbolTable lore -> a -> VName + -> Either (Certificates, a, SubExp, SubExp, IntType) (a, VName) +isIota vtable chunk arr + | Just (BasicOp (Iota _ x s it), cs) <- ST.lookupExp arr vtable = + Left (cs, chunk, x, s, it) + | otherwise = + Right (chunk, arr) + +-- If a kernel produces something invariant to the kernel, turn it +-- into a replicate. +removeInvariantKernelResults :: TopDownRuleOp (Wise Kernels) +removeInvariantKernelResults vtable (Pattern [] kpes) attr + (Kernel desc space ts (KernelBody _ kstms kres)) = do + (ts', kpes', kres') <- + unzip3 <$> filterM checkForInvarianceResult (zip3 ts kpes kres) + + -- Check if we did anything at all. + when (kres == kres') + cannotSimplify + + addStm $ Let (Pattern [] kpes') attr $ Op $ Kernel desc space ts' $ + mkWiseKernelBody () kstms kres' + where isInvariant Constant{} = True + isInvariant (Var v) = isJust $ ST.lookup v vtable + + num_threads = spaceNumThreads space + space_dims = map snd $ spaceDimensions space + + checkForInvarianceResult (_, pe, ThreadsReturn threads se) + | isInvariant se = + case threads of + AllThreads -> do + letBindNames_ [patElemName pe] $ BasicOp $ + Replicate (Shape [num_threads]) se + return False + ThreadsInSpace -> do + let rep a d = BasicOp . Replicate (Shape [d]) <$> letSubExp "rep" a + letBindNames_ [patElemName pe] =<< + foldM rep (BasicOp (SubExp se)) (reverse space_dims) + return False + _ -> return True + checkForInvarianceResult _ = + return True +removeInvariantKernelResults _ _ _ _ = cannotSimplify + +-- Some kernel results can be moved outside the kernel, which can +-- simplify further analysis. +distributeKernelResults :: BottomUpRuleOp (Wise Kernels) +distributeKernelResults (vtable, used) + (Pattern [] kpes) attr (Kernel desc kspace kts (KernelBody _ kstms kres)) = do + -- Iterate through the bindings. For each, we check whether it is + -- in kres and can be moved outside. If so, we remove it from kres + -- and kpes and make it a binding outside. + (kpes', kts', kres', kstms_rev) <- localScope (scopeOfKernelSpace kspace) $ + foldM distribute (kpes, kts, kres, []) kstms + + when (kpes' == kpes) + cannotSimplify + + addStm $ Let (Pattern [] kpes') attr $ + Op $ Kernel desc kspace kts' $ mkWiseKernelBody () (stmsFromList $ reverse kstms_rev) kres' + where + free_in_kstms = fold $ fmap freeInStm kstms + + distribute (kpes', kts', kres', kstms_rev) bnd + | Let (Pattern [] [pe]) _ (BasicOp (Index arr slice)) <- bnd, + kspace_slice <- map (DimFix . Var . fst) $ spaceDimensions kspace, + kspace_slice `isPrefixOf` slice, + remaining_slice <- drop (length kspace_slice) slice, + all (isJust . flip ST.lookup vtable) $ S.toList $ + freeIn arr <> freeIn remaining_slice, + Just (kpe, kpes'', kts'', kres'') <- isResult kpes' kts' kres' pe = do + let outer_slice = map (\(_, d) -> DimSlice + (constant (0::Int32)) + d + (constant (1::Int32))) $ + spaceDimensions kspace + index kpe' = letBind_ (Pattern [] [kpe']) $ BasicOp $ Index arr $ + outer_slice <> remaining_slice + if patElemName kpe `UT.isConsumed` used + then do precopy <- newVName $ baseString (patElemName kpe) <> "_precopy" + index kpe { patElemName = precopy } + letBind_ (Pattern [] [kpe]) $ BasicOp $ Copy precopy + else index kpe + return (kpes'', kts'', kres'', + if patElemName pe `S.member` free_in_kstms + then bnd : kstms_rev + else kstms_rev) + + distribute (kpes', kts', kres', kstms_rev) bnd = + return (kpes', kts', kres', bnd : kstms_rev) + + isResult kpes' kts' kres' pe = + case partition matches $ zip3 kpes' kts' kres' of + ([(kpe,_,_)], kpes_and_kres) + | (kpes'', kts'', kres'') <- unzip3 kpes_and_kres -> + Just (kpe, kpes'', kts'', kres'') + _ -> Nothing + where matches (_, _, kre) = kre == ThreadsReturn ThreadsInSpace (Var $ patElemName pe) +distributeKernelResults _ _ _ _ = cannotSimplify + +simplifyKnownIterationStream :: TopDownRuleOp (Wise InKernel) +-- Remove GroupStreams over single-element arrays. Not much to stream +-- here, and no information to exploit. +simplifyKnownIterationStream _ pat _ (GroupStream (Constant v) _ lam accs arrs) + | oneIsh v = do + let GroupStreamLambda chunk_size chunk_offset acc_params arr_params body = lam + + letBindNames_ [chunk_size] $ BasicOp $ SubExp $ constant (1::Int32) + + letBindNames_ [chunk_offset] $ BasicOp $ SubExp $ constant (0::Int32) + + forM_ (zip acc_params accs) $ \(p,a) -> + letBindNames_ [paramName p] $ BasicOp $ SubExp a + + forM_ (zip arr_params arrs) $ \(p,a) -> + letBindNames_ [paramName p] $ BasicOp $ Index a $ + fullSlice (paramType p) + [DimSlice (Var chunk_offset) (Var chunk_size) (constant (1::Int32))] + + res <- bodyBind body + forM_ (zip (patternElements pat) res) $ \(pe,r) -> + letBindNames_ [patElemName pe] $ BasicOp $ SubExp r +simplifyKnownIterationStream _ _ _ _ = cannotSimplify + +removeUnusedStreamInputs :: TopDownRuleOp (Wise InKernel) +removeUnusedStreamInputs _ pat _ (GroupStream w maxchunk lam accs arrs) + | (used,unused) <- partition (isUsed . paramName . fst) $ zip arr_params arrs, + not $ null unused = do + let (arr_params', arrs') = unzip used + lam' = GroupStreamLambda chunk_size chunk_offset acc_params arr_params' body + letBind_ pat $ Op $ GroupStream w maxchunk lam' accs arrs' + where GroupStreamLambda chunk_size chunk_offset acc_params arr_params body = lam + + isUsed = (`S.member` freeInBody body) +removeUnusedStreamInputs _ _ _ _ = cannotSimplify + +inKernelRules :: RuleBook (Wise InKernel) +inKernelRules = standardRules <> + ruleBook [RuleOp fuseStreamIota, + RuleOp simplifyKnownIterationStream, + RuleOp removeUnusedStreamInputs] [] diff --git a/src/Futhark/Representation/Kernels/Sizes.hs b/src/Futhark/Representation/Kernels/Sizes.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Kernels/Sizes.hs @@ -0,0 +1,27 @@ +module Futhark.Representation.Kernels.Sizes + ( SizeClass (..), KernelPath ) + where + +import Futhark.Util.Pretty +import Language.Futhark.Core (VName) +import Futhark.Representation.AST.Pretty () + +-- | An indication of which comparisons have been performed to get to +-- this point, as well as the result of each comparison. +type KernelPath = [(VName, Bool)] + +-- | The class of some kind of configurable size. Each class may +-- impose constraints on the valid values. +data SizeClass = SizeThreshold KernelPath + | SizeGroup + | SizeNumGroups + | SizeTile + deriving (Eq, Ord, Show) + +instance Pretty SizeClass where + ppr (SizeThreshold path) = text $ "threshold (" ++ unwords (map pStep path) ++ ")" + where pStep (v, True) = pretty v + pStep (v, False) = '!' : pretty v + ppr SizeGroup = text "group_size" + ppr SizeNumGroups = text "num_groups" + ppr SizeTile = text "tile_size" diff --git a/src/Futhark/Representation/Primitive.hs b/src/Futhark/Representation/Primitive.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Primitive.hs @@ -0,0 +1,1074 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE LambdaCase #-} +-- | Definitions of primitive types, the values that inhabit these +-- types, and operations on these values. A primitive value can also +-- be called a scalar. +-- +-- Essentially, this module describes the subset of the (internal) +-- Futhark language that operates on primitive types. +module Futhark.Representation.Primitive + ( -- * Types + IntType (..), allIntTypes + , FloatType (..), allFloatTypes + , PrimType (..), allPrimTypes + + -- * Values + , IntValue(..) + , intValue, intValueType, valueIntegral + , FloatValue(..) + , floatValue, floatValueType + , PrimValue(..) + , primValueType + , blankPrimValue + + -- * Operations + , UnOp (..), allUnOps + , BinOp (..), allBinOps + , ConvOp (..), allConvOps + , CmpOp (..), allCmpOps + + -- ** Unary Operations + , doUnOp + , doComplement + , doAbs, doFAbs + , doSSignum, doUSignum + + -- ** Binary Operations + , doBinOp + , doAdd, doMul, doSDiv, doSMod + , doPow + + -- ** Conversion Operations + , doConvOp + , doZExt, doSExt + , doFPConv + , doFPToUI, doFPToSI + , doUIToFP, doSIToFP + , intToInt64, intToWord64 + + -- * Comparison Operations + , doCmpOp + , doCmpEq + , doCmpUlt, doCmpUle + , doCmpSlt, doCmpSle + , doFCmpLt, doFCmpLe + + -- * Type Of + , binOpType + , unOpType + , cmpOpType + , convOpType + + -- * Primitive functions + , primFuns + + -- * Utility + , zeroIsh + , oneIsh + , negativeIsh + , primBitSize + , primByteSize + , commutativeBinOp + + -- * Prettyprinting + , convOpFun + , prettySigned + ) + where + +import Control.Applicative +import Data.Binary.IEEE754 (floatToWord, wordToFloat, doubleToWord, wordToDouble) +import Data.Bits +import Data.Int (Int16, Int32, Int64, Int8) +import qualified Data.Map as M +import Data.Word + +import Prelude + +import Futhark.Util.Pretty +import Futhark.Util (roundFloat, roundDouble) + +-- | An integer type, ordered by size. Note that signedness is not a +-- property of the type, but a property of the operations performed on +-- values of these types. +data IntType = Int8 + | Int16 + | Int32 + | Int64 + deriving (Eq, Ord, Show, Enum, Bounded) + +instance Pretty IntType where + ppr Int8 = text "i8" + ppr Int16 = text "i16" + ppr Int32 = text "i32" + ppr Int64 = text "i64" + +-- | A list of all integer types. +allIntTypes :: [IntType] +allIntTypes = [minBound..maxBound] + +-- | A floating point type. +data FloatType = Float32 + | Float64 + deriving (Eq, Ord, Show, Enum, Bounded) + +instance Pretty FloatType where + ppr Float32 = text "f32" + ppr Float64 = text "f64" + +-- | A list of all floating-point types. +allFloatTypes :: [FloatType] +allFloatTypes = [minBound..maxBound] + +-- | Low-level primitive types. +data PrimType = IntType IntType + | FloatType FloatType + | Bool + | Cert + deriving (Eq, Ord, Show) + +instance Enum PrimType where + toEnum 0 = IntType Int8 + toEnum 1 = IntType Int16 + toEnum 2 = IntType Int32 + toEnum 3 = IntType Int64 + toEnum 4 = FloatType Float32 + toEnum 5 = FloatType Float64 + toEnum 6 = Bool + toEnum _ = Cert + + fromEnum (IntType Int8) = 0 + fromEnum (IntType Int16) = 1 + fromEnum (IntType Int32) = 2 + fromEnum (IntType Int64) = 3 + fromEnum (FloatType Float32) = 4 + fromEnum (FloatType Float64) = 5 + fromEnum Bool = 6 + fromEnum Cert = 7 + +instance Bounded PrimType where + minBound = IntType Int8 + maxBound = Cert + +instance Pretty PrimType where + ppr (IntType t) = ppr t + ppr (FloatType t) = ppr t + ppr Bool = text "bool" + ppr Cert = text "cert" + +-- | A list of all primitive types. +allPrimTypes :: [PrimType] +allPrimTypes = map IntType allIntTypes ++ + map FloatType allFloatTypes ++ + [Bool, Cert] + +-- | An integer value. +data IntValue = Int8Value !Int8 + | Int16Value !Int16 + | Int32Value !Int32 + | Int64Value !Int64 + deriving (Eq, Ord, Show) + +instance Pretty IntValue where + ppr (Int8Value v) = text $ show v ++ "i8" + ppr (Int16Value v) = text $ show v ++ "i16" + ppr (Int32Value v) = text $ show v ++ "i32" + ppr (Int64Value v) = text $ show v ++ "i64" + +-- | Create an 'IntValue' from a type and an 'Integer'. +intValue :: Integral int => IntType -> int -> IntValue +intValue Int8 = Int8Value . fromIntegral +intValue Int16 = Int16Value . fromIntegral +intValue Int32 = Int32Value . fromIntegral +intValue Int64 = Int64Value . fromIntegral + +intValueType :: IntValue -> IntType +intValueType Int8Value{} = Int8 +intValueType Int16Value{} = Int16 +intValueType Int32Value{} = Int32 +intValueType Int64Value{} = Int64 + +-- | Convert an 'IntValue' to any 'Integral' type. +valueIntegral ::Integral int => IntValue -> int +valueIntegral (Int8Value v) = fromIntegral v +valueIntegral (Int16Value v) = fromIntegral v +valueIntegral (Int32Value v) = fromIntegral v +valueIntegral (Int64Value v) = fromIntegral v + +-- | A floating-point value. +data FloatValue = Float32Value !Float + | Float64Value !Double + deriving (Eq, Ord, Show) + + +instance Pretty FloatValue where + ppr (Float32Value v) + | isInfinite v, v >= 0 = text "f32.inf" + | isInfinite v, v < 0 = text "-f32.inf" + | isNaN v = text "f32.nan" + | otherwise = text $ show v ++ "f32" + ppr (Float64Value v) + | isInfinite v, v >= 0 = text "f64.inf" + | isInfinite v, v < 0 = text "-f64.inf" + | isNaN v = text "f64.nan" + | otherwise = text $ show v ++ "f64" + +-- | Create a 'FloatValue' from a type and a 'Rational'. +floatValue :: Real num => FloatType -> num -> FloatValue +floatValue Float32 = Float32Value . fromRational . toRational +floatValue Float64 = Float64Value . fromRational . toRational + +floatValueType :: FloatValue -> FloatType +floatValueType Float32Value{} = Float32 +floatValueType Float64Value{} = Float64 + +-- | Non-array values. +data PrimValue = IntValue !IntValue + | FloatValue !FloatValue + | BoolValue !Bool + | Checked -- ^ The only value of type @cert@. + deriving (Eq, Ord, Show) + +instance Pretty PrimValue where + ppr (IntValue v) = ppr v + ppr (BoolValue True) = text "true" + ppr (BoolValue False) = text "false" + ppr (FloatValue v) = ppr v + ppr Checked = text "checked" + +-- | The type of a basic value. +primValueType :: PrimValue -> PrimType +primValueType (IntValue v) = IntType $ intValueType v +primValueType (FloatValue v) = FloatType $ floatValueType v +primValueType BoolValue{} = Bool +primValueType Checked = Cert + +-- | A "blank" value of the given primitive type - this is zero, or +-- whatever is close to it. Don't depend on this value, but use it +-- for e.g. creating arrays to be populated by do-loops. +blankPrimValue :: PrimType -> PrimValue +blankPrimValue (IntType Int8) = IntValue $ Int8Value 0 +blankPrimValue (IntType Int16) = IntValue $ Int16Value 0 +blankPrimValue (IntType Int32) = IntValue $ Int32Value 0 +blankPrimValue (IntType Int64) = IntValue $ Int64Value 0 +blankPrimValue (FloatType Float32) = FloatValue $ Float32Value 0.0 +blankPrimValue (FloatType Float64) = FloatValue $ Float64Value 0.0 +blankPrimValue Bool = BoolValue False +blankPrimValue Cert = Checked + +-- | Various unary operators. It is a bit ad-hoc what is a unary +-- operator and what is a built-in function. Perhaps these should all +-- go away eventually. +data UnOp = Not -- ^ E.g., @! True == False@. + | Complement IntType -- ^ E.g., @~(~1) = 1@. + | Abs IntType -- ^ @abs(-2) = 2@. + | FAbs FloatType -- ^ @fabs(-2.0) = 2.0@. + | SSignum IntType -- ^ Signed sign function: @ssignum(-2)@ = -1. + | USignum IntType -- ^ Unsigned sign function: @usignum(2)@ = 1. + deriving (Eq, Ord, Show) + +-- | Binary operators. These correspond closely to the binary operators in +-- LLVM. Most are parametrised by their expected input and output +-- types. +data BinOp = Add IntType -- ^ Integer addition. + | FAdd FloatType -- ^ Floating-point addition. + + | Sub IntType -- ^ Integer subtraction. + | FSub FloatType -- ^ Floating-point subtraction. + + | Mul IntType -- ^ Integer multiplication. + | FMul FloatType -- ^ Floating-point multiplication. + + | UDiv IntType + -- ^ Unsigned integer division. Rounds towards + -- negativity infinity. Note: this is different + -- from LLVM. + | SDiv IntType + -- ^ Signed integer division. Rounds towards + -- negativity infinity. Note: this is different + -- from LLVM. + | FDiv FloatType -- ^ Floating-point division. + + | UMod IntType + -- ^ Unsigned integer modulus; the countepart to 'UDiv'. + | SMod IntType + -- ^ Signed integer modulus; the countepart to 'SDiv'. + + | SQuot IntType + -- ^ Signed integer division. Rounds towards zero. + -- This corresponds to the @sdiv@ instruction in LLVM. + | SRem IntType + -- ^ Signed integer division. Rounds towards zero. + -- This corresponds to the @srem@ instruction in LLVM. + + | SMin IntType + -- ^ Returns the smallest of two signed integers. + | UMin IntType + -- ^ Returns the smallest of two unsigned integers. + | FMin FloatType + -- ^ Returns the smallest of two floating-point numbers. + | SMax IntType + -- ^ Returns the greatest of two signed integers. + | UMax IntType + -- ^ Returns the greatest of two unsigned integers. + | FMax FloatType + -- ^ Returns the greatest of two floating-point numbers. + + | Shl IntType -- ^ Left-shift. + | LShr IntType -- ^ Logical right-shift, zero-extended. + | AShr IntType -- ^ Arithmetic right-shift, sign-extended. + + | And IntType -- ^ Bitwise and. + | Or IntType -- ^ Bitwise or. + | Xor IntType -- ^ Bitwise exclusive-or. + + | Pow IntType -- ^ Integer exponentiation. + | FPow FloatType -- ^ Floating-point exponentiation. + + | LogAnd -- ^ Boolean and - not short-circuiting. + | LogOr -- ^ Boolean or - not short-circuiting. + deriving (Eq, Ord, Show) + +-- | Comparison operators are like 'BinOp's, but they return 'Bool's. +-- The somewhat ugly constructor names are straight out of LLVM. +data CmpOp = CmpEq PrimType -- ^ All types equality. + | CmpUlt IntType -- ^ Unsigned less than. + | CmpUle IntType -- ^ Unsigned less than or equal. + | CmpSlt IntType -- ^ Signed less than. + | CmpSle IntType -- ^ Signed less than or equal. + + -- Comparison operators for floating-point values. TODO: extend + -- this to handle NaNs and such, like the LLVM fcmp instruction. + | FCmpLt FloatType -- ^ Floating-point less than. + | FCmpLe FloatType -- ^ Floating-point less than or equal. + + -- Boolean comparison. + | CmpLlt -- ^ Boolean less than. + | CmpLle -- ^ Boolean less than or equal. + deriving (Eq, Ord, Show) + +-- | Conversion operators try to generalise the @from t0 x to t1@ +-- instructions from LLVM. +data ConvOp = ZExt IntType IntType + -- ^ Zero-extend the former integer type to the latter. + -- If the new type is smaller, the result is a + -- truncation. + | SExt IntType IntType + -- ^ Sign-extend the former integer type to the latter. + -- If the new type is smaller, the result is a + -- truncation. + | FPConv FloatType FloatType + -- ^ Convert value of the former floating-point type to + -- the latter. If the new type is smaller, the result + -- is a truncation. + | FPToUI FloatType IntType + -- ^ Convert a floating-point value to the nearest + -- unsigned integer (rounding towards zero). + | FPToSI FloatType IntType + -- ^ Convert a floating-point value to the nearest + -- signed integer (rounding towards zero). + | UIToFP IntType FloatType + -- ^ Convert an unsigned integer to a floating-point value. + | SIToFP IntType FloatType + -- ^ Convert a signed integer to a floating-point value. + | IToB IntType + -- ^ Convert an integer to a boolean value. Zero + -- becomes false; anything else is true. + | BToI IntType + -- ^ Convert a boolean to an integer. True is converted + -- to 1 and False to 0. + deriving (Eq, Ord, Show) + +-- | A list of all unary operators for all types. +allUnOps :: [UnOp] +allUnOps = Not : + map Complement [minBound..maxBound] ++ + map Abs [minBound..maxBound] ++ + map FAbs [minBound..maxBound] ++ + map SSignum [minBound..maxBound] ++ + map USignum [minBound..maxBound] + +-- | A list of all binary operators for all types. +allBinOps :: [BinOp] +allBinOps = concat [ map Add allIntTypes + , map FAdd allFloatTypes + , map Sub allIntTypes + , map FSub allFloatTypes + , map Mul allIntTypes + , map FMul allFloatTypes + , map UDiv allIntTypes + , map SDiv allIntTypes + , map FDiv allFloatTypes + , map UMod allIntTypes + , map SMod allIntTypes + , map SQuot allIntTypes + , map SRem allIntTypes + , map SMin allIntTypes + , map UMin allIntTypes + , map FMin allFloatTypes + , map SMax allIntTypes + , map UMax allIntTypes + , map FMax allFloatTypes + , map Shl allIntTypes + , map LShr allIntTypes + , map AShr allIntTypes + , map And allIntTypes + , map Or allIntTypes + , map Xor allIntTypes + , map Pow allIntTypes + , map FPow allFloatTypes + , [LogAnd, LogOr] + ] + +-- | A list of all comparison operators for all types. +allCmpOps :: [CmpOp] +allCmpOps = concat [ map CmpEq allPrimTypes + , map CmpUlt allIntTypes + , map CmpUle allIntTypes + , map CmpSlt allIntTypes + , map CmpSle allIntTypes + , map FCmpLt allFloatTypes + , map FCmpLe allFloatTypes + ] + +-- | A list of all conversion operators for all types. +allConvOps :: [ConvOp] +allConvOps = concat [ ZExt <$> allIntTypes <*> allIntTypes + , SExt <$> allIntTypes <*> allIntTypes + , FPConv <$> allFloatTypes <*> allFloatTypes + , FPToUI <$> allFloatTypes <*> allIntTypes + , FPToSI <$> allFloatTypes <*> allIntTypes + , UIToFP <$> allIntTypes <*> allFloatTypes + , SIToFP <$> allIntTypes <*> allFloatTypes + , IToB <$> allIntTypes + , BToI <$> allIntTypes + ] + +doUnOp :: UnOp -> PrimValue -> Maybe PrimValue +doUnOp Not (BoolValue b) = Just $ BoolValue $ not b +doUnOp Complement{} (IntValue v) = Just $ IntValue $ doComplement v +doUnOp Abs{} (IntValue v) = Just $ IntValue $ doAbs v +doUnOp FAbs{} (FloatValue v) = Just $ FloatValue $ doFAbs v +doUnOp SSignum{} (IntValue v) = Just $ IntValue $ doSSignum v +doUnOp USignum{} (IntValue v) = Just $ IntValue $ doUSignum v +doUnOp _ _ = Nothing + +-- | E.g., @~(~1) = 1@. +doComplement :: IntValue -> IntValue +doComplement v = intValue (intValueType v) $ complement $ intToInt64 v + +-- | @abs(-2) = 2@. +doAbs :: IntValue -> IntValue +doAbs v = intValue (intValueType v) $ abs $ intToInt64 v + +-- | @abs(-2.0) = 2.0@. +doFAbs :: FloatValue -> FloatValue +doFAbs v = floatValue (floatValueType v) $ abs $ floatToDouble v + +-- | @ssignum(-2)@ = -1. +doSSignum :: IntValue -> IntValue +doSSignum v = intValue (intValueType v) $ signum $ intToInt64 v + +-- | @usignum(-2)@ = -1. +doUSignum :: IntValue -> IntValue +doUSignum v = intValue (intValueType v) $ signum $ intToWord64 v + +doBinOp :: BinOp -> PrimValue -> PrimValue -> Maybe PrimValue +doBinOp Add{} = doIntBinOp doAdd +doBinOp FAdd{} = doFloatBinOp (+) (+) +doBinOp Sub{} = doIntBinOp doSub +doBinOp FSub{} = doFloatBinOp (-) (-) +doBinOp Mul{} = doIntBinOp doMul +doBinOp FMul{} = doFloatBinOp (*) (*) +doBinOp UDiv{} = doRiskyIntBinOp doUDiv +doBinOp SDiv{} = doRiskyIntBinOp doSDiv +doBinOp FDiv{} = doFloatBinOp (/) (/) +doBinOp UMod{} = doRiskyIntBinOp doUMod +doBinOp SMod{} = doRiskyIntBinOp doSMod +doBinOp SQuot{} = doRiskyIntBinOp doSQuot +doBinOp SRem{} = doRiskyIntBinOp doSRem +doBinOp SMin{} = doIntBinOp doSMin +doBinOp UMin{} = doIntBinOp doUMin +doBinOp FMin{} = doFloatBinOp min min +doBinOp SMax{} = doIntBinOp doSMax +doBinOp UMax{} = doIntBinOp doUMax +doBinOp FMax{} = doFloatBinOp max max +doBinOp Shl{} = doIntBinOp doShl +doBinOp LShr{} = doIntBinOp doLShr +doBinOp AShr{} = doIntBinOp doAShr +doBinOp And{} = doIntBinOp doAnd +doBinOp Or{} = doIntBinOp doOr +doBinOp Xor{} = doIntBinOp doXor +doBinOp Pow{} = doRiskyIntBinOp doPow +doBinOp FPow{} = doFloatBinOp (**) (**) +doBinOp LogAnd{} = doBoolBinOp (&&) +doBinOp LogOr{} = doBoolBinOp (||) + +doIntBinOp :: (IntValue -> IntValue -> IntValue) -> PrimValue -> PrimValue + -> Maybe PrimValue +doIntBinOp f (IntValue v1) (IntValue v2) = + Just $ IntValue $ f v1 v2 +doIntBinOp _ _ _ = Nothing + +doRiskyIntBinOp :: (IntValue -> IntValue -> Maybe IntValue) -> PrimValue -> PrimValue + -> Maybe PrimValue +doRiskyIntBinOp f (IntValue v1) (IntValue v2) = + IntValue <$> f v1 v2 +doRiskyIntBinOp _ _ _ = Nothing + +doFloatBinOp :: (Float -> Float -> Float) + -> (Double -> Double -> Double) + -> PrimValue -> PrimValue + -> Maybe PrimValue +doFloatBinOp f32 _ (FloatValue (Float32Value v1)) (FloatValue (Float32Value v2)) = + Just $ FloatValue $ Float32Value $ f32 v1 v2 +doFloatBinOp _ f64 (FloatValue (Float64Value v1)) (FloatValue (Float64Value v2)) = + Just $ FloatValue $ Float64Value $ f64 v1 v2 +doFloatBinOp _ _ _ _ = Nothing + +doBoolBinOp :: (Bool -> Bool -> Bool) -> PrimValue -> PrimValue + -> Maybe PrimValue +doBoolBinOp f (BoolValue v1) (BoolValue v2) = + Just $ BoolValue $ f v1 v2 +doBoolBinOp _ _ _ = Nothing + +-- | Integer addition. +doAdd :: IntValue -> IntValue -> IntValue +doAdd v1 v2 = intValue (intValueType v1) $ intToInt64 v1 + intToInt64 v2 + +-- | Integer subtraction. +doSub :: IntValue -> IntValue -> IntValue +doSub v1 v2 = intValue (intValueType v1) $ intToInt64 v1 - intToInt64 v2 + +-- | Integer multiplication. +doMul :: IntValue -> IntValue -> IntValue +doMul v1 v2 = intValue (intValueType v1) $ intToInt64 v1 * intToInt64 v2 + +-- | Unsigned integer division. Rounds towards +-- negativity infinity. Note: this is different +-- from LLVM. +doUDiv :: IntValue -> IntValue -> Maybe IntValue +doUDiv v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToWord64 v1 `div` intToWord64 v2 + +-- | Signed integer division. Rounds towards +-- negativity infinity. Note: this is different +-- from LLVM. +doSDiv :: IntValue -> IntValue -> Maybe IntValue +doSDiv v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `div` intToInt64 v2 + +-- | Unsigned integer modulus; the countepart to 'UDiv'. +doUMod :: IntValue -> IntValue -> Maybe IntValue +doUMod v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToWord64 v1 `mod` intToWord64 v2 + +-- | Signed integer modulus; the countepart to 'SDiv'. +doSMod :: IntValue -> IntValue -> Maybe IntValue +doSMod v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `mod` intToInt64 v2 + +-- | Signed integer division. Rounds towards zero. +-- This corresponds to the @sdiv@ instruction in LLVM. +doSQuot :: IntValue -> IntValue -> Maybe IntValue +doSQuot v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `quot` intToInt64 v2 + +-- | Signed integer division. Rounds towards zero. +-- This corresponds to the @srem@ instruction in LLVM. +doSRem :: IntValue -> IntValue -> Maybe IntValue +doSRem v1 v2 + | zeroIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 `rem` intToInt64 v2 + +-- | Minimum of two signed integers. +doSMin :: IntValue -> IntValue -> IntValue +doSMin v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `min` intToInt64 v2 + +-- | Minimum of two unsigned integers. +doUMin :: IntValue -> IntValue -> IntValue +doUMin v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `min` intToWord64 v2 + +-- | Maximum of two signed integers. +doSMax :: IntValue -> IntValue -> IntValue +doSMax v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `max` intToInt64 v2 + +-- | Maximum of two unsigned integers. +doUMax :: IntValue -> IntValue -> IntValue +doUMax v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `max` intToWord64 v2 + +-- | Left-shift. +doShl :: IntValue -> IntValue -> IntValue +doShl v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `shift` intToInt v2 + +-- | Logical right-shift, zero-extended. +doLShr :: IntValue -> IntValue -> IntValue +doLShr v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `shift` negate (intToInt v2) + +-- | Arithmetic right-shift, sign-extended. +doAShr :: IntValue -> IntValue -> IntValue +doAShr v1 v2 = intValue (intValueType v1) $ intToInt64 v1 `shift` negate (intToInt v2) + +-- | Bitwise and. +doAnd :: IntValue -> IntValue -> IntValue +doAnd v1 v2 = intValue (intValueType v1) $ intToWord64 v1 .&. intToWord64 v2 + +-- | Bitwise or. +doOr :: IntValue -> IntValue -> IntValue +doOr v1 v2 = intValue (intValueType v1) $ intToWord64 v1 .|. intToWord64 v2 + +-- | Bitwise exclusive-or. +doXor :: IntValue -> IntValue -> IntValue +doXor v1 v2 = intValue (intValueType v1) $ intToWord64 v1 `xor` intToWord64 v2 + +-- | Signed integer exponentatation. +doPow :: IntValue -> IntValue -> Maybe IntValue +doPow v1 v2 + | negativeIshInt v2 = Nothing + | otherwise = Just $ intValue (intValueType v1) $ intToInt64 v1 ^ intToInt64 v2 + +doConvOp :: ConvOp -> PrimValue -> Maybe PrimValue +doConvOp (ZExt _ to) (IntValue v) = Just $ IntValue $ doZExt v to +doConvOp (SExt _ to) (IntValue v) = Just $ IntValue $ doSExt v to +doConvOp (FPConv _ to) (FloatValue v) = Just $ FloatValue $ doFPConv v to +doConvOp (FPToUI _ to) (FloatValue v) = Just $ IntValue $ doFPToUI v to +doConvOp (FPToSI _ to) (FloatValue v) = Just $ IntValue $ doFPToSI v to +doConvOp (UIToFP _ to) (IntValue v) = Just $ FloatValue $ doUIToFP v to +doConvOp (SIToFP _ to) (IntValue v) = Just $ FloatValue $ doSIToFP v to +doConvOp (IToB _) (IntValue v) = Just $ BoolValue $ intToInt64 v /= 0 +doConvOp (BToI to) (BoolValue v) = Just $ IntValue $ intValue to $ if v then 1 else 0::Int +doConvOp _ _ = Nothing + +-- | Zero-extend the given integer value to the size of the given +-- type. If the type is smaller than the given value, the result is a +-- truncation. +doZExt :: IntValue -> IntType -> IntValue +doZExt (Int8Value x) t = intValue t $ toInteger (fromIntegral x :: Word8) +doZExt (Int16Value x) t = intValue t $ toInteger (fromIntegral x :: Word16) +doZExt (Int32Value x) t = intValue t $ toInteger (fromIntegral x :: Word32) +doZExt (Int64Value x) t = intValue t $ toInteger (fromIntegral x :: Word64) + +-- | Sign-extend the given integer value to the size of the given +-- type. If the type is smaller than the given value, the result is a +-- truncation. +doSExt :: IntValue -> IntType -> IntValue +doSExt (Int8Value x) t = intValue t $ toInteger x +doSExt (Int16Value x) t = intValue t $ toInteger x +doSExt (Int32Value x) t = intValue t $ toInteger x +doSExt (Int64Value x) t = intValue t $ toInteger x + +-- | Convert the former floating-point type to the latter. +doFPConv :: FloatValue -> FloatType -> FloatValue +doFPConv (Float32Value v) Float32 = Float32Value v +doFPConv (Float64Value v) Float32 = Float32Value $ fromRational $ toRational v +doFPConv (Float64Value v) Float64 = Float64Value v +doFPConv (Float32Value v) Float64 = Float64Value $ fromRational $ toRational v + +-- | Convert a floating-point value to the nearest +-- unsigned integer (rounding towards zero). +doFPToUI :: FloatValue -> IntType -> IntValue +doFPToUI v t = intValue t (truncate $ floatToDouble v :: Word64) + +-- | Convert a floating-point value to the nearest +-- signed integer (rounding towards zero). +doFPToSI :: FloatValue -> IntType -> IntValue +doFPToSI v t = intValue t (truncate $ floatToDouble v :: Word64) + +-- | Convert an unsigned integer to a floating-point value. +doUIToFP :: IntValue -> FloatType -> FloatValue +doUIToFP v t = floatValue t $ intToWord64 v + +-- | Convert a signed integer to a floating-point value. +doSIToFP :: IntValue -> FloatType -> FloatValue +doSIToFP v t = floatValue t $ intToInt64 v + +doCmpOp :: CmpOp -> PrimValue -> PrimValue -> Maybe Bool +doCmpOp CmpEq{} v1 v2 = Just $ v1 == v2 +doCmpOp CmpUlt{} (IntValue v1) (IntValue v2) = Just $ doCmpUlt v1 v2 +doCmpOp CmpUle{} (IntValue v1) (IntValue v2) = Just $ doCmpUle v1 v2 +doCmpOp CmpSlt{} (IntValue v1) (IntValue v2) = Just $ doCmpSlt v1 v2 +doCmpOp CmpSle{} (IntValue v1) (IntValue v2) = Just $ doCmpSle v1 v2 +doCmpOp FCmpLt{} (FloatValue v1) (FloatValue v2) = Just $ doFCmpLt v1 v2 +doCmpOp FCmpLe{} (FloatValue v1) (FloatValue v2) = Just $ doFCmpLe v1 v2 +doCmpOp CmpLlt{} (BoolValue v1) (BoolValue v2) = Just $ not v1 && v2 +doCmpOp CmpLle{} (BoolValue v1) (BoolValue v2) = Just $ not (v1 && not v2) +doCmpOp _ _ _ = Nothing + +-- | Compare any two primtive values for exact equality. +doCmpEq :: PrimValue -> PrimValue -> Bool +doCmpEq v1 v2 = v1 == v2 + +-- | Unsigned less than. +doCmpUlt :: IntValue -> IntValue -> Bool +doCmpUlt v1 v2 = intToWord64 v1 < intToWord64 v2 + +-- | Unsigned less than or equal. +doCmpUle :: IntValue -> IntValue -> Bool +doCmpUle v1 v2 = intToWord64 v1 <= intToWord64 v2 + +-- | Signed less than. +doCmpSlt :: IntValue -> IntValue -> Bool +doCmpSlt = (<) + +-- | Signed less than or equal. +doCmpSle :: IntValue -> IntValue -> Bool +doCmpSle = (<=) + +-- | Floating-point less than. +doFCmpLt :: FloatValue -> FloatValue -> Bool +doFCmpLt = (<) + +-- | Floating-point less than or equal. +doFCmpLe :: FloatValue -> FloatValue -> Bool +doFCmpLe = (<=) + +-- | Translate an 'IntValue' to 'Word64'. This is guaranteed to fit. +intToWord64 :: IntValue -> Word64 +intToWord64 (Int8Value v) = fromIntegral (fromIntegral v :: Word8) +intToWord64 (Int16Value v) = fromIntegral (fromIntegral v :: Word16) +intToWord64 (Int32Value v) = fromIntegral (fromIntegral v :: Word32) +intToWord64 (Int64Value v) = fromIntegral (fromIntegral v :: Word64) + +-- | Translate an 'IntValue' to 'Int64'. This is guaranteed to fit. +intToInt64 :: IntValue -> Int64 +intToInt64 (Int8Value v) = fromIntegral v +intToInt64 (Int16Value v) = fromIntegral v +intToInt64 (Int32Value v) = fromIntegral v +intToInt64 (Int64Value v) = fromIntegral v + +-- | Careful - there is no guarantee this will fit. +intToInt :: IntValue -> Int +intToInt = fromIntegral . intToInt64 + +floatToDouble :: FloatValue -> Double +floatToDouble (Float32Value v) = fromRational $ toRational v +floatToDouble (Float64Value v) = v + +-- | The result type of a binary operator. +binOpType :: BinOp -> PrimType +binOpType (Add t) = IntType t +binOpType (Sub t) = IntType t +binOpType (Mul t) = IntType t +binOpType (SDiv t) = IntType t +binOpType (SMod t) = IntType t +binOpType (SQuot t) = IntType t +binOpType (SRem t) = IntType t +binOpType (UDiv t) = IntType t +binOpType (UMod t) = IntType t +binOpType (SMin t) = IntType t +binOpType (UMin t) = IntType t +binOpType (FMin t) = FloatType t +binOpType (SMax t) = IntType t +binOpType (UMax t) = IntType t +binOpType (FMax t) = FloatType t +binOpType (Shl t) = IntType t +binOpType (LShr t) = IntType t +binOpType (AShr t) = IntType t +binOpType (And t) = IntType t +binOpType (Or t) = IntType t +binOpType (Xor t) = IntType t +binOpType (Pow t) = IntType t +binOpType (FPow t) = FloatType t +binOpType LogAnd = Bool +binOpType LogOr = Bool +binOpType (FAdd t) = FloatType t +binOpType (FSub t) = FloatType t +binOpType (FMul t) = FloatType t +binOpType (FDiv t) = FloatType t + +-- | The operand types of a comparison operator. +cmpOpType :: CmpOp -> PrimType +cmpOpType (CmpEq t) = t +cmpOpType (CmpSlt t) = IntType t +cmpOpType (CmpSle t) = IntType t +cmpOpType (CmpUlt t) = IntType t +cmpOpType (CmpUle t) = IntType t +cmpOpType (FCmpLt t) = FloatType t +cmpOpType (FCmpLe t) = FloatType t +cmpOpType CmpLlt = Bool +cmpOpType CmpLle = Bool + +-- | The operand and result type of a unary operator. +unOpType :: UnOp -> PrimType +unOpType (SSignum t) = IntType t +unOpType (USignum t) = IntType t +unOpType Not = Bool +unOpType (Complement t) = IntType t +unOpType (Abs t) = IntType t +unOpType (FAbs t) = FloatType t + +-- | The input and output types of a conversion operator. +convOpType :: ConvOp -> (PrimType, PrimType) +convOpType (ZExt from to) = (IntType from, IntType to) +convOpType (SExt from to) = (IntType from, IntType to) +convOpType (FPConv from to) = (FloatType from, FloatType to) +convOpType (FPToUI from to) = (FloatType from, IntType to) +convOpType (FPToSI from to) = (FloatType from, IntType to) +convOpType (UIToFP from to) = (IntType from, FloatType to) +convOpType (SIToFP from to) = (IntType from, FloatType to) +convOpType (IToB from) = (IntType from, Bool) +convOpType (BToI to) = (Bool, IntType to) + +-- | A mapping from names of primitive functions to their parameter +-- types, their result type, and a function for evaluating them. +primFuns :: M.Map String ([PrimType], PrimType, + [PrimValue] -> Maybe PrimValue) +primFuns = M.fromList + [ f32 "sqrt32" sqrt, f64 "sqrt64" sqrt + , f32 "log32" log, f64 "log64" log + , f32 "log10_32" (logBase 10), f64 "log10_64" (logBase 10) + , f32 "log2_32" (logBase 2), f64 "log2_64" (logBase 2) + , f32 "exp32" exp, f64 "exp64" exp + , f32 "sin32" sin, f64 "sin64" sin + , f32 "cos32" cos, f64 "cos64" cos + , f32 "tan32" tan, f64 "tan64" tan + , f32 "asin32" asin, f64 "asin64" asin + , f32 "acos32" acos, f64 "acos64" acos + , f32 "atan32" atan, f64 "atan64" atan + + , f32 "round32" roundFloat, f64 "round64" roundDouble + + , ("atan2_32", + ([FloatType Float32, FloatType Float32], FloatType Float32, + \case + [FloatValue (Float32Value x), FloatValue (Float32Value y)] -> + Just $ FloatValue $ Float32Value $ atan2 x y + _ -> Nothing)) + , ("atan2_64", + ([FloatType Float64, FloatType Float64], FloatType Float64, + \case + [FloatValue (Float64Value x), FloatValue (Float64Value y)] -> + Just $ FloatValue $ Float64Value $ atan2 x y + _ -> Nothing)) + + , ("isinf32", + ([FloatType Float32], Bool, + \case + [FloatValue (Float32Value x)] -> Just $ BoolValue $ isInfinite x + _ -> Nothing)) + , ("isinf64", + ([FloatType Float64], Bool, + \case + [FloatValue (Float64Value x)] -> Just $ BoolValue $ isInfinite x + _ -> Nothing)) + + , ("isnan32", + ([FloatType Float32], Bool, + \case + [FloatValue (Float32Value x)] -> Just $ BoolValue $ isNaN x + _ -> Nothing)) + , ("isnan64", + ([FloatType Float64], Bool, + \case + [FloatValue (Float64Value x)] -> Just $ BoolValue $ isNaN x + _ -> Nothing)) + + , ("to_bits32", + ([FloatType Float32], IntType Int32, + \case + [FloatValue (Float32Value x)] -> + Just $ IntValue $ Int32Value $ fromIntegral $ floatToWord x + _ -> Nothing)) + , ("to_bits64", + ([FloatType Float64], IntType Int64, + \case + [FloatValue (Float64Value x)] -> + Just $ IntValue $ Int64Value $ fromIntegral $ doubleToWord x + _ -> Nothing)) + + , ("from_bits32", + ([IntType Int32], FloatType Float32, + \case + [IntValue (Int32Value x)] -> + Just $ FloatValue $ Float32Value $ wordToFloat $ fromIntegral x + _ -> Nothing)) + , ("from_bits64", + ([IntType Int64], FloatType Float64, + \case + [IntValue (Int64Value x)] -> + Just $ FloatValue $ Float64Value $ wordToDouble $ fromIntegral x + _ -> Nothing)) + ] + where f32 s f = (s, ([FloatType Float32], FloatType Float32, f32PrimFun f)) + f64 s f = (s, ([FloatType Float64], FloatType Float64, f64PrimFun f)) + + f32PrimFun f [FloatValue (Float32Value x)] = + Just $ FloatValue $ Float32Value $ f x + f32PrimFun _ _ = Nothing + + f64PrimFun f [FloatValue (Float64Value x)] = + Just $ FloatValue $ Float64Value $ f x + f64PrimFun _ _ = Nothing + +-- | Is the given value kind of zero? +zeroIsh :: PrimValue -> Bool +zeroIsh (IntValue k) = zeroIshInt k +zeroIsh (FloatValue (Float32Value k)) = k == 0 +zeroIsh (FloatValue (Float64Value k)) = k == 0 +zeroIsh (BoolValue False) = True +zeroIsh _ = False + +-- | Is the given value kind of one? +oneIsh :: PrimValue -> Bool +oneIsh (IntValue (Int8Value k)) = k == 1 +oneIsh (IntValue (Int16Value k)) = k == 1 +oneIsh (IntValue (Int32Value k)) = k == 1 +oneIsh (IntValue (Int64Value k)) = k == 1 +oneIsh (FloatValue (Float32Value k)) = k == 1 +oneIsh (FloatValue (Float64Value k)) = k == 1 +oneIsh (BoolValue True) = True +oneIsh _ = False + +-- | Is the given value kind of negative? +negativeIsh :: PrimValue -> Bool +negativeIsh (IntValue k) = negativeIshInt k +negativeIsh (FloatValue (Float32Value k)) = k < 0 +negativeIsh (FloatValue (Float64Value k)) = k < 0 +negativeIsh (BoolValue _) = False +negativeIsh Checked = False + +-- | Is the given integer value kind of zero? +zeroIshInt :: IntValue -> Bool +zeroIshInt (Int8Value k) = k == 0 +zeroIshInt (Int16Value k) = k == 0 +zeroIshInt (Int32Value k) = k == 0 +zeroIshInt (Int64Value k) = k == 0 + +-- | Is the given integer value kind of negative? +negativeIshInt :: IntValue -> Bool +negativeIshInt (Int8Value k) = k < 0 +negativeIshInt (Int16Value k) = k < 0 +negativeIshInt (Int32Value k) = k < 0 +negativeIshInt (Int64Value k) = k < 0 + +-- | The size of a value of a given primitive type in bites. +primBitSize :: PrimType -> Int +primBitSize = (*8) . primByteSize + +-- | The size of a value of a given primitive type in eight-bit bytes. +primByteSize :: Num a => PrimType -> a +primByteSize (IntType t) = intByteSize t +primByteSize (FloatType t) = floatByteSize t +primByteSize Bool = 1 +primByteSize Cert = 1 + +-- | The size of a value of a given integer type in eight-bit bytes. +intByteSize :: Num a => IntType -> a +intByteSize Int8 = 1 +intByteSize Int16 = 2 +intByteSize Int32 = 4 +intByteSize Int64 = 8 + +-- | The size of a value of a given floating-point type in eight-bit bytes. +floatByteSize :: Num a => FloatType -> a +floatByteSize Float32 = 4 +floatByteSize Float64 = 8 + +-- | True if the given binary operator is commutative. +commutativeBinOp :: BinOp -> Bool +commutativeBinOp Add{} = True +commutativeBinOp FAdd{} = True +commutativeBinOp Mul{} = True +commutativeBinOp FMul{} = True +commutativeBinOp And{} = True +commutativeBinOp Or{} = True +commutativeBinOp Xor{} = True +commutativeBinOp LogOr{} = True +commutativeBinOp LogAnd{} = True +commutativeBinOp SMax{} = True +commutativeBinOp SMin{} = True +commutativeBinOp UMax{} = True +commutativeBinOp UMin{} = True +commutativeBinOp FMax{} = True +commutativeBinOp FMin{} = True +commutativeBinOp _ = False + +-- Prettyprinting instances + +instance Pretty BinOp where + ppr (Add t) = taggedI "add" t + ppr (FAdd t) = taggedF "fadd" t + ppr (Sub t) = taggedI "sub" t + ppr (FSub t) = taggedF "fsub" t + ppr (Mul t) = taggedI "mul" t + ppr (FMul t) = taggedF "fmul" t + ppr (UDiv t) = taggedI "udiv" t + ppr (UMod t) = taggedI "umod" t + ppr (SDiv t) = taggedI "sdiv" t + ppr (SMod t) = taggedI "smod" t + ppr (SQuot t) = taggedI "squot" t + ppr (SRem t) = taggedI "srem" t + ppr (FDiv t) = taggedF "fdiv" t + ppr (SMin t) = taggedI "smin" t + ppr (UMin t) = taggedI "umin" t + ppr (FMin t) = taggedF "fmin" t + ppr (SMax t) = taggedI "smax" t + ppr (UMax t) = taggedI "umax" t + ppr (FMax t) = taggedF "fmax" t + ppr (Shl t) = taggedI "shl" t + ppr (LShr t) = taggedI "lshr" t + ppr (AShr t) = taggedI "ashr" t + ppr (And t) = taggedI "and" t + ppr (Or t) = taggedI "or" t + ppr (Xor t) = taggedI "xor" t + ppr (Pow t) = taggedI "pow" t + ppr (FPow t) = taggedF "fpow" t + ppr LogAnd = text "logand" + ppr LogOr = text "logor" + +instance Pretty CmpOp where + ppr (CmpEq t) = text "eq_" <> ppr t + ppr (CmpUlt t) = taggedI "ult" t + ppr (CmpUle t) = taggedI "ule" t + ppr (CmpSlt t) = taggedI "slt" t + ppr (CmpSle t) = taggedI "sle" t + ppr (FCmpLt t) = taggedF "lt" t + ppr (FCmpLe t) = taggedF "le" t + ppr CmpLlt = text "llt" + ppr CmpLle = text "lle" + +instance Pretty ConvOp where + ppr op = convOp (convOpFun op) from to + where (from, to) = convOpType op + +instance Pretty UnOp where + ppr Not = text "!" + ppr (Abs t) = taggedI "abs" t + ppr (FAbs t) = taggedF "fabs" t + ppr (SSignum t) = taggedI "ssignum" t + ppr (USignum t) = taggedI "usignum" t + ppr (Complement t) = taggedI "complement" t + +convOpFun :: ConvOp -> String +convOpFun ZExt{} = "zext" +convOpFun SExt{} = "sext" +convOpFun FPConv{} = "fpconv" +convOpFun FPToUI{} = "fptoui" +convOpFun FPToSI{} = "fptosi" +convOpFun UIToFP{} = "uitofp" +convOpFun SIToFP{} = "sitofp" +convOpFun IToB{} = "itob" +convOpFun BToI{} = "btoi" + +taggedI :: String -> IntType -> Doc +taggedI s Int8 = text $ s ++ "8" +taggedI s Int16 = text $ s ++ "16" +taggedI s Int32 = text $ s ++ "32" +taggedI s Int64 = text $ s ++ "64" + +taggedF :: String -> FloatType -> Doc +taggedF s Float32 = text $ s ++ "32" +taggedF s Float64 = text $ s ++ "64" + +convOp :: (Pretty from, Pretty to) => String -> from -> to -> Doc +convOp s from to = text s <> text "_" <> ppr from <> text "_" <> ppr to + +-- | True if signed. Only makes a difference for integer types. +prettySigned :: Bool -> PrimType -> String +prettySigned True (IntType it) = 'u' : drop 1 (pretty it) +prettySigned _ t = pretty t diff --git a/src/Futhark/Representation/Ranges.hs b/src/Futhark/Representation/Ranges.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/Ranges.hs @@ -0,0 +1,189 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | A representation where all bindings are annotated with range +-- information. +module Futhark.Representation.Ranges + ( -- * The Lore definition + Ranges + , module Futhark.Representation.AST.Attributes.Ranges + -- * Module re-exports + , module Futhark.Representation.AST.Attributes + , module Futhark.Representation.AST.Traversals + , module Futhark.Representation.AST.Pretty + , module Futhark.Representation.AST.Syntax + -- * Adding ranges + , addRangesToPattern + , mkRangedLetStm + , mkRangedBody + , mkPatternRanges + , mkBodyRanges + -- * Removing ranges + , removeProgRanges + , removeFunDefRanges + , removeExpRanges + , removeBodyRanges + , removeStmRanges + , removeLambdaRanges + , removePatternRanges + ) +where + +import Control.Monad.Identity +import Control.Monad.Reader +import qualified Data.Set as S +import Data.Monoid ((<>)) +import Data.Foldable + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Attributes.Ranges +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Pretty +import Futhark.Analysis.Rephrase +import qualified Futhark.Util.Pretty as PP + +-- | The lore for the basic representation. +data Ranges lore + +instance (Annotations lore, CanBeRanged (Op lore)) => + Annotations (Ranges lore) where + type LetAttr (Ranges lore) = (Range, LetAttr lore) + type ExpAttr (Ranges lore) = ExpAttr lore + type BodyAttr (Ranges lore) = ([Range], BodyAttr lore) + type FParamAttr (Ranges lore) = FParamAttr lore + type LParamAttr (Ranges lore) = LParamAttr lore + type RetType (Ranges lore) = RetType lore + type BranchType (Ranges lore) = BranchType lore + type Op (Ranges lore) = OpWithRanges (Op lore) + +withoutRanges :: (HasScope (Ranges lore) m, Monad m) => + ReaderT (Scope lore) m a -> + m a +withoutRanges m = do + scope <- asksScope $ fmap unRange + runReaderT m scope + where unRange :: NameInfo (Ranges lore) -> NameInfo lore + unRange (LetInfo (_, x)) = LetInfo x + unRange (FParamInfo x) = FParamInfo x + unRange (LParamInfo x) = LParamInfo x + unRange (IndexInfo x) = IndexInfo x + +instance (Attributes lore, CanBeRanged (Op lore)) => + Attributes (Ranges lore) where + expTypesFromPattern = + withoutRanges . expTypesFromPattern . removePatternRanges + +instance RangeOf (Range, attr) where + rangeOf = fst + +instance RangesOf ([Range], attr) where + rangesOf = fst + +instance PrettyAnnot (PatElemT attr) => + PrettyAnnot (PatElemT (Range, attr)) where + + ppAnnot patelem = + range_annot <> inner_annot + where range_annot = + case fst . patElemAttr $ patelem of + (Nothing, Nothing) -> Nothing + range -> + Just $ PP.oneLine $ + PP.text "-- " <> PP.ppr (patElemName patelem) <> PP.text " range: " <> + PP.ppr range + inner_annot = ppAnnot $ fmap snd patelem + + +instance (PrettyLore lore, CanBeRanged (Op lore)) => PrettyLore (Ranges lore) where + ppExpLore attr = ppExpLore attr . removeExpRanges + +removeRanges :: CanBeRanged (Op lore) => Rephraser Identity (Ranges lore) lore +removeRanges = Rephraser { rephraseExpLore = return + , rephraseLetBoundLore = return . snd + , rephraseBodyLore = return . snd + , rephraseFParamLore = return + , rephraseLParamLore = return + , rephraseRetType = return + , rephraseBranchType = return + , rephraseOp = return . removeOpRanges + } + +removeProgRanges :: CanBeRanged (Op lore) => + Prog (Ranges lore) -> Prog lore +removeProgRanges = runIdentity . rephraseProg removeRanges + +removeFunDefRanges :: CanBeRanged (Op lore) => + FunDef (Ranges lore) -> FunDef lore +removeFunDefRanges = runIdentity . rephraseFunDef removeRanges + +removeExpRanges :: CanBeRanged (Op lore) => + Exp (Ranges lore) -> Exp lore +removeExpRanges = runIdentity . rephraseExp removeRanges + +removeBodyRanges :: CanBeRanged (Op lore) => + Body (Ranges lore) -> Body lore +removeBodyRanges = runIdentity . rephraseBody removeRanges + +removeStmRanges :: CanBeRanged (Op lore) => + Stm (Ranges lore) -> Stm lore +removeStmRanges = runIdentity . rephraseStm removeRanges + +removeLambdaRanges :: CanBeRanged (Op lore) => + Lambda (Ranges lore) -> Lambda lore +removeLambdaRanges = runIdentity . rephraseLambda removeRanges + +removePatternRanges :: PatternT (Range, a) + -> PatternT a +removePatternRanges = runIdentity . rephrasePattern (return . snd) + +addRangesToPattern :: (Attributes lore, CanBeRanged (Op lore)) => + Pattern lore -> Exp (Ranges lore) + -> Pattern (Ranges lore) +addRangesToPattern pat e = + uncurry Pattern $ mkPatternRanges pat e + +mkRangedBody :: BodyAttr lore -> Stms (Ranges lore) -> Result + -> Body (Ranges lore) +mkRangedBody innerlore bnds res = + Body (mkBodyRanges bnds res, innerlore) bnds res + +mkPatternRanges :: (Attributes lore, CanBeRanged (Op lore)) => + Pattern lore + -> Exp (Ranges lore) + -> ([PatElemT (Range, LetAttr lore)], + [PatElemT (Range, LetAttr lore)]) +mkPatternRanges pat e = + (map (`addRanges` unknownRange) $ patternContextElements pat, + zipWith addRanges (patternValueElements pat) ranges) + where addRanges patElem range = + let innerlore = patElemAttr patElem + in patElem `setPatElemLore` (range, innerlore) + ranges = expRanges e + +mkBodyRanges :: Stms lore -> Result -> [Range] +mkBodyRanges bnds = map $ removeUnknownBounds . rangeOf + where boundInBnds = + fold $ fmap (S.fromList . patternNames . stmPattern) bnds + removeUnknownBounds (lower,upper) = + (removeUnknownBound lower, + removeUnknownBound upper) + removeUnknownBound (Just bound) + | freeIn bound `intersects` boundInBnds = Nothing + | otherwise = Just bound + removeUnknownBound Nothing = + Nothing + +intersects :: Ord a => S.Set a -> S.Set a -> Bool +intersects a b = not $ S.null $ a `S.intersection` b + +mkRangedLetStm :: (Attributes lore, CanBeRanged (Op lore)) => + Pattern lore + -> ExpAttr lore + -> Exp (Ranges lore) + -> Stm (Ranges lore) +mkRangedLetStm pat explore e = + Let (addRangesToPattern pat e) (StmAux mempty explore) e diff --git a/src/Futhark/Representation/SOACS.hs b/src/Futhark/Representation/SOACS.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/SOACS.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +-- | A simple representation with SOACs and nested parallelism. +module Futhark.Representation.SOACS + ( -- * The Lore definition + SOACS + -- * Syntax types + , Prog + , Body + , Stm + , Pattern + , BasicOp + , Exp + , Lambda + , FunDef + , FParam + , LParam + , RetType + , PatElem + -- * Module re-exports + , module Futhark.Representation.AST.Attributes + , module Futhark.Representation.AST.Traversals + , module Futhark.Representation.AST.Pretty + , module Futhark.Representation.AST.Syntax + , module Futhark.Representation.SOACS.SOAC + , AST.LambdaT(Lambda) + , AST.BodyT(Body) + , AST.PatternT(Pattern) + , AST.PatElemT(PatElem) + , AST.ProgT(Prog) + , AST.ExpT(BasicOp) + , AST.FunDefT(FunDef) + , AST.ParamT(Param) + ) +where + +import Control.Monad + +import qualified Futhark.Representation.AST.Syntax as AST +import Futhark.Representation.AST.Syntax + hiding (Prog, BasicOp, Exp, Body, Stm, + Pattern, Lambda, FunDef, FParam, LParam, RetType, PatElem) +import Futhark.Representation.SOACS.SOAC +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Pretty +import Futhark.Binder +import Futhark.Construct +import qualified Futhark.TypeCheck as TypeCheck + +-- This module could be written much nicer if Haskell had functors +-- like Standard ML. Instead, we have to abuse the namespace/module +-- system. + +-- | The lore for the basic representation. +data SOACS + +instance Annotations SOACS where + type Op SOACS = SOAC SOACS + +instance Attributes SOACS where + expTypesFromPattern = return . expExtTypesFromPattern + +type Prog = AST.Prog SOACS +type BasicOp = AST.BasicOp SOACS +type Exp = AST.Exp SOACS +type Body = AST.Body SOACS +type Stm = AST.Stm SOACS +type Pattern = AST.Pattern SOACS +type Lambda = AST.Lambda SOACS +type FunDef = AST.FunDefT SOACS +type FParam = AST.FParam SOACS +type LParam = AST.LParam SOACS +type RetType = AST.RetType SOACS +type PatElem = AST.PatElem SOACS + +instance TypeCheck.Checkable SOACS where + checkExpLore = return + checkBodyLore = return + checkFParamLore _ = TypeCheck.checkType + checkLParamLore _ = TypeCheck.checkType + checkLetBoundLore _ = TypeCheck.checkType + checkRetType = mapM_ TypeCheck.checkExtType . retTypeValues + checkOp = typeCheckSOAC + matchPattern pat = TypeCheck.matchExtPattern pat <=< expExtType + primFParam name t = + return $ AST.Param name (AST.Prim t) + primLParam name t = + return $ AST.Param name (AST.Prim t) + matchReturnType = TypeCheck.matchExtReturnType . map fromDecl + matchBranchType = TypeCheck.matchExtBranchType + +instance Bindable SOACS where + mkBody = AST.Body () + mkExpPat ctx val _ = basicPattern ctx val + mkExpAttr _ _ = () + mkLetNames = simpleMkLetNames + +instance BinderOps SOACS where + mkExpAttrB = bindableMkExpAttrB + mkBodyB = bindableMkBodyB + mkLetNamesB = bindableMkLetNamesB + +instance PrettyLore SOACS where diff --git a/src/Futhark/Representation/SOACS/SOAC.hs b/src/Futhark/Representation/SOACS/SOAC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/SOACS/SOAC.hs @@ -0,0 +1,740 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE ConstraintKinds #-} +module Futhark.Representation.SOACS.SOAC + ( SOAC(..) + , StreamForm(..) + , ScremaForm(..) + , GenReduceOp(..) + , Scan + , Reduce + + , typeCheckSOAC + + -- * Utility + , getStreamOrder + , getStreamAccums + , scremaType + , soacType + + , mkIdentityLambda + , isIdentityLambda + , composeLambda + , nilFn + , scanomapSOAC + , redomapSOAC + , scanSOAC + , reduceSOAC + , mapSOAC + , isScanomapSOAC + , isRedomapSOAC + , isScanSOAC + , isReduceSOAC + , isMapSOAC + + , ppScrema + , ppGenReduce + + -- * Generic traversal + , SOACMapper(..) + , identitySOACMapper + , mapSOACM + ) + where + +import Control.Monad.Writer +import Control.Monad.Identity +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import Data.List + +import Futhark.Representation.AST +import qualified Futhark.Analysis.Alias as Alias +import qualified Futhark.Util.Pretty as PP +import Futhark.Util.Pretty (ppr, Doc, Pretty, parens, comma, (), commasep, text) +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Transform.Substitute +import Futhark.Transform.Rename +import Futhark.Optimise.Simplify.Lore +import Futhark.Representation.Ranges (Ranges, removeLambdaRanges) +import Futhark.Representation.AST.Attributes.Ranges +import Futhark.Representation.Aliases (Aliases, removeLambdaAliases) +import Futhark.Analysis.Usage +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.TypeCheck as TC +import Futhark.Analysis.Metrics +import qualified Futhark.Analysis.Range as Range +import Futhark.Construct +import Futhark.Util (maybeNth, chunks, splitAt3) + +data SOAC lore = + Stream SubExp (StreamForm lore) (LambdaT lore) [VName] + | Scatter SubExp (LambdaT lore) [VName] [(SubExp, Int, VName)] + -- Scatter + -- + -- + -- + -- is the length of each index array and value array, since they + -- all must be the same length for any fusion to make sense. If you have a + -- list of index-value array pairs of different sizes, you need to use + -- multiple writes instead. + -- + -- The lambda body returns the output in this manner: + -- + -- [index_0, index_1, ..., index_n, value_0, value_1, ..., value_n] + -- + -- This must be consistent along all Scatter-related optimisations. + | GenReduce SubExp [GenReduceOp lore] (LambdaT lore) [VName] + -- GenReduce + -- + -- The first SubExp is the length of the input arrays. The first + -- list describes the operations to perform. The 'LambdaT' is the + -- bucket function. Finally comes the input images. + | Screma SubExp (ScremaForm lore) [VName] + -- ^ A combination of scan, reduction, and map. The first + -- 'SubExp' is the size of the input arrays. The first + -- 'Lambda'/'SubExp' pair is for scan and its neutral elements. + -- The second is for the reduction. The final lambda is for the + -- map part, and finally comes the input arrays. + | CmpThreshold SubExp String + deriving (Eq, Ord, Show) + +data GenReduceOp lore = GenReduceOp { genReduceWidth :: SubExp + , genReduceDest :: [VName] + , genReduceNeutral :: [SubExp] + , genReduceOp :: LambdaT lore + } + deriving (Eq, Ord, Show) + +data StreamForm lore = + Parallel StreamOrd Commutativity (LambdaT lore) [SubExp] + | Sequential [SubExp] + deriving (Eq, Ord, Show) + +-- | The essential parts of a 'Screma' factored out (everything +-- except the input arrays). +data ScremaForm lore = ScremaForm + (Scan lore) + (Reduce lore) + (LambdaT lore) + deriving (Eq, Ord, Show) + +type Scan lore = (LambdaT lore, [SubExp]) +type Reduce lore = (Commutativity, LambdaT lore, [SubExp]) + +scremaType :: SubExp -> ScremaForm lore -> [Type] +scremaType w (ScremaForm (scan_lam, _scan_nes) (_, red_lam, _red_nes) map_lam) = + map (`arrayOfRow` w) scan_tps ++ red_tps ++ map (`arrayOfRow` w) map_tps + where scan_tps = lambdaReturnType scan_lam + red_tps = lambdaReturnType red_lam + map_tps = drop (length scan_tps + length red_tps) $ lambdaReturnType map_lam + +-- | Construct a lambda that takes parameters of the given types and +-- simply returns them unchanged. +mkIdentityLambda :: (Bindable lore, MonadFreshNames m) => + [Type] -> m (Lambda lore) +mkIdentityLambda ts = do + params <- mapM (newParam "x") ts + return Lambda { lambdaParams = params + , lambdaBody = mkBody mempty $ map (Var . paramName) params + , lambdaReturnType = ts } + +-- | Is the given lambda an identity lambda? +isIdentityLambda :: Lambda lore -> Bool +isIdentityLambda lam = bodyResult (lambdaBody lam) == + map (Var . paramName) (lambdaParams lam) + +composeLambda :: (Bindable lore, BinderOps lore, MonadFreshNames m, + HasScope somelore m, SameScope somelore lore) => + Lambda lore + -> Lambda lore + -> Lambda lore + -> m (Lambda lore) +composeLambda scan_fun red_fun map_fun = do + body <- runBodyBinder $ inScopeOf scan_fun $ inScopeOf red_fun $ inScopeOf map_fun $ do + mapM_ addStm $ bodyStms $ lambdaBody map_fun + let (scan_res, red_res, map_res) = splitAt3 n m $ bodyResult $ lambdaBody map_fun + + forM_ (zip scan_y_params scan_res) $ \(p,se) -> + letBindNames_ [paramName p] $ BasicOp $ SubExp se + forM_ (zip red_y_params red_res) $ \(p,se) -> + letBindNames_ [paramName p] $ BasicOp $ SubExp se + mapM_ addStm $ bodyStms $ lambdaBody scan_fun + mapM_ addStm $ bodyStms $ lambdaBody red_fun + + resultBodyM $ + bodyResult (lambdaBody scan_fun) ++ + bodyResult (lambdaBody red_fun) ++ + map_res + + return Lambda { lambdaParams = scan_x_params ++ red_x_params ++ lambdaParams map_fun + , lambdaBody = body + , lambdaReturnType = lambdaReturnType map_fun } + where n = length $ lambdaReturnType scan_fun + m = length $ lambdaReturnType red_fun + (scan_x_params, scan_y_params) = splitAt n $ lambdaParams scan_fun + (red_x_params, red_y_params) = splitAt m $ lambdaParams red_fun + +-- | A lambda with no parameters that returns no values. +nilFn :: Bindable lore => LambdaT lore +nilFn = Lambda mempty (mkBody mempty mempty) mempty + +isNilFn :: LambdaT lore -> Bool +isNilFn (Lambda ps body ts) = + null ps && null ts && + null (bodyStms body) && null (bodyResult body) + +scanomapSOAC :: Bindable lore => + Lambda lore -> [SubExp] -> Lambda lore -> ScremaForm lore +scanomapSOAC lam nes = ScremaForm (lam, nes) (mempty, nilFn, mempty) + +redomapSOAC :: Bindable lore => + Commutativity -> Lambda lore -> [SubExp] -> Lambda lore -> ScremaForm lore +redomapSOAC comm lam nes = ScremaForm (nilFn, mempty) (comm, lam, nes) + +scanSOAC :: (Bindable lore, MonadFreshNames m) => + Lambda lore -> [SubExp] -> m (ScremaForm lore) +scanSOAC lam nes = scanomapSOAC lam nes <$> mkIdentityLambda (lambdaReturnType lam) + +reduceSOAC :: (Bindable lore, MonadFreshNames m) => + Commutativity -> Lambda lore -> [SubExp] -> m (ScremaForm lore) +reduceSOAC comm lam nes = redomapSOAC comm lam nes <$> mkIdentityLambda (lambdaReturnType lam) + +mapSOAC :: Bindable lore => Lambda lore -> ScremaForm lore +mapSOAC = ScremaForm (nilFn, mempty) (mempty, nilFn, mempty) + +isScanomapSOAC :: ScremaForm lore -> Maybe (Lambda lore, [SubExp], Lambda lore) +isScanomapSOAC (ScremaForm (scan_lam, scan_nes) (_, _, red_nes) map_lam) = do + guard $ null red_nes + guard $ not $ null scan_nes + return (scan_lam, scan_nes, map_lam) + +isScanSOAC :: ScremaForm lore -> Maybe (Lambda lore, [SubExp]) +isScanSOAC form = do (scan_lam, scan_nes, map_lam) <- isScanomapSOAC form + guard $ isIdentityLambda map_lam + return (scan_lam, scan_nes) + +isRedomapSOAC :: ScremaForm lore -> Maybe (Commutativity, Lambda lore, [SubExp], Lambda lore) +isRedomapSOAC (ScremaForm (_, scan_nes) (comm, red_lam, red_nes) map_lam) = do + guard $ null scan_nes + guard $ not $ null red_nes + return (comm, red_lam, red_nes, map_lam) + +isReduceSOAC :: ScremaForm lore -> Maybe (Commutativity, Lambda lore, [SubExp]) +isReduceSOAC form = do (comm, red_lam, red_nes, map_lam) <- isRedomapSOAC form + guard $ isIdentityLambda map_lam + return (comm, red_lam, red_nes) + +isMapSOAC :: ScremaForm lore -> Maybe (Lambda lore) +isMapSOAC (ScremaForm (_, scan_nes) (_, _, red_nes) map_lam) = do + guard $ null scan_nes + guard $ null red_nes + return map_lam + +-- | Like 'Mapper', but just for 'SOAC's. +data SOACMapper flore tlore m = SOACMapper { + mapOnSOACSubExp :: SubExp -> m SubExp + , mapOnSOACLambda :: Lambda flore -> m (Lambda tlore) + , mapOnSOACVName :: VName -> m VName + } + +-- | A mapper that simply returns the SOAC verbatim. +identitySOACMapper :: Monad m => SOACMapper lore lore m +identitySOACMapper = SOACMapper { mapOnSOACSubExp = return + , mapOnSOACLambda = return + , mapOnSOACVName = return + } + +-- | Map a monadic action across the immediate children of a +-- SOAC. The mapping does not descend recursively into subexpressions +-- and is done left-to-right. +mapSOACM :: (Applicative m, Monad m) => + SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore) +mapSOACM tv (Stream size form lam arrs) = + Stream <$> mapOnSOACSubExp tv size <*> + mapOnStreamForm form <*> mapOnSOACLambda tv lam <*> + mapM (mapOnSOACVName tv) arrs + where mapOnStreamForm (Parallel o comm lam0 acc) = + Parallel <$> pure o <*> pure comm <*> + mapOnSOACLambda tv lam0 <*> + mapM (mapOnSOACSubExp tv) acc + mapOnStreamForm (Sequential acc) = + Sequential <$> mapM (mapOnSOACSubExp tv) acc +mapSOACM tv (Scatter len lam ivs as) = + Scatter + <$> mapOnSOACSubExp tv len + <*> mapOnSOACLambda tv lam + <*> mapM (mapOnSOACVName tv) ivs + <*> mapM (\(aw,an,a) -> (,,) <$> mapOnSOACSubExp tv aw <*> + pure an <*> mapOnSOACVName tv a) as +mapSOACM tv (GenReduce len ops bucket_fun imgs) = + GenReduce + <$> mapOnSOACSubExp tv len + <*> mapM (\(GenReduceOp e arrs nes op) -> + GenReduceOp <$> mapOnSOACSubExp tv e + <*> mapM (mapOnSOACVName tv) arrs + <*> mapM (mapOnSOACSubExp tv) nes + <*> mapOnSOACLambda tv op) ops + <*> mapOnSOACLambda tv bucket_fun + <*> mapM (mapOnSOACVName tv) imgs +mapSOACM tv (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) = + Screma <$> mapOnSOACSubExp tv w <*> + (ScremaForm <$> + ((,) <$> mapOnSOACLambda tv scan_lam <*> mapM (mapOnSOACSubExp tv) scan_nes) <*> + ((,,) comm <$> mapOnSOACLambda tv red_lam <*> mapM (mapOnSOACSubExp tv) red_nes) <*> + mapOnSOACLambda tv map_lam) + <*> mapM (mapOnSOACVName tv) arrs +mapSOACM tv (CmpThreshold what s) = CmpThreshold <$> mapOnSOACSubExp tv what <*> pure s + +instance Attributes lore => FreeIn (SOAC lore) where + freeIn = execWriter . mapSOACM free + where walk f x = tell (f x) >> return x + free = SOACMapper { mapOnSOACSubExp = walk freeIn + , mapOnSOACLambda = walk freeInLambda + , mapOnSOACVName = walk freeIn + } + +instance Attributes lore => Substitute (SOAC lore) where + substituteNames subst = + runIdentity . mapSOACM substitute + where substitute = + SOACMapper { mapOnSOACSubExp = return . substituteNames subst + , mapOnSOACLambda = return . substituteNames subst + , mapOnSOACVName = return . substituteNames subst + } + +instance Attributes lore => Rename (SOAC lore) where + rename = mapSOACM renamer + where renamer = SOACMapper rename rename rename + +soacType :: SOAC lore -> [Type] +soacType (Stream outersize form lam _) = + map (substNamesInType substs) rtp + where nms = map paramName $ take (1 + length accs) params + substs = M.fromList $ zip nms (outersize:accs) + Lambda params _ rtp = lam + accs = case form of + Parallel _ _ _ acc -> acc + Sequential acc -> acc +soacType (Scatter _w lam _ivs as) = + zipWith arrayOfRow val_ts ws + where val_ts = concatMap (take 1) $ chunks ns $ + drop (sum ns) $ lambdaReturnType lam + (ws, ns, _) = unzip3 as +soacType (GenReduce _len ops _bucket_fun _imgs) = do + op <- ops + map (`arrayOfRow` genReduceWidth op) (lambdaReturnType $ genReduceOp op) +soacType (Screma w form _arrs) = + scremaType w form +soacType CmpThreshold{} = [Prim Bool] + +instance TypedOp (SOAC lore) where + opType = pure . staticShapes . soacType + +instance (Attributes lore, Aliased lore) => AliasedOp (SOAC lore) where + opAliases = map (const mempty) . soacType + + -- Only map functions can consume anything. The operands to scan + -- and reduce functions are always considered "fresh". + consumedInOp (Screma _ (ScremaForm _ _ map_lam) arrs) = + S.map consumedArray $ consumedByLambda map_lam + where consumedArray v = fromMaybe v $ lookup v params_to_arrs + params_to_arrs = zip (map paramName $ lambdaParams map_lam) arrs + consumedInOp (Stream _ form lam arrs) = + S.fromList $ subExpVars $ + case form of Sequential accs -> + map (consumedArray accs) $ S.toList $ consumedByLambda lam + Parallel _ _ _ accs -> + map (consumedArray accs) $ S.toList $ consumedByLambda lam + where consumedArray accs v = fromMaybe (Var v) $ lookup v $ paramsToInput accs + -- Drop the chunk parameter, which cannot alias anything. + paramsToInput accs = zip + (map paramName $ drop 1 $ lambdaParams lam) + (accs++map Var arrs) + consumedInOp (Scatter _ _ _ as) = + S.fromList $ map (\(_, _, a) -> a) as + consumedInOp (GenReduce _ ops _ _) = + S.fromList $ concatMap genReduceDest ops + consumedInOp CmpThreshold{} = mempty + +mapGenReduceOp :: (LambdaT flore -> LambdaT tlore) + -> GenReduceOp flore -> GenReduceOp tlore +mapGenReduceOp f (GenReduceOp w dests nes lam) = + GenReduceOp w dests nes $ f lam + +instance (Attributes lore, + Attributes (Aliases lore), + CanBeAliased (Op lore)) => CanBeAliased (SOAC lore) where + type OpWithAliases (SOAC lore) = SOAC (Aliases lore) + + addOpAliases (Stream size form lam arr) = + Stream size (analyseStreamForm form) + (Alias.analyseLambda lam) arr + where analyseStreamForm (Parallel o comm lam0 acc) = + Parallel o comm (Alias.analyseLambda lam0) acc + analyseStreamForm (Sequential acc) = Sequential acc + addOpAliases (Scatter len lam ivs as) = + Scatter len (Alias.analyseLambda lam) ivs as + addOpAliases (GenReduce len ops bucket_fun imgs) = + GenReduce len (map (mapGenReduceOp Alias.analyseLambda) ops) + (Alias.analyseLambda bucket_fun) imgs + addOpAliases (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) = + Screma w (ScremaForm + (Alias.analyseLambda scan_lam, scan_nes) + (comm, Alias.analyseLambda red_lam, red_nes) + (Alias.analyseLambda map_lam)) + arrs + addOpAliases (CmpThreshold what s) = CmpThreshold what s + + removeOpAliases = runIdentity . mapSOACM remove + where remove = SOACMapper return (return . removeLambdaAliases) return + +instance Attributes lore => IsOp (SOAC lore) where + safeOp CmpThreshold{} = True + safeOp _ = False + cheapOp _ = True + +substNamesInType :: M.Map VName SubExp -> Type -> Type +substNamesInType _ tp@(Prim _) = tp +substNamesInType subs (Mem se space) = + Mem (substNamesInSubExp subs se) space +substNamesInType subs (Array btp shp u) = + let shp' = Shape $ map (substNamesInSubExp subs) (shapeDims shp) + in Array btp shp' u + +substNamesInSubExp :: M.Map VName SubExp -> SubExp -> SubExp +substNamesInSubExp _ e@(Constant _) = e +substNamesInSubExp subs (Var idd) = + M.findWithDefault (Var idd) idd subs + +instance (Ranged inner) => RangedOp (SOAC inner) where + opRanges op = replicate (length $ soacType op) unknownRange + +instance (Attributes lore, CanBeRanged (Op lore)) => CanBeRanged (SOAC lore) where + type OpWithRanges (SOAC lore) = SOAC (Ranges lore) + + removeOpRanges = runIdentity . mapSOACM remove + where remove = SOACMapper return (return . removeLambdaRanges) return + addOpRanges (Stream w form lam arr) = + Stream w + (Range.runRangeM $ analyseStreamForm form) + (Range.runRangeM $ Range.analyseLambda lam) + arr + where analyseStreamForm (Sequential acc) = + return $ Sequential acc + analyseStreamForm (Parallel o comm lam0 acc) = do + lam0' <- Range.analyseLambda lam0 + return $ Parallel o comm lam0' acc + addOpRanges (Scatter len lam ivs as) = + Scatter len (Range.runRangeM $ Range.analyseLambda lam) ivs as + addOpRanges (GenReduce len ops bucket_fun imgs) = + GenReduce len (map (mapGenReduceOp $ Range.runRangeM . Range.analyseLambda) ops) + (Range.runRangeM $ Range.analyseLambda bucket_fun) imgs + addOpRanges (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) = + Screma w (ScremaForm + (Range.runRangeM $ Range.analyseLambda scan_lam, scan_nes) + (comm, Range.runRangeM $ Range.analyseLambda red_lam, red_nes) + (Range.runRangeM $ Range.analyseLambda map_lam)) + arrs + addOpRanges (CmpThreshold what s) = CmpThreshold what s + +instance (Attributes lore, CanBeWise (Op lore)) => CanBeWise (SOAC lore) where + type OpWithWisdom (SOAC lore) = SOAC (Wise lore) + + removeOpWisdom = runIdentity . mapSOACM remove + where remove = SOACMapper return (return . removeLambdaWisdom) return + +instance Annotations lore => ST.IndexOp (SOAC lore) 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 + arr_indexes' = foldl expandPrimExpTable arr_indexes $ bodyStms $ lambdaBody lam + case se of + Var v -> M.lookup v arr_indexes' + _ -> Nothing + where lambdaAndSubExp (Screma _ (ScremaForm (_, scan_nes) (_, _, red_nes) map_lam) arrs) = + nthMapOut (length scan_nes + length red_nes) map_lam arrs + lambdaAndSubExp _ = + Nothing + + nthMapOut num_accs lam arrs = do + se <- maybeNth (num_accs+k) $ bodyResult $ lambdaBody lam + return (lam, se, drop num_accs $ lambdaParams lam, arrs) + + arrIndex p arr = do + (pe,cs) <- ST.index' arr [i] vtable + return (paramName p, (pe,cs)) + + expandPrimExpTable table stm + | [v] <- patternNames $ stmPattern stm, + Just (pe,cs) <- + runWriterT $ primExpFromExp (asPrimExp table) $ stmExp stm, + all (`ST.elem` vtable) (unCertificates $ stmCerts stm) = + M.insert v (pe, stmCerts stm <> cs) table + | otherwise = + table + + asPrimExp table v + | Just (e,cs) <- M.lookup v table = tell cs >> return e + | Just (Prim pt) <- ST.lookupType v vtable = + return $ LeafExp v pt + | otherwise = lift Nothing + indexOp _ _ _ _ = Nothing + +instance Aliased lore => UsageInOp (SOAC lore) where + usageInOp (Screma _ (ScremaForm _ _ f) arrs) = usageInLambda f arrs + usageInOp _ = mempty + +typeCheckSOAC :: TC.Checkable lore => SOAC (Aliases lore) -> TC.TypeM lore () +typeCheckSOAC (CmpThreshold what _) = TC.require [Prim int32] what +typeCheckSOAC (Stream size form lam arrexps) = do + let accexps = getStreamAccums form + TC.require [Prim int32] size + accargs <- mapM TC.checkArg accexps + arrargs <- mapM lookupType arrexps + _ <- TC.checkSOACArrayArgs size arrexps + let chunk = head $ lambdaParams lam + let asArg t = (t, mempty) + inttp = Prim int32 + lamarrs'= map (`setOuterSize` Var (paramName chunk)) arrargs + let acc_len= length accexps + let lamrtp = take acc_len $ lambdaReturnType lam + unless (map TC.argType accargs == lamrtp) $ + TC.bad $ TC.TypeError "Stream with inconsistent accumulator type in lambda." + -- check reduce's lambda, if any + _ <- case form of + Parallel _ _ lam0 _ -> do + let acct = map TC.argType accargs + outerRetType = lambdaReturnType lam0 + TC.checkLambda lam0 $ map TC.noArgAliases $ accargs ++ accargs + unless (acct == outerRetType) $ + TC.bad $ TC.TypeError $ + "Initial value is of type " ++ prettyTuple acct ++ + ", but stream's reduce lambda returns type " ++ prettyTuple outerRetType ++ "." + _ -> return () + -- just get the dflow of lambda on the fakearg, which does not alias + -- arr, so we can later check that aliases of arr are not used inside lam. + let fake_lamarrs' = map asArg lamarrs' + TC.checkLambda lam $ asArg inttp : accargs ++ fake_lamarrs' + +typeCheckSOAC (Scatter w lam ivs as) = do + -- Requirements: + -- + -- 0. @lambdaReturnType@ of @lam@ must be a list + -- [index types..., value types]. + -- + -- 1. The number of index types must be equal to the number of value types + -- and the number of writes to arrays in @as@. + -- + -- 2. Each index type must have the type i32. + -- + -- 3. Each array in @as@ and the value types must have the same type + -- + -- 4. Each array in @as@ is consumed. This is not really a check, but more + -- of a requirement, so that e.g. the source is not hoisted out of a + -- loop, which will mean it cannot be consumed. + -- + -- 5. Each of ivs must be an array matching a corresponding lambda + -- parameters. + -- + -- Code: + + -- First check the input size. + TC.require [Prim int32] w + + -- 0. + let (_as_ws, as_ns, _as_vs) = unzip3 as + rts = lambdaReturnType lam + rtsLen = length rts `div` 2 + rtsI = take rtsLen rts + rtsV = drop rtsLen rts + + -- 1. + unless (rtsLen == sum as_ns) + $ TC.bad $ TC.TypeError "Scatter: Uneven number of index types, value types, and arrays outputs." + + -- 2. + forM_ rtsI $ \rtI -> unless (Prim int32 == rtI) $ + TC.bad $ TC.TypeError "Scatter: Index return type must be i32." + + forM_ (zip (chunks as_ns rtsV) as) $ \(rtVs, (aw, _, a)) -> do + -- All lengths must have type i32. + TC.require [Prim int32] aw + + -- 3. + forM_ rtVs $ \rtV -> TC.requireI [rtV `arrayOfRow` aw] a + + -- 4. + TC.consume =<< TC.lookupAliases a + + -- 5. + arrargs <- TC.checkSOACArrayArgs w ivs + TC.checkLambda lam arrargs + +typeCheckSOAC (GenReduce len ops bucket_fun imgs) = do + TC.require [Prim int32] len + + -- Check the operators. + forM_ ops $ \(GenReduceOp dest_w dests nes op) -> do + nes' <- mapM TC.checkArg nes + TC.require [Prim int32] dest_w + + -- Operator type must match the type of neutral elements. + TC.checkLambda op $ map TC.noArgAliases $ nes' ++ nes' + let nes_t = map TC.argType nes' + unless (nes_t == lambdaReturnType op) $ + TC.bad $ TC.TypeError $ "Operator has return type " ++ + prettyTuple (lambdaReturnType op) ++ " but neutral element has type " ++ + prettyTuple nes_t + + -- Arrays must have proper type. + forM_ (zip nes_t dests) $ \(t, dest) -> do + TC.requireI [t `arrayOfRow` dest_w] dest + TC.consume =<< TC.lookupAliases dest + + -- Types of input arrays must equal parameter types for bucket function. + img' <- TC.checkSOACArrayArgs len imgs + TC.checkLambda bucket_fun img' + + -- Return type of bucket function must be an index for each + -- operation followed by the values to write. + nes_ts <- concat <$> mapM (mapM subExpType . genReduceNeutral) ops + let bucket_ret_t = replicate (length ops) (Prim int32) ++ nes_ts + unless (bucket_ret_t == lambdaReturnType bucket_fun) $ + TC.bad $ TC.TypeError $ "Bucket function has return type " ++ + prettyTuple (lambdaReturnType bucket_fun) ++ " but should have type " ++ + prettyTuple bucket_ret_t + +typeCheckSOAC (Screma w (ScremaForm (scan_lam, scan_nes) (_, red_lam, red_nes) map_lam) arrs) = do + TC.require [Prim int32] w + arrs' <- TC.checkSOACArrayArgs w arrs + scan_nes' <- mapM TC.checkArg scan_nes + red_nes' <- mapM TC.checkArg red_nes + TC.checkLambda map_lam $ map TC.noArgAliases arrs' + TC.checkLambda scan_lam $ map TC.noArgAliases $ scan_nes' ++ scan_nes' + TC.checkLambda red_lam $ map TC.noArgAliases $ red_nes' ++ red_nes' + let scan_t = map TC.argType scan_nes' + red_t = map TC.argType red_nes' + map_lam_ts = lambdaReturnType map_lam + + unless (scan_t == lambdaReturnType scan_lam) $ + TC.bad $ TC.TypeError $ "Scan function returns type " ++ + prettyTuple (lambdaReturnType scan_lam) ++ " but neutral element has type " ++ + prettyTuple scan_t + + unless (red_t == lambdaReturnType red_lam) $ + TC.bad $ TC.TypeError $ "Reduce function returns type " ++ + prettyTuple (lambdaReturnType red_lam) ++ " but neutral element has type " ++ + prettyTuple red_t + + unless (take (length scan_nes + length red_nes) map_lam_ts == + map TC.argType (scan_nes'++ red_nes')) $ + TC.bad $ TC.TypeError $ "Map function return type " ++ prettyTuple map_lam_ts ++ + " wrong for given scan and reduction functions." + +-- | Get Stream's accumulators as a sub-expression list +getStreamAccums :: StreamForm lore -> [SubExp] +getStreamAccums (Parallel _ _ _ accs) = accs +getStreamAccums (Sequential accs) = accs + +getStreamOrder :: StreamForm lore -> StreamOrd +getStreamOrder (Parallel o _ _ _) = o +getStreamOrder (Sequential _) = InOrder + +instance OpMetrics (Op lore) => OpMetrics (SOAC lore) where + opMetrics (Stream _ _ lam _) = + inside "Stream" $ lambdaMetrics lam + opMetrics (Scatter _len lam _ivs _as) = + inside "Scatter" $ lambdaMetrics lam + opMetrics (GenReduce _len ops bucket_fun _imgs) = + inside "GenReduce" $ mapM_ (lambdaMetrics . genReduceOp) ops >> lambdaMetrics bucket_fun + opMetrics (Screma _ (ScremaForm (scan_lam, _) (_, red_lam, _) map_lam) _) = + inside "Screma" $ + lambdaMetrics scan_lam >> lambdaMetrics red_lam >> lambdaMetrics map_lam + opMetrics CmpThreshold{} = seen "CmpThreshold" + +instance PrettyLore lore => PP.Pretty (SOAC lore) where + ppr (Stream size form lam arrs) = + case form of + Parallel o comm lam0 acc -> + let ord_str = if o == Disorder then "Per" else "" + comm_str = case comm of Commutative -> "Comm" + Noncommutative -> "" + in text ("streamPar"++ord_str++comm_str) <> + parens (ppr size <> comma ppr lam0 comma ppr lam + commasep ( PP.braces (commasep $ map ppr acc) : map ppr arrs )) + Sequential acc -> + text "streamSeq" <> + parens (ppr size <> comma ppr lam <> comma + commasep ( PP.braces (commasep $ map ppr acc) : map ppr arrs )) + ppr (Scatter len lam ivs as) = + ppSOAC "scatter" len [lam] (Just (map Var ivs)) (map (\(_,n,a) -> (n,a)) as) + ppr (GenReduce len ops bucket_fun imgs) = + ppGenReduce len ops bucket_fun imgs + ppr (Screma w (ScremaForm (scan_lam, scan_nes) (_, red_lam, red_nes) map_lam) arrs) + | isNilFn scan_lam, null scan_nes, + isNilFn red_lam, null red_nes = + text "map" <> parens (ppr w <> comma + ppr map_lam <> comma + commasep (map ppr arrs)) + + | isNilFn scan_lam, null scan_nes = + text "redomap" <> parens (ppr w <> comma + ppr red_lam <> comma + commasep (map ppr red_nes) <> comma + ppr map_lam <> comma + commasep (map ppr arrs)) + + | isNilFn red_lam, null red_nes = + text "scanomap" <> parens (ppr w <> comma + ppr scan_lam <> comma + commasep (map ppr scan_nes) <> comma + ppr map_lam <> comma + commasep (map ppr arrs)) + + ppr (Screma w form arrs) = ppScrema w form arrs + ppr (CmpThreshold what s) = text "cmpThreshold(" <> ppr what <> comma PP.<+> text (show s) <> text ")" + +ppScrema :: (PrettyLore lore, Pretty inp) => + SubExp -> ScremaForm lore -> [inp] -> Doc +ppScrema w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs = + text s <> parens (ppr w <> comma + ppr scan_lam <> comma + PP.braces (commasep $ map ppr scan_nes) + ppr red_lam <> comma + PP.braces (commasep $ map ppr red_nes) + ppr map_lam <> comma + commasep (map ppr arrs)) + where s = case comm of Noncommutative -> "screma" + Commutative -> "scremaComm" + +ppGenReduce :: (PrettyLore lore, Pretty inp) => + SubExp -> [GenReduceOp lore] -> Lambda lore -> [inp] -> Doc +ppGenReduce len ops bucket_fun imgs = + text "gen_reduce" <> + parens (ppr len <> comma + PP.braces (mconcat $ intersperse (comma <> PP.line) $ map ppOp ops) <> comma + ppr bucket_fun <> comma + commasep (map ppr imgs)) + where ppOp (GenReduceOp w dests nes op) = + ppr w <> comma <> PP.braces (commasep $ map ppr dests) <> comma + PP.braces (commasep $ map ppr nes) <> comma ppr op + +ppSOAC :: (Pretty fn, Pretty v) => + String -> SubExp -> [fn] -> Maybe [SubExp] -> [v] -> Doc +ppSOAC name size funs es as = + text name <> parens (ppr size <> comma + ppList funs + commasep (es' ++ map ppr as)) + where es' = maybe [] ((:[]) . ppTuple') es + +ppList :: Pretty a => [a] -> Doc +ppList as = case map ppr as of + [] -> mempty + a':as' -> foldl () (a' <> comma) $ map (<> comma) as' diff --git a/src/Futhark/Representation/SOACS/Simplify.hs b/src/Futhark/Representation/SOACS/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Representation/SOACS/Simplify.hs @@ -0,0 +1,491 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +module Futhark.Representation.SOACS.Simplify + ( simplifySOACS + , simplifyLambda + , simplifyStms + + , simpleSOACS + ) +where + +import Control.Monad +import Data.Foldable +import Data.Either +import Data.List +import Data.Maybe +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Futhark.Representation.SOACS +import qualified Futhark.Representation.AST as AST +import Futhark.Representation.AST.Attributes.Aliases +import qualified Futhark.Optimise.Simplify.Engine as Engine +import qualified Futhark.Optimise.Simplify as Simplify +import Futhark.Optimise.Simplify.Rules +import Futhark.MonadFreshNames +import Futhark.Optimise.Simplify.Rule +import Futhark.Optimise.Simplify.ClosedForm +import Futhark.Optimise.Simplify.Lore +import Futhark.Tools +import Futhark.Pass +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Analysis.DataDependencies +import Futhark.Transform.Rename +import Futhark.Util + +simpleSOACS :: Simplify.SimpleOps SOACS +simpleSOACS = Simplify.bindableSimpleOps simplifySOAC + +simplifySOACS :: Prog -> PassM Prog +simplifySOACS = Simplify.simplifyProg simpleSOACS soacRules blockers + where blockers = Engine.noExtraHoistBlockers { Engine.getArraySizes = getShapeNames } + +-- | Getting the roots of what to hoist, for now only variable +-- names that represent shapes/sizes. +getShapeNames :: (LetAttr lore ~ (VarWisdom, Type)) => + AST.Stm lore -> Names +getShapeNames bnd = + let tps1 = map patElemType $ patternElements $ stmPattern bnd + tps2 = map (snd . patElemAttr) $ patternElements $ stmPattern bnd + in S.fromList $ subExpVars $ concatMap arrayDims (tps1 ++ tps2) + +simplifyLambda :: (HasScope SOACS m, MonadFreshNames m) => + Lambda -> [Maybe VName] -> m Lambda +simplifyLambda = + Simplify.simplifyLambda simpleSOACS soacRules Engine.noExtraHoistBlockers + +simplifyStms :: (HasScope SOACS m, MonadFreshNames m) => + Stms SOACS -> m (Stms SOACS) +simplifyStms = + Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers + +simplifySOAC :: Simplify.SimplifyOp SOACS +simplifySOAC (CmpThreshold what s) = do + what' <- Engine.simplify what + return (CmpThreshold what' s, mempty) +simplifySOAC (Stream outerdim form lam arr) = do + outerdim' <- Engine.simplify outerdim + (form', form_hoisted) <- simplifyStreamForm form + arr' <- mapM Engine.simplify arr + (lam', lam_hoisted) <- Engine.simplifyLambda lam (map Just arr) + return (Stream outerdim' form' lam' arr', form_hoisted <> lam_hoisted) + where simplifyStreamForm (Parallel o comm lam0 acc) = do + acc' <- mapM Engine.simplify acc + (lam0', hoisted) <- Engine.simplifyLambda lam0 $ + replicate (length $ lambdaParams lam0) Nothing + return (Parallel o comm lam0' acc', hoisted) + simplifyStreamForm (Sequential acc) = do + acc' <- mapM Engine.simplify acc + return (Sequential acc', mempty) + +simplifySOAC (Scatter len lam ivs as) = do + len' <- Engine.simplify len + (lam', hoisted) <- Engine.simplifyLambda lam $ map Just ivs + ivs' <- mapM Engine.simplify ivs + as' <- mapM Engine.simplify as + return (Scatter len' lam' ivs' as', hoisted) + +simplifySOAC (GenReduce w ops bfun imgs) = do + w' <- Engine.simplify w + (ops', hoisted) <- fmap unzip $ forM ops $ \(GenReduceOp dests_w dests nes op) -> do + dests_w' <- Engine.simplify dests_w + dests' <- Engine.simplify dests + nes' <- mapM Engine.simplify nes + (op', hoisted) <- Engine.simplifyLambda op $ replicate (length $ lambdaParams op) Nothing + return (GenReduceOp dests_w' dests' nes' op', hoisted) + imgs' <- mapM Engine.simplify imgs + (bfun', bfun_hoisted) <- Engine.simplifyLambda bfun $ map Just imgs + return (GenReduce w' ops' bfun' imgs', mconcat hoisted <> bfun_hoisted) + +simplifySOAC (Screma w (ScremaForm (scan_lam, scan_nes) (comm, red_lam, red_nes) map_lam) arrs) = do + (scan_lam', scan_lam_hoisted) <- + Engine.simplifyLambda scan_lam $ replicate (length scan_nes) Nothing + (red_lam', red_lam_hoisted) <- + Engine.simplifyLambda red_lam $ replicate (length red_nes) Nothing + (map_lam', map_lam_hoisted) <- Engine.simplifyLambda map_lam $ map Just arrs + (,) <$> (Screma <$> Engine.simplify w <*> + (ScremaForm <$> + ((,) scan_lam' <$> Engine.simplify scan_nes) <*> + ((,,) comm red_lam' <$> Engine.simplify red_nes) <*> + pure map_lam') <*> + Engine.simplify arrs) <*> + pure (scan_lam_hoisted <> red_lam_hoisted <> map_lam_hoisted) + +instance BinderOps (Wise SOACS) where + mkExpAttrB = bindableMkExpAttrB + mkBodyB = bindableMkBodyB + mkLetNamesB = bindableMkLetNamesB + +fixLambdaParams :: (MonadBinder m, Bindable (Lore m), BinderOps (Lore m)) => + AST.Lambda (Lore m) -> [Maybe SubExp] -> m (AST.Lambda (Lore m)) +fixLambdaParams lam fixes = do + body <- runBodyBinder $ localScope (scopeOfLParams $ lambdaParams lam) $ do + zipWithM_ maybeFix (lambdaParams lam) fixes' + return $ lambdaBody lam + return lam { lambdaBody = body + , lambdaParams = map fst $ filter (isNothing . snd) $ + zip (lambdaParams lam) fixes' } + where fixes' = fixes ++ repeat Nothing + maybeFix p (Just x) = letBindNames_ [paramName p] $ BasicOp $ SubExp x + maybeFix _ Nothing = return () + +removeLambdaResults :: [Bool] -> AST.Lambda lore -> AST.Lambda lore +removeLambdaResults keep lam = lam { lambdaBody = lam_body' + , lambdaReturnType = ret } + where keep' :: [a] -> [a] + keep' = map snd . filter fst . zip (keep ++ repeat True) + lam_body = lambdaBody lam + lam_body' = lam_body { bodyResult = keep' $ bodyResult lam_body } + ret = keep' $ lambdaReturnType lam + +soacRules :: RuleBook (Wise SOACS) +soacRules = standardRules <> ruleBook topDownRules bottomUpRules + +topDownRules :: [TopDownRule (Wise SOACS)] +topDownRules = [RuleOp removeReplicateMapping, + RuleOp removeReplicateWrite, + RuleOp removeUnusedSOACInput, + RuleOp simplifyClosedFormReduce, + RuleOp simplifyKnownIterationSOAC, + RuleOp fuseConcatScatter + ] + +bottomUpRules :: [BottomUpRule (Wise SOACS)] +bottomUpRules = [RuleOp removeDeadMapping, + RuleOp removeDeadReduction, + RuleOp removeDeadWrite, + RuleBasicOp removeUnnecessaryCopy, + RuleOp liftIdentityMapping, + RuleOp removeDuplicateMapOutput, + RuleOp mapOpToOp + ] + +liftIdentityMapping :: BottomUpRuleOp (Wise SOACS) +liftIdentityMapping (_, usages) pat _ (Screma w form arrs) + | Just fun <- isMapSOAC form = do + let inputMap = M.fromList $ zip (map paramName $ lambdaParams fun) arrs + free = freeInBody $ lambdaBody fun + rettype = lambdaReturnType fun + ses = bodyResult $ lambdaBody fun + + freeOrConst (Var v) = v `S.member` free + freeOrConst Constant{} = True + + checkInvariance (outId, Var v, _) (invariant, mapresult, rettype') + | Just inp <- M.lookup v inputMap = + let e | patElemName outId `UT.isConsumed` usages + || inp `UT.isConsumed` usages = + Copy inp + | otherwise = + SubExp $ Var inp + in ((Pattern [] [outId], BasicOp e) : invariant, + mapresult, + rettype') + checkInvariance (outId, e, t) (invariant, mapresult, rettype') + | freeOrConst e = ((Pattern [] [outId], BasicOp $ Replicate (Shape [w]) e) : invariant, + mapresult, + rettype') + | otherwise = (invariant, + (outId, e) : mapresult, + t : rettype') + + case foldr checkInvariance ([], [], []) $ + zip3 (patternElements pat) ses rettype of + ([], _, _) -> cannotSimplify + (invariant, mapresult, rettype') -> do + let (pat', ses') = unzip mapresult + fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } + , lambdaReturnType = rettype' + } + mapM_ (uncurry letBind) invariant + letBindNames_ (map patElemName pat') $ Op $ Screma w (mapSOAC fun') arrs +liftIdentityMapping _ _ _ _ = cannotSimplify + +-- | Remove all arguments to the map that are simply replicates. +-- These can be turned into free variables instead. +removeReplicateMapping :: TopDownRuleOp (Wise SOACS) +removeReplicateMapping vtable pat _ (Screma w form arrs) + | Just fun <- isMapSOAC form, + Just (bnds, fun', arrs') <- removeReplicateInput vtable fun arrs = do + forM_ bnds $ \(vs,cs,e) -> certifying cs $ letBindNames vs e + letBind_ pat $ Op $ Screma w (mapSOAC fun') arrs' +removeReplicateMapping _ _ _ _ = cannotSimplify + +-- | Like 'removeReplicateMapping', but for 'Scatter'. +removeReplicateWrite :: TopDownRuleOp (Wise SOACS) +removeReplicateWrite vtable pat _ (Scatter len lam ivs as) + | Just (bnds, lam', ivs') <- removeReplicateInput vtable lam ivs = do + forM_ bnds $ \(vs,cs,e) -> certifying cs $ letBindNames vs e + letBind_ pat $ Op $ Scatter len lam' ivs' as +removeReplicateWrite _ _ _ _ = cannotSimplify + +removeReplicateInput :: Aliased lore => + ST.SymbolTable lore + -> AST.Lambda lore -> [VName] + -> Maybe ([([VName], Certificates, AST.Exp lore)], + AST.Lambda lore, [VName]) +removeReplicateInput vtable fun arrs + | not $ null parameterBnds = do + let (arr_params', arrs') = unzip params_and_arrs + fun' = fun { lambdaParams = acc_params <> arr_params' } + return (parameterBnds, fun', arrs') + | otherwise = Nothing + + where params = lambdaParams fun + (acc_params, arr_params) = + splitAt (length params - length arrs) params + (params_and_arrs, parameterBnds) = + partitionEithers $ zipWith isReplicateAndNotConsumed arr_params arrs + + isReplicateAndNotConsumed p v + | Just (BasicOp (Replicate (Shape (_:ds)) e), v_cs) <- + ST.lookupExp v vtable, + not $ paramName p `S.member` consumedByLambda fun = + Right ([paramName p], + v_cs, + case ds of + [] -> BasicOp $ SubExp e + _ -> BasicOp $ Replicate (Shape ds) e) + | otherwise = + Left (p, v) + +-- | Remove inputs that are not used inside the SOAC. +removeUnusedSOACInput :: TopDownRuleOp (Wise SOACS) +removeUnusedSOACInput _ pat _ (Screma w (ScremaForm scan reduce map_lam) arrs) + | (used,unused) <- partition usedInput params_and_arrs, + not (null unused) = do + let (used_params, used_arrs) = unzip used + map_lam' = map_lam { lambdaParams = used_params } + letBind_ pat $ Op $ Screma w (ScremaForm scan reduce map_lam') used_arrs + where params_and_arrs = zip (lambdaParams map_lam) arrs + used_in_body = freeInBody $ lambdaBody map_lam + usedInput (param, _) = paramName param `S.member` used_in_body +removeUnusedSOACInput _ _ _ _ = cannotSimplify + +removeDeadMapping :: BottomUpRuleOp (Wise SOACS) +removeDeadMapping (_, used) pat _ (Screma w form arrs) + | Just fun <- isMapSOAC form = + let ses = bodyResult $ lambdaBody fun + isUsed (bindee, _, _) = (`UT.used` used) $ patElemName bindee + (pat',ses', ts') = unzip3 $ filter isUsed $ + zip3 (patternElements pat) ses $ lambdaReturnType fun + fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } + , lambdaReturnType = ts' + } + in if pat /= Pattern [] pat' + then letBind_ (Pattern [] pat') $ Op $ Screma w (mapSOAC fun') arrs + else cannotSimplify +removeDeadMapping _ _ _ _ = cannotSimplify + +removeDuplicateMapOutput :: BottomUpRuleOp (Wise SOACS) +removeDuplicateMapOutput (_, used) pat _ (Screma w form arrs) + | Just fun <- isMapSOAC form = + let ses = bodyResult $ lambdaBody fun + ts = lambdaReturnType fun + pes = patternValueElements pat + ses_ts_pes = zip3 ses ts pes + (ses_ts_pes', copies) = + foldl checkForDuplicates (mempty,mempty) ses_ts_pes + in if null copies then cannotSimplify + else do + let (ses', ts', pes') = unzip3 ses_ts_pes' + pat' = Pattern [] pes' + fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = ses' } + , lambdaReturnType = ts' } + letBind_ pat' $ Op $ Screma w (mapSOAC fun') arrs + forM_ copies $ \(from,to) -> + if UT.isConsumed (patElemName to) used then + letBind_ (Pattern [] [to]) $ BasicOp $ Copy $ patElemName from + else + letBind_ (Pattern [] [to]) $ BasicOp $ SubExp $ Var $ patElemName from + where checkForDuplicates (ses_ts_pes',copies) (se,t,pe) + | Just (_,_,pe') <- find (\(x,_,_) -> x == se) ses_ts_pes' = + -- This subexp has been returned before, producing the + -- array pe'. + (ses_ts_pes', (pe', pe) : copies) + | otherwise = (ses_ts_pes' ++ [(se,t,pe)], copies) +removeDuplicateMapOutput _ _ _ _ = cannotSimplify + +-- Mapping some operations becomes an extension of that operation. +mapOpToOp :: BottomUpRuleOp (Wise SOACS) + +mapOpToOp (_, used) pat aux1 e + | Just (map_pe, cs, w, BasicOp (Reshape newshape reshape_arr), [p], [arr]) <- + isMapWithOp pat e, + paramName p == reshape_arr, + not $ UT.isConsumed (patElemName map_pe) used = do + let redim | isJust $ shapeCoercion newshape = DimCoercion w + | otherwise = DimNew w + certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ + BasicOp $ Reshape (redim : newshape) arr + + | Just (_, cs, _, + BasicOp (Concat d arr arrs dw), ps, outer_arr : outer_arrs) <- + isMapWithOp pat e, + (arr:arrs) == map paramName ps = + certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ + BasicOp $ Concat (d+1) outer_arr outer_arrs dw + + | Just (map_pe, cs, _, + BasicOp (Rearrange perm rearrange_arr), [p], [arr]) <- + isMapWithOp pat e, + paramName p == rearrange_arr, + not $ UT.isConsumed (patElemName map_pe) used = + certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ + BasicOp $ Rearrange (0 : map (1+) perm) arr + + | Just (map_pe, cs, _, BasicOp (Rotate rots rotate_arr), [p], [arr]) <- + isMapWithOp pat e, + paramName p == rotate_arr, + not $ UT.isConsumed (patElemName map_pe) used = + certifying (stmAuxCerts aux1 <> cs) $ letBind_ pat $ + BasicOp $ Rotate (intConst Int32 0 : rots) arr + +mapOpToOp _ _ _ _ = cannotSimplify + +isMapWithOp :: PatternT attr + -> SOAC (Wise SOACS) + -> Maybe (PatElemT attr, Certificates, SubExp, + AST.Exp (Wise SOACS), [ParamT Type], [VName]) +isMapWithOp pat e + | Pattern [] [map_pe] <- pat, + Screma w form arrs <- e, + Just map_lam <- isMapSOAC form, + [Let (Pattern [] [pe]) aux2 e'] <- + stmsToList $ bodyStms $ lambdaBody map_lam, + [Var r] <- bodyResult $ lambdaBody map_lam, + r == patElemName pe = + Just (map_pe, stmAuxCerts aux2, w, e', lambdaParams map_lam, arrs) + | otherwise = Nothing + +-- | Some of the results of a reduction (or really: Redomap) may be +-- dead. We remove them here. The trick is that we need to look at +-- the data dependencies to see that the "dead" result is not +-- actually used for computing one of the live ones. +removeDeadReduction :: BottomUpRuleOp (Wise SOACS) +removeDeadReduction (_, used) pat (StmAux cs _) (Screma w form arrs) + | Just (comm, redlam, nes, maplam) <- isRedomapSOAC form, + not $ all (`UT.used` used) $ patternNames pat, -- Quick/cheap check + + let redlam_deps = dataDependencies $ lambdaBody redlam, + let redlam_res = bodyResult $ lambdaBody redlam, + let redlam_params = lambdaParams redlam, + let used_after = map snd $ filter ((`UT.used` used) . patElemName . fst) $ + zip (patternElements pat) redlam_params, + let necessary = findNecessaryForReturned (`elem` used_after) + (zip redlam_params $ redlam_res <> redlam_res) redlam_deps, + let alive_mask = map ((`S.member` necessary) . paramName) redlam_params, + + not $ all (==True) alive_mask = do + + let fixDeadToNeutral lives ne = if lives then Nothing else Just ne + dead_fix = zipWith fixDeadToNeutral alive_mask nes + (used_pes, _, used_nes) = + unzip3 $ filter (\(_,x,_) -> paramName x `S.member` necessary) $ + zip3 (patternElements pat) redlam_params nes + + let maplam' = removeLambdaResults alive_mask maplam + redlam' <- removeLambdaResults alive_mask <$> fixLambdaParams redlam (dead_fix++dead_fix) + + certifying cs $ letBind_ (Pattern [] used_pes) $ + Op $ Screma w (redomapSOAC comm redlam' used_nes maplam') arrs + +removeDeadReduction _ _ _ _ = cannotSimplify + +-- | If we are writing to an array that is never used, get rid of it. +removeDeadWrite :: BottomUpRuleOp (Wise SOACS) +removeDeadWrite (_, used) pat _ (Scatter w fun arrs dests) = + let (i_ses, v_ses) = splitAt (length dests) $ bodyResult $ lambdaBody fun + (i_ts, v_ts) = splitAt (length dests) $ lambdaReturnType fun + isUsed (bindee, _, _, _, _, _) = (`UT.used` used) $ patElemName bindee + (pat', i_ses', v_ses', i_ts', v_ts', dests') = + unzip6 $ filter isUsed $ + zip6 (patternElements pat) i_ses v_ses i_ts v_ts dests + fun' = fun { lambdaBody = (lambdaBody fun) { bodyResult = i_ses' ++ v_ses' } + , lambdaReturnType = i_ts' ++ v_ts' + } + in if pat /= Pattern [] pat' + then letBind_ (Pattern [] pat') $ Op $ Scatter w fun' arrs dests' + else cannotSimplify +removeDeadWrite _ _ _ _ = cannotSimplify + +-- handles now concatenation of more than two arrays +fuseConcatScatter :: TopDownRuleOp (Wise SOACS) +fuseConcatScatter vtable pat _ (Scatter _ fun arrs dests) + | Just (ws@(w':_), xss, css) <- unzip3 <$> mapM isConcat arrs, + xivs <- transpose xss, + all (w'==) ws = do + let r = length xivs + fun2s <- mapM (\_ -> renameLambda fun) [1 .. r-1] + let fun_n = length $ lambdaReturnType fun + (fun_is, fun_vs) = unzip $ map (splitAt (fun_n `div` 2) . + bodyResult . lambdaBody ) (fun:fun2s) + (its, vts) = unzip $ replicate r $ + splitAt (fun_n `div` 2) $ lambdaReturnType fun + new_stmts = mconcat $ map (bodyStms . lambdaBody) (fun:fun2s) + let fun' = Lambda + { lambdaParams = mconcat $ map lambdaParams (fun:fun2s) + , lambdaBody = mkBody new_stmts $ + mix fun_is <> mix fun_vs + , lambdaReturnType = mix its <> mix vts + } + certifying (mconcat css) $ + letBind_ pat $ Op $ Scatter w' fun' (concat xivs) $ map (incWrites r) dests + where sizeOf :: VName -> Maybe SubExp + sizeOf x = arraySize 0 . ST.entryType <$> ST.lookup x vtable + mix = concat . transpose + incWrites r (w, n, a) = (w, n*r, a) -- ToDO: is it (n*r) or (n+r-1)?? + isConcat v = case ST.lookupExp v vtable of + Just (BasicOp (Concat 0 x ys _), cs) -> do + x_w <- sizeOf x + y_ws<- mapM sizeOf ys + guard $ all (x_w==) y_ws + return (x_w, x:ys, cs) + _ -> Nothing + +fuseConcatScatter _ _ _ _ = cannotSimplify + +simplifyClosedFormReduce :: TopDownRuleOp (Wise SOACS) +simplifyClosedFormReduce vtable pat _ (Screma _ form arrs) + | Just (_, red_fun, nes) <- isReduceSOAC form = + foldClosedForm (`ST.lookupExp` vtable) pat red_fun nes arrs +simplifyClosedFormReduce _ _ _ _ = cannotSimplify + +-- For now we just remove singleton SOACs. +simplifyKnownIterationSOAC :: (BinderOps lore, Op lore ~ SOAC lore) => + TopDownRuleOp lore +simplifyKnownIterationSOAC _ pat _ (Screma (Constant k) + (ScremaForm (scan_lam, scan_nes) + (_, red_lam, red_nes) + map_lam) + arrs) + | oneIsh k = do + zipWithM_ bindMapParam (lambdaParams map_lam) arrs + (to_scan, to_red, map_res) <- splitAt3 (length scan_nes) (length red_nes) <$> + bodyBind (lambdaBody map_lam) + scan_res <- eLambda scan_lam $ map eSubExp $ scan_nes ++ to_scan + red_res <- eLambda red_lam $ map eSubExp $ red_nes ++ to_red + + zipWithM_ bindArrayResult scan_pes scan_res + zipWithM_ bindResult red_pes red_res + zipWithM_ bindArrayResult map_pes map_res + + where (scan_pes, red_pes, map_pes) = splitAt3 (length scan_nes) (length red_nes) $ + patternElements pat + bindMapParam p a = do + a_t <- lookupType a + letBindNames_ [paramName p] $ + BasicOp $ Index a $ fullSlice a_t [DimFix $ constant (0::Int32)] + bindArrayResult pe se = + letBindNames_ [patElemName pe] $ + BasicOp $ ArrayLit [se] $ rowType $ patElemType pe + bindResult pe se = + letBindNames_ [patElemName pe] $ BasicOp $ SubExp se +simplifyKnownIterationSOAC _ _ _ _ = cannotSimplify diff --git a/src/Futhark/Test.hs b/src/Futhark/Test.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Test.hs @@ -0,0 +1,407 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | Facilities for reading Futhark test programs. A Futhark test +-- program is an ordinary Futhark program where an initial comment +-- block specifies input- and output-sets. +module Futhark.Test + ( testSpecFromFile + , testSpecsFromPaths + , valuesFromByteString + , getValues + , getValuesBS + , compareValues + , Mismatch + + , ProgramTest (..) + , StructureTest (..) + , StructurePipeline (..) + , WarningTest (..) + , TestAction (..) + , ExpectedError (..) + , InputOutputs (..) + , TestRun (..) + , ExpectedResult (..) + , Values (..) + , Value + ) + where + +import Control.Applicative +import qualified Data.ByteString.Lazy as BS +import Control.Monad +import Control.Monad.IO.Class +import qualified Data.Map.Strict as M +import Data.Char +import Data.Functor +import Data.Maybe +import Data.Foldable (foldl') +import Data.Semigroup +import qualified Data.Text as T +import qualified Data.Text.IO as T +import qualified Data.Text.Encoding as T +import Data.Void +import System.FilePath +import Codec.Compression.GZip +import Codec.Compression.Zlib.Internal (DecompressError) +import qualified Control.Exception.Base as E + +import Text.Megaparsec hiding (many, some) +import Text.Megaparsec.Char +import Text.Regex.TDFA + +import Prelude + +import Futhark.Analysis.Metrics +import Futhark.Util.Pretty (pretty, prettyText) +import Futhark.Test.Values +import Futhark.Util (directoryContents) + +-- | Description of a test to be carried out on a Futhark program. +-- The Futhark program is stored separately. +data ProgramTest = + ProgramTest { testDescription :: + T.Text + , testTags :: + [T.Text] + , testAction :: + TestAction + } + deriving (Show) + +-- | How to test a program. +data TestAction + = CompileTimeFailure ExpectedError + | RunCases [InputOutputs] [StructureTest] [WarningTest] + deriving (Show) + +-- | Input and output pairs for some entry point(s). +data InputOutputs = InputOutputs { iosEntryPoint :: T.Text + , iosTestRuns :: [TestRun] } + deriving (Show) + +-- | The error expected for a negative test. +data ExpectedError = AnyError + | ThisError T.Text Regex + +instance Show ExpectedError where + show AnyError = "AnyError" + show (ThisError r _) = "ThisError " ++ show r + +-- | How a program can be transformed. +data StructurePipeline = KernelsPipeline + | SOACSPipeline + | SequentialCpuPipeline + | GpuPipeline + deriving (Show) + +-- | A structure test specifies a compilation pipeline, as well as +-- metrics for the program coming out the other end. +data StructureTest = StructureTest StructurePipeline AstMetrics + deriving (Show) + +-- | A warning test requires that a warning matching the regular +-- expression is produced. The program must also compile succesfully. +data WarningTest = ExpectedWarning T.Text Regex + +instance Show WarningTest where + show (ExpectedWarning r _) = "ExpectedWarning " ++ T.unpack r + +-- | A condition for execution, input, and expected result. +data TestRun = TestRun + { runTags :: [String] + , runInput :: Values + , runExpectedResult :: ExpectedResult Values + , runIndex :: Int + , runDescription :: String + } + deriving (Show) + +-- | Several Values - either literally, or by reference to a file. +data Values = Values [Value] + | InFile FilePath + deriving (Show) + +-- | How a test case is expected to terminate. +data ExpectedResult values + = Succeeds (Maybe values) -- ^ Execution suceeds, with or without + -- expected result values. + | RunTimeFailure ExpectedError -- ^ Execution fails with this error. + deriving (Show) + +type Parser = Parsec Void T.Text + +lexeme :: Parser a -> Parser a +lexeme p = p <* space + +-- | Like 'lexeme', but does not consume trailing linebreaks. +lexeme' :: Parser a -> Parser a +lexeme' p = p <* many (oneOf (" \t" :: String)) + +lexstr :: T.Text -> Parser () +lexstr = void . try . lexeme . string + +braces :: Parser a -> Parser a +braces p = lexstr "{" *> p <* lexstr "}" + +parseNatural :: Parser Int +parseNatural = lexeme $ foldl' (\acc x -> acc * 10 + x) 0 . + map num <$> some digitChar + where num c = ord c - ord '0' + +parseDescription :: Parser T.Text +parseDescription = lexeme $ T.pack <$> (anySingle `manyTill` parseDescriptionSeparator) + +parseDescriptionSeparator :: Parser () +parseDescriptionSeparator = try (string descriptionSeparator >> + void (satisfy isSpace `manyTill` newline)) <|> eof + +descriptionSeparator :: T.Text +descriptionSeparator = "==" + +parseTags :: Parser [T.Text] +parseTags = lexstr "tags" *> braces (many parseTag) <|> pure [] + where parseTag = T.pack <$> lexeme (some $ satisfy constituent) + constituent c = not (isSpace c) && c /= '}' + +parseAction :: Parser TestAction +parseAction = CompileTimeFailure <$> (lexstr "error:" *> parseExpectedError) <|> + (RunCases <$> parseInputOutputs <*> + many parseExpectedStructure <*> many parseWarning) + +parseInputOutputs :: Parser [InputOutputs] +parseInputOutputs = do + entrys <- parseEntryPoints + cases <- parseRunCases + return $ map (`InputOutputs` cases) entrys + +parseEntryPoints :: Parser [T.Text] +parseEntryPoints = (lexstr "entry:" *> many entry <* space) <|> pure ["main"] + where constituent c = not (isSpace c) && c /= '}' + entry = lexeme' $ T.pack <$> some (satisfy constituent) + +parseRunTags :: Parser [String] +parseRunTags = many parseTag + where parseTag = try $ lexeme $ do s <- some $ satisfy isAlphaNum + guard $ s `notElem` ["input", "structure", "warning"] + return s + +parseRunCases :: Parser [TestRun] +parseRunCases = parseRunCases' (0::Int) + where parseRunCases' i = (:) <$> parseRunCase i <*> parseRunCases' (i+1) + <|> pure [] + parseRunCase i = do + tags <- parseRunTags + input <- parseInput + expr <- parseExpectedResult + return $ TestRun tags input expr i $ desc i input + + -- If the file is gzipped, we strip the 'gz' extension from + -- the dataset name. This makes it more convenient to rename + -- from 'foo.in' to 'foo.in.gz', as the reported dataset name + -- does not change (which would make comparisons to historical + -- data harder). + desc _ (InFile path) + | takeExtension path == ".gz" = dropExtension path + | otherwise = path + desc i (Values vs) = + -- Turn linebreaks into space. + "#" ++ show i ++ " (\"" ++ unwords (lines vs') ++ "\")" + where vs' = case unwords (map pretty vs) of + s | length s > 50 -> take 50 s ++ "..." + | otherwise -> s + + +parseExpectedResult :: Parser (ExpectedResult Values) +parseExpectedResult = + (Succeeds . Just <$> (lexstr "output" *> parseValues)) <|> + (RunTimeFailure <$> (lexstr "error:" *> parseExpectedError)) <|> + pure (Succeeds Nothing) + +parseExpectedError :: Parser ExpectedError +parseExpectedError = lexeme $ do + s <- T.strip <$> restOfLine + if T.null s + then return AnyError + -- blankCompOpt creates a regular expression that treats + -- newlines like ordinary characters, which is what we want. + else ThisError s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s) + +parseInput :: Parser Values +parseInput = lexstr "input" *> parseValues + +parseValues :: Parser Values +parseValues = do s <- parseBlock + case valuesFromByteString "input" $ BS.fromStrict $ T.encodeUtf8 s of + Left err -> fail err + Right vs -> return $ Values vs + <|> lexstr "@" *> lexeme (InFile . T.unpack <$> nextWord) + +parseBlock :: Parser T.Text +parseBlock = lexeme $ braces (T.pack <$> parseBlockBody 0) + +parseBlockBody :: Int -> Parser String +parseBlockBody n = do + c <- lookAhead anySingle + case (c,n) of + ('}', 0) -> return mempty + ('}', _) -> (:) <$> anySingle <*> parseBlockBody (n-1) + ('{', _) -> (:) <$> anySingle <*> parseBlockBody (n+1) + _ -> (:) <$> anySingle <*> parseBlockBody n + +restOfLine :: Parser T.Text +restOfLine = T.pack <$> (anySingle `manyTill` (void newline <|> eof)) + +nextWord :: Parser T.Text +nextWord = T.pack <$> (anySingle `manyTill` satisfy isSpace) + +parseWarning :: Parser WarningTest +parseWarning = lexstr "warning:" >> parseExpectedWarning + where parseExpectedWarning = lexeme $ do + s <- T.strip <$> restOfLine + ExpectedWarning s <$> makeRegexOptsM blankCompOpt defaultExecOpt (T.unpack s) + +parseExpectedStructure :: Parser StructureTest +parseExpectedStructure = + lexstr "structure" *> + (StructureTest <$> optimisePipeline <*> parseMetrics) + +optimisePipeline :: Parser StructurePipeline +optimisePipeline = lexstr "distributed" $> KernelsPipeline <|> + lexstr "gpu" $> GpuPipeline <|> + lexstr "cpu" $> SequentialCpuPipeline <|> + pure SOACSPipeline + +parseMetrics :: Parser AstMetrics +parseMetrics = braces $ fmap (AstMetrics . M.fromList) $ many $ + (,) <$> (T.pack <$> lexeme (some (satisfy constituent))) <*> parseNatural + where constituent c = isAlpha c || c == '/' + +testSpec :: Parser ProgramTest +testSpec = + ProgramTest <$> parseDescription <*> parseTags <*> parseAction + +parserState :: Int -> FilePath -> s -> State s +parserState line name t = + State { stateInput = t + , stateOffset = 0 + , statePosState = PosState + { pstateInput = t + , pstateOffset = 0 + , pstateSourcePos = SourcePos + { sourceName = name + , sourceLine = mkPos line + , sourceColumn = mkPos 3 } + , pstateTabWidth = defaultTabWidth + , pstateLinePrefix = "-- "} + } + + +readTestSpec :: Int -> String -> T.Text -> Either (ParseErrorBundle T.Text Void) ProgramTest +readTestSpec line name t = + snd $ runParser' (testSpec <* eof) $ parserState line name t + +readInputOutputs :: Int -> String -> T.Text -> Either (ParseErrorBundle T.Text Void) [InputOutputs] +readInputOutputs line name t = + snd $ runParser' (parseDescription *> space *> parseInputOutputs <* eof) $ + parserState line name t + +commentPrefix :: T.Text +commentPrefix = T.pack "--" + +-- | Read the test specification from the given Futhark program. +-- Note: will call 'error' on parse errors. +testSpecFromFile :: FilePath -> IO ProgramTest +testSpecFromFile path = do + blocks <- testBlocks <$> T.readFile path + let (first_spec_line, first_spec, rest_specs) = + case blocks of [] -> (0, mempty, []) + (n,s):ss -> (n, s, ss) + case readTestSpec (1+first_spec_line) path first_spec of + Left err -> error $ errorBundlePretty err + Right v -> foldM moreCases v rest_specs + + where moreCases test (lineno, cases) = + case readInputOutputs lineno path cases of + Left err -> error $ errorBundlePretty err + Right cases' -> + case testAction test of + RunCases old_cases structures warnings -> + return test { testAction = RunCases (old_cases ++ cases') structures warnings } + _ -> fail "Secondary test block provided, but primary test block specifies compilation error." + +testBlocks :: T.Text -> [(Int, T.Text)] +testBlocks = mapMaybe isTestBlock . commentBlocks + where isTestBlock (n,block) + | any ((" " <> descriptionSeparator) `T.isPrefixOf`) block = + Just (n, T.unlines block) + | otherwise = + Nothing + +commentBlocks :: T.Text -> [(Int, [T.Text])] +commentBlocks = commentBlocks' . zip [0..] . T.lines + where isComment = (commentPrefix `T.isPrefixOf`) + commentBlocks' ls = + let ls' = dropWhile (not . isComment . snd) ls + in case ls' of + [] -> [] + (n,_) : _ -> + let (block, ls'') = span (isComment . snd) ls' + block' = map (T.drop 2 . snd) block + in (n, block') : commentBlocks' ls'' + +-- | Read test specifications from the given path, which can be a file +-- or directory containing @.fut@ files and further directories. +-- Calls 'error' on parse errors, or if the given path name does not +-- name a file that exists. +testSpecsFromPath :: FilePath -> IO [(FilePath, ProgramTest)] +testSpecsFromPath path = do + programs <- testPrograms path + zip programs <$> mapM testSpecFromFile programs + +-- | Read test specifications from the given paths, which can be a +-- files or directories containing @.fut@ files and further +-- directories. Calls 'error' on parse errors, or if any of the +-- immediately passed path names do not name a file that exists. +testSpecsFromPaths :: [FilePath] -> IO [(FilePath, ProgramTest)] +testSpecsFromPaths = fmap concat . mapM testSpecsFromPath + +testPrograms :: FilePath -> IO [FilePath] +testPrograms dir = filter isFut <$> directoryContents dir + where isFut = (==".fut") . takeExtension + +-- | 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 srcname = + maybe (Left $ "Cannot parse values from " ++ srcname) Right . readValues + +-- | Get the actual core Futhark values corresponding to a 'Values' +-- specification. The 'FilePath' is the directory which file paths +-- are read relative to. +getValues :: MonadIO m => FilePath -> Values -> m [Value] +getValues _ (Values vs) = + return vs +getValues dir (InFile file) = do + s <- getValuesBS dir (InFile file) + case valuesFromByteString file' s of + Left e -> fail $ show e + Right vs -> return vs + where file' = dir file + +-- | Extract a pretty representation of some 'Values'. In the IO +-- monad because this might involve reading from a file. There is no +-- guarantee that the resulting byte string yields a readable value. +getValuesBS :: MonadIO m => FilePath -> Values -> m BS.ByteString +getValuesBS _ (Values vs) = + return $ BS.fromStrict $ T.encodeUtf8 $ T.unlines $ map prettyText vs +getValuesBS dir (InFile file) = + case takeExtension file of + ".gz" -> liftIO $ do + s <- E.try readAndDecompress + case s of + Left e -> fail $ show file ++ ": " ++ show (e :: DecompressError) + Right s' -> return s' + + _ -> liftIO $ BS.readFile file' + where file' = dir file + readAndDecompress = do s <- BS.readFile file' + E.evaluate $ decompress s diff --git a/src/Futhark/Test/Values.hs b/src/Futhark/Test/Values.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Test/Values.hs @@ -0,0 +1,542 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | 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. +module Futhark.Test.Values + ( Value + , valueType + + -- * Reading Values + , readValues + + -- * Comparing Values + , compareValues + , Mismatch + , explainMismatch + ) + where + +import Control.Monad +import Control.Monad.ST +import Data.Binary +import Data.Binary.Put +import Data.Binary.Get +import Data.Binary.IEEE754 +import qualified Data.ByteString.Lazy.Char8 as BS +import Data.Maybe +import Data.Int (Int8, Int16, Int32, Int64) +import Data.Char (isSpace, ord, chr) +import Data.Vector.Binary +import qualified Data.Vector.Unboxed.Mutable as UMVec +import qualified Data.Vector.Unboxed as UVec +import Data.Vector.Generic (freeze) +import Data.Loc (Pos(..)) + +import qualified Language.Futhark.Syntax as F +import Language.Futhark.Pretty() +import Futhark.Representation.Primitive (PrimValue) +import Language.Futhark.Parser.Lexer +import qualified Futhark.Util.Pretty as PP +import Futhark.Representation.AST.Attributes.Constants (IsValue(..)) +import Futhark.Representation.AST.Pretty () +import Futhark.Util.Pretty + +type STVector s = UMVec.STVector s +type Vector = UVec.Vector + +-- | An efficiently represented Futhark value. Use 'pretty' to get a +-- human-readable representation, and the instances of 'Get' and 'Put' +-- to obtain binary representations +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 putInt8 + put (Int16Value shape vs) = putBinaryValue " i16" shape vs putInt16le + put (Int32Value shape vs) = putBinaryValue " i32" shape vs putInt32le + put (Int64Value shape vs) = putBinaryValue " i64" shape vs putInt64le + put (Word8Value shape vs) = putBinaryValue " i8" shape vs putWord8 + put (Word16Value shape vs) = putBinaryValue " i16" shape vs putWord16le + put (Word32Value shape vs) = putBinaryValue " i32" shape vs putWord32le + put (Word64Value shape vs) = putBinaryValue " i64" shape vs putWord64le + put (Float32Value shape vs) = putBinaryValue " f32" shape vs putFloat32le + put (Float64Value shape vs) = putBinaryValue " f64" shape vs putFloat64le + put (BoolValue shape vs) = putBinaryValue " f64" shape vs $ putInt8 . boolToInt + where boolToInt True = 1 + boolToInt 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' = UVec.fromList shape + + case BS.unpack type_f of + " i8" -> get' (Int8Value shape') getInt8 num_elems + " i16" -> get' (Int16Value shape') getInt16le num_elems + " i32" -> get' (Int32Value shape') getInt32le num_elems + " i64" -> get' (Int64Value shape') getInt64le num_elems + " u8" -> get' (Word8Value shape') getWord8 num_elems + " u16" -> get' (Word16Value shape') getWord16le num_elems + " u32" -> get' (Word32Value shape') getWord32le num_elems + " u64" -> get' (Word64Value shape') getWord64le num_elems + " f32" -> get' (Float32Value shape') getFloat32le num_elems + " f64" -> get' (Float64Value shape') getFloat64le num_elems + "bool" -> get' (BoolValue shape') getBool num_elems + s -> fail $ "Cannot parse binary values of type " ++ show s + where getBool = (/=0) <$> getWord8 + + get' mk get_elem num_elems = + mk <$> genericGetVectorWith (pure num_elems) get_elem + +putBinaryValue :: UVec.Unbox a => + String -> Vector Int -> Vector a -> (a -> Put) -> Put +putBinaryValue tstr shape vs putv = do + putInt8 $ fromIntegral $ ord 'b' + putWord8 binaryFormatVersion + putWord8 $ fromIntegral $ UVec.length shape + mapM_ (putInt8 . fromIntegral . ord) tstr + mapM_ (putInt64le . fromIntegral) $ UVec.toList shape + mapM_ putv $ UVec.toList vs + +instance PP.Pretty Value where + ppr v | product (valueShape v) == 0 = + text "empty" <> + parens (dims <> text (valueElemType v)) + where dims = mconcat $ replicate (length (valueShape v)-1) $ text "[]" + ppr (Int8Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Int16Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Int32Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Int64Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Word8Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Word16Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Word32Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Word64Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Float32Value shape vs) = pprArray (UVec.toList shape) vs + ppr (Float64Value shape vs) = pprArray (UVec.toList shape) vs + ppr (BoolValue shape vs) = pprArray (UVec.toList shape) vs + +pprArray :: (UVec.Unbox a, F.IsPrimValue a) => [Int] -> UVec.Vector a -> Doc +pprArray [] vs = + ppr $ F.primValue $ UVec.head vs +pprArray (d:ds) vs = + brackets $ commasep $ map (pprArray ds . slice) [0..d-1] + where slice_size = product ds + slice i = UVec.slice (i*slice_size) slice_size vs + +-- | A textual description of the type of a value. Follows Futhark +-- type notation, and contains the exact dimension sizes if an array. +valueType :: Value -> String +valueType v = concatMap (\d -> "[" ++ show d ++ "]") (valueShape v) ++ + valueElemType v + +valueElemType :: Value -> String +valueElemType (Int8Value _ _) = "i8" +valueElemType (Int16Value _ _) = "i16" +valueElemType (Int32Value _ _) = "i32" +valueElemType (Int64Value _ _) = "i64" +valueElemType (Word8Value _ _) = "u8" +valueElemType (Word16Value _ _) = "u16" +valueElemType (Word32Value _ _) = "u32" +valueElemType (Word64Value _ _) = "u64" +valueElemType (Float32Value _ _) = "f32" +valueElemType (Float64Value _ _) = "f64" +valueElemType (BoolValue _ _) = "bool" + +valueShape :: Value -> [Int] +valueShape (Int8Value shape _) = UVec.toList shape +valueShape (Int16Value shape _) = UVec.toList shape +valueShape (Int32Value shape _) = UVec.toList shape +valueShape (Int64Value shape _) = UVec.toList shape +valueShape (Word8Value shape _) = UVec.toList shape +valueShape (Word16Value shape _) = UVec.toList shape +valueShape (Word32Value shape _) = UVec.toList shape +valueShape (Word64Value shape _) = UVec.toList shape +valueShape (Float32Value shape _) = UVec.toList shape +valueShape (Float64Value shape _) = UVec.toList shape +valueShape (BoolValue shape _) = UVec.toList shape + +-- The parser + +dropRestOfLine, dropSpaces :: BS.ByteString -> BS.ByteString +dropRestOfLine = BS.drop 1 . BS.dropWhile (/='\n') +dropSpaces t = case BS.dropWhile isSpace t of + t' | "--" `BS.isPrefixOf` t' -> dropSpaces $ dropRestOfLine t' + | otherwise -> t' + +type ReadValue v = BS.ByteString -> Maybe (v, BS.ByteString) + +symbol :: Char -> BS.ByteString -> Maybe BS.ByteString +symbol c t + | Just (c', t') <- BS.uncons t, c' == c = Just $ dropSpaces t' + | otherwise = Nothing + +lexeme :: BS.ByteString -> BS.ByteString -> Maybe BS.ByteString +lexeme l t + | l `BS.isPrefixOf` t = Just $ dropSpaces $ BS.drop (BS.length l) t + | otherwise = Nothing + +-- (Used elements, shape, elements, remaining input) +type State s v = (Int, Vector Int, STVector s v, BS.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 -> + readArrayElemsST (j+1) r rv (i, shape, arr, t') + | 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 UVec.// [(r-d, n)] + | old_n == n = Just shape + | otherwise = Nothing + where r = UVec.length shape + old_n = shape UVec.! (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 => + Int -> ReadValue v -> BS.ByteString -> Maybe (Vector Int, Vector v, BS.ByteString) +readRankedArrayOf r rv t = runST $ do + arr <- UMVec.new 1024 + ms <- readRankedArrayOfST r rv (0, UVec.replicate r (-1), arr, t) + case ms of + Just (i, shape, arr', t') -> do + arr'' <- freeze (UMVec.slice 0 i arr') + return $ Just (shape, 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) = BS.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) = BS.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) = BS.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) = BS.span constituent t + +readEmptyArrayOfRank :: Int -> BS.ByteString -> Maybe (Value, BS.ByteString) +readEmptyArrayOfRank r t + | Just t' <- symbol '[' t, + Just t'' <- symbol ']' t' = readEmptyArrayOfRank (r+1) t'' + | otherwise = do + (pt, t') <- readPrimType t + v <- case pt of + "i8" -> Just $ Int8Value (UVec.replicate r 0) UVec.empty + "i16" -> Just $ Int16Value (UVec.replicate r 0) UVec.empty + "i32" -> Just $ Int32Value (UVec.replicate r 0) UVec.empty + "i64" -> Just $ Int64Value (UVec.replicate r 0) UVec.empty + "u8" -> Just $ Word8Value (UVec.replicate r 0) UVec.empty + "u16" -> Just $ Word16Value (UVec.replicate r 0) UVec.empty + "u32" -> Just $ Word32Value (UVec.replicate r 0) UVec.empty + "u64" -> Just $ Word64Value (UVec.replicate r 0) UVec.empty + "f32" -> Just $ Float32Value (UVec.replicate r 0) UVec.empty + "f64" -> Just $ Float64Value (UVec.replicate r 0) UVec.empty + "bool" -> Just $ BoolValue (UVec.replicate r 0) UVec.empty + _ -> Nothing + return (v, t') + +readEmptyArray :: BS.ByteString -> Maybe (Value, BS.ByteString) +readEmptyArray t = do + t' <- symbol '(' =<< lexeme "empty" t + (v, t'') <- readEmptyArrayOfRank 1 t' + t''' <- symbol ')' t'' + return (v, t''') + +readValue :: BS.ByteString -> Maybe (Value, BS.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 :: BS.ByteString -> Maybe [Value] +readValues = readValues' . dropSpaces + where readValues' t + | BS.null t = Just [] + | otherwise = do (a, t') <- readValue t + (a:) <$> readValues' t' + +-- Comparisons + +-- | Two values differ in some way. +data Mismatch = PrimValueMismatch (Int,Int) PrimValue PrimValue + -- ^ The position the value number and a flat index + -- into the array. + | ArrayShapeMismatch Int [Int] [Int] + | TypeMismatch Int String String + | ValueCountMismatch Int Int + +instance Show Mismatch where + show (PrimValueMismatch (i,j) got expected) = + explainMismatch (i,j) "" got expected + show (ArrayShapeMismatch i got expected) = + explainMismatch i "array of shape " got expected + show (TypeMismatch i got expected) = + explainMismatch 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 :: (Show i, PP.Pretty a) => i -> String -> a -> a -> String +explainMismatch i what got expected = + "Value " ++ show 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] -> Maybe [Mismatch] +compareValues got expected + | n /= m = Just [ValueCountMismatch n m] + | otherwise = case catMaybes $ zipWith3 compareValue [0..] got expected of + [] -> Nothing + es -> Just es + where n = length got + m = length expected + + +compareValue :: Int -> Value -> Value -> Maybe 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 + _ -> + Just $ TypeMismatch i (valueElemType got_v) (valueElemType expected_v) + | otherwise = + Just $ ArrayShapeMismatch i (valueShape got_v) (valueShape expected_v) + where compareNum tol = compareGen $ compareElement tol + compareFloat tol = compareGen $ compareFloatElement tol + + compareGen cmp got expected = + foldl mplus Nothing $ + zipWith cmp (UVec.toList $ UVec.indexed got) (UVec.toList expected) + + compareElement tol (j, got) expected + | comparePrimValue tol got expected = Nothing + | otherwise = Just $ PrimValueMismatch (i,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,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 :: (Ord a, Fractional a, UVec.Unbox a) => Vector a -> a +tolerance = UVec.foldl tolerance' minTolerance + where tolerance' t v = max t $ minTolerance * v diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Tools.hs @@ -0,0 +1,195 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +module Futhark.Tools + ( + module Futhark.Construct + + , nonuniqueParams + , redomapToMapAndReduce + , scanomapToMapAndScan + , dissectScrema + , sequentialStreamWholeArray + + , partitionChunkedFoldParameters + , partitionChunkedKernelLambdaParameters + , partitionChunkedKernelFoldParameters + + -- * Primitive expressions + , module Futhark.Analysis.PrimExp.Convert + ) +where + +import Control.Monad.Identity +import Data.Semigroup ((<>)) + +import Futhark.Representation.AST +import Futhark.Representation.SOACS.SOAC +import Futhark.MonadFreshNames +import Futhark.Construct +import Futhark.Analysis.PrimExp.Convert +import Futhark.Util + +nonuniqueParams :: (MonadFreshNames m, Bindable lore, HasScope lore m, BinderOps lore) => + [LParam lore] -> m ([LParam lore], Stms lore) +nonuniqueParams params = runBinder $ forM params $ \param -> + if not $ primType $ paramType param then do + param_name <- newVName $ baseString (paramName param) ++ "_nonunique" + let param' = Param param_name $ paramType param + localScope (scopeOfLParams [param']) $ + letBindNames_ [paramName param] $ BasicOp $ Copy $ paramName param' + return param' + else + return param + +-- | Turns a binding of a @redomap@ into two seperate bindings, a +-- @map@ binding and a @reduce@ binding (returned in that order). +-- +-- Reuses the original pattern for the @reduce@, and creates a new +-- pattern with new 'Ident's for the result of the @map@. +-- +-- Only handles a 'Pattern' with an empty 'patternContextElements' +redomapToMapAndReduce :: (MonadFreshNames m, Bindable lore, + ExpAttr lore ~ (), Op lore ~ SOAC lore) => + Pattern lore + -> ( SubExp + , Commutativity + , LambdaT lore, LambdaT lore, [SubExp] + , [VName]) + -> m (Stm lore, Stm lore) +redomapToMapAndReduce (Pattern [] patelems) + (w, comm, redlam, map_lam, accs, arrs) = do + (map_pat, red_pat, red_args) <- + splitScanOrRedomap patelems w map_lam accs + let map_bnd = mkLet [] map_pat $ Op $ Screma w (mapSOAC map_lam) arrs + (nes, red_arrs) = unzip red_args + red_bnd <- Let red_pat (defAux ()) . Op <$> + (Screma w <$> reduceSOAC comm redlam nes <*> pure red_arrs) + return (map_bnd, red_bnd) +redomapToMapAndReduce _ _ = + error "redomapToMapAndReduce does not handle a non-empty 'patternContextElements'" + +-- | Like 'redomapToMapAndReduce', but for 'Scanomap'. +scanomapToMapAndScan :: (MonadFreshNames m, Bindable lore, + ExpAttr lore ~ (), Op lore ~ SOAC lore) => + Pattern lore + -> ( SubExp + , LambdaT lore, LambdaT lore, [SubExp] + , [VName]) + -> m (Stm lore, Stm lore) +scanomapToMapAndScan (Pattern [] patelems) (w, scanlam, map_lam, accs, arrs) = do + (map_pat, scan_pat, scan_args) <- + splitScanOrRedomap patelems w map_lam accs + let map_bnd = mkLet [] map_pat $ Op $ Screma w (mapSOAC map_lam) arrs + (nes, scan_arrs) = unzip scan_args + scan_bnd <- Let scan_pat (defAux ()) . Op <$> + (Screma w <$> scanSOAC scanlam nes <*> pure scan_arrs) + return (map_bnd, scan_bnd) +scanomapToMapAndScan _ _ = + error "scanomapToMapAndScan does not handle a non-empty 'patternContextElements'" + +splitScanOrRedomap :: (Typed attr, MonadFreshNames m) => + [PatElemT attr] + -> SubExp -> LambdaT lore -> [SubExp] + -> m ([Ident], PatternT attr, [(SubExp, VName)]) +splitScanOrRedomap patelems w map_lam accs = do + let (acc_patelems, arr_patelems) = splitAt (length accs) patelems + (acc_ts, _arr_ts) = splitAt (length accs) $ lambdaReturnType map_lam + map_accpat <- zipWithM accMapPatElem acc_patelems acc_ts + map_arrpat <- mapM arrMapPatElem arr_patelems + let map_pat = map_accpat ++ map_arrpat + red_args = zip accs $ map identName map_accpat + return (map_pat, Pattern [] acc_patelems, red_args) + where + accMapPatElem pe acc_t = + newIdent (baseString (patElemName pe) ++ "_map_acc") $ acc_t `arrayOfRow` w + arrMapPatElem = return . patElemIdent + +-- | Turn a Screma into simpler Scremas that are all simple scans, +-- reduces, and maps. This is used to handle Scremas that are so +-- complicated that we cannot directly generate efficient parallel +-- code for them. In essense, what happens is the opposite of +-- horisontal fusion. +dissectScrema :: (MonadBinder m, Op (Lore m) ~ SOAC (Lore m), + Bindable (Lore m)) => + Pattern (Lore m) -> SubExp -> ScremaForm (Lore m) -> [VName] + -> m () +dissectScrema pat w (ScremaForm (scan_lam, scan_nes) + (comm, red_lam, red_nes) + map_lam) arrs = do + let (scan_res, red_res, map_res) = splitAt3 (length scan_nes) (length red_nes) $ + patternNames pat + -- First we perform the Map, then we perform the Reduce, and finally + -- the Scan. + to_scan <- replicateM (length scan_nes) $ newVName "to_scan" + to_red <- replicateM (length red_nes) $ newVName "to_red" + letBindNames_ (to_scan <> to_red <> map_res) $ Op $ Screma w (mapSOAC map_lam) arrs + + reduce <- reduceSOAC comm red_lam red_nes + letBindNames_ red_res $ Op $ Screma w reduce to_red + + scan <- scanSOAC scan_lam scan_nes + letBindNames_ scan_res $ Op $ Screma w scan to_scan + +sequentialStreamWholeArray :: (MonadBinder m, Bindable (Lore m)) => + Pattern (Lore m) + -> SubExp -> [SubExp] + -> LambdaT (Lore m) -> [VName] + -> m () +sequentialStreamWholeArray pat w nes lam arrs = do + -- We just set the chunksize to w and inline the lambda body. There + -- is no difference between parallel and sequential streams here. + let (chunk_size_param, fold_params, arr_params) = + partitionChunkedFoldParameters (length nes) $ lambdaParams lam + + -- The chunk size is the full size of the array. + letBindNames_ [paramName chunk_size_param] $ BasicOp $ SubExp w + + -- The accumulator parameters are initialised to the neutral element. + forM_ (zip fold_params nes) $ \(p, ne) -> + letBindNames [paramName p] $ BasicOp $ SubExp ne + + -- Finally, the array parameters are set to the arrays (but reshaped + -- to make the types work out; this will be simplified rapidly). + forM_ (zip arr_params arrs) $ \(p, arr) -> + letBindNames [paramName p] $ BasicOp $ + Reshape (map DimCoercion $ arrayDims $ paramType p) arr + + -- Then we just inline the lambda body. + mapM_ addStm $ bodyStms $ lambdaBody lam + + -- The number of results in the body matches exactly the size (and + -- order) of 'pat', so we bind them up here, again with a reshape to + -- make the types work out. We also do a copy to ensure that the + -- result does not have any aliases (as the semantics of Stream + -- require). + forM_ (zip (patternElements pat) $ bodyResult $ lambdaBody lam) $ \(pe, se) -> + case (arrayDims $ patElemType pe, se) of + (dims, Var v) + | not $ null dims -> do + v_reshaped <- letExp (baseString v <> "_reshaped") $ + BasicOp $ Reshape (map DimCoercion dims) v + letBindNames_ [patElemName pe] $ BasicOp $ Copy v_reshaped + _ -> letBindNames_ [patElemName pe] $ BasicOp $ SubExp se + +partitionChunkedFoldParameters :: Int -> [Param attr] + -> (Param attr, [Param attr], [Param attr]) +partitionChunkedFoldParameters _ [] = + error "partitionChunkedFoldParameters: lambda takes no parameters" +partitionChunkedFoldParameters num_accs (chunk_param : params) = + let (acc_params, arr_params) = splitAt num_accs params + in (chunk_param, acc_params, arr_params) + +partitionChunkedKernelFoldParameters :: Int -> [Param attr] + -> (VName, Param attr, [Param attr], [Param attr]) +partitionChunkedKernelFoldParameters num_accs (i_param : chunk_param : params) = + let (acc_params, arr_params) = splitAt num_accs params + in (paramName i_param, chunk_param, acc_params, arr_params) +partitionChunkedKernelFoldParameters _ _ = + error "partitionChunkedKernelFoldParameters: lambda takes too few parameters" + +partitionChunkedKernelLambdaParameters :: [Param attr] + -> (VName, Param attr, [Param attr]) +partitionChunkedKernelLambdaParameters (i_param : chunk_param : params) = + (paramName i_param, chunk_param, params) +partitionChunkedKernelLambdaParameters _ = + error "partitionChunkedKernelLambdaParameters: lambda takes too few parameters" diff --git a/src/Futhark/Transform/CopyPropagate.hs b/src/Futhark/Transform/CopyPropagate.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Transform/CopyPropagate.hs @@ -0,0 +1,19 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ConstraintKinds #-} +-- | Perform copy propagation. This is done by invoking the +-- simplifier with no rules, so hoisting and dead-code elimination may +-- also take place. +module Futhark.Transform.CopyPropagate + (copyPropagateInStms) + where + +import Futhark.MonadFreshNames +import Futhark.Representation.AST +import Futhark.Optimise.Simplify + +-- | Run copy propagation. +copyPropagateInStms :: (MonadFreshNames m, SimplifiableLore lore, HasScope lore m) => + SimpleOps lore + -> Stms lore + -> m (Stms lore) +copyPropagateInStms simpl = simplifyStms simpl mempty noExtraHoistBlockers diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -0,0 +1,382 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} +-- | The code generator cannot handle the array combinators (@map@ and +-- friends), so this module was written to transform them into the +-- equivalent do-loops. The transformation is currently rather naive, +-- and - it's certainly worth considering when we can express such +-- transformations in-place. +module Futhark.Transform.FirstOrderTransform + ( transformFunDef + + , Transformer + , transformStmRecursively + , transformLambda + , transformSOAC + , transformBody + + -- * Utility + , doLoopMapAccumL + , doLoopMapAccumL' + ) + where + +import Control.Monad.Except +import Control.Monad.State +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import qualified Futhark.Representation.AST as AST +import Futhark.Representation.SOACS +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Representation.Aliases (Aliases, removeLambdaAliases) +import Futhark.Representation.AST.Attributes.Aliases +import Futhark.Util (chunks, splitAt3) + +transformFunDef :: (MonadFreshNames m, Bindable tolore, BinderOps tolore, + LetAttr SOACS ~ LetAttr tolore, + CanBeAliased (Op tolore)) => + FunDef -> m (AST.FunDef tolore) +transformFunDef (FunDef entry fname rettype params body) = do + (body',_) <- modifyNameSource $ runState $ runBinderT m mempty + return $ FunDef entry fname rettype params body' + where m = localScope (scopeOfFParams params) $ insertStmsM $ transformBody body + +-- | The constraints that a monad must uphold in order to be used for +-- first-order transformation. +type Transformer m = (MonadBinder m, + Bindable (Lore m), BinderOps (Lore m), + LocalScope (Lore m) m, + LetAttr SOACS ~ LetAttr (Lore m), + LParamAttr SOACS ~ LParamAttr (Lore m), + CanBeAliased (Op (Lore m))) + +transformBody :: Transformer m => + Body -> m (AST.Body (Lore m)) +transformBody (Body () bnds res) = insertStmsM $ do + mapM_ transformStmRecursively bnds + return $ resultBody res + +-- | First transform any nested 'Body' or 'Lambda' elements, then +-- apply 'transformSOAC' if the expression is a SOAC. +transformStmRecursively :: Transformer m => + Stm -> m () + +transformStmRecursively (Let pat aux (Op soac)) = + certifying (stmAuxCerts aux) $ + transformSOAC pat =<< mapSOACM soacTransform soac + where soacTransform = identitySOACMapper { mapOnSOACLambda = transformLambda } + +transformStmRecursively (Let pat aux e) = + certifying (stmAuxCerts aux) $ + letBind_ pat =<< mapExpM transform e + where transform = identityMapper { mapOnBody = \scope -> localScope scope . transformBody + , mapOnRetType = return + , mapOnBranchType = return + , mapOnFParam = return + , mapOnLParam = return + , mapOnOp = fail "Unhandled Op in first order transform" + } + +-- | 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. +transformSOAC :: Transformer m => + AST.Pattern (Lore m) + -> SOAC (Lore m) + -> m () + +transformSOAC pat CmpThreshold{} = + letBind_ pat $ BasicOp $ SubExp $ constant False -- close enough + +transformSOAC pat (Screma w form@(ScremaForm (scan_lam, scan_nes) + (_, red_lam, red_nes) + map_lam) arrs) = do + let (scan_arr_ts, _red_ts, map_arr_ts) = + splitAt3 (length scan_nes) (length red_nes) $ scremaType w form + scan_arrs <- resultArray scan_arr_ts + map_arrs <- resultArray map_arr_ts + + -- We construct a loop that contains several groups of merge + -- parameters: + -- + -- (0) scan accumulator. + -- (1) scan results. + -- (2) reduce results (and accumulator). + -- (3) map results. + -- + -- Inside the loop, the parameters to map_lam become for-in + -- parameters. + + scanacc_params <- mapM (newParam "scanacc" . flip toDecl Nonunique) $ lambdaReturnType scan_lam + scanout_params <- mapM (newParam "scanout" . flip toDecl Unique) scan_arr_ts + redout_params <- mapM (newParam "redout" . flip toDecl Nonunique) $ lambdaReturnType red_lam + mapout_params <- mapM (newParam "mapout" . flip toDecl Unique) map_arr_ts + + let merge = concat [zip scanacc_params scan_nes, + zip scanout_params $ map Var scan_arrs, + zip redout_params red_nes, + zip mapout_params $ map Var map_arrs] + i <- newVName "i" + let loopform = ForLoop i Int32 w [] + + loop_body <- runBodyBinder $ + localScope (scopeOfFParams $ map fst merge) $ + inScopeOf loopform $ do + + forM_ (zip (lambdaParams map_lam) arrs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames_ [paramName p] $ BasicOp $ Index arr $ + fullSlice arr_t [DimFix $ Var i] + + -- Insert the statements of the lambda. We have taken care to + -- ensure that the parameters are bound at this point. + mapM_ addStm $ bodyStms $ lambdaBody map_lam + -- Split into scan results, reduce results, and map results. + let (scan_res, red_res, map_res) = + splitAt3 (length scan_nes) (length red_nes) $ + bodyResult $ lambdaBody map_lam + + scan_res' <- eLambda scan_lam $ map (pure . BasicOp . SubExp) $ + map (Var . paramName) scanacc_params ++ scan_res + red_res' <- eLambda red_lam $ map (pure . BasicOp . SubExp) $ + map (Var . paramName) redout_params ++ red_res + + -- Write the scan accumulator to the scan result arrays. + scan_outarrs <- letwith (map paramName scanout_params) (pexp (Var i)) $ + map (BasicOp . SubExp) scan_res' + + -- Write the map results to the map result arrays. + map_outarrs <- letwith (map paramName mapout_params) (pexp (Var i)) $ + map (BasicOp . SubExp) map_res + + return $ resultBody $ concat [scan_res', + map Var scan_outarrs, + red_res', + map Var map_outarrs] + + -- We need to discard the final scan accumulators, as they are not + -- bound in the original pattern. + pat' <- discardPattern (map paramType scanacc_params) pat + letBind_ pat' $ DoLoop [] merge loopform loop_body + +transformSOAC pat (Stream w form lam arrs) = + sequentialStreamWholeArray pat w nes lam arrs + where nes = getStreamAccums form + +transformSOAC pat (Scatter len lam ivs as) = do + iter <- newVName "write_iter" + + let (_as_ws, as_ns, as_vs) = unzip3 as + ts <- mapM lookupType as_vs + asOuts <- mapM (newIdent "write_out") ts + + let ivsLen = length (lambdaReturnType lam) `div` 2 + + -- Scatter is in-place, so we use the input array as the output array. + let merge = loopMerge asOuts $ map Var as_vs + loopBody <- runBodyBinder $ + localScope (M.insert iter (IndexInfo Int32) $ + scopeOfFParams $ map fst merge) $ do + ivs' <- forM ivs $ \iv -> do + iv_t <- lookupType iv + letSubExp "write_iv" $ BasicOp $ Index iv $ fullSlice iv_t [DimFix $ Var iter] + ivs'' <- bindLambda lam (map (BasicOp . SubExp) ivs') + + let indexes = chunks as_ns $ take ivsLen ivs'' + values = chunks as_ns $ drop ivsLen ivs'' + + ress <- forM (zip3 indexes values (map identName asOuts)) $ \(indexes', values', arr) -> do + let saveInArray arr' (indexCur, valueCur) = + letExp "write_out" =<< eWriteArray arr' [eSubExp indexCur] (eSubExp valueCur) + + foldM saveInArray arr $ zip indexes' values' + return $ resultBody (map Var ress) + letBind_ pat $ DoLoop [] merge (ForLoop iter Int32 len []) loopBody + +transformSOAC pat (GenReduce len ops bucket_fun imgs) = do + iter <- newVName "iter" + + -- Bind arguments to parameters for the merge-variables. + hists_ts <- mapM lookupType $ concatMap genReduceDest ops + hists_out <- mapM (newIdent "dests") hists_ts + let merge = loopMerge hists_out $ concatMap (map Var . genReduceDest) ops + + -- Bind lambda-bodies for operators. + loopBody <- runBodyBinder $ + localScope (M.insert iter (IndexInfo Int32) $ + scopeOfFParams $ map fst merge) $ do + + -- Bind images to parameters of bucket function. + imgs' <- forM imgs $ \img -> do + img_t <- lookupType img + letSubExp "pixel" $ BasicOp $ Index img $ fullSlice img_t [DimFix $ Var iter] + imgs'' <- bindLambda bucket_fun $ map (BasicOp . SubExp) imgs' + + -- Split out values from bucket function. + let lens = length ops + inds = take lens imgs'' + vals = chunks (map (length . lambdaReturnType . genReduceOp) ops) $ drop lens imgs'' + hists_out' = chunks (map (length . lambdaReturnType . genReduceOp) ops) $ + map identName hists_out + + -- Read values from histograms. + h_vals <- forM (zip inds hists_out') $ \(idx, hist) -> + forM hist $ \arr -> do + arr_t <- lookupType arr + letSubExp "read_hist" $ BasicOp $ Index arr $ fullSlice arr_t [DimFix idx] + + -- Apply operators. + h_vals' <- forM (zip3 (map genReduceOp ops) vals h_vals) $ \(op, ne_val, h_val) -> + bindLambda op $ map (BasicOp . SubExp) $ ne_val ++ h_val + + -- Write values back to histograms. + ress <- forM (zip3 inds h_vals' hists_out') $ \(idx, val, hist) -> + forM (zip val hist) $ \(v, arr) -> + letExp "write_hist" =<< eWriteArray arr [eSubExp idx] (eSubExp v) + + return $ resultBody $ map Var $ concat ress + -- Wrap up the above into a for-loop. + letBind_ pat $ DoLoop [] merge (ForLoop iter Int32 len []) loopBody + +-- | Recursively first-order-transform a lambda. +transformLambda :: (MonadFreshNames m, + Bindable lore, BinderOps lore, + LocalScope somelore m, + SameScope somelore lore, + LetAttr SOACS ~ LetAttr lore, + CanBeAliased (Op lore)) => + Lambda -> m (AST.Lambda lore) +transformLambda (Lambda params body rettype) = do + body' <- runBodyBinder $ + localScope (scopeOfLParams params) $ + transformBody body + return $ Lambda params body' rettype + +newFold :: Transformer m => + String -> [(SubExp,Type)] -> [VName] + -> m ([Ident], [SubExp], [Ident]) +newFold what accexps_and_types arrexps = do + initacc <- mapM copyIfArray acc_exps + acc <- mapM (newIdent "acc") acc_types + arrts <- mapM lookupType arrexps + inarrs <- mapM (newIdent $ what ++ "_inarr") arrts + return (acc, initacc, inarrs) + where (acc_exps, acc_types) = unzip accexps_and_types + +copyIfArray :: Transformer m => + SubExp -> m SubExp +copyIfArray (Constant v) = return $ Constant v +copyIfArray (Var v) = Var <$> copyIfArrayName v + +copyIfArrayName :: Transformer m => + VName -> m VName +copyIfArrayName v = do + t <- lookupType v + case t of + Array {} -> letExp (baseString v ++ "_first_order_copy") $ BasicOp $ Copy v + _ -> return v + +index :: (HasScope lore m, Monad m) => + [VName] -> SubExp -> m [AST.Exp lore] +index arrs i = forM arrs $ \arr -> do + arr_t <- lookupType arr + return $ BasicOp $ Index arr $ fullSlice arr_t [DimFix i] + +resultArray :: Transformer m => [Type] -> m [VName] +resultArray = mapM oneArray + where oneArray t = letExp "result" $ BasicOp $ Scratch (elemType t) (arrayDims t) + +letwith :: Transformer m => + [VName] -> m (AST.Exp (Lore m)) -> [AST.Exp (Lore m)] + -> m [VName] +letwith ks i vs = do + vs' <- letSubExps "values" vs + i' <- letSubExp "i" =<< i + let update k v = do + k_t <- lookupType k + letInPlace "lw_dest" k (fullSlice k_t [DimFix i']) $ BasicOp $ SubExp v + zipWithM update ks vs' + +pexp :: Applicative f => SubExp -> f (AST.Exp lore) +pexp = pure . BasicOp . SubExp + +bindLambda :: Transformer m => + AST.Lambda (Lore m) -> [AST.Exp (Lore m)] + -> m [SubExp] +bindLambda (Lambda params body _) args = do + forM_ (zip params args) $ \(param, arg) -> + if primType $ paramType param + then letBindNames [paramName param] arg + else letBindNames [paramName param] =<< eCopy (pure arg) + bodyBind body + +loopMerge :: [Ident] -> [SubExp] -> [(Param DeclType, SubExp)] +loopMerge vars = loopMerge' $ zip vars $ repeat Unique + +loopMerge' :: [(Ident,Uniqueness)] -> [SubExp] -> [(Param DeclType, SubExp)] +loopMerge' vars vals = [ (Param pname $ toDecl ptype u, val) + | ((Ident pname ptype, u),val) <- zip vars vals ] + +discardPattern :: (MonadFreshNames m, LetAttr (Lore m) ~ LetAttr SOACS) => + [Type] -> AST.Pattern (Lore m) -> m (AST.Pattern (Lore m)) +discardPattern discard pat = do + discard_pat <- basicPattern [] <$> mapM (newIdent "discard") discard + return $ discard_pat <> pat + +-- | Turn a Haskell-style mapAccumL into a sequential do-loop. This +-- is the guts of transforming a 'Redomap'. +doLoopMapAccumL :: (LocalScope (Lore m) m, MonadBinder m, + Bindable (Lore m), BinderOps (Lore m), + LetAttr (Lore m) ~ Type, + CanBeAliased (Op (Lore m))) => + SubExp + -> AST.Lambda (Aliases (Lore m)) + -> [SubExp] + -> [VName] + -> [VName] + -> m (AST.Exp (Lore m)) +doLoopMapAccumL width innerfun accexps arrexps mapout_arrs = do + (merge, i, loopbody) <- + doLoopMapAccumL' width innerfun accexps arrexps mapout_arrs + return $ DoLoop [] merge (ForLoop i Int32 width []) loopbody + +doLoopMapAccumL' :: (LocalScope (Lore m) m, MonadBinder m, + Bindable (Lore m), BinderOps (Lore m), + LetAttr (Lore m) ~ Type, + CanBeAliased (Op (Lore m))) => + SubExp + -> AST.Lambda (Aliases (Lore m)) + -> [SubExp] + -> [VName] + -> [VName] + -> m ([(AST.FParam (Lore m), SubExp)], VName, AST.Body (Lore m)) +doLoopMapAccumL' width innerfun accexps arrexps mapout_arrs = do + i <- newVName "i" + -- for the MAP part + let acc_num = length accexps + let res_tps = lambdaReturnType innerfun + let map_arr_tps = drop acc_num res_tps + let res_ts = [ arrayOf t (Shape [width]) NoUniqueness + | t <- map_arr_tps ] + let accts = map paramType $ fst $ splitAt acc_num $ lambdaParams innerfun + outarrs <- mapM (newIdent "mapaccum_outarr") res_ts + -- for the REDUCE part + (acc, initacc, inarrs) <- newFold "mapaccum" (zip accexps accts) arrexps + let consumed = consumedInBody $ lambdaBody innerfun + withUniqueness p | identName p `S.member` consumed = (p, Unique) + | p `elem` outarrs = (p, Unique) + | otherwise = (p, Nonunique) + merge = loopMerge' (map withUniqueness $ inarrs++acc++outarrs) + (map Var arrexps++initacc++map Var mapout_arrs) + loopbody <- runBodyBinder $ localScope (scopeOfFParams $ map fst merge) $ do + accxis<- bindLambda (removeLambdaAliases innerfun) . + (map (BasicOp . SubExp . Var . identName) acc ++) =<< + index (map identName inarrs) (Var i) + let (acc', xis) = splitAt acc_num accxis + dests <- letwith (map identName outarrs) (pexp (Var i)) $ + map (BasicOp . SubExp) xis + return $ resultBody (map (Var . identName) inarrs ++ acc' ++ map Var dests) + return (merge, i, loopbody) diff --git a/src/Futhark/Transform/Rename.hs b/src/Futhark/Transform/Rename.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Transform/Rename.hs @@ -0,0 +1,324 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +-- | This module provides facilities for transforming Futhark programs such +-- that names are unique, via the 'renameProg' function. +-- Additionally, the module also supports adding integral \"tags\" to +-- names (incarnated as the 'ID' type), in order to support more +-- efficient comparisons and renamings. This is done by 'tagProg'. +-- The intent is that you call 'tagProg' once at some early stage, +-- then use 'renameProg' from then on. Functions are also provided +-- for removing the tags again from expressions, patterns and typs. +module Futhark.Transform.Rename + ( + -- * Renaming programs + renameProg + -- * Renaming parts of a program. + -- + -- These all require execution in a 'MonadFreshNames' environment. + , renameExp + , renameStm + , renameBody + , renameLambda + , renameFun + , renamePattern + -- * Renaming annotations + , RenameM + , substituteRename + , bindingForRename + , renamingStms + , Rename (..) + , Renameable + ) + where + +import Control.Monad.State +import Control.Monad.Reader +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import Data.Semigroup ((<>)) + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Attributes.Patterns +import Futhark.FreshNames +import Futhark.MonadFreshNames (MonadFreshNames(..), modifyNameSource) +import Futhark.Transform.Substitute + +runRenamer :: RenameM a -> VNameSource -> (a, VNameSource) +runRenamer m src = runReader (runStateT m src) env + where env = RenameEnv M.empty newName + +-- | Rename variables such that each is unique. The semantics of the +-- program are unaffected, under the assumption that the program was +-- correct to begin with. In particular, the renaming may make an +-- invalid program valid. +renameProg :: (Renameable lore, MonadFreshNames m) => + Prog lore -> m (Prog lore) +renameProg prog = modifyNameSource $ + runRenamer $ Prog <$> mapM rename (progFunctions prog) + +-- | Rename bound variables such that each is unique. The semantics +-- of the expression is unaffected, under the assumption that the +-- expression was correct to begin with. Any free variables are left +-- untouched. +renameExp :: (Renameable lore, MonadFreshNames m) => + Exp lore -> m (Exp lore) +renameExp = modifyNameSource . runRenamer . rename + +-- | Rename bound variables such that each is unique. The semantics +-- of the binding is unaffected, under the assumption that the +-- 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) +renameStm binding = do + e <- renameExp $ stmExp binding + return binding { stmExp = e } + +-- | Rename bound variables such that each is unique. The semantics +-- 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) +renameBody = modifyNameSource . runRenamer . rename + +-- | Rename bound variables such that each is unique. The semantics +-- of the lambda is unaffected, under the assumption that the body was +-- 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) +renameLambda = modifyNameSource . runRenamer . rename + +-- | Rename bound variables such that each is unique. The semantics +-- of the function is unaffected, under the assumption that the body +-- was correct to begin with. Any free variables are left untouched. +-- Note in particular that the parameters of the lambda are renamed. +renameFun :: (Renameable lore, MonadFreshNames m) => + FunDef lore -> m (FunDef lore) +renameFun = modifyNameSource . runRenamer . rename + +-- | Produce an equivalent pattern but with each pattern element given +-- a new name. +renamePattern :: (Rename attr, MonadFreshNames m) => + PatternT attr -> m (PatternT attr) +renamePattern = modifyNameSource . runRenamer . rename' + where rename' pat = bind (patternNames pat) $ rename pat + +data RenameEnv = RenameEnv { + envNameMap :: M.Map VName VName + , envNameFn :: VNameSource -> VName -> (VName, VNameSource) + } + +-- | The monad in which renaming is performed. +type RenameM = StateT VNameSource (Reader RenameEnv) + +-- | Produce a map of the substitutions that should be performed by +-- the renamer. +renamerSubstitutions :: RenameM Substitutions +renamerSubstitutions = lift $ asks envNameMap + +-- | Perform a renaming using the 'Substitute' instance. This only +-- works if the argument does not itself perform any name binding, but +-- it can save on boilerplate for simple types. +substituteRename :: Substitute a => a -> RenameM a +substituteRename x = do + substs <- renamerSubstitutions + return $ substituteNames substs x + +-- | Return a fresh, unique name. The @VName@ is prepended to the +-- name. +new :: VName -> RenameM VName +new k = do (k', src') <- asks envNameFn <*> get <*> pure k + put src' + return k' + +-- | Members of class 'Rename' can be uniquely renamed. +class Rename a where + -- | Rename the given value such that it does not contain shadowing, + -- and has incorporated any substitutions present in the 'RenameM' + -- environment. + rename :: a -> RenameM a + +instance Rename VName where + rename name = fromMaybe name <$> + asks (M.lookup name . envNameMap) + +instance Rename a => Rename [a] where + rename = mapM rename + +instance (Rename a, Rename b) => Rename (a,b) where + rename (a,b) = (,) <$> rename a <*> rename b + +instance (Rename a, Rename b, Rename c) => Rename (a,b,c) where + rename (a,b,c) = do + a' <- rename a + b' <- rename b + c' <- rename c + return (a',b',c') + +instance Rename a => Rename (Maybe a) where + rename = maybe (return Nothing) (fmap Just . rename) + +instance Rename Bool where + rename = return + +instance Rename Ident where + rename (Ident name tp) = do + name' <- rename name + tp' <- rename tp + return $ Ident name' tp' + +-- | Create a bunch of new names and bind them for substitution. +bindingForRename :: [VName] -> RenameM a -> RenameM a +bindingForRename = bind + +bind :: [VName] -> RenameM a -> RenameM a +bind vars body = do + vars' <- mapM new vars + -- This works because map union prefers elements from left + -- operand. + local (bind' vars') body + where bind' vars' env = env { envNameMap = M.fromList (zip vars vars') + `M.union` envNameMap env } + +-- | 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 stms m = descend mempty stms + where descend stms' rem_stms = case stmsHead rem_stms of + Nothing -> m stms' + Just (stm, rem_stms') -> bind (patternNames $ stmPattern stm) $ do + stm' <- rename stm + descend (stms' <> oneStm stm') rem_stms' + +instance Renameable lore => Rename (FunDef lore) where + rename (FunDef entry fname ret params body) = + bind (map paramName params) $ do + params' <- mapM rename params + body' <- rename body + ret' <- rename ret + return $ FunDef entry fname ret' params' body' + +instance Rename SubExp where + rename (Var v) = Var <$> rename v + rename (Constant v) = return $ Constant v + +instance Rename attr => Rename (ParamT attr) where + rename (Param name attr) = Param <$> rename name <*> rename attr + +instance Rename attr => Rename (PatternT attr) where + rename (Pattern context values) = Pattern <$> rename context <*> rename values + +instance Rename attr => Rename (PatElemT attr) where + rename (PatElem ident attr) = PatElem <$> rename ident <*> rename attr + +instance Rename Certificates where + rename (Certificates cs) = Certificates <$> rename cs + +instance Rename attr => Rename (StmAux attr) where + rename (StmAux cs attr) = + StmAux <$> rename cs <*> rename attr + +instance Renameable lore => Rename (Body lore) where + rename (Body attr stms res) = do + attr' <- rename attr + renamingStms stms $ \stms' -> + Body attr' stms' <$> rename res + +instance Renameable lore => Rename (Stm lore) where + rename (Let pat elore e) = Let <$> rename pat <*> rename elore <*> rename e + +instance Renameable lore => Rename (Exp lore) where + rename (DoLoop ctx val form loopbody) = do + let (ctxparams, ctxinit) = unzip ctx + (valparams, valinit) = unzip val + ctxinit' <- mapM rename ctxinit + valinit' <- mapM rename valinit + case form of + ForLoop loopvar it boundexp loop_vars -> do + let (loop_params, loop_arrs) = unzip loop_vars + boundexp' <- rename boundexp + loop_arrs' <- rename loop_arrs + bind (map paramName (ctxparams++valparams) ++ + map paramName loop_params) $ do + ctxparams' <- mapM rename ctxparams + valparams' <- mapM rename valparams + loop_params' <- mapM rename loop_params + bind [loopvar] $ do + loopvar' <- rename loopvar + loopbody' <- rename loopbody + return $ DoLoop + (zip ctxparams' ctxinit') (zip valparams' valinit') + (ForLoop loopvar' it boundexp' $ + zip loop_params' loop_arrs') loopbody' + WhileLoop cond -> + bind (map paramName $ ctxparams++valparams) $ do + ctxparams' <- mapM rename ctxparams + valparams' <- mapM rename valparams + loopbody' <- rename loopbody + cond' <- rename cond + return $ DoLoop + (zip ctxparams' ctxinit') (zip valparams' valinit') + (WhileLoop cond') loopbody' + rename e = mapExpM mapper e + where mapper = Mapper { + mapOnBody = const rename + , mapOnSubExp = rename + , mapOnVName = rename + , mapOnCertificates = rename + , mapOnRetType = rename + , mapOnBranchType = rename + , mapOnFParam = rename + , mapOnLParam = rename + , mapOnOp = rename + } + +instance Rename shape => + Rename (TypeBase shape u) where + rename (Array et size u) = do + size' <- rename size + return $ Array et size' u + rename (Prim et) = return $ Prim et + rename (Mem e space) = Mem <$> rename e <*> pure space + +instance Renameable lore => Rename (Lambda lore) where + rename (Lambda params body ret) = + bind (map paramName params) $ do + params' <- mapM rename params + body' <- rename body + ret' <- mapM rename ret + return $ Lambda params' body' ret' + +instance Rename Names where + rename = fmap S.fromList . mapM rename . S.toList + +instance Rename Rank where + rename = return + +instance Rename d => Rename (ShapeBase d) where + rename (Shape l) = Shape <$> mapM rename l + +instance Rename ExtSize where + rename (Free se) = Free <$> rename se + rename (Ext x) = return $ Ext x + +instance Rename () where + rename = return + +instance Rename d => Rename (DimIndex d) where + 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 (LetAttr lore), + Rename (ExpAttr lore), + Rename (BodyAttr lore), + Rename (FParamAttr lore), + Rename (LParamAttr lore), + Rename (RetType lore), + Rename (BranchType lore), + Rename (Op lore)) diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Transform/Substitute.hs @@ -0,0 +1,190 @@ +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +-- | +-- +-- This module contains facilities for replacing variable names in +-- syntactic constructs. +module Futhark.Transform.Substitute + (Substitutions, + Substitute(..), + Substitutable) + where + +import Control.Monad.Identity +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Traversals +import Futhark.Representation.AST.Attributes.Scope +import Futhark.Analysis.PrimExp + +-- | The substitutions to be made are given by a mapping from names to +-- names. +type Substitutions = M.Map VName VName + +-- | A type that is an instance of this class supports substitution of +-- any names contained within. +class Substitute a where + -- | @substituteNames m e@ replaces the variable names in @e@ with + -- new names, based on the mapping in @m@. It is assumed that all + -- names in @e@ are unique, i.e. there is no shadowing. + substituteNames :: M.Map VName VName -> a -> a + +instance Substitute a => Substitute [a] where + substituteNames substs = map $ substituteNames substs + +instance Substitute (Stm lore) => Substitute (Stms lore) where + substituteNames substs = fmap $ substituteNames substs + +instance (Substitute a, Substitute b) => Substitute (a,b) where + substituteNames substs (x,y) = + (substituteNames substs x, substituteNames substs y) + +instance (Substitute a, Substitute b, Substitute c) => Substitute (a,b,c) where + substituteNames substs (x,y,z) = + (substituteNames substs x, + substituteNames substs y, + substituteNames substs z) + +instance (Substitute a, Substitute b, Substitute c, Substitute d) => Substitute (a,b,c,d) where + substituteNames substs (x,y,z,u) = + (substituteNames substs x, + substituteNames substs y, + substituteNames substs z, + substituteNames substs u) + +instance Substitute a => Substitute (Maybe a) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitute Bool where + substituteNames = flip const + +instance Substitute VName where + substituteNames substs k = M.findWithDefault k k substs + +instance Substitute SubExp where + substituteNames substs (Var v) = Var $ substituteNames substs v + substituteNames _ (Constant v) = Constant v + +instance Substitutable lore => Substitute (Exp lore) where + substituteNames substs = mapExp $ replace substs + +instance Substitute attr => Substitute (PatElemT attr) where + substituteNames substs (PatElem ident attr) = + PatElem (substituteNames substs ident) (substituteNames substs attr) + +instance Substitute attr => Substitute (StmAux attr) where + substituteNames substs (StmAux cs attr) = + StmAux (substituteNames substs cs) (substituteNames substs attr) + +instance Substitute attr => Substitute (ParamT attr) where + substituteNames substs (Param name attr) = + Param + (substituteNames substs name) + (substituteNames substs attr) + +instance Substitute attr => Substitute (PatternT attr) where + substituteNames substs (Pattern context values) = + Pattern (substituteNames substs context) (substituteNames substs values) + +instance Substitute Certificates where + substituteNames substs (Certificates cs) = + Certificates $ substituteNames substs cs + +instance Substitutable lore => Substitute (Stm lore) where + substituteNames substs (Let pat annot e) = + Let + (substituteNames substs pat) + (substituteNames substs annot) + (substituteNames substs e) + +instance Substitutable lore => Substitute (Body lore) where + substituteNames substs (Body attr stms res) = + Body + (substituteNames substs attr) + (substituteNames substs stms) + (substituteNames substs res) + +replace :: Substitutable lore => M.Map VName VName -> Mapper lore lore Identity +replace substs = Mapper { + mapOnVName = return . substituteNames substs + , mapOnSubExp = return . substituteNames substs + , mapOnBody = const $ return . substituteNames substs + , mapOnCertificates = return . substituteNames substs + , mapOnRetType = return . substituteNames substs + , mapOnBranchType = return . substituteNames substs + , mapOnFParam = return . substituteNames substs + , mapOnLParam = return . substituteNames substs + , mapOnOp = return . substituteNames substs + } + +instance Substitute Rank where + substituteNames _ = id + +instance Substitute () where + substituteNames _ = id + +instance Substitute d => Substitute (ShapeBase d) where + substituteNames substs (Shape es) = + Shape $ map (substituteNames substs) es + +instance Substitute d => Substitute (Ext d) where + substituteNames substs (Free x) = Free $ substituteNames substs x + substituteNames _ (Ext x) = Ext x + +instance Substitute Names where + substituteNames = S.map . substituteNames + +instance Substitute shape => Substitute (TypeBase shape u) where + substituteNames _ (Prim et) = Prim et + substituteNames substs (Array et sz u) = + Array et (substituteNames substs sz) u + substituteNames substs (Mem sz space) = + Mem (substituteNames substs sz) space + +instance Substitutable lore => Substitute (Lambda lore) where + substituteNames substs (Lambda params body rettype) = + Lambda + (substituteNames substs params) + (substituteNames substs body) + (map (substituteNames substs) rettype) + +instance Substitute Ident where + substituteNames substs v = + v { identName = substituteNames substs $ identName v + , identType = substituteNames substs $ identType v + } + +instance Substitute d => Substitute (DimChange d) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitute d => Substitute (DimIndex d) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitute v => Substitute (PrimExp v) where + substituteNames substs = fmap $ substituteNames substs + +instance Substitutable lore => Substitute (NameInfo lore) where + substituteNames subst (LetInfo attr) = + LetInfo $ substituteNames subst attr + substituteNames subst (FParamInfo attr) = + FParamInfo $ substituteNames subst attr + substituteNames subst (LParamInfo attr) = + LParamInfo $ substituteNames subst attr + substituteNames _ (IndexInfo it) = + IndexInfo it + +-- | Lores in which all annotations support name +-- substitution. +type Substitutable lore = (Annotations lore, + Substitute (ExpAttr lore), + Substitute (BodyAttr lore), + Substitute (LetAttr lore), + Substitute (FParamAttr lore), + Substitute (LParamAttr lore), + Substitute (RetType lore), + Substitute (BranchType lore), + Substitute (Op lore)) diff --git a/src/Futhark/TypeCheck.hs b/src/Futhark/TypeCheck.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/TypeCheck.hs @@ -0,0 +1,1098 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, TypeFamilies, ScopedTypeVariables #-} +-- | The type checker checks whether the program is type-consistent. +module Futhark.TypeCheck + ( -- * Interface + checkProg + , TypeError (..) + , ErrorCase (..) + + -- * Extensionality + , TypeM + , bad + , context + , message + , Checkable (..) + , lookupVar + , lookupAliases + , Occurences + , UsageMap + , usageMap + , collectOccurences + , subCheck + + -- * Checkers + , require + , requireI + , requirePrimExp + , checkSubExp + , checkExp + , checkStms + , checkStm + , checkType + , checkExtType + , matchExtPattern + , matchExtReturnType + , matchExtBranchType + , argType + , argAliases + , noArgAliases + , checkArg + , checkSOACArrayArgs + , checkLambda + , checkFun' + , checkLambdaParams + , checkBody + , checkLambdaBody + , consume + , consumeOnlyParams + , binding + ) + where + +import Control.Parallel.Strategies +import Control.Monad.Reader +import Control.Monad.Writer +import Control.Monad.State +import Control.Monad.RWS.Strict +import Data.List +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.Maybe +import qualified Data.Semigroup as Sem + +import Futhark.Analysis.PrimExp +import Futhark.Construct (instantiateShapes) +import Futhark.Representation.Aliases +import Futhark.Analysis.Alias +import Futhark.Util +import Futhark.Util.Pretty (Pretty, prettyDoc, indent, ppr, text, (<+>), align) + +-- | Information about an error during type checking. The 'Show' +-- instance for this type produces a human-readable description. +data ErrorCase lore = + TypeError String + | UnexpectedType (Exp lore) Type [Type] + | ReturnTypeError Name [ExtType] [ExtType] + | DupDefinitionError Name + | DupParamError Name VName + | DupPatternError VName + | InvalidPatternError (Pattern (Aliases lore)) [ExtType] (Maybe String) + | UnknownVariableError VName + | UnknownFunctionError Name + | ParameterMismatch (Maybe Name) [Type] [Type] + | SlicingError Int Int + | BadAnnotation String Type Type + | ReturnAliased Name VName + | UniqueReturnAliased Name + | NotAnArray VName Type + | PermutationError [Int] Int (Maybe VName) + +instance Checkable lore => Show (ErrorCase lore) where + show (TypeError msg) = + "Type error:\n" ++ msg + show (UnexpectedType e _ []) = + "Type of expression\n" ++ + prettyDoc 160 (indent 2 $ ppr e) ++ + "\ncannot have any type - possibly a bug in the type checker." + show (UnexpectedType e t ts) = + "Type of expression\n" ++ + prettyDoc 160 (indent 2 $ ppr e) ++ + "\nmust be one of " ++ intercalate ", " (map pretty ts) ++ ", but is " ++ + pretty t ++ "." + show (ReturnTypeError fname rettype bodytype) = + "Declaration of function " ++ nameToString fname ++ + " declares return type\n " ++ prettyTuple rettype ++ + "\nBut body has type\n " ++ prettyTuple bodytype + show (DupDefinitionError name) = + "Duplicate definition of function " ++ nameToString name ++ "" + show (DupParamError funname paramname) = + "Parameter " ++ pretty paramname ++ + " mentioned multiple times in argument list of function " ++ + nameToString funname ++ "." + show (DupPatternError name) = + "Variable " ++ pretty name ++ " bound twice in pattern." + show (InvalidPatternError pat t desc) = + "Pattern " ++ pretty pat ++ + " cannot match value of type " ++ prettyTuple t ++ end + where end = case desc of Nothing -> "." + Just desc' -> ":\n" ++ desc' + show (UnknownVariableError name) = + "Use of unknown variable " ++ pretty name ++ "." + show (UnknownFunctionError fname) = + "Call of unknown function " ++ nameToString fname ++ "." + show (ParameterMismatch fname expected got) = + "In call of " ++ fname' ++ ":\n" ++ + "expecting " ++ show nexpected ++ " argument(s) of type(s) " ++ + expected' ++ ", but got " ++ show ngot ++ + " arguments of types " ++ intercalate ", " (map pretty got) ++ "." + where (nexpected, expected') = + (length expected, intercalate ", " $ map pretty expected) + ngot = length got + fname' = maybe "anonymous function" (("function "++) . nameToString) fname + show (SlicingError dims got) = + show got ++ " indices given, but type of indexee has " ++ show dims ++ " dimension(s)." + show (BadAnnotation desc expected got) = + "Annotation of \"" ++ desc ++ "\" type of expression is " ++ pretty expected ++ + ", but derived to be " ++ pretty got ++ "." + show (ReturnAliased fname name) = + "Unique return value of function " ++ nameToString fname ++ + " is aliased to " ++ pretty name ++ ", which is not consumed." + show (UniqueReturnAliased fname) = + "A unique tuple element of return value of function " ++ + nameToString fname ++ " is aliased to some other tuple component." + show (NotAnArray e t) = + "The expression " ++ pretty e ++ + " is expected to be an array, but is " ++ pretty t ++ "." + show (PermutationError perm rank name) = + "The permutation (" ++ intercalate ", " (map show perm) ++ + ") is not valid for array " ++ name' ++ "of rank " ++ show rank ++ "." + where name' = maybe "" ((++" ") . pretty) name + +-- | A type error. +data TypeError lore = Error [String] (ErrorCase lore) + +instance Checkable lore => Show (TypeError lore) where + show (Error [] err) = + show err + show (Error msgs err) = + intercalate "\n" msgs ++ "\n" ++ show err + +-- | A tuple of a return type and a list of parameters, possibly +-- named. +type FunBinding lore = ([RetType (Aliases lore)], [FParam (Aliases lore)]) + +type VarBinding lore = NameInfo (Aliases lore) + +data Usage = Consumed + | Observed + deriving (Eq, Ord, Show) + +data Occurence = Occurence { observed :: Names + , consumed :: Names + } + deriving (Eq, Show) + +observation :: Names -> Occurence +observation = flip Occurence S.empty + +consumption :: Names -> Occurence +consumption = Occurence S.empty + +nullOccurence :: Occurence -> Bool +nullOccurence occ = S.null (observed occ) && S.null (consumed occ) + +type Occurences = [Occurence] + +type UsageMap = M.Map VName [Usage] + +usageMap :: Occurences -> UsageMap +usageMap = foldl comb M.empty + where comb m (Occurence obs cons) = + let m' = S.foldl' (ins Observed) m obs + in S.foldl' (ins Consumed) m' cons + ins v m k = M.insertWith (++) k [v] m + +allConsumed :: Occurences -> Names +allConsumed = S.unions . map consumed + +seqOccurences :: Occurences -> Occurences -> Occurences +seqOccurences occurs1 occurs2 = + filter (not . nullOccurence) (map filt occurs1) ++ occurs2 + where filt occ = + occ { observed = observed occ `S.difference` postcons } + postcons = allConsumed occurs2 + +altOccurences :: Occurences -> Occurences -> Occurences +altOccurences occurs1 occurs2 = + filter (not . nullOccurence) (map filt occurs1) ++ occurs2 + where filt occ = + occ { consumed = consumed occ `S.difference` postcons + , observed = observed occ `S.difference` postcons } + postcons = allConsumed occurs2 + +unOccur :: Names -> Occurences -> Occurences +unOccur to_be_removed = filter (not . nullOccurence) . map unOccur' + where unOccur' occ = + occ { observed = observed occ `S.difference` to_be_removed + , consumed = consumed occ `S.difference` to_be_removed + } + +-- | The 'Consumption' data structure is used to keep track of which +-- variables have been consumed, as well as whether a violation has been detected. +data Consumption = ConsumptionError String + | Consumption Occurences + deriving (Show) + +instance Sem.Semigroup Consumption where + ConsumptionError e <> _ = ConsumptionError e + _ <> ConsumptionError e = ConsumptionError e + Consumption o1 <> Consumption o2 + | v:_ <- S.toList $ consumed_in_o1 `S.intersection` used_in_o2 = + ConsumptionError $ "Variable " <> pretty v <> " referenced after being consumed." + | otherwise = + Consumption $ o1 `seqOccurences` o2 + where consumed_in_o1 = mconcat $ map consumed o1 + used_in_o2 = mconcat $ map consumed o2 <> map observed o2 + +instance Monoid Consumption where + mempty = Consumption mempty + mappend = (Sem.<>) + +-- | The environment contains a variable table and a function table. +-- Type checking happens with access to this environment. The +-- 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) + , envContext :: [String] + } + +-- | The type checker runs in this monad. +newtype TypeM lore a = TypeM (RWST + (Env lore) -- Reader + Consumption -- Writer + Names -- State + (Either (TypeError lore)) -- Inner monad + a) + deriving (Monad, Functor, Applicative, + MonadReader (Env lore), + MonadWriter Consumption, + MonadState Names) + +instance Checkable lore => + HasScope (Aliases lore) (TypeM lore) where + lookupType = fmap typeOf . lookupVar + askScope = asks $ M.fromList . mapMaybe varType . M.toList . envVtable + where varType (name, attr) = Just (name, attr) + +runTypeM :: Env lore -> TypeM lore a + -> Either (TypeError lore) (a, Consumption) +runTypeM env (TypeM m) = evalRWST m env mempty + +bad :: ErrorCase lore -> TypeM lore a +bad e = do + messages <- asks envContext + TypeM $ lift $ Left $ Error (reverse messages) e + +-- | Add information about what is being type-checked to the current +-- context. Liberal use of this combinator makes it easier to track +-- type errors, as the strings are added to type errors signalled via +-- 'bad'. +context :: String + -> TypeM lore a + -> TypeM lore a +context s = local $ \env -> env { envContext = s : envContext env} + +message :: Pretty a => + String -> a -> String +message s x = prettyDoc 80 $ + text s <+> align (ppr x) + +-- | Mark a name as bound. If the name has been bound previously in +-- the program, report a type error. +bound :: VName -> TypeM lore () +bound name = do already_seen <- gets $ S.member name + when already_seen $ + bad $ TypeError $ "Name " ++ pretty name ++ " bound twice" + modify $ S.insert name + +occur :: Occurences -> TypeM lore () +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 => + VName -> TypeM lore () +observe name = do + attr <- lookupVar name + unless (primType $ typeOf attr) $ + occur [observation $ S.insert name $ aliases attr] + +-- | Proclaim that we have written to the given variable. +consume :: Names -> TypeM lore () +consume als = occur [consumption als] + +collectOccurences :: TypeM lore a -> TypeM lore (a, Occurences) +collectOccurences m = pass $ do + (x, c) <- listen m + o <- checkConsumption c + return ((x, o), const mempty) + +checkConsumption :: Consumption -> TypeM lore Occurences +checkConsumption (ConsumptionError e) = bad $ TypeError e +checkConsumption (Consumption os) = return os + +alternative :: TypeM lore a -> TypeM lore b -> TypeM lore (a,b) +alternative m1 m2 = pass $ do + (x, c1) <- listen m1 + (y, c2) <- listen m2 + os1 <- checkConsumption c1 + os2 <- checkConsumption c2 + let usage = Consumption $ os1 `altOccurences` os2 + return ((x, y), const usage) + +-- | Permit consumption of only the specified names. If one of these +-- 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 consumable m = do + (x, os) <- collectOccurences m + tell . Consumption =<< mapM inspect os + return x + where inspect o = do + new_consumed <- mconcat <$> mapM wasConsumed (S.toList $ consumed o) + return o { consumed = new_consumed } + wasConsumed v + | Just als <- lookup v consumable = return als + | otherwise = + bad $ TypeError $ + unlines [pretty v ++ " was invalidly consumed.", + what ++ " can be consumed here."] + what | null consumable = "Nothing" + | otherwise = "Only " ++ intercalate ", " (map (pretty . fst) consumable) + +-- | Given the immediate aliases, compute the full transitive alias +-- set (including the immediate aliases). +expandAliases :: Names -> Env lore -> Names +expandAliases names env = names `S.union` aliasesOfAliases + where aliasesOfAliases = mconcat . map look . S.toList $ names + look k = case M.lookup k $ envVtable env of + Just (LetInfo (als, _)) -> unNames als + _ -> mempty + +binding :: Checkable lore => + Scope (Aliases lore) + -> TypeM lore a + -> TypeM lore a +binding bnds = check . local (`bindVars` bnds) + where bindVars = M.foldlWithKey' bindVar + boundnames = M.keys bnds + boundnameset = S.fromList boundnames + + bindVar env name (LetInfo (Names' als, attr)) = + let als' | primType (typeOf attr) = mempty + | otherwise = expandAliases als env + inedges = S.toList als' + update (LetInfo (Names' thesenames, thisattr)) = + LetInfo (Names' $ S.insert name thesenames, thisattr) + update b = b + in env { envVtable = + M.insert name (LetInfo (Names' als', attr)) $ + adjustSeveral update inedges $ + envVtable env + } + bindVar env name attr = + env { envVtable = M.insert name attr $ envVtable env } + + adjustSeveral f = flip $ foldl $ flip $ M.adjust f + + -- Check whether the bound variables have been used correctly + -- within their scope. + check m = do + mapM_ bound $ M.keys bnds + (a, os) <- collectOccurences m + tell $ Consumption $ unOccur boundnameset os + return a + +lookupVar :: VName -> TypeM lore (NameInfo (Aliases lore)) +lookupVar name = do + bnd <- asks $ M.lookup name . envVtable + case bnd of + Nothing -> bad $ UnknownVariableError name + Just attr -> return attr + +lookupAliases :: Checkable lore => VName -> TypeM lore Names +lookupAliases name = do + info <- lookupVar name + return $ if primType $ typeOf info + then mempty + else S.insert name $ aliases info + +aliases :: NameInfo (Aliases lore) -> Names +aliases (LetInfo (als, _)) = unNames als +aliases _ = mempty + +subExpAliasesM :: Checkable lore => SubExp -> TypeM lore Names +subExpAliasesM Constant{} = return mempty +subExpAliasesM (Var v) = lookupAliases v + +lookupFun :: Checkable lore => + Name + -> [SubExp] + -> TypeM lore ([RetType lore], [DeclType]) +lookupFun fname args = do + bnd <- asks $ M.lookup fname . envFtable + case bnd of + Nothing -> bad $ UnknownFunctionError fname + Just (ftype, params) -> do + argts <- mapM subExpType args + case applyRetType ftype params $ zip args argts of + Nothing -> + bad $ ParameterMismatch (Just fname) (map paramType params) argts + Just rt -> + return (rt, map paramDeclType params) + +-- | @checkAnnotation loc s t1 t2@ checks if @t2@ is equal to +-- @t1@. If not, a 'BadAnnotation' is raised. +checkAnnotation :: String -> Type -> Type + -> TypeM lore () +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 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 ts ident = require ts $ Var ident + +checkArrIdent :: Checkable lore => + VName -> TypeM lore Type +checkArrIdent v = do + t <- lookupType v + case t of + Array{} -> return t + _ -> bad $ NotAnArray v t + +-- | Type check a program containing arbitrary type information, +-- yielding either a type error or a program with complete type +-- information. +checkProg :: Checkable lore => + Prog lore -> Either (TypeError lore) () +checkProg prog = do + let typeenv = Env { envVtable = M.empty + , envFtable = mempty + , envContext = [] + } + let onFunction ftable fun = + fmap fst $ runTypeM typeenv $ + local (\env -> env { envFtable = ftable }) $ + checkFun fun + (ftable, _) <- runTypeM typeenv buildFtable + sequence_ $ parMap rpar (onFunction ftable) $ progFunctions prog' + where + prog' = aliasAnalysis prog + buildFtable = do table <- initialFtable prog' + foldM expand table $ progFunctions prog' + expand ftable (FunDef _ name ret params _) + | M.member name ftable = + bad $ DupDefinitionError name + | otherwise = + return $ M.insert name (ret,params) ftable + +-- The prog argument is just to disambiguate the lore. +initialFtable :: Checkable lore => + Prog (Aliases lore) -> TypeM lore (M.Map Name (FunBinding lore)) +initialFtable _ = fmap M.fromList $ mapM addBuiltin $ M.toList builtInFunctions + where addBuiltin (fname, (t, ts)) = do + ps <- mapM (primFParam name) ts + return (fname, ([primRetType t], ps)) + name = VName (nameFromString "x") 0 + +checkFun :: Checkable lore => + FunDef (Aliases lore) -> TypeM lore () +checkFun (FunDef _ fname rettype params body) = + context ("In function " ++ nameToString fname) $ + checkFun' (fname, + retTypeValues rettype, + funParamsToNameInfos params, + body) consumable $ do + checkFunParams params + checkRetType rettype + checkFunBody rettype body + where consumable = [ (paramName param, mempty) + | param <- params + , unique $ paramDeclType param + ] + +funParamsToNameInfos :: [FParam lore] + -> [(VName, NameInfo (Aliases lore))] +funParamsToNameInfos = map nameTypeAndLore + where nameTypeAndLore fparam = (paramName fparam, + FParamInfo $ paramAttr fparam) + +checkFunParams :: Checkable lore => + [FParam lore] -> TypeM lore () +checkFunParams = mapM_ $ \param -> + context ("In function parameter " ++ pretty param) $ + checkFParamLore (paramName param) (paramAttr param) + +checkLambdaParams :: Checkable lore => + [LParam lore] -> TypeM lore () +checkLambdaParams = mapM_ $ \param -> + context ("In lambda parameter " ++ pretty param) $ + checkLParamLore (paramName param) (paramAttr param) + +checkFun' :: Checkable lore => + (Name, + [DeclExtType], + [(VName, NameInfo (Aliases lore))], + BodyT (Aliases lore)) + -> [(VName, Names)] + -> TypeM lore () + -> TypeM lore () +checkFun' (fname, rettype, params, body) consumable check = do + checkNoDuplicateParams + binding (M.fromList params) $ + consumeOnlyParams consumable $ do + check + checkReturnAlias $ bodyAliases body + where param_names = map fst params + + checkNoDuplicateParams = foldM_ expand [] param_names + + expand seen pname + | Just _ <- find (==pname) seen = + bad $ DupParamError fname pname + | otherwise = + return $ pname : seen + + -- | Check that unique return values do not alias a + -- non-consumed parameter. + checkReturnAlias = + foldM_ checkReturnAlias' S.empty . returnAliasing rettype + + checkReturnAlias' seen (Unique, names) + | any (`S.member` S.map snd seen) $ S.toList names = + bad $ UniqueReturnAliased fname + | otherwise = do + consume names + return $ seen `S.union` tag Unique names + checkReturnAlias' seen (Nonunique, names) + | any (`S.member` seen) $ S.toList $ tag Unique names = + bad $ UniqueReturnAliased fname + | otherwise = return $ seen `S.union` tag Nonunique names + + tag u = S.map $ \name -> (u, name) + + returnAliasing expected got = + reverse $ + zip (reverse (map uniqueness expected) ++ repeat Nonunique) $ + reverse got + +subCheck :: forall lore newlore a. + (Checkable newlore, + RetType lore ~ RetType newlore, + LetAttr lore ~ LetAttr newlore, + FParamAttr lore ~ FParamAttr newlore, + LParamAttr lore ~ LParamAttr newlore) => + TypeM newlore a -> + TypeM lore a +subCheck m = do + typeenv <- asks newEnv + case runTypeM typeenv m of + Left err -> bad $ TypeError $ show err + Right (x, cons) -> tell cons >> return x + where newEnv :: Env lore -> Env newlore + newEnv (Env vtable ftable ctx) = + Env (M.map coerceVar vtable) ftable ctx + coerceVar (LetInfo x) = LetInfo x + coerceVar (FParamInfo x) = FParamInfo x + coerceVar (LParamInfo x) = LParamInfo x + coerceVar (IndexInfo it) = IndexInfo it + +checkSubExp :: Checkable lore => SubExp -> TypeM lore Type +checkSubExp (Constant val) = + return $ Prim $ primValueType val +checkSubExp (Var ident) = context ("In subexp " ++ pretty ident) $ do + observe ident + lookupType ident + +checkStms :: Checkable lore => + Stms (Aliases lore) -> TypeM lore a + -> TypeM lore a +checkStms origbnds m = delve $ stmsToList origbnds + where delve (stm@(Let pat _ e):bnds) = do + context ("In expression of statement " ++ pretty pat) $ + checkExp e + checkStm stm $ + delve bnds + delve [] = + m + +checkResult :: Checkable lore => + Result -> TypeM lore () +checkResult = mapM_ checkSubExp + +checkFunBody :: Checkable lore => + [RetType lore] + -> Body (Aliases lore) + -> TypeM lore () +checkFunBody rt (Body (_,lore) bnds res) = do + checkStms bnds $ do + context "When checking body result" $ checkResult res + context "When matching declared return type to result of body" $ + matchReturnType rt res + checkBodyLore lore + +checkLambdaBody :: Checkable lore => + [Type] -> Body (Aliases lore) -> TypeM lore () +checkLambdaBody ret (Body (_,lore) bnds res) = do + checkStms bnds $ checkLambdaResult ret res + checkBodyLore lore + +checkLambdaResult :: Checkable lore => + [Type] -> Result -> TypeM lore () +checkLambdaResult ts es + | length ts /= length es = + bad $ TypeError $ + "Lambda has return type " ++ prettyTuple ts ++ + " describing " ++ show (length ts) ++ " values, but body returns " ++ + show (length es) ++ " values: " ++ prettyTuple es + | otherwise = forM_ (zip ts es) $ \(t, e) -> do + et <- checkSubExp e + unless (et == t) $ + bad $ TypeError $ + "Subexpression " ++ pretty e ++ " has type " ++ pretty et ++ + " but expected " ++ pretty t + +checkBody :: Checkable lore => + Body (Aliases lore) -> TypeM lore () +checkBody (Body (_,lore) bnds res) = do + checkStms bnds $ checkResult res + checkBodyLore lore + +checkBasicOp :: Checkable lore => + BasicOp (Aliases lore) -> TypeM lore () + +checkBasicOp (SubExp es) = + void $ checkSubExp es + +checkBasicOp (Opaque es) = + void $ checkSubExp es + +checkBasicOp (ArrayLit [] _) = + return () + +checkBasicOp (ArrayLit (e:es') t) = do + let check elemt eleme = do + elemet <- checkSubExp eleme + unless (elemet == elemt) $ + bad $ TypeError $ pretty elemet ++ + " is not of expected type " ++ pretty elemt ++ "." + et <- checkSubExp e + + -- Compare that type with the one given for the array literal. + checkAnnotation "array-element" t et + + mapM_ (check et) es' + +checkBasicOp (UnOp op e) = require [Prim $ unOpType op] e + +checkBasicOp (BinOp op e1 e2) = checkBinOpArgs (binOpType op) e1 e2 + +checkBasicOp (CmpOp op e1 e2) = checkCmpOp op e1 e2 + +checkBasicOp (ConvOp op e) = require [Prim $ fst $ convOpType op] e + +checkBasicOp (Index ident idxes) = do + vt <- lookupType ident + observe ident + when (arrayRank vt /= length idxes) $ + bad $ SlicingError (arrayRank vt) (length idxes) + mapM_ checkDimIndex idxes + +checkBasicOp (Update src idxes se) = do + src_t <- checkArrIdent src + when (arrayRank src_t /= length idxes) $ + bad $ SlicingError (arrayRank src_t) (length idxes) + + se_aliases <- subExpAliasesM se + when (src `S.member` se_aliases) $ + bad $ TypeError "The target of an Update must not alias the value to be written." + + mapM_ checkDimIndex idxes + require [Prim (elemType src_t) `arrayOfShape` Shape (sliceDims idxes)] se + consume =<< lookupAliases src + +checkBasicOp (Iota e x s et) = do + require [Prim int32] e + require [Prim $ IntType et] x + require [Prim $ IntType et] s + +checkBasicOp (Replicate (Shape dims) valexp) = do + mapM_ (require [Prim int32]) dims + void $ checkSubExp valexp + +checkBasicOp (Repeat shapes innershape v) = do + v_t <- lookupType v + mapM_ (mapM_ (require [Prim int32]) . shapeDims) $ innershape : shapes + unless (length shapes == arrayRank v_t) $ + bad $ TypeError "Incorrect number of shapes in repeat." + +checkBasicOp (Scratch _ shape) = + mapM_ checkSubExp shape + +checkBasicOp (Reshape newshape arrexp) = do + rank <- arrayRank <$> checkArrIdent arrexp + mapM_ (require [Prim int32] . newDim) newshape + zipWithM_ (checkDimChange rank) newshape [0..] + where checkDimChange _ (DimNew _) _ = + return () + checkDimChange rank (DimCoercion se) i + | i >= rank = + bad $ TypeError $ + "Asked to coerce dimension " ++ show i ++ " to " ++ pretty se ++ + ", but array " ++ pretty arrexp ++ " has only " ++ pretty rank ++ " dimensions" + | otherwise = + return () + +checkBasicOp (Rearrange perm arr) = do + arrt <- lookupType arr + let rank = arrayRank arrt + when (length perm /= rank || sort perm /= [0..rank-1]) $ + bad $ PermutationError perm rank $ Just arr + +checkBasicOp (Rotate rots arr) = do + arrt <- lookupType arr + let rank = arrayRank arrt + mapM_ (require [Prim int32]) rots + when (length rots /= rank) $ + bad $ TypeError $ "Cannot rotate " ++ show (length rots) ++ + " dimensions of " ++ show rank ++ "-dimensional array." + +checkBasicOp (Concat i arr1exp arr2exps ressize) = do + arr1t <- checkArrIdent arr1exp + arr2ts <- mapM checkArrIdent arr2exps + let success = all (== (dropAt i 1 $ arrayDims arr1t)) $ + map (dropAt i 1 . arrayDims) arr2ts + unless success $ + bad $ TypeError $ + "Types of arguments to concat do not match. Got " ++ + pretty arr1t ++ " and " ++ intercalate ", " (map pretty arr2ts) + require [Prim int32] ressize + +checkBasicOp (Copy e) = + void $ checkArrIdent e + +checkBasicOp (Manifest perm arr) = + checkBasicOp $ Rearrange perm arr -- Basically same thing! + +checkBasicOp (Assert e _ _) = + require [Prim Bool] e + +checkBasicOp (Partition _ flags arrs) = do + flagst <- lookupType flags + unless (rowType flagst == Prim int32) $ + bad $ TypeError $ "Flag array has type " ++ pretty flagst ++ "." + forM_ arrs $ \arr -> do + arrt <- lookupType arr + unless (arrayRank arrt > 0) $ + bad $ TypeError $ + "Array argument " ++ pretty arr ++ + " to partition has type " ++ pretty arrt ++ "." + +checkExp :: Checkable lore => + Exp (Aliases lore) -> TypeM lore () + +checkExp (BasicOp op) = checkBasicOp op + +checkExp (If e1 e2 e3 info) = do + require [Prim Bool] e1 + _ <- checkBody e2 `alternative` checkBody e3 + context "in true branch" $ matchBranchType (ifReturns info) e2 + context "in false branch" $ matchBranchType (ifReturns info) e3 + +checkExp (Apply fname args rettype_annot _) = do + (rettype_derived, paramtypes) <- lookupFun fname $ map fst args + argflows <- mapM (checkArg . fst) args + when (rettype_derived /= rettype_annot) $ + bad $ TypeError $ "Expected apply result type " ++ pretty rettype_derived + ++ " but annotation is " ++ pretty rettype_annot + checkFuncall (Just fname) paramtypes argflows + +checkExp (DoLoop ctxmerge valmerge form loopbody) = do + let merge = ctxmerge ++ valmerge + (mergepat, mergeexps) = unzip merge + mergeargs <- mapM checkArg mergeexps + + binding (scopeOf form) $ do + case form of + ForLoop loopvar it boundexp loopvars -> do + iparam <- primFParam loopvar $ IntType it + let funparams = iparam : mergepat + paramts = map paramDeclType funparams + + forM_ loopvars $ \(p,a) -> do + a_t <- lookupType a + observe a + case peelArray 1 a_t of + Just a_t_r -> do + checkLParamLore (paramName p) $ paramAttr p + unless (a_t_r `subtypeOf` typeOf (paramAttr p)) $ + bad $ TypeError $ "Loop parameter " ++ pretty p ++ + " not valid for element of " ++ pretty a ++ ", which has row type " ++ pretty a_t_r + _ -> bad $ TypeError $ "Cannot loop over " ++ pretty a ++ + " of type " ++ pretty a_t + + boundarg <- checkArg boundexp + checkFuncall Nothing paramts $ boundarg : mergeargs + + WhileLoop cond -> do + case find ((==cond) . paramName . fst) merge of + Just (condparam,_) -> + unless (paramType condparam == Prim Bool) $ + bad $ TypeError $ + "Conditional '" ++ pretty cond ++ "' of while-loop is not boolean, but " ++ + pretty (paramType condparam) ++ "." + Nothing -> + bad $ TypeError $ + "Conditional '" ++ pretty cond ++ "' of while-loop is not a merge varible." + let funparams = mergepat + paramts = map paramDeclType funparams + checkFuncall Nothing paramts mergeargs + + let rettype = map paramDeclType mergepat + consumable = [ (paramName param, mempty) + | param <- mergepat, + unique $ paramDeclType param + ] + + context "Inside the loop body" $ + checkFun' (nameFromString "", + staticShapes rettype, + funParamsToNameInfos mergepat, + loopbody) consumable $ do + checkFunParams mergepat + checkBody loopbody + + let rettype_ext = existentialiseExtTypes (map paramName mergepat) $ + staticShapes $ map fromDecl rettype + + bodyt <- extendedScope (traverse subExpType $ bodyResult loopbody) $ + scopeOf $ bodyStms loopbody + + case instantiateShapes (`maybeNth` bodyResult loopbody) rettype_ext of + Nothing -> bad $ ReturnTypeError (nameFromString "") + (staticShapes $ map fromDecl rettype) (staticShapes bodyt) + Just rettype' -> + unless (bodyt `subtypesOf` rettype') $ + bad $ ReturnTypeError (nameFromString "") + (staticShapes rettype') (staticShapes bodyt) + +checkExp (Op op) = checkOp op + +checkSOACArrayArgs :: Checkable lore => + SubExp -> [VName] -> TypeM lore [Arg] +checkSOACArrayArgs width vs = + forM vs $ \v -> do + (vt, v') <- checkSOACArrayArg v + let argSize = arraySize 0 vt + unless (argSize == width) $ + bad $ TypeError $ + "SOAC argument " ++ pretty v ++ " has outer size " ++ + pretty argSize ++ ", but width of SOAC is " ++ + pretty width + return v' + where checkSOACArrayArg ident = do + (t, als) <- checkArg $ Var ident + case peelArray 1 t of + Nothing -> bad $ TypeError $ + "SOAC argument " ++ pretty ident ++ " is not an array" + Just rt -> return (t, (rt, als)) + +checkType :: Checkable lore => + TypeBase Shape u -> TypeM lore () +checkType = mapM_ checkSubExp . arrayDims + +checkExtType :: Checkable lore => + TypeBase ExtShape u + -> TypeM lore () +checkExtType = mapM_ checkExtDim . shapeDims . arrayShape + where checkExtDim (Free se) = void $ checkSubExp se + checkExtDim (Ext _) = return () + +checkCmpOp :: Checkable lore => + CmpOp -> SubExp -> SubExp + -> TypeM lore () +checkCmpOp (CmpEq t) x y = do + require [Prim t] x + require [Prim t] y +checkCmpOp (CmpUlt t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (CmpUle t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (CmpSlt t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (CmpSle t) x y = checkBinOpArgs (IntType t) x y +checkCmpOp (FCmpLt t) x y = checkBinOpArgs (FloatType t) x y +checkCmpOp (FCmpLe t) x y = checkBinOpArgs (FloatType t) x y +checkCmpOp CmpLlt x y = checkBinOpArgs Bool x y +checkCmpOp CmpLle x y = checkBinOpArgs Bool x y + +checkBinOpArgs :: Checkable lore => + PrimType -> SubExp -> SubExp -> TypeM lore () +checkBinOpArgs t e1 e2 = do + require [Prim t] e1 + require [Prim t] e2 + +checkPatElem :: Checkable lore => + PatElemT (LetAttr lore) -> TypeM lore () +checkPatElem (PatElem name attr) = checkLetBoundLore name attr + +checkDimIndex :: Checkable lore => + DimIndex SubExp -> TypeM lore () +checkDimIndex (DimFix i) = require [Prim int32] i +checkDimIndex (DimSlice i n s) = mapM_ (require [Prim int32]) [i,n,s] + +checkStm :: Checkable lore => + Stm (Aliases lore) + -> TypeM lore a + -> TypeM lore a +checkStm stm@(Let pat (StmAux (Certificates cs) (_,attr)) e) m = do + mapM_ (requireI [Prim Cert]) cs + checkExpLore attr + context ("When matching\n" ++ message " " pat ++ "\nwith\n" ++ message " " e) $ + matchPattern pat e + binding (scopeOf stm) $ do + mapM_ checkPatElem (patternElements $ removePatternAliases pat) + m + +matchExtPattern :: Checkable lore => + Pattern (Aliases lore) -> [ExtType] -> TypeM lore () +matchExtPattern pat ts = + unless (expExtTypesFromPattern pat == ts) $ + bad $ InvalidPatternError pat ts Nothing + +matchExtReturnType :: Checkable lore => + [ExtType] -> Result -> TypeM lore () +matchExtReturnType rettype res = do + ts <- mapM subExpType res + matchExtReturns rettype res ts + +matchExtBranchType :: Checkable lore => + [ExtType] -> Body (Aliases lore) -> TypeM lore () +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 rettype res ts = do + let problem :: TypeM lore a + problem = bad $ TypeError $ unlines [ "Type annotation is" + , " " ++ prettyTuple rettype + , "But result returns type" + , " " ++ prettyTuple ts ] + + let (ctx_res, val_res) = splitFromEnd (length rettype) res + (ctx_ts, val_ts) = splitFromEnd (length rettype) ts + + unless (length val_res == length rettype) problem + + let ctx_vals = zip ctx_res ctx_ts + instantiateExt i = case maybeNth i ctx_vals of + Just (se, Prim (IntType Int32)) -> return se + _ -> problem + + rettype' <- instantiateShapes instantiateExt rettype + + unless (rettype' == val_ts) problem + +validApply :: ArrayShape shape => + [TypeBase shape Uniqueness] + -> [TypeBase shape NoUniqueness] + -> Bool +validApply expected got = + length got == length expected && + and (zipWith subtypeOf + (map rankShaped got) + (map (fromDecl . rankShaped) expected)) + +type Arg = (Type, Names) + +argType :: Arg -> Type +argType (t, _) = t + +-- | Remove all aliases from the 'Arg'. +argAliases :: Arg -> Names +argAliases (_, als) = als + +noArgAliases :: Arg -> Arg +noArgAliases (t, _) = (t, mempty) + +checkArg :: Checkable lore => + SubExp -> TypeM lore Arg +checkArg arg = do argt <- checkSubExp arg + als <- subExpAliasesM arg + return (argt, als) + +checkFuncall :: Maybe Name + -> [DeclType] -> [Arg] + -> TypeM lore () +checkFuncall fname paramts args = do + let argts = map argType args + unless (validApply paramts argts) $ + bad $ ParameterMismatch fname + (map fromDecl paramts) $ + map argType args + forM_ (zip (map diet paramts) args) $ \(d, (_, als)) -> + occur [consumption (consumeArg als d)] + where consumeArg als Consume = als + consumeArg _ Observe = mempty + +checkLambda :: Checkable lore => + Lambda (Aliases lore) -> [Arg] -> TypeM lore () +checkLambda (Lambda params body rettype) args = do + let fname = nameFromString "" + if length params == length args then do + checkFuncall Nothing + (map ((`toDecl` Nonunique) . paramType) params) args + let consumable = zip (map paramName params) (map argAliases args) + checkFun' (fname, + staticShapes $ map (`toDecl` Nonunique) rettype, + [ (paramName param, + LParamInfo $ paramAttr param) + | param <- params ], + body) consumable $ do + checkLambdaParams params + mapM_ checkType rettype + checkLambdaBody rettype body + else bad $ TypeError $ "Anonymous function defined with " ++ show (length params) ++ " parameters, but expected to take " ++ show (length args) ++ " arguments." + +checkPrimExp :: Checkable lore => PrimExp VName -> TypeM lore () +checkPrimExp ValueExp{} = return () +checkPrimExp (LeafExp v pt) = requireI [Prim pt] v +checkPrimExp (BinOpExp op x y) = do requirePrimExp (binOpType op) x + requirePrimExp (binOpType op) y +checkPrimExp (CmpOpExp op x y) = do requirePrimExp (cmpOpType op) x + requirePrimExp (cmpOpType op) y +checkPrimExp (UnOpExp op x) = requirePrimExp (unOpType op) x +checkPrimExp (ConvOpExp op x) = requirePrimExp (fst $ convOpType op) x +checkPrimExp (FunExp h args t) = do + (h_ts, h_ret, _) <- maybe (bad $ TypeError $ "Unknown function: " ++ h) + return $ M.lookup h primFuns + when (length h_ts /= length args) $ + bad $ TypeError $ "Function expects " ++ show (length h_ts) ++ + " parameters, but given " ++ show (length args) ++ " arguments." + when (h_ret /= t) $ + bad $ TypeError $ "Function return annotation is " ++ pretty t ++ + ", but expected " ++ pretty h_ret + zipWithM_ requirePrimExp h_ts args + +requirePrimExp :: Checkable lore => PrimType -> PrimExp VName -> TypeM lore () +requirePrimExp t e = context ("in PrimExp " ++ pretty e) $ do + checkPrimExp e + unless (primExpType e == t) $ bad $ TypeError $ + pretty e ++ " must have type " ++ pretty t + +-- | The class of lores that can be type-checked. +class (Attributes lore, CanBeAliased (Op lore)) => Checkable lore where + checkExpLore :: ExpAttr lore -> TypeM lore () + checkBodyLore :: BodyAttr lore -> TypeM lore () + checkFParamLore :: VName -> FParamAttr lore -> TypeM lore () + checkLParamLore :: VName -> LParamAttr lore -> TypeM lore () + checkLetBoundLore :: VName -> LetAttr lore -> TypeM lore () + checkRetType :: [RetType lore] -> TypeM lore () + checkOp :: OpWithAliases (Op lore) -> TypeM lore () + matchPattern :: Pattern (Aliases lore) -> Exp (Aliases lore) -> TypeM lore () + primFParam :: VName -> PrimType -> TypeM lore (FParam (Aliases lore)) + primLParam :: VName -> PrimType -> TypeM lore (LParam (Aliases lore)) + matchReturnType :: [RetType lore] -> Result -> TypeM lore () + matchBranchType :: [BranchType lore] -> Body (Aliases lore) -> TypeM lore () diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Util.hs @@ -0,0 +1,255 @@ +-- | Non-Futhark-specific utilities. If you find yourself writing +-- general functions on generic data structures, consider putting them +-- here. +-- +-- Sometimes it is also preferable to copy a small function rather +-- than introducing a large dependency. In this case, make sure to +-- note where you got it from (and make sure that the license is +-- compatible). +module Futhark.Util + (mapAccumLM, + chunk, + chunks, + dropAt, + takeLast, + dropLast, + mapEither, + maybeNth, + maybeHead, + splitFromEnd, + splitAt3, + splitAt4, + focusNth, + unixEnvironment, + isEnvVarSet, + runProgramWithExitCode, + directoryContents, + roundFloat, + roundDouble, + fromPOSIX, + toPOSIX, + trim, + zEncodeString + ) + where + +import Numeric +import Control.Exception +import Data.Char +import Data.List +import Data.Either +import Data.Maybe +import System.Environment +import System.IO.Unsafe +import qualified System.Directory.Tree as Dir +import System.Process +import System.Exit +import qualified System.FilePath.Posix as Posix +import qualified System.FilePath as Native + +-- | Like 'mapAccumL', but monadic. +mapAccumLM :: Monad m => + (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y]) +mapAccumLM _ acc [] = return (acc, []) +mapAccumLM f acc (x:xs) = do + (acc', x') <- f acc x + (acc'', xs') <- mapAccumLM f acc' xs + return (acc'', x':xs') + +-- | @chunk n a@ splits @a@ into @n@-size-chunks. If the length of +-- @a@ is not divisible by @n@, the last chunk will have fewer than +-- @n@ elements (but it will never be empty). +chunk :: Int -> [a] -> [[a]] +chunk _ [] = [] +chunk n xs = + let (bef,aft) = splitAt n xs + in bef : chunk n aft + +-- | @chunks ns a@ splits @a@ into chunks determined by the elements +-- of @ns@. It must hold that @sum ns == length a@, or the resulting +-- list may contain too few chunks, or not all elements of @a@. +chunks :: [Int] -> [a] -> [[a]] +chunks [] _ = [] +chunks (n:ns) xs = + let (bef,aft) = splitAt n xs + in bef : chunks ns aft + +-- | @dropAt i n@ drops @n@ elements starting at element @i@. +dropAt :: Int -> Int -> [a] -> [a] +dropAt i n xs = take i xs ++ drop (i+n) xs + +-- | @takeLast n l@ takes the last @n@ elements of @l@. +takeLast :: Int -> [a] -> [a] +takeLast n = reverse . take n . reverse + +-- | @dropLast n l@ drops the last @n@ elements of @l@. +dropLast :: Int -> [a] -> [a] +dropLast n = reverse . drop n . reverse + +-- | A combination of 'map' and 'partitionEithers'. +mapEither :: (a -> Either b c) -> [a] -> ([b], [c]) +mapEither f l = partitionEithers $ map f l + +-- | Return the list element at the given index, if the index is valid. +maybeNth :: Integral int => int -> [a] -> Maybe a +maybeNth i l + | i >= 0, v:_ <- genericDrop i l = Just v + | otherwise = Nothing + +-- | Return the first element of the list, if it exists. +maybeHead :: [a] -> Maybe a +maybeHead [] = Nothing +maybeHead (x:_) = Just x + +-- | Like 'splitAt', but from the end. +splitFromEnd :: Int -> [a] -> ([a], [a]) +splitFromEnd i l = splitAt (length l - i) l + +-- | Like 'splitAt', but produces three lists. +splitAt3 :: Int -> Int -> [a] -> ([a], [a], [a]) +splitAt3 n m l = + let (xs, l') = splitAt n l + (ys, zs) = splitAt m l' + in (xs, ys, zs) + +-- | Like 'splitAt', but produces four lists. +splitAt4 :: Int -> Int -> Int -> [a] -> ([a], [a], [a], [a]) +splitAt4 n m k l = + let (xs, l') = splitAt n l + (ys, l'') = splitAt m l' + (zs, vs) = splitAt k l'' + in (xs, ys, zs, vs) + +-- | Return the list element at the given index, if the index is +-- valid, along with the elements before and after. +focusNth :: Integral int => int -> [a] -> Maybe ([a], a, [a]) +focusNth i xs + | (bef, x:aft) <- genericSplitAt i xs = Just (bef, x, aft) + | otherwise = Nothing + +{-# NOINLINE unixEnvironment #-} +-- | The Unix environment when the Futhark compiler started. +unixEnvironment :: [(String,String)] +unixEnvironment = unsafePerformIO getEnvironment + +-- Is an environment variable set to 0 or 1? If 0, return False; if 1, True; +-- otherwise the default value. +isEnvVarSet :: String -> Bool -> Bool +isEnvVarSet name default_val = fromMaybe default_val $ do + val <- lookup name unixEnvironment + case val of + "0" -> return False + "1" -> return True + _ -> Nothing + +-- | Like 'readProcessWithExitCode', but also wraps exceptions when +-- the indicated binary cannot be launched, or some other exception is +-- thrown. +runProgramWithExitCode :: FilePath -> [String] -> String + -> IO (Either IOException (ExitCode, String, String)) +runProgramWithExitCode exe args inp = + (Right <$> readProcessWithExitCode exe args inp) + `catch` \e -> return (Left e) + +-- | Every non-directory file contained in a directory tree. +directoryContents :: FilePath -> IO [FilePath] +directoryContents dir = do + _ Dir.:/ tree <- Dir.readDirectoryWith return dir + case Dir.failures tree of + Dir.Failed _ err : _ -> throw err + _ -> return $ mapMaybe isFile $ Dir.flattenDir tree + where isFile (Dir.File _ path) = Just path + isFile _ = Nothing + +foreign import ccall "nearbyint" c_nearbyint :: Double -> Double +foreign import ccall "nearbyintf" c_nearbyintf :: Float -> Float + +-- | Round a single-precision floating point number correctly. +roundFloat :: Float -> Float +roundFloat = c_nearbyintf + +-- | Round a double-precision floating point number correctly. +roundDouble :: Double -> Double +roundDouble = c_nearbyint + +-- | Turn a POSIX filepath into a filepath for the native system. +toPOSIX :: Native.FilePath -> Posix.FilePath +toPOSIX = Posix.joinPath . Native.splitDirectories + +-- | Some bad operating systems do not use forward slash as +-- directory separator - this is where we convert Futhark includes +-- (which always use forward slash) to native paths. +fromPOSIX :: Posix.FilePath -> Native.FilePath +fromPOSIX = Native.joinPath . Posix.splitDirectories + +-- | Remove leading and trailing whitespace from a string. Not an +-- efficient implementation! +trim :: String -> String +trim = reverse . dropWhile isSpace . reverse . dropWhile isSpace + +-- Z-encoding from https://ghc.haskell.org/trac/ghc/wiki/Commentary/Compiler/SymbolNames +-- +-- Slightly simplified as we do not need it to deal with tuples and +-- the like. +-- +-- (c) The University of Glasgow, 1997-2006 + + +type UserString = String -- As the user typed it +type EncodedString = String -- Encoded form + +zEncodeString :: UserString -> EncodedString +zEncodeString "" = "" +zEncodeString (c:cs) = encodeDigitChar c ++ concatMap encodeChar cs + +unencodedChar :: Char -> Bool -- True for chars that don't need encoding +unencodedChar 'Z' = False +unencodedChar 'z' = False +unencodedChar '_' = True +unencodedChar c = isAsciiLower c + || isAsciiUpper c + || isDigit c + +-- If a digit is at the start of a symbol then we need to encode it. +-- Otherwise names like 9pH-0.1 give linker errors. +encodeDigitChar :: Char -> EncodedString +encodeDigitChar c | isDigit c = encodeAsUnicodeCharar c + | otherwise = encodeChar c + +encodeChar :: Char -> EncodedString +encodeChar c | unencodedChar c = [c] -- Common case first + +-- Constructors +encodeChar '(' = "ZL" -- Needed for things like (,), and (->) +encodeChar ')' = "ZR" -- For symmetry with ( +encodeChar '[' = "ZM" +encodeChar ']' = "ZN" +encodeChar ':' = "ZC" +encodeChar 'Z' = "ZZ" + +-- Variables +encodeChar 'z' = "zz" +encodeChar '&' = "za" +encodeChar '|' = "zb" +encodeChar '^' = "zc" +encodeChar '$' = "zd" +encodeChar '=' = "ze" +encodeChar '>' = "zg" +encodeChar '#' = "zh" +encodeChar '.' = "zi" +encodeChar '<' = "zl" +encodeChar '-' = "zm" +encodeChar '!' = "zn" +encodeChar '+' = "zp" +encodeChar '\'' = "zq" +encodeChar '\\' = "zr" +encodeChar '/' = "zs" +encodeChar '*' = "zt" +encodeChar '_' = "zu" +encodeChar '%' = "zv" +encodeChar c = encodeAsUnicodeCharar c + +encodeAsUnicodeCharar :: Char -> EncodedString +encodeAsUnicodeCharar c = 'z' : if isDigit (head hex_str) then hex_str + else '0':hex_str + where hex_str = showHex (ord c) "U" diff --git a/src/Futhark/Util/IntegralExp.hs b/src/Futhark/Util/IntegralExp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Util/IntegralExp.hs @@ -0,0 +1,75 @@ +-- | It is occasionally useful to define generic functions that can +-- not only compute their result as an integer, but also as a symbolic +-- expression in the form of an AST. +-- +-- There are some Haskell hacks for this - it is for example not hard +-- to define an instance of 'Num' that constructs an AST. However, +-- this falls down for some other interesting classes, like +-- 'Integral', which requires both the problematic method +-- 'fromInteger', and also that the type is an instance of 'Enum'. +-- +-- We can always just define hobbled instances that call 'error' for +-- those methods that are impractical, but this is ugly. +-- +-- Hence, this module defines similes to standard Haskell numeric +-- typeclasses that have been modified to make generic functions +-- slightly easier to write. +module Futhark.Util.IntegralExp + ( IntegralExp (..) + , Wrapped (..) + , quotRoundingUp + ) + where + +import Data.Int + +class Num e => IntegralExp e where + quot :: e -> e -> e + rem :: e -> e -> e + div :: e -> e -> e + mod :: e -> e -> e + sgn :: e -> Maybe Int + + fromInt8 :: Int8 -> e + fromInt16 :: Int16 -> e + fromInt32 :: Int32 -> e + fromInt64 :: Int64 -> e + +-- | This wrapper allows you to use a type that is an instance of the +-- true class whenever the simile class is required. +newtype Wrapped a = Wrapped { wrappedValue :: a } + deriving (Eq, Ord, Show) + +liftOp :: (a -> a) + -> Wrapped a -> Wrapped a +liftOp op (Wrapped x) = Wrapped $ op x + +liftOp2 :: (a -> a -> a) + -> Wrapped a -> Wrapped a -> Wrapped a +liftOp2 op (Wrapped x) (Wrapped y) = Wrapped $ x `op` y + +instance Num a => Num (Wrapped a) where + (+) = liftOp2 (Prelude.+) + (-) = liftOp2 (Prelude.-) + (*) = liftOp2 (Prelude.*) + abs = liftOp Prelude.abs + signum = liftOp Prelude.signum + fromInteger = Wrapped . Prelude.fromInteger + negate = liftOp Prelude.negate + +instance Integral a => IntegralExp (Wrapped a) where + quot = liftOp2 Prelude.quot + rem = liftOp2 Prelude.rem + div = liftOp2 Prelude.div + mod = liftOp2 Prelude.mod + sgn = Just . fromIntegral . signum . toInteger . wrappedValue + + fromInt8 = fromInteger . toInteger + fromInt16 = fromInteger . toInteger + fromInt32 = fromInteger . toInteger + fromInt64 = fromInteger . toInteger + +-- | Like 'quot', but rounds up. +quotRoundingUp :: IntegralExp num => num -> num -> num +quotRoundingUp x y = + (x + y - 1) `Futhark.Util.IntegralExp.quot` y diff --git a/src/Futhark/Util/Log.hs b/src/Futhark/Util/Log.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Util/Log.hs @@ -0,0 +1,63 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE OverloadedStrings #-} +-- | Opaque type for an operations log that provides fast O(1) +-- appends. +module Futhark.Util.Log + ( Log + , toText + , ToLog (..) + , MonadLogger (..) + ) + +where + +import Control.Monad.Writer +import qualified Control.Monad.RWS.Strict +import qualified Control.Monad.RWS.Lazy +import qualified Data.Text as T +import qualified Data.DList as DL +import qualified Data.Semigroup as Sem + +newtype Log = Log { unLog :: DL.DList T.Text } + +instance Sem.Semigroup Log where + Log l1 <> Log l2 = Log $ l1 <> l2 + +instance Monoid Log where + mappend = (Sem.<>) + mempty = Log mempty + +-- | Transform a log into text. Every log entry becomes its own line +-- (or possibly more, in case of multi-line entries). +toText :: Log -> T.Text +toText = T.intercalate "\n" . DL.toList . unLog + +-- | Typeclass for things that can be turned into a single-entry log. +class ToLog a where + toLog :: a -> Log + +instance ToLog String where + toLog = Log . DL.singleton . T.pack + +instance ToLog T.Text where + toLog = Log . DL.singleton + +-- | Typeclass for monads that support logging. +class (Applicative m, Monad m) => MonadLogger m where + -- | Add one log entry. + logMsg :: ToLog a => a -> m () + logMsg = addLog . toLog + + -- | Append an entire log. + addLog :: Log -> m () + +instance (Applicative m, Monad m) => MonadLogger (WriterT Log m) where + addLog = tell + +instance (Applicative m, Monad m) => MonadLogger (Control.Monad.RWS.Lazy.RWST r Log s m) where + addLog = tell + +instance (Applicative m, Monad m) => MonadLogger (Control.Monad.RWS.Strict.RWST r Log s m) where + addLog = tell diff --git a/src/Futhark/Util/Options.hs b/src/Futhark/Util/Options.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Util/Options.hs @@ -0,0 +1,84 @@ +-- | Common code for parsing command line options based on getopt. +module Futhark.Util.Options + ( FunOptDescr + , mainWithOptions + , commonOptions + ) where + +import System.Environment +import Control.Monad.IO.Class +import System.IO +import System.Exit +import System.Console.GetOpt + +import Futhark.Version + +-- | A command line option that either purely updates a configuration, +-- or performs an IO action (and stops). +type FunOptDescr cfg = OptDescr (Either (IO ()) (cfg -> cfg)) + +-- | Generate a main action that parses the given command line options +-- (while always adding 'commonOptions'). +mainWithOptions :: cfg + -> [FunOptDescr cfg] + -> String + -> ([String] -> cfg -> Maybe (IO ())) + -> IO () +mainWithOptions emptyConfig commandLineOptions usage f = do + args <- getArgs + case getOpt' Permute commandLineOptions' args of + (opts, nonopts, [], []) -> + case applyOpts opts of + Right config + | Just m <- f nonopts config -> m + | otherwise -> invalid nonopts [] [] + Left m -> m + (_, nonopts, unrecs, errs) -> invalid nonopts unrecs errs + where applyOpts opts = do fs <- sequence opts + return $ foldl (.) id (reverse fs) emptyConfig + + invalid nonopts unrecs errs = do help <- helpStr usage commandLineOptions' + badOptions help nonopts errs unrecs + + commandLineOptions' = + commonOptions usage commandLineOptions ++ commandLineOptions + +helpStr :: String -> [OptDescr a] -> IO String +helpStr usage opts = do + prog <- getProgName + + let header = unlines ["Usage: " ++ prog ++ " " ++ usage, "Options:"] + return $ usageInfo header opts + +badOptions :: String -> [String] -> [String] -> [String] -> IO () +badOptions usage nonopts errs unrecs = do + mapM_ (errput . ("Junk argument: " ++)) nonopts + mapM_ (errput . ("Unrecognised argument: " ++)) unrecs + hPutStr stderr $ concat errs ++ usage + exitWith $ ExitFailure 1 + +-- | Short-hand for 'liftIO . hPutStrLn stderr' +errput :: MonadIO m => String -> m () +errput = liftIO . hPutStrLn stderr + +-- | Common definitions for @-v@ and @-h@, given the list of all other +-- options. +commonOptions :: String -> [FunOptDescr cfg] -> [FunOptDescr cfg] +commonOptions usage options = + [ Option "V" ["version"] + (NoArg $ Left $ do header + exitSuccess) + "Print version information and exit." + + , Option "h" ["help"] + (NoArg $ Left $ do header + putStrLn "" + putStrLn =<< helpStr usage (commonOptions usage [] ++ options) + exitSuccess) + "Print help and exit." + ] + where header = do + putStrLn $ "Futhark " ++ versionString + putStrLn "Copyright (C) DIKU, University of Copenhagen, released under the ISC license." + putStrLn "This is free software: you are free to change and redistribute it." + putStrLn "There is NO WARRANTY, to the extent permitted by law." diff --git a/src/Futhark/Util/Pretty.hs b/src/Futhark/Util/Pretty.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Util/Pretty.hs @@ -0,0 +1,69 @@ +-- | A re-export of the prettyprinting library, along with a convenience function. +module Futhark.Util.Pretty + ( module Text.PrettyPrint.Mainland + , module Text.PrettyPrint.Mainland.Class + , pretty + , prettyDoc + , prettyTuple + , prettyText + , prettyOneLine + + , apply + , oneLine + , annot + , nestedBlock + ) + where + +import Data.Text (Text) +import qualified Data.Text.Lazy as LT + +import Text.PrettyPrint.Mainland hiding (pretty) +import Text.PrettyPrint.Mainland.Class +import qualified Text.PrettyPrint.Mainland as PP + +-- | Prettyprint a value, wrapped to 80 characters. +pretty :: Pretty a => a -> String +pretty = PP.pretty 80 . ppr + +-- | Prettyprint a value to a 'Text', wrapped to 80 characters. +prettyText :: Pretty a => a -> Text +prettyText = LT.toStrict . PP.prettyLazyText 80 . ppr + +-- | Prettyprint a value without any width restriction. +prettyOneLine :: Pretty a => a -> String +prettyOneLine = ($"") . displayS . renderCompact . oneLine . ppr + +-- | Re-export of 'PP.pretty'. +prettyDoc :: Int -> Doc -> String +prettyDoc = PP.pretty + +ppTuple' :: Pretty a => [a] -> Doc +ppTuple' ets = braces $ commasep $ map ppr ets + +-- | Prettyprint a list enclosed in curly braces. +prettyTuple :: Pretty a => [a] -> String +prettyTuple = PP.pretty 80 . ppTuple' + +-- | The document @'apply' ds@ separates @ds@ with commas and encloses them with +-- parentheses. +apply :: [Doc] -> Doc +apply = parens . commasep . map align + +-- | Make sure that the given document is printed on just a single line. +oneLine :: PP.Doc -> PP.Doc +oneLine s = PP.text $ PP.displayS (PP.renderCompact s) "" + +-- | Stack and prepend a list of 'Doc's to another 'Doc', separated by +-- a linebreak. If the list is empty, the second 'Doc' will be +-- returned without a preceding linebreak. +annot :: [Doc] -> Doc -> Doc +annot [] s = s +annot l s = stack l s + +-- | Surround the given document with enclosers and add linebreaks and +-- indents. +nestedBlock :: String -> String -> Doc -> Doc +nestedBlock pre post body = text pre + PP.indent 2 body + text post diff --git a/src/Futhark/Util/Table.hs b/src/Futhark/Util/Table.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Util/Table.hs @@ -0,0 +1,53 @@ +-- | Basic table building for prettier futhark-test output. +module Futhark.Util.Table + ( buildTable + , mkEntry + , Entry + ) where + +import Data.List +import System.Console.ANSI + +data RowTemplate = RowTemplate [Int] Int deriving (Show) + +-- | A table entry. Consists of the content as well a list of +-- SGR commands to color/stylelize the entry. +type Entry = (String, [SGR]) + +-- | Makes a table entry with the default SGR mode. +mkEntry :: String -> (String, [SGR]) +mkEntry s = (s, []) + +color :: [SGR] -> String -> String +color sgr s = setSGRCode sgr ++ s ++ setSGRCode [Reset] + +buildRowTemplate :: [[Entry]] -> Int -> RowTemplate + +buildRowTemplate rows = RowTemplate widths + where widths = map (maximum . map (length . fst)) . transpose $ rows + +buildRow :: RowTemplate -> [Entry] -> String +buildRow (RowTemplate widths pad) entries = cells ++ "\n" + where bar = "\x2502" + cells = concatMap buildCell (zip entries widths) ++ bar + buildCell ((entry, sgr), width) = + let padding = width - length entry + pad + in bar ++ " " ++ color sgr entry ++ replicate padding ' ' + +buildSep :: Char -> Char -> Char -> RowTemplate -> String +buildSep lCorner rCorner sep (RowTemplate widths pad) = + corners . concatMap cellFloor $ widths + where cellFloor width = replicate (width + pad + 1) '\x2500' ++ [sep] + corners [] = "" + corners s = [lCorner] ++ init s ++ [rCorner] + +-- | Builds a table from a list of entries and a padding amount that +-- determines padding from the right side of the widest entry in each column. +buildTable :: [[Entry]] -> Int -> String +buildTable rows pad = buildTop template ++ sepRows ++ buildBottom template + where sepRows = intercalate (buildFloor template) builtRows + builtRows = map (buildRow template) rows + template = buildRowTemplate rows pad + buildTop rt = buildSep '\x250C' '\x2510' '\x252C' rt ++ "\n" + buildFloor rt = buildSep '\x251C' '\x2524' '\x253C' rt ++ "\n" + buildBottom = buildSep '\x2514' '\x2518' '\x2534' diff --git a/src/Futhark/Version.hs b/src/Futhark/Version.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Version.hs @@ -0,0 +1,34 @@ +{-# LANGUAGE TemplateHaskell #-} +-- | This module exports version information about the Futhark +-- compiler. +module Futhark.Version + ( + version + , versionString + ) + where + +import Data.Version +import Development.GitRev + +import qualified Paths_futhark + +-- | The version of Futhark that we are using. This is equivalent to +-- the version defined in the .cabal file. +version :: Version +version = Paths_futhark.version + +-- | The version of Futhark that we are using, as a 'String' +versionString :: String +versionString = showVersion version ++ "\n" ++ gitversion + where + gitversion = concat ["git: " + , branch + , take 7 $(gitHash) + , " (", $(gitCommitDate), ")" + , dirty + ] + branch | $(gitBranch) == "master" = "" + | otherwise = $(gitBranch) ++ " @ " + dirty | $(gitDirtyTracked) = " [modified]" + | otherwise = "" diff --git a/src/Language/Futhark.hs b/src/Language/Futhark.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark.hs @@ -0,0 +1,73 @@ +-- | Re-export the external Futhark modules for convenience. +module Language.Futhark + ( module Language.Futhark.Syntax + , module Language.Futhark.Attributes + , module Language.Futhark.Pretty + + , Ident, DimIndex, Exp, Pattern + , ModExp, ModParam, SigExp, ModBind, SigBind + , ValBind, Dec, Spec, Prog + , TypeBind, TypeDecl + , StructTypeArg, ArrayElemType + , TypeParam + ) + where + +import Language.Futhark.Syntax +import Language.Futhark.Attributes +import Language.Futhark.Pretty + +-- | An identifier with type- and aliasing information. +type Ident = IdentBase Info VName + +-- | An index with type information. +type DimIndex = DimIndexBase Info VName + +-- | An expression with type information. +type Exp = ExpBase Info VName + +-- | A pattern with type information. +type Pattern = PatternBase Info VName + +-- | An constant declaration with type information. +type ValBind = ValBindBase Info VName + +-- | A type declaration with type information +type TypeDecl = TypeDeclBase Info VName + +-- | A type binding with type information. +type TypeBind = TypeBindBase Info VName + +-- | A type-checked module binding. +type ModBind = ModBindBase Info VName + +-- | A type-checked module type binding. +type SigBind = SigBindBase Info VName + +-- | A type-checked module expression. +type ModExp = ModExpBase Info VName + +-- | A type-checked module parameter. +type ModParam = ModParamBase Info VName + +-- | A type-checked module type expression. +type SigExp = SigExpBase Info VName + +-- | A type-checked declaration. +type Dec = DecBase Info VName + +-- | A type-checked specification. +type Spec = SpecBase Info VName + +-- | An Futhark program with type information. +type Prog = ProgBase Info VName + +-- | A known type arg with shape annotations but no aliasing information. +type StructTypeArg = TypeArg (DimDecl VName) () + +-- | A type-checked type parameter. +type TypeParam = TypeParamBase VName + +-- | A known array element type with no shape annotations, but aliasing +-- information. +type ArrayElemType = ArrayElemTypeBase () Names diff --git a/src/Language/Futhark/Attributes.hs b/src/Language/Futhark/Attributes.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Attributes.hs @@ -0,0 +1,1037 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +-- | This module provides various simple ways to query and manipulate +-- fundamental Futhark terms, such as types and values. The intent is to +-- keep "Futhark.Language.Syntax" simple, and put whatever embellishments +-- we need here. +module Language.Futhark.Attributes + ( + -- * Various + Intrinsic(..) + , intrinsics + , maxIntrinsicTag + , namesToPrimTypes + , qualName + , qualify + , typeName + , valueType + , leadingOperator + , progImports + , decImports + , progModuleTypes + , identifierReference + , identifierReferences + + -- * Queries on expressions + , typeOf + + -- * Queries on patterns and params + , patIdentSet + , patternType + , patternStructType + , patternParam + , patternNoShapeAnnotations + , patternOrderZero + , patternDimNames + + -- * Queries on types + , uniqueness + , unique + , recordArrayElemUniqueness + , aliases + , diet + , arrayRank + , nestedDims + , returnType + , concreteType + , orderZero + , unfoldFunType + , foldFunType + , typeVars + , typeDimNames + + -- * Operations on types + , rank + , peelArray + , stripArray + , arrayOf + , arrayOfWithAliases + , toStructural + , toStruct + , fromStruct + , setAliases + , addAliases + , setUniqueness + , modifyShapeAnnotations + , setArrayShape + , removeShapeAnnotations + , vacuousShapeAnnotations + , typeToRecordArrayElem + , typeToRecordArrayElem' + , recordArrayElemToType + , tupleRecord + , isTupleRecord + , areTupleFields + , tupleFieldNames + , sortFields + , isTypeParam + + -- | Values of these types are produces by the parser. They use + -- unadorned names and have no type information, apart from that + -- which is syntactically required. + , NoInfo(..) + , UncheckedType + , UncheckedTypeExp + , UncheckedArrayElemType + , UncheckedIdent + , UncheckedTypeDecl + , UncheckedDimIndex + , UncheckedExp + , UncheckedModExp + , UncheckedSigExp + , UncheckedTypeParam + , UncheckedPattern + , UncheckedValBind + , UncheckedDec + , UncheckedProg + ) + where + +import Control.Monad.Writer +import Data.Char +import Data.Foldable +import qualified Data.Map.Strict as M +import qualified Data.Set as S +import Data.List +import Data.Loc +import Data.Maybe +import Data.Ord +import Data.Bifunctor +import Data.Bifoldable + +import Prelude + +import Futhark.Util.Pretty + +import Language.Futhark.Syntax +import qualified Futhark.Representation.Primitive as Primitive + +-- | Return the dimensionality of a type. For non-arrays, this is +-- zero. For a one-dimensional array it is one, for a two-dimensional +-- it is two, and so forth. +arrayRank :: TypeBase dim as -> Int +arrayRank = shapeRank . arrayShape + +-- | Return the shape of a type - for non-arrays, this is 'mempty'. +arrayShape :: TypeBase dim as -> ShapeDecl dim +arrayShape (Array _ ds _) = ds +arrayShape _ = mempty + +-- | Return any shape declarations in the type, with duplicates +-- removed. +nestedDims :: TypeBase (DimDecl VName) as -> [DimDecl VName] +nestedDims t = + case t of Array a ds _ -> nub $ arrayNestedDims a <> shapeDims ds + Record fs -> nub $ fold $ fmap nestedDims fs + Prim{} -> mempty + TypeVar _ _ _ targs -> concatMap typeArgDims targs + Arrow _ v t1 t2 -> filter (notV v) $ nestedDims t1 <> nestedDims t2 + where arrayNestedDims ArrayPrimElem{} = + mempty + arrayNestedDims (ArrayPolyElem _ targs _) = + concatMap typeArgDims targs + arrayNestedDims (ArrayRecordElem ts) = + fold (fmap recordArrayElemNestedDims ts) + + recordArrayElemNestedDims (RecordArrayArrayElem a ds _) = + arrayNestedDims a <> shapeDims ds + recordArrayElemNestedDims (RecordArrayElem et) = + arrayNestedDims et + + typeArgDims (TypeArgDim d _) = [d] + typeArgDims (TypeArgType at _) = nestedDims at + + notV Nothing = const True + notV (Just v) = (/=NamedDim (qualName v)) + +-- | Set the dimensions of an array. If the given type is not an +-- array, return the type unchanged. +setArrayShape :: TypeBase dim as -> ShapeDecl dim -> TypeBase dim as +setArrayShape (Array t _ u) ds = Array t ds u +setArrayShape t _ = t + +-- | Change the shape of a type to be just the 'Rank'. +removeShapeAnnotations :: TypeBase dim as -> TypeBase () as +removeShapeAnnotations = modifyShapeAnnotations $ const () + +-- | Change all size annotations to be 'AnyDim'. +vacuousShapeAnnotations :: TypeBase dim as -> TypeBase (DimDecl vn) as +vacuousShapeAnnotations = modifyShapeAnnotations $ const AnyDim + +-- | Change the size annotations of a type. +modifyShapeAnnotations :: (oldshape -> newshape) + -> TypeBase oldshape as + -> TypeBase newshape as +modifyShapeAnnotations f = bimap f id + +-- | Return the uniqueness of a type. +uniqueness :: TypeBase shape as -> Uniqueness +uniqueness (Array _ _ u) = u +uniqueness (TypeVar _ u _ _) = u +uniqueness _ = Nonunique + +recordArrayElemUniqueness :: RecordArrayElemTypeBase shape as -> Uniqueness +recordArrayElemUniqueness RecordArrayElem{} = Nonunique +recordArrayElemUniqueness (RecordArrayArrayElem _ _ u) = u + +-- | @unique t@ is 'True' if the type of the argument is unique. +unique :: TypeBase shape as -> Bool +unique = (==Unique) . uniqueness + +-- | Return the set of all variables mentioned in the aliasing of a +-- type. +aliases :: Monoid as => TypeBase shape as -> as +aliases = bifoldMap (const mempty) id + +-- | @diet t@ returns a description of how a function parameter of +-- type @t@ might consume its argument. +diet :: TypeBase shape as -> Diet +diet (Record ets) = RecordDiet $ fmap diet ets +diet (Prim _) = Observe +diet TypeVar{} = Observe +diet (Arrow _ _ t1 t2) = FuncDiet (diet t1) (diet t2) +diet (Array _ _ Unique) = Consume +diet (Array _ _ Nonunique) = Observe + +-- | @t `maskAliases` d@ removes aliases (sets them to 'mempty') from +-- the parts of @t@ that are denoted as 'Consumed' by the 'Diet' @d@. +maskAliases :: Monoid as => + TypeBase shape as + -> Diet + -> TypeBase shape as +maskAliases t Consume = t `setAliases` mempty +maskAliases t Observe = t +maskAliases (Record ets) (RecordDiet ds) = + Record $ M.intersectionWith maskAliases ets ds +maskAliases t FuncDiet{} = t +maskAliases _ _ = error "Invalid arguments passed to maskAliases." + +-- | Convert any type to one that has rank information, no alias +-- information, and no embedded names. +toStructural :: TypeBase dim as + -> TypeBase () () +toStructural = removeNames . removeShapeAnnotations + +-- | Remove aliasing information from a type. +toStruct :: TypeBase dim as + -> TypeBase dim () +toStruct t = t `setAliases` () + +-- | Replace no aliasing with an empty alias set. +fromStruct :: TypeBase dim as + -> TypeBase dim Names +fromStruct t = t `setAliases` S.empty + +-- | @peelArray n t@ returns the type resulting from peeling the first +-- @n@ array dimensions from @t@. Returns @Nothing@ if @t@ has less +-- than @n@ dimensions. +peelArray :: Int -> TypeBase dim as -> Maybe (TypeBase dim as) +peelArray 0 t = Just t +peelArray n (Array (ArrayPrimElem et _) shape _) + | shapeRank shape == n = + Just $ Prim et +peelArray n (Array (ArrayPolyElem et targs als) shape u) + | shapeRank shape == n = + Just $ TypeVar als u et targs +peelArray n (Array (ArrayRecordElem ts) shape u) + | shapeRank shape == n = + Just $ Record $ fmap asType ts + where asType (RecordArrayElem (ArrayPrimElem bt _)) = Prim bt + asType (RecordArrayElem (ArrayPolyElem bt targs als)) = TypeVar als u bt targs + asType (RecordArrayElem (ArrayRecordElem ts')) = Record $ fmap asType ts' + asType (RecordArrayArrayElem et e_shape _) = Array et e_shape u +peelArray n (Array et shape u) = do + shape' <- stripDims n shape + return $ Array et shape' u +peelArray _ _ = Nothing + +-- | Remove names from a type - this involves removing all size +-- annotations from arrays, as well as all aliasing. +removeNames :: TypeBase dim as + -> TypeBase () () +removeNames = flip setAliases () . removeShapeAnnotations + +-- | @arrayOf t s u@ constructs an array type. The convenience +-- compared to using the 'Array' constructor directly is that @t@ can +-- itself be an array. If @t@ is an @n@-dimensional array, and @s@ is +-- a list of length @n@, the resulting type is of an @n+m@ dimensions. +-- The uniqueness of the new array will be @u@, no matter the +-- uniqueness of @t@. The function returns 'Nothing' in case an +-- attempt is made to create an array of functions. +arrayOf :: Monoid as => + TypeBase dim as + -> ShapeDecl dim + -> Uniqueness + -> Maybe (TypeBase dim as) +arrayOf t = arrayOfWithAliases t mempty + +arrayOfWithAliases :: Monoid as => + TypeBase dim as + -> as + -> ShapeDecl dim + -> Uniqueness + -> Maybe (TypeBase dim as) +arrayOfWithAliases (Array et shape1 _) as shape2 u = + Just $ Array et (shape2 <> shape1) u `setAliases` as +arrayOfWithAliases (Prim et) as shape u = + Just $ Array (ArrayPrimElem et as) shape u +arrayOfWithAliases (TypeVar _ _ x targs) as shape u = + Just $ Array (ArrayPolyElem x targs as) shape u +arrayOfWithAliases (Record ts) as shape u = do + ts' <- traverse (typeToRecordArrayElem' as) ts + return $ Array (ArrayRecordElem ts') shape u +arrayOfWithAliases Arrow{} _ _ _ = Nothing + +typeToRecordArrayElem :: Monoid as => + TypeBase dim as + -> Maybe (RecordArrayElemTypeBase dim as) +typeToRecordArrayElem = typeToRecordArrayElem' mempty + +typeToRecordArrayElem' :: Monoid as => + as -> TypeBase dim as + -> Maybe (RecordArrayElemTypeBase dim as) +typeToRecordArrayElem' as (Prim bt) = + Just $ RecordArrayElem $ ArrayPrimElem bt as +typeToRecordArrayElem' as (TypeVar t_as _ bt targs) = + Just $ RecordArrayElem $ ArrayPolyElem bt targs (as <> t_as) +typeToRecordArrayElem' as (Record ts') = + RecordArrayElem . ArrayRecordElem <$> + traverse (typeToRecordArrayElem' as) ts' +typeToRecordArrayElem' _ (Array et shape u) = + Just $ RecordArrayArrayElem et shape u +typeToRecordArrayElem' _ Arrow{} = Nothing + +recordArrayElemToType :: Monoid as => + RecordArrayElemTypeBase dim as + -> (TypeBase dim as, as) +recordArrayElemToType (RecordArrayElem et) = arrayElemToType et +recordArrayElemToType (RecordArrayArrayElem et shape u) = (Array et shape u, mempty) + +arrayElemToType :: Monoid as => ArrayElemTypeBase dim as -> (TypeBase dim as, as) +arrayElemToType (ArrayPrimElem bt als) = (Prim bt, als) +arrayElemToType (ArrayPolyElem bt targs als) = (TypeVar als Nonunique bt targs, als) +arrayElemToType (ArrayRecordElem ts) = + let ts' = fmap recordArrayElemToType ts + in (Record $ fmap fst ts', foldMap snd ts') + +-- | @stripArray n t@ removes the @n@ outermost layers of the array. +-- Essentially, it is the type of indexing an array of type @t@ with +-- @n@ indexes. +stripArray :: Monoid as => Int -> TypeBase dim as -> TypeBase dim as +stripArray n (Array et shape u) + | Just shape' <- stripDims n shape = + Array et shape' u + | otherwise = fst (arrayElemToType et) `setUniqueness` u +stripArray _ t = t + +-- | Create a record type corresponding to a tuple with the given +-- element types. +tupleRecord :: [TypeBase dim as] -> TypeBase dim as +tupleRecord = Record . M.fromList . zip tupleFieldNames + +isTupleRecord :: TypeBase dim as -> Maybe [TypeBase dim as] +isTupleRecord (Record fs) = areTupleFields fs +isTupleRecord _ = Nothing + +areTupleFields :: M.Map Name a -> Maybe [a] +areTupleFields fs = + let fs' = sortFields fs + in if and $ zipWith (==) (map fst fs') tupleFieldNames + then Just $ map snd fs' + else Nothing + +-- | Increasing field names for a tuple (starts at 1). +tupleFieldNames :: [Name] +tupleFieldNames = map (nameFromString . show) [(1::Int)..] + +-- | Sort fields by their name; taking care to sort numeric fields by +-- their numeric value. This ensures that tuples and tuple-like +-- records match. +sortFields :: M.Map Name a -> [(Name,a)] +sortFields l = map snd $ sortOn fst $ zip (map (fieldish . fst) l') l' + where l' = M.toList l + fieldish s = case reads $ nameToString s of + [(x, "")] -> Left (x::Int) + _ -> Right s + +isTypeParam :: TypeParamBase vn -> Bool +isTypeParam TypeParamType{} = True +isTypeParam TypeParamDim{} = False + + +-- | Set the uniqueness attribute of a type. If the type is a tuple, +-- the uniqueness of its components will be modified. +setUniqueness :: TypeBase dim as -> Uniqueness -> TypeBase dim as +setUniqueness (Array et shape _) u = + Array (setArrayElemUniqueness et u) shape u +setUniqueness (TypeVar als _ t targs) u = + TypeVar als u t targs +setUniqueness (Record ets) u = + Record $ fmap (`setUniqueness` u) ets +setUniqueness t _ = t + +setArrayElemUniqueness :: ArrayElemTypeBase dim as + -> Uniqueness -> ArrayElemTypeBase dim as +setArrayElemUniqueness (ArrayPrimElem bt as) _ = + ArrayPrimElem bt as +setArrayElemUniqueness (ArrayPolyElem v args as) _ = + ArrayPolyElem v args as +setArrayElemUniqueness (ArrayRecordElem r) u = + ArrayRecordElem $ fmap set r + where set (RecordArrayElem et) = + RecordArrayElem $ setArrayElemUniqueness et u + set (RecordArrayArrayElem et shape e_u) = + RecordArrayArrayElem (setArrayElemUniqueness et u) shape e_u + +-- | @t \`setAliases\` als@ returns @t@, but with @als@ substituted for +-- any already present aliasing. +setAliases :: TypeBase dim asf -> ast -> TypeBase dim ast +setAliases t = addAliases t . const + +-- | @t \`addAliases\` f@ returns @t@, but with any already present +-- aliasing replaced by @f@ applied to that aliasing. +addAliases :: TypeBase dim asf -> (asf -> ast) + -> TypeBase dim ast +addAliases t f = bimap id f t + +intValueType :: IntValue -> IntType +intValueType Int8Value{} = Int8 +intValueType Int16Value{} = Int16 +intValueType Int32Value{} = Int32 +intValueType Int64Value{} = Int64 + +floatValueType :: FloatValue -> FloatType +floatValueType Float32Value{} = Float32 +floatValueType Float64Value{} = Float64 + +-- | The type of a basic value. +primValueType :: PrimValue -> PrimType +primValueType (SignedValue v) = Signed $ intValueType v +primValueType (UnsignedValue v) = Unsigned $ intValueType v +primValueType (FloatValue v) = FloatType $ floatValueType v +primValueType BoolValue{} = Bool + +valueType :: Value -> TypeBase () () +valueType (PrimValue bv) = Prim $ primValueType bv +valueType (ArrayValue _ t) = t + +-- | Construct a 'ShapeDecl' with the given number of zero-information +-- dimensions. +rank :: Int -> ShapeDecl () +rank n = ShapeDecl $ replicate n () + +-- | The type of an Futhark term. The aliasing will refer to itself, if +-- the term is a non-tuple-typed variable. +typeOf :: ExpBase Info VName -> CompType +typeOf (Literal val _) = Prim $ primValueType val +typeOf (IntLit _ (Info t) _) = fromStruct t +typeOf (FloatLit _ (Info t) _) = fromStruct t +typeOf (Parens e _) = typeOf e +typeOf (QualParens _ e _) = typeOf e +typeOf (TupLit es _) = tupleRecord $ map typeOf es +typeOf (RecordLit fs _) = + -- Reverse, because M.unions is biased to the left. + Record $ M.unions $ reverse $ map record fs + where record (RecordFieldExplicit name e _) = M.singleton name $ typeOf e + record (RecordFieldImplicit name (Info t) _) = + M.singleton (baseName name) $ t `addAliases` S.insert name +typeOf (ArrayLit _ (Info t) _) = t +typeOf (Range _ _ _ (Info t) _) = t +typeOf (BinOp _ _ _ _ (Info t) _) = removeShapeAnnotations t +typeOf (Project _ _ (Info t) _) = t +typeOf (If _ _ _ (Info t) _) = t +typeOf (Var qn (Info t) _) = removeShapeAnnotations t `addAliases` S.insert (qualLeaf qn) +typeOf (Ascript e _ _) = typeOf e +typeOf (Apply _ _ _ (Info t) _) = removeShapeAnnotations t +typeOf (Negate e _) = typeOf e +typeOf (LetPat _ _ _ body _) = typeOf body +typeOf (LetFun _ _ body _) = typeOf body +typeOf (LetWith _ _ _ _ body _) = typeOf body +typeOf (Index _ _ (Info t) _) = t +typeOf (Update e _ _ _) = typeOf e `setAliases` mempty +typeOf (RecordUpdate _ _ _ (Info t) _) = removeShapeAnnotations t +typeOf (Zip _ _ _ (Info t) _) = t +typeOf (Unzip _ ts _) = + tupleRecord $ map unInfo ts +typeOf (Unsafe e _) = typeOf e +typeOf (Assert _ e _ _) = typeOf e +typeOf (Map _ _ (Info t) _) = t `setUniqueness` Unique +typeOf (Reduce _ _ _ arr _) = + stripArray 1 (typeOf arr) `setAliases` mempty +typeOf (GenReduce hist _ _ _ _ _) = + typeOf hist `setAliases` mempty `setUniqueness` Unique +typeOf (Scan _ _ arr _) = typeOf arr `setAliases` mempty `setUniqueness` Unique +typeOf (Filter _ arr _) = typeOf arr `setAliases` mempty `setUniqueness` Unique +typeOf (Partition _ _ arr _) = + tupleRecord [typeOf arr `setAliases` mempty `setUniqueness` Unique, + Array (ArrayPrimElem (Signed Int32) mempty) (rank 1) Unique] +typeOf (Stream _ lam _ _) = + rettype (typeOf lam) `setUniqueness` Unique + where rettype (Arrow _ _ _ t) = rettype t + rettype t = t +typeOf (DoLoop _ pat _ _ _ _) = patternType pat +typeOf (Lambda _ params _ _ (Info (als, t)) _) = + removeShapeAnnotations (foldr (uncurry (Arrow ()) . patternParam) t params) + `setAliases` als +typeOf (OpSection _ (Info t) _) = + removeShapeAnnotations t +typeOf (OpSectionLeft _ _ _ (_, Info pt2) (Info ret) _) = + removeShapeAnnotations $ foldFunType [fromStruct pt2] ret +typeOf (OpSectionRight _ _ _ (Info pt1, _) (Info ret) _) = + removeShapeAnnotations $ foldFunType [fromStruct pt1] ret +typeOf (ProjectSection _ (Info t) _) = + removeShapeAnnotations t +typeOf (IndexSection _ (Info t) _) = + removeShapeAnnotations t + +foldFunType :: Monoid as => [TypeBase dim as] -> TypeBase dim as -> TypeBase dim as +foldFunType ps ret = foldr (Arrow mempty Nothing) ret ps + +-- | Extract the parameter types and return type from a type. +-- If the type is not an arrow type, the list of parameter types is empty. +unfoldFunType :: TypeBase dim as -> ([TypeBase dim as], TypeBase dim as) +unfoldFunType (Arrow _ _ t1 t2) = let (ps, r) = unfoldFunType t2 + in (t1 : ps, r) +unfoldFunType t = ([], t) + +-- | The type names mentioned in a type. +typeVars :: Monoid as => TypeBase dim as -> Names +typeVars t = + case t of + Prim{} -> mempty + TypeVar _ _ tn targs -> + mconcat $ typeVarFree tn : map typeArgFree targs + Arrow _ _ t1 t2 -> typeVars t1 <> typeVars t2 + Record fields -> foldMap typeVars fields + Array ArrayPrimElem{} _ _ -> mempty + Array (ArrayPolyElem tn targs _) _ _ -> + mconcat $ typeVarFree tn : map typeArgFree targs + Array (ArrayRecordElem fields) _ _ -> + foldMap (typeVars . fst . recordArrayElemToType) fields + where typeVarFree = S.singleton . typeLeaf + typeArgFree (TypeArgType ta _) = typeVars ta + typeArgFree TypeArgDim{} = mempty + +-- | The result of applying the arguments of the given types to a +-- function with the given return type, consuming its parameters with +-- the given diets. +returnType :: TypeBase dim () + -> [Diet] + -> [CompType] + -> TypeBase dim Names +returnType (Array et shape Unique) _ _ = + Array (bimap id (const mempty) et) shape Unique +returnType (Array et shape Nonunique) ds args = + Array (arrayElemReturnType et ds args) shape Nonunique +returnType (Record fs) ds args = + Record $ fmap (\et -> returnType et ds args) fs +returnType (Prim t) _ _ = Prim t +returnType (TypeVar () Unique t targs) _ _ = + TypeVar mempty Unique t $ map (bimap id (const mempty)) targs +returnType (TypeVar () Nonunique t targs) ds args = + TypeVar als Nonunique t $ map (\arg -> typeArgReturnType arg ds args) targs + where als = mconcat $ map aliases $ zipWith maskAliases args ds +returnType (Arrow _ v t1 t2) ds args = + Arrow als v (bimap id (const mempty) t1) (returnType t2 ds args) + where als = foldMap aliases $ zipWith maskAliases args ds + +typeArgReturnType :: TypeArg shape () -> [Diet] -> [CompType] + -> TypeArg shape Names +typeArgReturnType (TypeArgDim v loc) _ _ = + TypeArgDim v loc +typeArgReturnType (TypeArgType t loc) ds args = + TypeArgType (returnType t ds args) loc + +arrayElemReturnType :: ArrayElemTypeBase dim () + -> [Diet] + -> [CompType] + -> ArrayElemTypeBase dim Names +arrayElemReturnType (ArrayPrimElem bt ()) ds args = + ArrayPrimElem bt als + where als = mconcat $ map aliases $ zipWith maskAliases args ds +arrayElemReturnType (ArrayPolyElem bt targs ()) ds args = + ArrayPolyElem bt (map (\arg -> typeArgReturnType arg ds args) targs) als + where als = mconcat $ map aliases $ zipWith maskAliases args ds +arrayElemReturnType (ArrayRecordElem et) ds args = + ArrayRecordElem $ fmap (\t -> recordArrayElemReturnType t ds args) et + +recordArrayElemReturnType :: RecordArrayElemTypeBase dim () + -> [Diet] + -> [CompType] + -> RecordArrayElemTypeBase dim Names +recordArrayElemReturnType (RecordArrayElem et) ds args = + RecordArrayElem $ arrayElemReturnType et ds args +recordArrayElemReturnType (RecordArrayArrayElem et shape u) ds args = + RecordArrayArrayElem (arrayElemReturnType et ds args) shape u + +-- | Is the type concrete, i.e, without any type variables or function arrows? +concreteType :: TypeBase f vn -> Bool +concreteType Prim{} = True +concreteType TypeVar{} = False +concreteType Arrow{} = False +concreteType (Record ts) = all concreteType ts +concreteType (Array at _ _) = concreteArrayType at + where concreteArrayType ArrayPrimElem{} = True + concreteArrayType ArrayPolyElem{} = False + concreteArrayType (ArrayRecordElem ts) = all concreteRecordArrayElem ts + + concreteRecordArrayElem (RecordArrayElem et) = concreteArrayType et + concreteRecordArrayElem (RecordArrayArrayElem et _ _) = concreteArrayType et + +-- | @orderZero t@ is 'True' if the argument type has order 0, i.e., it is not +-- a function type, does not contain a function type as a subcomponent, and may +-- not be instantiated with a function type. +orderZero :: TypeBase dim as -> Bool +orderZero (Prim _) = True +orderZero Array{} = True +orderZero (Record fs) = all orderZero $ M.elems fs +orderZero TypeVar{} = True +orderZero Arrow{} = False + +-- | Extract all the shape names that occur in a given pattern. +patternDimNames :: PatternBase Info VName -> Names +patternDimNames (TuplePattern ps _) = foldMap patternDimNames ps +patternDimNames (RecordPattern fs _) = foldMap (patternDimNames . snd) fs +patternDimNames (PatternParens p _) = patternDimNames p +patternDimNames (Id _ (Info tp) _) = typeDimNames tp +patternDimNames (Wildcard (Info tp) _) = typeDimNames tp +patternDimNames (PatternAscription p (TypeDecl _ (Info t)) _) = + patternDimNames p <> typeDimNames t + +-- | Extract all the shape names that occur in a given type. +typeDimNames :: TypeBase (DimDecl VName) als -> Names +typeDimNames = foldMap dimName . nestedDims + where dimName :: DimDecl VName -> Names + dimName (NamedDim qn) = S.singleton $ qualLeaf qn + dimName _ = mempty + +-- | @patternOrderZero pat@ is 'True' if all of the types in the given pattern +-- have order 0. +patternOrderZero :: PatternBase Info vn -> Bool +patternOrderZero pat = case pat of + TuplePattern ps _ -> all patternOrderZero ps + RecordPattern fs _ -> all (patternOrderZero . snd) fs + PatternParens p _ -> patternOrderZero p + Id _ (Info t) _ -> orderZero t + Wildcard (Info t) _ -> orderZero t + PatternAscription p _ _ -> patternOrderZero p + +-- | The set of identifiers bound in a pattern. +patIdentSet :: (Functor f, Ord vn) => PatternBase f vn -> S.Set (IdentBase f vn) +patIdentSet (Id v t loc) = S.singleton $ Ident v (removeShapeAnnotations <$> t) loc +patIdentSet (PatternParens p _) = patIdentSet p +patIdentSet (TuplePattern pats _) = mconcat $ map patIdentSet pats +patIdentSet (RecordPattern fs _) = mconcat $ map (patIdentSet . snd) fs +patIdentSet Wildcard{} = mempty +patIdentSet (PatternAscription p _ _) = patIdentSet p + +-- | The type of values bound by the pattern. +patternType :: PatternBase Info VName -> CompType +patternType (Wildcard (Info t) _) = removeShapeAnnotations t +patternType (PatternParens p _) = patternType p +patternType (Id _ (Info t) _) = removeShapeAnnotations t +patternType (TuplePattern pats _) = tupleRecord $ map patternType pats +patternType (RecordPattern fs _) = Record $ patternType <$> M.fromList fs +patternType (PatternAscription p _ _) = patternType p + +-- | The type matched by the pattern, including shape declarations if present. +patternStructType :: PatternBase Info VName -> StructType +patternStructType (PatternAscription p _ _) = patternStructType p +patternStructType (PatternParens p _) = patternStructType p +patternStructType (Id _ (Info t) _) = t `setAliases` () +patternStructType (TuplePattern ps _) = tupleRecord $ map patternStructType ps +patternStructType (RecordPattern fs _) = Record $ patternStructType <$> M.fromList fs +patternStructType (Wildcard (Info t) _) = vacuousShapeAnnotations $ toStruct t + +-- | When viewed as a function parameter, does this pattern correspond +-- to a named parameter of some type? +patternParam :: PatternBase Info VName -> (Maybe VName, StructType) +patternParam (PatternParens p _) = + patternParam p +patternParam (PatternAscription (Id v _ _) td _) = + (Just v, unInfo $ expandedType td) +patternParam p = + (Nothing, patternStructType p) + +-- | Remove all shape annotations from a pattern, leaving them unnamed +-- instead. +patternNoShapeAnnotations :: PatternBase Info VName -> PatternBase Info VName +patternNoShapeAnnotations (PatternAscription p (TypeDecl te (Info t)) loc) = + PatternAscription (patternNoShapeAnnotations p) + (TypeDecl te $ Info $ vacuousShapeAnnotations t) loc +patternNoShapeAnnotations (PatternParens p loc) = + PatternParens (patternNoShapeAnnotations p) loc +patternNoShapeAnnotations (Id v (Info t) loc) = + Id v (Info $ vacuousShapeAnnotations t) loc +patternNoShapeAnnotations (TuplePattern ps loc) = + TuplePattern (map patternNoShapeAnnotations ps) loc +patternNoShapeAnnotations (RecordPattern ps loc) = + RecordPattern (map (fmap patternNoShapeAnnotations) ps) loc +patternNoShapeAnnotations (Wildcard (Info t) loc) = + Wildcard (Info (vacuousShapeAnnotations t)) loc + +-- | Names of primitive types to types. This is only valid if no +-- shadowing is going on, but useful for tools. +namesToPrimTypes :: M.Map Name PrimType +namesToPrimTypes = M.fromList + [ (nameFromString $ pretty t, t) | + t <- Bool : + map Signed [minBound..maxBound] ++ + map Unsigned [minBound..maxBound] ++ + map FloatType [minBound..maxBound] ] + +-- | The nature of something predefined. These can either be +-- monomorphic or overloaded. An overloaded builtin is a list valid +-- types it can be instantiated with, to the parameter and result +-- type, with 'Nothing' representing the overloaded parameter type. +data Intrinsic = IntrinsicMonoFun [PrimType] PrimType + | IntrinsicOverloadedFun [PrimType] [Maybe PrimType] (Maybe PrimType) + | IntrinsicPolyFun [TypeParamBase VName] [TypeBase () ()] (TypeBase () ()) + | IntrinsicType PrimType + | IntrinsicEquality -- Special cased. + | IntrinsicOpaque + +-- | A map of all built-ins. +intrinsics :: M.Map VName Intrinsic +intrinsics = M.fromList $ zipWith namify [10..] $ + + map primFun (M.toList Primitive.primFuns) ++ + + [ ("~", IntrinsicOverloadedFun + (map Signed [minBound..maxBound] ++ + map Unsigned [minBound..maxBound]) + [Nothing] Nothing) + , ("!", IntrinsicMonoFun [Bool] Bool)] ++ + + [("opaque", IntrinsicOpaque)] ++ + + map unOpFun Primitive.allUnOps ++ + + map binOpFun Primitive.allBinOps ++ + + map cmpOpFun Primitive.allCmpOps ++ + + map convOpFun Primitive.allConvOps ++ + + map signFun Primitive.allIntTypes ++ + + map unsignFun Primitive.allIntTypes ++ + + map intrinsicType (map Signed [minBound..maxBound] ++ + map Unsigned [minBound..maxBound] ++ + map FloatType [minBound..maxBound] ++ + [Bool]) ++ + + -- The reason for the loop formulation is to ensure that we + -- get a missing case warning if we forget a case. + mapMaybe mkIntrinsicBinOp [minBound..maxBound] ++ + + [("flatten", IntrinsicPolyFun [tp_a] + [Array (ArrayPolyElem tv_a' [] ()) (rank 2) Nonunique] $ + Array (ArrayPolyElem tv_a' [] ()) (rank 1) Nonunique), + ("unflatten", IntrinsicPolyFun [tp_a] + [Prim $ Signed Int32, + Prim $ Signed Int32, + Array (ArrayPolyElem tv_a' [] ()) (rank 1) Nonunique] $ + Array (ArrayPolyElem tv_a' [] ()) (rank 2) Nonunique), + + ("concat", IntrinsicPolyFun [tp_a] + [arr_a, arr_a] uarr_a), + ("rotate", IntrinsicPolyFun [tp_a] + [Prim $ Signed Int32, arr_a] arr_a), + ("transpose", IntrinsicPolyFun [tp_a] [arr_a] arr_a), + + ("cmp_threshold", IntrinsicPolyFun [] + [Prim $ Signed Int32, + Array (ArrayPrimElem (Signed Int32) ()) (rank 1) Nonunique] $ + Prim Bool), + + ("scatter", IntrinsicPolyFun [tp_a] + [Array (ArrayPolyElem tv_a' [] ()) (rank 1) Unique, + Array (ArrayPrimElem (Signed Int32) ()) (rank 1) Nonunique, + Array (ArrayPolyElem tv_a' [] ()) (rank 1) Nonunique] $ + Array (ArrayPolyElem tv_a' [] ()) (rank 1) Unique), + + ("zip", IntrinsicPolyFun [tp_a, tp_b] [arr_a, arr_b] arr_a_b), + ("unzip", IntrinsicPolyFun [tp_a, tp_b] [arr_a_b] t_arr_a_arr_b), + + ("gen_reduce", IntrinsicPolyFun [tp_a] + [uarr_a, + t_a `arr` (t_a `arr` t_a), + t_a, + Array (ArrayPrimElem (Signed Int32) ()) (rank 1) Nonunique, + arr_a] + uarr_a), + + ("map", IntrinsicPolyFun [tp_a, tp_b] [t_a `arr` t_b, arr_a] uarr_b), + + ("reduce", IntrinsicPolyFun [tp_a] + [t_a `arr` (t_a `arr` t_a), t_a, arr_a] t_a), + + ("reduce_comm", IntrinsicPolyFun [tp_a] + [t_a `arr` (t_a `arr` t_a), t_a, arr_a] t_a), + + ("scan", IntrinsicPolyFun [tp_a] + [t_a `arr` (t_a `arr` t_a), t_a, arr_a] uarr_a), + + ("partition", + IntrinsicPolyFun [tp_a] + [Prim (Signed Int32), t_a `arr` Prim (Signed Int32), arr_a] $ + tupleRecord [uarr_a, Array (ArrayPrimElem (Signed Int32) ()) (rank 1) Unique]), + + ("stream_map", + IntrinsicPolyFun [tp_a, tp_b] [arr_a `arr` arr_b, arr_a] uarr_b), + + ("stream_map_per", + IntrinsicPolyFun [tp_a, tp_b] [arr_a `arr` arr_b, arr_a] uarr_b), + + ("stream_red", + IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), arr_a `arr` t_b, arr_a] t_b), + + ("stream_red_per", + IntrinsicPolyFun [tp_a, tp_b] [t_b `arr` (t_b `arr` t_b), arr_a `arr` t_b, arr_a] t_b), + + + ("trace", IntrinsicPolyFun [tp_a] [t_a] t_a), + ("break", IntrinsicPolyFun [tp_a] [t_a] t_a)] + + where tv_a = VName (nameFromString "a") 0 + tv_a' = typeName tv_a + t_a = TypeVar () Nonunique tv_a' [] + arr_a = Array (ArrayPolyElem tv_a' [] ()) (rank 1) Nonunique + uarr_a = Array (ArrayPolyElem tv_a' [] ()) (rank 1) Unique + tp_a = TypeParamType Unlifted tv_a noLoc + + tv_b = VName (nameFromString "b") 1 + tv_b' = typeName tv_b + t_b = TypeVar () Nonunique tv_b' [] + arr_b = Array (ArrayPolyElem tv_b' [] ()) (rank 1) Nonunique + uarr_b = Array (ArrayPolyElem tv_b' [] ()) (rank 1) Unique + tp_b = TypeParamType Unlifted tv_b noLoc + + arr_a_b = Array (ArrayRecordElem (M.fromList $ zip tupleFieldNames + [RecordArrayElem $ ArrayPolyElem tv_a' [] (), + RecordArrayElem $ ArrayPolyElem tv_b' [] ()])) + (rank 1) Nonunique + t_arr_a_arr_b = Record $ M.fromList $ zip tupleFieldNames [arr_a, arr_b] + + arr = Arrow mempty Nothing + + namify i (k,v) = (VName (nameFromString k) i, v) + + primFun (name, (ts,t, _)) = + (name, IntrinsicMonoFun (map unPrim ts) $ unPrim t) + + unOpFun bop = (pretty bop, IntrinsicMonoFun [t] t) + where t = unPrim $ Primitive.unOpType bop + + binOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] t) + where t = unPrim $ Primitive.binOpType bop + + cmpOpFun bop = (pretty bop, IntrinsicMonoFun [t, t] Bool) + where t = unPrim $ Primitive.cmpOpType bop + + convOpFun cop = (pretty cop, IntrinsicMonoFun [unPrim ft] $ unPrim tt) + where (ft, tt) = Primitive.convOpType cop + + signFun t = ("sign_" ++ pretty t, IntrinsicMonoFun [Unsigned t] $ Signed t) + + unsignFun t = ("unsign_" ++ pretty t, IntrinsicMonoFun [Signed t] $ Unsigned t) + + unPrim (Primitive.IntType t) = Signed t + unPrim (Primitive.FloatType t) = FloatType t + unPrim Primitive.Bool = Bool + unPrim Primitive.Cert = Bool + + intrinsicType t = (pretty t, IntrinsicType t) + + anyIntType = map Signed [minBound..maxBound] ++ + map Unsigned [minBound..maxBound] + anyNumberType = anyIntType ++ + map FloatType [minBound..maxBound] + anyPrimType = Bool : anyNumberType + + mkIntrinsicBinOp :: BinOp -> Maybe (String, Intrinsic) + mkIntrinsicBinOp op = do op' <- intrinsicBinOp op + return (pretty op, op') + + binOp ts = Just $ IntrinsicOverloadedFun ts [Nothing, Nothing] Nothing + ordering = Just $ IntrinsicOverloadedFun anyPrimType [Nothing, Nothing] (Just Bool) + + intrinsicBinOp Plus = binOp anyNumberType + intrinsicBinOp Minus = binOp anyNumberType + intrinsicBinOp Pow = binOp anyNumberType + intrinsicBinOp Times = binOp anyNumberType + intrinsicBinOp Divide = binOp anyNumberType + intrinsicBinOp Mod = binOp anyNumberType + intrinsicBinOp Quot = binOp anyIntType + intrinsicBinOp Rem = binOp anyIntType + intrinsicBinOp ShiftR = binOp anyIntType + intrinsicBinOp ShiftL = binOp anyIntType + intrinsicBinOp Band = binOp anyIntType + intrinsicBinOp Xor = binOp anyIntType + intrinsicBinOp Bor = binOp anyIntType + intrinsicBinOp LogAnd = Just $ IntrinsicMonoFun [Bool,Bool] Bool + intrinsicBinOp LogOr = Just $ IntrinsicMonoFun [Bool,Bool] Bool + intrinsicBinOp Equal = Just IntrinsicEquality + intrinsicBinOp NotEqual = Just IntrinsicEquality + intrinsicBinOp Less = ordering + intrinsicBinOp Leq = ordering + intrinsicBinOp Greater = ordering + intrinsicBinOp Geq = ordering + intrinsicBinOp _ = Nothing + +-- | The largest tag used by an intrinsic - this can be used to +-- determine whether a 'VName' refers to an intrinsic or a user-defined name. +maxIntrinsicTag :: Int +maxIntrinsicTag = maximum $ map baseTag $ M.keys intrinsics + +-- | Create a name with no qualifiers from a name. +qualName :: v -> QualName v +qualName = QualName [] + +-- | Add another qualifier (at the head) to a qualified name. +qualify :: v -> QualName v -> QualName v +qualify k (QualName ks v) = QualName (k:ks) v + +-- | Create a type name name with no qualifiers from a 'VName'. +typeName :: VName -> TypeName +typeName = typeNameFromQualName . qualName + +-- | The modules imported by a Futhark program. +progImports :: ProgBase f vn -> [(String,SrcLoc)] +progImports = concatMap decImports . progDecs + +-- | The modules imported by a single declaration. +decImports :: DecBase f vn -> [(String,SrcLoc)] +decImports (OpenDec x _ _) = modExpImports x +decImports (ModDec md) = modExpImports $ modExp md +decImports SigDec{} = [] +decImports TypeDec{} = [] +decImports ValDec{} = [] +decImports (LocalDec d _) = decImports d + +modExpImports :: ModExpBase f vn -> [(String,SrcLoc)] +modExpImports ModVar{} = [] +modExpImports (ModParens p _) = modExpImports p +modExpImports (ModImport f _ loc) = [(f,loc)] +modExpImports (ModDecs ds _) = concatMap decImports ds +modExpImports (ModApply _ me _ _ _) = modExpImports me +modExpImports (ModAscript me _ _ _) = modExpImports me +modExpImports ModLambda{} = [] + +-- | The set of module types used in any exported (non-local) +-- declaration. +progModuleTypes :: Ord vn => ProgBase f vn -> S.Set vn +progModuleTypes = mconcat . map onDec . progDecs + where onDec (OpenDec x _ _) = onModExp x + onDec (ModDec md) = + maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md) + onDec SigDec{} = mempty + onDec TypeDec{} = mempty + onDec ValDec{} = mempty + onDec (LocalDec _ _) = mempty + + onModExp ModVar{} = mempty + onModExp (ModParens p _) = onModExp p + onModExp ModImport {} = mempty + onModExp (ModDecs ds _) = mconcat $ map onDec ds + onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2 + onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se + onModExp (ModLambda p r me _) = + onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me + + onModParam = onSigExp . modParamType + + onSigExp (SigVar v _) = S.singleton $ qualLeaf v + onSigExp (SigParens e _) = onSigExp e + onSigExp SigSpecs{} = mempty + onSigExp (SigWith e _ _) = onSigExp e + onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2 + +-- | Extract a leading @((name, namespace, file), remainder)@ from a +-- documentation comment string. These are formatted as +-- \`name\`\@namespace[\@file]. Let us hope that this pattern does not occur +-- anywhere else. +identifierReference :: String -> Maybe ((String, String, Maybe FilePath), String) +identifierReference ('`' : s) + | (identifier, '`' : '@' : s') <- break (=='`') s, + (namespace, s'') <- span isAlpha s', + not $ null namespace = + case s'' of + '@' : '"' : s''' + | (file, '"' : s'''') <- span (/= '"') s''' -> + Just ((identifier, namespace, Just file), s'''') + _ -> Just ((identifier, namespace, Nothing), s'') + +identifierReference _ = Nothing + +-- | Find all the identifier references in a string. +identifierReferences :: String -> [(String, String, Maybe FilePath)] +identifierReferences [] = [] +identifierReferences s + | Just (ref, s') <- identifierReference s = + ref : identifierReferences s' +identifierReferences (_:s') = + identifierReferences s' + +-- | Given an operator name, return the operator that determines its +-- syntactical properties. +leadingOperator :: Name -> BinOp +leadingOperator s = maybe Backtick snd $ find ((`isPrefixOf` s') . fst) $ + sortBy (flip $ comparing $ length . fst) $ + zip (map pretty operators) operators + where s' = nameToString s + operators :: [BinOp] + operators = [minBound..maxBound::BinOp] + +-- | A type with no aliasing information but shape annotations. +type UncheckedType = TypeBase (ShapeDecl Name) () + +type UncheckedTypeExp = TypeExp Name + +-- | An array element type with no aliasing information. +type UncheckedArrayElemType = ArrayElemTypeBase (ShapeDecl Name) () + +-- | A type declaration with no expanded type. +type UncheckedTypeDecl = TypeDeclBase NoInfo Name + +-- | An identifier with no type annotations. +type UncheckedIdent = IdentBase NoInfo Name + +-- | An index with no type annotations. +type UncheckedDimIndex = DimIndexBase NoInfo Name + +-- | An expression with no type annotations. +type UncheckedExp = ExpBase NoInfo Name + +-- | A module expression with no type annotations. +type UncheckedModExp = ModExpBase NoInfo Name + +-- | A module type expression with no type annotations. +type UncheckedSigExp = SigExpBase NoInfo Name + +-- | A type parameter with no type annotations. +type UncheckedTypeParam = TypeParamBase Name + +-- | A pattern with no type annotations. +type UncheckedPattern = PatternBase NoInfo Name + +-- | A function declaration with no type annotations. +type UncheckedValBind = ValBindBase NoInfo Name + +-- | A declaration with no type annotations. +type UncheckedDec = DecBase NoInfo Name + +-- | A Futhark program with no type annotations. +type UncheckedProg = ProgBase NoInfo Name diff --git a/src/Language/Futhark/Core.hs b/src/Language/Futhark/Core.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Core.hs @@ -0,0 +1,139 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +-- | This module contains very basic definitions for Futhark - so basic, +-- that they can be shared between the internal and external +-- representation. +module Language.Futhark.Core + ( Uniqueness(..) + , StreamOrd(..) + , Commutativity(..) + + -- * Location utilities + , locStr + + -- * Name handling + , Name + , nameToString + , nameFromString + , nameToText + , nameFromText + , VName(..) + , baseTag + , baseName + , baseString + , pretty + -- * Special identifiers + , defaultEntryPoint + + -- * Integer re-export + , Int8, Int16, Int32, Int64 + , Word8, Word16, Word32, Word64 + ) + +where + +import Data.Int (Int8, Int16, Int32, Int64) +import Data.String +import Data.Word (Word8, Word16, Word32, Word64) +import Data.Loc +import qualified Data.Semigroup as Sem +import qualified Data.Text as T + +import Futhark.Util.Pretty + +-- | The uniqueness attribute of a type. This essentially indicates +-- whether or not in-place modifications are acceptable. With respect +-- to ordering, 'Unique' is greater than 'Nonunique'. +data Uniqueness = Nonunique -- ^ May have references outside current function. + | Unique -- ^ No references outside current function. + deriving (Eq, Ord, Show) + +instance Sem.Semigroup Uniqueness where + (<>) = min + +instance Monoid Uniqueness where + mempty = Unique + mappend = (Sem.<>) + +instance Pretty Uniqueness where + ppr Unique = star + ppr Nonunique = empty + +data StreamOrd = InOrder + | Disorder + deriving (Eq, Ord, Show) + +-- | 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 Sem.Semigroup Commutativity where + (<>) = min + +instance Monoid Commutativity where + mempty = Commutative + mappend = (Sem.<>) + +-- | The name of the default program entry point (main). +defaultEntryPoint :: Name +defaultEntryPoint = nameFromString "main" + +-- | The abstract (not really) type representing names in the Futhark +-- compiler. 'String's, being lists of characters, are very slow, +-- while 'T.Text's are based on byte-arrays. +newtype Name = Name T.Text + deriving (Show, Eq, Ord, IsString, Sem.Semigroup) + +instance Pretty Name where + ppr = text . nameToString + +-- | Convert a name to the corresponding list of characters. +nameToString :: Name -> String +nameToString (Name t) = T.unpack t + +-- | Convert a list of characters to the corresponding name. +nameFromString :: String -> Name +nameFromString = Name . T.pack + +-- | Convert a name to the corresponding 'T.Text'. +nameToText :: Name -> T.Text +nameToText (Name t) = t + +-- | Convert a 'T.Text' to the corresponding name. +nameFromText :: T.Text -> Name +nameFromText = Name + +-- | A human-readable location string, of the form +-- @filename:lineno:columnno@. +locStr :: SrcLoc -> String +locStr (SrcLoc NoLoc) = "unknown location" +locStr (SrcLoc (Loc (Pos file line1 col1 _) (Pos _ line2 col2 _))) = + -- Assume that both positions are in the same file (what would the + -- alternative mean?) + file ++ ":" ++ show line1 ++ ":" ++ show col1 + ++ "-" ++ show line2 ++ ":" ++ show col2 + +-- | A name tagged with some integer. Only the integer is used in +-- comparisons, no matter the type of @vn@. +data VName = VName !Name !Int + deriving (Show) + +-- | Return the tag contained in the 'VName'. +baseTag :: VName -> Int +baseTag (VName _ tag) = tag + +-- | Return the name contained in the 'VName'. +baseName :: VName -> Name +baseName (VName vn _) = vn + +-- | Return the base 'Name' converted to a string. +baseString :: VName -> String +baseString = nameToString . baseName + +instance Eq VName where + VName _ x == VName _ y = x == y + +instance Ord VName where + VName _ x `compare` VName _ y = x `compare` y diff --git a/src/Language/Futhark/Futlib.hs b/src/Language/Futhark/Futlib.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Futlib.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} +-- | The Futhark basis library embedded embedded as strings read during +-- compilation of the Futhark compiler. The advantage is that the +-- standard library can be accessed without reading it from disk, thus +-- saving users from include path headaches. +module Language.Futhark.Futlib (futlib, prelude) where + +import Data.FileEmbed +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import qualified System.FilePath.Posix as Posix + +import Futhark.Util (toPOSIX) + +-- | Futlib embedded as 'T.Text' values, one for every file. +futlib :: [(Posix.FilePath, T.Text)] +futlib = map fixup futlib_bs + where futlib_bs = $(embedDir "futlib") + fixup (path, s) = ("/futlib" Posix. toPOSIX path, T.decodeUtf8 s) + +-- The files intended to be implicitly imported into every Futhark +-- program. Make sure it does not depend on anything too big to be +-- serialised efficiently. +prelude :: [String] +prelude = map ("/futlib/"++) ["prelude"] diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Interpreter.hs @@ -0,0 +1,1153 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Language.Futhark.Interpreter + ( Ctx(..) + , Env(..) + , InterpreterError + , initialCtx + , interpretExp + , interpretDec + , interpretImport + , interpretFunction + , ExtOp(..) + , typeEnv + , Value (ValuePrim, ValueArray, ValueRecord) + , mkArray + , fromTuple + , isEmptyArray + ) where + +import Control.Monad.Free.Church +import Control.Monad.Except +import Control.Monad.Reader +import qualified Control.Monad.Fail as Fail +import Data.Array +import Data.Bifunctor (bimap) +import Data.List hiding (break) +import Data.Maybe +import qualified Data.Map as M +import qualified Data.Semigroup as Sem +import Data.Monoid +import Data.Loc + +import Language.Futhark hiding (Value) +import Futhark.Representation.Primitive (intValue, floatValue) +import qualified Futhark.Representation.Primitive as P +import qualified Language.Futhark.Semantic as T + +import Futhark.Util.Pretty hiding (apply, bool, stack) +import Futhark.Util (chunk, splitFromEnd, maybeHead) + +import Prelude hiding (mod, break) + +data ExtOp a = ExtOpTrace SrcLoc String a + | ExtOpBreak [SrcLoc] Ctx T.Env a + | ExtOpError InterpreterError + +instance Functor ExtOp where + fmap f (ExtOpTrace w s x) = ExtOpTrace w s $ f x + fmap f (ExtOpBreak w ctx env x) = ExtOpBreak w ctx env $ f x + fmap _ (ExtOpError err) = ExtOpError err + +data StackFrame = StackFrame { stackFrameSrcLoc :: SrcLoc + , stackFrameEnv :: Env + } + +type Stack = [StackFrame] + +-- | The monad in which evaluation takes place. +newtype EvalM a = EvalM (ReaderT (Stack, M.Map FilePath Env) + (F ExtOp) a) + deriving (Monad, Applicative, Functor, + MonadFree ExtOp, + MonadReader (Stack, M.Map FilePath Env)) + +instance Fail.MonadFail EvalM where + fail = error + +runEvalM :: M.Map FilePath Env -> EvalM a -> F ExtOp a +runEvalM imports (EvalM m) = runReaderT m (mempty, imports) + +stacking :: SrcLoc -> Env -> EvalM a -> EvalM a +stacking loc env = local $ \(ss, imports) -> + if isNoLoc loc then (ss, imports) else (StackFrame loc env:ss, imports) + where isNoLoc :: SrcLoc -> Bool + isNoLoc = (==NoLoc) . locOf + +stacktrace :: EvalM [SrcLoc] +stacktrace = asks $ map stackFrameSrcLoc . reverse . fst + +stacktraceTop :: EvalM SrcLoc +stacktraceTop = fromMaybe noLoc . maybeHead <$> stacktrace + +lookupImport :: FilePath -> EvalM (Maybe Env) +lookupImport f = asks $ M.lookup f . snd + +-- | A fully evaluated Futhark value. +data Value = ValuePrim !PrimValue + | ValueArray !(Array Int Value) + | ValueRecord (M.Map Name Value) + | ValueFun (Value -> EvalM Value) + +instance Eq Value where + ValuePrim x == ValuePrim y = x == y + ValueArray x == ValueArray y = x == y + ValueRecord x == ValueRecord y = x == y + _ == _ = False + +prettyRecord :: Pretty a => M.Map Name a -> Doc +prettyRecord m + | Just vs <- areTupleFields m = + parens $ commasep $ map ppr vs + | otherwise = + braces $ commasep $ map field $ M.toList m + where field (k, v) = ppr k <+> equals <+> ppr v + +instance Pretty Value where + ppr (ValuePrim v) = ppr v + ppr (ValueArray a) = + let elements = elems a -- [Value] + (x:_) = elements + separator = case x of + (ValueArray _) -> comma <> line + _ -> comma <> space + in brackets $ cat $ punctuate separator (map ppr elements) + + ppr (ValueRecord m) = prettyRecord m + ppr ValueFun{} = text "#" + +-- | Create an array value; failing if that would result in an +-- irregular array. +mkArray :: [Value] -> Maybe Value +mkArray vs = + case vs of [] -> Just $ toArray' vs + v:_ | all ((==valueShape v) . valueShape) vs -> Just $ toArray' vs + | otherwise -> Nothing + +-- | A shape is a tree to accomodate the case of records. +data Shape = ShapeDim Int32 Shape + | ShapeLeaf + | ShapeRecord (M.Map Name Shape) + deriving (Eq, Show) + +instance Pretty Shape where + ppr ShapeLeaf = mempty + ppr (ShapeDim d s) = brackets (ppr d) <> ppr s + ppr (ShapeRecord m) = prettyRecord m + +emptyShape :: Shape -> Bool +emptyShape ShapeLeaf = False +emptyShape (ShapeDim d s) = d == 0 || emptyShape s +emptyShape (ShapeRecord fs) = any emptyShape fs + +valueShape :: Value -> Shape +valueShape (ValueArray arr) = ShapeDim (arrayLength arr) $ + case elems arr of + [] -> ShapeLeaf + v:_ -> valueShape v +valueShape (ValueRecord fs) = ShapeRecord $ M.map valueShape fs +valueShape _ = ShapeLeaf + +isEmptyArray :: Value -> Bool +isEmptyArray = emptyShape . valueShape + +arrayLength :: Integral int => Array Int Value -> int +arrayLength = fromIntegral . (+1) . snd . bounds + +toTuple :: [Value] -> Value +toTuple = ValueRecord . M.fromList . zip tupleFieldNames + +fromTuple :: Value -> Maybe [Value] +fromTuple (ValueRecord m) = areTupleFields m +fromTuple _ = Nothing + +asInteger :: Value -> Integer +asInteger (ValuePrim (SignedValue v)) = P.valueIntegral v +asInteger (ValuePrim (UnsignedValue v)) = + toInteger (P.valueIntegral (P.doZExt v Int64) :: Word64) +asInteger v = error $ "Unexpectedly not an integer: " ++ pretty v + +asInt :: Value -> Int +asInt = fromIntegral . asInteger + +asSigned :: Value -> IntValue +asSigned (ValuePrim (SignedValue v)) = v +asSigned v = error $ "Unexpected not a signed integer: " ++ pretty v + +asInt32 :: Value -> Int32 +asInt32 = fromIntegral . asInteger + +asBool :: Value -> Bool +asBool (ValuePrim (BoolValue x)) = x +asBool v = error $ "Unexpectedly not an integer: " ++ pretty v + +lookupInEnv :: (Env -> M.Map VName x) + -> QualName VName -> Env -> Maybe x +lookupInEnv onEnv qv env = f env $ qualQuals qv + where f m (q:qs) = + case M.lookup q $ envTerm m of + Just (TermModule (Module mod)) -> f mod qs + _ -> Nothing + f m [] = M.lookup (qualLeaf qv) $ onEnv m + +lookupVar :: QualName VName -> Env -> Maybe TermBinding +lookupVar = lookupInEnv envTerm + +lookupType :: QualName VName -> Env -> Maybe T.TypeBinding +lookupType = lookupInEnv envType + +-- | A TermValue with a 'Nothing' type annotation is an intrinsic. +data TermBinding = TermValue (Maybe T.BoundV) Value + | TermModule Module + +data Module = Module Env + | ModuleFun (Module -> EvalM Module) + +data Env = Env { envTerm :: M.Map VName TermBinding + , envType :: M.Map VName T.TypeBinding + } + +instance Monoid Env where + mempty = Env mempty mempty + mappend = (Sem.<>) + +instance Sem.Semigroup Env where + Env vm1 tm1 <> Env vm2 tm2 = Env (vm1 <> vm2) (tm1 <> tm2) + +newtype InterpreterError = InterpreterError String + +valEnv :: M.Map VName (Maybe T.BoundV, Value) -> Env +valEnv m = Env { envTerm = M.map (uncurry TermValue) m + , envType = mempty + } + +modEnv :: M.Map VName Module -> Env +modEnv m = Env { envTerm = M.map TermModule m + , envType = mempty + } + +instance Show InterpreterError where + show (InterpreterError s) = s + +bad :: SrcLoc -> Env -> String -> EvalM a +bad loc env s = stacking loc env $ do + ss <- map locStr <$> stacktrace + liftF $ ExtOpError $ InterpreterError $ "Error at " ++ intercalate " -> " ss ++ ": " ++ s + +trace :: Value -> EvalM () +trace v = do + top <- stacktraceTop + liftF $ ExtOpTrace top (pretty v) () + +typeEnv :: Env -> T.Env +typeEnv env = + -- FIXME: some shadowing issues are probably not right here. + let valMap (TermValue (Just t) _) = Just t + valMap _ = Nothing + vtable = M.mapMaybe valMap $ envTerm env + nameMap k | k `M.member` vtable = Just ((T.Term, baseName k), qualName k) + | otherwise = Nothing + in mempty { T.envNameMap = M.fromList $ mapMaybe nameMap $ M.keys $ envTerm env + , T.envVtable = vtable } + +break :: EvalM () +break = do + -- We don't want the env of the function that is calling + -- intrinsics.break, since that is just going to be the boring + -- wrapper function (intrinsics are never called directly). + -- This is why we go a step up the stack. + stack <- asks $ drop 1 . fst + case stack of + [] -> return () + top:_ -> do + let env = stackFrameEnv top + imports <- asks snd + liftF $ ExtOpBreak + (map stackFrameSrcLoc $ reverse stack) + (Ctx env imports) (typeEnv env) () + +fromArray :: Value -> [Value] +fromArray (ValueArray as) = elems as +fromArray v = error $ "Expected array value, but found: " ++ pretty v + +-- | This is where we enforce the regularity constraint for arrays. +toArray :: [Value] -> EvalM Value +toArray = maybe (bad noLoc mempty "irregular array") return . mkArray + +toArray' :: [Value] -> Value +toArray' vs = ValueArray (listArray (0, length vs - 1) vs) + +apply :: SrcLoc -> Env -> Value -> Value -> EvalM Value +apply loc env (ValueFun f) v = stacking loc env $ f v +apply _ _ f _ = error $ "Cannot apply non-function: " ++ pretty f + +apply2 :: SrcLoc -> Env -> Value -> Value -> Value -> EvalM Value +apply2 loc env f x y = stacking loc env $ do f' <- apply noLoc mempty f x + apply noLoc mempty f' y + +matchPattern :: Env -> Pattern -> Value + -> EvalM (M.Map VName (Maybe T.BoundV, Value)) +matchPattern env = matchPattern' env mempty + +matchPattern' :: Env -> M.Map VName (Maybe T.BoundV, Value) + -> Pattern -> Value + -> EvalM (M.Map VName (Maybe T.BoundV, Value)) +matchPattern' _ m (Id v (Info t) _) val = + pure $ M.insert v (Just $ T.BoundV [] $ toStruct t, val) m +matchPattern' env m (PatternParens p _) val = + matchPattern' env m p val +matchPattern' env m (TuplePattern ps _) (ValueRecord vs) = + foldM (\m' (p,v) -> matchPattern' env m' p v) m $ + zip ps (map snd $ sortFields vs) +matchPattern' env m (RecordPattern ps _) (ValueRecord vs) = + foldM (\m' (p,v) -> matchPattern' env m' p v) m $ + zip (map snd $ sortFields $ M.fromList ps) (map snd $ sortFields vs) +matchPattern' _ m Wildcard{} _ = pure m +matchPattern' env m (PatternAscription pat td loc) v = do + t <- evalType env $ unInfo $ expandedType td + case matchValueToType env m t v of + Left err -> bad loc env err + Right m' -> matchPattern' env m' pat v +matchPattern' _ _ pat v = + error $ "matchPattern': missing case for " ++ pretty pat ++ " and " ++ pretty v + +-- | For matching size annotations (the actual type will have been +-- verified by the type checker). It is assumed that previously +-- unbound names are in binding position here. +matchValueToType :: Env -> M.Map VName (Maybe T.BoundV, Value) + -> StructType + -> Value + -> Either String (M.Map VName (Maybe T.BoundV, Value)) + +-- Empty arrays always match. +matchValueToType env m t@(Array _ (ShapeDecl ds@(d:_)) _) val@(ValueArray arr) + | any zeroDim ds, emptyShape (valueShape val) = + Right $ m <> mconcat (map namedAreZero ds) + + | otherwise = + case d of + NamedDim v + | Just x <- look v -> + if x == arr_n + then continue m + else wrong $ "`" <> pretty v <> "` (" <> pretty x <> ")" + | otherwise -> + continue $ M.insert (qualLeaf v) + (Just $ T.BoundV [] $ Prim $ Signed Int32, + ValuePrim $ SignedValue $ Int32Value arr_n) + m + AnyDim -> continue m + ConstDim x + | fromIntegral x == arr_n -> continue m + | otherwise -> wrong $ pretty x + where arr_n = arrayLength arr + + look v + | Just (TermValue _ (ValuePrim (SignedValue (Int32Value x)))) <- + lookupVar v env = Just x + | Just (_, ValuePrim (SignedValue (Int32Value x))) <- + M.lookup (qualLeaf v) m = Just x + | otherwise = Nothing + + continue m' = case elems arr of + [] -> return m' + v:_ -> matchValueToType env m' (stripArray 1 t) v + + wrong x = Left $ "Size annotation " <> x <> + " does not match observed size " <> pretty arr_n <> "." + + zeroDim (NamedDim v) = isNothing (look v) || Just 0 == look v + zeroDim AnyDim = True + zeroDim (ConstDim x) = x == 0 + + namedAreZero (NamedDim v) + | isNothing $ look v = + M.singleton (qualLeaf v) (Just $ T.BoundV [] $ Prim $ Signed Int32, + ValuePrim $ SignedValue $ Int32Value 0) + | otherwise = + mempty + namedAreZero _ = mempty + +matchValueToType env m (Record fs) (ValueRecord arr) = + foldM (\m' (t, v) -> matchValueToType env m' t v) m $ + M.intersectionWith (,) fs arr + +matchValueToType _ m _ _ = return m + +data Indexing = IndexingFix Int32 + | IndexingSlice (Maybe Int32) (Maybe Int32) (Maybe Int32) + +instance Pretty Indexing where + ppr (IndexingFix i) = ppr i + ppr (IndexingSlice i j (Just s)) = + maybe mempty ppr i <> text ":" <> + maybe mempty ppr j <> text ":" <> + ppr s + ppr (IndexingSlice i (Just j) s) = + maybe mempty ppr i <> text ":" <> + ppr j <> + maybe mempty ((text ":" <>) . ppr) s + ppr (IndexingSlice i Nothing Nothing) = + maybe mempty ppr i <> text ":" + +indexesFor :: Maybe Int32 -> Maybe Int32 -> Maybe Int32 + -> Array Int Value -> Maybe [Int] +indexesFor start end stride arr + | (start', end', stride') <- slice, + end' == start' || signum' (end' - start') == signum' stride', + stride' /= 0, + is <- [start', start'+stride' .. end'-signum stride'], + all inBounds is = + Just $ map fromIntegral is + | otherwise = + Nothing + where n = arrayLength arr + + inBounds i = i >= 0 && i < n + + slice = + case (start, end, stride) of + (Just start', _, _) -> + let end' = fromMaybe n end + in (start', end', fromMaybe 1 stride) + (Nothing, Just end', _) -> + let start' = 0 + in (start', end', fromMaybe 1 stride) + (Nothing, Nothing, Just stride') -> + (if stride' > 0 then 0 else n-1, + if stride' > 0 then n else -1, + stride') + (Nothing, Nothing, Nothing) -> + (0, n, 1) + +-- | 'signum', but with 0 as 1. +signum' :: (Eq p, Num p) => p -> p +signum' x = if x == 0 then 1 else signum x + +indexArray :: [Indexing] -> Value -> Maybe Value +indexArray (IndexingFix i:is) (ValueArray arr) + | i >= 0, i < n = + indexArray is $ arr ! fromIntegral i + | otherwise = + Nothing + where n = arrayLength arr +indexArray (IndexingSlice start end stride:is) (ValueArray arr) = do + js <- indexesFor start end stride arr + toArray' <$> mapM (indexArray is . (arr!)) js +indexArray _ v = Just v + +updateArray :: [Indexing] -> Value -> Value -> Maybe Value +updateArray (IndexingFix i:is) (ValueArray arr) v + | i >= 0, i < n = do + v' <- updateArray is (arr ! i') v + Just $ ValueArray $ arr // [(i', v')] + | otherwise = + Nothing + where n = arrayLength arr + i' = fromIntegral i +updateArray (IndexingSlice start end stride:is) (ValueArray arr) (ValueArray v) = do + arr_is <- indexesFor start end stride arr + guard $ length arr_is == arrayLength v + let update arr' (i, v') = do + x <- updateArray is (arr!i) v' + return $ arr' // [(i, x)] + fmap ValueArray $ foldM update arr $ zip arr_is $ elems v +updateArray _ _ v = Just v + +evalDimIndex :: Env -> DimIndex -> EvalM Indexing +evalDimIndex env (DimFix x) = + IndexingFix . asInt32 <$> eval env x +evalDimIndex env (DimSlice start end stride) = + IndexingSlice <$> traverse (fmap asInt32 . eval env) start + <*> traverse (fmap asInt32 . eval env) end + <*> traverse (fmap asInt32 . eval env) stride + +evalIndex :: SrcLoc -> Env -> [Indexing] -> Value -> EvalM Value +evalIndex loc env is arr = do + let oob = bad loc env $ "Index [" <> intercalate ", " (map pretty is) <> + "] out of bounds for array of shape " <> + pretty (valueShape arr) <> "." + maybe oob return $ indexArray is arr + +evalTermVar :: Env -> QualName VName -> EvalM Value +evalTermVar env qv = + case lookupVar qv env of + Just (TermValue _ v) -> return v + _ -> error $ "`" <> pretty qv <> "` is not bound to a value." + +-- | Expand type based on information that was not available at +-- type-checking time (the structure of abstract types). +evalType :: Env -> StructType -> EvalM StructType +evalType _ (Prim pt) = return $ Prim pt +evalType env (Record fs) = Record <$> traverse (evalType env) fs +evalType env (Arrow () p t1 t2) = + Arrow () p <$> evalType env t1 <*> evalType env t2 +evalType env t@(Array _ shape u) = do + let et = stripArray (shapeRank shape) t + et' <- evalType env et + shape' <- traverse evalDim shape + return $ + fromMaybe (error "Cannot construct array after substitution") $ + arrayOf et' shape' u + where evalDim (NamedDim qn) + | Just (TermValue _ (ValuePrim (SignedValue (Int32Value x)))) <- + lookupVar qn env = + return $ ConstDim $ fromIntegral x + evalDim d = return d +evalType env t@(TypeVar () _ tn args) = + case lookupType (qualNameFromTypeName tn) env of + Just (T.TypeAbbr _ ps t') -> do + (substs, types) <- mconcat <$> zipWithM matchPtoA ps args + let onDim (NamedDim v) = fromMaybe (NamedDim v) $ M.lookup (qualLeaf v) substs + onDim d = d + if null ps then return $ bimap onDim id t' + else evalType (Env mempty types <> env) $ bimap onDim id t' + Nothing -> return t + where matchPtoA (TypeParamDim p _) (TypeArgDim (NamedDim qv) _) = + return (M.singleton p $ NamedDim qv, mempty) + matchPtoA (TypeParamDim p _) (TypeArgDim (ConstDim k) _) = + return (M.singleton p $ ConstDim k, mempty) + matchPtoA (TypeParamType l p _) (TypeArgType t' _) = do + t'' <- evalType env t' + return (mempty, M.singleton p $ T.TypeAbbr l [] t'') + matchPtoA _ _ = return mempty + +eval :: Env -> Exp -> EvalM Value + +eval _ (Literal v _) = return $ ValuePrim v + +eval env (Parens e _ ) = eval env e + +eval env (QualParens _ e _ ) = eval env e + +eval env (TupLit vs _) = toTuple <$> mapM (eval env) vs + +eval env (RecordLit fields _) = + ValueRecord . M.fromList <$> mapM evalField fields + where evalField (RecordFieldExplicit k e _) = do + v <- eval env e + return (k, v) + evalField (RecordFieldImplicit k t loc) = do + v <- eval env $ Var (qualName k) (vacuousShapeAnnotations <$> t) loc + return (baseName k, v) + +eval env (ArrayLit vs _ _) = toArray =<< mapM (eval env) vs + +eval env (Range start maybe_second end (Info t) _) = do + start' <- asInteger <$> eval env start + maybe_second' <- traverse (fmap asInteger . eval env) maybe_second + (end', dir) <- case end of + DownToExclusive e -> (,-1) . (+1) . asInteger <$> eval env e + ToInclusive e -> (,maybe 1 (signum . subtract start') maybe_second') . + asInteger <$> eval env e + UpToExclusive e -> (,1) . subtract 1 . asInteger <$> eval env e + + let second = fromMaybe (start' + dir) maybe_second' + step = second - start' + if step == 0 || dir /= signum step then toArray [] + else toArray $ map toInt [start',second..end'] + + where toInt = + case stripArray 1 t of + Prim (Signed t') -> + ValuePrim . SignedValue . intValue t' + Prim (Unsigned t') -> + ValuePrim . UnsignedValue . intValue t' + _ -> error $ "Nonsensical range type: " ++ show t + +eval env (Var qv _ _) = evalTermVar env qv + +eval env (Ascript e td loc) = do + v <- eval env e + t <- evalType env $ unInfo $ expandedType td + case matchValueToType env mempty t v of + Right _ -> return v + Left _ -> bad loc env $ "Value `" <> pretty v <> "` cannot match shape of type `" <> + pretty (declaredType td) <> "` (`" <> pretty t <> "`)." + +eval env (LetPat _ p e body _) = do + v <- eval env e + p_env <- valEnv <$> matchPattern env p v + eval (p_env <> env) body + +eval env (LetFun f (tparams, pats, _, Info ret, fbody) body loc) = do + v <- eval env $ Lambda tparams pats fbody Nothing (Info (mempty, ret)) loc + let ftype = T.BoundV [] $ foldr (uncurry (Arrow ()) . patternParam) ret pats + eval (valEnv (M.singleton f (Just ftype, v)) <> env) body + +eval _ (IntLit v (Info t) _) = + case t of + Prim (Signed it) -> + return $ ValuePrim $ SignedValue $ intValue it v + Prim (Unsigned it) -> + return $ ValuePrim $ UnsignedValue $ intValue it v + Prim (FloatType ft) -> + return $ ValuePrim $ FloatValue $ floatValue ft v + _ -> error $ "eval: nonsensical type for integer literal: " ++ pretty t + +eval _ (FloatLit v (Info t) _) = + case t of + Prim (FloatType ft) -> + return $ ValuePrim $ FloatValue $ floatValue ft v + _ -> error $ "eval: nonsensical type for float literal: " ++ pretty t + +eval env (BinOp op op_t (x, _) (y, _) _ loc) + | baseString (qualLeaf op) == "&&" = do + x' <- asBool <$> eval env x + if x' + then eval env y + else return $ ValuePrim $ BoolValue False + | baseString (qualLeaf op) == "||" = do + x' <- asBool <$> eval env x + if x' + then return $ ValuePrim $ BoolValue True + else eval env y + | otherwise = do + op' <- eval env $ Var op op_t loc + x' <- eval env x + y' <- eval env y + apply2 loc env op' x' y' + +eval env (If cond e1 e2 _ _) = do + cond' <- asBool <$> eval env cond + if cond' then eval env e1 else eval env e2 + +eval env (Apply f x _ _ loc) = do + f' <- eval env f + x' <- eval env x + apply loc env f' x' + +eval env (Negate e _) = do + ev <- eval env e + ValuePrim <$> case ev of + ValuePrim (SignedValue (Int8Value v)) -> return $ SignedValue $ Int8Value (-v) + ValuePrim (SignedValue (Int16Value v)) -> return $ SignedValue $ Int16Value (-v) + ValuePrim (SignedValue (Int32Value v)) -> return $ SignedValue $ Int32Value (-v) + ValuePrim (SignedValue (Int64Value v)) -> return $ SignedValue $ Int64Value (-v) + ValuePrim (UnsignedValue (Int8Value v)) -> return $ UnsignedValue $ Int8Value (-v) + ValuePrim (UnsignedValue (Int16Value v)) -> return $ UnsignedValue $ Int16Value (-v) + ValuePrim (UnsignedValue (Int32Value v)) -> return $ UnsignedValue $ Int32Value (-v) + ValuePrim (UnsignedValue (Int64Value v)) -> return $ UnsignedValue $ Int64Value (-v) + ValuePrim (FloatValue (Float32Value v)) -> return $ FloatValue $ Float32Value (-v) + ValuePrim (FloatValue (Float64Value v)) -> return $ FloatValue $ Float64Value (-v) + _ -> fail $ "Cannot negate " ++ pretty ev + +eval env (Index e is _ loc) = do + is' <- mapM (evalDimIndex env) is + arr <- eval env e + evalIndex loc env is' arr + +eval env (Update src is v loc) = + maybe oob return =<< + updateArray <$> mapM (evalDimIndex env) is <*> eval env src <*> eval env v + where oob = bad loc env "Bad update" + +eval env (RecordUpdate src all_fs v _ _) = + update <$> eval env src <*> pure all_fs <*> eval env v + where update _ [] v' = v' + update (ValueRecord src') (f:fs) v' + | Just f_v <- M.lookup f src' = + ValueRecord $ M.insert f (update f_v fs v') src' + update _ _ _ = error "eval RecordUpdate: invalid value." + +eval env (LetWith dest src is v body loc) = do + dest' <- maybe oob return =<< + updateArray <$> mapM (evalDimIndex env) is <*> + evalTermVar env (qualName $ identName src) <*> eval env v + let t = T.BoundV [] $ vacuousShapeAnnotations $ toStruct $ unInfo $ identType dest + eval (valEnv (M.singleton (identName dest) (Just t, dest')) <> env) body + where oob = bad loc env "Bad update" + +-- We treat zero-parameter lambdas as simply an expression to +-- evaluate immediately. Note that this is *not* the same as a lambda +-- that takes an empty tuple '()' as argument! Zero-parameter lambdas +-- can never occur in a well-formed Futhark program, but they are +-- convenient in the interpreter. +eval env (Lambda _ [] body _ (Info (_, t)) loc) = do + v <- eval env body + case (t, v) of + (Arrow _ _ _ rt, ValueFun f) -> + return $ ValueFun $ \arg -> do r <- f arg + rt' <- evalType env rt + match rt' r + _ -> match t v + where match vt v = + case matchValueToType env mempty vt v of + Right _ -> return v + Left _ -> bad loc env $ "Value `" <> pretty v <> + "` cannot match type `" <> pretty vt <> "`." + +eval env (Lambda tparams (p:ps) body mrd (Info (als, ret)) loc) = + return $ ValueFun $ \v -> do + p_env <- valEnv <$> matchPattern env p v + eval (p_env <> env) $ Lambda tparams ps body mrd (Info (als, ret)) loc + +eval env (OpSection qv _ _) = evalTermVar env qv + +eval env (OpSectionLeft qv _ e _ _ loc) = + join $ apply loc env <$> evalTermVar env qv <*> eval env e + +eval env (OpSectionRight qv _ e _ _ loc) = do + f <- evalTermVar env qv + y <- eval env e + return $ ValueFun $ \x -> join $ apply loc env <$> apply loc env f x <*> pure y + +eval env (IndexSection is _ loc) = do + is' <- mapM (evalDimIndex env) is + return $ ValueFun $ evalIndex loc env is' + +eval _ (ProjectSection ks _ _) = return $ ValueFun $ flip (foldM walk) ks + where walk (ValueRecord fs) f + | Just v' <- M.lookup f fs = return v' + walk _ _ = fail "Value does not have expected field." + +eval env (DoLoop _ pat init_e form body _) = do + init_v <- eval env init_e + case form of For iv bound -> do + bound' <- asSigned <$> eval env bound + forLoop (identName iv) bound' (zero bound') init_v + ForIn in_pat in_e -> do + in_vs <- fromArray <$> eval env in_e + foldM (forInLoop in_pat) init_v in_vs + While cond -> + whileLoop cond init_v + where withLoopParams v = (<>env) . valEnv <$> matchPattern env pat v + + inc = (`P.doAdd` Int64Value 1) + zero = (`P.doMul` Int64Value 0) + + forLoop iv bound i v + | i == bound = return v + | otherwise = do + env' <- withLoopParams v + forLoop iv bound (inc i) =<< + eval (valEnv (M.singleton iv (Just $ T.BoundV [] $ Prim $ Signed Int32, + ValuePrim (SignedValue i))) <> env') body + + whileLoop cond v = do + env' <- withLoopParams v + continue <- asBool <$> eval env' cond + if continue + then whileLoop cond =<< eval env' body + else return v + + forInLoop in_pat v in_v = do + env' <- withLoopParams v + pat_env <- matchPattern env' in_pat in_v + eval (valEnv pat_env <> env') body + +eval env (Project f e _ _) = do + v <- eval env e + case v of + ValueRecord fs | Just v' <- M.lookup f fs -> return v' + _ -> fail "Value does not have expected field." + +eval env (Unsafe e _) = eval env e + +eval env (Assert what e (Info s) loc) = do + cond <- asBool <$> eval env what + unless cond $ bad loc env s + eval env e + +eval _ e = error $ "eval not yet: " ++ show e + +substituteInModule :: M.Map VName VName -> Module -> Module +substituteInModule substs = onModule + where + rev_substs = reverseSubstitutions substs + replace v = fromMaybe v $ M.lookup v rev_substs + replaceQ v = maybe v qualName $ M.lookup (qualLeaf v) rev_substs + replaceM f m = M.fromList $ do + (k, v) <- M.toList m + return (replace k, f v) + onModule (Module (Env terms types)) = + Module $ Env (replaceM onTerm terms) (replaceM onType types) + onModule (ModuleFun f) = + ModuleFun $ \m -> onModule <$> f (substituteInModule rev_substs m) + onTerm (TermValue t v) = TermValue t v + onTerm (TermModule m) = TermModule $ onModule m + onType (T.TypeAbbr l ps t) = T.TypeAbbr l ps $ bimap onDim id t + onDim (NamedDim v) = NamedDim $ replaceQ v + onDim (ConstDim x) = ConstDim x + onDim AnyDim = AnyDim + +reverseSubstitutions :: M.Map VName VName -> M.Map VName VName +reverseSubstitutions = M.fromList . map (uncurry $ flip (,)) . M.toList + +evalModuleVar :: Env -> QualName VName -> EvalM Module +evalModuleVar env qv = + case lookupVar qv env of + Just (TermModule m) -> return m + _ -> error $ "`" <> pretty qv <> "` is not bound to a module." + +evalModExp :: Env -> ModExp -> EvalM Module + +evalModExp _ (ModImport _ (Info f) _) = do + f' <- lookupImport f + case f' of Nothing -> error $ "Unknown import " ++ show f + Just m -> return $ Module m + +evalModExp env (ModDecs ds _) = do + Env terms types <- foldM evalDec env ds + -- Remove everything that was present in the original Env. + return $ Module $ Env (terms `M.difference` envTerm env) + (types `M.difference` envType env) + +evalModExp env (ModVar qv _) = + evalModuleVar env qv + +evalModExp env (ModAscript me _ (Info substs) _) = + substituteInModule substs <$> evalModExp env me + +evalModExp env (ModParens me _) = evalModExp env me + +evalModExp env (ModLambda p ret e loc) = + return $ ModuleFun $ \am -> do + let env' = env { envTerm = M.insert (modParamName p) (TermModule am) $ envTerm env } + evalModExp env' $ case ret of + Nothing -> e + Just (se, rsubsts) -> ModAscript e se rsubsts loc + +evalModExp env (ModApply f e (Info psubst) (Info rsubst) _) = do + ModuleFun f' <- evalModExp env f + e' <- evalModExp env e + substituteInModule rsubst <$> f' (substituteInModule psubst e') + +evalDec :: Env -> Dec -> EvalM Env + +evalDec env (ValDec (ValBind _ v _ (Info t) tps ps def _ loc)) = do + t' <- evalType env t + let ftype = T.BoundV tps $ foldr (uncurry (Arrow ()) . patternParam) t' ps + val <- eval env $ Lambda tps ps def Nothing (Info (mempty, t')) loc + return $ valEnv (M.singleton v (Just ftype, val)) <> env + +evalDec env (OpenDec me (Info _) _) = do + Module me' <- evalModExp env me + return $ me' <> env + +evalDec env (LocalDec d _) = evalDec env d +evalDec env SigDec{} = return env +evalDec env (TypeDec (TypeBind v ps t _ _)) = do + t' <- evalType env $ unInfo $ expandedType t + let abbr = T.TypeAbbr Lifted ps t' + return env { envType = M.insert v abbr $ envType env } +evalDec env (ModDec (ModBind v ps ret body _ loc)) = do + mod <- evalModExp env $ wrapInLambda ps + return $ modEnv (M.singleton v mod) <> env + where wrapInLambda [] = case ret of + Just (se, substs) -> ModAscript body se substs loc + Nothing -> body + wrapInLambda [p] = ModLambda p ret body loc + wrapInLambda (p:ps') = ModLambda p Nothing (wrapInLambda ps') loc + +data Ctx = Ctx { ctxEnv :: Env + , ctxImports :: M.Map FilePath Env + } + +-- | The initial environment contains definitions of the various intrinsic functions. +initialCtx :: Ctx +initialCtx = + Ctx (Env (M.insert (VName (nameFromString "intrinsics") 0) + (TermModule (Module $ Env terms types)) terms) + types) + mempty + where + terms = M.mapMaybeWithKey (const . def . baseString) intrinsics + types = M.mapMaybeWithKey (const . tdef . baseString) intrinsics + + sintOp f = [ (getS, putS, P.doBinOp (f Int8)) + , (getS, putS, P.doBinOp (f Int16)) + , (getS, putS, P.doBinOp (f Int32)) + , (getS, putS, P.doBinOp (f Int64))] + uintOp f = [ (getU, putU, P.doBinOp (f Int8)) + , (getU, putU, P.doBinOp (f Int16)) + , (getU, putU, P.doBinOp (f Int32)) + , (getU, putU, P.doBinOp (f Int64))] + intOp f = sintOp f ++ uintOp f + floatOp f = [ (getF, putF, P.doBinOp (f Float32)) + , (getF, putF, P.doBinOp (f Float64))] + arithOp f g = Just $ bopDef $ intOp f ++ floatOp g + + flipCmps = map (\(f, g, h) -> (f, g, flip h)) + sintCmp f = [ (getS, Just . BoolValue, P.doCmpOp (f Int8)) + , (getS, Just . BoolValue, P.doCmpOp (f Int16)) + , (getS, Just . BoolValue, P.doCmpOp (f Int32)) + , (getS, Just . BoolValue, P.doCmpOp (f Int64))] + uintCmp f = [ (getU, Just . BoolValue, P.doCmpOp (f Int8)) + , (getU, Just . BoolValue, P.doCmpOp (f Int16)) + , (getU, Just . BoolValue, P.doCmpOp (f Int32)) + , (getU, Just . BoolValue, P.doCmpOp (f Int64))] + floatCmp f = [ (getF, Just . BoolValue, P.doCmpOp (f Float32)) + , (getF, Just . BoolValue, P.doCmpOp (f Float64))] + boolCmp f = [ (getB, Just . BoolValue, P.doCmpOp f) ] + + getV (SignedValue x) = Just $ P.IntValue x + getV (UnsignedValue x) = Just $ P.IntValue x + getV (FloatValue x) = Just $ P.FloatValue x + getV (BoolValue x) = Just $ P.BoolValue x + putV (P.IntValue x) = SignedValue x + putV (P.FloatValue x) = FloatValue x + putV (P.BoolValue x) = BoolValue x + putV P.Checked = BoolValue True + + getS (SignedValue x) = Just $ P.IntValue x + getS _ = Nothing + putS (P.IntValue x) = Just $ SignedValue x + putS _ = Nothing + + getU (UnsignedValue x) = Just $ P.IntValue x + getU _ = Nothing + putU (P.IntValue x) = Just $ UnsignedValue x + putU _ = Nothing + + getF (FloatValue x) = Just $ P.FloatValue x + getF _ = Nothing + putF (P.FloatValue x) = Just $ FloatValue x + putF _ = Nothing + + getB (BoolValue x) = Just $ P.BoolValue x + getB _ = Nothing + + fun1 f = + TermValue Nothing $ ValueFun $ \x -> f x + fun2 f = + TermValue Nothing $ ValueFun $ \x -> return $ ValueFun $ \y -> f x y + fun2t f = + TermValue Nothing $ ValueFun $ \v -> + case fromTuple v of Just [x,y] -> f x y + _ -> error $ "Expected pair; got: " ++ pretty v + fun3t f = + TermValue Nothing $ ValueFun $ \v -> + case fromTuple v of Just [x,y,z] -> f x y z + _ -> error $ "Expected triple; got: " ++ pretty v + + fun5t f = + TermValue Nothing $ ValueFun $ \v -> + case fromTuple v of Just [x,y,z,a,b] -> f x y z a b + _ -> error $ "Expected quintuple; got: " ++ pretty v + + bopDef fs = fun2 $ \x y -> + case (x, y) of + (ValuePrim x', ValuePrim y') + | Just z <- msum $ map (`bopDef'` (x', y')) fs -> + return $ ValuePrim z + _ -> + bad noLoc mempty $ "Cannot apply operator to arguments `" <> + pretty x <> "` and `" <> pretty y <> "`." + where bopDef' (valf, retf, op) (x, y) = do + x' <- valf x + y' <- valf y + retf =<< op x' y' + + unopDef fs = fun1 $ \x -> + case x of + (ValuePrim x') + | Just r <- msum $ map (`unopDef'` x') fs -> + return $ ValuePrim r + _ -> + bad noLoc mempty $ "Cannot apply function to argument `" <> + pretty x <> "`." + where unopDef' (valf, retf, op) x = do + x' <- valf x + retf =<< op x' + + def "~" = Just $ unopDef [ (getS, putS, P.doUnOp $ P.Complement Int8) + , (getS, putS, P.doUnOp $ P.Complement Int16) + , (getS, putS, P.doUnOp $ P.Complement Int32) + , (getS, putS, P.doUnOp $ P.Complement Int64) + , (getU, putU, P.doUnOp $ P.Complement Int8) + , (getU, putU, P.doUnOp $ P.Complement Int16) + , (getU, putU, P.doUnOp $ P.Complement Int32) + , (getU, putU, P.doUnOp $ P.Complement Int64)] + def "!" = Just $ fun1 $ return . ValuePrim . BoolValue . not . asBool + + def "+" = arithOp P.Add P.FAdd + def "-" = arithOp P.Sub P.FSub + def "*" = arithOp P.Mul P.FMul + def "**" = arithOp P.Pow P.FPow + def "/" = Just $ bopDef $ sintOp P.SDiv ++ uintOp P.UDiv ++ floatOp P.FDiv + def "%" = Just $ bopDef $ sintOp P.SMod ++ uintOp P.UMod + def "//" = Just $ bopDef $ sintOp P.SQuot ++ uintOp P.UDiv + def "%%" = Just $ bopDef $ sintOp P.SRem ++ uintOp P.UMod + def "^" = Just $ bopDef $ intOp P.Xor + def "&" = Just $ bopDef $ intOp P.And + def "|" = Just $ bopDef $ intOp P.Or + def ">>" = Just $ bopDef $ sintOp P.AShr ++ uintOp P.LShr + def "<<" = Just $ bopDef $ intOp P.Shl + def ">>>" = Just $ bopDef $ sintOp P.LShr ++ uintOp P.LShr + def "==" = Just $ fun2 $ \xs ys -> return $ ValuePrim $ BoolValue $ xs == ys + def "!=" = Just $ fun2 $ \xs ys -> return $ ValuePrim $ BoolValue $ xs /= ys + + -- The short-circuiting is handled directly in 'eval'; these cases + -- are only used when partially applying and such. + def "&&" = Just $ fun2 $ \x y -> + return $ ValuePrim $ BoolValue $ asBool x && asBool y + def "||" = Just $ fun2 $ \x y -> + return $ ValuePrim $ BoolValue $ asBool x || asBool y + + def "<" = Just $ bopDef $ + sintCmp P.CmpSlt ++ uintCmp P.CmpUlt ++ + floatCmp P.FCmpLt ++ boolCmp P.CmpLlt + def ">" = Just $ bopDef $ flipCmps $ + sintCmp P.CmpSlt ++ uintCmp P.CmpUlt ++ + floatCmp P.FCmpLt ++ boolCmp P.CmpLlt + def "<=" = Just $ bopDef $ + sintCmp P.CmpSle ++ uintCmp P.CmpUle ++ + floatCmp P.FCmpLe ++ boolCmp P.CmpLle + def ">=" = Just $ bopDef $ flipCmps $ + sintCmp P.CmpSle ++ uintCmp P.CmpUle ++ + floatCmp P.FCmpLe ++ boolCmp P.CmpLle + + def s + | Just bop <- find ((s==) . pretty) P.allBinOps = + Just $ bopDef [(getV, Just . putV, P.doBinOp bop)] + | Just cop <- find ((s==) . pretty) P.allConvOps = + Just $ unopDef [(getV, Just . putV, P.doConvOp cop)] + | Just unop <- find ((s==) . pretty) P.allUnOps = + Just $ unopDef [(getV, Just . putV, P.doUnOp unop)] + | Just unop <- find ((s==) . pretty) P.allCmpOps = + Just $ bopDef [(getV, bool, P.doCmpOp unop)] + + | Just (pts, _, f) <- M.lookup s P.primFuns = + case length pts of + 1 -> Just $ unopDef [(getV, Just . putV, f . pure)] + _ -> Just $ bopDef [(getV, Just . putV, \x y -> f [x,y])] + + | "sign_" `isPrefixOf` s = + Just $ fun1 $ \x -> + case x of (ValuePrim (UnsignedValue x')) -> + return $ ValuePrim $ SignedValue x' + _ -> error $ "Cannot sign: " ++ pretty x + | "unsign_" `isPrefixOf` s = + Just $ fun1 $ \x -> + case x of (ValuePrim (SignedValue x')) -> + return $ ValuePrim $ UnsignedValue x' + _ -> error $ "Cannot unsign: " ++ pretty x + where bool = Just . BoolValue + + def "map" = Just $ fun2t $ \f xs -> + toArray =<< mapM (apply noLoc mempty f) (fromArray xs) + + def s | "reduce" `isPrefixOf` s = Just $ fun3t $ \f ne xs -> + foldM (apply2 noLoc mempty f) ne $ fromArray xs + + def "scan" = Just $ fun3t $ \f ne xs -> do + let next (out, acc) x = do + x' <- apply2 noLoc mempty f acc x + return (x':out, x') + toArray . reverse . fst =<< foldM next ([], ne) (fromArray xs) + + def s | "stream_map" `isPrefixOf` s = + Just $ fun2t $ apply noLoc mempty + + def s | "stream_red" `isPrefixOf` s = + Just $ fun3t $ \_ f xs -> apply noLoc mempty f xs + + def "scatter" = Just $ fun3t $ \arr is vs -> + case arr of + ValueArray arr' -> + return $ ValueArray $ foldl' update arr' + (zip (map asInt $ fromArray is) (fromArray vs)) + _ -> + error $ "scatter expects array, but got: " ++ pretty arr + where update arr' (i, v) = + if i >= 0 && i < arrayLength arr' + then arr' // [(i, v)] else arr' + + def "gen_reduce" = Just $ fun5t $ \arr fun _ is vs -> + case arr of + ValueArray arr' -> + ValueArray <$> foldM (update fun) arr' + (zip (map asInt $ fromArray is) (fromArray vs)) + _ -> + error $ "gen_reduce expects array, but got: " ++ pretty arr + where update fun arr' (i, v) = + if i >= 0 && i < arrayLength arr' + then do + v' <- apply2 noLoc mempty fun (arr' ! i) v + return $ arr' // [(i, v')] + else return arr' + + def "partition" = Just $ fun3t $ \k f xs -> + let next outs x = do + i <- asInt <$> apply noLoc mempty f x + return $ insertAt i x outs + pack parts = + toTuple [toArray' $ concat parts, + toArray' $ + map (ValuePrim . SignedValue . Int32Value . genericLength) parts] + in pack . map reverse <$> + foldM next (replicate (asInt k) []) (fromArray xs) + where insertAt 0 x (l:ls) = (x:l):ls + insertAt i x (l:ls) = l:insertAt (i-1) x ls + insertAt _ _ ls = ls + + def "cmp_threshold" = Just $ fun2t $ \_ _ -> + return $ ValuePrim $ BoolValue True + + def "unzip" = Just $ fun1 $ \x -> + toTuple <$> listPair (unzip $ map (fromPair . fromTuple) $ fromArray x) + where fromPair (Just [x,y]) = (x,y) + fromPair l = error $ "Not a pair: " ++ pretty l + listPair (xs, ys) = do + xs' <- toArray xs + ys' <- toArray ys + return [xs', ys'] + + def "zip" = Just $ fun2t $ \xs ys -> + toArray $ map toTuple $ transpose [fromArray xs, fromArray ys] + + def "concat" = Just $ fun2t $ \xs ys -> + toArray $ fromArray xs ++ fromArray ys + + def "transpose" = Just $ fun1 $ + (toArray <=< mapM toArray) . transpose . map fromArray . fromArray + + def "rotate" = Just $ fun2t $ \i xs -> + if asInt i > 0 + then let (bef, aft) = splitAt (asInt i) $ fromArray xs + in toArray $ aft ++ bef + else let (bef, aft) = splitFromEnd (-asInt i) $ fromArray xs + in toArray $ aft ++ bef + + def "flatten" = Just $ fun1 $ + toArray . concatMap fromArray . fromArray + + def "unflatten" = Just $ fun3t $ \_ m xs -> + toArray =<< mapM toArray (chunk (asInt m) $ fromArray xs) + + def "opaque" = Just $ fun1 return + + def "trace" = Just $ fun1 $ \v -> trace v >> return v + + def "break" = Just $ fun1 $ \v -> do + break + return v + + def s | nameFromString s `M.member` namesToPrimTypes = Nothing + + def s = error $ "Missing intrinsic: " ++ s + + tdef s = do + t <- nameFromString s `M.lookup` namesToPrimTypes + return $ T.TypeAbbr Unlifted [] $ Prim t + +interpretExp :: Ctx -> Exp -> F ExtOp Value +interpretExp ctx e = runEvalM (ctxImports ctx) $ eval (ctxEnv ctx) e + +interpretDec :: Ctx -> Dec -> F ExtOp Ctx +interpretDec ctx d = do + env <- runEvalM (ctxImports ctx) $ evalDec (ctxEnv ctx) d + return ctx { ctxEnv = env } + +interpretImport :: Ctx -> (FilePath, Prog) -> F ExtOp Ctx +interpretImport ctx (fp, prog) = do + env <- runEvalM (ctxImports ctx) $ foldM evalDec (ctxEnv ctx) $ progDecs prog + return ctx { ctxImports = M.insert fp env $ ctxImports ctx } + +-- | Execute the named function on the given arguments; will fail +-- horribly if these are ill-typed. +interpretFunction :: Ctx -> VName -> [Value] -> F ExtOp Value +interpretFunction ctx fname vs = runEvalM (ctxImports ctx) $ do + f <- evalTermVar (ctxEnv ctx) $ qualName fname + foldM (apply noLoc mempty) f vs diff --git a/src/Language/Futhark/Parser.hs b/src/Language/Futhark/Parser.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Parser.hs @@ -0,0 +1,56 @@ +-- | Interface to the Futhark parser. +module Language.Futhark.Parser + ( parseFuthark + , parseExp + , parseType + + , parseValue + , parseValues + + , parseDecOrExpIncrM + + , ParseError (..) + + , scanTokensText + , L(..) + , Token(..) + ) + where + +import qualified Data.Text as T + +import Language.Futhark.Syntax +import Language.Futhark.Attributes +import Language.Futhark.Parser.Parser +import Language.Futhark.Parser.Lexer + +-- | Parse an entire Futhark program from the given 'T.Text', using +-- the 'FilePath' as the source name for error messages. +parseFuthark :: FilePath -> T.Text + -> Either ParseError UncheckedProg +parseFuthark = parse prog + +-- | Parse an Futhark expression from the given 'String', using the +-- 'FilePath' as the source name for error messages. +parseExp :: FilePath -> T.Text + -> Either ParseError UncheckedExp +parseExp = parse expression + +-- | Parse an Futhark type from the given 'String', using the +-- 'FilePath' as the source name for error messages. +parseType :: FilePath -> T.Text + -> Either ParseError UncheckedTypeExp +parseType = parse futharkType + +-- | Parse any Futhark value from the given 'String', using the 'FilePath' +-- as the source name for error messages. +parseValue :: FilePath -> T.Text + -> Either ParseError Value +parseValue = parse anyValue + +-- | Parse several Futhark values (separated by anything) from the given +-- 'String', using the 'FilePath' as the source name for error +-- messages. +parseValues :: FilePath -> T.Text + -> Either ParseError [Value] +parseValues = parse anyValues diff --git a/src/Language/Futhark/Parser/Lexer.x b/src/Language/Futhark/Parser/Lexer.x new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Parser/Lexer.x @@ -0,0 +1,397 @@ +{ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE BangPatterns #-} +{-# OPTIONS_GHC -w #-} +-- | The Futhark lexer. Takes a string, produces a list of tokens with position information. +module Language.Futhark.Parser.Lexer + ( Token(..) + , L(..) + , scanTokens + , scanTokensText + ) where + +import qualified Data.ByteString.Lazy as BS +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import Data.Char (ord, toLower) +import Data.Loc hiding (L) +import Data.Int (Int8, Int16, Int32, Int64) +import Data.Word (Word8) +import Data.Bits +import Data.Function (fix) +import Data.List +import Data.Monoid + +import Language.Futhark.Core (Int8, Int16, Int32, Int64, + Word8, Word16, Word32, Word64, + Name, nameFromText, nameToText) +import Language.Futhark.Attributes (leadingOperator) +import Language.Futhark.Syntax (BinOp(..)) + +} + +%wrapper "monad-bytestring" + +@charlit = ($printable#['\\]|\\($printable|[0-9]+)) +@stringcharlit = ($printable#[\"\\]|\\($printable|[0-9]+)|\n) +@hexlit = 0[xX][0-9a-fA-F][0-9a-fA-F_]* +@declit = [0-9][0-9_]* +@binlit = 0[bB][01][01_]* +@romlit = 0[rR][IVXLCDM][IVXLCDM_]* +@intlit = @hexlit|@binlit|@declit|@romlit +@reallit = (([0-9][0-9_]*("."[0-9][0-9_]*)?))([eE][\+\-]?[0-9]+)? +@hexreallit = 0[xX][0-9a-fA-F][0-9a-fA-F_]*"."[0-9a-fA-F][0-9a-fA-F_]*([pP][\+\-]?[0-9]+) + +@field = [a-zA-Z0-9] [a-zA-Z0-9_]* + +@identifier = [a-zA-Z] [a-zA-Z0-9_']* | "_" [a-zA-Z0-9] [a-zA-Z0-9_']* +@qualidentifier = (@identifier ".")+ @identifier + +@unop = ("!"|"~") +@qualunop = (@identifier ".")+ @unop + +@opchar = ("+"|"-"|"*"|"/"|"%"|"="|"!"|">"|"<"|"|"|"&"|"^"|".") +@binop = ("+"|"-"|"*"|"/"|"%"|"="|"!"|">"|"<"|"|"|"&"|"^") @opchar* +@qualbinop = (@identifier ".")+ @binop + +@space = [\ \t\f\v] +@doc = "-- |"[^\n]*(\n@space*"--"[^\n]*)* + +tokens :- + + $white+ ; + @doc { tokenM $ return . DOC . T.unpack . T.unlines . + map (T.drop 3 . T.stripStart) . + T.split (== '\n') . ("--"<>) . + T.drop 4 } + "--"[^\n]* ; + "=" { tokenC EQU } + "(" { tokenC LPAR } + ")" { tokenC RPAR } + ")[" { tokenC RPAR_THEN_LBRACKET } + "[" { tokenC LBRACKET } + "]" { tokenC RBRACKET } + "{" { tokenC LCURLY } + "}" { tokenC RCURLY } + "," { tokenC COMMA } + "_" { tokenC UNDERSCORE } + "->" { tokenC RIGHT_ARROW } + "<-" { tokenC LEFT_ARROW } + ":" { tokenC COLON } + "." { tokenC DOT } + "\" { tokenC BACKSLASH } + "'" { tokenC APOSTROPHE } + "'^" { tokenC APOSTROPHE_THEN_HAT } + "`" { tokenC BACKTICK } + "#" { tokenC HASH } + "..<" { tokenC TWO_DOTS_LT } + "..>" { tokenC TWO_DOTS_GT } + "..." { tokenC THREE_DOTS } + ".." { tokenC TWO_DOTS } + + @intlit i8 { tokenM $ return . I8LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') } + @intlit i16 { tokenM $ return . I16LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') } + @intlit i32 { tokenM $ return . I32LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') } + @intlit i64 { tokenM $ return . I64LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='i') } + @intlit u8 { tokenM $ return . U8LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') } + @intlit u16 { tokenM $ return . U16LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') } + @intlit u32 { tokenM $ return . U32LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') } + @intlit u64 { tokenM $ return . U64LIT . readIntegral . T.filter (/= '_') . T.takeWhile (/='u') } + @intlit { tokenM $ return . INTLIT . readIntegral . T.filter (/= '_') } + + @reallit f32 { tokenM $ fmap F32LIT . tryRead "f32" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') } + @reallit f64 { tokenM $ fmap F64LIT . tryRead "f64" . suffZero . T.filter (/= '_') . T.takeWhile (/='f') } + @reallit { tokenM $ fmap FLOATLIT . tryRead "f64" . suffZero . T.filter (/= '_') } + @hexreallit f32 { tokenM $ fmap F32LIT . readHexRealLit "f32" . suffZero . T.filter (/= '_') . fst . T.breakOn (T.pack "f32") } + @hexreallit f64 { tokenM $ fmap F64LIT . readHexRealLit "f64" . suffZero . T.filter (/= '_') . fst . T.breakOn (T.pack "f64") } + @hexreallit { tokenM $ fmap FLOATLIT . readHexRealLit "f64" . suffZero . T.filter (/= '_') . fst . T.breakOn (T.pack "f64") } + "'" @charlit "'" { tokenM $ fmap CHARLIT . tryRead "char" } + \" @stringcharlit* \" { tokenM $ fmap STRINGLIT . tryRead "string" } + + @identifier { tokenS keyword } + @identifier "[" { tokenM $ fmap INDEXING . indexing . T.takeWhile (/='[') } + @qualidentifier "[" { tokenM $ fmap (uncurry QUALINDEXING) . mkQualId . T.takeWhile (/='[') } + @identifier "." "(" { tokenM $ fmap (QUALPAREN []) . indexing . T.init . T.takeWhile (/='(') } + @qualidentifier "." "(" { tokenM $ fmap (uncurry QUALPAREN) . mkQualId . T.init . T.takeWhile (/='(') } + + @unop { tokenS $ UNOP . nameFromText } + @qualunop { tokenM $ fmap (uncurry QUALUNOP) . mkQualId } + + @binop { tokenM $ return . symbol [] . nameFromText } + @qualbinop { tokenM $ \s -> do (qs,k) <- mkQualId s; return (symbol qs k) } +{ + +keyword :: T.Text -> Token +keyword s = + case s of + "true" -> TRUE + "false" -> FALSE + "if" -> IF + "then" -> THEN + "else" -> ELSE + "let" -> LET + "loop" -> LOOP + "in" -> IN + "val" -> VAL + "for" -> FOR + "do" -> DO + "with" -> WITH + "local" -> LOCAL + "open" -> OPEN + "include" -> INCLUDE + "import" -> IMPORT + "type" -> TYPE + "entry" -> ENTRY + "module" -> MODULE + "while" -> WHILE + "unsafe" -> UNSAFE + "assert" -> ASSERT + + _ -> ID $ nameFromText s + +indexing :: T.Text -> Alex Name +indexing s = case keyword s of + ID v -> return v + _ -> fail $ "Cannot index keyword '" ++ T.unpack s ++ "'." + +mkQualId :: T.Text -> Alex ([Name], Name) +mkQualId s = case reverse $ T.splitOn "." s of + [] -> fail "mkQualId: no components" + k:qs -> return (map nameFromText (reverse qs), nameFromText k) + +-- | Suffix a zero if the last character is dot. +suffZero :: T.Text -> T.Text +suffZero s = if T.last s == '.' then s <> "0" else s + +tryRead :: Read a => String -> T.Text -> Alex a +tryRead desc s = case reads s' of + [(x, "")] -> return x + _ -> fail $ "Invalid " ++ desc ++ " literal: `" ++ T.unpack s ++ "'." + where s' = T.unpack s + +readIntegral :: Integral a => T.Text -> a +readIntegral s + | "0x" `T.isPrefixOf` s || "0X" `T.isPrefixOf` s = + T.foldl (another hex_digits) 0 (T.drop 2 s) + | "0b" `T.isPrefixOf` s || "0b" `T.isPrefixOf` s = + T.foldl (another binary_digits) 0 (T.drop 2 s) + | "0r" `T.isPrefixOf` s = + fromRoman (T.drop 2 s) + | otherwise = + T.foldl (another decimal_digits) 0 s + where another digits acc c = acc * base + maybe 0 fromIntegral (elemIndex (toLower c) digits) + where base = fromIntegral $ length digits + + binary_digits = ['0', '1'] + decimal_digits = ['0'..'9'] + hex_digits = decimal_digits ++ ['a'..'f'] + +tokenC v = tokenS $ const v + +tokenS f = tokenM $ return . f + +type Lexeme a = ((Int, Int, Int), (Int, Int, Int), a) + +tokenM :: (T.Text -> Alex a) + -> (AlexPosn, Char, ByteString.ByteString, Int64) + -> Int64 + -> Alex (Lexeme a) +tokenM f (AlexPn addr line col, _, s, _) len = do + x <- f $ T.decodeUtf8 $ BS.toStrict s' + return (pos, advance pos s', x) + where pos = (line, col, addr) + s' = BS.take len s + +advance :: (Int, Int, Int) -> ByteString.ByteString -> (Int, Int, Int) +advance orig_pos = foldl' advance' orig_pos . init . ByteString.unpack + where advance' (!line, !col, !addr) c + | c == nl = (line + 1, 1, addr + 1) + | otherwise = (line, col + 1, addr + 1) + nl = fromIntegral $ ord '\n' + +symbol :: [Name] -> Name -> Token +symbol [] q + | nameToText q == "*" = ASTERISK + | nameToText q == "-" = NEGATE + | nameToText q == "<" = LTH + | nameToText q == "^" = HAT + | otherwise = SYMBOL (leadingOperator q) [] q +symbol qs q = SYMBOL (leadingOperator q) qs q + + +romanNumerals :: Integral a => [(T.Text,a)] +romanNumerals = reverse + [ ("I", 1) + , ("IV", 4) + , ("V", 5) + , ("IX", 9) + , ("X", 10) + , ("XL", 40) + , ("L", 50) + , ("XC", 90) + , ("C", 100) + , ("CD", 400) + , ("D", 500) + , ("CM", 900) + , ("M", 1000) + ] + +fromRoman :: Integral a => T.Text -> a +fromRoman s = + case find ((`T.isPrefixOf` s) . fst) romanNumerals of + Nothing -> 0 + Just (d,n) -> n+fromRoman (T.drop (T.length d) s) + +fromHexRealLit :: RealFloat a => T.Text -> Maybe a +fromHexRealLit s = + let num = (T.drop 2 s) in + -- extract number into integer, fractional and (optional) exponent + let comps = (T.split (\x -> x == '.' || x == 'p' || x == 'P') num) in + case comps of + [i, f, p] -> + let int_part = readIntegral (T.pack ("0x" ++ (T.unpack i))) + frac_part = readIntegral (T.pack ("0x" ++ (T.unpack f))) + exponent = if ((T.pack "-") `T.isPrefixOf` p) + then -1 * (readIntegral p) + else readIntegral p + + frac_len = T.length f + frac_val = (fromIntegral frac_part) / (16.0 ** (fromIntegral frac_len)) + total_val = ((fromIntegral int_part) + frac_val) * (2.0 ** (fromIntegral exponent)) in + Just (total_val) + _ -> Nothing + +readHexRealLit :: RealFloat a => String -> T.Text -> Alex a +readHexRealLit desc s = + case fromHexRealLit s of + Just (n) -> return n + Nothing -> fail $ "Invalid " ++ desc ++ " literal: " ++ T.unpack s + +alexGetPosn :: Alex (Int, Int, Int) +alexGetPosn = Alex $ \s -> + let (AlexPn off line col) = alex_pos s + in Right (s, (line, col, off)) + +alexEOF = do + posn <- alexGetPosn + return (posn, posn, EOF) + +-- | A value tagged with a source location. +data L a = L SrcLoc a deriving (Show) + +instance Eq a => Eq (L a) where + L _ x == L _ y = x == y + +instance Located (L a) where + locOf (L (SrcLoc loc) _) = loc + +-- | A lexical token. It does not itself contain position +-- information, so in practice the parser will consume tokens tagged +-- with a source position. +data Token = ID Name + | INDEXING Name + | QUALINDEXING [Name] Name + | QUALPAREN [Name] Name + | UNOP Name + | QUALUNOP [Name] Name + | SYMBOL BinOp [Name] Name + + | INTLIT Integer + | STRINGLIT String + | I8LIT Int8 + | I16LIT Int16 + | I32LIT Int32 + | I64LIT Int64 + | U8LIT Word8 + | U16LIT Word16 + | U32LIT Word32 + | U64LIT Word64 + | FLOATLIT Double + | F32LIT Float + | F64LIT Double + | CHARLIT Char + + | COLON + | BACKSLASH + | APOSTROPHE + | APOSTROPHE_THEN_HAT + | BACKTICK + | HASH + | DOT + | TWO_DOTS + | TWO_DOTS_LT + | TWO_DOTS_GT + | THREE_DOTS + | LPAR + | RPAR + | RPAR_THEN_LBRACKET + | LBRACKET + | RBRACKET + | LCURLY + | RCURLY + | COMMA + | UNDERSCORE + | RIGHT_ARROW + | LEFT_ARROW + + | EQU + | ASTERISK + | NEGATE + | LTH + | HAT + + | IF + | THEN + | ELSE + | LET + | LOOP + | IN + | FOR + | DO + | WITH + | UNSAFE + | ASSERT + | TRUE + | FALSE + | WHILE + | INCLUDE + | IMPORT + | ENTRY + | TYPE + | MODULE + | VAL + | OPEN + | LOCAL + + | DOC String + + | EOF + + deriving (Show, Eq, Ord) + +runAlex' :: AlexPosn -> ByteString.ByteString -> Alex a -> Either String a +runAlex' start_pos input__ (Alex f) = + case f (AlexState { alex_pos = start_pos + , alex_bpos = 0 + , alex_inp = input__ + , alex_chr = '\n' + , alex_scd = 0}) of Left msg -> Left msg + Right ( _, a ) -> Right a + +scanTokensText :: Pos -> T.Text -> Either String ([L Token], Pos) +scanTokensText pos = scanTokens pos . BS.fromStrict . T.encodeUtf8 + +scanTokens :: Pos -> BS.ByteString -> Either String ([L Token], Pos) +scanTokens (Pos file start_line start_col start_off) str = + runAlex' (AlexPn start_off start_line start_col) str $ do + fix $ \loop -> do + tok <- alexMonadScan + case tok of + (start, end, EOF) -> + return ([], posnToPos end) + (start, end, t) -> do + (rest, endpos) <- loop + return (L (pos start end) t : rest, endpos) + where pos start end = SrcLoc $ Loc (posnToPos start) (posnToPos end) + posnToPos (line, col, off) = Pos file line col off +} diff --git a/src/Language/Futhark/Parser/Parser.y b/src/Language/Futhark/Parser/Parser.y new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Parser/Parser.y @@ -0,0 +1,1094 @@ +{ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE OverloadedStrings #-} +-- | Futhark parser written with Happy. +module Language.Futhark.Parser.Parser + ( prog + , expression + , futharkType + , anyValue + , anyValues + + , ParserMonad + , parse + , ParseError(..) + , parseDecOrExpIncrM + ) + where + +import Control.Monad +import Control.Monad.Trans +import Control.Monad.Except +import Control.Monad.Reader +import Control.Monad.Trans.State +import Control.Arrow +import Data.Array +import qualified Data.Text as T +import Data.Char (ord) +import Data.Maybe (fromMaybe, fromJust) +import Data.Loc hiding (L) -- Lexer has replacements. +import qualified Data.Map.Strict as M +import Data.Monoid + +import Language.Futhark.Syntax hiding (ID) +import Language.Futhark.Attributes +import Language.Futhark.Pretty +import Language.Futhark.Parser.Lexer + +} + +%name prog Prog +%name futharkType TypeExp +%name expression Exp +%name declaration Dec +%name anyValue Value +%name anyValues CatValues + +%tokentype { L Token } +%error { parseError } +%monad { ParserMonad } +%lexer { lexer } { L _ EOF } + +%token + if { L $$ IF } + then { L $$ THEN } + else { L $$ ELSE } + let { L $$ LET } + loop { L $$ LOOP } + in { L $$ IN } + + id { L _ (ID _) } + 'id[' { L _ (INDEXING _) } + + 'qid[' { L _ (QUALINDEXING _ _) } + + 'qid.(' { L _ (QUALPAREN _ _) } + + unop { L _ (UNOP _) } + qunop { L _ (QUALUNOP _ _) } + + intlit { L _ (INTLIT _) } + i8lit { L _ (I8LIT _) } + i16lit { L _ (I16LIT _) } + i32lit { L _ (I32LIT _) } + i64lit { L _ (I64LIT _) } + u8lit { L _ (U8LIT _) } + u16lit { L _ (U16LIT _) } + u32lit { L _ (U32LIT _) } + u64lit { L _ (U64LIT _) } + floatlit { L _ (FLOATLIT _) } + f32lit { L _ (F32LIT _) } + f64lit { L _ (F64LIT _) } + stringlit { L _ (STRINGLIT _) } + charlit { L _ (CHARLIT _) } + + '#' { L $$ HASH } + '..' { L $$ TWO_DOTS } + '...' { L $$ THREE_DOTS } + '..<' { L $$ TWO_DOTS_LT } + '..>' { L $$ TWO_DOTS_GT } + '=' { L $$ EQU } + + '*' { L $$ ASTERISK } + '-' { L $$ NEGATE } + '<' { L $$ LTH } + '^' { L $$ HAT } + + '+...' { L _ (SYMBOL Plus _ _) } + '-...' { L _ (SYMBOL Minus _ _) } + '*...' { L _ (SYMBOL Times _ _) } + '/...' { L _ (SYMBOL Divide _ _) } + '%...' { L _ (SYMBOL Mod _ _) } + '//...' { L _ (SYMBOL Quot _ _) } + '%%...' { L _ (SYMBOL Rem _ _) } + '==...' { L _ (SYMBOL Equal _ _) } + '!=...' { L _ (SYMBOL NotEqual _ _) } + '<...' { L _ (SYMBOL Less _ _) } + '>...' { L _ (SYMBOL Greater _ _) } + '<=...' { L _ (SYMBOL Leq _ _) } + '>=...' { L _ (SYMBOL Geq _ _) } + '**...' { L _ (SYMBOL Pow _ _) } + '<<...' { L _ (SYMBOL ShiftL _ _) } + '>>...' { L _ (SYMBOL ShiftR _ _) } + '|>...' { L _ (SYMBOL PipeRight _ _) } + '<|...' { L _ (SYMBOL PipeLeft _ _) } + '|...' { L _ (SYMBOL Bor _ _) } + '&...' { L _ (SYMBOL Band _ _) } + '^...' { L _ (SYMBOL Xor _ _) } + '||...' { L _ (SYMBOL LogOr _ _) } + '&&...' { L _ (SYMBOL LogAnd _ _) } + + '(' { L $$ LPAR } + ')' { L $$ RPAR } + ')[' { L $$ RPAR_THEN_LBRACKET } + '{' { L $$ LCURLY } + '}' { L $$ RCURLY } + '[' { L $$ LBRACKET } + ']' { L $$ RBRACKET } + ',' { L $$ COMMA } + '_' { L $$ UNDERSCORE } + '\\' { L $$ BACKSLASH } + '\'' { L $$ APOSTROPHE } + '\'^' { L $$ APOSTROPHE_THEN_HAT } + '`' { L $$ BACKTICK } + entry { L $$ ENTRY } + '->' { L $$ RIGHT_ARROW } + '<-' { L $$ LEFT_ARROW } + ':' { L $$ COLON } + '.' { L $$ DOT } + for { L $$ FOR } + do { L $$ DO } + with { L $$ WITH } + unsafe { L $$ UNSAFE } + assert { L $$ ASSERT } + true { L $$ TRUE } + false { L $$ FALSE } + while { L $$ WHILE } + include { L $$ INCLUDE } + import { L $$ IMPORT } + type { L $$ TYPE } + module { L $$ MODULE } + val { L $$ VAL } + open { L $$ OPEN } + local { L $$ LOCAL } + doc { L _ (DOC _) } + +%left bottom +%left ifprec letprec unsafe +%left ',' +%left ':' +%right '...' '..<' '..>' '..' +%left '`' +%right '->' +%left with +%left '=' '<-' +%left '|>...' +%right '<|...' +%left '||...' +%left '&&...' +%left '<=...' '>=...' '>...' '<' '<...' '==...' '!=...' +%left '&...' '^...' '^' '|...' +%left '<<...' '>>...' +%left '+...' '-...' '-' +%left '*...' '*' '/...' '%...' '//...' '%%...' +%left '**...' +%left juxtprec +%left indexprec +%% + +-- The main parser. + +Doc :: { DocComment } + : doc { let L loc (DOC s) = $1 in DocComment s loc } + +-- Three cases to avoid ambiguities. +Prog :: { UncheckedProg } + -- File begins with a file comment, followed by a Dec with a comment. + : Doc Doc Dec_ Decs { Prog (Just $1) (addDoc $2 $3 : $4) } + -- File begins with a file comment, followed by a Dec with no comment. + | Doc Dec_ Decs { Prog (Just $1) ($2 : $3) } + -- File begins with a dec with no comment. + | Dec_ Decs { Prog Nothing ($1 : $2) } +; + +Dec :: { UncheckedDec } + : Dec_ { $1 } + | Doc Dec_ { addDoc $1 $2 } + +Decs :: { [UncheckedDec] } + : { [] } + | Dec Decs { $1 : $2 } + +Dec_ :: { UncheckedDec } + : Val { ValDec $1 } + | TypeAbbr { TypeDec $1 } + | SigBind { SigDec $1 } + | ModBind { ModDec $1 } + | open ModExp + { OpenDec $2 NoInfo $1 } + | import stringlit + { let L loc (STRINGLIT s) = $2 in LocalDec (OpenDec (ModImport s NoInfo loc) NoInfo $1) (srcspan $1 $>) } + | local Dec { LocalDec $2 (srcspan $1 $>) } +; + +SigExp :: { UncheckedSigExp } + : QualName { let (v, loc) = $1 in SigVar v loc } + | '{' Specs '}' { SigSpecs $2 (srcspan $1 $>) } + | SigExp with TypeRef { SigWith $1 $3 (srcspan $1 $>) } + | '(' SigExp ')' { SigParens $2 (srcspan $1 $>) } + | '(' id ':' SigExp ')' '->' SigExp + { let L _ (ID name) = $2 + in SigArrow (Just name) $4 $7 (srcspan $1 $>) } + | SigExp '->' SigExp { SigArrow Nothing $1 $3 (srcspan $1 $>) } + +TypeRef :: { TypeRefBase NoInfo Name } + : QualName TypeParams '=' TypeExpTerm + { TypeRef (fst $1) $2 (TypeDecl $4 NoInfo) (srcspan (snd $1) $>) } + +SigBind :: { SigBindBase NoInfo Name } + : module type id '=' SigExp + { let L _ (ID name) = $3 + in SigBind name $5 Nothing (srcspan $1 $>) } + +ModExp :: { UncheckedModExp } + : ModExp ':' SigExp + { ModAscript $1 $3 NoInfo (srcspan $1 $>) } + | '\\' ModParam maybeAscription(SimpleSigExp) '->' ModExp + { ModLambda $2 (fmap (,NoInfo) $3) $5 (srcspan $1 $>) } + | import stringlit + { let L _ (STRINGLIT s) = $2 in ModImport s NoInfo (srcspan $1 $>) } + | ModExpApply + { $1 } + | ModExpAtom + { $1 } + + +ModExpApply :: { UncheckedModExp } + : ModExpAtom ModExpAtom %prec juxtprec + { ModApply $1 $2 NoInfo NoInfo (srcspan $1 $>) } + | ModExpApply ModExpAtom %prec juxtprec + { ModApply $1 $2 NoInfo NoInfo (srcspan $1 $>) } + +ModExpAtom :: { UncheckedModExp } + : '(' ModExp ')' + { ModParens $2 (srcspan $1 $>) } + | QualName + { let (v, loc) = $1 in ModVar v loc } + | '{' Decs '}' { ModDecs $2 (srcspan $1 $>) } + +SimpleSigExp :: { UncheckedSigExp } + : QualName { let (v, loc) = $1 in SigVar v loc } + | '(' SigExp ')' { $2 } + +ModBind :: { ModBindBase NoInfo Name } + : module id ModParams maybeAscription(SigExp) '=' ModExp + { let L floc (ID fname) = $2; + in ModBind fname $3 (fmap (,NoInfo) $4) $6 Nothing (srcspan $1 $>) + } + +ModParam :: { ModParamBase NoInfo Name } + : '(' id ':' SigExp ')' { let L _ (ID name) = $2 in ModParam name $4 NoInfo (srcspan $1 $>) } + +ModParams :: { [ModParamBase NoInfo Name] } + : ModParam ModParams { $1 : $2 } + | { [] } + +Spec :: { SpecBase NoInfo Name } + : val id TypeParams ':' TypeExpDecl + { let L loc (ID name) = $2 + in ValSpec name $3 $5 Nothing (srcspan $1 $>) } + | val BindingBinOp ':' TypeExpDecl + { ValSpec $2 [] $4 Nothing (srcspan $1 $>) } + | val BindingUnOp ':' TypeExpDecl + { ValSpec $2 [] $4 Nothing (srcspan $1 $>) } + | TypeAbbr + { TypeAbbrSpec $1 } + | type id TypeParams + { let L _ (ID name) = $2 + in TypeSpec Unlifted name $3 Nothing (srcspan $1 $>) } + | type 'id[' id ']' TypeParams + { let L _ (INDEXING name) = $2; L ploc (ID pname) = $3 + in TypeSpec Unlifted name (TypeParamDim pname ploc : $5) Nothing (srcspan $1 $>) } + | type '^' id TypeParams + { let L _ (ID name) = $3 + in TypeSpec Lifted name $4 Nothing (srcspan $1 $>) } + | type '^' 'id[' id ']' TypeParams + { let L _ (INDEXING name) = $3; L ploc (ID pname) = $4 + in TypeSpec Lifted name (TypeParamDim pname ploc : $6) Nothing (srcspan $1 $>) } + | module id ':' SigExp + { let L _ (ID name) = $2 + in ModSpec name $4 Nothing (srcspan $1 $>) } + | include SigExp + { IncludeSpec $2 (srcspan $1 $>) } + | Doc Spec + { addDocSpec $1 $2 } + +Specs :: { [SpecBase NoInfo Name] } + : Spec Specs { $1 : $2 } + | { [] } + +TypeParam :: { TypeParamBase Name } + : '[' id ']' { let L _ (ID name) = $2 in TypeParamDim name (srcspan $1 $>) } + | '\'' id { let L _ (ID name) = $2 in TypeParamType Unlifted name (srcspan $1 $>) } + | '\'^' id { let L _ (ID name) = $2 in TypeParamType Lifted name (srcspan $1 $>) } + +TypeParams :: { [TypeParamBase Name] } + : TypeParam TypeParams { $1 : $2 } + | { [] } + +TypeParams1 :: { (TypeParamBase Name, [TypeParamBase Name]) } + : TypeParam TypeParams { ($1, $2) } + +UnOp :: { (QualName Name, SrcLoc) } + : qunop { let L loc (QUALUNOP qs v) = $1 in (QualName qs v, loc) } + | unop { let L loc (UNOP v) = $1 in (qualName v, loc) } + +-- Note that this production does not include Minus, but does include +-- operator sections. +BinOp :: { QualName Name } + : '+...' { binOpName $1 } + | '-...' { binOpName $1 } + | '*...' { binOpName $1 } + | '*' { qualName (nameFromString "*") } + | '/...' { binOpName $1 } + | '%...' { binOpName $1 } + | '//...' { binOpName $1 } + | '%%...' { binOpName $1 } + | '==...' { binOpName $1 } + | '!=...' { binOpName $1 } + | '<...' { binOpName $1 } + | '<=...' { binOpName $1 } + | '>...' { binOpName $1 } + | '>=...' { binOpName $1 } + | '&&...' { binOpName $1 } + | '||...' { binOpName $1 } + | '**...' { binOpName $1 } + | '^...' { binOpName $1 } + | '^' { qualName (nameFromString "^") } + | '&...' { binOpName $1 } + | '|...' { binOpName $1 } + | '>>...' { binOpName $1 } + | '<<...' { binOpName $1 } + | '<|...' { binOpName $1 } + | '|>...' { binOpName $1 } + | '<' { qualName (nameFromString "<") } + | '`' QualName '`' { fst $2 } + +BindingUnOp :: { Name } + : UnOp {% let (QualName qs name, _) = $1 in do + unless (null qs) $ fail "Cannot use a qualified name in binding position." + return name } + +BindingBinOp :: { Name } + : BinOp {% let QualName qs name = $1 in do + unless (null qs) $ fail "Cannot use a qualified name in binding position." + return name } + | '-' { nameFromString "-" } + +BindingId :: { (Name, SrcLoc) } + : id { let L loc (ID name) = $1 in (name, loc) } + | '(' BindingBinOp ')' { ($2, $1) } + | '(' BindingUnOp ')' { ($2, $1) } + +Val :: { ValBindBase NoInfo Name } +Val : let BindingId TypeParams FunParams maybeAscription(TypeExpDecl) '=' Exp + { let (name, _) = $2 + in ValBind (name==defaultEntryPoint) name (fmap declaredType $5) NoInfo + $3 $4 $7 Nothing (srcspan $1 $>) + } + + | entry BindingId TypeParams FunParams maybeAscription(TypeExpDecl) '=' Exp + { let (name, loc) = $2 + in ValBind True name (fmap declaredType $5) NoInfo + $3 $4 $7 Nothing (srcspan $1 $>) } + + | let FunParam BindingBinOp FunParam maybeAscription(TypeExpDecl) '=' Exp + { ValBind False $3 (fmap declaredType $5) NoInfo [] [$2,$4] $7 Nothing (srcspan $1 $>) + } + + | let BindingUnOp TypeParams FunParams maybeAscription(TypeExpDecl) '=' Exp + { let name = $2 + in ValBind (name==defaultEntryPoint) name (fmap declaredType $5) NoInfo + $3 $4 $7 Nothing (srcspan $1 $>) + } + +TypeExpDecl :: { TypeDeclBase NoInfo Name } + : TypeExp %prec bottom { TypeDecl $1 NoInfo } + +TypeAbbr :: { TypeBindBase NoInfo Name } +TypeAbbr : type id TypeParams '=' TypeExpDecl + { let L _ (ID name) = $2 + in TypeBind name $3 $5 Nothing (srcspan $1 $>) } + | type 'id[' id ']' TypeParams '=' TypeExpDecl + { let L loc (INDEXING name) = $2; L ploc (ID pname) = $3 + in TypeBind name (TypeParamDim pname ploc:$5) $7 Nothing (srcspan $1 $>) } + +TypeExp :: { UncheckedTypeExp } + : '(' id ':' TypeExp ')' '->' TypeExp + { let L _ (ID v) = $2 in TEArrow (Just v) $4 $7 (srcspan $1 $>) } + | TypeExpTerm '->' TypeExp + { TEArrow Nothing $1 $3 (srcspan $1 $>) } + | TypeExpTerm { $1 } + + +TypeExpTerm :: { UncheckedTypeExp } + : '*' TypeExpTerm + { TEUnique $2 (srcspan $1 $>) } + | '[' DimDecl ']' TypeExpTerm %prec indexprec + { TEArray $4 (fst $2) (srcspan $1 $>) } + | '[' ']' TypeExpTerm %prec indexprec + { TEArray $3 AnyDim (srcspan $1 $>) } + | TypeExpApply { $1 } + + -- Errors + | '[' DimDecl ']' %prec bottom + {% parseErrorAt (srcspan $1 $>) $ Just $ + unlines ["missing array row type.", + "Did you mean []" ++ pretty (fst $2) ++ "?"] + } + +TypeExpApply :: { UncheckedTypeExp } + : TypeExpApply TypeArg + { TEApply $1 $2 (srcspan $1 $>) } + | 'id[' DimDecl ']' + { let L loc (INDEXING v) = $1 + in TEApply (TEVar (qualName v) loc) (TypeArgExpDim (fst $2) loc) (srcspan $1 $>) } + | 'qid[' DimDecl ']' + { let L loc (QUALINDEXING qs v) = $1 + in TEApply (TEVar (QualName qs v) loc) (TypeArgExpDim (fst $2) loc) (srcspan $1 $>) } + | TypeExpAtom + { $1 } + +TypeExpAtom :: { UncheckedTypeExp } + : '(' TypeExp ')' { $2 } + | '(' ')' { TETuple [] (srcspan $1 $>) } + | '(' TypeExp ',' TupleTypes ')' { TETuple ($2:$4) (srcspan $1 $>) } + | '{' '}' { TERecord [] (srcspan $1 $>) } + | '{' FieldTypes1 '}' { TERecord $2 (srcspan $1 $>) } + | QualName { TEVar (fst $1) (snd $1) } + +TypeArg :: { TypeArgExp Name } + : '[' DimDecl ']' { TypeArgExpDim (fst $2) (srcspan $1 $>) } + | '[' ']' { TypeArgExpDim AnyDim (srcspan $1 $>) } + | TypeExpAtom { TypeArgExpType $1 } + +FieldType :: { (Name, UncheckedTypeExp) } +FieldType : FieldId ':' TypeExp { (fst $1, $3) } + +FieldTypes1 :: { [(Name, UncheckedTypeExp)] } +FieldTypes1 : FieldType { [$1] } + | FieldType ',' FieldTypes1 { $1 : $3 } + +TupleTypes :: { [UncheckedTypeExp] } + : TypeExp { [$1] } + | TypeExp ',' TupleTypes { $1 : $3 } + +DimDecl :: { (DimDecl Name, SrcLoc) } + : QualName + { (NamedDim (fst $1), snd $1) } + | intlit + { let L loc (INTLIT n) = $1 + in (ConstDim (fromIntegral n), loc) } + + -- Errors + | '#' {% parseErrorAt (srclocOf $1) $ Just $ + unlines ["found implicit size quantification.", + "This is no longer supported. Use explicit size parameters."] + } + + +FunParam :: { PatternBase NoInfo Name } +FunParam : InnerPattern { $1 } + +FunParams1 :: { (PatternBase NoInfo Name, [PatternBase NoInfo Name]) } +FunParams1 : FunParam { ($1, []) } + | FunParam FunParams1 { ($1, fst $2 : snd $2) } + +FunParams :: { [PatternBase NoInfo Name] } +FunParams : { [] } + | FunParam FunParams { $1 : $2 } + +QualName :: { (QualName Name, SrcLoc) } + : id FieldAccesses + { let L vloc (ID v) = $1 in + foldl (\(QualName qs v', loc) (y, yloc) -> + (QualName (qs ++ [v']) y, srcspan loc yloc)) + (qualName v, vloc) $2 } + +-- Expressions are divided into several layers. The first distinction +-- (between Exp and Exp2) is to factor out ascription, which we do not +-- permit inside array indices operations (there is an ambiguity with +-- array slices). +Exp :: { UncheckedExp } + : Exp ':' TypeExpDecl { Ascript $1 $3 (srcspan $1 $>) } + | Exp2 %prec ':' { $1 } + +Exp2 :: { UncheckedExp } + : if Exp then Exp else Exp %prec ifprec + { If $2 $4 $6 NoInfo (srcspan $1 $>) } + + | loop TypeParams Pattern LoopForm do Exp %prec ifprec + {% fmap (\t -> DoLoop $2 $3 t $4 $6 (srcspan $1 $>)) (patternExp $3) } + + | loop TypeParams Pattern '=' Exp LoopForm do Exp %prec ifprec + { DoLoop $2 $3 $5 $6 $8 (srcspan $1 $>) } + + | LetExp %prec letprec { $1 } + + | unsafe Exp2 { Unsafe $2 (srcspan $1 $>) } + | assert Atom Atom { Assert $2 $3 NoInfo (srcspan $1 $>) } + + | Exp2 '+...' Exp2 { binOp $1 $2 $3 } + | Exp2 '-...' Exp2 { binOp $1 $2 $3 } + | Exp2 '-' Exp2 { binOp $1 (L $2 (SYMBOL Minus [] (nameFromString "-"))) $3 } + | Exp2 '*...' Exp2 { binOp $1 $2 $3 } + | Exp2 '*' Exp2 { binOp $1 (L $2 (SYMBOL Times [] (nameFromString "*"))) $3 } + | Exp2 '/...' Exp2 { binOp $1 $2 $3 } + | Exp2 '%...' Exp2 { binOp $1 $2 $3 } + | Exp2 '//...' Exp2 { binOp $1 $2 $3 } + | Exp2 '%%...' Exp2 { binOp $1 $2 $3 } + | Exp2 '**...' Exp2 { binOp $1 $2 $3 } + | Exp2 '>>...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<<...' Exp2 { binOp $1 $2 $3 } + | Exp2 '&...' Exp2 { binOp $1 $2 $3 } + | Exp2 '|...' Exp2 { binOp $1 $2 $3 } + | Exp2 '&&...' Exp2 { binOp $1 $2 $3 } + | Exp2 '||...' Exp2 { binOp $1 $2 $3 } + | Exp2 '^...' Exp2 { binOp $1 $2 $3 } + | Exp2 '^' Exp2 { binOp $1 (L $2 (SYMBOL Xor [] (nameFromString "^"))) $3 } + | Exp2 '==...' Exp2 { binOp $1 $2 $3 } + | Exp2 '!=...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<=...' Exp2 { binOp $1 $2 $3 } + | Exp2 '>...' Exp2 { binOp $1 $2 $3 } + | Exp2 '>=...' Exp2 { binOp $1 $2 $3 } + | Exp2 '|>...' Exp2 { binOp $1 $2 $3 } + | Exp2 '<|...' Exp2 { binOp $1 $2 $3 } + + | Exp2 '<' Exp2 { binOp $1 (L $2 (SYMBOL Less [] (nameFromString "<"))) $3 } + | Exp2 '`' QualName '`' Exp2 { BinOp (fst $3) NoInfo ($1, NoInfo) ($5, NoInfo) NoInfo (srclocOf $1) } + + | Exp2 '...' Exp2 { Range $1 Nothing (ToInclusive $3) NoInfo (srcspan $1 $>) } + | Exp2 '..<' Exp2 { Range $1 Nothing (UpToExclusive $3) NoInfo (srcspan $1 $>) } + | Exp2 '..>' Exp2 { Range $1 Nothing (DownToExclusive $3) NoInfo (srcspan $1 $>) } + | Exp2 '..' Exp2 '...' Exp2 { Range $1 (Just $3) (ToInclusive $5) NoInfo (srcspan $1 $>) } + | Exp2 '..' Exp2 '..<' Exp2 { Range $1 (Just $3) (UpToExclusive $5) NoInfo (srcspan $1 $>) } + | Exp2 '..' Exp2 '..>' Exp2 { Range $1 (Just $3) (DownToExclusive $5) NoInfo (srcspan $1 $>) } + | Exp2 '..' Atom {% twoDotsRange $2 } + | Atom '..' Exp2 {% twoDotsRange $2 } + | '-' Exp2 + { Negate $2 $1 } + + | Exp2 with '[' DimIndices ']' '=' Exp2 + { Update $1 $4 $7 (srcspan $1 $>) } + + | Exp2 with FieldAccesses_ '=' Exp2 + { RecordUpdate $1 (map fst $3) $5 NoInfo (srcspan $1 $>) } + + | Exp2 with FieldAccesses_ '<-' Exp2 + { RecordUpdate $1 (map fst $3) $5 NoInfo (srcspan $1 $>) } + | Exp2 with '[' DimIndices ']' '<-' Exp2 + { Update $1 $4 $7 (srcspan $1 $>) } + + | '\\' TypeParams FunParams1 maybeAscription(TypeExpTerm) '->' Exp + { Lambda $2 (fst $3 : snd $3) $6 (fmap (flip TypeDecl NoInfo) $4) NoInfo (srcspan $1 $>) } + + | Apply { $1 } + +Apply :: { UncheckedExp } + : Apply Atom %prec juxtprec + { Apply $1 $2 NoInfo NoInfo (srcspan $1 $>) } + | UnOp Atom %prec juxtprec + { Apply (Var (fst $1) NoInfo (snd $1)) $2 NoInfo NoInfo (srcspan (snd $1) $>) } + | Atom %prec juxtprec + { $1 } + +Atom :: { UncheckedExp } +Atom : PrimLit { Literal (fst $1) (snd $1) } + | intlit { let L loc (INTLIT x) = $1 in IntLit x NoInfo loc } + | floatlit { let L loc (FLOATLIT x) = $1 in FloatLit x NoInfo loc } + | stringlit { let L loc (STRINGLIT s) = $1 in + ArrayLit (map (flip Literal loc . SignedValue . Int32Value . fromIntegral . ord) s) NoInfo loc } + | '(' Exp ')' FieldAccesses + { foldl (\x (y, _) -> Project y x NoInfo (srclocOf x)) + (Parens $2 (srcspan $1 $3)) + $4 } + | '(' Exp ')[' DimIndices ']' { Index (Parens $2 $1) $4 NoInfo (srcspan $1 $>) } + | '(' Exp ',' Exps1 ')' { TupLit ($2 : fst $4 : snd $4) (srcspan $1 $>) } + | '(' ')' { TupLit [] (srcspan $1 $>) } + | '[' Exps1 ']' { ArrayLit (fst $2:snd $2) NoInfo (srcspan $1 $>) } + | '[' ']' { ArrayLit [] NoInfo (srcspan $1 $>) } + + | QualVarSlice FieldAccesses + { let (v,slice,loc) = $1 + in foldl (\x (y, _) -> Project y x NoInfo (srclocOf x)) + (Index (Var v NoInfo loc) slice NoInfo loc) + $2 } + | QualName + { Var (fst $1) NoInfo (snd $1) } + | '{' Fields '}' { RecordLit $2 (srcspan $1 $>) } + | 'qid.(' Exp ')' + { let L loc (QUALPAREN qs name) = $1 in QualParens (QualName qs name) $2 loc } + + -- Operator sections. + | '(' UnOp ')' + { Var (fst $2) NoInfo (srcspan (snd $2) $>) } + | '(' '-' ')' + { OpSection (qualName (nameFromString "-")) NoInfo (srcspan $1 $>) } + | '(' Exp2 '-' ')' + { OpSectionLeft (qualName (nameFromString "-")) + NoInfo $2 (NoInfo, NoInfo) NoInfo (srcspan $1 $>) } + | '(' BinOp Exp2 ')' + { OpSectionRight $2 NoInfo $3 (NoInfo, NoInfo) NoInfo (srcspan $1 $>) } + | '(' Exp2 BinOp ')' + { OpSectionLeft $3 NoInfo $2 (NoInfo, NoInfo) NoInfo (srcspan $1 $>) } + | '(' BinOp ')' + { OpSection $2 NoInfo (srcspan $1 $>) } + + | '(' FieldAccess FieldAccesses ')' + { ProjectSection (map fst ($2:$3)) NoInfo (srcspan $1 $>) } + + | '(' '.' '[' DimIndices ']' ')' + { IndexSection $4 NoInfo (srcspan $1 $>) } + + +PrimLit :: { (PrimValue, SrcLoc) } + : true { (BoolValue True, $1) } + | false { (BoolValue False, $1) } + + | i8lit { let L loc (I8LIT num) = $1 in (SignedValue $ Int8Value num, loc) } + | i16lit { let L loc (I16LIT num) = $1 in (SignedValue $ Int16Value num, loc) } + | i32lit { let L loc (I32LIT num) = $1 in (SignedValue $ Int32Value num, loc) } + | i64lit { let L loc (I64LIT num) = $1 in (SignedValue $ Int64Value num, loc) } + + | u8lit { let L loc (U8LIT num) = $1 in (UnsignedValue $ Int8Value $ fromIntegral num, loc) } + | u16lit { let L loc (U16LIT num) = $1 in (UnsignedValue $ Int16Value $ fromIntegral num, loc) } + | u32lit { let L loc (U32LIT num) = $1 in (UnsignedValue $ Int32Value $ fromIntegral num, loc) } + | u64lit { let L loc (U64LIT num) = $1 in (UnsignedValue $ Int64Value $ fromIntegral num, loc) } + + | f32lit { let L loc (F32LIT num) = $1 in (FloatValue $ Float32Value num, loc) } + | f64lit { let L loc (F64LIT num) = $1 in (FloatValue $ Float64Value num, loc) } + + | charlit { let L loc (CHARLIT char) = $1 + in (SignedValue $ Int32Value $ fromIntegral $ ord char, loc) } + +Exps1 :: { (UncheckedExp, [UncheckedExp]) } + : Exps1_ { case reverse (snd $1 : fst $1) of + [] -> (snd $1, []) + y:ys -> (y, ys) } + +Exps1_ :: { ([UncheckedExp], UncheckedExp) } + : Exps1_ ',' Exp { (snd $1 : fst $1, $3) } + | Exp { ([], $1) } + +FieldAccess :: { (Name, SrcLoc) } + : '.' FieldId { (fst $2, srcspan $1 (snd $>)) } + +FieldAccesses :: { [(Name, SrcLoc)] } + : FieldAccess FieldAccesses { $1 : $2 } + | { [] } + +FieldAccesses_ :: { [(Name, SrcLoc)] } + : FieldId FieldAccesses { (fst $1, snd $1) : $2 } + +Field :: { FieldBase NoInfo Name } + : FieldId '=' Exp { RecordFieldExplicit (fst $1) $3 (srcspan (snd $1) $>) } + | id { let L loc (ID s) = $1 in RecordFieldImplicit s NoInfo loc } + +Fields :: { [FieldBase NoInfo Name] } + : Fields1 { $1 } + | { [] } + +Fields1 :: { [FieldBase NoInfo Name] } + : Field ',' Fields1 { $1 : $3 } + | Field { [$1] } + +LetExp :: { UncheckedExp } + : let Pattern '=' Exp LetBody + { LetPat [] $2 $4 $5 (srcspan $1 $>) } + | let TypeParams1 Pattern '=' Exp LetBody + { LetPat (fst $2 : snd $2) $3 $5 $6 (srcspan $1 $>) } + + | let id TypeParams FunParams1 maybeAscription(TypeExpDecl) '=' Exp LetBody + { let L _ (ID name) = $2 + in LetFun name ($3, fst $4 : snd $4, (fmap declaredType $5), NoInfo, $7) $8 (srcspan $1 $>) } + + | let VarSlice '=' Exp LetBody + { let (v,slice,loc) = $2; ident = Ident v NoInfo loc + in LetWith ident ident slice $4 $5 (srcspan $1 $>) } + +LetBody :: { UncheckedExp } + : in Exp %prec letprec { $2 } + | LetExp %prec letprec { $1 } + +LoopForm :: { LoopFormBase NoInfo Name } +LoopForm : for VarId '<' Exp + { For $2 $4 } + | for Pattern in Exp + { ForIn $2 $4 } + | while Exp + { While $2 } + +VarSlice :: { (Name, [UncheckedDimIndex], SrcLoc) } + : 'id[' DimIndices ']' + { let L _ (INDEXING v) = $1 + in (v, $2, srcspan $1 $>) } + +QualVarSlice :: { (QualName Name, [UncheckedDimIndex], SrcLoc) } + : VarSlice + { let (x, y, z) = $1 in (qualName x, y, z) } + | 'qid[' DimIndices ']' + { let L _ (QUALINDEXING qs v) = $1 in (QualName qs v, $2, srcspan $1 $>) } + +DimIndex :: { UncheckedDimIndex } + : Exp2 { DimFix $1 } + | Exp2 ':' Exp2 { DimSlice (Just $1) (Just $3) Nothing } + | Exp2 ':' { DimSlice (Just $1) Nothing Nothing } + | ':' Exp2 { DimSlice Nothing (Just $2) Nothing } + | ':' { DimSlice Nothing Nothing Nothing } + | Exp2 ':' Exp2 ':' Exp2 { DimSlice (Just $1) (Just $3) (Just $5) } + | ':' Exp2 ':' Exp2 { DimSlice Nothing (Just $2) (Just $4) } + | Exp2 ':' ':' Exp2 { DimSlice (Just $1) Nothing (Just $4) } + | ':' ':' Exp2 { DimSlice Nothing Nothing (Just $3) } + +DimIndices :: { [UncheckedDimIndex] } + : { [] } + | DimIndices1 { fst $1 : snd $1 } + +DimIndices1 :: { (UncheckedDimIndex, [UncheckedDimIndex]) } + : DimIndex { ($1, []) } + | DimIndex ',' DimIndices1 { ($1, fst $3 : snd $3) } + +VarId :: { IdentBase NoInfo Name } +VarId : id { let L loc (ID name) = $1 in Ident name NoInfo loc } + +FieldId :: { (Name, SrcLoc) } + : id { let L loc (ID name) = $1 in (name, loc) } + | intlit { let L loc (INTLIT n) = $1 in (nameFromString (show n), loc) } + +Pattern :: { PatternBase NoInfo Name } +Pattern : InnerPattern ':' TypeExpDecl { PatternAscription $1 $3 (srcspan $1 $>) } + | InnerPattern { $1 } + +Patterns1 :: { [PatternBase NoInfo Name] } + : Pattern { [$1] } + | Pattern ',' Patterns1 { $1 : $3 } + +InnerPattern :: { PatternBase NoInfo Name } +InnerPattern : id { let L loc (ID name) = $1 in Id name NoInfo loc } + | '(' BindingBinOp ')' { Id $2 NoInfo (srcspan $1 $>) } + | '(' BindingUnOp ')' { Id $2 NoInfo (srcspan $1 $>) } + | '_' { Wildcard NoInfo $1 } + | '(' ')' { TuplePattern [] (srcspan $1 $>) } + | '(' Pattern ')' { PatternParens $2 (srcspan $1 $>) } + | '(' Pattern ',' Patterns1 ')' { TuplePattern ($2:$4) (srcspan $1 $>) } + | '{' FieldPatterns '}' { RecordPattern $2 (srcspan $1 $>) } + +FieldPattern :: { (Name, PatternBase NoInfo Name) } + : FieldId '=' Pattern + { (fst $1, $3) } + | FieldId ':' TypeExpDecl + { (fst $1, PatternAscription (Id (fst $1) NoInfo (snd $1)) $3 (srcspan (snd $1) $>)) } + | FieldId + { (fst $1, Id (fst $1) NoInfo (snd $1)) } + +FieldPatterns :: { [(Name, PatternBase NoInfo Name)] } + : FieldPatterns1 { $1 } + | { [] } + +FieldPatterns1 :: { [(Name, PatternBase NoInfo Name)] } + : FieldPattern ',' FieldPatterns1 { $1 : $3 } + | FieldPattern { [$1] } + + +maybeAscription(p) : ':' p { Just $2 } + | { Nothing } + +Value :: { Value } +Value : IntValue { $1 } + | FloatValue { $1 } + | StringValue { $1 } + | BoolValue { $1 } + | ArrayValue { $1 } + +CatValues :: { [Value] } +CatValues : Value CatValues { $1 : $2 } + | { [] } + +PrimType :: { PrimType } + : id {% let L _ (ID s) = $1 in primTypeFromName s } + +IntValue :: { Value } + : SignedLit { PrimValue (SignedValue (fst $1)) } + | '-' SignedLit { PrimValue (SignedValue (intNegate (fst $2))) } + | UnsignedLit { PrimValue (UnsignedValue (fst $1)) } + +FloatValue :: { Value } + : FloatLit { PrimValue (FloatValue (fst $1)) } + | '-' FloatLit { PrimValue (FloatValue (floatNegate (fst $2))) } + +StringValue :: { Value } +StringValue : stringlit { let L pos (STRINGLIT s) = $1 in + ArrayValue (arrayFromList $ map (PrimValue . SignedValue . Int32Value . fromIntegral . ord) s) $ Prim $ Signed Int32 } + +BoolValue :: { Value } +BoolValue : true { PrimValue $ BoolValue True } + | false { PrimValue $ BoolValue False } + +SignedLit :: { (IntValue, SrcLoc) } + : i8lit { let L loc (I8LIT num) = $1 in (Int8Value num, loc) } + | i16lit { let L loc (I16LIT num) = $1 in (Int16Value num, loc) } + | i32lit { let L loc (I32LIT num) = $1 in (Int32Value num, loc) } + | i64lit { let L loc (I64LIT num) = $1 in (Int64Value num, loc) } + | intlit { let L loc (INTLIT num) = $1 in (Int32Value $ fromInteger num, loc) } + | charlit { let L loc (CHARLIT char) = $1 in (Int32Value $ fromIntegral $ ord char, loc) } + +UnsignedLit :: { (IntValue, SrcLoc) } + : u8lit { let L pos (U8LIT num) = $1 in (Int8Value $ fromIntegral num, pos) } + | u16lit { let L pos (U16LIT num) = $1 in (Int16Value $ fromIntegral num, pos) } + | u32lit { let L pos (U32LIT num) = $1 in (Int32Value $ fromIntegral num, pos) } + | u64lit { let L pos (U64LIT num) = $1 in (Int64Value $ fromIntegral num, pos) } + +FloatLit :: { (FloatValue, SrcLoc) } + : f32lit { let L loc (F32LIT num) = $1 in (Float32Value num, loc) } + | f64lit { let L loc (F64LIT num) = $1 in (Float64Value num, loc) } + | QualName {% let (qn, loc) = $1 in + if qn == QualName [nameFromString "f32"] (nameFromString "inf") + then return (Float32Value (1/0), loc) + else if qn == QualName [nameFromString "f32"] (nameFromString "nan") + then return (Float32Value (0/0), loc) + else if qn == QualName [nameFromString "f64"] (nameFromString "inf") + then return (Float64Value (1/0), loc) + else if qn == QualName [nameFromString "f64"] (nameFromString "nan") + then return (Float64Value (0/0), loc) + else parseErrorAt (snd $1) Nothing } + | floatlit { let L loc (FLOATLIT num) = $1 in (Float64Value num, loc) } + +ArrayValue :: { Value } +ArrayValue : '[' Value ']' + {% return $ ArrayValue (arrayFromList [$2]) $ toStruct $ valueType $2 + } + | '[' Value ',' Values ']' + {% case combArrayElements $2 $4 of + Left e -> throwError e + Right v -> return $ ArrayValue (arrayFromList $ $2:$4) $ valueType v + } + | id '(' PrimType ')' + {% ($1 `mustBe` "empty") >> return (ArrayValue (listArray (0,-1) []) (Prim $3)) } + | id '(' RowType ')' + {% ($1 `mustBe` "empty") >> return (ArrayValue (listArray (0,-1) []) $3) } + + -- Errors + | '[' ']' + {% emptyArrayError $1 } + +RowType :: { TypeBase () () } +RowType : '[' ']' RowType { fromJust $ arrayOf $3 (rank 1) Nonunique } + | '[' ']' PrimType { fromJust $ arrayOf (Prim $3) (rank 1) Nonunique } + +Values :: { [Value] } +Values : Value ',' Values { $1 : $3 } + | Value { [$1] } + | { [] } + +{ + +addDoc :: DocComment -> UncheckedDec -> UncheckedDec +addDoc doc (ValDec val) = ValDec (val { valBindDoc = Just doc }) +addDoc doc (TypeDec tp) = TypeDec (tp { typeDoc = Just doc }) +addDoc doc (SigDec sig) = SigDec (sig { sigDoc = Just doc }) +addDoc doc (ModDec mod) = ModDec (mod { modDoc = Just doc }) +addDoc _ dec = dec + +addDocSpec :: DocComment -> SpecBase NoInfo Name -> SpecBase NoInfo Name +addDocSpec doc (TypeAbbrSpec tpsig) = TypeAbbrSpec (tpsig { typeDoc = Just doc }) +addDocSpec doc val@(ValSpec {}) = val { specDoc = Just doc } +addDocSpec doc (TypeSpec l name ps _ loc) = TypeSpec l name ps (Just doc) loc +addDocSpec doc (ModSpec name se _ loc) = ModSpec name se (Just doc) loc +addDocSpec _ spec = spec + +reverseNonempty :: (a, [a]) -> (a, [a]) +reverseNonempty (x, l) = + case reverse (x:l) of + x':rest -> (x', rest) + [] -> (x, []) + +mustBe (L loc (ID got)) expected + | nameToString got == expected = return () +mustBe (L loc _) expected = + parseErrorAt loc $ Just $ + "Only the keyword '" ++ expected ++ "' may appear here." + +data ParserEnv = ParserEnv { + parserFile :: FilePath + } + +type ParserMonad a = + ExceptT String ( + StateT ParserEnv ( + StateT ([L Token], Pos) ReadLineMonad)) a + +data ReadLineMonad a = Value a + | GetLine (Maybe T.Text -> ReadLineMonad a) + +readLineFromMonad :: ReadLineMonad (Maybe T.Text) +readLineFromMonad = GetLine Value + +instance Monad ReadLineMonad where + return = Value + Value x >>= f = f x + GetLine g >>= f = GetLine $ \s -> g s >>= f + +instance Functor ReadLineMonad where + f `fmap` m = do x <- m + return $ f x + +instance Applicative ReadLineMonad where + (<*>) = ap + +getLinesFromM :: Monad m => m T.Text -> ReadLineMonad a -> m a +getLinesFromM _ (Value x) = return x +getLinesFromM fetch (GetLine f) = do + s <- fetch + getLinesFromM fetch $ f $ Just s + +getLinesFromTexts :: [T.Text] -> ReadLineMonad a -> Either String a +getLinesFromTexts _ (Value x) = Right x +getLinesFromTexts (x : xs) (GetLine f) = getLinesFromTexts xs $ f $ Just x +getLinesFromTexts [] (GetLine f) = getLinesFromTexts [] $ f Nothing + +getNoLines :: ReadLineMonad a -> Either String a +getNoLines (Value x) = Right x +getNoLines (GetLine f) = getNoLines $ f Nothing + +combArrayElements :: Value + -> [Value] + -> Either String Value +combArrayElements t ts = foldM comb t ts + where comb x y + | valueType x == valueType y = Right x + | otherwise = Left $ "Elements " ++ pretty x ++ " and " ++ + pretty y ++ " cannot exist in same array." + +arrayFromList :: [a] -> Array Int a +arrayFromList l = listArray (0, length l-1) l + +patternExp :: UncheckedPattern -> ParserMonad UncheckedExp +patternExp (Id v _ loc) = return $ Var (qualName v) NoInfo loc +patternExp (TuplePattern pats loc) = TupLit <$> (mapM patternExp pats) <*> return loc +patternExp (Wildcard _ loc) = parseErrorAt loc $ Just "cannot have wildcard here." +patternExp (PatternAscription pat _ _) = patternExp pat +patternExp (PatternParens pat _) = patternExp pat +patternExp (RecordPattern fs loc) = RecordLit <$> mapM field fs <*> pure loc + where field (name, pat) = RecordFieldExplicit name <$> patternExp pat <*> pure loc + +eof :: Pos -> L Token +eof pos = L (SrcLoc $ Loc pos pos) EOF + +binOpName (L _ (SYMBOL _ qs op)) = QualName qs op + +binOp x (L _ (SYMBOL _ qs op)) y = + BinOp (QualName qs op) NoInfo (x, NoInfo) (y, NoInfo) NoInfo $ + srcspan x y + +getTokens :: ParserMonad ([L Token], Pos) +getTokens = lift $ lift get + +putTokens :: ([L Token], Pos) -> ParserMonad () +putTokens = lift . lift . put + +primTypeFromName :: Name -> ParserMonad PrimType +primTypeFromName s = maybe boom return $ M.lookup s namesToPrimTypes + where boom = fail $ "No type named " ++ nameToString s + +getFilename :: ParserMonad FilePath +getFilename = lift $ gets parserFile + +intNegate :: IntValue -> IntValue +intNegate (Int8Value v) = Int8Value (-v) +intNegate (Int16Value v) = Int16Value (-v) +intNegate (Int32Value v) = Int32Value (-v) +intNegate (Int64Value v) = Int64Value (-v) + +floatNegate :: FloatValue -> FloatValue +floatNegate (Float32Value v) = Float32Value (-v) +floatNegate (Float64Value v) = Float64Value (-v) + +readLine :: ParserMonad (Maybe T.Text) +readLine = lift $ lift $ lift readLineFromMonad + +lexer :: (L Token -> ParserMonad a) -> ParserMonad a +lexer cont = do + (ts, pos) <- getTokens + case ts of + [] -> do + ended <- lift $ runExceptT $ cont $ eof pos + case ended of + Right x -> return x + Left parse_e -> do + line <- readLine + ts' <- + case line of Nothing -> throwError parse_e + Just line' -> return $ scanTokensText (advancePos pos '\n') line' + (ts'', pos') <- + case ts' of Right x -> return x + Left lex_e -> throwError lex_e + case ts'' of + [] -> cont $ eof pos + xs -> do + putTokens (xs, pos') + lexer cont + (x : xs) -> do + putTokens (xs, pos) + cont x + +parseError :: L Token -> ParserMonad a +parseError (L loc EOF) = + parseErrorAt (srclocOf loc) $ Just "unexpected end of file." +parseError (L loc DOC{}) = + parseErrorAt (srclocOf loc) $ + Just "documentation comments ('-- |') are only permitted when preceding declarations." +parseError tok = parseErrorAt (srclocOf tok) Nothing + +parseErrorAt :: SrcLoc -> Maybe String -> ParserMonad a +parseErrorAt loc Nothing = throwError $ "Error at " ++ locStr loc ++ ": Parse error." +parseErrorAt loc (Just s) = throwError $ "Error at " ++ locStr loc ++ ": " ++ s + +emptyArrayError :: SrcLoc -> ParserMonad a +emptyArrayError loc = + parseErrorAt loc $ + Just "write empty arrays as 'empty(t)', for element type 't'.\n" + +twoDotsRange :: SrcLoc -> ParserMonad a +twoDotsRange loc = parseErrorAt loc $ Just "use '...' for ranges, not '..'.\n" + +--- Now for the parser interface. + +-- | A parse error. Use 'show' to get a human-readable description. +data ParseError = ParseError String + +instance Show ParseError where + show (ParseError s) = s + +parseInMonad :: ParserMonad a -> FilePath -> T.Text + -> ReadLineMonad (Either ParseError a) +parseInMonad p file program = + either (Left . ParseError) Right <$> either (return . Left) + (evalStateT (evalStateT (runExceptT p) env)) + (scanTokensText (Pos file 1 1 0) program) + where env = ParserEnv file + +parseIncremental :: ParserMonad a -> FilePath -> T.Text + -> Either ParseError a +parseIncremental p file program = + either (Left . ParseError) id + $ getLinesFromTexts (T.lines program) + $ parseInMonad p file mempty + +parse :: ParserMonad a -> FilePath -> T.Text + -> Either ParseError a +parse p file program = + either (Left . ParseError) id + $ getNoLines $ parseInMonad p file program + +-- | Parse an Futhark expression incrementally from monadic actions, using the +-- 'FilePath' as the source name for error messages. +parseExpIncrM :: Monad m => + m T.Text -> FilePath -> T.Text + -> m (Either ParseError UncheckedExp) +parseExpIncrM fetch file program = + getLinesFromM fetch $ parseInMonad expression file program + +-- | Parse either an expression or a declaration incrementally; +-- favouring declarations in case of ambiguity. +parseDecOrExpIncrM :: Monad m => + m T.Text -> FilePath -> T.Text + -> m (Either ParseError (Either UncheckedDec UncheckedExp)) +parseDecOrExpIncrM fetch file input = + case parseInMonad declaration file input of + Value Left{} -> fmap Right <$> parseExpIncrM fetch file input + Value (Right d) -> return $ Right $ Left d + GetLine c -> do + l <- fetch + parseDecOrExpIncrM fetch file $ input <> "\n" <> l +} diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Pretty.hs @@ -0,0 +1,471 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +-- | Futhark prettyprinter. This module defines 'Pretty' instances +-- for the AST defined in "Language.Futhark.Syntax". +module Language.Futhark.Pretty + ( pretty + , prettyTuple + , leadingOperator + , IsName(..) + , prettyName + , Annot + ) +where + +import Control.Monad +import Data.Array +import Data.Functor +import qualified Data.Map.Strict as M +import Data.List +import Data.Maybe +import Data.Monoid +import Data.Ord +import Data.Word + +import Prelude + +import Futhark.Util.Pretty +import Futhark.Util + +import Language.Futhark.Syntax +import Language.Futhark.Attributes + +commastack :: [Doc] -> Doc +commastack = align . stack . punctuate comma + +-- | A class for types that are variable names in the Futhark source +-- language. This is used instead of a mere 'Pretty' instance because +-- in the compiler frontend we want to print VNames differently +-- depending on whether the FUTHARK_COMPILER_DEBUGGING environment +-- variable is set, yet in the backend we want to always print VNames +-- with the tag. To avoid erroneously using the 'Pretty' instance for +-- VNames, we in fact only define it inside the modules for the core +-- language (as an orphan instance). +class IsName v where + pprName :: v -> Doc + +-- | Depending on the environment variable FUTHARK_COMPILER_DEBUGGING, +-- VNames are printed as either the name with an internal tag, or just +-- the base name. +instance IsName VName where + pprName | isEnvVarSet "FUTHARK_COMPILER_DEBUGGING" False = + \(VName vn i) -> ppr vn <> text "_" <> text (show i) + | otherwise = ppr . baseName + +instance IsName Name where + pprName = ppr + +prettyName :: IsName v => v -> String +prettyName = prettyDoc 80 . pprName + +-- | Class for type constructors that represent annotations. Used in +-- the prettyprinter to either print the original AST, or the computed +-- attribute. +class Annot f where + unAnnot :: f a -> Maybe a + +instance Annot NoInfo where + unAnnot = const Nothing + +instance Annot Info where + unAnnot = Just . unInfo + +pprAnnot :: (Annot f, Pretty a, Pretty b) => a -> f b -> Doc +pprAnnot a b = maybe (ppr a) ppr $ unAnnot b + +instance Pretty Value where + ppr (PrimValue bv) = ppr bv + ppr (ArrayValue a t) + | [] <- elems a = text "empty" <> parens (ppr t) + | Array{} <- t = brackets $ commastack $ map ppr $ elems a + | otherwise = brackets $ commasep $ map ppr $ elems a + +instance Pretty PrimValue where + ppr (UnsignedValue (Int8Value v)) = + text (show (fromIntegral v::Word8)) <> text "u8" + ppr (UnsignedValue (Int16Value v)) = + text (show (fromIntegral v::Word16)) <> text "u16" + ppr (UnsignedValue (Int32Value v)) = + text (show (fromIntegral v::Word32)) <> text "u32" + ppr (UnsignedValue (Int64Value v)) = + text (show (fromIntegral v::Word64)) <> text "u64" + ppr (SignedValue v) = ppr v + ppr (BoolValue True) = text "true" + ppr (BoolValue False) = text "false" + ppr (FloatValue v) = ppr v + +instance IsName vn => Pretty (DimDecl vn) where + ppr AnyDim = mempty + ppr (NamedDim v) = ppr v + ppr (ConstDim n) = ppr n + + +instance IsName vn => Pretty (ShapeDecl (DimDecl vn)) where + ppr (ShapeDecl ds) = mconcat (map (brackets . ppr) ds) + +instance Pretty (ShapeDecl ()) where + ppr (ShapeDecl ds) = mconcat $ replicate (length ds) $ text "[]" + +instance Pretty (ShapeDecl dim) => Pretty (RecordArrayElemTypeBase dim as) where + ppr (RecordArrayElem et) = ppr et + ppr (RecordArrayArrayElem et shape u) = + ppr u <> ppr shape <> ppr et + +instance Pretty (ShapeDecl dim) => Pretty (ArrayElemTypeBase dim as) where + ppr (ArrayPrimElem pt _) = ppr pt + ppr (ArrayPolyElem v args _) = + ppr (qualNameFromTypeName v) <+> spread (map ppr args) + ppr (ArrayRecordElem fs) + | Just ts <- areTupleFields fs = + parens (commasep $ map ppr ts) + | otherwise = + braces (commasep $ map ppField $ M.toList fs) + where ppField (name, t) = text (nameToString name) <> colon <+> ppr t + +instance Pretty (ShapeDecl dim) => Pretty (TypeBase dim as) where + ppr = pprPrec 0 + pprPrec _ (Prim et) = ppr et + pprPrec _ (TypeVar _ u et targs) = + ppr u <> ppr (qualNameFromTypeName et) <+> spread (map ppr targs) + pprPrec _ (Array at shape u) = ppr u <> ppr shape <> ppr at + pprPrec _ (Record fs) + | Just ts <- areTupleFields fs = + parens $ commasep $ map ppr ts + | otherwise = + braces $ commasep $ map ppField $ M.toList fs + where ppField (name, t) = text (nameToString name) <> colon <+> ppr t + pprPrec p (Arrow _ (Just v) t1 t2) = + parensIf (p > 0) $ + parens (pprName v <> colon <+> ppr t1) <+> text "->" <+> ppr t2 + pprPrec p (Arrow _ Nothing t1 t2) = + parensIf (p > 0) $ pprPrec 1 t1 <+> text "->" <+> ppr t2 + +instance Pretty (ShapeDecl dim) => Pretty (TypeArg dim as) where + ppr (TypeArgDim d _) = ppr $ ShapeDecl [d] + ppr (TypeArgType t _) = ppr t + +instance (Eq vn, IsName vn) => Pretty (TypeExp vn) where + ppr (TEUnique t _) = text "*" <> ppr t + ppr (TEArray at d _) = ppr (ShapeDecl [d]) <> ppr at + ppr (TETuple ts _) = parens $ commasep $ map ppr ts + ppr (TERecord fs _) = braces $ commasep $ map ppField fs + where ppField (name, t) = text (nameToString name) <> colon <+> ppr t + ppr (TEVar name _) = ppr name + ppr (TEApply t arg _) = ppr t <+> ppr arg + ppr (TEArrow (Just v) t1 t2 _) = parens v' <+> text "->" <+> ppr t2 + where v' = pprName v <> colon <+> ppr t1 + ppr (TEArrow Nothing t1 t2 _) = ppr t1 <+> text "->" <+> ppr t2 + +instance (Eq vn, IsName vn) => Pretty (TypeArgExp vn) where + ppr (TypeArgExpDim d _) = ppr $ ShapeDecl [d] + ppr (TypeArgExpType d) = ppr d + +instance (Eq vn, IsName vn, Annot f) => Pretty (TypeDeclBase f vn) where + ppr x = pprAnnot (declaredType x) (expandedType x) + +instance IsName vn => Pretty (QualName vn) where + ppr (QualName names name) = + mconcat $ punctuate (text ".") $ map pprName names ++ [pprName name] + +instance IsName vn => Pretty (IdentBase f vn) where + ppr = pprName . identName + +hasArrayLit :: ExpBase ty vn -> Bool +hasArrayLit ArrayLit{} = True +hasArrayLit (TupLit es2 _) = any hasArrayLit es2 +hasArrayLit _ = False + +instance (Eq vn, IsName vn, Annot f) => Pretty (DimIndexBase f vn) where + ppr (DimFix e) = ppr e + ppr (DimSlice i j (Just s)) = + maybe mempty ppr i <> text ":" <> + maybe mempty ppr j <> text ":" <> + ppr s + ppr (DimSlice i (Just j) s) = + maybe mempty ppr i <> text ":" <> + ppr j <> + maybe mempty ((text ":" <>) . ppr) s + ppr (DimSlice i Nothing Nothing) = + maybe mempty ppr i <> text ":" + +instance (Eq vn, IsName vn, Annot f) => Pretty (ExpBase f vn) where + ppr = pprPrec (-1) + pprPrec _ (Var name _ _) = ppr name + pprPrec _ (Parens e _) = align $ parens $ ppr e + pprPrec _ (QualParens v e _) = ppr v <> text "." <> align (parens $ ppr e) + pprPrec _ (Ascript e t _) = pprPrec 0 e <> colon <+> pprPrec 0 t + pprPrec _ (Literal v _) = ppr v + pprPrec _ (IntLit v _ _) = ppr v + pprPrec _ (FloatLit v _ _) = ppr v + pprPrec _ (TupLit es _) + | any hasArrayLit es = parens $ commastack $ map ppr es + | otherwise = parens $ commasep $ map ppr es + pprPrec _ (RecordLit fs _) + | any fieldArray fs = braces $ commastack $ map ppr fs + | otherwise = braces $ commasep $ map ppr fs + where fieldArray (RecordFieldExplicit _ e _) = hasArrayLit e + fieldArray RecordFieldImplicit{} = False + pprPrec _ (ArrayLit es _ _) = + brackets $ commasep $ map ppr es + pprPrec p (Range start maybe_step end _ _) = + parensIf (p /= -1) $ ppr start <> + maybe mempty ((text ".." <>) . ppr) maybe_step <> + case end of + DownToExclusive end' -> text "..>" <> ppr end' + ToInclusive end' -> text "..." <> ppr end' + UpToExclusive end' -> text "..<" <> ppr end' + pprPrec p (BinOp bop _ (x,_) (y,_) _ _) = prettyBinOp p bop x y + pprPrec _ (Project k e _ _) = ppr e <> text "." <> ppr k + pprPrec _ (If c t f _ _) = text "if" <+> ppr c + text "then" <+> align (ppr t) + text "else" <+> align (ppr f) + pprPrec p (Apply f arg _ _ _) = + parensIf (p >= 10) $ ppr f <+> pprPrec 10 arg + pprPrec _ (Negate e _) = text "-" <> ppr e + pprPrec p (LetPat tparams pat e body _) = + parensIf (p /= -1) $ align $ + text "let" <+> align (spread $ map ppr tparams ++ [ppr pat]) <+> + (if linebreak + then equals indent 2 (ppr e) + else equals <+> align (ppr e)) + (case body of LetPat{} -> ppr body + _ -> text "in" <+> ppr body) + where linebreak = case e of + Map{} -> True + Reduce{} -> True + GenReduce{} -> True + Filter{} -> True + Scan{} -> True + DoLoop{} -> True + LetPat{} -> True + LetWith{} -> True + If{} -> True + ArrayLit{} -> False + _ -> hasArrayLit e + pprPrec _ (LetFun fname (tparams, params, retdecl, rettype, e) body _) = + text "let" <+> pprName fname <+> spread (map ppr tparams ++ map ppr params) <> + retdecl' <+> equals indent 2 (ppr e) <+> text "in" + ppr body + where retdecl' = case (ppr <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of + Just rettype' -> text ":" <+> rettype' + Nothing -> mempty + pprPrec _ (LetWith dest src idxs ve body _) + | dest == src = + text "let" <+> ppr dest <> list (map ppr idxs) <+> + equals <+> align (ppr ve) <+> + text "in" ppr body + | otherwise = + text "let" <+> ppr dest <+> equals <+> ppr src <+> + text "with" <+> brackets (commasep (map ppr idxs)) <+> + text "<-" <+> align (ppr ve) <+> + text "in" ppr body + pprPrec _ (Update src idxs ve _) = + ppr src <+> text "with" <+> + brackets (commasep (map ppr idxs)) <+> + text "<-" <+> align (ppr ve) + pprPrec _ (RecordUpdate src fs ve _ _) = + ppr src <+> text "with" <+> + mconcat (intersperse (text ".") (map ppr fs)) <+> + text "<-" <+> align (ppr ve) + pprPrec _ (Index e idxs _ _) = + pprPrec 9 e <> brackets (commasep (map ppr idxs)) + pprPrec _ (Map lam a _ _) = ppSOAC "map" [lam] [a] + pprPrec _ (Reduce Commutative lam e a _) = ppSOAC "reduce_comm" [lam] [e, a] + pprPrec _ (Reduce Noncommutative lam e a _) = ppSOAC "reduce" [lam] [e, a] + pprPrec _ (GenReduce hist op ne bfun img _) = + ppSOAC "gen_reduce" [op, bfun] [hist, ne, img] -- do this manually + pprPrec _ (Stream form lam arr _) = + case form of + MapLike o -> + let ord_str = if o == Disorder then "_per" else "" + in text ("stream_map"++ord_str) <> + ppr lam pprPrec 10 arr + RedLike o comm lam0 -> + let ord_str = if o == Disorder then "_per" else "" + comm_str = case comm of Commutative -> "_comm" + Noncommutative -> "" + in text ("stream_red"++ord_str++comm_str) <> + ppr lam0 ppr lam pprPrec 10 arr + pprPrec _ (Scan lam e a _) = ppSOAC "scan" [lam] [e, a] + pprPrec _ (Filter lam a _) = ppSOAC "filter" [lam] [a] + pprPrec _ (Partition k lam a _) = text "partition" <+> ppr k <+> spread (map (pprPrec 10) [lam, a]) + pprPrec _ (Zip 0 e es _ _) = text "zip" <+> spread (map (pprPrec 10) (e:es)) + pprPrec _ (Zip i e es _ _) = text "zip@" <> ppr i <+> spread (map (pprPrec 10) (e:es)) + pprPrec _ (Unzip e _ _) = text "unzip" <+> pprPrec (-1) e + pprPrec _ (Unsafe e _) = text "unsafe" <+> pprPrec (-1) e + pprPrec _ (Assert e1 e2 _ _) = text "assert" <+> pprPrec 10 e1 <+> pprPrec 10 e2 + pprPrec p (Lambda tparams params body ascript _ _) = + parensIf (p /= -1) $ + text "\\" <> spread (map ppr tparams ++ map ppr params) <> + ppAscription ascript <+> + text "->" indent 2 (ppr body) + pprPrec _ (OpSection binop _ _) = + parens $ ppr binop + pprPrec _ (OpSectionLeft binop _ x _ _ _) = + parens $ ppr x <+> ppr binop + pprPrec _ (OpSectionRight binop _ x _ _ _) = + parens $ ppr binop <+> ppr x + pprPrec _ (ProjectSection fields _ _) = + parens $ mconcat $ map p fields + where p name = text "." <> ppr name + pprPrec _ (IndexSection idxs _ _) = + parens $ text "." <> brackets (commasep (map ppr idxs)) + pprPrec _ (DoLoop tparams pat initexp form loopbody _) = + text "loop" <+> spread (map ppr tparams ++ [ppr pat]) <+> + equals <+> ppr initexp <+> ppr form <+> text "do" + indent 2 (ppr loopbody) + +instance (Eq vn, IsName vn, Annot f) => Pretty (FieldBase f vn) where + ppr (RecordFieldExplicit name e _) = ppr name <> equals <> ppr e + ppr (RecordFieldImplicit name _ _) = pprName name + +instance (Eq vn, IsName vn, Annot f) => Pretty (LoopFormBase f vn) where + ppr (For i ubound) = + text "for" <+> ppr i <+> text "<" <+> align (ppr ubound) + ppr (ForIn x e) = + text "for" <+> ppr x <+> text "in" <+> ppr e + ppr (While cond) = + text "while" <+> ppr cond + +instance (Eq vn, IsName vn, Annot f) => Pretty (PatternBase f vn) where + ppr (PatternAscription p t _) = ppr p <> text ":" <+> ppr t + ppr (PatternParens p _) = parens $ ppr p + ppr (Id v t _) = case unAnnot t of + Just t' -> parens $ pprName v <> colon <+> ppr t' + Nothing -> pprName v + ppr (TuplePattern pats _) = parens $ commasep $ map ppr pats + ppr (RecordPattern fs _) = braces $ commasep $ map ppField fs + where ppField (name, t) = text (nameToString name) <> equals <> ppr t + ppr (Wildcard t _) = case unAnnot t of + Just t' -> parens $ text "_" <> colon <+> ppr t' + Nothing -> text "_" + +ppAscription :: (Eq vn, IsName vn, Annot f) => Maybe (TypeDeclBase f vn) -> Doc +ppAscription Nothing = mempty +ppAscription (Just t) = text ":" <> ppr t + +instance (Eq vn, IsName vn, Annot f) => Pretty (ProgBase f vn) where + ppr = stack . punctuate line . map ppr . progDecs + +instance (Eq vn, IsName vn, Annot f) => Pretty (DecBase f vn) where + ppr (ValDec dec) = ppr dec + ppr (TypeDec dec) = ppr dec + ppr (SigDec sig) = ppr sig + ppr (ModDec sd) = ppr sd + ppr (OpenDec x _ _) = text "open" <+> ppr x + ppr (LocalDec dec _) = text "local" <+> ppr dec + +instance (Eq vn, IsName vn, Annot f) => Pretty (ModExpBase f vn) where + ppr (ModVar v _) = ppr v + ppr (ModParens e _) = parens $ ppr e + ppr (ModImport v _ _) = text "import" <+> ppr (show v) + ppr (ModDecs ds _) = nestedBlock "{" "}" (stack $ punctuate line $ map ppr ds) + ppr (ModApply f a _ _ _) = parens $ ppr f <+> parens (ppr a) + ppr (ModAscript me se _ _) = ppr me <> colon <+> ppr se + ppr (ModLambda param maybe_sig body _) = + text "\\" <> ppr param <> maybe_sig' <+> + text "->" indent 2 (ppr body) + where maybe_sig' = case maybe_sig of Nothing -> mempty + Just (sig, _) -> colon <+> ppr sig + +instance (Eq vn, IsName vn, Annot f) => Pretty (TypeBindBase f vn) where + ppr (TypeBind name params usertype _ _) = + text "type" <+> pprName name <+> spread (map ppr params) <+> equals <+> ppr usertype + +instance (Eq vn, IsName vn) => Pretty (TypeParamBase vn) where + ppr (TypeParamDim name _) = brackets $ pprName name + ppr (TypeParamType Unlifted name _) = text "'" <> pprName name + ppr (TypeParamType Lifted name _) = text "'^" <> pprName name + +instance (Eq vn, IsName vn, Annot f) => Pretty (ValBindBase f vn) where + ppr (ValBind entry name retdecl rettype tparams args body _ _) = + text fun <+> pprName name <+> + spread (map ppr tparams ++ map ppr args) <> retdecl' <> text " =" + indent 2 (ppr body) + where fun | entry = "entry" + | otherwise = "let" + retdecl' = case (ppr <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of + Just rettype' -> text ":" <+> rettype' + Nothing -> mempty + +instance (Eq vn, IsName vn, Annot f) => Pretty (SpecBase f vn) where + ppr (TypeAbbrSpec tpsig) = ppr tpsig + ppr (TypeSpec Unlifted name ps _ _) = text "type" <+> pprName name <+> spread (map ppr ps) + ppr (TypeSpec Lifted name ps _ _) = text "type^" <+> pprName name <+> spread (map ppr ps) + ppr (ValSpec name tparams vtype _ _) = + text "val" <+> pprName name <+> spread (map ppr tparams) <> colon <+> ppr vtype + ppr (ModSpec name sig _ _) = + text "module" <+> pprName name <> colon <+> ppr sig + ppr (IncludeSpec e _) = + text "include" <+> ppr e + +instance (Eq vn, IsName vn, Annot f) => Pretty (SigExpBase f vn) where + ppr (SigVar v _) = ppr v + ppr (SigParens e _) = parens $ ppr e + ppr (SigSpecs ss _) = nestedBlock "{" "}" (stack $ punctuate line $ map ppr ss) + ppr (SigWith s (TypeRef v ps td _) _) = + ppr s <+> text "with" <+> ppr v <+> spread (map ppr ps) <> text " =" <+> ppr td + ppr (SigArrow (Just v) e1 e2 _) = + parens (pprName v <> colon <+> ppr e1) <+> text "->" <+> ppr e2 + ppr (SigArrow Nothing e1 e2 _) = + ppr e1 <+> text "->" <+> ppr e2 + +instance (Eq vn, IsName vn, Annot f) => Pretty (SigBindBase f vn) where + ppr (SigBind name e _ _) = + text "module type" <+> pprName name <+> equals <+> ppr e + +instance (Eq vn, IsName vn, Annot f) => Pretty (ModParamBase f vn) where + ppr (ModParam pname psig _ _) = + parens (pprName pname <> colon <+> ppr psig) + +instance (Eq vn, IsName vn, Annot f) => Pretty (ModBindBase f vn) where + ppr (ModBind name ps sig e _ _) = + text "module" <+> pprName name <+> spread (map ppr ps) <+> sig' <> text " =" <+> ppr e + where sig' = case sig of Nothing -> mempty + Just (s,_) -> colon <+> ppr s <> text " " + +prettyBinOp :: (Eq vn, IsName vn, Annot f) => + Int -> QualName vn -> ExpBase f vn -> ExpBase f vn -> Doc +prettyBinOp p bop x y = parensIf (p > symPrecedence) $ + pprPrec symPrecedence x <+/> + bop' <+> + pprPrec symRPrecedence y + where bop' = case leading of Backtick -> text "`" <> ppr bop <> text "`" + _ -> ppr bop + leading = leadingOperator $ nameFromString $ pretty $ pprName $ qualLeaf bop + symPrecedence = precedence leading + symRPrecedence = rprecedence leading + precedence PipeRight = -1 + precedence PipeLeft = -1 + precedence LogAnd = 0 + precedence LogOr = 0 + precedence Band = 1 + precedence Bor = 1 + precedence Xor = 1 + precedence Equal = 2 + precedence NotEqual = 2 + precedence Less = 2 + precedence Leq = 2 + precedence Greater = 2 + precedence Geq = 2 + precedence ShiftL = 3 + precedence ShiftR = 3 + precedence Plus = 4 + precedence Minus = 4 + precedence Times = 5 + precedence Divide = 5 + precedence Mod = 5 + precedence Quot = 5 + precedence Rem = 5 + precedence Pow = 6 + precedence Backtick = 9 + rprecedence Minus = 10 + rprecedence Divide = 10 + rprecedence op = precedence op + +ppSOAC :: (Eq vn, IsName vn, Pretty fn, Annot f) => + String -> [fn] -> [ExpBase f vn] -> Doc +ppSOAC name funs es = + text name <+> align (spread (map (parens . ppr) funs) + spread (map (pprPrec 10) es)) diff --git a/src/Language/Futhark/Semantic.hs b/src/Language/Futhark/Semantic.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Semantic.hs @@ -0,0 +1,140 @@ +-- | Definitions of various semantic objects (*not* the Futhark +-- semantics themselves). +module Language.Futhark.Semantic + ( ImportName + , mkInitialImport + , mkImportFrom + , includeToFilePath + , includeToString + + , FileModule(..) + , Imports + + , Namespace(..) + , Env(..) + , TySet + , FunSig(..) + , NameMap + , BoundV(..) + , Mod(..) + , TypeBinding(..) + , MTy(..) + ) +where + +import Data.Semigroup ((<>)) +import Data.Loc +import qualified Data.Map.Strict as M +import qualified Data.Semigroup as Sem +import qualified System.FilePath.Posix as Posix +import qualified System.FilePath as Native + +import Language.Futhark +import Futhark.Util (dropLast, toPOSIX, fromPOSIX) + +-- | Canonical reference to a Futhark code file. Does not include the +-- @.fut@ extension. This is most often a path relative to the +-- current working directory of the compiler. +data ImportName = ImportName Posix.FilePath SrcLoc + deriving (Eq, Ord, Show) + +instance Located ImportName where + locOf (ImportName _ loc) = locOf loc + +-- | Create an import name immediately from a file path specified by +-- the user. +mkInitialImport :: Native.FilePath -> ImportName +mkInitialImport s = ImportName (Posix.normalise $ toPOSIX s) noLoc + +-- | We resolve '..' paths here and assume that no shenanigans are +-- going on with symbolic links. If there is, too bad. Don't do +-- that. +mkImportFrom :: ImportName -> String -> SrcLoc -> ImportName +mkImportFrom (ImportName includer _) includee + | Posix.isAbsolute includee = ImportName includee + | otherwise = ImportName $ Posix.normalise $ Posix.joinPath $ includer' ++ includee' + where (dotdots, includee') = span ("../"==) $ Posix.splitPath includee + includer_parts = init $ Posix.splitPath includer + includer' + | length dotdots > length includer_parts = + replicate (length dotdots - length includer_parts) "../" + | otherwise = + dropLast (length dotdots) includer_parts + +-- | Create a @.fut@ file corresponding to an 'ImportName'. +includeToFilePath :: ImportName -> Native.FilePath +includeToFilePath (ImportName s _) = fromPOSIX $ Posix.normalise s Posix.<.> "fut" + +-- | Produce a human-readable canonicalized string from an +-- 'ImportName'. +includeToString :: ImportName -> String +includeToString (ImportName s _) = Posix.normalise $ Posix.makeRelative "/" s + +-- | The result of type checking some file. Can be passed to further +-- invocations of the type checker. +data FileModule = FileModule { fileAbs :: TySet -- ^ Abstract types. + , fileEnv :: Env + , fileProg :: Prog + } + +-- | A mapping from import names to imports. The ordering is significant. +type Imports = [(String, FileModule)] + +-- | The space inhabited by a name. +data Namespace = Term -- ^ Functions and values. + | Type + | Signature + deriving (Eq, Ord, Show, Enum) + +-- | A mapping of abstract types to their liftedness. +type TySet = M.Map (QualName VName) Liftedness + +-- | Representation of a module, which is either a plain environment, +-- or a parametric module ("functor" in SML). +data Mod = ModEnv Env + | ModFun FunSig + deriving (Show) + +-- | A parametric functor consists of a set of abstract types, the +-- environment of its parameter, and the resulting module type. +data FunSig = FunSig { funSigAbs :: TySet + , funSigMod :: Mod + , funSigMty :: MTy + } + deriving (Show) + +-- | Representation of a module type. +data MTy = MTy { mtyAbs :: TySet + -- ^ Abstract types in the module type. + , mtyMod :: Mod + } + deriving (Show) + +-- | A binding from a name to its definition as a type. +data TypeBinding = TypeAbbr Liftedness [TypeParam] StructType + deriving (Eq, Show) + +-- | Type parameters, list of parameter types (optinally named), and +-- return type. The type parameters are in scope in both parameter +-- types and the return type. Non-functional values have only a +-- return type. +data BoundV = BoundV [TypeParam] StructType + deriving (Show) + +type NameMap = M.Map (Namespace, Name) (QualName VName) + +-- | Modules produces environment with this representation. +data Env = Env { envVtable :: M.Map VName BoundV + , envTypeTable :: M.Map VName TypeBinding + , envSigTable :: M.Map VName MTy + , envModTable :: M.Map VName Mod + , envNameMap :: NameMap + } deriving (Show) + +instance Sem.Semigroup Env where + Env vt1 tt1 st1 mt1 nt1 <> Env vt2 tt2 st2 mt2 nt2 = + Env (vt1<>vt2) (tt1<>tt2) (st1<>st2) (mt1<>mt2) (nt1<>nt2) + +instance Monoid Env where + mempty = Env mempty mempty mempty mempty mempty + mappend = (Sem.<>) diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Syntax.hs @@ -0,0 +1,1019 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE StandaloneDeriving #-} +-- | This is an ever-changing syntax representation for Futhark. Some +-- types, such as @Exp@, are parametrised by type and name +-- representation. See the @https://futhark.readthedocs.org@ for a +-- language reference, or this module may be a little hard to +-- understand. +module Language.Futhark.Syntax + ( + module Language.Futhark.Core + + -- * Types + , Uniqueness(..) + , IntType(..) + , FloatType(..) + , PrimType(..) + , ArrayDim (..) + , DimDecl (..) + , ShapeDecl (..) + , shapeRank + , stripDims + , unifyShapes + , TypeName(..) + , typeNameFromQualName + , qualNameFromTypeName + , TypeBase(..) + , TypeArg(..) + , TypeExp(..) + , TypeArgExp(..) + , RecordArrayElemTypeBase(..) + , ArrayElemTypeBase(..) + , CompType + , PatternType + , StructType + , Diet(..) + , TypeDeclBase (..) + + -- * Values + , IntValue(..) + , FloatValue(..) + , PrimValue(..) + , IsPrimValue(..) + , Value(..) + + -- * Abstract syntax tree + , BinOp (..) + , IdentBase (..) + , Inclusiveness(..) + , DimIndexBase(..) + , ExpBase(..) + , FieldBase(..) + , LoopFormBase (..) + , PatternBase(..) + , StreamForm(..) + + -- * Module language + , SpecBase(..) + , SigExpBase(..) + , TypeRefBase(..) + , SigBindBase(..) + , ModExpBase(..) + , ModBindBase(..) + , ModParamBase(..) + + -- * Definitions + , DocComment(..) + , ValBindBase(..) + , Liftedness(..) + , TypeBindBase(..) + , TypeParamBase(..) + , typeParamName + , ProgBase(..) + , DecBase(..) + + -- * Miscellaneous + , NoInfo(..) + , Info(..) + , Names + , QualName(..) + ) + where + +import Control.Applicative +import Control.Monad +import Data.Array +import Data.Bifoldable +import Data.Bifunctor +import Data.Bitraversable +import Data.Foldable +import Data.Loc +import qualified Data.Map.Strict as M +import Data.Monoid +import Data.Ord +import qualified Data.Set as S +import Data.Traversable +import qualified Data.Semigroup as Sem +import Prelude + +import Futhark.Representation.Primitive (FloatType (..), + FloatValue (..), + IntType (..), IntValue (..)) +import Futhark.Util.Pretty +import Language.Futhark.Core + +-- | Convenience class for deriving 'Show' instances for the AST. +class (Show vn, + Show (f VName), + Show (f Diet), + Show (f String), + Show (f [VName]), + Show (f PatternType), + Show (f CompType), + Show (f (TypeBase () ())), + Show (f Int), + Show (f [TypeBase () ()]), + Show (f StructType), + Show (f (Names, StructType)), + Show (f ([TypeBase () ()], PatternType)), + Show (f (M.Map VName VName)), + Show (f [RecordArrayElemTypeBase () Names]), + Show (f Uniqueness), + Show (f ([CompType], CompType))) => Showable f vn where + +-- | No information functor. Usually used for placeholder type- or +-- aliasing information. +data NoInfo a = NoInfo + deriving (Eq, Ord, Show) + +instance Show vn => Showable NoInfo vn where +instance Functor NoInfo where + fmap _ NoInfo = NoInfo +instance Foldable NoInfo where + foldr _ b NoInfo = b +instance Traversable NoInfo where + traverse _ NoInfo = pure NoInfo + +-- | Some information. The dual to 'NoInfo' +newtype Info a = Info { unInfo :: a } + deriving (Eq, Ord, Show) + +instance Show vn => Showable Info vn where +instance Functor Info where + fmap f (Info x) = Info $ f x +instance Foldable Info where + foldr f b (Info x) = f x b +instance Traversable Info where + traverse f (Info x) = Info <$> f x + +-- | Low-level primitive types. +data PrimType = Signed IntType + | Unsigned IntType + | FloatType FloatType + | Bool + deriving (Eq, Ord, Show) + +-- | Non-array values. +data PrimValue = SignedValue !IntValue + | UnsignedValue !IntValue + | FloatValue !FloatValue + | BoolValue !Bool + deriving (Eq, Ord, Show) + +class IsPrimValue v where + primValue :: v -> PrimValue + +instance IsPrimValue Int where + primValue = SignedValue . Int32Value . fromIntegral + +instance IsPrimValue Int8 where + primValue = SignedValue . Int8Value +instance IsPrimValue Int16 where + primValue = SignedValue . Int16Value +instance IsPrimValue Int32 where + primValue = SignedValue . Int32Value +instance IsPrimValue Int64 where + primValue = SignedValue . Int64Value + +instance IsPrimValue Word8 where + primValue = UnsignedValue . Int8Value . fromIntegral +instance IsPrimValue Word16 where + primValue = UnsignedValue . Int16Value . fromIntegral +instance IsPrimValue Word32 where + primValue = UnsignedValue . Int32Value . fromIntegral +instance IsPrimValue Word64 where + primValue = UnsignedValue . Int64Value . fromIntegral + +instance IsPrimValue Float where + primValue = FloatValue . Float32Value + +instance IsPrimValue Double where + primValue = FloatValue . Float64Value + +instance IsPrimValue Bool where + primValue = BoolValue + +class (Eq dim, Ord dim) => ArrayDim dim where + -- | @unifyDims x y@ combines @x@ and @y@ to contain their maximum + -- common information, and fails if they conflict. + unifyDims :: dim -> dim -> Maybe dim + +instance ArrayDim () where + unifyDims () () = Just () + +-- | Declaration of a dimension size. +data DimDecl vn = NamedDim (QualName vn) + -- ^ The size of the dimension is this name, which + -- must be in scope. In a return type, this will + -- give rise to an assertion. + | ConstDim Int + -- ^ The size is a constant. + | AnyDim + -- ^ No dimension declaration. + deriving (Eq, Ord, Show) + +instance Functor DimDecl where + fmap = fmapDefault + +instance Foldable DimDecl where + foldMap = foldMapDefault + +instance Traversable DimDecl where + traverse f (NamedDim qn) = NamedDim <$> traverse f qn + traverse _ (ConstDim x) = pure $ ConstDim x + traverse _ AnyDim = pure AnyDim + +instance (Eq vn, Ord vn) => ArrayDim (DimDecl vn) where + unifyDims AnyDim y = Just y + unifyDims x AnyDim = Just x + unifyDims (NamedDim x) (NamedDim y) | x == y = Just $ NamedDim x + unifyDims (ConstDim x) (ConstDim y) | x == y = Just $ ConstDim x + unifyDims _ _ = Nothing + +-- | The size of an array type is a list of its dimension sizes. If +-- 'Nothing', that dimension is of a (statically) unknown size. +newtype ShapeDecl dim = ShapeDecl { shapeDims :: [dim] } + deriving (Eq, Ord, Show) + +instance Foldable ShapeDecl where + foldr f x (ShapeDecl ds) = foldr f x ds + +instance Traversable ShapeDecl where + traverse f (ShapeDecl ds) = ShapeDecl <$> traverse f ds + +instance Functor ShapeDecl where + fmap f (ShapeDecl ds) = ShapeDecl $ map f ds + +instance Sem.Semigroup (ShapeDecl dim) where + ShapeDecl l1 <> ShapeDecl l2 = ShapeDecl $ l1 ++ l2 + +instance Monoid (ShapeDecl dim) where + mempty = ShapeDecl [] + mappend = (Sem.<>) + +-- | The number of dimensions contained in a shape. +shapeRank :: ShapeDecl dim -> Int +shapeRank = length . shapeDims + +-- | @stripDims n shape@ strips the outer @n@ dimensions from +-- @shape@, returning 'Nothing' if this would result in zero or +-- fewer dimensions. +stripDims :: Int -> ShapeDecl dim -> Maybe (ShapeDecl dim) +stripDims i (ShapeDecl l) + | i < length l = Just $ ShapeDecl $ drop i l + | otherwise = Nothing + + +-- | @unifyShapes x y@ combines @x@ and @y@ to contain their maximum +-- common information, and fails if they conflict. +unifyShapes :: ArrayDim dim => ShapeDecl dim -> ShapeDecl dim -> Maybe (ShapeDecl dim) +unifyShapes (ShapeDecl xs) (ShapeDecl ys) = do + guard $ length xs == length ys + ShapeDecl <$> zipWithM unifyDims xs ys + +-- | A type name consists of qualifiers (for error messages) and a +-- 'VName' (for equality checking). +data TypeName = TypeName { typeQuals :: [VName], typeLeaf :: VName } + deriving (Show) + +instance Eq TypeName where + TypeName _ x == TypeName _ y = x == y + +instance Ord TypeName where + TypeName _ x `compare` TypeName _ y = x `compare` y + +typeNameFromQualName :: QualName VName -> TypeName +typeNameFromQualName (QualName qs x) = TypeName qs x + +qualNameFromTypeName :: TypeName -> QualName VName +qualNameFromTypeName (TypeName qs x) = QualName qs x + +-- | Types that can be elements of tuple-arrays. +data RecordArrayElemTypeBase dim as = + RecordArrayElem (ArrayElemTypeBase dim as) + | RecordArrayArrayElem (ArrayElemTypeBase dim as) (ShapeDecl dim) Uniqueness + deriving (Eq, Show) + +instance Bitraversable RecordArrayElemTypeBase where + bitraverse f g (RecordArrayElem t) = RecordArrayElem <$> bitraverse f g t + bitraverse f g (RecordArrayArrayElem a shape u) = + RecordArrayArrayElem <$> bitraverse f g a <*> traverse f shape <*> pure u + +instance Bifunctor RecordArrayElemTypeBase where + bimap = bimapDefault + +instance Bifoldable RecordArrayElemTypeBase where + bifoldMap = bifoldMapDefault + +data ArrayElemTypeBase dim as = + ArrayPrimElem PrimType as + | ArrayPolyElem TypeName [TypeArg dim as] as + | ArrayRecordElem (M.Map Name (RecordArrayElemTypeBase dim as)) + deriving (Eq, Show) + +instance Bitraversable ArrayElemTypeBase where + bitraverse _ g (ArrayPrimElem t as) = + ArrayPrimElem t <$> g as + bitraverse f g (ArrayPolyElem t args as) = + ArrayPolyElem t <$> traverse (bitraverse f g) args <*> g as + bitraverse f g (ArrayRecordElem fs) = + ArrayRecordElem <$> traverse (bitraverse f g) fs + +instance Bifunctor ArrayElemTypeBase where + bimap = bimapDefault + +instance Bifoldable ArrayElemTypeBase where + bifoldMap = bifoldMapDefault + +-- | An expanded Futhark type is either an array, a prim type, a +-- tuple, or a type variable. When comparing types for equality with +-- '==', aliases are ignored, but dimensions much match. Function +-- parameter names are ignored. +data TypeBase dim as = Prim PrimType + | Array (ArrayElemTypeBase dim as) (ShapeDecl dim) Uniqueness + | Record (M.Map Name (TypeBase dim as)) + | TypeVar as Uniqueness TypeName [TypeArg dim as] + | Arrow as (Maybe VName) (TypeBase dim as) (TypeBase dim as) + -- ^ The aliasing corresponds to the lexical + -- closure of the function. + deriving (Show) + +instance (Eq dim, Eq as) => Eq (TypeBase dim as) where + Prim x1 == Prim y1 = x1 == y1 + Array x1 y1 z1 == Array x2 y2 z2 = x1 == x2 && y1 == y2 && z1 == z2 + Record x1 == Record x2 = x1 == x2 + TypeVar _ u1 x1 y1 == TypeVar _ u2 x2 y2 = u1 == u2 && x1 == x2 && y1 == y2 + Arrow _ _ x1 y1 == Arrow _ _ x2 y2 = x1 == x2 && y1 == y2 + _ == _ = False + +instance Bitraversable TypeBase where + bitraverse _ _ (Prim t) = pure $ Prim t + bitraverse f g (Array a shape u) = + Array <$> bitraverse f g a <*> traverse f shape <*> pure u + bitraverse f g (Record fs) = Record <$> traverse (bitraverse f g) fs + bitraverse f g (TypeVar als u t args) = + TypeVar <$> g als <*> pure u <*> pure t <*> traverse (bitraverse f g) args + bitraverse f g (Arrow als v t1 t2) = + Arrow <$> g als <*> pure v <*> bitraverse f g t1 <*> bitraverse f g t2 + +instance Bifunctor TypeBase where + bimap = bimapDefault + +instance Bifoldable TypeBase where + bifoldMap = bifoldMapDefault + +data TypeArg dim as = TypeArgDim dim SrcLoc + | TypeArgType (TypeBase dim as) SrcLoc + deriving (Eq, Show) + +instance Bitraversable TypeArg where + bitraverse f _ (TypeArgDim v loc) = TypeArgDim <$> f v <*> pure loc + bitraverse f g (TypeArgType t loc) = TypeArgType <$> bitraverse f g t <*> pure loc + +instance Bifunctor TypeArg where + bimap = bimapDefault + +instance Bifoldable TypeArg where + bifoldMap = bifoldMapDefault + +-- | A type with aliasing information and no shape annotations, used +-- for describing the type of a computation. +type CompType = TypeBase () Names + +-- | A type with aliasing information and shape annotations, used for +-- describing the type of a pattern. +type PatternType = TypeBase (DimDecl VName) Names + +-- | An unstructured type with type variables and possibly shape +-- declarations - this is what the user types in the source program. +data TypeExp vn = TEVar (QualName vn) SrcLoc + | TETuple [TypeExp vn] SrcLoc + | TERecord [(Name, TypeExp vn)] SrcLoc + | TEArray (TypeExp vn) (DimDecl vn) SrcLoc + | TEUnique (TypeExp vn) SrcLoc + | TEApply (TypeExp vn) (TypeArgExp vn) SrcLoc + | TEArrow (Maybe vn) (TypeExp vn) (TypeExp vn) SrcLoc + deriving (Eq, Show) + +instance Located (TypeExp vn) where + locOf (TEArray _ _ loc) = locOf loc + locOf (TETuple _ loc) = locOf loc + locOf (TERecord _ loc) = locOf loc + locOf (TEVar _ loc) = locOf loc + locOf (TEUnique _ loc) = locOf loc + locOf (TEApply _ _ loc) = locOf loc + locOf (TEArrow _ _ _ loc) = locOf loc + +data TypeArgExp vn = TypeArgExpDim (DimDecl vn) SrcLoc + | TypeArgExpType (TypeExp vn) + deriving (Eq, Show) + +instance Located (TypeArgExp vn) where + locOf (TypeArgExpDim _ loc) = locOf loc + locOf (TypeArgExpType t) = locOf t + +-- | A "structural" type with shape annotations and no aliasing +-- information, used for declarations. +type StructType = TypeBase (DimDecl VName) () + +-- | A declaration of the type of something. +data TypeDeclBase f vn = + TypeDecl { declaredType :: TypeExp vn + -- ^ The type declared by the user. + , expandedType :: f StructType + -- ^ The type deduced by the type checker. + } +deriving instance Showable f vn => Show (TypeDeclBase f vn) + +instance Located (TypeDeclBase f vn) where + locOf = locOf . declaredType + +-- | Information about which parts of a value/type are consumed. +data Diet = RecordDiet (M.Map Name Diet) -- ^ Consumes these fields in the record. + | FuncDiet Diet Diet + -- ^ A function that consumes its argument(s) like this. + -- The final 'Diet' should always be 'Observe', as there + -- is no way for a function to consume its return value. + | Consume -- ^ Consumes this value. + | Observe -- ^ Only observes value in this position, does + -- not consume. + deriving (Eq, Show) + +-- | Simple Futhark values. Values are fully evaluated and their type +-- is always unambiguous. +data Value = PrimValue !PrimValue + | ArrayValue !(Array Int Value) (TypeBase () ()) + -- ^ It is assumed that the array is 0-indexed. The type + -- is the full type. + deriving (Eq, Show) + +-- | An identifier consists of its name and the type of the value +-- bound to the identifier. +data IdentBase f vn = Ident { identName :: vn + , identType :: f CompType + , identSrcLoc :: SrcLoc + } +deriving instance Showable f vn => Show (IdentBase f vn) + +instance Eq vn => Eq (IdentBase ty vn) where + x == y = identName x == identName y + +instance Ord vn => Ord (IdentBase ty vn) where + compare = comparing identName + +instance Located (IdentBase ty vn) where + locOf = locOf . identSrcLoc + +-- | Default binary operators. +data BinOp = Backtick + -- ^ A pseudo-operator standing in for any normal + -- identifier used as an operator (they all have the + -- same fixity). + -- Binary Ops for Numbers + | Plus + | Minus + | Pow + | Times + | Divide + | Mod + | Quot + | Rem + | ShiftR + | ShiftL + | Band + | Xor + | Bor + | LogAnd + | LogOr + -- Relational Ops for all primitive types at least + | Equal + | NotEqual + | Less + | Leq + | Greater + | Geq + -- Some functional ops. + | PipeRight -- ^ @|>@ + | PipeLeft -- ^ @<|@ + -- Misc + deriving (Eq, Ord, Show, Enum, Bounded) + +-- | Whether a bound for an end-point of a 'DimSlice' or a range +-- literal is inclusive or exclusive. +data Inclusiveness a = DownToExclusive a + | ToInclusive a -- ^ May be "down to" if step is negative. + | UpToExclusive a + deriving (Eq, Ord, Show) + +instance Located a => Located (Inclusiveness a) where + locOf (DownToExclusive x) = locOf x + locOf (ToInclusive x) = locOf x + locOf (UpToExclusive x) = locOf x + +instance Functor Inclusiveness where + fmap = fmapDefault + +instance Foldable Inclusiveness where + foldMap = foldMapDefault + +instance Traversable Inclusiveness where + traverse f (DownToExclusive x) = DownToExclusive <$> f x + traverse f (ToInclusive x) = ToInclusive <$> f x + traverse f (UpToExclusive x) = UpToExclusive <$> f x + +-- | An indexing of a single dimension. +data DimIndexBase f vn = DimFix (ExpBase f vn) + | DimSlice (Maybe (ExpBase f vn)) + (Maybe (ExpBase f vn)) + (Maybe (ExpBase f vn)) +deriving instance Showable f vn => Show (DimIndexBase f vn) + +-- | A name qualified with a breadcrumb of module accesses. +data QualName vn = QualName { qualQuals :: ![vn] + , qualLeaf :: !vn + } + deriving (Eq, Ord, Show) + +instance Functor QualName where + fmap = fmapDefault + +instance Foldable QualName where + foldMap = foldMapDefault + +instance Traversable QualName where + traverse f (QualName qs v) = QualName <$> traverse f qs <*> f v + +-- | The Futhark expression language. +-- +-- In a value of type @Exp f vn@, annotations are wrapped in the +-- functor @f@, and all names are of type @vn@. +-- +-- This allows us to encode whether or not the expression has been +-- type-checked in the Haskell type of the expression. Specifically, +-- the parser will produce expressions of type @Exp 'NoInfo' 'Name'@, +-- and the type checker will convert these to @Exp 'Info' 'VName'@, in +-- which type information is always present and all names are unique. +data ExpBase f vn = + Literal PrimValue SrcLoc + + | IntLit Integer (f (TypeBase () ())) SrcLoc + -- ^ A polymorphic integral literal. + + | FloatLit Double (f (TypeBase () ())) SrcLoc + -- ^ A polymorphic decimal literal. + + | Parens (ExpBase f vn) SrcLoc + -- ^ A parenthesized expression. + + | QualParens (QualName vn) (ExpBase f vn) SrcLoc + + | TupLit [ExpBase f vn] SrcLoc + -- ^ Tuple literals, e.g., @{1+3, {x, y+z}}@. + + | RecordLit [FieldBase f vn] SrcLoc + -- ^ Record literals, e.g. @{x=2,y=3,z}@. + + | ArrayLit [ExpBase f vn] (f CompType) SrcLoc + -- ^ Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@. + -- Second arg is the row type of the rows of the array. + + | Range (ExpBase f vn) (Maybe (ExpBase f vn)) (Inclusiveness (ExpBase f vn)) (f CompType) SrcLoc + + | Var (QualName vn) (f PatternType) SrcLoc + + | Ascript (ExpBase f vn) (TypeDeclBase f vn) SrcLoc + -- ^ Type ascription: @e : t@. + + | LetPat [TypeParamBase vn] (PatternBase f vn) (ExpBase f vn) (ExpBase f vn) SrcLoc + + | LetFun vn ([TypeParamBase vn], [PatternBase f vn], Maybe (TypeExp vn), f StructType, ExpBase f vn) + (ExpBase f vn) SrcLoc + + | If (ExpBase f vn) (ExpBase f vn) (ExpBase f vn) (f CompType) SrcLoc + + | Apply (ExpBase f vn) (ExpBase f vn) (f Diet) (f PatternType) SrcLoc + + | Negate (ExpBase f vn) SrcLoc + -- ^ Numeric negation (ugly special case; Haskell did it first). + + | Lambda [TypeParamBase vn] [PatternBase f vn] (ExpBase f vn) + (Maybe (TypeDeclBase f vn)) (f (Names, StructType)) SrcLoc + + | OpSection (QualName vn) (f PatternType) SrcLoc + -- ^ @+@; first two types are operands, third is result. + | OpSectionLeft (QualName vn) (f PatternType) + (ExpBase f vn) (f StructType, f StructType) (f PatternType) SrcLoc + -- ^ @2+@; first type is operand, second is result. + | OpSectionRight (QualName vn) (f PatternType) + (ExpBase f vn) (f StructType, f StructType) (f PatternType) SrcLoc + -- ^ @+2@; first type is operand, second is result. + | ProjectSection [Name] (f PatternType) SrcLoc + -- ^ Field projection as a section: @(.x.y.z)@. + | IndexSection [DimIndexBase f vn] (f PatternType) SrcLoc + -- ^ Array indexing as a section: @(.[i,j])@. + + | DoLoop + [TypeParamBase vn] + (PatternBase f vn) -- Merge variable pattern + (ExpBase f vn) -- Initial values of merge variables. + (LoopFormBase f vn) -- Do or while loop. + (ExpBase f vn) -- Loop body. + SrcLoc + + | BinOp (QualName vn) (f PatternType) + (ExpBase f vn, f StructType) (ExpBase f vn, f StructType) + (f PatternType) SrcLoc + + | Project Name (ExpBase f vn) (f CompType) SrcLoc + + -- Primitive array operations + | LetWith (IdentBase f vn) (IdentBase f vn) + [DimIndexBase f vn] (ExpBase f vn) + (ExpBase f vn) SrcLoc + + | Index (ExpBase f vn) [DimIndexBase f vn] (f CompType) SrcLoc + + | Update (ExpBase f vn) [DimIndexBase f vn] (ExpBase f vn) SrcLoc + + | RecordUpdate (ExpBase f vn) [Name] (ExpBase f vn) (f PatternType) SrcLoc + + -- Second-Order Array Combinators accept curried and + -- anonymous functions as first params. + | Map (ExpBase f vn) (ExpBase f vn) (f CompType) SrcLoc + -- ^ @map (+1) [1, 2, ..., n] = [2, 3, ..., n+1]@. + + | Reduce Commutativity (ExpBase f vn) (ExpBase f vn) (ExpBase f vn) SrcLoc + -- ^ @reduce (+) 0 ([1,2,...,n]) = (0+1+2+...+n)@. + + | GenReduce (ExpBase f vn) (ExpBase f vn) (ExpBase f vn) + (ExpBase f vn) (ExpBase f vn) SrcLoc + -- ^ @gen_reduce [1,1,1] (+) 0 [1,1,1] [1,1,1] = [4,1,1]@ + + | Scan (ExpBase f vn) (ExpBase f vn) (ExpBase f vn) SrcLoc + -- ^ @scan (+) 0 ([ 1, 2, 3 ]) = [ 1, 3, 6 ]@. + + | Filter (ExpBase f vn) (ExpBase f vn) SrcLoc + -- ^ Return those elements of the array that satisfy the + -- predicate. + + | Partition Int (ExpBase f vn) (ExpBase f vn) SrcLoc + -- ^ @partition k f a@, where @f@ returns an integer, + -- returns a tuple @(a', is)@ that describes a + -- partitioning of @a@ into @n@ equivalence classes. + -- Here, @a'@ is a re-ordering of @a@, and @is@ is an + -- array of @k@ offsets into @a'@. + + | Stream (StreamForm f vn) (ExpBase f vn) (ExpBase f vn) SrcLoc + -- ^ Streaming: intuitively, this gives a size-parameterized + -- composition for SOACs that cannot be fused, e.g., due to scan. + -- For example, assuming @A : [int], f : int->int, g : real->real@, + -- the code: @let x = map(f,A) in let y = scan(op+,0,x) in map(g,y)@ + -- can be re-written (streamed) in the source-Futhark language as: + -- @let (acc, z) =@ + -- @ stream (fn (int,[real]) (real chunk, real acc, [int] a) =>@ + -- @ let x = map (f, A )@ + -- @ let y0= scan(op +, 0, x )@ + -- @ let y = map (op +(acc), y0)@ + -- @ ( acc+y0[chunk-1], map(g, y) )@ + -- @ ) 0 A@ + -- where (i) @chunk@ is a symbolic int denoting the chunk + -- size, (ii) @0@ is the initial value of the accumulator, + -- which allows the streaming of @scan@. + -- Finally, the unnamed function (@fn...@) implements the a fold that: + -- computes the accumulator of @scan@, as defined inside its body, AND + -- implicitly concatenates each of the result arrays across + -- the iteration space. + -- In essence, sequential codegen can choose chunk = 1 and thus + -- eliminate the SOACs on the outermost level, while parallel codegen + -- may choose the maximal chunk size that still satisfies the memory + -- requirements of the device. + + | Zip Int (ExpBase f vn) [ExpBase f vn] (f CompType) SrcLoc + -- ^ Conventional zip taking nonzero arrays as arguments. + -- All arrays must have the exact same length. + + | Unzip (ExpBase f vn) [f CompType] SrcLoc + -- ^ Unzip that can unzip to tuples of arbitrary size. + -- The types are the elements of the tuple. + + | Unsafe (ExpBase f vn) SrcLoc + -- ^ Explore the Danger Zone and elide safety checks on + -- array operations and other assertions during execution + -- of this expression. Make really sure the code is + -- correct. + + | Assert (ExpBase f vn) (ExpBase f vn) (f String) SrcLoc + -- ^ Fail if the first expression does not return true, + -- and return the value of the second expression if it + -- does. + +deriving instance Showable f vn => Show (ExpBase f vn) + +data StreamForm f vn = MapLike StreamOrd + | RedLike StreamOrd Commutativity (ExpBase f vn) +deriving instance Showable f vn => Show (StreamForm f vn) + +instance Located (ExpBase f vn) where + locOf (Literal _ loc) = locOf loc + locOf (IntLit _ _ loc) = locOf loc + locOf (FloatLit _ _ loc) = locOf loc + locOf (Parens _ loc) = locOf loc + locOf (QualParens _ _ loc) = locOf loc + locOf (TupLit _ pos) = locOf pos + locOf (RecordLit _ pos) = locOf pos + locOf (Project _ _ _ pos) = locOf pos + locOf (ArrayLit _ _ pos) = locOf pos + locOf (Range _ _ _ _ pos) = locOf pos + locOf (BinOp _ _ _ _ _ pos) = locOf pos + locOf (If _ _ _ _ pos) = locOf pos + locOf (Var _ _ loc) = locOf loc + locOf (Ascript _ _ loc) = locOf loc + locOf (Negate _ pos) = locOf pos + locOf (Apply _ _ _ _ pos) = locOf pos + locOf (LetPat _ _ _ _ pos) = locOf pos + locOf (LetFun _ _ _ loc) = locOf loc + locOf (LetWith _ _ _ _ _ pos) = locOf pos + locOf (Index _ _ _ loc) = locOf loc + locOf (Update _ _ _ pos) = locOf pos + locOf (RecordUpdate _ _ _ _ pos) = locOf pos + locOf (Map _ _ _ loc) = locOf loc + locOf (Reduce _ _ _ _ pos) = locOf pos + locOf (GenReduce _ _ _ _ _ pos) = locOf pos + locOf (Zip _ _ _ _ loc) = locOf loc + locOf (Unzip _ _ pos) = locOf pos + locOf (Scan _ _ _ pos) = locOf pos + locOf (Filter _ _ pos) = locOf pos + locOf (Partition _ _ _ loc) = locOf loc + locOf (Lambda _ _ _ _ _ loc) = locOf loc + locOf (OpSection _ _ loc) = locOf loc + locOf (OpSectionLeft _ _ _ _ _ loc) = locOf loc + locOf (OpSectionRight _ _ _ _ _ loc) = locOf loc + locOf (ProjectSection _ _ loc) = locOf loc + locOf (IndexSection _ _ loc) = locOf loc + locOf (DoLoop _ _ _ _ _ pos) = locOf pos + locOf (Stream _ _ _ pos) = locOf pos + locOf (Unsafe _ loc) = locOf loc + locOf (Assert _ _ _ loc) = locOf loc + +-- | An entry in a record literal. +data FieldBase f vn = RecordFieldExplicit Name (ExpBase f vn) SrcLoc + | RecordFieldImplicit vn (f CompType) SrcLoc + +deriving instance Showable f vn => Show (FieldBase f vn) + +instance Located (FieldBase f vn) where + locOf (RecordFieldExplicit _ _ loc) = locOf loc + locOf (RecordFieldImplicit _ _ loc) = locOf loc + +-- | Whether the loop is a @for@-loop or a @while@-loop. +data LoopFormBase f vn = For (IdentBase f vn) (ExpBase f vn) + | ForIn (PatternBase f vn) (ExpBase f vn) + | While (ExpBase f vn) +deriving instance Showable f vn => Show (LoopFormBase f vn) + +-- | A pattern as used most places where variables are bound (function +-- parameters, @let@ expressions, etc). +data PatternBase f vn = TuplePattern [PatternBase f vn] SrcLoc + | RecordPattern [(Name, PatternBase f vn)] SrcLoc + | PatternParens (PatternBase f vn) SrcLoc + | Id vn (f PatternType) SrcLoc + | Wildcard (f PatternType) SrcLoc -- Nothing, i.e. underscore. + | PatternAscription (PatternBase f vn) (TypeDeclBase f vn) SrcLoc +deriving instance Showable f vn => Show (PatternBase f vn) + +instance Located (PatternBase f vn) where + locOf (TuplePattern _ loc) = locOf loc + locOf (RecordPattern _ loc) = locOf loc + locOf (PatternParens _ loc) = locOf loc + locOf (Id _ _ loc) = locOf loc + locOf (Wildcard _ loc) = locOf loc + locOf (PatternAscription _ _ loc) = locOf loc + +-- | Documentation strings, including source location. +data DocComment = DocComment String SrcLoc + deriving (Show) + +instance Located DocComment where + locOf (DocComment _ loc) = locOf loc + +-- | Function Declarations +data ValBindBase f vn = ValBind { valBindEntryPoint :: Bool + -- ^ True if this function is an entry point. + , valBindName :: vn + , valBindRetDecl :: Maybe (TypeExp vn) + , valBindRetType :: f StructType + , valBindTypeParams :: [TypeParamBase vn] + , valBindParams :: [PatternBase f vn] + , valBindBody :: ExpBase f vn + , valBindDoc :: Maybe DocComment + , valBindLocation :: SrcLoc + } +deriving instance Showable f vn => Show (ValBindBase f vn) + +instance Located (ValBindBase f vn) where + locOf = locOf . valBindLocation + +-- | Type Declarations +data TypeBindBase f vn = TypeBind { typeAlias :: vn + , typeParams :: [TypeParamBase vn] + , typeExp :: TypeDeclBase f vn + , typeDoc :: Maybe DocComment + , typeBindLocation :: SrcLoc + } +deriving instance Showable f vn => Show (TypeBindBase f vn) + +instance Located (TypeBindBase f vn) where + locOf = locOf . typeBindLocation + +-- | The liftedness of a type parameter. By the @Ord@ instance, +-- @Unlifted@ is less than @Lifted@. +data Liftedness = Unlifted -- ^ May only be instantiated with a zero-order type. + | Lifted -- ^ May be instantiated to a functional type. + deriving (Eq, Ord, Show) + +data TypeParamBase vn = TypeParamDim vn SrcLoc + -- ^ A type parameter that must be a size. + | TypeParamType Liftedness vn SrcLoc + -- ^ A type parameter that must be a type. + deriving (Eq, Show) + +instance Functor TypeParamBase where + fmap = fmapDefault + +instance Foldable TypeParamBase where + foldMap = foldMapDefault + +instance Traversable TypeParamBase where + traverse f (TypeParamDim v loc) = TypeParamDim <$> f v <*> pure loc + traverse f (TypeParamType l v loc) = TypeParamType l <$> f v <*> pure loc + +instance Located (TypeParamBase vn) where + locOf (TypeParamDim _ loc) = locOf loc + locOf (TypeParamType _ _ loc) = locOf loc + +typeParamName :: TypeParamBase vn -> vn +typeParamName (TypeParamDim v _) = v +typeParamName (TypeParamType _ v _) = v + +data SpecBase f vn = ValSpec { specName :: vn + , specTypeParams :: [TypeParamBase vn] + , specType :: TypeDeclBase f vn + , specDoc :: Maybe DocComment + , specLocation :: SrcLoc + } + | TypeAbbrSpec (TypeBindBase f vn) + | TypeSpec Liftedness vn [TypeParamBase vn] (Maybe DocComment) SrcLoc -- ^ Abstract type. + | ModSpec vn (SigExpBase f vn) (Maybe DocComment) SrcLoc + | IncludeSpec (SigExpBase f vn) SrcLoc +deriving instance Showable f vn => Show (SpecBase f vn) + +instance Located (SpecBase f vn) where + locOf (ValSpec _ _ _ _ loc) = locOf loc + locOf (TypeAbbrSpec tbind) = locOf tbind + locOf (TypeSpec _ _ _ _ loc) = locOf loc + locOf (ModSpec _ _ _ loc) = locOf loc + locOf (IncludeSpec _ loc) = locOf loc + +data SigExpBase f vn = SigVar (QualName vn) SrcLoc + | SigParens (SigExpBase f vn) SrcLoc + | SigSpecs [SpecBase f vn] SrcLoc + | SigWith (SigExpBase f vn) (TypeRefBase f vn) SrcLoc + | SigArrow (Maybe vn) (SigExpBase f vn) (SigExpBase f vn) SrcLoc +deriving instance Showable f vn => Show (SigExpBase f vn) + +-- | A type refinement. +data TypeRefBase f vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeDeclBase f vn) SrcLoc +deriving instance Showable f vn => Show (TypeRefBase f vn) + +instance Located (TypeRefBase f vn) where + locOf (TypeRef _ _ _ loc) = locOf loc + +instance Located (SigExpBase f vn) where + locOf (SigVar _ loc) = locOf loc + locOf (SigParens _ loc) = locOf loc + locOf (SigSpecs _ loc) = locOf loc + locOf (SigWith _ _ loc) = locOf loc + locOf (SigArrow _ _ _ loc) = locOf loc + +data SigBindBase f vn = SigBind { sigName :: vn + , sigExp :: SigExpBase f vn + , sigDoc :: Maybe DocComment + , sigLoc :: SrcLoc + } +deriving instance Showable f vn => Show (SigBindBase f vn) + +instance Located (SigBindBase f vn) where + locOf = locOf . sigLoc + +data ModExpBase f vn = ModVar (QualName vn) SrcLoc + | ModParens (ModExpBase f vn) SrcLoc + | ModImport FilePath (f FilePath) SrcLoc + -- ^ The contents of another file as a module. + | ModDecs [DecBase f vn] SrcLoc + | ModApply (ModExpBase f vn) (ModExpBase f vn) (f (M.Map VName VName)) (f (M.Map VName VName)) SrcLoc + -- ^ Functor application. + | ModAscript (ModExpBase f vn) (SigExpBase f vn) (f (M.Map VName VName)) SrcLoc + | ModLambda (ModParamBase f vn) + (Maybe (SigExpBase f vn, f (M.Map VName VName))) + (ModExpBase f vn) + SrcLoc +deriving instance Showable f vn => Show (ModExpBase f vn) + +instance Located (ModExpBase f vn) where + locOf (ModVar _ loc) = locOf loc + locOf (ModParens _ loc) = locOf loc + locOf (ModImport _ _ loc) = locOf loc + locOf (ModDecs _ loc) = locOf loc + locOf (ModApply _ _ _ _ loc) = locOf loc + locOf (ModAscript _ _ _ loc) = locOf loc + locOf (ModLambda _ _ _ loc) = locOf loc + +data ModBindBase f vn = + ModBind { modName :: vn + , modParams :: [ModParamBase f vn] + , modSignature :: Maybe (SigExpBase f vn, f (M.Map VName VName)) + , modExp :: ModExpBase f vn + , modDoc :: Maybe DocComment + , modLocation :: SrcLoc + } +deriving instance Showable f vn => Show (ModBindBase f vn) + +instance Located (ModBindBase f vn) where + locOf = locOf . modLocation + +data ModParamBase f vn = ModParam { modParamName :: vn + , modParamType :: SigExpBase f vn + , modParamAbs :: f [VName] + , modParamLocation :: SrcLoc + } +deriving instance Showable f vn => Show (ModParamBase f vn) + +instance Located (ModParamBase f vn) where + locOf = locOf . modParamLocation + +-- | A top-level binding. +data DecBase f vn = ValDec (ValBindBase f vn) + | TypeDec (TypeBindBase f vn) + | SigDec (SigBindBase f vn) + | ModDec (ModBindBase f vn) + | OpenDec (ModExpBase f vn) (f [VName]) SrcLoc + | LocalDec (DecBase f vn) SrcLoc +deriving instance Showable f vn => Show (DecBase f vn) + +instance Located (DecBase f vn) where + locOf (ValDec d) = locOf d + locOf (TypeDec d) = locOf d + locOf (SigDec d) = locOf d + locOf (ModDec d) = locOf d + locOf (OpenDec _ _ loc) = locOf loc + locOf (LocalDec _ loc) = locOf loc + +-- | The program described by a single Futhark file. May depend on +-- other files. +data ProgBase f vn = Prog { progDoc :: Maybe DocComment + , progDecs :: [DecBase f vn] + } +deriving instance Showable f vn => Show (ProgBase f vn) + +-- | A set of names. +type Names = S.Set VName + +--- Some prettyprinting definitions are here because we need them in +--- the Attributes module. + +instance Pretty PrimType where + ppr (Unsigned Int8) = text "u8" + ppr (Unsigned Int16) = text "u16" + ppr (Unsigned Int32) = text "u32" + ppr (Unsigned Int64) = text "u64" + ppr (Signed t) = ppr t + ppr (FloatType t) = ppr t + ppr Bool = text "bool" + +instance Pretty BinOp where + ppr Backtick = text "``" + ppr Plus = text "+" + ppr Minus = text "-" + ppr Pow = text "**" + ppr Times = text "*" + ppr Divide = text "/" + ppr Mod = text "%" + ppr Quot = text "//" + ppr Rem = text "%%" + ppr ShiftR = text ">>" + ppr ShiftL = text "<<" + ppr Band = text "&" + ppr Xor = text "^" + ppr Bor = text "|" + ppr LogAnd = text "&&" + ppr LogOr = text "||" + ppr Equal = text "==" + ppr NotEqual = text "!=" + ppr Less = text "<" + ppr Leq = text "<=" + ppr Greater = text ">" + ppr Geq = text ">=" + ppr PipeLeft = text "<|" + ppr PipeRight = text "|>" diff --git a/src/Language/Futhark/Traversals.hs b/src/Language/Futhark/Traversals.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Traversals.hs @@ -0,0 +1,315 @@ +{-# LANGUAGE FlexibleInstances #-} +-- | +-- +-- Functions for generic traversals across Futhark syntax trees. The +-- motivation for this module came from dissatisfaction with rewriting +-- the same trivial tree recursions for every module. A possible +-- alternative would be to use normal \"Scrap your +-- boilerplate\"-techniques, but these are rejected for two reasons: +-- +-- * They are too slow. +-- +-- * More importantly, they do not tell you whether you have missed +-- some cases. +-- +-- Instead, this module defines various traversals of the Futhark syntax +-- tree. The implementation is rather tedious, but the interface is +-- easy to use. +-- +-- A traversal of the Futhark syntax tree is expressed as a tuple of +-- functions expressing the operations to be performed on the various +-- types of nodes. +module Language.Futhark.Traversals + ( ASTMapper(..) + , ASTMappable(..) + ) where + +import qualified Data.Set as S + +import Language.Futhark.Syntax + +-- | 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 ASTMapper m = ASTMapper { + mapOnExp :: ExpBase Info VName -> m (ExpBase Info VName) + , mapOnName :: VName -> m VName + , mapOnQualName :: QualName VName -> m (QualName VName) + , mapOnType :: TypeBase () () -> m (TypeBase () ()) + , mapOnCompType :: CompType -> m CompType + , mapOnStructType :: StructType -> m StructType + , mapOnPatternType :: PatternType -> m PatternType + } + +class ASTMappable x where + -- | Map a monadic action across the immediate children of an + -- object. Importantly, the 'astMap' action is not invoked for + -- the object itself, and the mapping does not descend recursively + -- into subexpressions. The mapping is done left-to-right. + astMap :: Monad m => ASTMapper m -> x -> m x + +instance ASTMappable (ExpBase Info VName) where + astMap tv (Var name t loc) = + Var <$> mapOnQualName tv name <*> traverse (mapOnPatternType tv) t <*> + pure loc + astMap _ (Literal val loc) = + pure $ Literal val loc + astMap tv (IntLit val t loc) = + IntLit val <$> traverse (mapOnType tv) t <*> pure loc + astMap tv (FloatLit val t loc) = + FloatLit val <$> traverse (mapOnType tv) t <*> pure loc + astMap tv (Parens e loc) = + Parens <$> mapOnExp tv e <*> pure loc + astMap tv (QualParens name e loc) = + QualParens <$> mapOnQualName tv name <*> mapOnExp tv e <*> pure loc + astMap tv (TupLit els loc) = + TupLit <$> mapM (mapOnExp tv) els <*> pure loc + astMap tv (RecordLit fields loc) = + RecordLit <$> astMap tv fields <*> pure loc + astMap tv (ArrayLit els t loc) = + ArrayLit <$> mapM (mapOnExp tv) els <*> traverse (mapOnCompType tv) t <*> pure loc + astMap tv (Range start next end t loc) = + Range <$> mapOnExp tv start <*> traverse (mapOnExp tv) next <*> + traverse (mapOnExp tv) end <*> traverse (mapOnCompType tv) t <*> pure loc + astMap tv (Ascript e tdecl loc) = + Ascript <$> mapOnExp tv e <*> astMap tv tdecl <*> pure loc + astMap tv (BinOp fname t (x,xt) (y,yt) (Info rt) loc) = + BinOp <$> mapOnQualName tv fname <*> traverse (mapOnPatternType tv) t <*> + ((,) <$> mapOnExp tv x <*> traverse (mapOnStructType tv) xt) <*> + ((,) <$> mapOnExp tv y <*> traverse (mapOnStructType tv) yt) <*> + (Info <$> mapOnPatternType tv rt) <*> pure loc + astMap tv (Negate x loc) = + Negate <$> mapOnExp tv x <*> pure loc + astMap tv (If c texp fexp t loc) = + If <$> mapOnExp tv c <*> mapOnExp tv texp <*> mapOnExp tv fexp <*> + traverse (mapOnCompType tv) t <*> pure loc + astMap tv (Apply f arg d (Info t) loc) = + Apply <$> mapOnExp tv f <*> mapOnExp tv arg <*> + pure d <*> (Info <$> mapOnPatternType tv t) <*> + pure loc + astMap tv (LetPat tparams pat e body loc) = + LetPat <$> mapM (astMap tv) tparams <*> + astMap tv pat <*> mapOnExp tv e <*> + mapOnExp tv body <*> pure loc + astMap tv (LetFun name (fparams, params, ret, t, e) body loc) = + LetFun <$> mapOnName tv name <*> + ((,,,,) <$> mapM (astMap tv) fparams <*> mapM (astMap tv) params <*> + traverse (astMap tv) ret <*> traverse (mapOnStructType tv) t <*> + mapOnExp tv e) <*> + mapOnExp tv body <*> pure loc + astMap tv (LetWith dest src idxexps vexp body loc) = + pure LetWith <*> + astMap tv dest <*> astMap tv src <*> + mapM (astMap tv) idxexps <*> mapOnExp tv vexp <*> + mapOnExp tv body <*> pure loc + astMap tv (Update src slice v loc) = + Update <$> mapOnExp tv src <*> mapM (astMap tv) slice <*> + mapOnExp tv v <*> pure loc + astMap tv (RecordUpdate src fs v (Info t) loc) = + RecordUpdate <$> mapOnExp tv src <*> pure fs <*> + mapOnExp tv v <*> (Info <$> mapOnPatternType tv t) <*> pure loc + astMap tv (Project field e t loc) = + Project field <$> mapOnExp tv e <*> traverse (mapOnCompType tv) t <*> pure loc + astMap tv (Index arr idxexps t loc) = + pure Index <*> + astMap tv arr <*> + mapM (astMap tv) idxexps <*> + traverse (mapOnCompType tv) t <*> + pure loc + astMap tv (Map fun e t loc) = + Map <$> mapOnExp tv fun <*> mapOnExp tv e <*> + traverse (mapOnCompType tv) t <*> pure loc + astMap tv (Reduce comm fun startexp arrexp loc) = + Reduce comm <$> mapOnExp tv fun <*> + mapOnExp tv startexp <*> mapOnExp tv arrexp <*> pure loc + astMap tv (GenReduce hist op ne bfun img loc) = + GenReduce <$> mapOnExp tv hist <*> mapOnExp tv op <*> mapOnExp tv ne + <*> mapOnExp tv bfun <*> mapOnExp tv img <*> pure loc + astMap tv (Zip i e es t loc) = + Zip i <$> mapOnExp tv e <*> mapM (mapOnExp tv) es <*> + traverse (mapOnCompType tv) t <*> pure loc + astMap tv (Unzip e ts loc) = + Unzip <$> mapOnExp tv e <*> mapM (traverse $ mapOnCompType tv) ts <*> pure loc + astMap tv (Unsafe e loc) = + Unsafe <$> mapOnExp tv e <*> pure loc + astMap tv (Assert e1 e2 desc loc) = + Assert <$> mapOnExp tv e1 <*> mapOnExp tv e2 <*> pure desc <*> pure loc + astMap tv (Scan fun startexp arrexp loc) = + pure Scan <*> mapOnExp tv fun <*> + mapOnExp tv startexp <*> mapOnExp tv arrexp <*> + pure loc + astMap tv (Filter fun arrexp loc) = + pure Filter <*> mapOnExp tv fun <*> mapOnExp tv arrexp <*> pure loc + astMap tv (Partition k fun arrexp loc) = + Partition k <$> mapOnExp tv fun <*> mapOnExp tv arrexp <*> pure loc + astMap tv (Stream form fun arr loc) = + pure Stream <*> mapOnStreamForm form <*> mapOnExp tv fun <*> + mapOnExp tv arr <*> pure loc + where mapOnStreamForm (MapLike o) = pure $ MapLike o + mapOnStreamForm (RedLike o comm lam) = + RedLike o comm <$> mapOnExp tv lam + astMap tv (Lambda tparams params body ret t loc) = + Lambda <$> mapM (astMap tv) tparams <*> mapM (astMap tv) params <*> + astMap tv body <*> traverse (astMap tv) ret <*> + traverse (traverse $ mapOnStructType tv) t <*> pure loc + astMap tv (OpSection name t loc) = + OpSection <$> mapOnQualName tv name <*> + traverse (mapOnPatternType tv) t <*> pure loc + astMap tv (OpSectionLeft name t arg (t1a, t1b) t2 loc) = + OpSectionLeft <$> mapOnQualName tv name <*> + traverse (mapOnPatternType tv) t <*> mapOnExp tv arg <*> + ((,) <$> traverse (mapOnStructType tv) t1a <*> + traverse (mapOnStructType tv) t1b) <*> + traverse (mapOnPatternType tv) t2 <*> pure loc + astMap tv (OpSectionRight name t arg (t1a, t1b) t2 loc) = + OpSectionRight <$> mapOnQualName tv name <*> + traverse (mapOnPatternType tv) t <*> mapOnExp tv arg <*> + ((,) <$> traverse (mapOnStructType tv) t1a <*> + traverse (mapOnStructType tv) t1b) <*> + traverse (mapOnPatternType tv) t2 <*> pure loc + astMap tv (ProjectSection fields t loc) = + ProjectSection fields <$> traverse (mapOnPatternType tv) t <*> pure loc + astMap tv (IndexSection idxs t loc) = + IndexSection <$> mapM (astMap tv) idxs <*> + traverse (mapOnPatternType tv) t <*> pure loc + astMap tv (DoLoop tparams mergepat mergeexp form loopbody loc) = + DoLoop <$> mapM (astMap tv) tparams <*> astMap tv mergepat <*> + mapOnExp tv mergeexp <*> astMap tv form <*> + mapOnExp tv loopbody <*> pure loc + +instance ASTMappable (LoopFormBase Info VName) where + astMap tv (For i bound) = For <$> astMap tv i <*> astMap tv bound + astMap tv (ForIn pat e) = ForIn <$> astMap tv pat <*> astMap tv e + astMap tv (While e) = While <$> astMap tv e + +instance ASTMappable (TypeExp VName) where + astMap tv (TEVar qn loc) = TEVar <$> mapOnQualName tv qn <*> pure loc + astMap tv (TETuple ts loc) = TETuple <$> traverse (astMap tv) ts <*> pure loc + astMap tv (TERecord ts loc) = + TERecord <$> traverse (traverse $ astMap tv) ts <*> pure loc + astMap tv (TEArray te dim loc) = + TEArray <$> astMap tv te <*> astMap tv dim <*> pure loc + astMap tv (TEUnique t loc) = TEUnique <$> astMap tv t <*> pure loc + astMap tv (TEApply t1 t2 loc) = + TEApply <$> astMap tv t1 <*> astMap tv t2 <*> pure loc + astMap tv (TEArrow v t1 t2 loc) = + TEArrow v <$> astMap tv t1 <*> astMap tv t2 <*> pure loc + +instance ASTMappable (TypeArgExp VName) where + astMap tv (TypeArgExpDim dim loc) = + TypeArgExpDim <$> astMap tv dim <*> pure loc + astMap tv (TypeArgExpType te) = + TypeArgExpType <$> astMap tv te + +instance ASTMappable (DimDecl VName) where + astMap tv (NamedDim vn) = NamedDim <$> mapOnQualName tv vn + astMap _ (ConstDim k) = pure $ ConstDim k + astMap _ AnyDim = pure AnyDim + +instance ASTMappable (TypeParamBase VName) where + astMap = traverse . mapOnName + +instance ASTMappable (DimIndexBase Info VName) where + astMap tv (DimFix j) = DimFix <$> astMap tv j + astMap tv (DimSlice i j stride) = + DimSlice <$> + maybe (return Nothing) (fmap Just . astMap tv) i <*> + maybe (return Nothing) (fmap Just . astMap tv) j <*> + maybe (return Nothing) (fmap Just . astMap tv) stride + +instance ASTMappable Names where + astMap tv = fmap S.fromList . traverse (mapOnName tv) . S.toList + +type TypeTraverser f t dim1 als1 dim2 als2 = + (TypeName -> f TypeName) -> (dim1 -> f dim2) -> (als1 -> f als2) -> + t dim1 als1 -> f (t dim2 als2) + +traverseType :: Applicative f => + TypeTraverser f TypeBase dim1 als1 dims als2 +traverseType _ _ _ (Prim t) = pure $ Prim t +traverseType f g h (Array et shape u) = + Array <$> traverseArrayElemType f g h et <*> traverse g shape <*> pure u +traverseType f g h (Record fs) = Record <$> traverse (traverseType f g h) fs +traverseType f g h (TypeVar als u t args) = + TypeVar <$> h als <*> pure u <*> f t <*> traverse (traverseTypeArg f g h) args +traverseType f g h (Arrow als v t1 t2) = + Arrow <$> h als <*> pure v <*> traverseType f g h t1 <*> traverseType f g h t2 + +traverseArrayElemType :: Applicative f => + TypeTraverser f ArrayElemTypeBase dim1 als1 dim2 als2 +traverseArrayElemType _ _ h (ArrayPrimElem t as) = + ArrayPrimElem t <$> h as +traverseArrayElemType f g h (ArrayPolyElem t args as) = + ArrayPolyElem <$> f t <*> traverse (traverseTypeArg f g h) args <*> h as +traverseArrayElemType f g h (ArrayRecordElem fs) = + ArrayRecordElem <$> traverse (traverseRecordArrayElemType f g h) fs + +traverseRecordArrayElemType :: Applicative f => + TypeTraverser f RecordArrayElemTypeBase dim1 als1 dim2 als2 +traverseRecordArrayElemType f g h (RecordArrayElem et) = + RecordArrayElem <$> traverseArrayElemType f g h et +traverseRecordArrayElemType f g h (RecordArrayArrayElem et shape u) = + RecordArrayArrayElem <$> traverseArrayElemType f g h et <*> + traverse g shape <*> pure u + +traverseTypeArg :: Applicative f => + TypeTraverser f TypeArg dim1 als1 dim2 als2 +traverseTypeArg _ g _ (TypeArgDim d loc) = TypeArgDim <$> g d <*> pure loc +traverseTypeArg f g h (TypeArgType t loc) = TypeArgType <$> traverseType f g h t <*> pure loc + +instance ASTMappable (TypeBase () ()) where + astMap tv = traverseType f pure pure + where f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName + +instance ASTMappable CompType where + astMap tv = traverseType f pure (astMap tv) + where f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName + +instance ASTMappable StructType where + astMap tv = traverseType f (astMap tv) pure + where f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName + +instance ASTMappable PatternType where + astMap tv = traverseType f (astMap tv) (astMap tv) + where f = fmap typeNameFromQualName . mapOnQualName tv . qualNameFromTypeName + +instance ASTMappable (TypeDeclBase Info VName) where + astMap tv (TypeDecl dt (Info et)) = + TypeDecl <$> astMap tv dt <*> (Info <$> mapOnStructType tv et) + +instance ASTMappable (IdentBase Info VName) where + astMap tv (Ident name (Info t) loc) = + Ident <$> mapOnName tv name <*> (Info <$> mapOnCompType tv t) <*> pure loc + +instance ASTMappable (PatternBase Info VName) where + astMap tv (Id name (Info t) loc) = + Id <$> mapOnName tv name <*> (Info <$> mapOnPatternType tv t) <*> pure loc + astMap tv (TuplePattern pats loc) = + TuplePattern <$> mapM (astMap tv) pats <*> pure loc + astMap tv (RecordPattern fields loc) = + RecordPattern <$> mapM (traverse $ astMap tv) fields <*> pure loc + astMap tv (PatternParens pat loc) = + PatternParens <$> astMap tv pat <*> pure loc + astMap tv (PatternAscription pat t loc) = + PatternAscription <$> astMap tv pat <*> astMap tv t <*> pure loc + astMap tv (Wildcard (Info t) loc) = + Wildcard <$> (Info <$> mapOnPatternType tv t) <*> pure loc + +instance ASTMappable (FieldBase Info VName) where + astMap tv (RecordFieldExplicit name e loc) = + RecordFieldExplicit name <$> mapOnExp tv e <*> pure loc + astMap tv (RecordFieldImplicit name t loc) = + RecordFieldImplicit <$> mapOnName tv name + <*> traverse (mapOnCompType tv) t <*> pure loc + +instance ASTMappable a => ASTMappable (Info a) where + astMap tv = traverse $ astMap tv + +instance ASTMappable a => ASTMappable [a] where + astMap tv = traverse $ astMap tv + +instance (ASTMappable a, ASTMappable b) => ASTMappable (a,b) where + astMap tv (x,y) = (,) <$> astMap tv x <*> astMap tv y + +instance (ASTMappable a, ASTMappable b, ASTMappable c) => ASTMappable (a,b,c) where + astMap tv (x,y,z) = (,,) <$> astMap tv x <*> astMap tv y <*> astMap tv z diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker.hs @@ -0,0 +1,902 @@ +{-# LANGUAGE FlexibleContexts, TupleSections #-} +-- | The type checker checks whether the program is type-consistent +-- and adds type annotations and various other elaborations. The +-- program does not need to have any particular properties for the +-- type checker to function; in particular it does not need unique +-- names. +module Language.Futhark.TypeChecker + ( checkProg + , checkExp + , checkDec + , TypeError + , Warnings + , initialEnv + ) + where + +import Control.Monad.Except +import Control.Monad.Writer +import Data.List +import Data.Loc +import Data.Maybe +import Data.Either +import Data.Ord +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Prelude hiding (abs, mod) + +import Language.Futhark +import Language.Futhark.Semantic +import Futhark.FreshNames hiding (newName) +import Language.Futhark.TypeChecker.Monad +import Language.Futhark.TypeChecker.Terms +import Language.Futhark.TypeChecker.Unify (doUnification) +import Language.Futhark.TypeChecker.Types + +--- The main checker + +-- | Type check a program containing no type information, yielding +-- either a type error or a program with complete type information. +-- Accepts a mapping from file names (excluding extension) to +-- previously type checker results. The 'FilePath' is used to resolve +-- relative @import@s. +checkProg :: Imports + -> VNameSource + -> ImportName + -> UncheckedProg + -> Either TypeError (FileModule, Warnings, VNameSource) +checkProg files src name prog = + runTypeM initialEnv files' name src $ checkProgM prog + where files' = M.map fileEnv $ M.fromList files + +-- | Type check a single expression containing no type information, +-- yielding either a type error or the same expression annotated with +-- type information. See also 'checkProg'. +checkExp :: Imports + -> VNameSource + -> Env + -> UncheckedExp + -> Either TypeError Exp +checkExp files src env e = do + (e', _, _) <- runTypeM env files' (mkInitialImport "") src $ + checkOneExp e + return e' + where files' = M.map fileEnv $ M.fromList files + +-- | Type check a single declaration containing no type information, +-- yielding either a type error or the same expression annotated with +-- type information along the Env produced by that declaration. See +-- also 'checkProg'. +checkDec :: Imports + -> VNameSource + -> Env + -> ImportName + -> UncheckedDec + -> Either TypeError (Env, Dec, VNameSource) +checkDec files src env name d = do + ((env', d'), _, src') <- runTypeM env files' name src $ do + (_, env', d') <- checkOneDec d + return (env' <> env, d') + return (env', d', src') + where files' = M.map fileEnv $ M.fromList files + +-- | An initial environment for the type checker, containing +-- intrinsics and such. +initialEnv :: Env +initialEnv = intrinsicsModule + { envModTable = initialModTable + , envNameMap = M.insert + (Term, nameFromString "intrinsics") + (qualName intrinsics_v) + topLevelNameMap + } + where initialTypeTable = M.fromList $ mapMaybe addIntrinsicT $ M.toList intrinsics + initialModTable = M.singleton intrinsics_v (ModEnv intrinsicsModule) + + intrinsics_v = VName (nameFromString "intrinsics") 0 + + intrinsicsModule = Env mempty initialTypeTable mempty mempty intrinsicsNameMap + + addIntrinsicT (name, IntrinsicType t) = + Just (name, TypeAbbr Unlifted [] $ Prim t) + addIntrinsicT _ = + Nothing + +checkProgM :: UncheckedProg -> TypeM FileModule +checkProgM (Prog doc decs) = do + checkForDuplicateDecs decs + (abs, env, decs') <- checkDecs decs + return (FileModule abs env $ Prog doc decs') + +dupDefinitionError :: MonadTypeChecker m => + Namespace -> Name -> SrcLoc -> SrcLoc -> m a +dupDefinitionError space name pos1 pos2 = + throwError $ TypeError pos1 $ + "Duplicate definition of " ++ ppSpace space ++ " " ++ + nameToString name ++ ". Previously defined at " ++ locStr pos2 + +checkForDuplicateDecs :: [DecBase NoInfo Name] -> TypeM () +checkForDuplicateDecs = + foldM_ (flip f) mempty + where check namespace name loc known = + case M.lookup (namespace, name) known of + Just loc' -> + dupDefinitionError namespace name loc loc' + _ -> return $ M.insert (namespace, name) loc known + + f (ValDec (ValBind _ name _ _ _ _ _ _ loc)) = + check Term name loc + + f (TypeDec (TypeBind name _ _ _ loc)) = + check Type name loc + + f (SigDec (SigBind name _ _ loc)) = + check Signature name loc + + f (ModDec (ModBind name _ _ _ _ loc)) = + check Term name loc + + f OpenDec{} = return + + f LocalDec{} = return + +bindingTypeParams :: [TypeParam] -> TypeM a -> TypeM a +bindingTypeParams tparams = localEnv env + where env = mconcat $ map typeParamEnv tparams + + typeParamEnv (TypeParamDim v _) = + mempty { envVtable = + M.singleton v $ BoundV [] (Prim (Signed Int32)) } + typeParamEnv (TypeParamType l v _) = + mempty { envTypeTable = + M.singleton v $ TypeAbbr l [] $ TypeVar () Nonunique (typeName v) [] } + +checkSpecs :: [SpecBase NoInfo Name] -> TypeM (TySet, Env, [SpecBase Info VName]) + +checkSpecs [] = return (mempty, mempty, []) + +checkSpecs (ValSpec name tparams vtype doc loc : specs) = + bindSpaced [(Term, name)] $ do + name' <- checkName Term name loc + (tparams', rettype') <- + checkTypeParams tparams $ \tparams' -> bindingTypeParams tparams' $ do + (vtype', _) <- checkTypeDecl vtype + return (tparams', vtype') + + let binding = BoundV tparams' $ unInfo $ expandedType rettype' + valenv = + mempty { envVtable = M.singleton name' binding + , envNameMap = M.singleton (Term, name) $ qualName name' + } + (abstypes, env, specs') <- localEnv valenv $ checkSpecs specs + return (abstypes, + env <> valenv, + ValSpec name' tparams' rettype' doc loc : specs') + +checkSpecs (TypeAbbrSpec tdec : specs) = + bindSpaced [(Type, typeAlias tdec)] $ do + (tenv, tdec') <- checkTypeBind tdec + (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs + return (abstypes, + tenv <> env, + TypeAbbrSpec tdec' : specs') + +checkSpecs (TypeSpec l name ps doc loc : specs) = + checkTypeParams ps $ \ps' -> + bindSpaced [(Type, name)] $ do + name' <- checkName Type name loc + let tenv = mempty + { envNameMap = + M.singleton (Type, name) $ qualName name' + , envTypeTable = + M.singleton name' $ TypeAbbr l ps' $ + TypeVar () Nonunique (typeName name') $ map typeParamToArg ps' + } + (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs + return (M.insert (qualName name') l abstypes, + tenv <> env, + TypeSpec l name' ps' doc loc : specs') + +checkSpecs (ModSpec name sig doc loc : specs) = + bindSpaced [(Term, name)] $ do + name' <- checkName Term name loc + (mty, sig') <- checkSigExp sig + let senv = mempty { envNameMap = M.singleton (Term, name) $ qualName name' + , envModTable = M.singleton name' $ mtyMod mty + } + (abstypes, env, specs') <- localEnv senv $ checkSpecs specs + return (M.mapKeys (qualify name') (mtyAbs mty) <> abstypes, + senv <> env, + ModSpec name' sig' doc loc : specs') + +checkSpecs (IncludeSpec e loc : specs) = do + (e_abs, e_env, e') <- checkSigExpToEnv e + + mapM_ (warnIfShadowing . fmap baseName) $ M.keys e_abs + + (abstypes, env, specs') <- localEnv e_env $ checkSpecs specs + return (e_abs <> abstypes, + e_env <> env, + IncludeSpec e' loc : specs') + where warnIfShadowing qn = + (lookupType loc qn >> warnAbout qn) + `catchError` \_ -> return () + warnAbout qn = + warn loc $ "Inclusion shadows type `" ++ pretty qn ++ "`." + +checkSigExp :: SigExpBase NoInfo Name -> TypeM (MTy, SigExpBase Info VName) +checkSigExp (SigParens e loc) = do + (mty, e') <- checkSigExp e + return (mty, SigParens e' loc) +checkSigExp (SigVar name loc) = do + (name', mty) <- lookupMTy loc name + (mty', _) <- newNamesForMTy mty + return (mty', SigVar name' loc) +checkSigExp (SigSpecs specs loc) = do + checkForDuplicateSpecs specs + (abstypes, env, specs') <- checkSpecs specs + return (MTy abstypes $ ModEnv env, SigSpecs specs' loc) +checkSigExp (SigWith s (TypeRef tname ps td trloc) loc) = do + (s_abs, s_env, s') <- checkSigExpToEnv s + checkTypeParams ps $ \ps' -> do + (td', _) <- bindingTypeParams ps' $ checkTypeDecl td + (tname', s_abs', s_env') <- refineEnv loc s_abs s_env tname ps' $ unInfo $ expandedType td' + return (MTy s_abs' $ ModEnv s_env', SigWith s' (TypeRef tname' ps' td' trloc) loc) +checkSigExp (SigArrow maybe_pname e1 e2 loc) = do + (MTy s_abs e1_mod, e1') <- checkSigExp e1 + (env_for_e2, maybe_pname') <- + case maybe_pname of + Just pname -> bindSpaced [(Term, pname)] $ do + pname' <- checkName Term pname loc + return (mempty { envNameMap = M.singleton (Term, pname) $ qualName pname' + , envModTable = M.singleton pname' e1_mod + }, + Just pname') + Nothing -> + return (mempty, Nothing) + (e2_mod, e2') <- localEnv env_for_e2 $ checkSigExp e2 + return (MTy mempty $ ModFun $ FunSig s_abs e1_mod e2_mod, + SigArrow maybe_pname' e1' e2' loc) + +checkSigExpToEnv :: SigExpBase NoInfo Name -> TypeM (TySet, Env, SigExpBase Info VName) +checkSigExpToEnv e = do + (MTy abs mod, e') <- checkSigExp e + case mod of + ModEnv env -> return (abs, env, e') + ModFun{} -> unappliedFunctor $ srclocOf e + +checkSigBind :: SigBindBase NoInfo Name -> TypeM (Env, SigBindBase Info VName) +checkSigBind (SigBind name e doc loc) = do + (env, e') <- checkSigExp e + bindSpaced [(Signature, name)] $ do + name' <- checkName Signature name loc + return (mempty { envSigTable = M.singleton name' env + , envNameMap = M.singleton (Signature, name) (qualName name') + }, + SigBind name' e' doc loc) + +checkModExp :: ModExpBase NoInfo Name -> TypeM (MTy, ModExpBase Info VName) +checkModExp (ModParens e loc) = do + (mty, e') <- checkModExp e + return (mty, ModParens e' loc) +checkModExp (ModDecs decs loc) = do + checkForDuplicateDecs decs + (abstypes, env, decs') <- checkDecs decs + return (MTy abstypes $ ModEnv env, + ModDecs decs' loc) +checkModExp (ModVar v loc) = do + (v', env) <- lookupMod loc v + when (baseName (qualLeaf v') == nameFromString "intrinsics" && + baseTag (qualLeaf v') <= maxIntrinsicTag) $ + throwError $ TypeError loc "The 'intrinsics' module may not be used in module expressions." + return (MTy mempty env, ModVar v' loc) +checkModExp (ModImport name NoInfo loc) = do + (name', env) <- lookupImport loc name + return (MTy mempty $ ModEnv env, + ModImport name (Info name') loc) +checkModExp (ModApply f e NoInfo NoInfo loc) = do + (f_mty, f') <- checkModExp f + case mtyMod f_mty of + ModFun functor -> do + (e_mty, e') <- checkModExp e + (mty, psubsts, rsubsts) <- applyFunctor loc functor e_mty + return (mty, ModApply f' e' (Info psubsts) (Info rsubsts) loc) + _ -> + throwError $ TypeError loc "Cannot apply non-parametric module." +checkModExp (ModAscript me se NoInfo loc) = do + (me_mod, me') <- checkModExp me + (se_mty, se') <- checkSigExp se + match_subst <- badOnLeft $ matchMTys me_mod se_mty loc + return (se_mty, ModAscript me' se' (Info match_subst) loc) +checkModExp (ModLambda param maybe_fsig_e body_e loc) = + withModParam param $ \param' param_abs param_mod -> do + (maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc + return (MTy mempty $ ModFun $ FunSig param_abs param_mod mty, + ModLambda param' maybe_fsig_e' body_e' loc) + +checkModExpToEnv :: ModExpBase NoInfo Name -> TypeM (TySet, Env, ModExpBase Info VName) +checkModExpToEnv e = do + (MTy abs mod, e') <- checkModExp e + case mod of + ModEnv env -> return (abs, env, e') + ModFun{} -> unappliedFunctor $ srclocOf e + +withModParam :: ModParamBase NoInfo Name + -> (ModParamBase Info VName -> TySet -> Mod -> TypeM a) + -> TypeM a +withModParam (ModParam pname psig_e NoInfo loc) m = do + (MTy p_abs p_mod, psig_e') <- checkSigExp psig_e + bindSpaced [(Term, pname)] $ do + pname' <- checkName Term pname loc + let in_body_env = mempty { envModTable = M.singleton pname' p_mod } + localEnv in_body_env $ + m (ModParam pname' psig_e' (Info $ map qualLeaf $ M.keys p_abs) loc) p_abs p_mod + +withModParams :: [ModParamBase NoInfo Name] + -> ([(ModParamBase Info VName, TySet, Mod)] -> TypeM a) + -> TypeM a +withModParams [] m = m [] +withModParams (p:ps) m = + withModParam p $ \p' pabs pmod -> + withModParams ps $ \ps' -> m $ (p',pabs,pmod) : ps' + +checkModBody :: Maybe (SigExpBase NoInfo Name) + -> ModExpBase NoInfo Name + -> SrcLoc + -> TypeM (Maybe (SigExp, Info (M.Map VName VName)), + ModExp, MTy) +checkModBody maybe_fsig_e body_e loc = do + (body_mty, body_e') <- checkModExp body_e + case maybe_fsig_e of + Nothing -> + return (Nothing, body_e', body_mty) + Just fsig_e -> do + (fsig_mty, fsig_e') <- checkSigExp fsig_e + fsig_subst <- badOnLeft $ matchMTys body_mty fsig_mty loc + return (Just (fsig_e', Info fsig_subst), body_e', fsig_mty) + +applyFunctor :: SrcLoc + -> FunSig + -> MTy + -> TypeM (MTy, + M.Map VName VName, + M.Map VName VName) +applyFunctor applyloc (FunSig p_abs p_mod body_mty) a_mty = do + p_subst <- badOnLeft $ matchMTys a_mty (MTy p_abs p_mod) applyloc + + -- Apply type abbreviations from a_mty to body_mty. + let a_abbrs = mtyTypeAbbrs a_mty + let type_subst = M.mapMaybe (fmap TypeSub . (`M.lookup` a_abbrs)) p_subst + let body_mty' = substituteTypesInMTy type_subst body_mty + (body_mty'', body_subst) <- newNamesForMTy body_mty' + return (body_mty'', p_subst, body_subst) + +checkModBind :: ModBindBase NoInfo Name -> TypeM (TySet, Env, ModBindBase Info VName) +checkModBind (ModBind name [] maybe_fsig_e e doc loc) = do + (maybe_fsig_e', e', mty) <- checkModBody (fst <$> maybe_fsig_e) e loc + bindSpaced [(Term, name)] $ do + name' <- checkName Term name loc + return (mtyAbs mty, + mempty { envModTable = M.singleton name' $ mtyMod mty + , envNameMap = M.singleton (Term, name) $ qualName name' + }, + ModBind name' [] maybe_fsig_e' e' doc loc) +checkModBind (ModBind name (p:ps) maybe_fsig_e body_e doc loc) = do + (params', maybe_fsig_e', body_e', funsig) <- + withModParam p $ \p' p_abs p_mod -> + withModParams ps $ \params_stuff -> do + let (ps', ps_abs, ps_mod) = unzip3 params_stuff + (maybe_fsig_e', body_e', mty) <- checkModBody (fst <$> maybe_fsig_e) body_e loc + let addParam (x,y) mty' = MTy mempty $ ModFun $ FunSig x y mty' + return (p' : ps', maybe_fsig_e', body_e', + FunSig p_abs p_mod $ foldr addParam mty $ zip ps_abs ps_mod) + bindSpaced [(Term, name)] $ do + name' <- checkName Term name loc + return (mempty, + mempty { envModTable = + M.singleton name' $ ModFun funsig + , envNameMap = + M.singleton (Term, name) $ qualName name' + }, + ModBind name' params' maybe_fsig_e' body_e' doc loc) + +checkForDuplicateSpecs :: [SpecBase NoInfo Name] -> TypeM () +checkForDuplicateSpecs = + foldM_ (flip f) mempty + where check namespace name loc known = + case M.lookup (namespace, name) known of + Just loc' -> + dupDefinitionError namespace name loc loc' + _ -> return $ M.insert (namespace, name) loc known + + f (ValSpec name _ _ _ loc) = + check Term name loc + + f (TypeAbbrSpec (TypeBind name _ _ _ loc)) = + check Type name loc + + f (TypeSpec _ name _ _ loc) = + check Type name loc + + f (ModSpec name _ _ loc) = + check Term name loc + + f IncludeSpec{} = + return + +checkTypeBind :: TypeBindBase NoInfo Name + -> TypeM (Env, TypeBindBase Info VName) +checkTypeBind (TypeBind name ps td doc loc) = + checkTypeParams ps $ \ps' -> do + (td', l) <- bindingTypeParams ps' $ checkTypeDecl td + bindSpaced [(Type, name)] $ do + name' <- checkName Type name loc + return (mempty { envTypeTable = + M.singleton name' $ TypeAbbr l ps' $ unInfo $ expandedType td', + envNameMap = + M.singleton (Type, name) $ qualName name' + }, + TypeBind name' ps' td' doc loc) + +checkValBind :: ValBindBase NoInfo Name -> TypeM (Env, ValBind) +checkValBind (ValBind entry fname maybe_tdecl NoInfo tparams params body doc loc) = do + (fname', tparams', params', maybe_tdecl', rettype, body') <- + checkFunDef (fname, maybe_tdecl, tparams, params, body, loc) + + when (entry && any isTypeParam tparams') $ + throwError $ TypeError loc "Entry point functions may not be polymorphic." + + when (entry && singleTuplePattern params') $ + warn loc "This entry point accepts a *single* tuple-typed parameter, *not* multiple parameters.\nThis will be an error in the future." + + let (rettype_params, rettype') = unfoldFunType rettype + when (entry && (any (not . patternOrderZero) params' || + any (not . orderZero) rettype_params || + not (orderZero rettype'))) $ + throwError $ TypeError loc "Entry point functions may not be higher-order." + + return (mempty { envVtable = + M.singleton fname' $ + BoundV tparams' $ foldr (uncurry (Arrow ()) . patternParam) rettype params' + , envNameMap = + M.singleton (Term, fname) $ qualName fname' + }, + ValBind entry fname' maybe_tdecl' (Info rettype) tparams' params' body' doc loc) + +singleTuplePattern :: [Pattern] -> Bool +singleTuplePattern [TuplePattern _ _] = True +singleTuplePattern _ = False + +checkOneDec :: DecBase NoInfo Name -> TypeM (TySet, Env, DecBase Info VName) +checkOneDec (ModDec struct) = do + (abs, modenv, struct') <- checkModBind struct + return (abs, modenv, ModDec struct') + +checkOneDec (SigDec sig) = do + (sigenv, sig') <- checkSigBind sig + return (mempty, sigenv, SigDec sig') + +checkOneDec (TypeDec tdec) = do + (tenv, tdec') <- checkTypeBind tdec + return (mempty, tenv, TypeDec tdec') + +checkOneDec (OpenDec x NoInfo loc) = do + (x_abs, x_env, x') <- checkModExpToEnv x + let names = S.toList $ allNamesInEnv x_env + return (x_abs, + x_env, + OpenDec x' (Info names) loc) + +checkOneDec (LocalDec d loc) = do + (abstypes, env, d') <- checkOneDec d + return (abstypes, env, LocalDec d' loc) + +checkOneDec (ValDec vb) = do + (env, vb') <- checkValBind vb + return (mempty, env, ValDec vb') + +checkDecs :: [DecBase NoInfo Name] -> TypeM (TySet, Env, [DecBase Info VName]) +checkDecs (LocalDec d loc:ds) = do + (d_abstypes, d_env, d') <- checkOneDec d + (ds_abstypes, ds_env, ds') <- localEnv d_env $ checkDecs ds + return (d_abstypes <> ds_abstypes, + ds_env, + LocalDec d' loc : ds') + +checkDecs (d:ds) = do + (d_abstypes, d_env, d') <- checkOneDec d + (ds_abstypes, ds_env, ds') <- localEnv d_env $ checkDecs ds + return (d_abstypes <> ds_abstypes, + ds_env <> d_env, + d' : ds') + +checkDecs [] = + return (mempty, mempty, []) + +--- Signature matching + +-- Return new renamed/abstracted env, as well as a mapping from +-- names in the signature to names in the new env. This is used for +-- functor application. The first env is the module env, and the +-- second the env it must match. +matchMTys :: MTy -> MTy -> SrcLoc + -> Either TypeError (M.Map VName VName) +matchMTys = matchMTys' mempty + where + matchMTys' :: TypeSubs -> MTy -> MTy -> SrcLoc + -> Either TypeError (M.Map VName VName) + + matchMTys' _ (MTy _ ModFun{}) (MTy _ ModEnv{}) loc = + Left $ TypeError loc "Cannot match parametric module with non-paramatric module type." + + matchMTys' _ (MTy _ ModEnv{}) (MTy _ ModFun{}) loc = + Left $ TypeError loc "Cannot match non-parametric module with paramatric module type." + + matchMTys' old_abs_subst_to_type (MTy mod_abs mod) (MTy sig_abs sig) loc = do + -- Check that abstract types in 'sig' have an implementation in + -- 'mod'. This also gives us a substitution that we use to check + -- the types of values. + abs_substs <- resolveAbsTypes mod_abs mod sig_abs loc + + let abs_subst_to_type = old_abs_subst_to_type <> + M.map (TypeSub . snd) abs_substs + abs_name_substs = M.map (qualLeaf . fst) abs_substs + substs <- matchMods abs_subst_to_type mod sig loc + return (substs <> abs_name_substs) + + matchMods :: TypeSubs -> Mod -> Mod -> SrcLoc + -> Either TypeError (M.Map VName VName) + matchMods _ ModEnv{} ModFun{} loc = + Left $ TypeError loc "Cannot match non-parametric module with paramatric module type." + matchMods _ ModFun{} ModEnv{} loc = + Left $ TypeError loc "Cannot match parametric module with non-paramatric module type." + + matchMods abs_subst_to_type (ModEnv mod) (ModEnv sig) loc = + matchEnvs abs_subst_to_type mod sig loc + + matchMods old_abs_subst_to_type + (ModFun (FunSig mod_abs mod_pmod mod_mod)) + (ModFun (FunSig sig_abs sig_pmod sig_mod)) + loc = do + abs_substs <- resolveAbsTypes mod_abs mod_pmod sig_abs loc + let abs_subst_to_type = old_abs_subst_to_type <> + M.map (TypeSub . snd) abs_substs + abs_name_substs = M.map (qualLeaf . fst) abs_substs + pmod_substs <- matchMods abs_subst_to_type mod_pmod sig_pmod loc + mod_substs <- matchMTys' abs_subst_to_type mod_mod sig_mod loc + return (pmod_substs <> mod_substs <> abs_name_substs) + + matchEnvs :: TypeSubs + -> Env -> Env -> SrcLoc + -> Either TypeError (M.Map VName VName) + matchEnvs abs_subst_to_type env sig loc = do + -- XXX: we only want to create substitutions for visible names. + -- This must be wrong in some cases. Probably we need to + -- rethink how we do shadowing for module types. + let visible = S.fromList $ map qualLeaf $ M.elems $ envNameMap sig + isVisible name = name `S.member` visible + + -- Check that all values are defined correctly, substituting the + -- abstract types first. + val_substs <- fmap M.fromList $ forM (M.toList $ envVtable sig) $ \(name, spec_bv) -> do + let spec_bv' = substituteTypesInBoundV abs_subst_to_type spec_bv + case findBinding envVtable Term (baseName name) env of + Just (name', bv) -> matchVal loc name spec_bv' name' bv + _ -> missingVal loc (baseName name) + + -- Check that all type abbreviations are correctly defined. + abbr_name_substs <- fmap M.fromList $ + forM (filter (isVisible . fst) $ M.toList $ + envTypeTable sig) $ \(name, TypeAbbr _ spec_ps spec_t) -> + case findBinding envTypeTable Type (baseName name) env of + Just (name', TypeAbbr _ ps t) -> + matchTypeAbbr loc abs_subst_to_type val_substs name spec_ps spec_t name' ps t + Nothing -> missingType loc $ baseName name + + -- Check for correct modules. + mod_substs <- fmap M.unions $ forM (M.toList $ envModTable sig) $ \(name, modspec) -> + case findBinding envModTable Term (baseName name) env of + Just (name', mod) -> + M.insert name name' <$> matchMods abs_subst_to_type mod modspec loc + Nothing -> + missingMod loc $ baseName name + + return $ val_substs <> mod_substs <> abbr_name_substs + + matchTypeAbbr :: SrcLoc -> TypeSubs -> M.Map VName VName + -> VName -> [TypeParam] -> StructType + -> VName -> [TypeParam] -> StructType + -> Either TypeError (VName, VName) + matchTypeAbbr loc abs_subst_to_type val_substs spec_name spec_ps spec_t name ps t = do + -- We have to create substitutions for the type parameters, too. + unless (length spec_ps == length ps) nomatch + param_substs <- mconcat <$> zipWithM matchTypeParam spec_ps ps + let val_substs' = M.map (DimSub . NamedDim . qualName) val_substs + spec_t' = substituteTypes (val_substs'<>param_substs<>abs_subst_to_type) spec_t + if spec_t' == t + then return (spec_name, name) + else nomatch + where nomatch = mismatchedType loc (M.keys abs_subst_to_type) + (baseName spec_name) (spec_ps, spec_t) (ps, t) + + matchTypeParam (TypeParamDim x _) (TypeParamDim y _) = + pure $ M.singleton x $ DimSub $ NamedDim $ qualName y + matchTypeParam (TypeParamType Unlifted x _) (TypeParamType Unlifted y _) = + pure $ M.singleton x $ TypeSub $ TypeAbbr Unlifted [] $ + TypeVar () Nonunique (typeName y) [] + matchTypeParam (TypeParamType _ x _) (TypeParamType Lifted y _) = + pure $ M.singleton x $ TypeSub $ TypeAbbr Lifted [] $ + TypeVar () Nonunique (typeName y) [] + matchTypeParam _ _ = + nomatch + + matchVal :: SrcLoc + -> VName -> BoundV + -> VName -> BoundV + -> Either TypeError (VName, VName) + matchVal loc spec_name spec_t name t + | matchFunBinding loc spec_t t = return (spec_name, name) + matchVal loc spec_name spec_v _ v = + Left $ TypeError loc $ "Value `" ++ baseString spec_name ++ "` specified as type " ++ + ppValBind spec_v ++ " in signature, but has " ++ ppValBind v ++ " in structure." + + matchFunBinding :: SrcLoc -> BoundV -> BoundV -> Bool + matchFunBinding loc (BoundV _ orig_spec_t) (BoundV tps orig_t) = + -- Would be nice if we could propagate the actual error here. + case doUnification loc tps + (toStructural orig_spec_t) (toStructural orig_t) of + Left _ -> False + Right t -> t `subtypeOf` toStructural orig_spec_t + + missingType loc name = + Left $ TypeError loc $ + "Module does not define a type named " ++ pretty name ++ "." + + missingVal loc name = + Left $ TypeError loc $ + "Module does not define a value named " ++ pretty name ++ "." + + missingMod loc name = + Left $ TypeError loc $ + "Module does not define a module named " ++ pretty name ++ "." + + mismatchedType loc abs name spec_t env_t = + Left $ TypeError loc $ + unlines ["Module defines", + indent $ ppTypeAbbr abs name env_t, + "but module type requires", + indent $ ppTypeAbbr abs name spec_t] + + indent = intercalate "\n" . map (" "++) . lines + + resolveAbsTypes :: TySet -> Mod -> TySet -> SrcLoc + -> Either TypeError (M.Map VName (QualName VName, TypeBinding)) + resolveAbsTypes mod_abs mod sig_abs loc = do + let abs_mapping = M.fromList $ zip + (map (fmap baseName . fst) $ M.toList mod_abs) (M.toList mod_abs) + fmap M.fromList $ forM (M.toList sig_abs) $ \(name, name_l) -> + case findTypeDef (fmap baseName name) mod of + Just (name', TypeAbbr mod_l ps t) + | Unlifted <- name_l, + not (orderZero t) || mod_l == Lifted -> + mismatchedLiftedness loc (map qualLeaf $ M.keys mod_abs) name (ps, t) + | Just (abs_name, _) <- M.lookup (fmap baseName name) abs_mapping -> + return (qualLeaf name, (abs_name, TypeAbbr name_l ps t)) + | otherwise -> + return (qualLeaf name, (name', TypeAbbr name_l ps t)) + _ -> + missingType loc $ fmap baseName name + + mismatchedLiftedness loc abs name mod_t = + Left $ TypeError loc $ + unlines ["Module defines", + indent $ ppTypeAbbr abs name mod_t, + "but module type requires this type to be non-functional."] + + ppValBind (BoundV tps t) = unwords $ map pretty tps ++ [pretty t] + + ppTypeAbbr abs name (ps, t) = + "type " ++ unwords (pretty name : map pretty ps) ++ t' + where t' = case t of + TypeVar () _ tn args + | typeLeaf tn `elem` abs, + map typeParamToArg ps == args -> "" + _ -> " = " ++ pretty t + +findBinding :: (Env -> M.Map VName v) + -> Namespace -> Name + -> Env + -> Maybe (VName, v) +findBinding table namespace name the_env = do + QualName _ name' <- M.lookup (namespace, name) $ envNameMap the_env + (name',) <$> M.lookup name' (table the_env) + +findTypeDef :: QualName Name -> Mod -> Maybe (QualName VName, TypeBinding) +findTypeDef _ ModFun{} = Nothing +findTypeDef (QualName [] name) (ModEnv the_env) = do + (name', tb) <- findBinding envTypeTable Type name the_env + return (qualName name', tb) +findTypeDef (QualName (q:qs) name) (ModEnv the_env) = do + (q', q_mod) <- findBinding envModTable Term q the_env + (QualName qs' name', tb) <- findTypeDef (QualName qs name) q_mod + return (QualName (q':qs') name', tb) + +typeParamToArg :: TypeParam -> StructTypeArg +typeParamToArg (TypeParamDim v ploc) = + TypeArgDim (NamedDim $ qualName v) ploc +typeParamToArg (TypeParamType _ v ploc) = + TypeArgType (TypeVar () Nonunique (typeName v) []) ploc + +substituteTypesInMod :: TypeSubs -> Mod -> Mod +substituteTypesInMod substs (ModEnv e) = + ModEnv $ substituteTypesInEnv substs e +substituteTypesInMod substs (ModFun (FunSig abs mod mty)) = + ModFun $ FunSig abs (substituteTypesInMod substs mod) (substituteTypesInMTy substs mty) + +substituteTypesInMTy :: TypeSubs -> MTy -> MTy +substituteTypesInMTy substs (MTy abs mod) = MTy abs $ substituteTypesInMod substs mod + +substituteTypesInEnv :: TypeSubs -> Env -> Env +substituteTypesInEnv substs env = + env { envVtable = M.map (substituteTypesInBoundV substs) $ envVtable env + , envTypeTable = M.mapWithKey subT $ envTypeTable env + , envModTable = M.map (substituteTypesInMod substs) $ envModTable env + } + where subT name _ + | Just (TypeSub (TypeAbbr l ps t)) <- M.lookup name substs = TypeAbbr l ps t + subT _ (TypeAbbr l ps t) = TypeAbbr l ps $ substituteTypes substs t + +allNamesInMTy :: MTy -> S.Set VName +allNamesInMTy (MTy abs mod) = + S.fromList (map qualLeaf $ M.keys abs) <> allNamesInMod mod + +allNamesInMod :: Mod -> S.Set VName +allNamesInMod (ModEnv env) = allNamesInEnv env +allNamesInMod ModFun{} = mempty + +-- All names defined anywhere in the env. +allNamesInEnv :: Env -> S.Set VName +allNamesInEnv (Env vtable ttable stable modtable _names) = + S.fromList (M.keys vtable ++ M.keys ttable ++ + M.keys stable ++ M.keys modtable) <> + mconcat (map allNamesInMTy (M.elems stable) ++ + map allNamesInMod (M.elems modtable) ++ + map allNamesInType (M.elems ttable)) + where allNamesInType (TypeAbbr _ ps _) = S.fromList $ map typeParamName ps + +newNamesForMTy :: MTy -> TypeM (MTy, M.Map VName VName) +newNamesForMTy orig_mty = do + -- Create unique renames for the module type. + pairs <- forM (S.toList $ allNamesInMTy orig_mty) $ \v -> do + v' <- newName v + return (v, v') + let substs = M.fromList pairs + rev_substs = M.fromList $ map (uncurry $ flip (,)) pairs + + return (substituteInMTy substs orig_mty, rev_substs) + + where + substituteInMTy :: M.Map VName VName -> MTy -> MTy + substituteInMTy substs (MTy mty_abs mty_mod) = + MTy (M.mapKeys (fmap substitute) mty_abs) (substituteInMod mty_mod) + where + substituteInEnv (Env vtable ttable _stable modtable names) = + let vtable' = substituteInMap substituteInBinding vtable + ttable' = substituteInMap substituteInTypeBinding ttable + mtable' = substituteInMap substituteInMod modtable + in Env { envVtable = vtable' + , envTypeTable = ttable' + , envSigTable = mempty + , envModTable = mtable' + , envNameMap = M.map (fmap substitute) names + } + + substitute v = + fromMaybe v $ M.lookup v substs + + substituteInMap f m = + let (ks, vs) = unzip $ M.toList m + in M.fromList $ + zip (map (\k -> fromMaybe k $ M.lookup k substs) ks) + (map f vs) + + substituteInBinding (BoundV ps t) = + BoundV (map substituteInTypeParam ps) (substituteInType t) + + substituteInMod (ModEnv env) = + ModEnv $ substituteInEnv env + substituteInMod (ModFun funsig) = + ModFun $ substituteInFunSig funsig + + substituteInFunSig (FunSig abs mod mty) = + FunSig (M.mapKeys (fmap substitute) abs) + (substituteInMod mod) (substituteInMTy substs mty) + + substituteInTypeBinding (TypeAbbr l ps t) = + TypeAbbr l (map substituteInTypeParam ps) $ substituteInType t + + substituteInTypeParam (TypeParamDim p loc) = + TypeParamDim (substitute p) loc + substituteInTypeParam (TypeParamType l p loc) = + TypeParamType l (substitute p) loc + + substituteInType :: StructType -> StructType + substituteInType (TypeVar () u (TypeName qs v) targs) = + TypeVar () u (TypeName (map substitute qs) $ substitute v) $ map substituteInTypeArg targs + substituteInType (Prim t) = + Prim t + substituteInType (Record ts) = + Record $ fmap substituteInType ts + substituteInType (Array (ArrayPrimElem t ()) shape u) = + Array (ArrayPrimElem t ()) (substituteInShape shape) u + substituteInType (Array (ArrayPolyElem (TypeName qs v) targs ()) shape u) = + Array (ArrayPolyElem + (TypeName (map substitute qs) $ substitute v) + (map substituteInTypeArg targs) ()) + (substituteInShape shape) u + substituteInType (Array (ArrayRecordElem ts) shape u) = + let ts' = fmap (substituteInType . fst . recordArrayElemToType) ts + in case arrayOf (Record ts') (substituteInShape shape) u of + Just t' -> t' + _ -> error "substituteInType: Cannot create array after substitution." + substituteInType (Arrow als v t1 t2) = + Arrow als v (substituteInType t1) (substituteInType t2) + + substituteInShape (ShapeDecl ds) = + ShapeDecl $ map substituteInDim ds + substituteInDim (NamedDim (QualName qs v)) = + NamedDim $ QualName (map substitute qs) $ substitute v + substituteInDim d = d + + substituteInTypeArg (TypeArgDim (NamedDim (QualName qs v)) loc) = + TypeArgDim (NamedDim $ QualName (map substitute qs) $ substitute v) loc + substituteInTypeArg (TypeArgDim (ConstDim x) loc) = + TypeArgDim (ConstDim x) loc + substituteInTypeArg (TypeArgDim AnyDim loc) = + TypeArgDim AnyDim loc + substituteInTypeArg (TypeArgType t loc) = + TypeArgType (substituteInType t) loc + +mtyTypeAbbrs :: MTy -> M.Map VName TypeBinding +mtyTypeAbbrs (MTy _ mod) = modTypeAbbrs mod + +modTypeAbbrs :: Mod -> M.Map VName TypeBinding +modTypeAbbrs (ModEnv env) = + envTypeAbbrs env +modTypeAbbrs (ModFun (FunSig _ mod mty)) = + modTypeAbbrs mod <> mtyTypeAbbrs mty + +envTypeAbbrs :: Env -> M.Map VName TypeBinding +envTypeAbbrs env = + envTypeTable env <> + (mconcat . map modTypeAbbrs . M.elems . envModTable) env + +-- | Refine the given type name in the given env. +refineEnv :: SrcLoc -> TySet -> Env -> QualName Name -> [TypeParam] -> StructType + -> TypeM (QualName VName, TySet, Env) +refineEnv loc tset env tname ps t + | Just (tname', TypeAbbr l cur_ps (TypeVar () _ (TypeName qs v) _)) <- + findTypeDef tname (ModEnv env), + QualName (qualQuals tname') v `M.member` tset = + if paramsMatch cur_ps ps then + return (tname', + QualName qs v `M.delete` tset, + substituteTypesInEnv + (M.fromList [(qualLeaf tname', + TypeSub $ TypeAbbr l cur_ps t), + (v, TypeSub $ TypeAbbr l ps t)]) + env) + else throwError $ TypeError loc $ "Cannot refine a type having " <> + tpMsg ps <> " with a type having " <> tpMsg cur_ps <> "." + | otherwise = + throwError $ TypeError loc $ + pretty tname ++ " is not an abstract type in the module type." + where tpMsg [] = "no type parameters" + tpMsg xs = "type parameters " <> unwords (map pretty xs) + +paramsMatch :: [TypeParam] -> [TypeParam] -> Bool +paramsMatch ps1 ps2 = length ps1 == length ps2 && all match (zip ps1 ps2) + where match (TypeParamType l1 _ _, TypeParamType l2 _ _) = l1 <= l2 + match (TypeParamDim _ _, TypeParamDim _ _) = True + match _ = False diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Monad.hs @@ -0,0 +1,382 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts, TupleSections #-} +-- | Main monad in which the type checker runs, as well as ancillary +-- data definitions. +module Language.Futhark.TypeChecker.Monad + ( TypeM + , runTypeM + , askEnv + , askRootEnv + , localTmpEnv + , checkQualNameWithEnv + , bindSpaced + , qualifyTypeVars + , getType + + , TypeError(..) + , unexpectedType + , undefinedType + , unappliedFunctor + , unknownVariableError + , underscoreUse + , functionIsNotValue + + , BreadCrumb(..) + , MonadBreadCrumbs(..) + , typeError + + , MonadTypeChecker(..) + , checkName + , badOnLeft + + , module Language.Futhark.Warnings + + , Env(..) + , TySet + , FunSig(..) + , ImportTable + , NameMap + , BoundV(..) + , Mod(..) + , TypeBinding(..) + , MTy(..) + + , anySignedType + , anyUnsignedType + , anyIntType + , anyFloatType + , anyNumberType + , anyPrimType + + , Namespace(..) + , intrinsicsNameMap + , topLevelNameMap + , ppSpace + ) +where + +import Control.Monad.Except +import Control.Monad.Reader +import Control.Monad.Writer +import Control.Monad.State +import Control.Monad.RWS.Strict +import Control.Monad.Identity +import Data.List +import Data.Loc +import Data.Maybe +import Data.Either +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Prelude hiding (mapM, mod) + +import Language.Futhark +import Language.Futhark.Semantic +import Language.Futhark.Traversals +import Language.Futhark.Warnings +import Futhark.FreshNames hiding (newName) +import qualified Futhark.FreshNames + +-- | Information about an error during type checking. The 'Show' +-- instance for this type produces a human-readable description. +data TypeError = TypeError SrcLoc String + +unexpectedType :: MonadTypeChecker m => SrcLoc -> TypeBase () () -> [TypeBase () ()] -> m a +unexpectedType loc _ [] = + throwError $ TypeError loc $ + "Type of expression at " ++ locStr loc ++ + "cannot have any type - possibly a bug in the type checker." +unexpectedType loc t ts = + throwError $ TypeError loc $ + "Type of expression at " ++ locStr loc ++ " must be one of " ++ + intercalate ", " (map pretty ts) ++ ", but is " ++ + pretty t ++ "." + +undefinedType :: MonadTypeChecker m => SrcLoc -> QualName Name -> m a +undefinedType loc name = + throwError $ TypeError loc $ + "Unknown type " ++ pretty name ++ "." + +functionIsNotValue :: MonadTypeChecker m => SrcLoc -> QualName Name -> m a +functionIsNotValue loc name = + throwError $ TypeError loc $ + "Attempt to use function " ++ pretty name ++ " as value at " ++ locStr loc ++ "." + +unappliedFunctor :: MonadTypeChecker m => SrcLoc -> m a +unappliedFunctor loc = + throwError $ TypeError loc "Cannot have parametric module here." + +unknownVariableError :: MonadTypeChecker m => + Namespace -> QualName Name -> SrcLoc -> m a +unknownVariableError space name loc = + throwError $ TypeError loc $ + "Unknown " ++ ppSpace space ++ " " ++ pretty name + +underscoreUse :: MonadTypeChecker m => + SrcLoc -> QualName Name -> m a +underscoreUse loc name = + throwError $ TypeError loc $ + "Use of " ++ pretty name ++ ": variables prefixed with underscore must not be accessed." + +instance Show TypeError where + show (TypeError pos msg) = + "Error at " ++ locStr pos ++ ":\n" ++ msg + +type ImportTable = M.Map String Env + +data Context = Context { contextEnv :: Env + , contextRootEnv :: Env + , contextImportTable :: ImportTable + , contextImportName :: ImportName + } + +-- | The type checker runs in this monad. +newtype TypeM a = TypeM (RWST + Context -- Reader + Warnings -- Writer + VNameSource -- State + (Except TypeError) -- Inner monad + a) + deriving (Monad, Functor, Applicative, + MonadReader Context, + MonadWriter Warnings, + MonadState VNameSource, + MonadError TypeError) + +runTypeM :: Env -> ImportTable -> ImportName -> VNameSource + -> TypeM a + -> Either TypeError (a, Warnings, VNameSource) +runTypeM env imports fpath src (TypeM m) = do + (x, src', ws) <- runExcept $ runRWST m (Context env env imports fpath) src + return (x, ws, src') + +askEnv, askRootEnv :: TypeM Env +askEnv = asks contextEnv +askRootEnv = asks contextRootEnv + +localTmpEnv :: Env -> TypeM a -> TypeM a +localTmpEnv env = local $ \ctx -> + ctx { contextEnv = env <> contextEnv ctx } + +-- | A piece of information that describes what process the type +-- checker currently performing. This is used to give better error +-- messages. +data BreadCrumb = MatchingTypes (TypeBase () ()) (TypeBase () ()) + | MatchingFields Name + +instance Show BreadCrumb where + show (MatchingTypes t1 t2) = + "When matching type\n" ++ indent (pretty t1) ++ + "\nwith\n" ++ indent (pretty t2) + where indent = intercalate "\n" . map (" "++) . lines + show (MatchingFields field) = + "When matching types of record field `" ++ pretty field ++ "`." + +-- | Tracking breadcrumbs to give a kind of "stack trace" in errors. +class Monad m => MonadBreadCrumbs m where + breadCrumb :: BreadCrumb -> m a -> m a + breadCrumb _ m = m + + getBreadCrumbs :: m [BreadCrumb] + getBreadCrumbs = return [] + +typeError :: (MonadError TypeError m, MonadBreadCrumbs m) => + SrcLoc -> String -> m a +typeError loc s = do + bc <- getBreadCrumbs + let bc' | null bc = "" + | otherwise = "\n" ++ unlines (map show bc) + throwError $ TypeError loc $ s ++ bc' + +class MonadError TypeError m => MonadTypeChecker m where + warn :: SrcLoc -> String -> m () + + newName :: VName -> m VName + newID :: Name -> m VName + + bindNameMap :: NameMap -> m a -> m a + localEnv :: Env -> m a -> m a + + checkQualName :: Namespace -> QualName Name -> SrcLoc -> m (QualName VName) + + lookupType :: SrcLoc -> QualName Name -> m (QualName VName, [TypeParam], StructType, Liftedness) + lookupMod :: SrcLoc -> QualName Name -> m (QualName VName, Mod) + lookupMTy :: SrcLoc -> QualName Name -> m (QualName VName, MTy) + lookupImport :: SrcLoc -> FilePath -> m (FilePath, Env) + lookupVar :: SrcLoc -> QualName Name -> m (QualName VName, CompType) + +checkName :: MonadTypeChecker m => Namespace -> Name -> SrcLoc -> m VName +checkName space name loc = qualLeaf <$> checkQualName space (qualName name) loc + +bindSpaced :: MonadTypeChecker m => [(Namespace, Name)] -> m a -> m a +bindSpaced names body = do + names' <- mapM (newID . snd) names + let mapping = M.fromList (zip names $ map qualName names') + bindNameMap mapping body + +instance MonadTypeChecker TypeM where + warn loc problem = tell $ singleWarning loc problem + + newName s = do src <- get + let (s', src') = Futhark.FreshNames.newName src s + put src' + return s' + + newID s = newName $ VName s 0 + + bindNameMap m = local $ \ctx -> + let env = contextEnv ctx + in ctx { contextEnv = env { envNameMap = m <> envNameMap env } } + + localEnv env = local $ \ctx -> + let env' = env <> contextEnv ctx + in ctx { contextEnv = env', contextRootEnv = env' } + + checkQualName space name loc = snd <$> checkQualNameWithEnv space name loc + + lookupType loc qn = do + outer_env <- askRootEnv + (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Type qn loc + case M.lookup name $ envTypeTable scope of + Nothing -> undefinedType loc qn + Just (TypeAbbr l ps def) -> return (qn', ps, qualifyTypeVars outer_env mempty qs def, l) + + lookupMod loc qn = do + (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Term qn loc + case M.lookup name $ envModTable scope of + Nothing -> unknownVariableError Term qn loc + Just m -> return (qn', m) + + lookupMTy loc qn = do + (scope, qn'@(QualName _ name)) <- checkQualNameWithEnv Signature qn loc + (qn',) <$> maybe explode return (M.lookup name $ envSigTable scope) + where explode = unknownVariableError Signature qn loc + + lookupImport loc file = do + imports <- asks contextImportTable + my_path <- asks contextImportName + let canonical_import = includeToString $ mkImportFrom my_path file loc + case M.lookup canonical_import imports of + Nothing -> throwError $ TypeError loc $ + unlines ["Unknown import \"" ++ canonical_import ++ "\"", + "Known: " ++ intercalate ", " (M.keys imports)] + Just scope -> return (canonical_import, scope) + + lookupVar loc qn = do + outer_env <- askRootEnv + (env, qn'@(QualName qs name)) <- checkQualNameWithEnv Term qn loc + case M.lookup name $ envVtable env of + Nothing -> unknownVariableError Term qn loc + Just (BoundV _ t) + | "_" `isPrefixOf` baseString name -> underscoreUse loc qn + | otherwise -> + case getType t of + Left{} -> throwError $ TypeError loc $ + "Attempt to use function " ++ baseString name ++ " as value." + Right t' -> return (qn', removeShapeAnnotations $ fromStruct $ + qualifyTypeVars outer_env mempty qs t') + +-- | Extract from a type either a function type comprising a list of +-- parameter types and a return type, or a first-order type. +getType :: TypeBase dim as + -> Either ([(Maybe VName, TypeBase dim as)], TypeBase dim as) + (TypeBase dim as) +getType (Arrow _ v t1 t2) = + case getType t2 of + Left (ps, r) -> Left ((v, t1) : ps, r) + Right _ -> Left ([(v, t1)], t2) +getType t = Right t + +checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TypeM (Env, QualName VName) +checkQualNameWithEnv space qn@(QualName quals name) loc = do + env <- askEnv + descend env quals + where descend scope [] + | Just name' <- M.lookup (space, name) $ envNameMap scope = + return (scope, name') + | otherwise = + unknownVariableError space qn loc + + descend scope (q:qs) + | Just (QualName _ q') <- M.lookup (Term, q) $ envNameMap scope, + Just res <- M.lookup q' $ envModTable scope = + case res of + ModEnv q_scope -> do + (scope', QualName qs' name') <- descend q_scope qs + return (scope', QualName (q':qs') name') + ModFun{} -> unappliedFunctor loc + | otherwise = + unknownVariableError space qn loc + +-- Try to prepend qualifiers to the type names such that they +-- represent how to access the type in some scope. +qualifyTypeVars :: ASTMappable t => Env -> [VName] -> [VName] -> t -> t +qualifyTypeVars outer_env except qs = runIdentity . astMap mapper + where mapper = ASTMapper { mapOnExp = pure + , mapOnName = pure + , mapOnQualName = pure . qual + , mapOnType = pure + , mapOnCompType = pure + , mapOnStructType = pure + , mapOnPatternType = pure + } + qual (QualName orig_qs name) + | name `elem` except || + reachable orig_qs name outer_env = QualName orig_qs name + | otherwise = QualName (qs<>orig_qs) name + + reachable [] name env = + isJust $ find matches $ M.elems (envTypeTable env) + where matches (TypeAbbr _ [] (TypeVar _ _ (TypeName x_qs name') [])) = + null x_qs && name == name' + matches _ = False + + reachable (q:qs') name env + | Just (ModEnv env') <- M.lookup q $ envModTable env = + reachable qs' name env' + | otherwise = False + +badOnLeft :: MonadTypeChecker m => Either TypeError a -> m a +badOnLeft = either throwError return + +anySignedType :: [PrimType] +anySignedType = map Signed [minBound .. maxBound] + +anyUnsignedType :: [PrimType] +anyUnsignedType = map Unsigned [minBound .. maxBound] + +anyIntType :: [PrimType] +anyIntType = anySignedType ++ anyUnsignedType + +anyFloatType :: [PrimType] +anyFloatType = map FloatType [minBound .. maxBound] + +anyNumberType :: [PrimType] +anyNumberType = anyIntType ++ anyFloatType + +anyPrimType :: [PrimType] +anyPrimType = Bool : anyIntType ++ anyFloatType + +--- Name handling + +ppSpace :: Namespace -> String +ppSpace Term = "name" +ppSpace Type = "type" +ppSpace Signature = "module type" + +intrinsicsNameMap :: NameMap +intrinsicsNameMap = M.fromList $ map mapping $ M.toList intrinsics + where mapping (v, IntrinsicType{}) = ((Type, baseName v), QualName [mod] v) + mapping (v, _) = ((Term, baseName v), QualName [mod] v) + mod = VName (nameFromString "intrinsics") 0 + +topLevelNameMap :: NameMap +topLevelNameMap = M.filterWithKey (\k _ -> atTopLevel k) intrinsicsNameMap + where atTopLevel :: (Namespace, Name) -> Bool + atTopLevel (Type, _) = True + atTopLevel (Term, v) = v `S.member` (type_names <> binop_names <> unop_names <> fun_names) + where type_names = S.fromList $ map (nameFromString . pretty) anyPrimType + binop_names = S.fromList $ map (nameFromString . pretty) + [minBound..(maxBound::BinOp)] + unop_names = S.fromList $ map nameFromString ["~", "!"] + fun_names = S.fromList $ map nameFromString ["shape"] + atTopLevel _ = False diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -0,0 +1,1664 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleContexts #-} +-- | Facilities for type-checking Futhark terms. Checking a term +-- requires a little more context to track uniqueness and such. +-- +-- Type inference is implemented through a variation of +-- Hindley-Milner. The main complication is supporting the rich +-- number of built-in language constructs, as well as uniqueness +-- types. This is mostly done in an ad hoc way, and many programs +-- will require the programmer to fall back on type annotations. +module Language.Futhark.TypeChecker.Terms + ( checkOneExp + , checkFunDef + ) +where + +import Control.Monad.Except +import Control.Monad.State +import Control.Monad.RWS +import qualified Control.Monad.Fail as Fail +import Data.List +import Data.Loc +import Data.Maybe +import qualified Data.Semigroup as Sem +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Prelude hiding (mod) + +import Language.Futhark +import Language.Futhark.Traversals +import Language.Futhark.TypeChecker.Monad hiding (BoundV, checkQualNameWithEnv) +import Language.Futhark.TypeChecker.Types hiding (checkTypeDecl) +import Language.Futhark.TypeChecker.Unify +import qualified Language.Futhark.TypeChecker.Types as Types +import qualified Language.Futhark.TypeChecker.Monad as TypeM +import Futhark.Util.Pretty (Pretty) + +--- Uniqueness + +data Usage = Consumed SrcLoc + | Observed SrcLoc + deriving (Eq, Ord, Show) + +-- | The consumption set is a Maybe so we can distinguish whether a +-- consumption took place, but the variable went out of scope since, +-- or no consumption at all took place. +data Occurence = Occurence { observed :: Names + , consumed :: Maybe Names + , location :: SrcLoc + } + deriving (Eq, Show) + +instance Located Occurence where + locOf = locOf . location + +observation :: Names -> SrcLoc -> Occurence +observation = flip Occurence Nothing + +consumption :: Names -> SrcLoc -> Occurence +consumption = Occurence S.empty . Just + +-- | A null occurence is one that we can remove without affecting +-- anything. +nullOccurence :: Occurence -> Bool +nullOccurence occ = S.null (observed occ) && isNothing (consumed occ) + +-- | A seminull occurence is one that does not contain references to +-- any variables in scope. The big difference is that a seminull +-- occurence may denote a consumption, as long as the array that was +-- consumed is now out of scope. +seminullOccurence :: Occurence -> Bool +seminullOccurence occ = S.null (observed occ) && maybe True S.null (consumed occ) + +type Occurences = [Occurence] + +type UsageMap = M.Map VName [Usage] + +usageMap :: Occurences -> UsageMap +usageMap = foldl comb M.empty + where comb m (Occurence obs cons loc) = + let m' = S.foldl' (ins $ Observed loc) m obs + in S.foldl' (ins $ Consumed loc) m' $ fromMaybe mempty cons + ins v m k = M.insertWith (++) k [v] m + +combineOccurences :: MonadTypeChecker m => VName -> Usage -> Usage -> m Usage +combineOccurences _ (Observed loc) (Observed _) = return $ Observed loc +combineOccurences name (Consumed wloc) (Observed rloc) = + useAfterConsume (baseName name) rloc wloc +combineOccurences name (Observed rloc) (Consumed wloc) = + useAfterConsume (baseName name) rloc wloc +combineOccurences name (Consumed loc1) (Consumed loc2) = + consumeAfterConsume (baseName name) (max loc1 loc2) (min loc1 loc2) + +checkOccurences :: MonadTypeChecker m => Occurences -> m () +checkOccurences = void . M.traverseWithKey comb . usageMap + where comb _ [] = return () + comb name (u:us) = foldM_ (combineOccurences name) u us + +allObserved :: Occurences -> Names +allObserved = S.unions . map observed + +allConsumed :: Occurences -> Names +allConsumed = S.unions . map (fromMaybe mempty . consumed) + +allOccuring :: Occurences -> Names +allOccuring occs = allConsumed occs <> allObserved occs + +anyConsumption :: Occurences -> Maybe Occurence +anyConsumption = find (isJust . consumed) + +seqOccurences :: Occurences -> Occurences -> Occurences +seqOccurences occurs1 occurs2 = + filter (not . nullOccurence) $ map filt occurs1 ++ occurs2 + where filt occ = + occ { observed = observed occ `S.difference` postcons } + postcons = allConsumed occurs2 + +altOccurences :: Occurences -> Occurences -> Occurences +altOccurences occurs1 occurs2 = + filter (not . nullOccurence) $ map filt1 occurs1 ++ map filt2 occurs2 + where filt1 occ = + occ { consumed = S.difference <$> consumed occ <*> pure cons2 + , observed = observed occ `S.difference` cons2 } + filt2 occ = + occ { consumed = consumed occ + , observed = observed occ `S.difference` cons1 } + cons1 = allConsumed occurs1 + cons2 = allConsumed occurs2 + +--- Scope management + +data ValBinding = BoundV [TypeParam] PatternType + -- ^ Aliases in parameters indicate the lexical + -- closure. + | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType) + | EqualityF + | OpaqueF + | WasConsumed SrcLoc + deriving (Show) + +-- | Type checking happens with access to this environment. The +-- tables will be extended during type-checking as bindings come into +-- scope. +data TermScope = TermScope { scopeVtable :: M.Map VName ValBinding + , scopeTypeTable :: M.Map VName TypeBinding + , scopeNameMap :: NameMap + , scopeBreadCrumbs :: [BreadCrumb] + -- ^ Most recent first. + } deriving (Show) + +instance Sem.Semigroup TermScope where + TermScope vt1 tt1 nt1 bc1 <> TermScope vt2 tt2 nt2 bc2 = + TermScope (vt2 `M.union` vt1) (tt2 `M.union` tt1) (nt2 `M.union` nt1) (bc1 <> bc2) + +instance Monoid TermScope where + mempty = TermScope mempty mempty mempty mempty + mappend = (Sem.<>) + +envToTermScope :: Env -> TermScope +envToTermScope env = TermScope vtable (envTypeTable env) (envNameMap env) mempty + where vtable = M.map valBinding $ envVtable env + valBinding (TypeM.BoundV tps v) = BoundV tps $ v `setAliases` mempty + +constraintTypeVars :: Constraints -> Names +constraintTypeVars = mconcat . map f . M.elems + where f (Constraint t _) = typeVars t + f _ = mempty + +overloadedTypeVars :: Constraints -> Names +overloadedTypeVars = mconcat . map f . M.elems + where f (HasFields fs _) = mconcat $ map typeVars $ M.elems fs + f _ = mempty + +-- | Get the type of an expression, with all type variables +-- substituted. Never call 'typeOf' directly (except in a few +-- carefully inspected locations)! +expType :: Exp -> TermTypeM CompType +expType = normaliseType . typeOf + +-- | The state is a set of constraints and a counter for generating +-- type names. This is distinct from the usual counter we use for +-- generating unique names, as these will be user-visible. +type TermTypeState = (Constraints, Int) + +newtype TermTypeM a = TermTypeM (RWST + TermScope + Occurences + TermTypeState + TypeM + a) + deriving (Monad, Functor, Applicative, + MonadReader TermScope, + MonadWriter Occurences, + MonadState TermTypeState, + MonadError TypeError) + +instance Fail.MonadFail TermTypeM where + fail = typeError noLoc . ("unknown failure (likely a bug): "++) + +instance MonadUnify TermTypeM where + getConstraints = gets fst + putConstraints x = modify $ \s -> (x, snd s) + + newTypeVar loc desc = do + i <- incCounter + v <- newID $ nameFromString $ desc ++ show i + modifyConstraints $ M.insert v $ NoConstraint Nothing loc + return $ TypeVar mempty Nonunique (typeName v) [] + +instance MonadBreadCrumbs TermTypeM where + breadCrumb bc = local $ \env -> + env { scopeBreadCrumbs = bc : scopeBreadCrumbs env } + getBreadCrumbs = asks scopeBreadCrumbs + +runTermTypeM :: TermTypeM a -> TypeM (a, Occurences) +runTermTypeM (TermTypeM m) = do + initial_scope <- (initialTermScope <>) <$> (envToTermScope <$> askEnv) + evalRWST m initial_scope (mempty, 0) + +liftTypeM :: TypeM a -> TermTypeM a +liftTypeM = TermTypeM . lift + +incCounter :: TermTypeM Int +incCounter = do (x, i) <- get + put (x, i+1) + return i + +initialTermScope :: TermScope +initialTermScope = TermScope initialVtable mempty topLevelNameMap mempty + where initialVtable = M.fromList $ mapMaybe addIntrinsicF $ M.toList intrinsics + + funF ts t = foldr (Arrow mempty Nothing . Prim) (Prim t) ts + + addIntrinsicF (name, IntrinsicMonoFun ts t) = + Just (name, BoundV [] $ funF ts t) + addIntrinsicF (name, IntrinsicOverloadedFun ts pts rts) = + Just (name, OverloadedF ts pts rts) + addIntrinsicF (name, IntrinsicPolyFun tvs pts rt) = + Just (name, BoundV tvs $ + fromStruct $ vacuousShapeAnnotations $ + Arrow mempty Nothing pts' rt) + where pts' = case pts of [pt] -> pt + _ -> tupleRecord pts + addIntrinsicF (name, IntrinsicEquality) = + Just (name, EqualityF) + addIntrinsicF (name, IntrinsicOpaque) = + Just (name, OpaqueF) + addIntrinsicF _ = Nothing + +instance MonadTypeChecker TermTypeM where + warn loc problem = liftTypeM $ warn loc problem + newName = liftTypeM . newName + newID = liftTypeM . newID + + checkQualName space name loc = snd <$> checkQualNameWithEnv space name loc + + bindNameMap m = local $ \scope -> + scope { scopeNameMap = m <> scopeNameMap scope } + + localEnv env (TermTypeM m) = do + cur_state <- get + cur_scope <- ask + let cur_scope' = + cur_scope { scopeNameMap = scopeNameMap cur_scope `M.difference` envNameMap env } + (x,new_state,occs) <- liftTypeM $ localTmpEnv env $ + runRWST m cur_scope' cur_state + tell occs + put new_state + return x + + lookupType loc qn = do + outer_env <- liftTypeM askRootEnv + (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Type qn loc + case M.lookup name $ scopeTypeTable scope of + Nothing -> undefinedType loc qn + Just (TypeAbbr l ps def) -> + return (qn', ps, qualifyTypeVars outer_env (map typeParamName ps) qs def, l) + + lookupMod loc name = liftTypeM $ TypeM.lookupMod loc name + lookupMTy loc name = liftTypeM $ TypeM.lookupMTy loc name + lookupImport loc name = liftTypeM $ TypeM.lookupImport loc name + + lookupVar loc qn = do + outer_env <- liftTypeM askRootEnv + (scope, qn'@(QualName qs name)) <- checkQualNameWithEnv Term qn loc + + t <- case M.lookup name $ scopeVtable scope of + Nothing -> unknownVariableError Term qn loc + + Just (WasConsumed wloc) -> useAfterConsume (baseName name) loc wloc + + Just (BoundV tparams t) + | "_" `isPrefixOf` baseString name -> underscoreUse loc qn + | otherwise -> do + (tnames, t') <- instantiateTypeScheme loc tparams t + let qual = qualifyTypeVars outer_env tnames qs + qual . removeShapeAnnotations <$> normaliseType t' + + Just OpaqueF -> do + argtype <- newTypeVar loc "t" + return $ Arrow mempty Nothing argtype argtype + + Just EqualityF -> do + argtype <- newTypeVar loc "t" + equalityType loc argtype + return $ Arrow mempty Nothing argtype $ + Arrow mempty Nothing argtype $ Prim Bool + + Just (OverloadedF ts pts rt) -> do + argtype <- newTypeVar loc "t" + mustBeOneOf ts loc argtype + let (pts', rt') = instOverloaded argtype pts rt + return $ fromStruct $ foldr (Arrow mempty Nothing) rt' pts' + + observe $ Ident name (Info t) loc + return (qn', t) + + where instOverloaded argtype pts rt = + (map (maybe (toStruct argtype) Prim) pts, + maybe (toStruct argtype) Prim rt) + +checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) +checkQualNameWithEnv space qn@(QualName [q] _) loc + | nameToString q == "intrinsics" = do + -- Check if we are referring to the magical intrinsics + -- module. + (_, QualName _ q') <- liftTypeM $ TypeM.checkQualNameWithEnv Term (qualName q) loc + if baseTag q' <= maxIntrinsicTag + then checkIntrinsic space qn loc + else checkReallyQualName space qn loc +checkQualNameWithEnv space qn@(QualName quals name) loc = do + scope <- ask + case quals of + [] | Just name' <- M.lookup (space, name) $ scopeNameMap scope -> + return (scope, name') + _ -> checkReallyQualName space qn loc + +checkIntrinsic :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) +checkIntrinsic space qn@(QualName _ name) loc + | Just v <- M.lookup (space, name) intrinsicsNameMap = do + scope <- ask + return (scope, v) + | otherwise = + unknownVariableError space qn loc + +checkReallyQualName :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) +checkReallyQualName space qn loc = do + (env, name') <- liftTypeM $ TypeM.checkQualNameWithEnv space qn loc + return (envToTermScope env, name') + +-- | Wrap 'checkTypeDecl' to also perform an observation of every size +-- in the type. +checkTypeDecl :: TypeDeclBase NoInfo Name -> TermTypeM (TypeDeclBase Info VName) +checkTypeDecl tdecl = do + (tdecl', _) <- Types.checkTypeDecl tdecl + mapM_ observeDim $ nestedDims $ unInfo $ expandedType tdecl' + return tdecl' + where observeDim (NamedDim v) = observe $ Ident (qualLeaf v) (Info $ Prim $ Signed Int32) noLoc + observeDim _ = return () + +-- | Instantiate a type scheme with fresh type variables for its type +-- parameters. Returns the names of the fresh type variables, the instance +-- list, and the instantiated type. +instantiateTypeScheme :: SrcLoc -> [TypeParam] -> PatternType + -> TermTypeM ([VName], PatternType) +instantiateTypeScheme loc tparams t = do + let tparams' = filter isTypeParam tparams + tnames = map typeParamName tparams' + (fresh_tnames, inst_list) <- unzip <$> mapM (instantiateTypeParam loc) tparams' + let substs = M.fromList $ zip tnames $ + map vacuousShapeAnnotations inst_list + t' = substTypesAny (`M.lookup` substs) t + return (fresh_tnames, t') + +-- | Create a new type name and insert it (unconstrained) in the +-- substitution map. +instantiateTypeParam :: Monoid as => SrcLoc -> TypeParam -> TermTypeM (VName, TypeBase dim as) +instantiateTypeParam loc tparam = do + i <- incCounter + v <- newID $ nameFromString $ baseString (typeParamName tparam) ++ show i + modifyConstraints $ M.insert v $ NoConstraint (Just l) loc + return (v, TypeVar mempty Nonunique (typeName v) []) + where l = case tparam of TypeParamType x _ _ -> x + _ -> Lifted + +newArrayType :: SrcLoc -> String -> Int -> TermTypeM (TypeBase () (), TypeBase () ()) +newArrayType loc desc r = do + v <- newID $ nameFromString desc + modifyConstraints $ M.insert v $ NoConstraint Nothing loc + return (Array (ArrayPolyElem (typeName v) [] ()) + (ShapeDecl $ replicate r ()) Nonunique, + TypeVar () Nonunique (typeName v) []) + +--- Errors + +useAfterConsume :: MonadTypeChecker m => Name -> SrcLoc -> SrcLoc -> m a +useAfterConsume name rloc wloc = + throwError $ TypeError rloc $ + "Variable " ++ pretty name ++ " used," ++ + "but previously consumed at " ++ locStr wloc ++ ". (Possibly through aliasing)" + +consumeAfterConsume :: MonadTypeChecker m => Name -> SrcLoc -> SrcLoc -> m a +consumeAfterConsume name loc1 loc2 = + throwError $ TypeError loc2 $ + "Variable " ++ pretty name ++ " previously consumed at " ++ locStr loc1 ++ "." + +badLetWithValue :: MonadTypeChecker m => SrcLoc -> m a +badLetWithValue loc = + throwError $ TypeError loc + "New value for elements in let-with shares data with source array. This is illegal, as it prevents in-place modification." + +returnAliased :: MonadTypeChecker m => Name -> Name -> SrcLoc -> m () +returnAliased fname name loc = + throwError $ TypeError loc $ + "Unique return value of function " ++ nameToString fname ++ + " is aliased to " ++ pretty name ++ ", which is not consumed." + +uniqueReturnAliased :: MonadTypeChecker m => Name -> SrcLoc -> m a +uniqueReturnAliased fname loc = + throwError $ TypeError loc $ + "A unique tuple element of return value of function " ++ + nameToString fname ++ " is aliased to some other tuple component." + +--- Basic checking + +-- | Determine if two types are identical, ignoring uniqueness. +-- Causes a 'TypeError' if they fail to match, and otherwise returns +-- one of them. +unifyExpTypes :: Exp -> Exp -> TermTypeM CompType +unifyExpTypes e1 e2 = do + e1_t <- expType e1 + e2_t <- expType e2 + unify (srclocOf e2) (toStruct e1_t) (toStruct e2_t) + return $ unifyTypeAliases e1_t e2_t + +-- | Assumes that the two types have already been unified. +unifyTypeAliases :: CompType -> CompType -> CompType +unifyTypeAliases t1 t2 = + case (t1, t2) of + (Array et1 shape1 u1, Array et2 _ u2) -> + Array (unifyArrayElems et1 et2) shape1 $ min u1 u2 + (Record f1, Record f2) -> + Record $ M.intersectionWith unifyTypeAliases f1 f2 + (TypeVar als1 u v targs1, TypeVar als2 _ _ targs2) -> + TypeVar (als1 <> als2) u v $ zipWith unifyTypeArg targs1 targs2 + _ -> t1 + where unifyArrayElems (ArrayPrimElem pt1 als1) (ArrayPrimElem _ als2) = + ArrayPrimElem pt1 $ als1 <> als2 + unifyArrayElems (ArrayPolyElem v targs1 als1) (ArrayPolyElem _ targs2 als2) = + ArrayPolyElem v (zipWith unifyTypeArg targs1 targs2) $ als1 <> als2 + unifyArrayElems (ArrayRecordElem fields1) (ArrayRecordElem fields2) = + ArrayRecordElem $ M.intersectionWith unifyRecordArray fields1 fields2 + unifyArrayElems x _ = x + + unifyRecordArray (RecordArrayElem at1) (RecordArrayElem at2) = + RecordArrayElem $ unifyArrayElems at1 at2 + unifyRecordArray (RecordArrayArrayElem at1 shape1 u1) (RecordArrayArrayElem at2 _ u2) = + RecordArrayArrayElem (unifyArrayElems at1 at2) shape1 $ min u1 u2 + unifyRecordArray x _ = x + + unifyTypeArg (TypeArgType t1' loc) (TypeArgType t2' _) = + TypeArgType (unifyTypeAliases t1' t2') loc + unifyTypeArg a _ = a + +--- General binding. + +data InferredType = NoneInferred + | Ascribed PatternType + + +checkPattern' :: UncheckedPattern -> InferredType + -> TermTypeM Pattern + +checkPattern' (PatternParens p loc) t = + PatternParens <$> checkPattern' p t <*> pure loc + +checkPattern' (Id name NoInfo loc) (Ascribed t) = do + name' <- checkName Term name loc + let t' = case t of Record{} -> t + _ -> t `addAliases` S.insert name' + return $ Id name' (Info t') loc +checkPattern' (Id name NoInfo loc) NoneInferred = do + name' <- checkName Term name loc + t <- newTypeVar loc "t" + return $ Id name' (Info t) loc + +checkPattern' (Wildcard _ loc) (Ascribed t) = + return $ Wildcard (Info $ t `setUniqueness` Nonunique) loc +checkPattern' (Wildcard NoInfo loc) NoneInferred = do + t <- newTypeVar loc "t" + return $ Wildcard (Info t) loc + +checkPattern' (TuplePattern ps loc) (Ascribed t) + | Just ts <- isTupleRecord t, length ts == length ps = + TuplePattern <$> zipWithM checkPattern' ps (map Ascribed ts) <*> pure loc +checkPattern' p@(TuplePattern ps loc) (Ascribed t) = do + ps_t <- replicateM (length ps) (newTypeVar loc "t") + unify loc (tupleRecord ps_t) $ toStructural t + t' <- normaliseType t + checkPattern' p $ Ascribed t' +checkPattern' (TuplePattern ps loc) NoneInferred = + TuplePattern <$> mapM (`checkPattern'` NoneInferred) ps <*> pure loc + +checkPattern' (RecordPattern p_fs loc) (Ascribed (Record t_fs)) + | sort (map fst p_fs) == sort (M.keys t_fs) = + RecordPattern . M.toList <$> check <*> pure loc + where check = traverse (uncurry checkPattern') $ M.intersectionWith (,) + (M.fromList p_fs) (fmap Ascribed t_fs) +checkPattern' p@(RecordPattern fields loc) (Ascribed t) = do + fields' <- traverse (const $ newTypeVar loc "t") $ M.fromList fields + + when (sort (M.keys fields') /= sort (map fst fields)) $ + typeError loc $ "Duplicate fields in record pattern " ++ pretty p + + unify loc (Record fields') $ toStructural t + t' <- normaliseType t + checkPattern' p $ Ascribed t' +checkPattern' (RecordPattern fs loc) NoneInferred = + RecordPattern . M.toList <$> traverse (`checkPattern'` NoneInferred) (M.fromList fs) <*> pure loc + +checkPattern' (PatternAscription p (TypeDecl t NoInfo) loc) maybe_outer_t = do + (t', st, _) <- checkTypeExp t + + let st' = fromStruct st + case maybe_outer_t of + Ascribed outer_t -> do + unify loc (toStructural st) (toStructural outer_t) + + -- We also have to make sure that uniqueness and shapes match. + -- This is done explicitly, because they are ignored by + -- unification. + st'' <- normaliseType st' + outer_t' <- normaliseType outer_t + case unifyTypesU unifyUniqueness st' outer_t' of + Just outer_t'' -> + PatternAscription <$> checkPattern' p (Ascribed outer_t'') <*> + pure (TypeDecl t' (Info st)) <*> pure loc + Nothing -> + typeError loc $ "Cannot match type `" ++ pretty outer_t' ++ "' with expected type `" ++ + pretty st'' ++ "'." + + NoneInferred -> + PatternAscription <$> checkPattern' p (Ascribed st') <*> + pure (TypeDecl t' (Info st)) <*> pure loc + where unifyUniqueness u1 u2 = if u2 `subuniqueOf` u1 then Just u1 else Nothing + +bindPatternNames :: PatternBase NoInfo Name -> TermTypeM a -> TermTypeM a +bindPatternNames = bindSpaced . map asTerm . S.toList . patIdentSet + where asTerm v = (Term, identName v) + +checkPattern :: UncheckedPattern -> InferredType -> (Pattern -> TermTypeM a) + -> TermTypeM a +checkPattern p t m = do + checkForDuplicateNames [p] + bindPatternNames p $ + m =<< checkPattern' p t + +binding :: [Ident] -> TermTypeM a -> TermTypeM a +binding bnds = check . local (`bindVars` bnds) + where bindVars :: TermScope -> [Ident] -> TermScope + bindVars = foldl bindVar + + bindVar :: TermScope -> Ident -> TermScope + bindVar scope (Ident name (Info tp) _) = + let inedges = S.toList $ aliases tp + update (BoundV tparams tp') + -- If 'name' is record-typed, don't alias the components + -- to 'name', because records have no identity beyond + -- their components. + | Record _ <- tp = BoundV tparams tp' + | otherwise = BoundV tparams (tp' `addAliases` S.insert name) + update b = b + in scope { scopeVtable = M.insert name (BoundV [] $ vacuousShapeAnnotations tp) $ + adjustSeveral update inedges $ + scopeVtable scope + } + + adjustSeveral f = flip $ foldl $ flip $ M.adjust f + + -- Check whether the bound variables have been used correctly + -- within their scope. + check m = do + (a, usages) <- collectBindingsOccurences m + checkOccurences usages + + mapM_ (checkIfUsed usages) bnds + + return a + + -- Collect and remove all occurences in @bnds@. This relies + -- on the fact that no variables shadow any other. + collectBindingsOccurences m = pass $ do + (x, usage) <- listen m + let (relevant, rest) = split usage + return ((x, relevant), const rest) + where split = unzip . + map (\occ -> + let (obs1, obs2) = divide $ observed occ + occ_cons = divide <$> consumed occ + con1 = fst <$> occ_cons + con2 = snd <$> occ_cons + in (occ { observed = obs1, consumed = con1 }, + occ { observed = obs2, consumed = con2 })) + names = S.fromList $ map identName bnds + divide s = (s `S.intersection` names, s `S.difference` names) + +bindingTypes :: [(VName, (TypeBinding, Constraint))] -> TermTypeM a -> TermTypeM a +bindingTypes types m = do + modifyConstraints (<>M.map snd (M.fromList types)) + local extend m + where extend scope = scope { + scopeTypeTable = M.map fst (M.fromList types) <> scopeTypeTable scope + } + +bindingTypeParams :: [TypeParam] -> TermTypeM a -> TermTypeM a +bindingTypeParams tparams = binding (mapMaybe typeParamIdent tparams) . + bindingTypes (mapMaybe typeParamType tparams) + where typeParamType (TypeParamType l v loc) = + Just (v, (TypeAbbr l [] (TypeVar () Nonunique (typeName v) []), + ParamType l loc)) + typeParamType TypeParamDim{} = + Nothing + +typeParamIdent :: TypeParam -> Maybe Ident +typeParamIdent (TypeParamDim v loc) = + Just $ Ident v (Info (Prim (Signed Int32))) loc +typeParamIdent _ = Nothing + +bindingIdent :: IdentBase NoInfo Name -> CompType -> (Ident -> TermTypeM a) + -> TermTypeM a +bindingIdent (Ident v NoInfo vloc) t m = + bindSpaced [(Term, v)] $ do + v' <- checkName Term v vloc + let ident = Ident v' (Info t) vloc + binding [ident] $ m ident + +bindingPatternGroup :: [UncheckedTypeParam] + -> [(UncheckedPattern, InferredType)] + -> ([TypeParam] -> [Pattern] -> TermTypeM a) -> TermTypeM a +bindingPatternGroup tps orig_ps m = do + checkForDuplicateNames $ map fst orig_ps + checkTypeParams tps $ \tps' -> bindingTypeParams tps' $ do + let descend ps' ((p,t):ps) = + checkPattern p t $ \p' -> + binding (S.toList $ patIdentSet p') $ descend (p':ps') ps + descend ps' [] = do + -- Perform an observation of every type parameter. This + -- prevents unused-name warnings for otherwise unused + -- dimensions. + mapM_ observe $ mapMaybe typeParamIdent tps' + let ps'' = reverse ps' + checkShapeParamUses tps' ps'' + + m tps' ps'' + + descend [] orig_ps + +bindingPattern :: [UncheckedTypeParam] + -> PatternBase NoInfo Name -> InferredType + -> ([TypeParam] -> Pattern -> TermTypeM a) -> TermTypeM a +bindingPattern tps p t m = do + checkForDuplicateNames [p] + checkTypeParams tps $ \tps' -> bindingTypeParams tps' $ + checkPattern p t $ \p' -> binding (S.toList $ patIdentSet p') $ do + -- Perform an observation of every declared dimension. This + -- prevents unused-name warnings for otherwise unused dimensions. + mapM_ observe $ patternDims p' + checkShapeParamUses tps' [p'] + + m tps' p' + +-- | Ensure that every shape parameter is used in positive position at +-- least once before being used in negative position. +checkShapeParamUses :: [TypeParam] -> [Pattern] -> TermTypeM () +checkShapeParamUses tps ps = do + pos_uses <- foldM checkShapePositions [] ps + mapM_ (checkUsed pos_uses) tps + where checkShapePositions pos_uses p = do + let (pos, neg) = patternUses p + pos_uses' = pos <> pos_uses + forM_ neg (\pv -> unless (pv `elem` pos_uses') $ + typeError (srclocOf p) $ "Shape parameter " ++ + pretty (baseName pv) ++ " must first be given in " ++ + "a positive position (non-functional parameter).") + return pos_uses' + checkUsed uses (TypeParamDim pv loc) + | pv `elem` uses = return () + | otherwise = + typeError loc $ "Size parameter " ++ + pretty (baseName pv) ++ " not used in any value parameters." + checkUsed _ _ = return () + +-- | Return the shapes used in a given pattern in postive and negative +-- position, respectively. +patternUses :: Pattern -> ([VName], [VName]) +patternUses Id{} = mempty +patternUses Wildcard{} = mempty +patternUses (PatternParens p _) = patternUses p +patternUses (TuplePattern ps _) = foldMap patternUses ps +patternUses (RecordPattern fs _) = foldMap (patternUses . snd) fs +patternUses (PatternAscription p (TypeDecl declte _) _) = + patternUses p <> typeExpUses declte + where typeExpUses (TEVar _ _) = mempty + typeExpUses (TETuple tes _) = foldMap typeExpUses tes + typeExpUses (TERecord fs _) = foldMap (typeExpUses . snd) fs + typeExpUses (TEArray te d _) = typeExpUses te <> dimDeclUses d + typeExpUses (TEUnique te _) = typeExpUses te + typeExpUses (TEApply te targ _) = typeExpUses te <> typeArgUses targ + typeExpUses (TEArrow _ t1 t2 _) = + let (pos, neg) = typeExpUses t1 <> typeExpUses t2 + in (mempty, pos <> neg) + typeArgUses (TypeArgExpDim d _) = dimDeclUses d + typeArgUses (TypeArgExpType te) = typeExpUses te + + dimDeclUses (NamedDim v) = ([qualLeaf v], []) + dimDeclUses _ = mempty + +noTypeParamsPermitted :: [UncheckedTypeParam] -> TermTypeM () +noTypeParamsPermitted ps = + case mapMaybe typeParamLoc ps of + loc:_ -> typeError loc "Type parameters are not permitted here." + [] -> return () + where typeParamLoc (TypeParamDim _ _) = Nothing + typeParamLoc tparam = Just $ srclocOf tparam + +patternDims :: Pattern -> [Ident] +patternDims (PatternParens p _) = patternDims p +patternDims (TuplePattern pats _) = concatMap patternDims pats +patternDims (PatternAscription p (TypeDecl _ (Info t)) _) = + patternDims p <> mapMaybe (dimIdent (srclocOf p)) (nestedDims t) + where dimIdent _ AnyDim = Nothing + dimIdent _ (ConstDim _) = Nothing + dimIdent _ NamedDim{} = Nothing +patternDims _ = [] + +--- Main checkers + +-- | @require ts e@ causes a 'TypeError' if @expType e@ is not one of +-- the types in @ts@. Otherwise, simply returns @e@. +require :: [PrimType] -> Exp -> TermTypeM Exp +require ts e = do mustBeOneOf ts (srclocOf e) . toStruct =<< expType e + return e + +unifies :: TypeBase () () -> Exp -> TermTypeM Exp +unifies t e = do + unify (srclocOf e) t =<< toStruct <$> expType e + return e + +checkExp :: UncheckedExp -> TermTypeM Exp + +checkExp (Literal val loc) = + return $ Literal val loc + +checkExp (IntLit val NoInfo loc) = do + t <- newTypeVar loc "t" + mustBeOneOf anyNumberType loc t + return $ IntLit val (Info t) loc + +checkExp (FloatLit val NoInfo loc) = do + t <- newTypeVar loc "t" + mustBeOneOf anyFloatType loc t + return $ FloatLit val (Info t) loc + +checkExp (TupLit es loc) = + TupLit <$> mapM checkExp es <*> pure loc + +checkExp (RecordLit fs loc) = do + fs' <- evalStateT (mapM checkField fs) mempty + + return $ RecordLit fs' loc + where checkField (RecordFieldExplicit f e rloc) = do + errIfAlreadySet f rloc + modify $ M.insert f rloc + RecordFieldExplicit f <$> lift (checkExp e) <*> pure rloc + checkField (RecordFieldImplicit name NoInfo rloc) = do + errIfAlreadySet name rloc + (QualName _ name', t) <- lift $ lookupVar rloc $ qualName name + modify $ M.insert name rloc + return $ RecordFieldImplicit name' (Info t) rloc + + errIfAlreadySet f rloc = do + maybe_sloc <- gets $ M.lookup f + case maybe_sloc of + Just sloc -> + lift $ typeError rloc $ "Field '" ++ pretty f ++ + " previously defined at " ++ locStr sloc ++ "." + Nothing -> return () + +checkExp (ArrayLit all_es _ loc) = + -- Construct the result type and unify all elements with it. We + -- only create a type variable for empty arrays; otherwise we use + -- the type of the first element. This significantly cuts down on + -- the number of type variables generated for pathologically large + -- multidimensional array literals. + case all_es of + [] -> do et <- newTypeVar loc "t" + t <- arrayOfM loc et (rank 1) Unique + return $ ArrayLit [] (Info t) loc + e:es -> do + e' <- checkExp e + et <- expType e' + es' <- mapM (unifies (toStructural et) <=< checkExp) es + et' <- normaliseType et + t <- arrayOfM loc et' (rank 1) Unique + return $ ArrayLit (e':es') (Info t) loc + +checkExp (Range start maybe_step end NoInfo loc) = do + start' <- require anyIntType =<< checkExp start + start_t <- toStructural <$> expType start' + maybe_step' <- case maybe_step of + Nothing -> return Nothing + Just step -> do + let warning = warn loc "First and second element of range are identical, this will produce an empty array." + case (start, step) of + (Literal x _, Literal y _) -> when (x == y) warning + (Var x_name _ _, Var y_name _ _) -> when (x_name == y_name) warning + _ -> return () + Just <$> (unifies start_t =<< checkExp step) + + end' <- case end of + DownToExclusive e -> DownToExclusive <$> (unifies start_t =<< checkExp e) + UpToExclusive e -> UpToExclusive <$> (unifies start_t =<< checkExp e) + ToInclusive e -> ToInclusive <$> (unifies start_t =<< checkExp e) + + t <- arrayOfM loc start_t (rank 1) Unique + + return $ Range start' maybe_step' end' (Info (t `setAliases` mempty)) loc + +checkExp (Ascript e decl loc) = do + decl' <- checkTypeDecl decl + e' <- checkExp e + t <- toStruct <$> expType e' + let decl_t = removeShapeAnnotations $ unInfo $ expandedType decl' + unify loc decl_t t + + -- We also have to make sure that uniqueness matches. This is done + -- explicitly, because uniqueness is ignored by unification. + t' <- normaliseType t + decl_t' <- normaliseType decl_t + unless (t' `subtypeOf` decl_t') $ + typeError loc $ "Type \"" ++ pretty t' ++ " is not a subtype of \"" ++ + pretty decl_t' ++ "\"." + + return $ Ascript e' decl' loc + +checkExp (BinOp op NoInfo (e1,_) (e2,_) NoInfo loc) = do + (op', ftype) <- lookupVar loc op + (e1', e1_arg) <- checkArg e1 + (e2', e2_arg) <- checkArg e2 + + (p1_t, rt) <- checkApply loc ftype e1_arg + (p2_t, rt') <- checkApply loc (removeShapeAnnotations rt) e2_arg + + return $ BinOp op' (Info (vacuousShapeAnnotations ftype)) + (e1', Info $ toStruct p1_t) (e2', Info $ toStruct p2_t) + (Info rt') loc + +checkExp (Project k e NoInfo loc) = do + e' <- checkExp e + t <- expType e' + kt <- mustHaveField loc k t + return $ Project k e' (Info kt) loc + +checkExp (If e1 e2 e3 _ loc) = + sequentially checkCond $ \e1' _ -> do + ((e2', e3'), dflow) <- tapOccurences $ checkExp e2 `alternative` checkExp e3 + brancht <- unifyExpTypes e2' e3' + let t' = addAliases brancht (`S.difference` allConsumed dflow) + zeroOrderType loc "returned from branch" t' + return $ If e1' e2' e3' (Info t') loc + where checkCond = do + e1' <- checkExp e1 + unify (srclocOf e1') (Prim Bool) . toStruct =<< expType e1' + return e1' + +checkExp (Parens e loc) = + Parens <$> checkExp e <*> pure loc + +checkExp (QualParens modname e loc) = do + (modname',mod) <- lookupMod loc modname + case mod of + ModEnv env -> localEnv (qualifyEnv modname' env) $ do + e' <- checkExp e + return $ QualParens modname' e' loc + ModFun{} -> + typeError loc $ "Module " ++ pretty modname ++ " is a parametric module." + where qualifyEnv modname' env = + env { envNameMap = M.map (qualify' modname') $ envNameMap env } + qualify' modname' (QualName qs name) = + QualName (qualQuals modname' ++ [qualLeaf modname'] ++ qs) name + +checkExp (Var qn NoInfo loc) = do + -- The qualifiers of a variable is divided into two parts: first a + -- possibly-empty sequence of module qualifiers, followed by a + -- possible-empty sequence of record field accesses. We use scope + -- information to perform the split, by taking qualifiers off the + -- end until we find a module. + + (qn', t, fields) <- findRootVar (qualQuals qn) (qualLeaf qn) + foldM checkField (Var qn' (Info (vacuousShapeAnnotations t)) loc) fields + where findRootVar qs name = + (whenFound <$> lookupVar loc (QualName qs name)) `catchError` notFound qs name + + whenFound (qn', t) = (qn', t, []) + + notFound qs name err + | null qs = throwError err + | otherwise = do + (qn', t, fields) <- findRootVar (init qs) (last qs) `catchError` + const (throwError err) + return (qn', t, fields++[name]) + + checkField e k = do + t <- expType e + kt <- mustHaveField loc k t + return $ Project k e (Info kt) loc + +checkExp (Negate arg loc) = do + arg' <- require anyNumberType =<< checkExp arg + return $ Negate arg' loc + +checkExp (Apply e1 e2 NoInfo NoInfo loc) = do + e1' <- checkExp e1 + (e2', arg) <- checkArg e2 + t <- expType e1' + (t1, rt) <- checkApply loc t arg + return $ Apply e1' e2' (Info $ diet t1) (Info rt) loc + +checkExp (LetPat tparams pat e body loc) = do + noTypeParamsPermitted tparams + sequentially (checkExp e) $ \e' e_occs -> do + -- Not technically an ascription, but we want the pattern to have + -- exactly the type of 'e'. + t <- expType e' + case anyConsumption e_occs of + Just c -> + let msg = "of value computed with consumption at " ++ locStr (location c) + in zeroOrderType loc msg t + _ -> return () + bindingPattern tparams pat (Ascribed $ vacuousShapeAnnotations t) $ \tparams' pat' -> do + body' <- checkExp body + return $ LetPat tparams' pat' e' body' loc + +checkExp (LetFun name (tparams, params, maybe_retdecl, NoInfo, e) body loc) = + sequentially (checkFunDef' (name, maybe_retdecl, tparams, params, e, loc)) $ + \(name', tparams', params', maybe_retdecl', rettype, e') closure -> do + + let ftype = foldr (uncurry (Arrow ()) . patternParam) rettype params' + entry = BoundV tparams' $ ftype `setAliases` allOccuring closure + bindF scope = scope { scopeVtable = M.insert name' entry $ scopeVtable scope + , scopeNameMap = M.insert (Term, name) (qualName name') $ + scopeNameMap scope } + body' <- local bindF $ checkExp body + + return $ LetFun name' (tparams', params', maybe_retdecl', Info rettype, e') body' loc + +checkExp (LetWith dest src idxes ve body pos) = do + (t, _) <- newArrayType (srclocOf src) "src" $ length idxes + let elemt = stripArray (length $ filter isFix idxes) t + sequentially (checkIdent src) $ \src' _ -> do + let src'' = Var (qualName $ identName src') + (vacuousShapeAnnotations <$> identType src') + (srclocOf src) + void $ unifies t src'' + + unless (unique $ unInfo $ identType src') $ + typeError pos $ "Source '" ++ pretty (identName src) ++ + "' has type " ++ pretty (unInfo $ identType src') ++ ", which is not unique" + + idxes' <- mapM checkDimIndex idxes + sequentially (unifies elemt =<< checkExp ve) $ \ve' _ -> do + ve_t <- expType ve' + when (identName src' `S.member` aliases ve_t) $ + badLetWithValue pos + + bindingIdent dest (unInfo (identType src') `setAliases` S.empty) $ \dest' -> do + body' <- consuming src' $ checkExp body + return $ LetWith dest' src' idxes' ve' body' pos + where isFix DimFix{} = True + isFix _ = False + +checkExp (Update src idxes ve loc) = do + (t, _) <- newArrayType (srclocOf src) "src" $ length idxes + let elemt = stripArray (length $ filter isFix idxes) t + sequentially (checkExp ve >>= unifies elemt) $ \ve' _ -> + sequentially (checkExp src >>= unifies t) $ \src' _ -> do + + idxes' <- mapM checkDimIndex idxes + + src_t <- expType src' + unless (unique src_t) $ + typeError loc $ "Source '" ++ pretty src ++ + "' has type " ++ pretty src_t ++ ", which is not unique" + + let src_als = aliases src_t + ve_t <- expType ve' + unless (S.null $ src_als `S.intersection` aliases ve_t) $ badLetWithValue loc + + consume loc src_als + return $ Update src' idxes' ve' loc + where isFix DimFix{} = True + isFix _ = False + +checkExp (RecordUpdate src fields ve NoInfo loc) = do + src' <- checkExp src + ve' <- checkExp ve + a <- expType src' + r <- foldM (flip $ mustHaveField loc) a fields + unify loc (toStruct r) . toStruct =<< expType ve' + return $ RecordUpdate src' fields ve' + (Info $ vacuousShapeAnnotations $ fromStruct a) loc + +checkExp (Index e idxes NoInfo loc) = do + (t, _) <- newArrayType (srclocOf e) "e" $ length idxes + e' <- unifies t =<< checkExp e + idxes' <- mapM checkDimIndex idxes + t' <- stripArray (length $ filter isFix idxes) <$> normaliseType (typeOf e') + return $ Index e' idxes' (Info t') loc + where isFix DimFix{} = True + isFix _ = False + +checkExp (Zip i e es NoInfo loc) = do + let checkInput inp = do (arr_t, _) <- newArrayType (srclocOf e) "e" (1+i) + unifies arr_t =<< checkExp inp + e' <- checkInput e + es' <- mapM checkInput es + + ts <- forM (e':es') $ \arr_e -> do + arr_e_t <- expType arr_e + case typeToRecordArrayElem' (aliases arr_e_t) =<< peelArray (i+1) arr_e_t of + Just t -> return t + Nothing -> typeError (srclocOf arr_e) $ + "Expected array with at least " ++ show (1+i) ++ + " dimensions, but got " ++ pretty arr_e_t ++ "." + + let u = mconcat $ map (uniqueness . typeOf) $ e':es' + t = Array (ArrayRecordElem $ M.fromList $ zip tupleFieldNames ts) + (rank (1+i)) u + return $ Zip i e' es' (Info t) loc + +checkExp (Unzip e _ loc) = do + e' <- checkExp e + e_t <- expType e' + case e_t of + Array (ArrayRecordElem fs) shape u + | Just ets <- map (componentType shape u) <$> areTupleFields fs -> + return $ Unzip e' (map Info ets) loc + t -> + typeError loc $ + "Argument to unzip is not an array of tuples, but " ++ + pretty t ++ "." + where componentType shape u et = + case et of + RecordArrayElem et' -> + Array et' shape u + RecordArrayArrayElem et' et_shape et_u -> + Array et' (shape <> et_shape) (u `max` et_u) + +checkExp (Unsafe e loc) = + Unsafe <$> checkExp e <*> pure loc + +checkExp (Assert e1 e2 NoInfo loc) = do + e1' <- require [Bool] =<< checkExp e1 + e2' <- checkExp e2 + return $ Assert e1' e2' (Info (pretty e1)) loc + +checkExp Map{} = error "Map nodes should not appear in source program" +checkExp Reduce{} = error "Reduce nodes should not appear in source program" +checkExp GenReduce{} = error "GenReduce nodes should not appear in source program" +checkExp Scan{} = error "Scan nodes should not appear in source program" +checkExp Filter{} = error "Filter nodes should not appear in source program" +checkExp Partition{} = error "Partition nodes should not appear in source program" +checkExp Stream{} = error "Stream nodes should not appear in source program" + +checkExp (Lambda tparams params body maybe_retdecl NoInfo loc) = + removeSeminullOccurences $ + bindingPatternGroup tparams (zip params $ repeat NoneInferred) $ \tparams' params' -> do + maybe_retdecl' <- traverse checkTypeDecl maybe_retdecl + (body', closure) <- tapOccurences $ noUnique $ + checkFunBody body (unInfo . expandedType <$> maybe_retdecl') loc + (maybe_retdecl'', rettype) <- case maybe_retdecl' of + Just retdecl'@(TypeDecl _ (Info st)) -> return (Just retdecl', st) + Nothing -> do + body_t <- expType body' + return (Nothing, vacuousShapeAnnotations $ toStruct body_t) + return $ Lambda tparams' params' body' maybe_retdecl'' + (Info (allOccuring closure, rettype)) loc + +checkExp (OpSection op _ loc) = do + (op', ftype) <- lookupVar loc op + return $ OpSection op' (Info (vacuousShapeAnnotations ftype)) loc + +checkExp (OpSectionLeft op _ e _ _ loc) = do + (op', ftype) <- lookupVar loc op + (e', e_arg) <- checkArg e + (t1, rt) <- checkApply loc ftype e_arg + case rt of + Arrow _ _ t2 rettype -> + return $ OpSectionLeft op' (Info (vacuousShapeAnnotations ftype)) e' + (Info $ toStruct t1, Info $ toStruct t2) (Info rettype) loc + _ -> typeError loc $ + "Operator section with invalid operator of type " ++ pretty ftype + +checkExp (OpSectionRight op _ e _ _ loc) = do + (op', ftype) <- lookupVar loc op + (e', e_arg) <- checkArg e + case ftype of + Arrow as1 m1 t1 (Arrow as2 m2 t2 ret) -> do + (t2', Arrow _ _ t1' rettype) <- + checkApply loc (Arrow as2 m2 t2 (Arrow as1 m1 t1 ret)) e_arg + return $ OpSectionRight op' (Info (vacuousShapeAnnotations ftype)) e' + (Info $ toStruct t1', Info $ toStruct t2') (Info rettype) loc + _ -> typeError loc $ + "Operator section with invalid operator of type " ++ pretty ftype + +checkExp (ProjectSection fields NoInfo loc) = do + a <- newTypeVar loc "a" + b <- foldM (flip $ mustHaveField loc) a fields + return $ ProjectSection fields (Info $ Arrow mempty Nothing a b) loc + +checkExp (IndexSection idxes NoInfo loc) = do + (t, _) <- newArrayType loc "e" (length idxes) + idxes' <- mapM checkDimIndex idxes + let t' = stripArray (length $ filter isFix idxes) t + return $ IndexSection idxes' (Info $ vacuousShapeAnnotations $ fromStruct $ + Arrow mempty Nothing t t') loc + where isFix DimFix{} = True + isFix _ = False + +checkExp (DoLoop tparams mergepat mergeexp form loopbody loc) = + sequentially (checkExp mergeexp) $ \mergeexp' _ -> do + + noTypeParamsPermitted tparams + + zeroOrderType (srclocOf mergeexp) "used as loop variable" (typeOf mergeexp') + + merge_t <- do + merge_t <- expType mergeexp' + return $ Ascribed $ vacuousShapeAnnotations $ merge_t `setAliases` mempty + + -- First we do a basic check of the loop body to figure out which of + -- the merge parameters are being consumed. For this, we first need + -- to check the merge pattern, which requires the (initial) merge + -- expression. + -- + -- Play a little with occurences to ensure it does not look like + -- none of the merge variables are being used. + ((tparams', mergepat', form', loopbody'), bodyflow) <- + case form of + For i uboundexp -> do + uboundexp' <- require anySignedType =<< checkExp uboundexp + bound_t <- expType uboundexp' + bindingIdent i bound_t $ \i' -> + noUnique $ bindingPattern tparams mergepat merge_t $ + \tparams' mergepat' -> onlySelfAliasing $ tapOccurences $ do + loopbody' <- checkExp loopbody + return (tparams', + mergepat', + For i' uboundexp', + loopbody') + + ForIn xpat e -> do + (arr_t, _) <- newArrayType (srclocOf e) "e" 1 + e' <- unifies arr_t =<< checkExp e + t <- expType e' + case t of + _ | Just t' <- peelArray 1 t -> + bindingPattern [] xpat (Ascribed $ vacuousShapeAnnotations t') $ \_ xpat' -> + noUnique $ bindingPattern tparams mergepat merge_t $ + \tparams' mergepat' -> onlySelfAliasing $ tapOccurences $ do + loopbody' <- checkExp loopbody + return (tparams', + mergepat', + ForIn xpat' e', + loopbody') + | otherwise -> + typeError (srclocOf e) $ + "Iteratee of a for-in loop must be an array, but expression has type " ++ pretty t + + While cond -> + noUnique $ bindingPattern tparams mergepat merge_t $ \tparams' mergepat' -> + onlySelfAliasing $ tapOccurences $ + sequentially (unifies (Prim Bool) =<< checkExp cond) $ \cond' _ -> do + loopbody' <- checkExp loopbody + return (tparams', + mergepat', + While cond', + loopbody') + + mergepat'' <- do + loop_t <- expType loopbody' + convergePattern mergepat' (allConsumed bodyflow) loop_t (srclocOf loopbody') + + let consumeMerge (Id _ (Info pt) ploc) mt + | unique pt = consume ploc $ aliases mt + consumeMerge (TuplePattern pats _) t | Just ts <- isTupleRecord t = + zipWithM_ consumeMerge pats ts + consumeMerge (PatternParens pat _) t = + consumeMerge pat t + consumeMerge (PatternAscription pat _ _) t = + consumeMerge pat t + consumeMerge _ _ = + return () + consumeMerge mergepat'' =<< expType mergeexp' + return $ DoLoop tparams' mergepat'' mergeexp' form' loopbody' loc + + where + convergePattern pat body_cons body_t body_loc = do + let consumed_merge = S.map identName (patIdentSet pat) `S.intersection` + body_cons + uniquePat (Wildcard (Info t) wloc) = + Wildcard (Info $ t `setUniqueness` Nonunique) wloc + uniquePat (PatternParens p ploc) = + PatternParens (uniquePat p) ploc + uniquePat (Id name (Info t) iloc) + | name `S.member` consumed_merge = + let t' = t `setUniqueness` Unique `setAliases` mempty + in Id name (Info t') iloc + | otherwise = + let t' = case t of Record{} -> t + _ -> t `setUniqueness` Nonunique + in Id name (Info t') iloc + uniquePat (TuplePattern pats ploc) = + TuplePattern (map uniquePat pats) ploc + uniquePat (RecordPattern fs ploc) = + RecordPattern (map (fmap uniquePat) fs) ploc + uniquePat (PatternAscription p t ploc) = + PatternAscription p t ploc + + -- Make the pattern unique where needed. + pat' = uniquePat pat + + -- Now check that the loop returned the right type. + unify body_loc (toStruct body_t) $ toStruct $ patternType pat' + body_t' <- normaliseType body_t + pat_t <- normaliseType $ patternType pat' + unless (body_t' `subtypeOf` pat_t) $ + unexpectedType body_loc + (toStructural body_t') + [toStructural pat_t] + + -- Check that the new values of consumed merge parameters do not + -- alias something bound outside the loop, AND that anything + -- returned for a unique merge parameter does not alias anything + -- else returned. + bound_outside <- asks $ S.fromList . M.keys . scopeVtable + let checkMergeReturn (Id pat_v (Info pat_v_t) _) t + | unique pat_v_t, + v:_ <- S.toList $ aliases t `S.intersection` bound_outside = + lift $ typeError loc $ "Loop return value corresponding to merge parameter " ++ + prettyName pat_v ++ " aliases " ++ prettyName v ++ "." + | otherwise = do + (cons,obs) <- get + unless (S.null $ aliases t `S.intersection` cons) $ + lift $ typeError loc $ "Loop return value for merge parameter " ++ + prettyName pat_v ++ " aliases other consumed merge parameter." + when (unique pat_v_t && + not (S.null (aliases t `S.intersection` (cons<>obs)))) $ + lift $ typeError loc $ "Loop return value for consuming merge parameter " ++ + prettyName pat_v ++ " aliases previously returned value." ++ show (aliases t, cons, obs) + if unique pat_v_t + then put (cons<>aliases t, obs) + else put (cons, obs<>aliases t) + checkMergeReturn (TuplePattern pats _) t | Just ts <- isTupleRecord t = + zipWithM_ checkMergeReturn pats ts + checkMergeReturn _ _ = + return () + (pat_cons, _) <- execStateT (checkMergeReturn pat' body_t') (mempty, mempty) + let body_cons' = body_cons <> pat_cons + if body_cons' == body_cons && patternType pat' == patternType pat + then return pat' + else convergePattern pat' body_cons' body_t' body_loc + +checkIdent :: IdentBase NoInfo Name -> TermTypeM Ident +checkIdent (Ident name _ loc) = do + (QualName _ name', vt) <- lookupVar loc (qualName name) + return $ Ident name' (Info vt) loc + +checkDimIndex :: DimIndexBase NoInfo Name -> TermTypeM DimIndex +checkDimIndex (DimFix i) = + DimFix <$> (unifies (Prim $ Signed Int32) =<< checkExp i) +checkDimIndex (DimSlice i j s) = + DimSlice + <$> maybe (return Nothing) (fmap Just . unifies (Prim $ Signed Int32) <=< checkExp) i + <*> maybe (return Nothing) (fmap Just . unifies (Prim $ Signed Int32) <=< checkExp) j + <*> maybe (return Nothing) (fmap Just . unifies (Prim $ Signed Int32) <=< checkExp) s + +sequentially :: TermTypeM a -> (a -> Occurences -> TermTypeM b) -> TermTypeM b +sequentially m1 m2 = do + (a, m1flow) <- collectOccurences m1 + (b, m2flow) <- collectOccurences $ m2 a m1flow + occur $ m1flow `seqOccurences` m2flow + return b + +type Arg = (CompType, Occurences, SrcLoc) + +argType :: Arg -> CompType +argType (t, _, _) = t + +checkArg :: UncheckedExp -> TermTypeM (Exp, Arg) +checkArg arg = do + (arg', dflow) <- collectOccurences $ checkExp arg + arg_t <- expType arg' + return (arg', (arg_t, dflow, srclocOf arg')) + +checkApply :: SrcLoc -> CompType -> Arg + -> TermTypeM (PatternType, PatternType) +checkApply loc (Arrow as _ tp1 tp2) (argtype, dflow, argloc) = do + unify argloc (toStruct tp1) (toStruct argtype) + + -- Perform substitutions of instantiated variables in the types. + tp1' <- normaliseType tp1 + tp2' <- normaliseType tp2 + argtype' <- normaliseType argtype + + occur [observation as loc] + + checkOccurences dflow + occurs <- consumeArg argloc argtype' (diet tp1') + + case anyConsumption dflow of + Just c -> + let msg = "of value computed with consumption at " ++ locStr (location c) + in zeroOrderType argloc msg tp1 + _ -> return () + + occur $ dflow `seqOccurences` occurs + + return (vacuousShapeAnnotations tp1', + vacuousShapeAnnotations $ + returnType (toStruct tp2') [diet tp1'] [argtype']) + +checkApply loc tfun@TypeVar{} arg = do + tv <- newTypeVar loc "b" + unify loc (toStruct tfun) $ Arrow mempty Nothing (toStruct (argType arg)) tv + constraints <- getConstraints + checkApply loc (applySubst (`lookupSubst` constraints) tfun) arg + +checkApply loc ftype arg = + typeError loc $ + "Attempt to apply an expression of type " ++ pretty ftype ++ + " to an argument of type " ++ pretty (argType arg) ++ "." + +consumeArg :: SrcLoc -> CompType -> Diet -> TermTypeM [Occurence] +consumeArg loc (Record ets) (RecordDiet ds) = + concat . M.elems <$> traverse (uncurry $ consumeArg loc) (M.intersectionWith (,) ets ds) +consumeArg loc (Array _ _ Nonunique) Consume = + typeError loc "Consuming parameter passed non-unique argument." +consumeArg loc (Arrow _ _ t1 _) (FuncDiet d _) + | not $ contravariantArg t1 d = + typeError loc "Non-consuming higher-order parameter passed consuming argument." + where contravariantArg (Array _ _ Unique) Observe = + False + contravariantArg (TypeVar _ Unique _ _) Observe = + False + contravariantArg (Record ets) (RecordDiet ds) = + and (M.intersectionWith contravariantArg ets ds) + contravariantArg (Arrow _ _ tp tr) (FuncDiet dp dr) = + contravariantArg tp dp && contravariantArg tr dr + contravariantArg _ _ = + True +consumeArg loc (Arrow _ _ _ t2) (FuncDiet _ pd) = + consumeArg loc t2 pd +consumeArg loc at Consume = return [consumption (aliases at) loc] +consumeArg loc at _ = return [observation (aliases at) loc] + +checkOneExp :: UncheckedExp -> TypeM Exp +checkOneExp e = fmap fst . runTermTypeM $ do + e' <- checkExp e + fixOverloadedTypes + updateExpTypes e' + +-- | Type-check a top-level (or module-level) function definition. +checkFunDef :: (Name, Maybe UncheckedTypeExp, + [UncheckedTypeParam], [UncheckedPattern], + UncheckedExp, SrcLoc) + -> TypeM (VName, [TypeParam], [Pattern], Maybe (TypeExp VName), StructType, Exp) +checkFunDef f = fmap fst $ runTermTypeM $ do + (fname, tparams, params, maybe_retdecl, rettype, body) <- checkFunDef' f + + -- Since this is a top-level function, we also resolve overloaded + -- types, using either defaults or complaining about ambiguities. + fixOverloadedTypes + + -- Then replace all inferred types in the body and parameters. + body' <- updateExpTypes body + params' <- updateExpTypes params + maybe_retdecl' <- traverse updateExpTypes maybe_retdecl + rettype' <- normaliseType rettype + + return (fname, tparams, params', maybe_retdecl', rettype', body') + +-- | This is "fixing" as in "setting them", not "correcting them". We +-- only make very conservative fixing. +fixOverloadedTypes :: TermTypeM () +fixOverloadedTypes = getConstraints >>= mapM_ fixOverloaded . M.toList + where fixOverloaded (v, Overloaded ots loc) + | Signed Int32 `elem` ots = do + unify loc (TypeVar () Nonunique (typeName v) []) $ Prim $ Signed Int32 + warn loc "Defaulting ambiguous type to `i32`." + | FloatType Float64 `elem` ots = do + unify loc (TypeVar () Nonunique (typeName v) []) $ Prim $ FloatType Float64 + warn loc "Defaulting ambiguous type to `f64`." + | otherwise = + typeError loc $ + unlines ["Type is ambiguous (could be one of " ++ intercalate ", " (map pretty ots) ++ ").", + "Add a type annotation to disambiguate the type."] + + fixOverloaded (_, NoConstraint _ loc) = + typeError loc $ unlines ["Type of expression is ambiguous.", + "Add a type annotation to disambiguate the type."] + + fixOverloaded (_, Equality loc) = + typeError loc $ unlines ["Type is ambiguous (must be equality type).", + "Add a type annotation to disambiguate the type."] + + fixOverloaded (_, HasFields fs loc) = + typeError loc $ unlines ["Type is ambiguous (must be record with fields {" ++ fs' ++ "}).", + "Add a type annotation to disambiguate the type."] + where fs' = intercalate ", " $ map field $ M.toList fs + field (l, t) = pretty l ++ ": " ++ pretty t + + fixOverloaded _ = return () + +checkFunDef' :: (Name, Maybe UncheckedTypeExp, + [UncheckedTypeParam], [UncheckedPattern], + UncheckedExp, SrcLoc) + -> TermTypeM (VName, [TypeParam], [Pattern], Maybe (TypeExp VName), StructType, Exp) +checkFunDef' (fname, maybe_retdecl, tparams, params, body, loc) = noUnique $ do + when (nameToString fname == "&&") $ + typeError loc "The && operator may not be redefined." + when (nameToString fname == "||") $ + typeError loc "The || operator may not be redefined." + + then_substs <- getConstraints + + bindingPatternGroup tparams (zip params $ repeat NoneInferred) $ \tparams' params' -> do + maybe_retdecl' <- traverse checkTypeExp maybe_retdecl + + body' <- checkFunBody body ((\(_,t,_)->t) <$> maybe_retdecl') (maybe loc srclocOf maybe_retdecl) + + params'' <- updateExpTypes params' + body_t <- expType body' + + (maybe_retdecl'', rettype) <- case maybe_retdecl' of + Just (retdecl', retdecl_type, _) -> do + let rettype_structural = toStructural retdecl_type + checkReturnAlias rettype_structural params'' body_t + return (Just retdecl', retdecl_type) + Nothing -> return (Nothing, vacuousShapeAnnotations $ toStruct body_t) + + -- Candidates for let-generalisation are those type variables that + --- + -- (1) were not known before we checked this function, and + -- + -- (2) are not used in the (new) definition of any type variables + -- known before we checked this function. + -- + -- (3) are not referenced from an overloaded type (for example, + -- are the element types of an incompletely resolved record type). + -- This is a bit more restrictive than I'd like, and SML for + -- example does not have this restriction. + now_substs <- getConstraints + let then_type_variables = S.fromList $ M.keys then_substs + then_type_constraints = constraintTypeVars $ + M.filterWithKey (\k _ -> k `S.member` then_type_variables) now_substs + keep_type_variables = then_type_variables <> + then_type_constraints <> + overloadedTypeVars now_substs + + let new_substs = M.filterWithKey (\k _ -> not (k `S.member` keep_type_variables)) now_substs + tparams'' <- closeOverTypes new_substs tparams' $ + rettype : map patternStructType params'' + + -- We keep those type variables that were not closed over by + -- let-generalisation. + modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` map typeParamName tparams'' + + bindSpaced [(Term, fname)] $ do + fname' <- checkName Term fname loc + return (fname', tparams'', params'', maybe_retdecl'', rettype, body') + + where -- | Check that unique return values do not alias a + -- non-consumed parameter. + checkReturnAlias rettp params' = + foldM_ (checkReturnAlias' params') S.empty . returnAliasing rettp + checkReturnAlias' params' seen (Unique, names) + | any (`S.member` S.map snd seen) $ S.toList names = + uniqueReturnAliased fname loc + | otherwise = do + notAliasingParam params' names + return $ seen `S.union` tag Unique names + checkReturnAlias' _ seen (Nonunique, names) + | any (`S.member` seen) $ S.toList $ tag Unique names = + uniqueReturnAliased fname loc + | otherwise = return $ seen `S.union` tag Nonunique names + + notAliasingParam params' names = + forM_ params' $ \p -> + let consumedNonunique p' = + not (unique $ unInfo $ identType p') && (identName p' `S.member` names) + in case find consumedNonunique $ S.toList $ patIdentSet p of + Just p' -> + returnAliased fname (baseName $ identName p') loc + Nothing -> + return () + + tag u = S.map $ \name -> (u, name) + + returnAliasing (Record ets1) (Record ets2) = + concat $ M.elems $ M.intersectionWith returnAliasing ets1 ets2 + returnAliasing expected got = [(uniqueness expected, aliases got)] + +checkFunBody :: ExpBase NoInfo Name + -> Maybe StructType + -> SrcLoc + -> TermTypeM Exp +checkFunBody body maybe_rettype _loc = do + body' <- checkExp body + + -- Unify body return type with return annotation, if one exists. + case maybe_rettype of + Just rettype -> do + let rettype_structural = toStructural rettype + void $ unifies rettype_structural body' + -- We also have to make sure that uniqueness matches. This is done + -- explicitly, because uniqueness is ignored by unification. + rettype_structural' <- normaliseType rettype_structural + body_t <- expType body' + unless (body_t `subtypeOf` rettype_structural') $ + typeError (srclocOf body) $ "Body type `" ++ pretty body_t ++ + "' is not a subtype of annotated type `" ++ + pretty rettype_structural' ++ "'." + + Nothing -> return () + + return body' + +-- | Find at all type variables in the given type that are covered by +-- the constraints, and produce type parameters that close over them. +-- Produce an error if the given list of type parameters is non-empty, +-- yet does not cover all type variables in the type. +closeOverTypes :: Constraints -> [TypeParam] -> [StructType] -> TermTypeM [TypeParam] +closeOverTypes substs tparams ts = + case tparams of + [] -> fmap catMaybes $ mapM closeOver $ M.toList substs' + _ -> do mapM_ checkClosedOver $ M.toList substs' + return tparams + where substs' = M.filterWithKey (\k _ -> k `S.member` visible) substs + visible = mconcat (map typeVars ts) + + checkClosedOver (k, v) + | not (canBeClosedOver v) || + k `elem` map typeParamName tparams = return () + | otherwise = + typeError (srclocOf v) $ + unlines ["Type variable `" ++ prettyName k ++ + "' not closed over by type parameters " ++ + intercalate ", " (map pretty tparams) ++ ".", + "This is usually because a parameter needs a type annotation."] + + canBeClosedOver NoConstraint{} = True + canBeClosedOver _ = False + + closeOver (k, NoConstraint (Just Unlifted) loc) = return $ Just $ TypeParamType Unlifted k loc + closeOver (k, NoConstraint _ loc) = return $ Just $ TypeParamType Lifted k loc + closeOver (_, _) = return Nothing + +--- Consumption + +occur :: Occurences -> TermTypeM () +occur = tell + +-- | Proclaim that we have made read-only use of the given variable. +observe :: Ident -> TermTypeM () +observe (Ident nm (Info t) loc) = + let als = nm `S.insert` aliases t + in occur [observation als loc] + +-- | Proclaim that we have written to the given variable. +consume :: SrcLoc -> Names -> TermTypeM () +consume loc als = occur [consumption als loc] + +-- | Proclaim that we have written to the given variable, and mark +-- accesses to it and all of its aliases as invalid inside the given +-- computation. +consuming :: Ident -> TermTypeM a -> TermTypeM a +consuming (Ident name (Info t) loc) m = do + consume loc $ name `S.insert` aliases t + local consume' m + where consume' scope = + scope { scopeVtable = M.insert name (WasConsumed loc) $ scopeVtable scope } + +collectOccurences :: TermTypeM a -> TermTypeM (a, Occurences) +collectOccurences m = pass $ do + (x, dataflow) <- listen m + return ((x, dataflow), const mempty) + +tapOccurences :: TermTypeM a -> TermTypeM (a, Occurences) +tapOccurences = listen + +removeSeminullOccurences :: TermTypeM a -> TermTypeM a +removeSeminullOccurences = censor $ filter $ not . seminullOccurence + +checkIfUsed :: Occurences -> Ident -> TermTypeM () +checkIfUsed occs v + | not $ identName v `S.member` allOccuring occs, + not $ "_" `isPrefixOf` prettyName (identName v) = + warn (srclocOf v) $ "Unused variable '"++pretty (baseName $ identName v)++"'." + | otherwise = + return () + +alternative :: TermTypeM a -> TermTypeM b -> TermTypeM (a,b) +alternative m1 m2 = pass $ do + (x, occurs1) <- listen m1 + (y, occurs2) <- listen m2 + checkOccurences occurs1 + checkOccurences occurs2 + let usage = occurs1 `altOccurences` occurs2 + return ((x, y), const usage) + +-- | Make all bindings nonunique. +noUnique :: TermTypeM a -> TermTypeM a +noUnique = local (\scope -> scope { scopeVtable = M.map set $ scopeVtable scope}) + where set (BoundV tparams t) = BoundV tparams $ t `setUniqueness` Nonunique + set (OverloadedF ts pts rt) = OverloadedF ts pts rt + set EqualityF = EqualityF + set OpaqueF = OpaqueF + set (WasConsumed loc) = WasConsumed loc + +onlySelfAliasing :: TermTypeM a -> TermTypeM a +onlySelfAliasing = local (\scope -> scope { scopeVtable = M.mapWithKey set $ scopeVtable scope}) + where set k (BoundV tparams t) = BoundV tparams $ t `addAliases` S.intersection (S.singleton k) + set _ (OverloadedF ts pts rt) = OverloadedF ts pts rt + set _ EqualityF = EqualityF + set _ OpaqueF = OpaqueF + set _ (WasConsumed loc) = WasConsumed loc + +arrayOfM :: (Pretty (ShapeDecl dim), Monoid as) => + SrcLoc + -> TypeBase dim as -> ShapeDecl dim -> Uniqueness + -> TermTypeM (TypeBase dim as) +arrayOfM loc t shape u = do + zeroOrderType loc "used in array" t + maybe nope return $ arrayOf t shape u + where nope = typeError loc $ + "Cannot form an array with elements of type " ++ pretty t + +-- | Perform substitutions of instantiated variables on the type +-- annotations (including the instance lists) of an expression, or +-- something else. +updateExpTypes :: ASTMappable e => e -> TermTypeM e +updateExpTypes e = do + constraints <- getConstraints + let look = (`lookupSubst` constraints) + tv = ASTMapper { mapOnExp = astMap tv + , mapOnName = pure + , mapOnQualName = pure + , mapOnType = pure . applySubst look + , mapOnCompType = pure . applySubst look + , mapOnStructType = pure . applySubst look + , mapOnPatternType = pure . applySubst look + } + astMap tv e diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -0,0 +1,422 @@ +{-# LANGUAGE FlexibleContexts, FlexibleInstances #-} +module Language.Futhark.TypeChecker.Types + ( checkTypeExp + , checkTypeDecl + + , unifyTypesU + , subtypeOf + , subuniqueOf + + , checkForDuplicateNames + , checkTypeParams + + , TypeSub(..) + , TypeSubs + , substituteTypes + , substituteTypesInBoundV + + , Substitutable(..) + , substTypesAny + ) +where + +import Control.Monad.Reader +import Control.Monad.Except +import Control.Monad.State +import Data.List +import Data.Loc +import Data.Maybe +import Data.Monoid ((<>)) +import qualified Data.Map.Strict as M + +import Language.Futhark +import Language.Futhark.TypeChecker.Monad + +-- | @unifyTypes uf t2 t2@ attempts to unify @t1@ and @t2@. If +-- unification cannot happen, 'Nothing' is returned, otherwise a type +-- that combines the aliasing of @t1@ and @t2@ is returned. +-- Uniqueness is unified with @uf@. +unifyTypesU :: (Monoid als, Eq als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) + -> TypeBase dim als -> TypeBase dim als -> Maybe (TypeBase dim als) +unifyTypesU _ (Prim t1) (Prim t2) + | t1 == t2 = Just $ Prim t1 + | otherwise = Nothing +unifyTypesU uf (TypeVar als1 u1 t1 targs1) (TypeVar als2 u2 t2 targs2) + | t1 == t2 = do + u3 <- uf u1 u2 + targs3 <- zipWithM (unifyTypeArgs uf) targs1 targs2 + Just $ TypeVar (als1 <> als2) u3 t1 targs3 + | otherwise = Nothing +unifyTypesU uf (Array et1 shape1 u1) (Array et2 shape2 u2) = + Array <$> unifyArrayElemTypes uf et1 et2 <*> + unifyShapes shape1 shape2 <*> uf u1 u2 +unifyTypesU uf (Record ts1) (Record ts2) + | length ts1 == length ts2, + sort (M.keys ts1) == sort (M.keys ts2) = + Record <$> traverse (uncurry (unifyTypesU uf)) + (M.intersectionWith (,) ts1 ts2) +unifyTypesU uf (Arrow as1 mn1 t1 t1') (Arrow as2 _ t2 t2') = + Arrow (as1 <> as2) mn1 <$> unifyTypesU (flip uf) t1 t2 <*> unifyTypesU uf t1' t2' +unifyTypesU _ _ _ = Nothing + +unifyTypeArgs :: (Monoid als, Eq als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) + -> TypeArg dim als -> TypeArg dim als -> Maybe (TypeArg dim als) +unifyTypeArgs _ (TypeArgDim d1 loc) (TypeArgDim d2 _) = + TypeArgDim <$> unifyDims d1 d2 <*> pure loc +unifyTypeArgs uf (TypeArgType t1 loc) (TypeArgType t2 _) = + TypeArgType <$> unifyTypesU uf t1 t2 <*> pure loc +unifyTypeArgs _ _ _ = + Nothing + +unifyArrayElemTypes :: (Monoid als, Eq als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) + -> ArrayElemTypeBase dim als + -> ArrayElemTypeBase dim als + -> Maybe (ArrayElemTypeBase dim als) +unifyArrayElemTypes _ (ArrayPrimElem bt1 als1) (ArrayPrimElem bt2 als2) + | bt1 == bt2 = + Just $ ArrayPrimElem bt1 (als1 <> als2) +unifyArrayElemTypes _ (ArrayPolyElem bt1 targs1 als1) (ArrayPolyElem bt2 targs2 als2) + | bt1 == bt2, targs1 == targs2 = + Just $ ArrayPolyElem bt1 targs1 (als1 <> als2) +unifyArrayElemTypes uf (ArrayRecordElem et1) (ArrayRecordElem et2) + | sort (M.keys et1) == sort (M.keys et2) = + ArrayRecordElem <$> + traverse (uncurry $ unifyRecordArrayElemTypes uf) (M.intersectionWith (,) et1 et2) +unifyArrayElemTypes _ _ _ = + Nothing + +unifyRecordArrayElemTypes :: (Monoid als, Eq als, ArrayDim dim) => + (Uniqueness -> Uniqueness -> Maybe Uniqueness) + -> RecordArrayElemTypeBase dim als + -> RecordArrayElemTypeBase dim als + -> Maybe (RecordArrayElemTypeBase dim als) +unifyRecordArrayElemTypes uf (RecordArrayElem et1) (RecordArrayElem et2) = + RecordArrayElem <$> unifyArrayElemTypes uf et1 et2 +unifyRecordArrayElemTypes uf (RecordArrayArrayElem et1 shape1 u1) (RecordArrayArrayElem et2 shape2 u2) = + RecordArrayArrayElem <$> unifyArrayElemTypes uf et1 et2 <*> + unifyShapes shape1 shape2 <*> uf u1 u2 +unifyRecordArrayElemTypes _ _ _ = + Nothing + +-- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal to +-- @y@), meaning @x@ is valid whenever @y@ is. +subtypeOf :: ArrayDim dim => + TypeBase dim as1 -> TypeBase dim as2 -> Bool +subtypeOf t1 t2 = isJust $ unifyTypesU unifyUniqueness (toStruct t1) (toStruct t2) + where unifyUniqueness u2 u1 = if u2 `subuniqueOf` u1 then Just u1 else Nothing + +-- | @x `subuniqueOf` y@ is true if @x@ is not less unique than @y@. +subuniqueOf :: Uniqueness -> Uniqueness -> Bool +subuniqueOf Nonunique Unique = False +subuniqueOf _ _ = True + +data Bindage = BoundAsVar | UsedFree + deriving (Show, Eq) + +checkTypeDecl :: MonadTypeChecker m => + TypeDeclBase NoInfo Name + -> m (TypeDeclBase Info VName, Liftedness) +checkTypeDecl (TypeDecl t NoInfo) = do + checkForDuplicateNamesInType t + (t', st, l) <- checkTypeExp t + return (TypeDecl t' $ Info st, l) + +checkTypeExp :: MonadTypeChecker m => + TypeExp Name + -> m (TypeExp VName, StructType, Liftedness) +checkTypeExp (TEVar name loc) = do + (name', ps, t, l) <- lookupType loc name + case ps of + [] -> return (TEVar name' loc, t, l) + _ -> throwError $ TypeError loc $ + "Type constructor " ++ pretty name ++ " used without any arguments." +checkTypeExp (TETuple ts loc) = do + (ts', ts_s, ls) <- unzip3 <$> mapM checkTypeExp ts + return (TETuple ts' loc, tupleRecord ts_s, foldl' max Unlifted ls) +checkTypeExp t@(TERecord fs loc) = do + -- Check for duplicate field names. + let field_names = map fst fs + unless (sort field_names == sort (nub field_names)) $ + throwError $ TypeError loc $ "Duplicate record fields in " ++ pretty t + + fs_ts_ls <- traverse checkTypeExp $ M.fromList fs + let fs' = fmap (\(x,_,_) -> x) fs_ts_ls + ts_s = fmap (\(_,y,_) -> y) fs_ts_ls + ls = fmap (\(_,_,z) -> z) fs_ts_ls + return (TERecord (M.toList fs') loc, Record ts_s, foldl' max Unlifted ls) +checkTypeExp (TEArray t d loc) = do + (t', st, l) <- checkTypeExp t + d' <- checkDimDecl d + case (l, arrayOf st (ShapeDecl [d']) Nonunique) of + (Unlifted, Just st') -> return (TEArray t' d' loc, st', Unlifted) + _ -> throwError $ TypeError loc $ + "Cannot create array with elements of type `" ++ pretty st ++ "` (might be functional)." + where checkDimDecl AnyDim = + return AnyDim + checkDimDecl (ConstDim k) = + return $ ConstDim k + checkDimDecl (NamedDim v) = + NamedDim <$> checkNamedDim loc v +checkTypeExp (TEUnique t loc) = do + (t', st, l) <- checkTypeExp t + unless (mayContainArray st) $ + warn loc $ "Declaring `" <> pretty st <> "` as unique has no effect." + return (TEUnique t' loc, st `setUniqueness` Unique, l) + where mayContainArray Prim{} = False + mayContainArray Array{} = True + mayContainArray (Record fs) = any mayContainArray fs + mayContainArray TypeVar{} = True + mayContainArray Arrow{} = False +checkTypeExp (TEArrow (Just v) t1 t2 loc) = do + (t1', st1, _) <- checkTypeExp t1 + bindSpaced [(Term, v)] $ do + v' <- checkName Term v loc + let env = mempty { envVtable = M.singleton v' $ BoundV [] st1 } + localEnv env $ do + (t2', st2, _) <- checkTypeExp t2 + return (TEArrow (Just v') t1' t2' loc, + Arrow mempty (Just v') st1 st2, + Lifted) +checkTypeExp (TEArrow Nothing t1 t2 loc) = do + (t1', st1, _) <- checkTypeExp t1 + (t2', st2, _) <- checkTypeExp t2 + return (TEArrow Nothing t1' t2' loc, + Arrow mempty Nothing st1 st2, + Lifted) +checkTypeExp ote@TEApply{} = do + (tname, tname_loc, targs) <- rootAndArgs ote + (tname', ps, t, l) <- lookupType tloc tname + if length ps /= length targs + then throwError $ TypeError tloc $ + "Type constructor " ++ pretty tname ++ " requires " ++ show (length ps) ++ + " arguments, but application at " ++ locStr tloc ++ " provides " ++ show (length targs) + else do + (targs', substs) <- unzip <$> zipWithM checkArgApply ps targs + return (foldl (\x y -> TEApply x y tloc) (TEVar tname' tname_loc) targs', + substituteTypes (mconcat substs) t, + l) + where tloc = srclocOf ote + + rootAndArgs :: MonadTypeChecker m => TypeExp Name -> m (QualName Name, SrcLoc, [TypeArgExp Name]) + rootAndArgs (TEVar qn loc) = return (qn, loc, []) + rootAndArgs (TEApply op arg _) = do (op', loc, args) <- rootAndArgs op + return (op', loc, args++[arg]) + rootAndArgs te' = throwError $ TypeError (srclocOf te') $ + "Type '" ++ pretty te' ++ "' is not a type constructor." + + checkArgApply (TypeParamDim pv _) (TypeArgExpDim (NamedDim v) loc) = do + v' <- checkNamedDim loc v + return (TypeArgExpDim (NamedDim v') loc, + M.singleton pv $ DimSub $ NamedDim v') + checkArgApply (TypeParamDim pv _) (TypeArgExpDim (ConstDim x) loc) = + return (TypeArgExpDim (ConstDim x) loc, + M.singleton pv $ DimSub $ ConstDim x) + checkArgApply (TypeParamDim pv _) (TypeArgExpDim AnyDim loc) = + return (TypeArgExpDim AnyDim loc, + M.singleton pv $ DimSub AnyDim) + + checkArgApply (TypeParamType l pv _) (TypeArgExpType te) = do + (te', st, _) <- checkTypeExp te + return (TypeArgExpType te', + M.singleton pv $ TypeSub $ TypeAbbr l [] st) + + checkArgApply p a = + throwError $ TypeError tloc $ "Type argument " ++ pretty a ++ + " not valid for a type parameter " ++ pretty p + + +checkNamedDim :: MonadTypeChecker m => + SrcLoc -> QualName Name -> m (QualName VName) +checkNamedDim loc v = do + (v', t) <- lookupVar loc v + case t of + Prim (Signed Int32) -> return v' + _ -> throwError $ TypeError loc $ + "Dimension declaration " ++ pretty v ++ + " should be of type `i32`." + +-- | Check for duplication of names inside a pattern group. Produces +-- a description of all names used in the pattern group. +checkForDuplicateNames :: MonadTypeChecker m => + [UncheckedPattern] -> m () +checkForDuplicateNames = (`evalStateT` mempty) . mapM_ check + where check (Id v _ loc) = seen v loc + check (PatternParens p _) = check p + check Wildcard{} = return () + check (TuplePattern ps _) = mapM_ check ps + check (RecordPattern fs _) = mapM_ (check . snd) fs + check (PatternAscription p _ _) = check p + + seen v loc = do + already <- gets $ M.lookup v + case already of + Just prev_loc -> + lift $ throwError $ TypeError loc $ + "Name " ++ pretty v ++ " also bound at " ++ locStr prev_loc + Nothing -> + modify $ M.insert v loc + +-- | Check whether the type contains arrow types that define the same +-- parameter. These might also exist further down, but that's not +-- really a problem - we mostly do this checking to help the user, +-- since it is likely an error, but it's easy to assign a semantics to +-- it (normal name shadowing). +checkForDuplicateNamesInType :: MonadTypeChecker m => + TypeExp Name -> m () +checkForDuplicateNamesInType = checkForDuplicateNames . pats + where pats (TEArrow (Just v) t1 t2 loc) = Id v NoInfo loc : pats t1 ++ pats t2 + pats (TEArrow Nothing t1 t2 _) = pats t1 ++ pats t2 + pats (TETuple ts _) = concatMap pats ts + pats (TERecord fs _) = concatMap (pats . snd) fs + pats (TEArray t _ _) = pats t + pats (TEUnique t _) = pats t + pats (TEApply t1 (TypeArgExpType t2) _) = pats t1 ++ pats t2 + pats (TEApply t1 TypeArgExpDim{} _) = pats t1 + pats TEVar{} = [] + +checkTypeParams :: MonadTypeChecker m => + [TypeParamBase Name] + -> ([TypeParamBase VName] -> m a) + -> m a +checkTypeParams ps m = + bindSpaced (map typeParamSpace ps) $ + m =<< evalStateT (mapM checkTypeParam ps) mempty + where typeParamSpace (TypeParamDim pv _) = (Term, pv) + typeParamSpace (TypeParamType _ pv _) = (Type, pv) + + checkParamName ns v loc = do + seen <- gets $ M.lookup (ns,v) + case seen of + Just prev -> + throwError $ TypeError loc $ + "Type parameter " ++ pretty v ++ " previously defined at " ++ locStr prev + Nothing -> do + modify $ M.insert (ns,v) loc + lift $ checkName ns v loc + + checkTypeParam (TypeParamDim pv loc) = + TypeParamDim <$> checkParamName Term pv loc <*> pure loc + checkTypeParam (TypeParamType l pv loc) = + TypeParamType l <$> checkParamName Type pv loc <*> pure loc + +data TypeSub = TypeSub TypeBinding + | DimSub (DimDecl VName) + deriving (Show) + +type TypeSubs = M.Map VName TypeSub + +substituteTypes :: TypeSubs -> StructType -> StructType +substituteTypes substs ot = case ot of + Array at shape u -> + fromMaybe nope $ arrayOf (substituteTypesInArrayElem at) (substituteInShape shape) u + Prim t -> Prim t + TypeVar () u v targs + | Just (TypeSub (TypeAbbr _ ps t)) <- + M.lookup (qualLeaf (qualNameFromTypeName v)) substs -> + applyType ps t (map substituteInTypeArg targs) + `setUniqueness` u + | otherwise -> TypeVar () u v $ map substituteInTypeArg targs + Record ts -> + Record $ fmap (substituteTypes substs) ts + Arrow als v t1 t2 -> + Arrow als v (substituteTypes substs t1) (substituteTypes substs t2) + where nope = error "substituteTypes: Cannot create array after substitution." + + substituteTypesInArrayElem (ArrayPrimElem t ()) = + Prim t + substituteTypesInArrayElem (ArrayPolyElem v targs ()) + | Just (TypeSub (TypeAbbr _ ps t)) <- + M.lookup (qualLeaf (qualNameFromTypeName v)) substs = + applyType ps t (map substituteInTypeArg targs) + | otherwise = + TypeVar () Nonunique v (map substituteInTypeArg targs) + substituteTypesInArrayElem (ArrayRecordElem ts) = + Record ts' + where ts' = fmap (substituteTypes substs . + fst . recordArrayElemToType) ts + + substituteInTypeArg (TypeArgDim d loc) = + TypeArgDim (substituteInDim d) loc + substituteInTypeArg (TypeArgType t loc) = + TypeArgType (substituteTypes substs t) loc + + substituteInShape (ShapeDecl ds) = + ShapeDecl $ map substituteInDim ds + + substituteInDim (NamedDim v) + | Just (DimSub d) <- M.lookup (qualLeaf v) substs = d + substituteInDim d = d + +substituteTypesInBoundV :: TypeSubs -> BoundV -> BoundV +substituteTypesInBoundV substs (BoundV tps t) = + BoundV tps (substituteTypes substs t) + +applyType :: [TypeParam] -> StructType -> [StructTypeArg] -> StructType +applyType ps t args = + substituteTypes substs t + where substs = M.fromList $ zipWith mkSubst ps args + -- We are assuming everything has already been type-checked for correctness. + mkSubst (TypeParamDim pv _) (TypeArgDim (NamedDim v) _) = + (pv, DimSub $ NamedDim v) + mkSubst (TypeParamDim pv _) (TypeArgDim (ConstDim x) _) = + (pv, DimSub $ ConstDim x) + mkSubst (TypeParamDim pv _) (TypeArgDim AnyDim _) = + (pv, DimSub AnyDim) + mkSubst (TypeParamType l pv _) (TypeArgType at _) = + (pv, TypeSub $ TypeAbbr l [] at) + mkSubst p a = + error $ "applyType mkSubst: cannot substitute " ++ pretty a ++ " for " ++ pretty p + +-- | Class of types which allow for substitution of types with no +-- annotations for type variable names. +class Substitutable a where + applySubst :: (VName -> Maybe (TypeBase () ())) -> a -> a + +instance Substitutable (TypeBase () ()) where + applySubst = substTypesAny + +instance Substitutable (TypeBase () Names) where + applySubst = substTypesAny . (fmap fromStruct.) + +instance Substitutable (TypeBase (DimDecl VName) ()) where + applySubst = substTypesAny . (fmap vacuousShapeAnnotations.) + +instance Substitutable (TypeBase (DimDecl VName) Names) where + applySubst = substTypesAny . (fmap (vacuousShapeAnnotations . fromStruct).) + +-- | Perform substitutions, from type names to types, on a type. Works +-- regardless of what shape and uniqueness information is attached to the type. +substTypesAny :: (ArrayDim dim, Monoid as) => + (VName -> Maybe (TypeBase dim as)) + -> TypeBase dim as -> TypeBase dim as +substTypesAny lookupSubst ot = case ot of + Prim t -> Prim t + Array et shape u -> fromMaybe nope $ + uncurry arrayOfWithAliases (subsArrayElem et) shape u + -- We only substitute for a type variable with no arguments, since + -- type parameters cannot have higher kind. + TypeVar _ u v [] + | Just t <- lookupSubst $ qualLeaf (qualNameFromTypeName v) -> + t `setUniqueness` u + TypeVar als u v targs -> TypeVar als u v $ map subsTypeArg targs + Record ts -> Record $ fmap (substTypesAny lookupSubst) ts + Arrow als v t1 t2 -> + Arrow als v (substTypesAny lookupSubst t1) (substTypesAny lookupSubst t2) + + where nope = error "substTypesAny: Cannot create array after substitution." + + subsArrayElem (ArrayPrimElem t as) = (Prim t, as) + subsArrayElem (ArrayPolyElem v [] as) + | Just t <- lookupSubst $ qualLeaf (qualNameFromTypeName v) = (t, as) + subsArrayElem (ArrayPolyElem v targs as) = + (TypeVar as Nonunique v (map subsTypeArg targs), as) + subsArrayElem (ArrayRecordElem ts) = + let ts' = fmap recordArrayElemToType ts + in (Record $ fmap (substTypesAny lookupSubst . fst) ts', foldMap snd ts') + + subsTypeArg (TypeArgType t loc) = + TypeArgType (substTypesAny lookupSubst t) loc + subsTypeArg t = t diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -0,0 +1,355 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Language.Futhark.TypeChecker.Unify + ( Constraint(..) + , Constraints + , lookupSubst + , MonadUnify(..) + , BreadCrumb(..) + , typeError + + , zeroOrderType + , mustHaveField + , mustBeOneOf + , equalityType + , normaliseType + + , unify + , doUnification + ) +where + +import Control.Monad.Except +import Control.Monad.State +import Data.List +import Data.Loc +import Data.Maybe +import qualified Data.Map.Strict as M +import qualified Data.Set as S + +import Prelude hiding (mod) + +import Language.Futhark +import Language.Futhark.TypeChecker.Monad hiding (BoundV, checkQualNameWithEnv) +import Language.Futhark.TypeChecker.Types hiding (checkTypeDecl) +import Futhark.Util.Pretty (Pretty) + +-- | Mapping from fresh type variables, instantiated from the type +-- schemes of polymorphic functions, to (possibly) specific types as +-- determined on application and the location of that application, or +-- a partial constraint on their type. +type Constraints = M.Map VName Constraint + +data Constraint = NoConstraint (Maybe Liftedness) SrcLoc + | ParamType Liftedness SrcLoc + | Constraint (TypeBase () ()) SrcLoc + | Overloaded [PrimType] SrcLoc + | HasFields (M.Map Name (TypeBase () ())) SrcLoc + | Equality SrcLoc + deriving Show + +instance Located Constraint where + locOf (NoConstraint _ loc) = locOf loc + locOf (ParamType _ loc) = locOf loc + locOf (Constraint _ loc) = locOf loc + locOf (Overloaded _ loc) = locOf loc + locOf (HasFields _ loc) = locOf loc + locOf (Equality loc) = locOf loc + +lookupSubst :: VName -> Constraints -> Maybe (TypeBase () ()) +lookupSubst v constraints = do Constraint t _ <- M.lookup v constraints + Just t + +class (MonadBreadCrumbs m, MonadError TypeError m) => MonadUnify m where + getConstraints :: m Constraints + putConstraints :: Constraints -> m () + modifyConstraints :: (Constraints -> Constraints) -> m () + modifyConstraints f = do + x <- getConstraints + putConstraints $ f x + + newTypeVar :: Monoid als => SrcLoc -> String -> m (TypeBase dim als) + +normaliseType :: (Substitutable a, MonadUnify m) => a -> m a +normaliseType t = do constraints <- getConstraints + return $ applySubst (`lookupSubst` constraints) t + +-- | Is the given type variable actually the name of an abstract type +-- or type parameter, which we cannot substitute? +isRigid :: VName -> Constraints -> Bool +isRigid v constraints = case M.lookup v constraints of + Nothing -> True + Just ParamType{} -> True + _ -> False + +-- | Unifies two types. +unify :: MonadUnify m => SrcLoc -> TypeBase () () -> TypeBase () () -> m () +unify loc orig_t1 orig_t2 = do + orig_t1' <- normaliseType orig_t1 + orig_t2' <- normaliseType orig_t2 + breadCrumb (MatchingTypes orig_t1' orig_t2') $ subunify orig_t1 orig_t2 + where + subunify t1 t2 = do + constraints <- getConstraints + + let isRigid' v = isRigid v constraints + t1' = applySubst (`lookupSubst` constraints) t1 + t2' = applySubst (`lookupSubst` constraints) t2 + + failure = + typeError loc $ "Couldn't match expected type `" ++ + pretty t1' ++ "' with actual type `" ++ pretty t2' ++ "'." + + case (t1', t2') of + _ | t1' == t2' -> return () + + (Record fs, + Record arg_fs) + | M.keys fs == M.keys arg_fs -> + forM_ (M.toList $ M.intersectionWith (,) fs arg_fs) $ \(k, (k_t1, k_t2)) -> + breadCrumb (MatchingFields k) $ subunify k_t1 k_t2 + + (TypeVar _ _ (TypeName _ tn) targs, + TypeVar _ _ (TypeName _ arg_tn) arg_targs) + | tn == arg_tn, length targs == length arg_targs -> + zipWithM_ unifyTypeArg targs arg_targs + + (TypeVar _ _ (TypeName [] v1) [], + TypeVar _ _ (TypeName [] v2) []) -> + case (isRigid' v1, isRigid' v2) of + (True, True) -> failure + (True, False) -> linkVarToType loc v2 t1' + (False, True) -> linkVarToType loc v1 t2' + (False, False) -> linkVarToType loc v1 t2' + + (TypeVar _ _ (TypeName [] v1) [], _) + | not $ isRigid' v1 -> + linkVarToType loc v1 t2' + (_, TypeVar _ _ (TypeName [] v2) []) + | not $ isRigid' v2 -> + linkVarToType loc v2 t1' + + (Arrow _ _ a1 b1, + Arrow _ _ a2 b2) -> do + subunify a1 a2 + subunify b1 b2 + + (Array{}, Array{}) + | Just t1'' <- peelArray 1 t1', + Just t2'' <- peelArray 1 t2' -> + subunify t1'' t2'' + + (_, _) -> failure + + where unifyTypeArg TypeArgDim{} TypeArgDim{} = return () + unifyTypeArg (TypeArgType t _) (TypeArgType arg_t _) = + subunify t arg_t + unifyTypeArg _ _ = typeError loc + "Cannot unify a type argument with a dimension argument (or vice versa)." + +applySubstInConstraint :: VName -> TypeBase () () -> Constraint -> Constraint +applySubstInConstraint vn tp (Constraint t loc) = + Constraint (applySubst (`M.lookup` M.singleton vn tp) t) loc +applySubstInConstraint vn tp (HasFields fs loc) = + HasFields (M.map (applySubst (`M.lookup` M.singleton vn tp)) fs) loc +applySubstInConstraint _ _ (NoConstraint l loc) = NoConstraint l loc +applySubstInConstraint _ _ (Overloaded ts loc) = Overloaded ts loc +applySubstInConstraint _ _ (Equality loc) = Equality loc +applySubstInConstraint _ _ (ParamType l loc) = ParamType l loc + +linkVarToType :: MonadUnify m => SrcLoc -> VName -> TypeBase () () -> m () +linkVarToType loc vn tp = do + constraints <- getConstraints + if vn `S.member` typeVars tp + then typeError loc $ "Occurs check: cannot instantiate " ++ + prettyName vn ++ " with " ++ pretty tp' + else do modifyConstraints $ M.insert vn $ Constraint tp' loc + modifyConstraints $ M.map $ applySubstInConstraint vn tp' + case M.lookup vn constraints of + Just (NoConstraint (Just Unlifted) unlift_loc) -> + zeroOrderType loc ("used at " ++ locStr unlift_loc) tp' + Just (Equality _) -> + equalityType loc tp' + Just (Overloaded ts old_loc) + | tp `notElem` map Prim ts -> + case tp' of + TypeVar _ _ (TypeName [] v) [] + | not $ isRigid v constraints -> linkVarToTypes loc v ts + _ -> + typeError loc $ "Cannot unify `" ++ prettyName vn ++ "' with type `" ++ + pretty tp ++ "' (must be one of " ++ intercalate ", " (map pretty ts) ++ + " due to use at " ++ locStr old_loc ++ ")." + Just (HasFields required_fields old_loc) -> + case tp of + Record tp_fields + | all (`M.member` tp_fields) $ M.keys required_fields -> + mapM_ (uncurry $ unify loc) $ M.elems $ + M.intersectionWith (,) required_fields tp_fields + TypeVar _ _ (TypeName [] v) [] + | not $ isRigid v constraints -> + modifyConstraints $ M.insert v $ + HasFields required_fields old_loc + _ -> + let required_fields' = + intercalate ", " $ map field $ M.toList required_fields + field (l, t) = pretty l ++ ": " ++ pretty t + in typeError loc $ + "Cannot unify `" ++ prettyName vn ++ "' with type `" ++ + pretty tp ++ "' (must be a record with fields {" ++ + required_fields' ++ + "} due to use at " ++ locStr old_loc ++ ")." + _ -> return () + where tp' = removeUniqueness tp + +removeUniqueness :: TypeBase dim as -> TypeBase dim as +removeUniqueness (Record ets) = + Record $ fmap removeUniqueness ets +removeUniqueness (Arrow als p t1 t2) = + Arrow als p (removeUniqueness t1) (removeUniqueness t2) +removeUniqueness t = t `setUniqueness` Nonunique + +mustBeOneOf :: MonadUnify m => [PrimType] -> SrcLoc -> TypeBase () () -> m () +mustBeOneOf [req_t] loc t = unify loc (Prim req_t) t +mustBeOneOf ts loc t = do + constraints <- getConstraints + let t' = applySubst (`lookupSubst` constraints) t + isRigid' v = isRigid v constraints + + case t' of + TypeVar _ _ (TypeName [] v) [] + | not $ isRigid' v -> linkVarToTypes loc v ts + + Prim pt | pt `elem` ts -> return () + + _ -> failure + + where failure = typeError loc $ "Cannot unify type \"" ++ pretty t ++ + "\" with any of " ++ intercalate "," (map pretty ts) ++ "." + +linkVarToTypes :: MonadUnify m => SrcLoc -> VName -> [PrimType] -> m () +linkVarToTypes loc vn ts = do + vn_constraint <- M.lookup vn <$> getConstraints + case vn_constraint of + Just (Overloaded vn_ts vn_loc) -> + case ts `intersect` vn_ts of + [] -> typeError loc $ "Type constrained to one of " ++ + intercalate "," (map pretty ts) ++ " but also one of " ++ + intercalate "," (map pretty vn_ts) ++ " at " ++ locStr vn_loc ++ "." + ts' -> modifyConstraints $ M.insert vn $ Overloaded ts' loc + + _ -> modifyConstraints $ M.insert vn $ Overloaded ts loc + +equalityType :: (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => + SrcLoc -> TypeBase dim as -> m () +equalityType loc t = do + unless (orderZero t) $ + typeError loc $ + "Type \"" ++ pretty t ++ "\" does not support equality." + mapM_ mustBeEquality $ typeVars t + where mustBeEquality vn = do + constraints <- getConstraints + case M.lookup vn constraints of + Just (Constraint (TypeVar _ _ (TypeName [] vn') []) _) -> + mustBeEquality vn' + Just (Constraint vn_t _) + | not $ orderZero vn_t -> + typeError loc $ "Type \"" ++ pretty t ++ + "\" does not support equality." + | otherwise -> return () + Just (NoConstraint _ _) -> + modifyConstraints $ M.insert vn (Equality loc) + Just (Overloaded _ _) -> + return () -- All primtypes support equality. + _ -> + typeError loc $ "Type " ++ pretty (prettyName vn) ++ + " does not support equality." + +zeroOrderType :: (MonadUnify m, Pretty (ShapeDecl dim), Monoid as) => + SrcLoc -> String -> TypeBase dim as -> m () +zeroOrderType loc desc t = do + unless (orderZero t) $ + typeError loc $ "Type " ++ desc ++ + " must not be functional, but is " ++ pretty t ++ "." + mapM_ mustBeZeroOrder . S.toList . typeVars $ t + where mustBeZeroOrder vn = do + constraints <- getConstraints + case M.lookup vn constraints of + Just (Constraint vn_t old_loc) + | not $ orderZero t -> + typeError loc $ "Type " ++ desc ++ + " must be non-function, but inferred to be " ++ + pretty vn_t ++ " at " ++ locStr old_loc ++ "." + Just (NoConstraint _ _) -> + modifyConstraints $ M.insert vn (NoConstraint (Just Unlifted) loc) + Just (ParamType Lifted ploc) -> + typeError loc $ "Type " ++ desc ++ + " must be non-function, but type parameter " ++ prettyName vn ++ " at " ++ + locStr ploc ++ " may be a function." + _ -> return () + +mustHaveField :: (MonadUnify m, Monoid as) => + SrcLoc -> Name -> TypeBase dim as -> m (TypeBase dim as) +mustHaveField loc l t = do + constraints <- getConstraints + l_type <- newTypeVar loc "t" + let l_type' = toStructural l_type + case t of + TypeVar _ _ (TypeName _ tn) [] + | Just NoConstraint{} <- M.lookup tn constraints -> do + modifyConstraints $ M.insert tn $ HasFields (M.singleton l l_type') loc + return l_type + | Just (HasFields fields _) <- M.lookup tn constraints -> do + case M.lookup l fields of + Just t' -> unify loc (toStructural t) t' + Nothing -> modifyConstraints $ M.insert tn $ + HasFields (M.insert l l_type' fields) loc + return l_type + Record fields + | Just t' <- M.lookup l fields -> do + unify loc l_type' (toStructural t') + return t' + | otherwise -> + throwError $ TypeError loc $ + "Attempt to access field '" ++ pretty l ++ "' of value of type " ++ + pretty (toStructural t) ++ "." + _ -> do unify loc (toStructural t) $ Record $ M.singleton l l_type' + return l_type + +-- Simple MonadUnify implementation. + +type UnifyMState = (Constraints, Int) + +newtype UnifyM a = UnifyM (StateT UnifyMState (Except TypeError) a) + deriving (Monad, Functor, Applicative, + MonadState UnifyMState, + MonadError TypeError) + +instance MonadUnify UnifyM where + getConstraints = gets fst + putConstraints x = modify $ \s -> (x, snd s) + + newTypeVar loc desc = do + i <- do (x, i) <- get + put (x, i+1) + return i + let v = VName (nameFromString $ desc ++ show i) 0 + modifyConstraints $ M.insert v $ NoConstraint Nothing loc + return $ TypeVar mempty Nonunique (typeName v) [] + +instance MonadBreadCrumbs UnifyM where + +-- | Perform a unification of two types outside a monadic context. +-- The type parameters are allowed to be instantiated (with +-- 'TypeParamDim ignored); all other types are considered rigid. +doUnification :: SrcLoc -> [TypeParam] + -> TypeBase () () -> TypeBase () () + -> Either TypeError (TypeBase () ()) +doUnification loc tparams t1 t2 = runUnifyM tparams $ do + unify loc t1 t2 + normaliseType t2 + +runUnifyM :: [TypeParam] -> UnifyM a -> Either TypeError a +runUnifyM tparams (UnifyM m) = runExcept $ evalStateT m (constraints, 0) + where constraints = M.fromList $ mapMaybe f tparams + f TypeParamDim{} = Nothing + f (TypeParamType l p loc) = Just (p, NoConstraint (Just l) loc) diff --git a/src/Language/Futhark/Warnings.hs b/src/Language/Futhark/Warnings.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/Warnings.hs @@ -0,0 +1,38 @@ +module Language.Futhark.Warnings + ( Warnings + , singleWarning + ) where + +import Data.Monoid +import Data.List +import Data.Loc +import qualified Data.Semigroup as Sem + +import Prelude + +import Language.Futhark.Core (locStr) + +-- | The warnings produced by the compiler. The 'Show' instance +-- produces a human-readable description. +newtype Warnings = Warnings [(SrcLoc, String)] deriving (Eq) + +instance Sem.Semigroup Warnings where + Warnings ws1 <> Warnings ws2 = Warnings $ ws1 <> ws2 + +instance Monoid Warnings where + mempty = Warnings mempty + mappend = (Sem.<>) + +instance Show Warnings where + show (Warnings []) = "" + show (Warnings ws) = + intercalate "\n\n" ws' ++ "\n" + where ws' = map showWarning $ sortOn (off . locOf . fst) ws + off NoLoc = 0 + off (Loc p _) = posCoff p + showWarning (loc, w) = + "Warning at " ++ locStr loc ++ ":\n" ++ + intercalate "\n" (map (" "<>) $ lines w) + +singleWarning :: SrcLoc -> String -> Warnings +singleWarning loc problem = Warnings [(loc, problem)] diff --git a/src/futhark-bench.hs b/src/futhark-bench.hs new file mode 100644 --- /dev/null +++ b/src/futhark-bench.hs @@ -0,0 +1,390 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts #-} +-- | Simple tool for benchmarking Futhark programs. Use the @--json@ +-- flag for machine-readable output. +module Main (main) where + +import Control.Concurrent +import Control.Monad +import Control.Monad.Except +import qualified Data.ByteString.Char8 as SBS +import qualified Data.ByteString.Lazy.Char8 as LBS +import Data.Either +import Data.Maybe +import Data.Semigroup ((<>)) +import Data.List +import qualified Data.Text as T +import qualified Data.Text.IO as T +import qualified Data.Text.Encoding as T +import System.Console.GetOpt +import System.FilePath +import System.Directory +import System.IO +import System.IO.Temp +import System.Timeout +import System.Process.ByteString (readProcessWithExitCode) +import System.Exit +import qualified Text.JSON as JSON +import Text.Printf +import Text.Regex.TDFA + +import Futhark.Test +import Futhark.Util.Options + +data BenchOptions = BenchOptions + { optCompiler :: String + , optRunner :: String + , optRuns :: Int + , optExtraOptions :: [String] + , optJSON :: Maybe FilePath + , optTimeout :: Int + , optSkipCompilation :: Bool + , optExcludeCase :: [String] + , optIgnoreFiles :: [Regex] + } + +initialBenchOptions :: BenchOptions +initialBenchOptions = BenchOptions "futhark-c" "" 10 [] Nothing (-1) False + ["nobench", "disable"] [] + +-- | The name we use for compiled programs. +binaryName :: FilePath -> FilePath +binaryName = dropExtension + +newtype RunResult = RunResult { runMicroseconds :: Int } +data DataResult = DataResult String (Either T.Text ([RunResult], T.Text)) +data BenchResult = BenchResult FilePath [DataResult] + +resultsToJSON :: [BenchResult] -> JSON.JSValue +resultsToJSON = JSON.JSObject . JSON.toJSObject . map benchResultToJSObject + where benchResultToJSObject + :: BenchResult + -> (String, JSON.JSValue) + benchResultToJSObject (BenchResult prog rs) = + (prog, JSON.JSObject $ JSON.toJSObject + [("datasets", JSON.JSObject $ JSON.toJSObject $ + map dataResultToJSObject rs)]) + dataResultToJSObject + :: DataResult + -> (String, JSON.JSValue) + dataResultToJSObject (DataResult desc (Left err)) = + (desc, JSON.showJSON err) + dataResultToJSObject (DataResult desc (Right (runtimes, progerr))) = + (desc, JSON.JSObject $ JSON.toJSObject + [("runtimes", JSON.showJSON $ map runMicroseconds runtimes), + ("stderr", JSON.showJSON progerr)]) + +fork :: (a -> IO b) -> a -> IO (MVar b) +fork f x = do cell <- newEmptyMVar + void $ forkIO $ do result <- f x + putMVar cell result + return cell + +pmapIO :: (a -> IO b) -> [a] -> IO [b] +pmapIO f elems = go elems [] + where + go [] res = return res + go xs res = do + numThreads <- getNumCapabilities + let (e,es) = splitAt numThreads xs + mvars <- mapM (fork f) e + result <- mapM takeMVar mvars + go es (result ++ res) + +runBenchmarks :: BenchOptions -> [FilePath] -> IO () +runBenchmarks opts paths = do + -- We force line buffering to ensure that we produce running output. + -- Otherwise, CI tools and the like may believe we are hung and kill + -- us. + hSetBuffering stdout LineBuffering + benchmarks <- filter (not . ignored . fst) <$> testSpecsFromPaths paths + (skipped_benchmarks, compiled_benchmarks) <- + partitionEithers <$> pmapIO (compileBenchmark opts) benchmarks + + when (anyFailedToCompile skipped_benchmarks) exitFailure + + results <- concat <$> mapM (runBenchmark opts) compiled_benchmarks + case optJSON opts of + Nothing -> return () + Just file -> writeFile file $ JSON.encode $ resultsToJSON results + when (anyFailed results) exitFailure + + where ignored f = any (`match` f) $ optIgnoreFiles opts + +anyFailed :: [BenchResult] -> Bool +anyFailed = any failedBenchResult + where failedBenchResult (BenchResult _ xs) = + any failedResult xs + failedResult (DataResult _ Left{}) = True + failedResult _ = False + +anyFailedToCompile :: [SkipReason] -> Bool +anyFailedToCompile = elem FailedToCompile + +data SkipReason = Skipped | FailedToCompile + deriving (Eq) + +compileBenchmark :: BenchOptions -> (FilePath, ProgramTest) + -> IO (Either SkipReason (FilePath, [InputOutputs])) +compileBenchmark opts (program, spec) = + case testAction spec of + RunCases cases _ _ | "nobench" `notElem` testTags spec, + "disable" `notElem` testTags spec, + any hasRuns cases -> + if optSkipCompilation opts + then do + exists <- doesFileExist $ binaryName program + if exists + then return $ Right (program, cases) + else do putStrLn $ binaryName program ++ " does not exist, but --skip-compilation passed." + return $ Left FailedToCompile + else do + putStr $ "Compiling " ++ program ++ "...\n" + (futcode, _, futerr) <- liftIO $ readProcessWithExitCode compiler + [program, "-o", binaryName program] "" + + case futcode of + ExitSuccess -> return $ Right (program, cases) + ExitFailure 127 -> do putStrLn $ "Failed:\n" ++ progNotFound compiler + return $ Left FailedToCompile + ExitFailure _ -> do putStrLn "Failed:\n" + SBS.putStrLn futerr + return $ Left FailedToCompile + _ -> + return $ Left Skipped + where compiler = optCompiler opts + + hasRuns (InputOutputs _ runs) = not $ null runs + +runBenchmark :: BenchOptions -> (FilePath, [InputOutputs]) -> IO [BenchResult] +runBenchmark opts (program, cases) = mapM forInputOutputs cases + where forInputOutputs (InputOutputs entry_name runs) = do + putStr $ "Results for " ++ program' ++ ":\n" + BenchResult program' . catMaybes <$> + mapM (runBenchmarkCase opts program entry_name pad_to) runs + where program' = if entry_name == "main" + then program + else program ++ ":" ++ T.unpack entry_name + + pad_to = foldl max 0 $ concatMap (map (length . runDescription) . iosTestRuns) cases + +reportResult :: [RunResult] -> IO () +reportResult [] = + print (0::Int) +reportResult results = do + let runtimes = map (fromIntegral . runMicroseconds) results + avg = sum runtimes / fromIntegral (length runtimes) + rel_dev = stddevp runtimes / mean runtimes :: Double + putStrLn $ printf "%10.2f" avg ++ "μs (avg. of " ++ show (length runtimes) ++ + " runs; RSD: " ++ printf "%.2f" rel_dev ++ ")" + +progNotFound :: String -> String +progNotFound s = s ++ ": command not found" + +type BenchM = ExceptT T.Text IO + +runBenchM :: BenchM a -> IO (Either T.Text a) +runBenchM = runExceptT + +io :: IO a -> BenchM a +io = liftIO + +runBenchmarkCase :: BenchOptions -> FilePath -> T.Text -> Int -> TestRun + -> IO (Maybe DataResult) +runBenchmarkCase _ _ _ _ (TestRun _ _ RunTimeFailure{} _ _) = + return Nothing -- Not our concern, we are not a testing tool. +runBenchmarkCase opts _ _ _ (TestRun tags _ _ _ _) + | any (`elem` tags) $ optExcludeCase opts = + return Nothing +runBenchmarkCase opts program entry pad_to (TestRun _ input_spec (Succeeds expected_spec) _ dataset_desc) = + -- We store the runtime in a temporary file. + withSystemTempFile "futhark-bench" $ \tmpfile h -> do + hClose h -- We will be writing and reading this ourselves. + input <- getValuesBS dir input_spec + let getValuesAndBS vs = do + vs' <- getValues dir vs + bs <- getValuesBS dir vs + return (LBS.toStrict bs, vs') + maybe_expected <- maybe (return Nothing) (fmap Just . getValuesAndBS) expected_spec + let options = optExtraOptions opts ++ ["-e", T.unpack entry, + "-t", tmpfile, + "-r", show $ optRuns opts, + "-b"] + + -- Report the dataset name before running the program, so that if an + -- error occurs it's easier to see where. + putStr $ "dataset " ++ dataset_desc ++ ": " ++ + replicate (pad_to - length dataset_desc) ' ' + hFlush stdout + + -- Explicitly prefixing the current directory is necessary for + -- readProcessWithExitCode to find the binary when binOutputf has + -- no program component. + let (to_run, to_run_args) + | null $ optRunner opts = ("." binaryName program, options) + | otherwise = (optRunner opts, binaryName program : options) + + run_res <- + timeout (optTimeout opts * 1000000) $ + readProcessWithExitCode to_run to_run_args $ + LBS.toStrict input + + fmap (Just . DataResult dataset_desc) $ runBenchM $ case run_res of + Just (progCode, output, progerr) -> + do + case maybe_expected of + Nothing -> + didNotFail program progCode $ T.decodeUtf8 progerr + Just expected -> + compareResult program expected =<< + runResult program progCode output progerr + runtime_result <- io $ T.readFile tmpfile + runtimes <- case mapM readRuntime $ T.lines runtime_result of + Just runtimes -> return $ map RunResult runtimes + Nothing -> itWentWrong $ "Runtime file has invalid contents:\n" <> runtime_result + + io $ reportResult runtimes + return (runtimes, T.decodeUtf8 progerr) + Nothing -> + itWentWrong $ T.pack $ "Execution exceeded " ++ show (optTimeout opts) ++ " seconds." + + where dir = takeDirectory program + + +readRuntime :: T.Text -> Maybe Int +readRuntime s = case reads $ T.unpack s of + [(runtime, _)] -> Just runtime + _ -> Nothing + +didNotFail :: FilePath -> ExitCode -> T.Text -> BenchM () +didNotFail _ ExitSuccess _ = + return () +didNotFail program (ExitFailure code) stderr_s = + itWentWrong $ T.pack $ program ++ " failed with error code " ++ show code ++ + " and output:\n" ++ T.unpack stderr_s + +itWentWrong :: (MonadError T.Text m, MonadIO m) => + T.Text -> m a +itWentWrong t = do + liftIO $ putStrLn $ T.unpack t + throwError t + +runResult :: (MonadError T.Text m, MonadIO m) => + FilePath + -> ExitCode + -> SBS.ByteString + -> SBS.ByteString + -> m (SBS.ByteString, [Value]) +runResult program ExitSuccess stdout_s _ = + case valuesFromByteString "stdout" $ LBS.fromStrict stdout_s of + Left e -> do + let actualf = program `replaceExtension` "actual" + liftIO $ SBS.writeFile actualf stdout_s + itWentWrong $ T.pack $ show e <> "\n(See " <> actualf <> ")" + Right vs -> return (stdout_s, vs) +runResult program (ExitFailure code) _ stderr_s = + itWentWrong $ T.pack $ program ++ " failed with error code " ++ show code ++ + " and output:\n" ++ T.unpack (T.decodeUtf8 stderr_s) + +compareResult :: (MonadError T.Text m, MonadIO m) => + FilePath -> (SBS.ByteString, [Value]) -> (SBS.ByteString, [Value]) + -> m () +compareResult program (expected_bs, expected_vs) (actual_bs, actual_vs) = + case compareValues actual_vs expected_vs of + Just mismatch -> do + let actualf = program `replaceExtension` "actual" + expectedf = program `replaceExtension` "expected" + liftIO $ SBS.writeFile actualf actual_bs + liftIO $ SBS.writeFile expectedf expected_bs + itWentWrong $ T.pack actualf <> " and " <> T.pack expectedf <> + " do not match:\n" <> T.pack (show mismatch) + Nothing -> + return () + +commandLineOptions :: [FunOptDescr BenchOptions] +commandLineOptions = [ + Option "r" ["runs"] + (ReqArg (\n -> + case reads n of + [(n', "")] | n' >= 0 -> + Right $ \config -> + config { optRuns = n' + } + _ -> + Left $ error $ "'" ++ n ++ "' is not a non-negative integer.") + "RUNS") + "Run each test case this many times." + , Option [] ["compiler"] + (ReqArg (\prog -> + Right $ \config -> config { optCompiler = prog }) + "PROGRAM") + "The compiler used (defaults to 'futhark-c')." + , Option [] ["runner"] + (ReqArg (\prog -> Right $ \config -> config { optRunner = prog }) "PROGRAM") + "The program used to run the Futhark-generated programs (defaults to nothing)." + , Option "p" ["pass-option"] + (ReqArg (\opt -> + Right $ \config -> + config { optExtraOptions = opt : optExtraOptions config }) + "OPT") + "Pass this option to programs being run." + , Option [] ["json"] + (ReqArg (\file -> + Right $ \config -> config { optJSON = Just file}) + "FILE") + "Scatter results in JSON format here." + , Option [] ["timeout"] + (ReqArg (\n -> + case reads n of + [(n', "")] + | n' < max_timeout -> + Right $ \config -> config { optTimeout = fromIntegral n' } + _ -> + Left $ error $ "'" ++ n ++ + "' is not an integer smaller than" ++ show max_timeout ++ ".") + "SECONDS") + "Number of seconds before a dataset is aborted." + , Option [] ["skip-compilation"] + (NoArg $ Right $ \config -> config { optSkipCompilation = True }) + "Use already compiled program." + , Option [] ["exclude-case"] + (ReqArg (\s -> Right $ \config -> + config { optExcludeCase = s : optExcludeCase config }) + "TAG") + "Do not run test cases with this tag." + , Option [] ["ignore-files"] + (ReqArg (\s -> Right $ \config -> + config { optIgnoreFiles = makeRegex s : optIgnoreFiles config }) + "REGEX") + "Ignore files matching this regular expression." + ] + where max_timeout :: Int + max_timeout = maxBound `div` 1000000 + +main :: IO () +main = mainWithOptions initialBenchOptions commandLineOptions "options... programs..." $ \progs config -> + Just $ runBenchmarks config progs + +--- The following extracted from hstats package by Marshall Beddoe: +--- https://hackage.haskell.org/package/hstats-0.3 + +-- | Numerically stable mean +mean :: Floating a => [a] -> a +mean x = fst $ foldl' (\(!m, !n) x' -> (m+(x'-m)/(n+1),n+1)) (0,0) x + +-- | Standard deviation of population +stddevp :: (Floating a) => [a] -> a +stddevp xs = sqrt $ pvar xs + +-- | Population variance +pvar :: (Floating a) => [a] -> a +pvar xs = centralMoment xs (2::Int) + +-- | Central moments +centralMoment :: (Floating b, Integral t) => [b] -> t -> b +centralMoment _ 1 = 0 +centralMoment xs r = sum (map (\x -> (x-m)^r) xs) / n + where + m = mean xs + n = fromIntegral $ length xs diff --git a/src/futhark-c.hs b/src/futhark-c.hs new file mode 100644 --- /dev/null +++ b/src/futhark-c.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE FlexibleContexts #-} +module Main (main) where + +import Control.Monad.IO.Class +import System.FilePath +import System.Exit + +import Futhark.Pipeline +import Futhark.Passes +import qualified Futhark.CodeGen.Backends.SequentialC as SequentialC +import Futhark.Util.Pretty (prettyText) +import Futhark.Compiler.CLI +import Futhark.Util + +main :: IO () +main = compilerMain () [] + "Compile sequential C" "Generate sequential C code from optimised Futhark program." + sequentialCpuPipeline $ \() mode outpath prog -> do + cprog <- either (`internalError` prettyText prog) return =<< + SequentialC.compileProg prog + let cpath = outpath `addExtension` "c" + hpath = outpath `addExtension` "h" + + case mode of + ToLibrary -> do + let (header, impl) = SequentialC.asLibrary cprog + liftIO $ writeFile hpath header + liftIO $ writeFile cpath impl + ToExecutable -> do + liftIO $ writeFile cpath $ SequentialC.asExecutable cprog + ret <- liftIO $ runProgramWithExitCode "gcc" + [cpath, "-O3", "-std=c99", "-lm", "-o", outpath] "" + case ret of + Left err -> + externalErrorS $ "Failed to run gcc: " ++ show err + Right (ExitFailure code, _, gccerr) -> + externalErrorS $ "gcc failed with code " ++ + show code ++ ":\n" ++ gccerr + Right (ExitSuccess, _, _) -> + return () diff --git a/src/futhark-cs.hs b/src/futhark-cs.hs new file mode 100644 --- /dev/null +++ b/src/futhark-cs.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE FlexibleContexts #-} +module Main (main) where + +import Control.Monad.IO.Class +import Data.Maybe (fromMaybe) +import System.FilePath +import System.Directory +import System.Exit +import System.Environment + +import Futhark.Pipeline +import Futhark.Passes +import qualified Futhark.CodeGen.Backends.SequentialCSharp as SequentialCS +import Futhark.Util.Pretty (prettyText) +import Futhark.Compiler.CLI +import Futhark.Util + +main :: IO () +main = compilerMain () [] + "Compile sequential C#" "Generate sequential C# code from optimised Futhark program." + sequentialCpuPipeline $ \() mode outpath prog -> do + mono_libs <- liftIO $ fromMaybe "." <$> lookupEnv "MONO_PATH" + let class_name = + case mode of ToLibrary -> Just $ takeBaseName outpath + ToExecutable -> Nothing + csprog <- either (`internalError` prettyText prog) return =<< + SequentialCS.compileProg class_name prog + + let cspath = outpath `addExtension` "cs" + liftIO $ writeFile cspath csprog + + case mode of + ToLibrary -> return () + ToExecutable -> do + ret <- liftIO $ runProgramWithExitCode "csc" + ["-out:" ++ outpath + , "-lib:"++mono_libs + , "-r:Cloo.clSharp.dll" + , "-r:Mono.Options.dll" + , cspath + , "/unsafe"] "" + case ret of + Left err -> + externalErrorS $ "Failed to run csc: " ++ show err + Right (ExitFailure code, cscwarn, cscerr) -> + externalErrorS $ "csc failed with code " ++ show code ++ ":\n" ++ cscerr ++ cscwarn + Right (ExitSuccess, _, _) -> liftIO $ do + perms <- liftIO $ getPermissions outpath + setPermissions outpath $ setOwnerExecutable True perms diff --git a/src/futhark-csopencl.hs b/src/futhark-csopencl.hs new file mode 100644 --- /dev/null +++ b/src/futhark-csopencl.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE FlexibleContexts #-} +module Main (main) where + +import Control.Monad.IO.Class +import Data.Maybe (fromMaybe) +import System.Directory +import System.Environment +import System.Exit +import System.FilePath + +import Futhark.Pipeline +import Futhark.Passes +import qualified Futhark.CodeGen.Backends.CSOpenCL as CSOpenCL +import Futhark.Util.Pretty (prettyText) +import Futhark.Compiler.CLI +import Futhark.Util + +main :: IO () +main = compilerMain () [] + "Compile OpenCL C#" "Generate OpenCL C# code from optimised Futhark program." + gpuPipeline $ \() mode outpath prog -> do + mono_libs <- liftIO $ fromMaybe "." <$> lookupEnv "MONO_PATH" + + let class_name = + case mode of ToLibrary -> Just $ takeBaseName outpath + ToExecutable -> Nothing + csprog <- either (`internalError` prettyText prog) return =<< + CSOpenCL.compileProg class_name prog + + let cspath = outpath `addExtension` "cs" + liftIO $ writeFile cspath csprog + + case mode of + ToLibrary -> return () + ToExecutable -> do + ret <- liftIO $ runProgramWithExitCode "csc" + ["-out:" ++ outpath, "-lib:"++mono_libs, "-r:Cloo.clSharp.dll,Mono.Options.dll", cspath, "/unsafe"] "" + case ret of + Left err -> + externalErrorS $ "Failed to run csc: " ++ show err + Right (ExitFailure code, cscwarn, cscerr) -> + externalErrorS $ "csc failed with code " ++ show code ++ ":\n" ++ cscerr ++ cscwarn + Right (ExitSuccess, _, _) -> liftIO $ do + perms <- liftIO $ getPermissions outpath + setPermissions outpath $ setOwnerExecutable True perms diff --git a/src/futhark-dataset.hs b/src/futhark-dataset.hs new file mode 100644 --- /dev/null +++ b/src/futhark-dataset.hs @@ -0,0 +1,348 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | Randomly generate Futhark input files containing values of a +-- specified type and shape. +module Main (main) where + +import Control.Arrow (first) +import Control.Monad +import Control.Monad.State +import qualified Data.Binary as Bin +import qualified Data.ByteString.Lazy as BS +import Data.Binary.IEEE754 +import Data.Binary.Put +import qualified Data.ByteString.Lazy as BL +import qualified Data.Map.Strict as M +import Data.List +import qualified Data.Text as T +import Data.Word + +import System.Console.GetOpt +import System.Random + +import Language.Futhark.Syntax +import Language.Futhark.Attributes (UncheckedTypeExp, namesToPrimTypes) +import Language.Futhark.Parser +import Language.Futhark.Pretty () + +import Futhark.Test.Values +import Futhark.Util.Options + +main :: IO () +main = mainWithOptions initialDataOptions commandLineOptions "options..." f + where f [] config + | null $ optOrders config = Just $ do + maybe_vs <- readValues <$> BS.getContents + case maybe_vs of + Nothing -> error "Malformed data on standard input." + Just vs -> + case format config of + Text -> mapM_ (putStrLn . pretty) vs + Binary -> mapM_ (BS.putStr . Bin.encode) vs + Type -> mapM_ (putStrLn . valueType) vs + | otherwise = + Just $ zipWithM_ ($) (optOrders config) $ map mkStdGen [optSeed config..] + f _ _ = + Nothing + +data OutputFormat = Text + | Binary + | Type + deriving (Eq, Ord, Show) + +data DataOptions = DataOptions + { optSeed :: Int + , optRange :: RandomConfiguration + , optOrders :: [StdGen -> IO ()] + , format :: OutputFormat + } + +initialDataOptions :: DataOptions +initialDataOptions = DataOptions 0 initialRandomConfiguration [] Text + +commandLineOptions :: [FunOptDescr DataOptions] +commandLineOptions = [ + Option "s" ["seed"] + (ReqArg (\n -> + case reads n of + [(n', "")] -> + Right $ \config -> config { optSeed = n' } + _ -> + Left $ error $ "'" ++ n ++ "' is not an integer.") + "SEED") + "The seed to use when initialising the RNG." + , Option "g" ["generate"] + (ReqArg (\t -> + case tryMakeGenerator t of + Right g -> + Right $ \config -> + config { optOrders = + optOrders config ++ + [g (optRange config) (format config)] + } + Left err -> + Left $ error err) + "TYPE") + "Generate a random value of this type." + , Option [] ["text"] + (NoArg $ Right $ \opts -> opts { format = Text }) + "Output data in text format (must precede --generate)." + , Option "b" ["binary"] + (NoArg $ Right $ \opts -> opts { format = Binary }) + "Output data in binary Futhark format (must precede --generate)." + , Option "t" ["type"] + (NoArg $ Right $ \opts -> opts { format = Type }) + "Output the type (textually) rather than the value (must precede --generate)." + , setRangeOption "i8" seti8Range + , setRangeOption "i16" seti16Range + , setRangeOption "i32" seti32Range + , setRangeOption "i64" seti64Range + , setRangeOption "u8" setu8Range + , setRangeOption "u16" setu16Range + , setRangeOption "u32" setu32Range + , setRangeOption "u64" setu64Range + , setRangeOption "f32" setf32Range + , setRangeOption "f64" setf64Range + ] + +setRangeOption :: Read a => String + -> (Range a -> RandomConfiguration -> RandomConfiguration) + -> FunOptDescr DataOptions +setRangeOption tname set = + Option "" [name] + (ReqArg (\b -> + let (lower,rest) = span (/=':') b + upper = drop 1 rest + in case (reads lower, reads upper) of + ([(lower', "")], [(upper', "")]) -> + Right $ \config -> + config { optRange = set (lower', upper') $ optRange config } + _ -> + Left $ error $ "Invalid bounds: " ++ b + ) + "MIN:MAX") $ + "Range of " ++ tname ++ " values." + where name = tname ++ "-bounds" + +tryMakeGenerator :: String -> Either String (RandomConfiguration -> OutputFormat -> StdGen -> IO ()) +tryMakeGenerator t = do + t' <- toSimpleType =<< either (Left . show) Right (parseType name (T.pack t)) + return $ \conf fmt stdgen -> do + let (v, _) = randomValue conf t' stdgen + case fmt of + Text -> printSimpleValueT v + Binary -> printSimpleValueB t' v + Type -> putStrLn t + where name = "option " ++ t + +data SimpleType = SimpleArray SimpleType Int + | SimplePrim PrimType + deriving (Show) + +toSimpleType :: UncheckedTypeExp -> Either String SimpleType +toSimpleType TETuple{} = Left "Cannot handle tuples yet." +toSimpleType TERecord{} = Left "Cannot handle records yet." +toSimpleType TEApply{} = Left "Cannot handle type applications yet." +toSimpleType TEArrow{} = Left "Cannot generate functions." +toSimpleType (TEUnique t _) = toSimpleType t +toSimpleType (TEArray t d _) = + SimpleArray <$> toSimpleType t <*> constantDim d + where constantDim (ConstDim k) = Right k + constantDim _ = Left "Array has non-constant dimension declaration." +toSimpleType (TEVar (QualName [] v) _) + | Just t <- M.lookup v namesToPrimTypes = Right $ SimplePrim t +toSimpleType (TEVar v _) = + Left $ "Unknown type " ++ pretty v + +data SimpleValue = SimpleArrayValue [SimpleValue] + | SimplePrimValue PrimValue + deriving (Show) + +-- Ordinary prettyprinting consumes too much memory, likely because it +-- manifests the string to print instead of doing it lazily, which is +-- a bad idea for giant values. This is likely because it tries to do +-- a good job with respect to line wrapping and the like. We opt to +-- do a bad job instead, but one that we can do much faster. +printSimpleValueT :: SimpleValue -> IO () +printSimpleValueT = (>>putStrLn "") . flip evalStateT 0 . p + where elements_per_line = 20 :: Int + + p (SimplePrimValue v) = do + maybeNewline + lift $ putStr $ pretty v + p (SimpleArrayValue []) = + lift $ putStr "[]" + p (SimpleArrayValue (v:vs)) = do + lift $ putStr "[" + p v + forM_ vs $ \v' -> do + lift $ putStr ", " + p v' + lift $ putStr "]" + + maybeNewline = do + i <- get + if i >= elements_per_line + then do lift $ putStrLn "" + put 0 + else put $ i + 1 + +binaryFormatVersion :: Int +binaryFormatVersion = 2 + +printSimpleValueB :: SimpleType -> SimpleValue -> IO () +printSimpleValueB st sv = + BL.putStr $ runPut $ printHeader >> pSimpleValue sv + + where + printHeader = do + Bin.put 'b' + putWord8 $ fromIntegral binaryFormatVersion + let dims = getDims st + putWord8 $ fromIntegral $ length dims + putElemType st + case sv of + SimplePrimValue _ -> return () + SimpleArrayValue _ -> mapM_ (putWord64le . fromIntegral) dims + + -- Simply calling @Bin.put (" i8" :: String)@ would cause a lot of bytes to + -- be written. Doing it this way will only write 4 bytes. + putElemType (SimplePrim (Signed Int8)) = mapM_ Bin.put (" i8" :: String) + putElemType (SimplePrim (Signed Int16)) = mapM_ Bin.put (" i16" :: String) + putElemType (SimplePrim (Signed Int32)) = mapM_ Bin.put (" i32" :: String) + putElemType (SimplePrim (Signed Int64)) = mapM_ Bin.put (" i64" :: String) + putElemType (SimplePrim (Unsigned Int8)) = mapM_ Bin.put (" u8" :: String) + putElemType (SimplePrim (Unsigned Int16)) = mapM_ Bin.put (" u16" :: String) + putElemType (SimplePrim (Unsigned Int32)) = mapM_ Bin.put (" u32" :: String) + putElemType (SimplePrim (Unsigned Int64)) = mapM_ Bin.put (" u64" :: String) + putElemType (SimplePrim (FloatType Float32)) = mapM_ Bin.put (" f32" :: String) + putElemType (SimplePrim (FloatType Float64)) = mapM_ Bin.put (" f64" :: String) + putElemType (SimplePrim Bool) = mapM_ Bin.put ("bool" :: String) + putElemType (SimpleArray ty _) = putElemType ty + + getDims (SimplePrim _) = [] + getDims (SimpleArray ty dim) = dim : getDims ty + + pSimpleValue :: SimpleValue -> Put + pSimpleValue (SimplePrimValue pv) = p pv + pSimpleValue (SimpleArrayValue svs) = mapM_ pSimpleValue svs + + p :: PrimValue -> Put + p (SignedValue (Int8Value v)) = putWord8 $ fromIntegral $ fromEnum v + p (SignedValue (Int16Value v)) = putWord16le $ fromIntegral $ fromEnum v + p (SignedValue (Int32Value v)) = putWord32le $ fromIntegral $ fromEnum v + p (SignedValue (Int64Value v)) = putWord64le $ fromIntegral $ fromEnum v + p (UnsignedValue (Int8Value v)) = putWord8 $ fromIntegral $ fromEnum v + p (UnsignedValue (Int16Value v)) = putWord16le $ fromIntegral $ fromEnum v + p (UnsignedValue (Int32Value v)) = putWord32le $ fromIntegral $ fromEnum v + p (UnsignedValue (Int64Value v)) = putWord64le $ fromIntegral $ fromEnum v + p (FloatValue (Float32Value v)) = putFloat32le v + p (FloatValue (Float64Value v)) = putFloat64le v + p (BoolValue v) = putWord8 $ if v then 1 else 0 + +-- | Closed interval, as in @System.Random@. +type Range a = (a, a) + +data RandomConfiguration = RandomConfiguration + { i8Range :: Range Int8 + , i16Range :: Range Int16 + , i32Range :: Range Int32 + , i64Range :: Range Int64 + , u8Range :: Range Word8 + , u16Range :: Range Word16 + , u32Range :: Range Word32 + , u64Range :: Range Word64 + , f32Range :: Range Float + , f64Range :: Range Double + } + +-- The following lines provide evidence about how Haskells record +-- system sucks. +seti8Range :: Range Int8 -> RandomConfiguration -> RandomConfiguration +seti8Range bounds config = config { i8Range = bounds } +seti16Range :: Range Int16 -> RandomConfiguration -> RandomConfiguration +seti16Range bounds config = config { i16Range = bounds } +seti32Range :: Range Int32 -> RandomConfiguration -> RandomConfiguration +seti32Range bounds config = config { i32Range = bounds } +seti64Range :: Range Int64 -> RandomConfiguration -> RandomConfiguration +seti64Range bounds config = config { i64Range = bounds } + +setu8Range :: Range Word8 -> RandomConfiguration -> RandomConfiguration +setu8Range bounds config = config { u8Range = bounds } +setu16Range :: Range Word16 -> RandomConfiguration -> RandomConfiguration +setu16Range bounds config = config { u16Range = bounds } +setu32Range :: Range Word32 -> RandomConfiguration -> RandomConfiguration +setu32Range bounds config = config { u32Range = bounds } +setu64Range :: Range Word64 -> RandomConfiguration -> RandomConfiguration +setu64Range bounds config = config { u64Range = bounds } + +setf32Range :: Range Float -> RandomConfiguration -> RandomConfiguration +setf32Range bounds config = config { f32Range = bounds } +setf64Range :: Range Double -> RandomConfiguration -> RandomConfiguration +setf64Range bounds config = config { f64Range = bounds } + +initialRandomConfiguration :: RandomConfiguration +initialRandomConfiguration = RandomConfiguration + (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) + (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) (minBound, maxBound) + (0.0, 1.0) (0.0, 1.0) + +randomValue :: RandomConfiguration -> SimpleType -> StdGen -> (SimpleValue, StdGen) +randomValue conf (SimplePrim (Signed Int8)) stdgen = + randomC conf i8Range stdgen +randomValue conf (SimplePrim (Signed Int16)) stdgen = + randomC conf i16Range stdgen +randomValue conf (SimplePrim (Signed Int32)) stdgen = + randomC conf i32Range stdgen +randomValue conf (SimplePrim (Signed Int64)) stdgen = + randomC conf i64Range stdgen + +randomValue conf (SimplePrim (Unsigned Int8)) stdgen = + randomC conf u8Range stdgen +randomValue conf (SimplePrim (Unsigned Int16)) stdgen = + randomC conf u16Range stdgen +randomValue conf (SimplePrim (Unsigned Int32)) stdgen = + randomC conf u32Range stdgen +randomValue conf (SimplePrim (Unsigned Int64)) stdgen = + randomC conf u64Range stdgen + +randomValue _ (SimplePrim Bool) stdgen = + first (SimplePrimValue . BoolValue) $ random stdgen + +randomValue conf (SimplePrim (FloatType Float32)) stdgen = + randomC conf f32Range stdgen +randomValue conf (SimplePrim (FloatType Float64)) stdgen = + randomC conf f64Range stdgen + +randomValue conf (SimpleArray t d) stdgen = + first SimpleArrayValue $ uncurry (flip (,)) $ + mapAccumL f stdgen [0..d-1] + where f stdgen' _ = uncurry (flip (,)) $ randomValue conf t stdgen' + +class ToFuthark a where + toFuthark :: a -> SimpleValue + +instance ToFuthark Int8 where + toFuthark = SimplePrimValue . SignedValue . Int8Value +instance ToFuthark Int16 where + toFuthark = SimplePrimValue . SignedValue . Int16Value +instance ToFuthark Int32 where + toFuthark = SimplePrimValue . SignedValue . Int32Value +instance ToFuthark Int64 where + toFuthark = SimplePrimValue . SignedValue . Int64Value +instance ToFuthark Word8 where + toFuthark = SimplePrimValue . UnsignedValue . Int8Value . fromIntegral +instance ToFuthark Word16 where + toFuthark = SimplePrimValue . UnsignedValue . Int16Value . fromIntegral +instance ToFuthark Word32 where + toFuthark = SimplePrimValue . UnsignedValue . Int32Value . fromIntegral +instance ToFuthark Word64 where + toFuthark = SimplePrimValue . UnsignedValue . Int64Value . fromIntegral +instance ToFuthark Float where + toFuthark = SimplePrimValue . FloatValue . Float32Value +instance ToFuthark Double where + toFuthark = SimplePrimValue . FloatValue . Float64Value + +randomC :: (ToFuthark a, Random a) => + RandomConfiguration -> (RandomConfiguration -> Range a) -> StdGen + -> (SimpleValue, StdGen) +randomC conf pick = first toFuthark . randomR (pick conf) diff --git a/src/futhark-doc.hs b/src/futhark-doc.hs new file mode 100644 --- /dev/null +++ b/src/futhark-doc.hs @@ -0,0 +1,105 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE OverloadedStrings #-} +module Main (main) where + +import Control.Monad.IO.Class (liftIO) +import Control.Monad.State +import Data.FileEmbed +import Data.List +import Data.Semigroup ((<>)) +import System.FilePath +import System.Directory (createDirectoryIfMissing) +import System.Console.GetOpt +import System.IO +import System.Exit +import qualified Data.Text.Lazy as T +import qualified Data.Text.Lazy.IO as T +import Text.Blaze.Html.Renderer.Text + +import Futhark.Doc.Generator +import Futhark.Compiler (readLibrary, dumpError, newFutharkConfig, Imports, fileProg) +import Futhark.Pipeline (runFutharkM, FutharkM, Verbosity(..)) +import Language.Futhark.Syntax (progDoc, DocComment(..)) +import Futhark.Util.Options +import Futhark.Util (directoryContents, trim) + +main :: IO () +main = mainWithOptions initialDocConfig commandLineOptions "options... -o outdir programs..." f + where f [dir] config = Just $ do + res <- runFutharkM (m config dir) Verbose + case res of + Left err -> liftIO $ do + dumpError newFutharkConfig err + exitWith $ ExitFailure 2 + Right () -> + return () + f _ _ = Nothing + + m :: DocConfig -> FilePath -> FutharkM () + m config dir = + case docOutput config of + Nothing -> liftIO $ do + hPutStrLn stderr "Must specify output directory with -o." + exitWith $ ExitFailure 1 + Just outdir -> do + files <- liftIO $ futFiles dir + when (docVerbose config) $ liftIO $ do + mapM_ (hPutStrLn stderr . ("Found source file "<>)) files + hPutStrLn stderr "Reading files..." + (_w, imports, _vns) <- readLibrary files + liftIO $ printDecs config outdir files $ nubBy sameImport imports + + sameImport (x, _) (y, _) = x == y + +futFiles :: FilePath -> IO [FilePath] +futFiles dir = filter isFut <$> directoryContents dir + where isFut = (==".fut") . takeExtension + +printDecs :: DocConfig -> FilePath -> [FilePath] -> Imports -> IO () +printDecs cfg dir files imports = do + let direct_imports = map (normalise . dropExtension) files + (file_htmls, _warnings) = renderFiles direct_imports $ + filter (not . ignored) imports + mapM_ (write . fmap renderHtml) file_htmls + write ("style.css", cssFile) + + where write :: (String, T.Text) -> IO () + write (name, content) = do let file = dir makeRelative "/" name + when (docVerbose cfg) $ + hPutStrLn stderr $ "Writing " <> file + createDirectoryIfMissing True $ takeDirectory file + T.writeFile file content + + -- Some files are not worth documenting; typically because + -- they contain tests. The current crude mechanism is to + -- recognise them by a file comment containing "ignore". + ignored (_, fm) = + case progDoc (fileProg fm) of + Just (DocComment s _) -> trim s == "ignore" + _ -> False + +cssFile :: T.Text +cssFile = $(embedStringFile "rts/futhark-doc/style.css") + +data DocConfig = DocConfig { docOutput :: Maybe FilePath + , docVerbose :: Bool + } + +initialDocConfig :: DocConfig +initialDocConfig = DocConfig { docOutput = Nothing + , docVerbose = False + } + +type DocOption = OptDescr (Either (IO ()) (DocConfig -> DocConfig)) + +commandLineOptions :: [DocOption] +commandLineOptions = [ Option "o" ["output-directory"] + (ReqArg (\dirname -> Right $ \config -> config { docOutput = Just dirname }) + "DIR") + "Directory in which to put generated documentation." + , Option "v" ["verbose"] + (NoArg $ Right $ \config -> config { docVerbose = True }) + "Print status messages on stderr." + ] diff --git a/src/futhark-opencl.hs b/src/futhark-opencl.hs new file mode 100644 --- /dev/null +++ b/src/futhark-opencl.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE FlexibleContexts #-} +module Main (main) where + +import Control.Monad.IO.Class +import System.FilePath +import System.Exit +import qualified System.Info + +import Futhark.Pipeline +import Futhark.Passes +import qualified Futhark.CodeGen.Backends.COpenCL as COpenCL +import Futhark.Util +import Futhark.Util.Pretty (prettyText) +import Futhark.Compiler.CLI + +main :: IO () +main = compilerMain () [] + "Compile OpenCL" "Generate OpenCL/C code from optimised Futhark program." + gpuPipeline $ \() mode outpath prog -> do + cprog <- either (`internalError` prettyText prog) return =<< + COpenCL.compileProg prog + let cpath = outpath `addExtension` "c" + hpath = outpath `addExtension` "h" + extra_options + | System.Info.os == "darwin" = + ["-framework", "OpenCL"] + | System.Info.os == "mingw32" = + ["-lOpenCL64"] + | otherwise = + ["-lOpenCL"] + + case mode of + ToLibrary -> do + let (header, impl) = COpenCL.asLibrary cprog + liftIO $ writeFile hpath header + liftIO $ writeFile cpath impl + ToExecutable -> do + liftIO $ writeFile cpath $ COpenCL.asExecutable cprog + ret <- liftIO $ runProgramWithExitCode "gcc" + ([cpath, "-O3", "-std=c99", "-lm", "-o", outpath] ++ extra_options) "" + case ret of + Left err -> + externalErrorS $ "Failed to run gcc: " ++ show err + Right (ExitFailure code, _, gccerr) -> + externalErrorS $ "gcc failed with code " ++ + show code ++ ":\n" ++ gccerr + Right (ExitSuccess, _, _) -> + return () diff --git a/src/futhark-pkg.hs b/src/futhark-pkg.hs new file mode 100644 --- /dev/null +++ b/src/futhark-pkg.hs @@ -0,0 +1,387 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Main (main) where + +import Control.Monad.IO.Class +import Control.Monad.State +import Control.Monad.Reader +import Data.Maybe +import qualified Data.Map as M +import qualified Data.Text as T +import qualified Data.Text.IO as T +import qualified Data.ByteString.Lazy as LBS +import Data.List +import Data.Monoid +import System.Directory +import System.FilePath +import qualified System.FilePath.Posix as Posix +import System.Environment +import System.Exit +import System.IO +import System.Console.GetOpt + +import qualified Codec.Archive.Zip as Zip +import Network.HTTP.Client +import Network.HTTP.Client.TLS + +import Prelude + +import Futhark.Util.Options +import Futhark.Pkg.Types +import Futhark.Pkg.Info +import Futhark.Pkg.Solve +import Futhark.Util (directoryContents) +import Futhark.Util.Log + +--- Installing packages + +installInDir :: BuildList -> FilePath -> PkgM () +installInDir (BuildList bl) dir = do + let putEntry from_dir pdir entry + -- The archive may contain all kinds of other stuff that we don't want. + | not (isInPkgDir from_dir $ Zip.eRelativePath entry) + || hasTrailingPathSeparator (Zip.eRelativePath entry) = return () + | otherwise = do + -- Since we are writing to paths indicated in a zipfile we + -- downloaded from the wild Internet, we are going to be a + -- little bit paranoid. Specifically, we want to avoid + -- writing outside of the 'lib/' directory. We do this by + -- bailing out if the path contains any '..' components. We + -- have to use System.FilePath.Posix, because the zip library + -- claims to encode filepaths with '/' directory seperators no + -- matter the host OS. + when (".." `elem` Posix.splitPath (Zip.eRelativePath entry)) $ + fail $ "Zip archive for " <> pdir <> " contains suspicuous path: " <> + Zip.eRelativePath entry + let f = pdir makeRelative from_dir (Zip.eRelativePath entry) + createDirectoryIfMissing True $ takeDirectory f + LBS.writeFile f $ Zip.fromEntry entry + + isInPkgDir from_dir f = + Posix.splitPath from_dir `isPrefixOf` Posix.splitPath f + + forM_ (M.toList bl) $ \(p, v) -> do + info <- lookupPackageRev p v + a <- downloadZipball $ pkgRevZipballUrl info + m <- getManifest $ pkgRevGetManifest info + + -- Compute the directory in the zipball that should contain the + -- package files. + let noPkgDir = fail $ "futhark.pkg for " ++ T.unpack p ++ "-" ++ + T.unpack (prettySemVer v) ++ " does not define a package path." + from_dir <- maybe noPkgDir (return . (pkgRevZipballDir info <>)) $ pkgDir m + + -- The directory in the local file system that will contain the + -- package files. + let pdir = dir T.unpack p + -- Remove any existing directory for this package. This is a bit + -- inefficient, as the likelihood that the old ``lib`` directory + -- already contains the correct version is rather high. We should + -- have a way to recognise this situation, and not download the + -- zipball in that case. + liftIO $ removePathForcibly pdir + liftIO $ createDirectoryIfMissing True pdir + + liftIO $ mapM_ (putEntry from_dir pdir) $ Zip.zEntries a + +libDir, libNewDir, libOldDir :: FilePath +(libDir, libNewDir, libOldDir) = ("lib", "lib~new", "lib~old") + +-- | Install the packages listed in the build list in the 'lib' +-- directory of the current working directory. Since we are touching +-- the file system, we are going to be very paranoid. In particular, +-- we want to avoid corrupting the 'lib' directory if something fails +-- along the way. +-- +-- The procedure is as follows: +-- +-- 1) Create a directory 'lib~new'. Delete an existing 'lib~new' if +-- necessary. +-- +-- 2) Populate 'lib~new' based on the build list. +-- +-- 3) Rename 'lib' to 'lib~old'. Delete an existing 'lib~old' if +-- necessary. +-- +-- 4) Rename 'lib~new' to 'lib' +-- +-- 5) If the current package has package path 'p', move 'lib~old/p' to +-- 'lib~new/p'. +-- +-- 6) Delete 'lib~old'. +-- +-- Since POSIX at least guarantees atomic renames, the only place this +-- can fail is between steps 3, 4, and 5. In that case, at least the +-- 'lib~old' will still exist and can be put back by the user. +installBuildList :: Maybe PkgPath -> BuildList -> PkgM () +installBuildList p bl = do + libdir_exists <- liftIO $ doesDirectoryExist libDir + + -- 1 + liftIO $ do removePathForcibly libNewDir + createDirectoryIfMissing False libNewDir + + -- 2 + installInDir bl libNewDir + + -- 3 + when libdir_exists $ liftIO $ do + removePathForcibly libOldDir + renameDirectory libDir libOldDir + + -- 4 + liftIO $ renameDirectory libNewDir libDir + + -- 5 + case pkgPathFilePath <$> p of + Just pfp | libdir_exists -> liftIO $ do + pkgdir_exists <- doesDirectoryExist $ libOldDir pfp + when pkgdir_exists $ do + -- Ensure the parent directories exist so that we can move the + -- package directory directly. + createDirectoryIfMissing True $ takeDirectory $ libDir pfp + renameDirectory (libOldDir pfp) (libDir pfp) + _ -> return () + + -- 6 + when libdir_exists $ liftIO $ removePathForcibly libOldDir + +getPkgManifest :: PkgM PkgManifest +getPkgManifest = do + file_exists <- liftIO $ doesFileExist futharkPkg + dir_exists <- liftIO $ doesDirectoryExist futharkPkg + + case (file_exists, dir_exists) of + (True, _) -> liftIO $ parsePkgManifestFromFile futharkPkg + (_, True) -> fail $ futharkPkg <> + " exists, but it is a directory! What in Odin's beard..." + _ -> liftIO $ do T.putStrLn $ T.pack futharkPkg <> " not found - pretending it's empty." + return $ newPkgManifest Nothing + +putPkgManifest :: PkgManifest -> PkgM () +putPkgManifest = liftIO . T.writeFile futharkPkg . prettyPkgManifest + +--- The CLI + +newtype PkgConfig = PkgConfig { pkgVerbose :: Bool } + +-- | The monad in which futhark-pkg runs. +newtype PkgM a = PkgM { unPkgM :: ReaderT PkgConfig (StateT (PkgRegistry PkgM) IO) a } + deriving (Functor, Applicative, MonadIO, MonadReader PkgConfig) + +instance Monad PkgM where + PkgM m >>= f = PkgM $ m >>= unPkgM . f + return = PkgM . return + fail s = liftIO $ do + prog <- getProgName + putStrLn $ prog ++ ": " ++ s + exitFailure + +instance MonadPkgRegistry PkgM where + putPkgRegistry = PkgM . put + getPkgRegistry = PkgM get + +instance MonadLogger PkgM where + addLog l = do + verbose <- asks pkgVerbose + when verbose $ liftIO $ T.hPutStr stderr $ toText l + +runPkgM :: PkgConfig -> PkgM a -> IO a +runPkgM cfg (PkgM m) = evalStateT (runReaderT m cfg) mempty + +cmdMain :: String -> ([String] -> PkgConfig -> Maybe (IO ())) -> IO () +cmdMain = mainWithOptions (PkgConfig False) options + where options = [ Option "v" ["verbose"] + (NoArg $ Right $ \cfg -> cfg { pkgVerbose = True }) + "Write running diagnostics to stderr."] + +doFmt :: IO () +doFmt = mainWithOptions () [] "fmt" $ \args () -> + case args of + [] -> Just $ do + m <- parsePkgManifestFromFile futharkPkg + T.writeFile futharkPkg $ prettyPkgManifest m + _ -> Nothing + +doCheck :: IO () +doCheck = cmdMain "check" $ \args cfg -> + case args of + [] -> Just $ runPkgM cfg $ do + m <- getPkgManifest + bl <- solveDeps $ pkgRevDeps m + + liftIO $ T.putStrLn "Dependencies chosen:" + liftIO $ T.putStr $ prettyBuildList bl + + case commented $ manifestPkgPath m of + Nothing -> return () + Just p -> do + let pdir = "lib" T.unpack p + + pdir_exists <- liftIO $ doesDirectoryExist pdir + + unless pdir_exists $ liftIO $ do + T.putStrLn $ "Problem: the directory " <> T.pack pdir <> " does not exist." + exitFailure + + anything <- liftIO $ any ((==".fut") . takeExtension) <$> + directoryContents ("lib" T.unpack p) + unless anything $ liftIO $ do + T.putStrLn $ "Problem: the directory " <> T.pack pdir <> " does not contain any .fut files." + exitFailure + _ -> Nothing + +doSync :: IO () +doSync = cmdMain "sync" $ \args cfg -> + case args of + [] -> Just $ runPkgM cfg $ do + m <- getPkgManifest + bl <- solveDeps $ pkgRevDeps m + installBuildList (commented $ manifestPkgPath m) bl + _ -> Nothing + +doAdd :: IO () +doAdd = cmdMain "add PKGPATH" $ \args cfg -> + case args of + [p, v] | Right v' <- parseVersion $ T.pack v -> Just $ runPkgM cfg $ doAdd' (T.pack p) v' + [p] -> Just $ runPkgM cfg $ + -- Look up the newest revision of the package. + doAdd' (T.pack p) =<< lookupNewestRev (T.pack p) + _ -> Nothing + + where + doAdd' p v = do + m <- getPkgManifest + + -- See if this package (and its dependencies) even exists. We + -- do this by running the solver with the dependencies already + -- in the manifest, plus this new one. The Monoid instance for + -- PkgRevDeps is left-biased, so we are careful to use the new + -- version for this package. + _ <- solveDeps $ PkgRevDeps (M.singleton p (v, Nothing)) <> pkgRevDeps m + + -- We either replace any existing occurence of package 'p', or + -- we add a new one. + p_info <- lookupPackageRev p v + let hash = case (_svMajor v, _svMinor v, _svPatch v) of + -- We do not perform hash-pinning for + -- (0,0,0)-versions, because these already embed a + -- specific revision ID into their version number. + (0, 0, 0) -> Nothing + _ -> Just $ pkgRevCommit p_info + req = Required p v hash + (m', prev_r) = addRequiredToManifest req m + + case prev_r of + Just prev_r' + | requiredPkgRev prev_r' == v -> + liftIO $ T.putStrLn $ "Package already at version " <> prettySemVer v <> "; nothing to do." + | otherwise -> + liftIO $ T.putStrLn $ "Replaced " <> p <> " " <> + prettySemVer (requiredPkgRev prev_r') <> " => " <> prettySemVer v <> "." + Nothing -> + liftIO $ T.putStrLn $ "Added new required package " <> p <> " " <> prettySemVer v <> "." + putPkgManifest m' + liftIO $ T.putStrLn "Remember to run 'futhark-pkg sync'." + +doRemove :: IO () +doRemove = cmdMain "remove PKGPATH" $ \args cfg -> + case args of + [p] -> Just $ runPkgM cfg $ doRemove' $ T.pack p + _ -> Nothing + where + doRemove' p = do + m <- getPkgManifest + case removeRequiredFromManifest p m of + Nothing -> liftIO $ do + T.putStrLn $ "No package " <> p <> " found in " <> T.pack futharkPkg <> "." + exitFailure + Just (m', r) -> do + putPkgManifest m' + liftIO $ T.putStrLn $ "Removed " <> p <> " " <> prettySemVer (requiredPkgRev r) <> "." + +doInit :: IO () +doInit = cmdMain "create PKGPATH" $ \args cfg -> + case args of + [p] -> Just $ runPkgM cfg $ doCreate' $ T.pack p + _ -> Nothing + where + doCreate' p = do + exists <- liftIO $ (||) <$> doesFileExist futharkPkg <*> doesDirectoryExist futharkPkg + when exists $ liftIO $ do + T.putStrLn $ T.pack futharkPkg <> " already exists." + exitFailure + + liftIO $ createDirectoryIfMissing True $ "lib" T.unpack p + liftIO $ T.putStrLn $ "Created directory " <> T.pack ("lib" T.unpack p) <> "." + + putPkgManifest $ newPkgManifest $ Just p + liftIO $ T.putStrLn $ "Wrote " <> T.pack futharkPkg <> "." + +doUpgrade :: IO () +doUpgrade = cmdMain "upgrade" $ \args cfg -> + case args of + [] -> Just $ runPkgM cfg $ do + m <- getPkgManifest + rs <- traverse (mapM (traverse upgrade)) $ manifestRequire m + putPkgManifest m { manifestRequire = rs } + _ -> Nothing + where upgrade req = do + v <- lookupNewestRev $ requiredPkg req + h <- pkgRevCommit <$> lookupPackageRev (requiredPkg req) v + + when (v /= requiredPkgRev req) $ + liftIO $ T.putStrLn $ "Upgraded " <> requiredPkg req <> " " <> + prettySemVer (requiredPkgRev req) <> " => " <> prettySemVer v <> "." + + return req { requiredPkgRev = v + , requiredHash = Just h } + +doVersions :: IO () +doVersions = cmdMain "versions PKGPATH" $ \args cfg -> + case args of + [p] -> Just $ runPkgM cfg $ doVersions' $ T.pack p + _ -> Nothing + where doVersions' = + mapM_ (liftIO . T.putStrLn . prettySemVer) . M.keys . pkgVersions + <=< lookupPackage + +main :: IO () +main = do + -- Ensure that we can make HTTPS requests. + setGlobalManager =<< newManager tlsManagerSettings + + -- Avoid Git asking for credentials. We prefer failure. + liftIO $ setEnv "GIT_TERMINAL_PROMPT" "0" + + args <- getArgs + let commands = [ ("add", + (doAdd, "Add another required package to futhark.pkg.")) + , ("check", + (doCheck, "Check that futhark.pkg is satisfiable.")) + , ("init", + (doInit, "Create a new futhark.pkg and a lib/ skeleton.")) + , ("fmt", + (doFmt, "Reformat futhark.pkg.")) + , ("sync", + (doSync, "Populate lib/ as specified by futhark.pkg.")) + , ("remove", + (doRemove, "Remove a required package from futhark.pkg.")) + , ("upgrade", + (doUpgrade, "Upgrade all packages to newest versions.")) + , ("versions", + (doVersions, "List available versions for a package.")) + ] + usage = "options... <" <> intercalate "|" (map fst commands) <> ">" + case args of + cmd : args' | Just (m, _) <- lookup cmd commands -> withArgs args' m + _ -> mainWithOptions () [] usage $ \_ () -> Just $ do + let k = maximum (map (length . fst) commands) + 3 + usageMsg $ T.unlines $ + [" ...:", "", "Commands:"] ++ + [ " " <> T.pack cmd <> T.pack (replicate (k - length cmd) ' ') <> desc + | (cmd, (_, desc)) <- commands ] + + where usageMsg s = do + T.putStrLn $ "Usage: futhark-pkg [--version] [--help] " <> s + exitFailure diff --git a/src/futhark-py.hs b/src/futhark-py.hs new file mode 100644 --- /dev/null +++ b/src/futhark-py.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE FlexibleContexts #-} +module Main (main) where + +import Control.Monad.IO.Class +import System.FilePath +import System.Directory + +import Futhark.Pipeline +import Futhark.Passes +import qualified Futhark.CodeGen.Backends.SequentialPython as SequentialPy +import Futhark.Util.Pretty (prettyText) +import Futhark.Compiler.CLI + +main :: IO () +main = compilerMain () [] + "Compile sequential Python" "Generate sequential Python code from optimised Futhark program." + sequentialCpuPipeline $ \() mode outpath prog -> do + let class_name = + case mode of ToLibrary -> Just $ takeBaseName outpath + ToExecutable -> Nothing + pyprog <- either (`internalError` prettyText prog) return =<< + SequentialPy.compileProg class_name prog + + case mode of + ToLibrary -> + liftIO $ writeFile (outpath `addExtension` "py") pyprog + ToExecutable -> liftIO $ do + writeFile outpath pyprog + perms <- liftIO $ getPermissions outpath + setPermissions outpath $ setOwnerExecutable True perms diff --git a/src/futhark-pyopencl.hs b/src/futhark-pyopencl.hs new file mode 100644 --- /dev/null +++ b/src/futhark-pyopencl.hs @@ -0,0 +1,30 @@ +{-# LANGUAGE FlexibleContexts #-} +module Main (main) where + +import Control.Monad.IO.Class +import System.FilePath +import System.Directory + +import Futhark.Pipeline +import Futhark.Passes +import qualified Futhark.CodeGen.Backends.PyOpenCL as PyOpenCL +import Futhark.Util.Pretty (prettyText) +import Futhark.Compiler.CLI + +main :: IO () +main = compilerMain () [] + "Compile PyOpenCL" "Generate Python + OpenCL code from optimised Futhark program." + gpuPipeline $ \() mode outpath prog -> do + let class_name = + case mode of ToLibrary -> Just $ takeBaseName outpath + ToExecutable -> Nothing + pyprog <- either (`internalError` prettyText prog) return =<< + PyOpenCL.compileProg class_name prog + + case mode of + ToLibrary -> + liftIO $ writeFile (outpath `addExtension` "py") pyprog + ToExecutable -> liftIO $ do + writeFile outpath pyprog + perms <- liftIO $ getPermissions outpath + setPermissions outpath $ setOwnerExecutable True perms diff --git a/src/futhark-test.hs b/src/futhark-test.hs new file mode 100644 --- /dev/null +++ b/src/futhark-test.hs @@ -0,0 +1,590 @@ +{-# LANGUAGE OverloadedStrings, FlexibleContexts, LambdaCase #-} +-- | This program is a convenience utility for running the Futhark +-- test suite, and its test programs. +module Main (main) where + +import Control.Applicative.Lift (runErrors, failure, Errors, Lift(..)) +import Control.Concurrent +import Control.Exception +import Control.Monad +import Control.Monad.Except hiding (throwError) +import qualified Control.Monad.Except as E +import qualified Data.ByteString as SBS +import qualified Data.ByteString.Lazy as LBS + +import Data.List +import Data.Semigroup ((<>)) +import qualified Data.Map.Strict as M +import qualified Data.Text as T +import qualified Data.Text.Encoding as T +import qualified Data.Text.IO as T +import System.Console.ANSI +import System.Process.ByteString (readProcessWithExitCode) +import System.Exit +import System.FilePath +import System.Console.GetOpt +import System.IO +import Text.Regex.TDFA + +import Futhark.Analysis.Metrics +import Futhark.Test +import Futhark.Util.Options +import Futhark.Util.Pretty (prettyText) +import Futhark.Util.Table + +--- Test execution + +type TestM = ExceptT [T.Text] IO + +-- Taken from transformers-0.5.5.0. +eitherToErrors :: Either e a -> Errors e a +eitherToErrors = either failure Pure + +throwError :: MonadError [e] m => e -> m a +throwError e = E.throwError [e] + +runTestM :: TestM () -> IO TestResult +runTestM = fmap (either Failure $ const Success) . runExceptT + +io :: IO a -> TestM a +io = liftIO + +context :: T.Text -> TestM a -> TestM a +context s = withExceptT $ + \case + [] -> [] + (e:es') -> (s <> ":\n" <> e):es' + +accErrors :: [TestM a] -> TestM [a] +accErrors tests = do + eithers <- lift $ mapM runExceptT tests + let errors = traverse eitherToErrors eithers + ExceptT $ return $ runErrors errors + +accErrors_ :: [TestM a] -> TestM () +accErrors_ = void . accErrors + +data TestResult = Success + | Failure [T.Text] + deriving (Eq, Show) + +data TestCase = TestCase { _testCaseMode :: TestMode + , testCaseProgram :: FilePath + , testCaseTest :: ProgramTest + , _testCasePrograms :: ProgConfig + } + deriving (Show) + +instance Eq TestCase where + x == y = testCaseProgram x == testCaseProgram y + +instance Ord TestCase where + x `compare` y = testCaseProgram x `compare` testCaseProgram y + +data RunResult = ErrorResult Int SBS.ByteString + | SuccessResult [Value] + +progNotFound :: T.Text -> T.Text +progNotFound s = s <> ": command not found" + +optimisedProgramMetrics :: ProgConfig -> StructurePipeline -> FilePath -> TestM AstMetrics +optimisedProgramMetrics programs pipeline program = + case pipeline of SOACSPipeline -> + check "-s" + KernelsPipeline -> + check "--kernels" + SequentialCpuPipeline -> + check "--cpu" + GpuPipeline -> + check "--gpu" + where check opt = do + (code, output, err) <- + io $ readProcessWithExitCode (configTypeChecker programs) [opt, "--metrics", program] "" + let output' = T.decodeUtf8 output + case code of + ExitSuccess + | [(m, [])] <- reads $ T.unpack output' -> return m + | otherwise -> throwError $ "Could not read metrics output:\n" <> output' + ExitFailure 127 -> throwError $ progNotFound $ T.pack $ configTypeChecker programs + ExitFailure _ -> throwError $ T.decodeUtf8 err + +testMetrics :: ProgConfig -> FilePath -> StructureTest -> TestM () +testMetrics programs program (StructureTest pipeline (AstMetrics expected)) = + context "Checking metrics" $ do + actual <- optimisedProgramMetrics programs pipeline program + accErrors_ $ map (ok actual) $ M.toList expected + where ok (AstMetrics metrics) (name, expected_occurences) = + case M.lookup name metrics of + Nothing + | expected_occurences > 0 -> + throwError $ name <> " should have occurred " <> T.pack (show expected_occurences) <> + " times, but did not occur at all in optimised program." + Just actual_occurences + | expected_occurences /= actual_occurences -> + throwError $ name <> " should have occurred " <> T.pack (show expected_occurences) <> + " times, but occured " <> T.pack (show actual_occurences) <> " times." + _ -> return () + +testWarnings :: [WarningTest] -> SBS.ByteString -> TestM () +testWarnings warnings futerr = accErrors_ $ map testWarning warnings + where testWarning (ExpectedWarning regex_s regex) + | not (match regex $ T.unpack $ T.decodeUtf8 futerr) = + throwError $ "Expected warning:\n " <> regex_s <> + "\nGot warnings:\n " <> T.decodeUtf8 futerr + | otherwise = return () + +runTestCase :: TestCase -> TestM () +runTestCase (TestCase mode program testcase progs) = + case testAction testcase of + + CompileTimeFailure expected_error -> do + let typeChecker = configTypeChecker progs + context ("Type-checking with " <> T.pack typeChecker) $ do + (code, _, err) <- + io $ readProcessWithExitCode typeChecker ["-t", program] "" + case code of + ExitSuccess -> throwError "Expected failure\n" + ExitFailure 127 -> throwError $ progNotFound $ T.pack typeChecker + ExitFailure 1 -> throwError $ T.decodeUtf8 err + ExitFailure _ -> checkError expected_error err + + RunCases _ _ warnings | mode == TypeCheck -> do + let typeChecker = configTypeChecker progs + options = ["-t", program] ++ configExtraCompilerOptions progs + context ("Type-checking with " <> T.pack typeChecker) $ do + (code, _, err) <- io $ readProcessWithExitCode typeChecker options "" + testWarnings warnings err + case code of + ExitSuccess -> return () + ExitFailure 127 -> throwError $ progNotFound $ T.pack typeChecker + ExitFailure _ -> throwError $ T.decodeUtf8 err + + RunCases ios structures warnings -> do + -- Compile up-front and reuse same executable for several entry points. + let compiler = configCompiler progs + interpreter = configInterpreter progs + extra_options = configExtraCompilerOptions progs + unless (mode == Interpreted) $ + context ("Compiling with " <> T.pack compiler) $ do + compileTestProgram extra_options compiler program warnings + mapM_ (testMetrics progs program) structures + unless (mode == Compile) $ + context "Running compiled program" $ + accErrors_ $ map (runCompiledEntry program progs) ios + unless (mode == Compile || mode == Compiled) $ + context ("Interpreting with " <> T.pack interpreter) $ + accErrors_ $ map (runInterpretedEntry interpreter program) ios + +runInterpretedEntry :: String -> FilePath -> InputOutputs -> TestM() +runInterpretedEntry futharki program (InputOutputs entry run_cases) = + let dir = takeDirectory program + runInterpretedCase run@(TestRun _ inputValues expectedResult index _) = + unless ("compiled" `elem` runTags run) $ + context ("Entry point: " <> entry + <> "; dataset: " <> T.pack (runDescription run)) $ do + + input <- T.unlines . map prettyText <$> getValues dir inputValues + expectedResult' <- getExpectedResult dir expectedResult + (code, output, err) <- + io $ readProcessWithExitCode futharki ["-e", T.unpack entry, program] $ + T.encodeUtf8 input + case code of + ExitFailure 127 -> throwError $ progNotFound $ T.pack futharki + + _ -> compareResult entry index program expectedResult' + =<< runResult program code output err + + in accErrors_ $ map runInterpretedCase run_cases + +runCompiledEntry :: FilePath -> ProgConfig -> InputOutputs -> TestM () +runCompiledEntry program progs (InputOutputs entry run_cases) = + -- Explicitly prefixing the current directory is necessary for + -- readProcessWithExitCode to find the binary when binOutputf has + -- no path component. + let binOutputf = dropExtension program + dir = takeDirectory program + binpath = "." binOutputf + entry_options = ["-e", T.unpack entry] + + runner = configRunner progs + extra_options = configExtraOptions progs + (to_run, to_run_args) + | null runner = (binpath, entry_options ++ extra_options) + | otherwise = (runner, binpath : entry_options ++ extra_options) + + runCompiledCase run@(TestRun _ inputValues expectedResult index _) = + context ("Entry point: " <> entry + <> "; dataset: " <> T.pack (runDescription run)) $ do + + input <- getValuesBS dir inputValues + expectedResult' <- getExpectedResult dir expectedResult + (progCode, output, progerr) <- + io $ readProcessWithExitCode to_run to_run_args $ LBS.toStrict input + compareResult entry index program expectedResult' + =<< runResult program progCode output progerr + + in context ("Running " <> T.pack (unwords $ binpath : entry_options ++ extra_options)) $ + accErrors_ $ map runCompiledCase run_cases + +checkError :: ExpectedError -> SBS.ByteString -> TestM () +checkError (ThisError regex_s regex) err + | not (match regex $ T.unpack $ T.decodeUtf8 err) = + throwError $ "Expected error:\n " <> regex_s <> + "\nGot error:\n " <> T.decodeUtf8 err +checkError _ _ = + return () + +runResult :: FilePath -> ExitCode -> SBS.ByteString -> SBS.ByteString -> TestM RunResult +runResult program ExitSuccess stdout_s _ = + case valuesFromByteString "stdout" $ LBS.fromStrict stdout_s of + Left e -> do + let actualf = program `addExtension` "actual" + io $ SBS.writeFile actualf stdout_s + throwError $ T.pack e <> "\n(See " <> T.pack actualf <> ")" + Right vs -> return $ SuccessResult vs +runResult _ (ExitFailure code) _ stderr_s = + return $ ErrorResult code stderr_s + +getExpectedResult :: MonadIO m => + FilePath -> ExpectedResult Values + -> m (ExpectedResult [Value]) +getExpectedResult dir (Succeeds (Just vals)) = Succeeds . Just <$> getValues dir vals +getExpectedResult _ (Succeeds Nothing) = return $ Succeeds Nothing +getExpectedResult _ (RunTimeFailure err) = return $ RunTimeFailure err + +compileTestProgram :: [String] -> String -> FilePath -> [WarningTest] -> TestM () +compileTestProgram extra_options futharkc program warnings = do + (futcode, _, futerr) <- io $ readProcessWithExitCode futharkc options "" + testWarnings warnings futerr + case futcode of + ExitFailure 127 -> throwError $ progNotFound $ T.pack futharkc + ExitFailure _ -> throwError $ T.decodeUtf8 futerr + ExitSuccess -> return () + where binOutputf = dropExtension program + options = [program, "-o", binOutputf] ++ extra_options + +compareResult :: T.Text -> Int -> FilePath -> ExpectedResult [Value] -> RunResult + -> TestM () +compareResult _ _ _ (Succeeds Nothing) SuccessResult{} = + return () +compareResult entry index program (Succeeds (Just expectedResult)) (SuccessResult actualResult) = + case compareValues actualResult expectedResult of + Just mismatches -> + let reportMismatch mismatch = do + let actualf = program <.> T.unpack entry <.> show index <.> "actual" + expectedf = program <.> T.unpack entry <.> show index <.> "expected" + io $ SBS.writeFile actualf $ + T.encodeUtf8 $ T.unlines $ map prettyText actualResult + io $ SBS.writeFile expectedf $ + T.encodeUtf8 $ T.unlines $ map prettyText expectedResult + throwError $ T.pack actualf <> " and " <> T.pack expectedf <> + " do not match:\n" <> T.pack (show mismatch) <> "\n" + in mapM_ reportMismatch mismatches + Nothing -> + return () +compareResult _ _ _ (RunTimeFailure expectedError) (ErrorResult _ actualError) = + checkError expectedError actualError +compareResult _ _ _ (Succeeds _) (ErrorResult code err) = + throwError $ "Program failed with error code " <> + T.pack (show code) <> " and stderr:\n " <> T.decodeUtf8 err +compareResult _ _ _ (RunTimeFailure f) (SuccessResult _) = + throwError $ "Program succeeded, but expected failure:\n " <> T.pack (show f) + +--- +--- Test manager +--- + +data TestStatus = TestStatus { testStatusRemain :: [TestCase] + , testStatusRun :: [TestCase] + , testStatusTotal :: Int + , testStatusFail :: Int + , testStatusPass :: Int + , testStatusRuns :: Int + , testStatusRunsRemain :: Int + , testStatusRunPass :: Int + , testStatusRunFail :: Int + } + +catching :: IO TestResult -> IO TestResult +catching m = m `catch` save + where save :: SomeException -> IO TestResult + save e = return $ Failure [T.pack $ show e] + +doTest :: TestCase -> IO TestResult +doTest = catching . runTestM . runTestCase + +makeTestCase :: TestConfig -> TestMode -> (FilePath, ProgramTest) -> TestCase +makeTestCase config mode (file, spec) = + TestCase mode file spec $ configPrograms config + +data ReportMsg = TestStarted TestCase + | TestDone TestCase TestResult + +runTest :: MVar TestCase -> MVar ReportMsg -> IO () +runTest testmvar resmvar = forever $ do + test <- takeMVar testmvar + putMVar resmvar $ TestStarted test + res <- doTest test + putMVar resmvar $ TestDone test res + +excludedTest :: TestConfig -> TestCase -> Bool +excludedTest config = + any (`elem` configExclude config) . testTags . testCaseTest + +statusTable :: TestStatus -> String +statusTable ts = buildTable rows 1 + where rows = + [ [ mkEntry "", passed, failed, mkEntry "remaining"] + , map mkEntry ["programs", passedProgs, failedProgs, remainProgs'] + , map mkEntry ["runs", passedRuns, failedRuns, remainRuns'] + ] + passed = ("passed", [SetColor Foreground Vivid Green]) + failed = ("failed", [SetColor Foreground Vivid Red]) + passedProgs = show $ testStatusPass ts + failedProgs = show $ testStatusFail ts + totalProgs = show $ testStatusTotal ts + totalRuns = show $ testStatusRuns ts + passedRuns = show $ testStatusRunPass ts + failedRuns = show $ testStatusRunFail ts + remainProgs = show . length $ testStatusRemain ts + remainProgs' = remainProgs ++ "/" ++ totalProgs + remainRuns = show $ testStatusRunsRemain ts + remainRuns' = remainRuns ++ "/" ++ totalRuns + +tableLines :: Int +tableLines = 1 + (length . lines $ blankTable) + where blankTable = statusTable $ TestStatus [] [] 0 0 0 0 0 0 0 + +spaceTable :: IO () +spaceTable = putStr $ replicate tableLines '\n' + +reportTable :: TestStatus -> IO () +reportTable ts = do + moveCursorToTableTop + putStrLn $ statusTable ts + clearLine + putStrLn $ atMostChars 60 running + where running = "Now testing: " ++ + (unwords . reverse . map testCaseProgram . testStatusRun) ts + +moveCursorToTableTop :: IO () +moveCursorToTableTop = cursorUpLine tableLines + +atMostChars :: Int -> String -> String +atMostChars n s | length s > n = take (n-3) s ++ "..." + | otherwise = s + +reportText :: TestStatus -> IO () +reportText ts = + putStr $ "(" ++ show (testStatusFail ts) ++ " failed, " ++ + show (testStatusPass ts) ++ " passed, " ++ + show num_remain ++ " to go).\n" + where num_remain = length $ testStatusRemain ts + +runTests :: TestConfig -> [FilePath] -> IO () +runTests config paths = do + -- We force line buffering to ensure that we produce running output. + -- Otherwise, CI tools and the like may believe we are hung and kill + -- us. + hSetBuffering stdout LineBuffering + + let mode = configTestMode config + all_tests <- map (makeTestCase config mode) <$> testSpecsFromPaths paths + testmvar <- newEmptyMVar + reportmvar <- newEmptyMVar + concurrency <- getNumCapabilities + replicateM_ concurrency $ forkIO $ runTest testmvar reportmvar + + let (excluded, included) = partition (excludedTest config) all_tests + _ <- forkIO $ mapM_ (putMVar testmvar) included + isTTY <- (&& not (configLineOutput config)) <$> hIsTerminalDevice stdout + + let report = if isTTY then reportTable else reportText + clear = if isTTY then clearFromCursorToScreenEnd else putStr "\n" + + numTestCases tc = + case testAction $ testCaseTest tc of + CompileTimeFailure _ -> 1 + RunCases ios sts wts -> (length . concat) (iosTestRuns <$> ios) + + length sts + length wts + + getResults ts + | null (testStatusRemain ts) = report ts >> return ts + | otherwise = do + report ts + msg <- takeMVar reportmvar + case msg of + TestStarted test -> do + unless isTTY $ + putStr $ "Started testing " <> testCaseProgram test <> " " + getResults $ ts {testStatusRun = test : testStatusRun ts} + TestDone test res -> do + let ts' = ts { testStatusRemain = test `delete` testStatusRemain ts + , testStatusRun = test `delete` testStatusRun ts + , testStatusRunsRemain = testStatusRunsRemain ts + - numTestCases test + } + case res of + Success -> do + let ts'' = ts' { testStatusRunPass = + testStatusRunPass ts' + numTestCases test + } + unless isTTY $ + putStr $ "Finished testing " <> testCaseProgram test <> " " + getResults $ ts'' { testStatusPass = testStatusPass ts + 1} + Failure s -> do + when isTTY moveCursorToTableTop + clear + T.putStrLn $ (T.pack (inRed $ testCaseProgram test) <> ":\n") <> T.concat s + when isTTY spaceTable + getResults $ ts' { testStatusFail = testStatusFail ts' + 1 + , testStatusRunPass = testStatusRunPass ts' + + numTestCases test - length s + + , testStatusRunFail = testStatusRunFail ts' + + length s + } + + when isTTY spaceTable + + ts <- getResults TestStatus { testStatusRemain = included + , testStatusRun = [] + , testStatusTotal = length included + , testStatusFail = 0 + , testStatusPass = 0 + , testStatusRuns = sum $ map numTestCases included + , testStatusRunsRemain = sum $ map numTestCases included + , testStatusRunPass = 0 + , testStatusRunFail = 0 + } + + -- Removes "Now testing" output. + when isTTY $ cursorUpLine 1 >> clearLine + + let excluded_str = if null excluded + then "" + else " (" ++ show (length excluded) ++ " program(s) excluded).\n" + putStr excluded_str + exitWith $ case testStatusFail ts of 0 -> ExitSuccess + _ -> ExitFailure 1 + +inRed :: String -> String +inRed s = setSGRCode [SetColor Foreground Vivid Red] ++ s ++ setSGRCode [Reset] + +--- +--- Configuration and command line parsing +--- + +data TestConfig = TestConfig + { configTestMode :: TestMode + , configPrograms :: ProgConfig + , configExclude :: [T.Text] + , configLineOutput :: Bool + } + +defaultConfig :: TestConfig +defaultConfig = TestConfig { configTestMode = Everything + , configExclude = [ "disable" ] + , configPrograms = + ProgConfig + { configCompiler = "futhark-c" + , configInterpreter = "futharki" + , configTypeChecker = "futhark" + , configRunner = "" + , configExtraOptions = [] + , configExtraCompilerOptions = [] + } + , configLineOutput = False + } + +data ProgConfig = ProgConfig + { configCompiler :: FilePath + , configInterpreter :: FilePath + , configTypeChecker :: FilePath + , configRunner :: FilePath + , configExtraCompilerOptions :: [String] + , configExtraOptions :: [String] + -- ^ Extra options passed to the programs being run. + } + deriving (Show) + +changeProgConfig :: (ProgConfig -> ProgConfig) -> TestConfig -> TestConfig +changeProgConfig f config = config { configPrograms = f $ configPrograms config } + +setCompiler :: FilePath -> ProgConfig -> ProgConfig +setCompiler compiler config = + config { configCompiler = compiler } + +setInterpreter :: FilePath -> ProgConfig -> ProgConfig +setInterpreter interpreter config = + config { configInterpreter = interpreter } + +setTypeChecker :: FilePath -> ProgConfig -> ProgConfig +setTypeChecker typeChecker config = + config { configTypeChecker = typeChecker } + +setRunner :: FilePath -> ProgConfig -> ProgConfig +setRunner runner config = + config { configRunner = runner } + +addCompilerOption :: String -> ProgConfig -> ProgConfig +addCompilerOption option config = + config { configExtraCompilerOptions = configExtraCompilerOptions config ++ [option] } + +addOption :: String -> ProgConfig -> ProgConfig +addOption option config = + config { configExtraOptions = configExtraOptions config ++ [option] } + +data TestMode = TypeCheck + | Compile + | Compiled + | Interpreted + | Everything + deriving (Eq, Show) + +commandLineOptions :: [FunOptDescr TestConfig] +commandLineOptions = [ + Option "t" ["typecheck"] + (NoArg $ Right $ \config -> config { configTestMode = TypeCheck }) + "Only perform type-checking" + , Option "i" ["interpreted"] + (NoArg $ Right $ \config -> config { configTestMode = Interpreted }) + "Only interpret" + , Option "c" ["compiled"] + (NoArg $ Right $ \config -> config { configTestMode = Compiled }) + "Only run compiled code" + , Option "C" ["compile"] + (NoArg $ Right $ \config -> config { configTestMode = Compile }) + "Only compile, do not run." + , Option [] ["no-terminal", "notty"] + (NoArg $ Right $ \config -> config { configLineOutput = True }) + "Provide simpler line-based output." + , Option [] ["typechecker"] + (ReqArg (Right . changeProgConfig . setTypeChecker) "PROGRAM") + "What to run for type-checking (defaults to 'futhark')." + , Option [] ["compiler"] + (ReqArg (Right . changeProgConfig . setCompiler) "PROGRAM") + "What to run for code generation (defaults to 'futhark-c')." + , Option [] ["interpreter"] + (ReqArg (Right . changeProgConfig . setInterpreter) "PROGRAM") + "What to run for interpretation (defaults to 'futharki')." + , Option [] ["runner"] + (ReqArg (Right . changeProgConfig . setRunner) "PROGRAM") + "The program used to run the Futhark-generated programs (defaults to nothing)." + , Option [] ["exclude"] + (ReqArg (\tag -> + Right $ \config -> + config { configExclude = T.pack tag : configExclude config }) + "TAG") + "Exclude test programs that define this tag." + , Option "p" ["pass-option"] + (ReqArg (Right . changeProgConfig . addOption) "OPT") + "Pass this option to programs being run." + , Option [] ["pass-compiler-option"] + (ReqArg (Right . changeProgConfig . addCompilerOption) "OPT") + "Pass this option to the compiler (or typechecker if in -t mode)." + ] + +main :: IO () +main = mainWithOptions defaultConfig commandLineOptions "options... programs..." $ \progs config -> + Just $ runTests config progs diff --git a/src/futhark.hs b/src/futhark.hs new file mode 100644 --- /dev/null +++ b/src/futhark.hs @@ -0,0 +1,399 @@ +{-# LANGUAGE RankNTypes #-} +-- | Futhark Compiler Driver +module Main (main) where + +import Data.Maybe +import Control.Category (id) +import Control.Monad +import Control.Monad.State +import Data.Semigroup ((<>)) +import qualified Data.Text.IO as T +import System.IO +import System.Exit +import System.Console.GetOpt + +import Prelude hiding (id) + +import Futhark.Pass +import Futhark.Actions +import Futhark.Compiler +import Language.Futhark.Parser (parseFuthark) +import Futhark.Util.Options +import Futhark.Pipeline +import qualified Futhark.Representation.SOACS as SOACS +import Futhark.Representation.SOACS (SOACS) +import qualified Futhark.Representation.Kernels as Kernels +import Futhark.Representation.Kernels (Kernels) +import qualified Futhark.Representation.ExplicitMemory as ExplicitMemory +import Futhark.Representation.ExplicitMemory (ExplicitMemory) +import Futhark.Representation.AST (Prog, pretty) +import Futhark.TypeCheck (Checkable) +import qualified Futhark.Util.Pretty as PP + +import Futhark.Internalise.Defunctorise as Defunctorise +import Futhark.Internalise.Monomorphise as Monomorphise +import Futhark.Internalise.Defunctionalise as Defunctionalise +import Futhark.Optimise.InliningDeadFun +import Futhark.Optimise.CSE +import Futhark.Optimise.Fusion +import Futhark.Pass.FirstOrderTransform +import Futhark.Pass.Simplify +import Futhark.Optimise.InPlaceLowering +import Futhark.Optimise.DoubleBuffer +import Futhark.Optimise.TileLoops +import Futhark.Optimise.Unstream +import Futhark.Pass.KernelBabysitting +import Futhark.Pass.ExtractKernels +import Futhark.Pass.ExpandAllocations +import Futhark.Pass.ExplicitAllocations +import Futhark.Passes + +-- | What to do with the program after it has been read. +data FutharkPipeline = PrettyPrint + -- ^ Just print it. + | TypeCheck + -- ^ Run the type checker; print type errors. + | Pipeline [UntypedPass] + -- ^ Run this pipeline. + | Defunctorise + -- ^ Partially evaluate away the module language. + | Monomorphise + -- ^ Defunctorise and monomorphise. + | Defunctionalise + -- ^ Defunctorise, monomorphise, and defunctionalise. + +data Config = Config { futharkConfig :: FutharkConfig + , futharkPipeline :: FutharkPipeline + -- ^ Nothing is distinct from a empty pipeline - + -- it means we don't even run the internaliser. + , futharkAction :: UntypedAction + } + + +-- | Get a Futhark pipeline from the configuration - an empty one if +-- none exists. +getFutharkPipeline :: Config -> [UntypedPass] +getFutharkPipeline = toPipeline . futharkPipeline + where toPipeline (Pipeline p) = p + toPipeline _ = [] + +data UntypedPassState = SOACS (Prog SOACS.SOACS) + | Kernels (Prog Kernels.Kernels) + | ExplicitMemory (Prog ExplicitMemory.ExplicitMemory) + +getSOACSProg :: UntypedPassState -> Maybe (Prog SOACS.SOACS) +getSOACSProg (SOACS prog) = Just prog +getSOACSProg _ = Nothing + +class Representation s where + -- | A human-readable description of the representation expected or + -- contained, usable for error messages. + representation :: s -> String + +instance Representation UntypedPassState where + representation (SOACS _) = "SOACS" + representation (Kernels _) = "Kernels" + representation (ExplicitMemory _) = "ExplicitMemory" + +instance PP.Pretty UntypedPassState where + ppr (SOACS prog) = PP.ppr prog + ppr (Kernels prog) = PP.ppr prog + ppr (ExplicitMemory prog) = PP.ppr prog + +newtype UntypedPass = UntypedPass (UntypedPassState + -> PipelineConfig + -> FutharkM UntypedPassState) + +data UntypedAction = SOACSAction (Action SOACS) + | KernelsAction (Action Kernels) + | ExplicitMemoryAction (Action ExplicitMemory) + | PolyAction (Action SOACS) (Action Kernels) (Action ExplicitMemory) + +untypedActionName :: UntypedAction -> String +untypedActionName (SOACSAction a) = actionName a +untypedActionName (KernelsAction a) = actionName a +untypedActionName (ExplicitMemoryAction a) = actionName a +untypedActionName (PolyAction a _ _) = actionName a + +instance Representation UntypedAction where + representation (SOACSAction _) = "SOACS" + representation (KernelsAction _) = "Kernels" + representation (ExplicitMemoryAction _) = "ExplicitMemory" + representation PolyAction{} = "" + +newConfig :: Config +newConfig = Config newFutharkConfig (Pipeline []) $ PolyAction printAction printAction printAction + +changeFutharkConfig :: (FutharkConfig -> FutharkConfig) + -> Config -> Config +changeFutharkConfig f cfg = cfg { futharkConfig = f $ futharkConfig cfg } + +type FutharkOption = FunOptDescr Config + +passOption :: String -> UntypedPass -> String -> [String] -> FutharkOption +passOption desc pass short long = + Option short long + (NoArg $ Right $ \cfg -> + cfg { futharkPipeline = Pipeline $ getFutharkPipeline cfg ++ [pass] }) + desc + +explicitMemoryProg :: String -> UntypedPassState -> FutharkM (Prog ExplicitMemory.ExplicitMemory) +explicitMemoryProg _ (ExplicitMemory prog) = + return prog +explicitMemoryProg name rep = + externalErrorS $ "Pass " ++ name ++ + " expects ExplicitMemory representation, but got " ++ representation rep + +soacsProg :: String -> UntypedPassState -> FutharkM (Prog SOACS.SOACS) +soacsProg _ (SOACS prog) = + return prog +soacsProg name rep = + externalErrorS $ "Pass " ++ name ++ + " expects SOACS representation, but got " ++ representation rep + +kernelsProg :: String -> UntypedPassState -> FutharkM (Prog Kernels.Kernels) +kernelsProg _ (Kernels prog) = + return prog +kernelsProg name rep = + externalErrorS $ + "Pass " ++ name ++" expects Kernels representation, but got " ++ representation rep + +typedPassOption :: (Checkable fromlore, Checkable tolore) => + (String -> UntypedPassState -> FutharkM (Prog fromlore)) + -> (Prog tolore -> UntypedPassState) + -> Pass fromlore tolore + -> String + -> FutharkOption +typedPassOption getProg putProg pass short = + passOption (passDescription pass) (UntypedPass perform) short long + where perform s config = do + prog <- getProg (passName pass) s + putProg <$> runPasses (onePass pass) config prog + + long = [passLongOption pass] + +soacsPassOption :: Pass SOACS SOACS -> String -> FutharkOption +soacsPassOption = + typedPassOption soacsProg SOACS + +kernelsPassOption :: Pass Kernels Kernels -> String -> FutharkOption +kernelsPassOption = + typedPassOption kernelsProg Kernels + +explicitMemoryPassOption :: Pass ExplicitMemory ExplicitMemory -> String -> FutharkOption +explicitMemoryPassOption = + typedPassOption explicitMemoryProg ExplicitMemory + +simplifyOption :: String -> FutharkOption +simplifyOption short = + passOption (passDescription pass) (UntypedPass perform) short long + where perform (SOACS prog) config = + SOACS <$> runPasses (onePass simplifySOACS) config prog + perform (Kernels prog) config = + Kernels <$> runPasses (onePass simplifyKernels) config prog + perform (ExplicitMemory prog) config = + ExplicitMemory <$> runPasses (onePass simplifyExplicitMemory) config prog + + long = [passLongOption pass] + pass = simplifySOACS + +cseOption :: String -> FutharkOption +cseOption short = + passOption (passDescription pass) (UntypedPass perform) short long + where perform (SOACS prog) config = + SOACS <$> runPasses (onePass $ performCSE True) config prog + perform (Kernels prog) config = + Kernels <$> runPasses (onePass $ performCSE True) config prog + perform (ExplicitMemory prog) config = + ExplicitMemory <$> runPasses (onePass $ performCSE False) config prog + + long = [passLongOption pass] + pass = performCSE True :: Pass SOACS SOACS + +pipelineOption :: (UntypedPassState -> Maybe (Prog fromlore)) + -> String + -> (Prog tolore -> UntypedPassState) + -> String + -> Pipeline fromlore tolore + -> String + -> [String] + -> FutharkOption +pipelineOption getprog repdesc repf desc pipeline = + passOption desc $ UntypedPass pipelinePass + where pipelinePass rep config = + case getprog rep of + Just prog -> + repf <$> runPasses pipeline config prog + Nothing -> + externalErrorS $ "Expected " ++ repdesc ++ " representation, but got " ++ + representation rep + +soacsPipelineOption :: String -> Pipeline SOACS SOACS -> String -> [String] + -> FutharkOption +soacsPipelineOption = pipelineOption getSOACSProg "SOACS" SOACS + +kernelsPipelineOption :: String -> Pipeline SOACS Kernels -> String -> [String] + -> FutharkOption +kernelsPipelineOption = pipelineOption getSOACSProg "Kernels" Kernels + +explicitMemoryPipelineOption :: String -> Pipeline SOACS ExplicitMemory -> String -> [String] + -> FutharkOption +explicitMemoryPipelineOption = pipelineOption getSOACSProg "ExplicitMemory" ExplicitMemory + +commandLineOptions :: [FutharkOption] +commandLineOptions = + [ Option "v" ["verbose"] + (OptArg (Right . changeFutharkConfig . incVerbosity) "FILE") + "Print verbose output on standard error; wrong program to FILE." + , Option [] ["Werror"] + (NoArg $ Right $ changeFutharkConfig $ \opts -> opts { futharkWerror = True }) + "Treat warnings as errors." + + , Option "t" ["type-check"] + (NoArg $ Right $ \opts -> + opts { futharkPipeline = TypeCheck }) + "Type-check the program and print errors on standard error." + + , Option [] ["pretty-print"] + (NoArg $ Right $ \opts -> + opts { futharkPipeline = PrettyPrint }) + "Parse and pretty-print the AST of the given program." + + , Option [] ["compile-imperative"] + (NoArg $ Right $ \opts -> + opts { futharkAction = ExplicitMemoryAction impCodeGenAction }) + "Translate program into the imperative IL and write it on standard output." + , Option [] ["compile-imperative-kernels"] + (NoArg $ Right $ \opts -> + opts { futharkAction = ExplicitMemoryAction kernelImpCodeGenAction }) + "Translate program into the imperative IL with kernels and write it on standard output." + , Option [] ["range-analysis"] + (NoArg $ Right $ \opts -> opts { futharkAction = PolyAction rangeAction rangeAction rangeAction }) + "Print the program with range annotations added." + , Option "p" ["print"] + (NoArg $ Right $ \opts -> opts { futharkAction = PolyAction printAction printAction printAction }) + "Prettyprint the resulting internal representation on standard output (default action)." + , Option "m" ["metrics"] + (NoArg $ Right $ \opts -> opts { futharkAction = PolyAction metricsAction metricsAction metricsAction }) + "Print AST metrics of the resulting internal representation on standard output." + , Option [] ["defunctorise"] + (NoArg $ Right $ \opts -> opts { futharkPipeline = Defunctorise }) + "Partially evaluate all module constructs and print the residual program." + , Option [] ["monomorphise"] + (NoArg $ Right $ \opts -> opts { futharkPipeline = Monomorphise }) + "Monomorphise the program." + , Option [] ["defunctionalise"] + (NoArg $ Right $ \opts -> opts { futharkPipeline = Defunctionalise }) + "Defunctionalise the program." + , typedPassOption soacsProg Kernels firstOrderTransform "f" + , soacsPassOption fuseSOACs "o" + , soacsPassOption inlineAndRemoveDeadFunctions [] + , kernelsPassOption inPlaceLowering [] + , kernelsPassOption babysitKernels [] + , kernelsPassOption tileLoops [] + , kernelsPassOption unstream [] + , typedPassOption soacsProg Kernels extractKernels [] + + , typedPassOption kernelsProg ExplicitMemory explicitAllocations "a" + + , explicitMemoryPassOption doubleBuffer [] + , explicitMemoryPassOption expandAllocations [] + + , cseOption [] + , simplifyOption "e" + + , soacsPipelineOption "Run the default optimised pipeline" + standardPipeline "s" ["standard"] + , kernelsPipelineOption "Run the default optimised kernels pipeline" + kernelsPipeline [] ["kernels"] + , explicitMemoryPipelineOption "Run the full GPU compilation pipeline" + gpuPipeline [] ["gpu"] + , explicitMemoryPipelineOption "Run the sequential CPU compilation pipeline" + sequentialCpuPipeline [] ["cpu"] + ] + +incVerbosity :: Maybe FilePath -> FutharkConfig -> FutharkConfig +incVerbosity file cfg = + cfg { futharkVerbose = (v, file `mplus` snd (futharkVerbose cfg)) } + where v = case fst $ futharkVerbose cfg of + NotVerbose -> Verbose + Verbose -> VeryVerbose + VeryVerbose -> VeryVerbose + +-- | Entry point. Non-interactive, except when reading interpreter +-- input from standard input. +main :: IO () +main = mainWithOptions newConfig commandLineOptions "options... program" compile + where compile [file] config = + Just $ do + res <- runFutharkM (m file config) $ + fst $ futharkVerbose $ futharkConfig config + case res of + Left err -> do + dumpError (futharkConfig config) err + exitWith $ ExitFailure 2 + Right () -> return () + compile _ _ = + Nothing + m file config = + case futharkPipeline config of + TypeCheck -> do + -- No pipeline; just read the program and type check + (warnings, _, _) <- readProgram file + liftIO $ hPutStr stderr $ show warnings + PrettyPrint -> liftIO $ do + maybe_prog <- parseFuthark file <$> T.readFile file + case maybe_prog of + Left err -> fail $ show err + Right prog-> putStrLn $ pretty prog + Defunctorise -> do + (_, imports, src) <- readProgram file + liftIO $ mapM_ (putStrLn . pretty) $ + evalState (Defunctorise.transformProg imports) src + Monomorphise -> do + (_, imports, src) <- readProgram file + liftIO $ mapM_ (putStrLn . pretty) $ flip evalState src $ + Defunctorise.transformProg imports + >>= Monomorphise.transformProg + Defunctionalise -> do + (_, imports, src) <- readProgram file + liftIO $ mapM_ (putStrLn . pretty) $ flip evalState src $ + Defunctorise.transformProg imports + >>= Monomorphise.transformProg + >>= Defunctionalise.transformProg + Pipeline{} -> do + prog <- runPipelineOnProgram (futharkConfig config) id file + runPolyPasses config prog + +runPolyPasses :: Config -> SOACS.Prog -> FutharkM () +runPolyPasses config prog = do + prog' <- foldM (runPolyPass pipeline_config) (SOACS prog) (getFutharkPipeline config) + case (prog', futharkAction config) of + (SOACS soacs_prog, SOACSAction action) -> + actionProcedure action soacs_prog + (Kernels kernels_prog, KernelsAction action) -> + actionProcedure action kernels_prog + (ExplicitMemory mem_prog, ExplicitMemoryAction action) -> + actionProcedure action mem_prog + + (SOACS soacs_prog, PolyAction soacs_action _ _) -> + actionProcedure soacs_action soacs_prog + (Kernels kernels_prog, PolyAction _ kernels_action _) -> + actionProcedure kernels_action kernels_prog + (ExplicitMemory mem_prog, PolyAction _ _ mem_action) -> + actionProcedure mem_action mem_prog + + (_, action) -> + externalErrorS $ "Action " <> + untypedActionName action <> + " expects " ++ representation action ++ " representation, but got " ++ + representation prog' ++ "." + where pipeline_config = + PipelineConfig { pipelineVerbose = fst (futharkVerbose $ futharkConfig config) > NotVerbose + , pipelineValidate = True + } + +runPolyPass :: PipelineConfig + -> UntypedPassState -> UntypedPass -> FutharkM UntypedPassState +runPolyPass pipeline_config s (UntypedPass f) = + f s pipeline_config diff --git a/src/futharki.hs b/src/futharki.hs new file mode 100644 --- /dev/null +++ b/src/futharki.hs @@ -0,0 +1,458 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +module Main (main) where + +import Control.Monad.Free.Church +import Control.Exception +import Data.Array +import Data.Char +import Data.List +import Data.Loc +import Data.Maybe +import Data.Version +import qualified Data.Map as M +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.State +import Control.Monad.Except +import Data.Semigroup ((<>)) +import qualified Data.Text as T +import qualified Data.Text.IO as T +import NeatInterpolation (text) +import System.Directory +import System.FilePath +import System.Exit +import System.Console.GetOpt +import System.IO +import qualified System.Console.Haskeline as Haskeline + +import Language.Futhark +import Language.Futhark.Parser hiding (EOF) +import qualified Language.Futhark.TypeChecker as T +import qualified Language.Futhark.Semantic as T +import Futhark.MonadFreshNames +import Futhark.Version +import Futhark.Compiler +import Futhark.Pipeline +import Futhark.Util.Options +import Futhark.Util (toPOSIX, maybeHead) + +import qualified Language.Futhark.Interpreter as I + +banner :: String +banner = unlines [ + "|// |\\ | |\\ |\\ /", + "|/ | \\ |\\ |\\ |/ /", + "| | \\ |/ | |\\ \\", + "| | \\ | | | \\ \\" + ] + +main :: IO () +main = reportingIOErrors $ + mainWithOptions interpreterConfig options "options... program" run + where run [prog] config = Just $ interpret config prog + run [] _ = Just repl + run _ _ = Nothing + +data StopReason = EOF | Stop | Exit | Load FilePath + +repl :: IO () +repl = do + putStr banner + putStrLn $ "Version " ++ showVersion version ++ "." + putStrLn "Copyright (C) DIKU, University of Copenhagen, released under the ISC license." + putStrLn "" + putStrLn "Run :help for a list of commands." + putStrLn "" + + let toploop s = do + (stop, s') <- runStateT (runExceptT $ runFutharkiM $ forever readEvalPrint) s + case stop of + Left Stop -> finish s' + Left EOF -> finish s' + Left Exit -> finish s' + Left (Load file) -> do + liftIO $ T.putStrLn $ "Loading " <> T.pack file + maybe_new_state <- + liftIO $ newFutharkiState (futharkiCount s) $ Just file + case maybe_new_state of + Right new_state -> toploop new_state + Left err -> do liftIO $ putStrLn err + toploop s' + Right _ -> return () + + finish s = do + quit <- confirmQuit + if quit then return () else toploop s + + maybe_init_state <- liftIO $ newFutharkiState 0 Nothing + case maybe_init_state of + Left err -> error $ "Failed to initialise intepreter state: " ++ err + Right init_state -> Haskeline.runInputT Haskeline.defaultSettings $ toploop init_state + + putStrLn "Leaving futharki." + +confirmQuit :: Haskeline.InputT IO Bool +confirmQuit = do + c <- Haskeline.getInputChar "Quit futharki? (y/n) " + case c of + Nothing -> return True -- EOF + Just 'y' -> return True + Just 'n' -> return False + _ -> confirmQuit + +interpret :: InterpreterConfig -> FilePath -> IO () +interpret config fp = do + pr <- newFutharkiState 0 $ Just fp + env <- case pr of Left err -> do hPutStrLn stderr err + exitFailure + Right env -> return env + + let entry = interpreterEntryPoint config + (tenv, ienv) = futharkiEnv env + vr <- parseValues "stdin" <$> T.getContents + + inps <- + case vr of + Left err -> do + hPutStrLn stderr $ "Error when reading input: " ++ show err + exitFailure + Right vs + | Just vs' <- mapM convertValue vs -> + return vs' + | otherwise -> do + hPutStrLn stderr "Error when reading input: irregular array." + exitFailure + + (fname, ret) <- + case M.lookup (T.Term, entry) $ T.envNameMap tenv of + Just fname + | Just (T.BoundV _ t) <- M.lookup (qualLeaf fname) $ T.envVtable tenv -> + return (fname, toStructural $ snd $ unfoldFunType t) + _ -> do hPutStrLn stderr $ "Invalid entry point: " ++ pretty entry + exitFailure + + r <- runInterpreter' $ I.interpretFunction ienv (qualLeaf fname) inps + case r of + Left err -> do hPrint stderr err + exitFailure + Right res -> + case (I.fromTuple res, isTupleRecord ret) of + (Just vs, Just ts) -> zipWithM_ putValue vs ts + _ -> putValue res ret + +putValue :: I.Value -> TypeBase () () -> IO () +putValue v t + | I.isEmptyArray v = + putStrLn $ "empty(" ++ pretty (stripArray 1 t) ++ ")" + | otherwise = putStrLn $ pretty v + +convertValue :: Value -> Maybe I.Value +convertValue (PrimValue p) = Just $ I.ValuePrim p +convertValue (ArrayValue arr _) = I.mkArray =<< mapM convertValue (elems arr) + +newtype InterpreterConfig = InterpreterConfig { interpreterEntryPoint :: Name } + +interpreterConfig :: InterpreterConfig +interpreterConfig = InterpreterConfig defaultEntryPoint + +options :: [FunOptDescr InterpreterConfig] +options = [ Option "e" ["entry-point"] + (ReqArg (\entry -> Right $ \config -> + config { interpreterEntryPoint = nameFromString entry }) + "NAME") + "The entry point to execute." + ] + +data FutharkiState = + FutharkiState { futharkiImports :: Imports + , futharkiNameSource :: VNameSource + , futharkiCount :: Int + , futharkiEnv :: (T.Env, I.Ctx) + , futharkiBreaking :: Maybe Loc + -- ^ Are we currently stopped at a breakpoint? + , futharkiSkipBreaks :: [Loc] + -- ^ Skip breakpoints at these locations. + , futharkiLoaded :: Maybe FilePath + -- ^ The currently loaded file. + } + +newFutharkiState :: Int -> Maybe FilePath -> IO (Either String FutharkiState) +newFutharkiState count maybe_file = runExceptT $ do + (imports, src, tenv, ienv) <- case maybe_file of + + Nothing -> do + -- Load the builtins through the type checker. + (_, imports, src) <- badOnLeft =<< runExceptT (readLibrary []) + -- Then into the interpreter. + ienv <- foldM (\ctx -> badOnLeft <=< runInterpreter' . I.interpretImport ctx) + I.initialCtx $ map (fmap fileProg) imports + + -- Then make the prelude available in the type checker. + (tenv, d, src') <- badOnLeft $ T.checkDec imports src T.initialEnv + (T.mkInitialImport ".") $ mkOpen "/futlib/prelude" + -- Then in the interpreter. + ienv' <- badOnLeft =<< runInterpreter' (I.interpretDec ienv d) + return (imports, src', tenv, ienv') + + Just file -> do + (_, imports, src) <- + badOnLeft =<< liftIO (runExceptT (readProgram file) + `Haskeline.catch` \(err::IOException) -> + return (Left (ExternalError (T.pack $ show err)))) + let imp = T.mkInitialImport "." + ienv1 <- foldM (\ctx -> badOnLeft <=< runInterpreter' . I.interpretImport ctx) I.initialCtx $ + map (fmap fileProg) imports + (tenv1, d1, src') <- badOnLeft $ T.checkDec imports src T.initialEnv imp $ + mkOpen "/futlib/prelude" + (tenv2, d2, src'') <- badOnLeft $ T.checkDec imports src' tenv1 imp $ + mkOpen $ toPOSIX $ dropExtension file + ienv2 <- badOnLeft =<< runInterpreter' (I.interpretDec ienv1 d1) + ienv3 <- badOnLeft =<< runInterpreter' (I.interpretDec ienv2 d2) + return (imports, src'', tenv2, ienv3) + + return FutharkiState { futharkiImports = imports + , futharkiNameSource = src + , futharkiCount = count + , futharkiEnv = (tenv, ienv) + , futharkiBreaking = Nothing + , futharkiSkipBreaks = mempty + , futharkiLoaded = maybe_file + } + where badOnLeft :: Show err => Either err a -> ExceptT String IO a + badOnLeft (Right x) = return x + badOnLeft (Left err) = throwError $ show err + +getPrompt :: FutharkiM String +getPrompt = do + i <- gets futharkiCount + return $ "[" ++ show i ++ "]> " + +mkOpen :: FilePath -> UncheckedDec +mkOpen f = OpenDec (ModImport f NoInfo noLoc) NoInfo noLoc + +-- The ExceptT part is more of a continuation, really. +newtype FutharkiM a = + FutharkiM { runFutharkiM :: ExceptT StopReason (StateT FutharkiState (Haskeline.InputT IO)) a } + deriving (Functor, Applicative, Monad, + MonadState FutharkiState, MonadIO, MonadError StopReason) + +readEvalPrint :: FutharkiM () +readEvalPrint = do + prompt <- getPrompt + line <- inputLine prompt + breaking <- gets futharkiBreaking + case T.uncons line of + Nothing + | isJust breaking -> throwError Stop + | otherwise -> return () + + Just (':', command) -> do + let (cmdname, rest) = T.break isSpace command + arg = T.dropWhileEnd isSpace $ T.dropWhile isSpace rest + case filter ((cmdname `T.isPrefixOf`) . fst) commands of + [] -> liftIO $ T.putStrLn $ "Unknown command '" <> cmdname <> "'" + [(_, (cmdf, _))] -> cmdf arg + matches -> liftIO $ T.putStrLn $ "Ambiguous command; could be one of " <> + mconcat (intersperse ", " (map fst matches)) + + _ -> do + -- Read a declaration or expression. + maybe_dec_or_e <- parseDecOrExpIncrM (inputLine " ") prompt line + + case maybe_dec_or_e of + Left err -> liftIO $ print err + Right (Left d) -> onDec d + Right (Right e) -> onExp e + modify $ \s -> s { futharkiCount = futharkiCount s + 1 } + where inputLine prompt = do + inp <- FutharkiM $ lift $ lift $ Haskeline.getInputLine prompt + case inp of + Just s -> return $ T.pack s + Nothing -> throwError EOF + +getIt :: FutharkiM (Imports, VNameSource, T.Env, I.Ctx) +getIt = do + imports <- gets futharkiImports + src <- gets futharkiNameSource + (tenv, ienv) <- gets futharkiEnv + return (imports, src, tenv, ienv) + +onDec :: UncheckedDec -> FutharkiM () +onDec d = do + (imports, src, tenv, ienv) <- getIt + cur_import <- T.mkInitialImport . fromMaybe "." <$> gets futharkiLoaded + + -- Most of the complexity here concerns the dealing with the fact + -- that 'import "foo"' is a declaration. We have to involve a lot + -- of machinery to load this external code before executing the + -- declaration itself. + let basis = Basis imports src ["/futlib/prelude"] + mkImport = uncurry $ T.mkImportFrom cur_import + imp_r <- runExceptT $ readImports basis (map mkImport $ decImports d) + + case imp_r of + Left e -> liftIO $ print e + Right (_, imports', src') -> + case T.checkDec imports' src' tenv cur_import d of + Left e -> liftIO $ print e + Right (tenv', d', src'') -> do + let new_imports = filter ((`notElem` map fst imports) . fst) imports' + int_r <- runInterpreter $ do + let onImport ienv' (s, imp) = + I.interpretImport ienv' (s, T.fileProg imp) + ienv' <- foldM onImport ienv new_imports + I.interpretDec ienv' d' + case int_r of + Left err -> liftIO $ print err + Right ienv' -> modify $ \s -> s { futharkiEnv = (tenv', ienv') + , futharkiImports = imports' + , futharkiNameSource = src'' + } + +onExp :: UncheckedExp -> FutharkiM () +onExp e = do + (imports, src, tenv, ienv) <- getIt + case showErr (T.checkExp imports src tenv e) of + Left err -> liftIO $ putStrLn err + Right e' -> do + r <- runInterpreter $ I.interpretExp ienv e' + case r of + Left err -> liftIO $ print err + Right v -> liftIO $ putStrLn $ pretty v + where showErr :: Show a => Either a b -> Either String b + showErr = either (Left . show) Right + +runInterpreter :: F I.ExtOp a -> FutharkiM (Either I.InterpreterError a) +runInterpreter m = runF m (return . Right) intOp + where + intOp (I.ExtOpError err) = + return $ Left err + intOp (I.ExtOpTrace w v c) = do + liftIO $ putStrLn $ "Trace at " ++ locStr w ++ ": " ++ v + c + intOp (I.ExtOpBreak w ctx tenv c) = do + s <- get + + -- Are we supposed to skip this breakpoint? + let loc = maybe noLoc locOf $ maybeHead w + + -- We do not want recursive breakpoints. It could work fine + -- technically, but is probably too confusing to be useful. + unless (isJust (futharkiBreaking s) || loc `elem` futharkiSkipBreaks s) $ do + liftIO $ putStrLn $ "Breaking at " ++ intercalate " -> " (map locStr w) ++ "." + liftIO $ putStrLn " to continue." + + -- Note the cleverness to preserve the Haskeline session (for + -- line history and such). + (stop, s') <- + FutharkiM $ lift $ lift $ + runStateT (runExceptT $ runFutharkiM $ forever readEvalPrint) + s { futharkiEnv = (tenv, ctx) + , futharkiCount = futharkiCount s + 1 + , futharkiBreaking = Just loc } + + case stop of + Left (Load file) -> throwError $ Load file + _ -> do liftIO $ putStrLn "Continuing..." + put s { futharkiCount = futharkiCount s' + , futharkiSkipBreaks = futharkiSkipBreaks s' <> futharkiSkipBreaks s } + + c + +runInterpreter' :: MonadIO m => F I.ExtOp a -> m (Either I.InterpreterError a) +runInterpreter' m = runF m (return . Right) intOp + where intOp (I.ExtOpError err) = return $ Left err + intOp (I.ExtOpTrace w v c) = do + liftIO $ putStrLn $ "Trace at " ++ locStr w ++ ": " ++ v + c + intOp (I.ExtOpBreak _ _ _ c) = c + +type Command = T.Text -> FutharkiM () + +loadCommand :: Command +loadCommand file = do + loaded <- gets futharkiLoaded + case (T.null file, loaded) of + (True, Just loaded') -> throwError $ Load loaded' + (True, Nothing) -> liftIO $ T.putStrLn "No file specified and no file previously loaded." + (False, _) -> throwError $ Load $ T.unpack file + +typeCommand :: Command +typeCommand e = do + prompt <- getPrompt + case parseExp prompt e of + Left err -> liftIO $ print err + Right e' -> do + imports <- gets futharkiImports + src <- gets futharkiNameSource + (tenv, _) <- gets futharkiEnv + case T.checkExp imports src tenv e' of + Left err -> liftIO $ print err + Right e'' -> liftIO $ putStrLn $ pretty e' <> " : " <> pretty (typeOf e'') + +unbreakCommand :: Command +unbreakCommand _ = do + breaking <- gets futharkiBreaking + case breaking of + Nothing -> liftIO $ putStrLn "Not currently stopped at a breakpoint." + Just loc -> do modify $ \s -> s { futharkiSkipBreaks = loc : futharkiSkipBreaks s } + throwError Stop + +pwdCommand :: Command +pwdCommand _ = liftIO $ putStrLn =<< getCurrentDirectory + +cdCommand :: Command +cdCommand dir + | T.null dir = liftIO $ putStrLn "Usage: ':cd '." + | otherwise = + liftIO $ setCurrentDirectory (T.unpack dir) + `Haskeline.catch` \(err::IOException) -> print err + +helpCommand :: Command +helpCommand _ = liftIO $ forM_ commands $ \(cmd, (_, desc)) -> do + T.putStrLn $ ":" <> cmd + T.putStrLn $ T.replicate (1+T.length cmd) "-" + T.putStr desc + T.putStrLn "" + T.putStrLn "" + +quitCommand :: Command +quitCommand _ = throwError Exit + +commands :: [(T.Text, (Command, T.Text))] +commands = [("load", (loadCommand, [text| +Load a Futhark source file. Usage: + + > :load foo.fut + +If the loading succeeds, any subsequentialy entered expressions entered +subsequently will have access to the definition (such as function definitions) +in the source file. + +Only one source file can be loaded at a time. Using the :load command a +second time will replace the previously loaded file. It will also replace +any declarations entered at the REPL. + +|])), + ("type", (typeCommand, [text| +Show the type of an expression. +|])), + ("unbreak", (unbreakCommand, [text| +Skip all future occurences of the current breakpoint. +|])), + ("pwd", (pwdCommand, [text| +Print the current working directory. +|])), + ("cd", (cdCommand, [text| +Change the current working directory. +|])), + ("help", (helpCommand, [text| +Print a list of commands and a description of their behaviour. +|])), + ("quit", (quitCommand, [text| +Quit futharki. +|]))] diff --git a/unittests/Futhark/Analysis/ScalExpTests.hs b/unittests/Futhark/Analysis/ScalExpTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Analysis/ScalExpTests.hs @@ -0,0 +1,108 @@ +{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, FlexibleContexts #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Analysis.ScalExpTests + ( tests + , parseScalExp + , parseScalExp' + ) +where + +import Test.Tasty + +import Control.Applicative +import Control.Monad.State +import qualified Data.Map as M +import Data.Void +import Text.Megaparsec hiding (token, (<|>), many, State) +import Control.Monad.Combinators.Expr +import Text.Megaparsec.Char +import qualified Text.Megaparsec.Char.Lexer as L + +import Futhark.Analysis.ScalExp +import Futhark.Representation.AST hiding (constant, SDiv) + +tests :: TestTree +tests = testGroup "ScalExpTests" [] + +parseScalExp :: String -> ScalExp +parseScalExp = parseScalExp' M.empty + +parseScalExp' :: M.Map String (Int, Type) -> String -> ScalExp +parseScalExp' m s = case evalState (runParserT expr ("string: " ++ s) s) (0, m) of + Left err -> error $ show err + Right e -> e + +type ParserState = (Int, M.Map String (Int, Type)) +type Parser = ParsecT Void String (State ParserState) + +newVar :: String -> Type -> Parser Ident +newVar s t = do (x, m) <- lift get + case M.lookup s m of + Just _ -> fail $ "Variable " ++ s ++ " double-declared." + Nothing -> do lift $ put (x+1, M.insert s (x,t) m) + return $ Ident (VName (nameFromString s) x) t + +knownVar :: String -> Parser Ident +knownVar s = do (_, m) <- lift get + case M.lookup s m of + Just (y,t) -> return $ Ident (VName (nameFromString s) y) t + Nothing -> fail $ "Undeclared variable " ++ s + +token :: String -> Parser () +token = L.lexeme space . void . string + +parens :: Parser a -> Parser a +parens = between (token "(") (token ")") + +identifier :: Parser Ident +identifier = do s <- (:) <$> letterChar <*> many alphaNumChar + varDecl s <|> knownVar s + where varDecl s = do + t <- parens $ + (token "int" >> pure (Prim $ IntType Int32)) <|> + (token "float32" >> pure (Prim $ FloatType Float32)) <|> + (token "float64" >> pure (Prim $ FloatType Float64)) <|> + (token "bool" >> pure (Prim Bool)) + newVar s t + +constant :: Parser ScalExp +constant = (token "True" >> pure (Val $ BoolValue True)) <|> + (token "False" >> pure (Val $ BoolValue True)) <|> + (Val . IntValue . Int32Value <$> integer) + where integer = L.lexeme space (L.signed space L.decimal) + +expr :: Parser ScalExp +expr = makeExprParser prim operators + +prim :: Parser ScalExp +prim = parens expr <|> + constant <|> + maxapp <|> + minapp <|> + (scalExpId <$> identifier) + where maxapp = token "max" >> MaxMin False <$> parens (expr `sepBy` comma) + minapp = token "min" >> MaxMin True <$> parens (expr `sepBy` comma) + comma = token "," + scalExpId (Ident name (Prim t)) = Id name t + scalExpId (Ident name t) = error $ + pretty name ++ " is of type " ++ pretty t ++ + " but supposed to be a ScalExp." + +operators :: [[Operator Parser ScalExp]] +operators = [ [Prefix (token "-" >> return SNeg)] + , [InfixL (token "*" >> return STimes)] + , [InfixL (token "pow" >> return SPow)] + , [InfixL (token "/" >> return SDiv)] + , [InfixL (token "+" >> return SPlus)] + , [InfixL (token "-" >> return SMinus)] + , [InfixL (token "<=" >> return leq)] + , [InfixL (token "<" >> return lth)] + , [InfixL (token ">=" >> return (flip leq))] + , [InfixL (token ">" >> return (flip lth))] + , [InfixL (token "&&" >> return SLogAnd)] + , [InfixL (token "||" >> return SLogOr)] + ] + where leq x y = + RelExp LEQ0 $ x `SMinus` y + lth x y = + RelExp LTH0 $ x `SMinus` y diff --git a/unittests/Futhark/Optimise/AlgSimplifyTests.hs b/unittests/Futhark/Optimise/AlgSimplifyTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Optimise/AlgSimplifyTests.hs @@ -0,0 +1,101 @@ +module Futhark.Optimise.AlgSimplifyTests ( tests ) +where + +import Test.Tasty +import Test.Tasty.HUnit + +import Data.List +import qualified Data.Map.Strict as M + +import Futhark.Representation.AST +import Futhark.Analysis.ScalExp +import Futhark.Analysis.ScalExpTests (parseScalExp') +import Futhark.Analysis.AlgSimplify + +tests :: TestTree +tests = testGroup "AlgSimplifyTests" $ constantFoldTests ++ suffCondTests + +constantFoldTests :: [TestTree] +constantFoldTests = + [ cfoldTest "2+2" "4" + , cfoldTest "2-2" "0" + , cfoldTest "2*3" "6" + , cfoldTest "6/3" "2" + + -- Simple cases over; let's try some variables. + , cfoldTest "0+x" "x" + , cfoldTest "x+x" "2*x" -- Sensitive to operand order + , cfoldTest "x-0" "x" + , cfoldTest "x-x" "0" + , cfoldTest "x/x" "1" + , cfoldTest "x/1" "x" + , cfoldTest "x/x" "1" + ] + where vars = declareVars [("x", int32)] + simplify'' e = simplify' vars e [] + scalExp = parseScalExp' vars + + cfoldTest input expected = + testCase ("constant-fold " ++ input) $ + simplify'' input @?= scalExp expected + +suffCondTests :: [TestTree] +suffCondTests = + [ + suffCondTest "5 VarInfo -> VName +lookupVarName s varinfo = case M.lookup s varinfo of + Nothing -> error $ "Unknown variable " ++ s + Just (x,_) -> VName (nameFromString s) x + +declareVars :: VarDecls -> VarInfo +declareVars = M.fromList . snd . mapAccumL declare 0 + where declare i (name, t) = (i+1, (name, (i, Prim t))) + +instantiateRanges :: VarInfo -> RangesRep' -> RangesRep +instantiateRanges varinfo r = + M.fromList $ snd $ mapAccumL fix 0 r + where fix i (name, lower,upper) = + (i+1, + (lookupVarName name varinfo, + (i, fixBound lower, fixBound upper))) + fixBound "" = Nothing + fixBound s = Just $ parseScalExp' varinfo s + +simplify' :: VarInfo -> String -> RangesRep' -> ScalExp +simplify' varinfo s r = simplify e r' + where e = parseScalExp' varinfo s + r' = instantiateRanges varinfo r + +mkSuffConds' :: VarInfo -> String -> RangesRep' -> [[ScalExp]] +mkSuffConds' varinfo s r = + case mkSuffConds e r' of + Left _ -> [[e]] + Right sc -> sc + where e = simplify (parseScalExp' varinfo s) r' + r' = instantiateRanges varinfo r diff --git a/unittests/Futhark/Pkg/SolveTests.hs b/unittests/Futhark/Pkg/SolveTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Pkg/SolveTests.hs @@ -0,0 +1,109 @@ +{-# LANGUAGE OverloadedStrings #-} +module Futhark.Pkg.SolveTests (tests) where + +import qualified Data.Map as M +import qualified Data.Text as T +import Data.Monoid + +import Test.Tasty +import Test.Tasty.HUnit + +import Futhark.Pkg.Types +import Futhark.Pkg.Solve + +import Prelude + +semverE :: T.Text -> SemVer +semverE s = case parseVersion s of + Left err -> error $ T.unpack s <> + " is not a valid version number: " <> + errorBundlePretty err + Right x -> x + +-- | A world of packages and interdependencies for testing the solver +-- without touching the outside world. +testEnv :: PkgRevDepInfo +testEnv = M.fromList $ concatMap frob + [ ("athas", [ ("foo", [ ("0.1.0", []) + , ("0.2.0", [("athas/bar", "1.0.0")]) + , ("0.3.0", [])]) + , ("foo@v2", [ ("2.0.0", [("athas/quux", "0.1.0")])]) + , ("bar", [ ("1.0.0", [])]) + , ("baz", [ ("0.1.0", [("athas/foo", "0.3.0")])]) + , ("quux", [ ("0.1.0", [ ("athas/foo", "0.2.0") + , ("athas/baz", "0.1.0") ])]) + , ("quux_perm", [ ("0.1.0", [ ("athas/baz", "0.1.0") + , ("athas/foo", "0.2.0")])]) + , ("x_bar", [ ("1.0.0", [("athas/bar", "1.0.0")])]) + , ("x_foo", [ ("1.0.0", [("athas/foo", "0.3.0")])]) + , ("tricky", [ ("1.0.0", [ ("athas/foo", "0.2.0") + , ("athas/x_foo", "1.0.0")])]) + ]) + + -- Some mutually recursive packages. + , ("nasty", [ ("foo", [ ("1.0.0", [("nasty/bar", "1.0.0")])]) + , ("bar", [ ("1.0.0", [("nasty/foo", "1.0.0")])])]) + ] + where frob (user, repos) = do + (repo, repo_revs) <- repos + (rev, deps) <- repo_revs + let rev' = semverE rev + onDep (dp, dv) = (dp, (semverE dv, Nothing)) + deps' = PkgRevDeps $ M.fromList $ map onDep deps + return ((user <> "/" <> repo, rev'), deps') + +newtype SolverRes = SolverRes BuildList + deriving (Eq) + +instance Show SolverRes where + show (SolverRes bl) = T.unpack $ prettyBuildList bl + +solverTest :: PkgPath -> T.Text -> Either T.Text [(PkgPath, T.Text)] -> TestTree +solverTest p v expected = + testCase (T.unpack $ p <> "-" <> prettySemVer v') $ + fmap SolverRes (solveDepsPure testEnv target) + @?= expected' + where target = PkgRevDeps $ M.singleton p (v', Nothing) + v' = semverE v + expected' = SolverRes . BuildList . M.fromList . map onRes <$> expected + onRes (dp, dv) = (dp, semverE dv) + +tests :: TestTree +tests = testGroup "SolveTests" + [ + solverTest "athas/foo" "0.1.0" $ + Right [ ("athas/foo", "0.1.0")] + + , solverTest "athas/foo" "0.2.0" $ + Right [ ("athas/foo", "0.2.0") + , ("athas/bar", "1.0.0")] + + , solverTest "athas/quux" "0.1.0" $ + Right [ ("athas/quux", "0.1.0") + , ("athas/foo", "0.3.0") + , ("athas/baz", "0.1.0")] + + , solverTest "athas/quux_perm" "0.1.0" $ + Right [ ("athas/quux_perm", "0.1.0") + , ("athas/foo", "0.3.0") + , ("athas/baz", "0.1.0")] + + , solverTest "athas/foo@v2" "2.0.0" $ + Right [ ("athas/foo@v2", "2.0.0") + , ("athas/quux", "0.1.0") + , ("athas/foo", "0.3.0") + , ("athas/baz", "0.1.0") + ] + + , solverTest "athas/foo@v3" "3.0.0" $ + Left "Unknown package/version: athas/foo@v3-3.0.0" + + , solverTest "nasty/foo" "1.0.0" $ + Right [ ("nasty/foo", "1.0.0") + , ("nasty/bar", "1.0.0")] + + , solverTest "athas/tricky" "1.0.0" $ + Right [ ("athas/tricky", "1.0.0") + , ("athas/foo", "0.3.0") + , ("athas/x_foo", "1.0.0")] + ] diff --git a/unittests/Futhark/Representation/AST/Attributes/RearrangeTests.hs b/unittests/Futhark/Representation/AST/Attributes/RearrangeTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/AST/Attributes/RearrangeTests.hs @@ -0,0 +1,55 @@ +module Futhark.Representation.AST.Attributes.RearrangeTests + ( tests ) + where + +import Control.Applicative + +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck + +import Prelude + +import Futhark.Representation.AST.Attributes.Rearrange + +tests :: TestTree +tests = testGroup "RearrangeTests" $ + isMapTransposeTests ++ + [isMapTransposeProp] + +isMapTransposeTests :: [TestTree] +isMapTransposeTests = + [ testCase (unwords ["isMapTranspose", show perm, "==", show dres]) $ + isMapTranspose perm @?= dres + | (perm, dres) <- [ ([0,1,4,5,2,3], Just (2,2,2)) + , ([1,0,4,5,2,3], Nothing) + , ([1,0], Just (0, 1, 1)) + , ([0,2,1], Just (1, 1, 1)) + , ([0,1,2], Nothing) + , ([1,0,2], Nothing) + ] + ] + +newtype Permutation = Permutation [Int] + deriving (Eq, Ord, Show) + +instance Arbitrary Permutation where + arbitrary = do + Positive n <- arbitrary + Permutation <$> shuffle [0..n-1] + +isMapTransposeProp :: TestTree +isMapTransposeProp = testProperty "isMapTranspose corresponds to a map of transpose" prop + where prop :: Permutation -> Bool + prop (Permutation perm) = + case isMapTranspose perm of + Nothing -> True + Just (r1, r2, r3) -> + and [r1 >= 0, + r2 > 0, + r3 > 0, + r1 + r2 + r3 == length perm, + let (mapped, notmapped) =splitAt r1 perm + (pretrans, posttrans) = splitAt r2 notmapped + in mapped ++ posttrans ++ pretrans == [0..length perm-1] + ] diff --git a/unittests/Futhark/Representation/AST/Attributes/ReshapeTests.hs b/unittests/Futhark/Representation/AST/Attributes/ReshapeTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/AST/Attributes/ReshapeTests.hs @@ -0,0 +1,93 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Representation.AST.Attributes.ReshapeTests + ( tests + ) + where + +import Control.Applicative + +import Test.Tasty +import Test.Tasty.HUnit +import Test.Tasty.QuickCheck + +import Prelude + +import Futhark.Representation.AST.Attributes.Reshape +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Attributes.Constants + +tests :: TestTree +tests = testGroup "ReshapeTests" $ + fuseReshapeTests ++ + informReshapeTests ++ + reshapeOuterTests ++ + reshapeInnerTests ++ + [ fuseReshapeProp + , informReshapeProp + ] + +fuseReshapeTests :: [TestTree] +fuseReshapeTests = + [ testCase (unwords ["fuseReshape ", show d1, show d2]) $ + fuseReshape (d1 :: ShapeChange Int) d2 @?= dres -- type signature to avoid warning + | (d1, d2, dres) <- [ ([DimCoercion 1], [DimNew 1], [DimCoercion 1]) + , ([DimNew 1], [DimCoercion 1], [DimNew 1]) + , ([DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2]) + , ([DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2]) + ] + ] + +informReshapeTests :: [TestTree] +informReshapeTests = + [ testCase (unwords ["informReshape ", show shape, show sc, show sc_res]) $ + informReshape (shape :: [Int]) sc @?= sc_res -- type signature to avoid warning + | (shape, sc, sc_res) <- + [ ([1, 2], [DimNew 1, DimNew 3], [DimCoercion 1, DimNew 3]) + , ([2, 2], [DimNew 1, DimNew 3], [DimNew 1, DimNew 3]) + ] + ] + +reshapeOuterTests :: [TestTree] +reshapeOuterTests = + [ testCase (unwords ["reshapeOuter", show sc, show n, show shape, "==", show sc_res]) $ + reshapeOuter (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res + | (sc, n, shape, sc_res) <- + [ ([DimNew 1], 1, [4, 3], [DimNew 1, DimCoercion 3]) + , ([DimNew 1], 2, [4, 3], [DimNew 1]) + , ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 2, DimNew 2, DimNew 3]) + , ([DimNew 2, DimNew 2], 2, [4, 3], [DimNew 2, DimNew 2]) + ] + ] + +reshapeInnerTests :: [TestTree] +reshapeInnerTests = + [ testCase (unwords ["reshapeInner", show sc, show n, show shape, "==", show sc_res]) $ + reshapeInner (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res + | (sc, n, shape, sc_res) <- + [ ([DimNew 1], 1, [4, 3], [DimCoercion 4, DimNew 1]) + , ([DimNew 1], 0, [4, 3], [DimNew 1]) + , ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 4, DimNew 2, DimNew 2]) + , ([DimNew 2, DimNew 2], 0, [4, 3], [DimNew 2, DimNew 2]) + ] + ] + +intShape :: [Int] -> Shape +intShape = Shape . map (intConst Int32 . toInteger) + +intShapeChange :: ShapeChange Int -> ShapeChange SubExp +intShapeChange = map (fmap $ intConst Int32 . toInteger) + +fuseReshapeProp :: TestTree +fuseReshapeProp = testProperty "fuseReshape result matches second argument" prop + where prop :: ShapeChange Int -> ShapeChange Int -> Bool + prop sc1 sc2 = map newDim (fuseReshape sc1 sc2) == map newDim sc2 + +informReshapeProp :: TestTree +informReshapeProp = testProperty "informReshape result matches second argument" prop + where prop :: [Int] -> ShapeChange Int -> Bool + prop sc1 sc2 = map newDim (informReshape sc1 sc2) == map newDim sc2 + +instance Arbitrary d => Arbitrary (DimChange d) where + arbitrary = oneof [ DimNew <$> arbitrary + , DimCoercion <$> arbitrary + ] diff --git a/unittests/Futhark/Representation/AST/AttributesTests.hs b/unittests/Futhark/Representation/AST/AttributesTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/AST/AttributesTests.hs @@ -0,0 +1,17 @@ +{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Representation.AST.AttributesTests + ( tests + ) +where + +import Test.Tasty + +import Futhark.Representation.AST.SyntaxTests () +import qualified Futhark.Representation.AST.Attributes.ReshapeTests +import qualified Futhark.Representation.AST.Attributes.RearrangeTests + +tests :: TestTree +tests = testGroup "AttributesTests" + [Futhark.Representation.AST.Attributes.ReshapeTests.tests, + Futhark.Representation.AST.Attributes.RearrangeTests.tests] diff --git a/unittests/Futhark/Representation/AST/Syntax/CoreTests.hs b/unittests/Futhark/Representation/AST/Syntax/CoreTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/AST/Syntax/CoreTests.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Representation.AST.Syntax.CoreTests + ( tests ) + where + +import Control.Applicative + +import Test.Tasty +import Test.Tasty.HUnit +import Test.QuickCheck + +import Prelude + +import Language.Futhark.CoreTests () +import Futhark.Representation.PrimitiveTests() +import Futhark.Representation.AST.Syntax.Core +import Futhark.Representation.AST.Pretty () + +tests :: TestTree +tests = testGroup "Internal CoreTests" subShapeTests + +subShapeTests :: [TestTree] +subShapeTests = + [ shape [free 1, free 2] `isSubShapeOf` shape [free 1, free 2] + , shape [free 1, free 3] `isNotSubShapeOf` shape [free 1, free 2] + , shape [free 1] `isNotSubShapeOf` shape [free 1, free 2] + , shape [free 1, free 2] `isSubShapeOf` shape [free 1, Ext 3] + , shape [Ext 1, Ext 2] `isNotSubShapeOf` shape [Ext 1, Ext 1] + , shape [Ext 1, Ext 1] `isSubShapeOf` shape [Ext 1, Ext 2] + ] + where shape :: [ExtSize] -> ExtShape + shape = Shape + + free :: Int -> ExtSize + free = Free . Constant . IntValue . Int32Value . fromIntegral + + isSubShapeOf shape1 shape2 = + subShapeTest shape1 shape2 True + isNotSubShapeOf shape1 shape2 = + subShapeTest shape1 shape2 False + + subShapeTest :: ExtShape -> ExtShape -> Bool -> TestTree + subShapeTest shape1 shape2 expected = + testCase ("subshapeOf " ++ pretty shape1 ++ " " ++ + pretty shape2 ++ " == " ++ + show expected) $ + shape1 `subShapeOf` shape2 @?= expected + +instance Arbitrary NoUniqueness where + arbitrary = pure NoUniqueness + +instance (Arbitrary shape, Arbitrary u) => Arbitrary (TypeBase shape u) where + arbitrary = + oneof [ Prim <$> arbitrary + , Array <$> arbitrary <*> arbitrary <*> arbitrary + ] + +instance Arbitrary Ident where + arbitrary = Ident <$> arbitrary <*> arbitrary + +instance Arbitrary Rank where + arbitrary = Rank <$> elements [1..9] + +instance Arbitrary Shape where + arbitrary = Shape . map intconst <$> listOf1 (elements [1..9]) + where intconst = Constant . IntValue . Int32Value diff --git a/unittests/Futhark/Representation/AST/SyntaxTests.hs b/unittests/Futhark/Representation/AST/SyntaxTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/AST/SyntaxTests.hs @@ -0,0 +1,7 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Representation.AST.SyntaxTests + () +where + +-- There isn't anything to test in this module. At some point, maybe +-- we can put some Arbitrary instances here. diff --git a/unittests/Futhark/Representation/PrimitiveTests.hs b/unittests/Futhark/Representation/PrimitiveTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/PrimitiveTests.hs @@ -0,0 +1,61 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Representation.PrimitiveTests + ( tests + , arbitraryPrimValOfType + ) + where + +import Control.Applicative + +import Test.QuickCheck +import Test.Tasty +import Test.Tasty.HUnit + +import Prelude + +import Futhark.Representation.Primitive + +tests :: TestTree +tests = testGroup "PrimitiveTests" propPrimValuesHaveRightType + +propPrimValuesHaveRightType :: [TestTree] +propPrimValuesHaveRightType = [ testCase (show t ++ " has blank of right type") $ + primValueType (blankPrimValue t) @?= t + | t <- [minBound..maxBound] + ] + +instance Arbitrary IntType where + arbitrary = elements [minBound..maxBound] + +instance Arbitrary FloatType where + arbitrary = elements [minBound..maxBound] + +instance Arbitrary PrimType where + arbitrary = elements [minBound..maxBound] + +instance Arbitrary IntValue where + arbitrary = oneof [ Int8Value <$> arbitrary + , Int16Value <$> arbitrary + , Int32Value <$> arbitrary + , Int64Value <$> arbitrary ] + +instance Arbitrary FloatValue where + arbitrary = oneof [ Float32Value <$> arbitrary + , Float64Value <$> arbitrary ] + +instance Arbitrary PrimValue where + arbitrary = oneof [ IntValue <$> arbitrary + , FloatValue <$> arbitrary + , BoolValue <$> arbitrary + , pure Checked + ] + +arbitraryPrimValOfType :: PrimType -> Gen PrimValue +arbitraryPrimValOfType (IntType Int8) = IntValue . Int8Value <$> arbitrary +arbitraryPrimValOfType (IntType Int16) = IntValue . Int16Value <$> arbitrary +arbitraryPrimValOfType (IntType Int32) = IntValue . Int32Value <$> arbitrary +arbitraryPrimValOfType (IntType Int64) = IntValue . Int64Value <$> arbitrary +arbitraryPrimValOfType (FloatType Float32) = FloatValue . Float32Value <$> arbitrary +arbitraryPrimValOfType (FloatType Float64) = FloatValue . Float32Value <$> arbitrary +arbitraryPrimValOfType Bool = BoolValue <$> arbitrary +arbitraryPrimValOfType Cert = return Checked diff --git a/unittests/Language/Futhark/CoreTests.hs b/unittests/Language/Futhark/CoreTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Language/Futhark/CoreTests.hs @@ -0,0 +1,15 @@ +{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Language.Futhark.CoreTests () +where + +import Test.QuickCheck + +import Language.Futhark.Core +import Futhark.Representation.PrimitiveTests() + +instance Arbitrary Name where + arbitrary = nameFromString <$> listOf1 (elements ['a'..'z']) + +instance Arbitrary VName where + arbitrary = VName <$> arbitrary <*> arbitrary diff --git a/unittests/Language/Futhark/SyntaxTests.hs b/unittests/Language/Futhark/SyntaxTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Language/Futhark/SyntaxTests.hs @@ -0,0 +1,39 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# LANGUAGE FlexibleInstances #-} +module Language.Futhark.SyntaxTests (tests) +where + +import Control.Applicative + +import Prelude + +import Test.QuickCheck +import Test.Tasty + +import Language.Futhark.Syntax + +import Futhark.Representation.PrimitiveTests() + +tests :: TestTree +tests = testGroup "Source SyntaxTests" [] + +instance Arbitrary BinOp where + arbitrary = elements [minBound..maxBound] + +instance Arbitrary Uniqueness where + arbitrary = elements [Unique, Nonunique] + +instance Arbitrary PrimType where + arbitrary = oneof [ Signed <$> arbitrary + , Unsigned <$> arbitrary + , FloatType <$> arbitrary + , pure Bool + ] + +instance Arbitrary PrimValue where + arbitrary = oneof [ SignedValue <$> arbitrary + , UnsignedValue <$> arbitrary + , FloatValue <$> arbitrary + , BoolValue <$> arbitrary + ] diff --git a/unittests/futhark_tests.hs b/unittests/futhark_tests.hs new file mode 100644 --- /dev/null +++ b/unittests/futhark_tests.hs @@ -0,0 +1,22 @@ +module Main (main) where + +import qualified Language.Futhark.SyntaxTests +import qualified Futhark.Representation.AST.Syntax.CoreTests +import qualified Futhark.Representation.AST.AttributesTests +import qualified Futhark.Optimise.AlgSimplifyTests +import qualified Futhark.Pkg.SolveTests + +import Test.Tasty + +allTests :: TestTree +allTests = + testGroup "" + [ Language.Futhark.SyntaxTests.tests + , Futhark.Representation.AST.AttributesTests.tests + , Futhark.Optimise.AlgSimplifyTests.tests + , Futhark.Representation.AST.Syntax.CoreTests.tests + , Futhark.Pkg.SolveTests.tests + ] + +main :: IO () +main = defaultMain allTests