diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -2,12 +2,19 @@ -- -- see: https://github.com/sol/hpack -- --- hash: d926568e0952a1cecfa1efa7cf687ddc7cd27d1c642a929d85ca5b0e8465b408 +-- hash: ca599dc882c25837255ce48c8e6ab8c5a9f1ff05234e05f6d73705f5cf9fe1c6 name: futhark -version: 0.9.1 +version: 0.10.1 synopsis: An optimising compiler for a functional, array-oriented language. -description: See the website at https://futhark-lang.org +description: Futhark is a small programming language designed to be compiled to + efficient parallel code. It is a statically typed, data-parallel, + and purely functional array language in the ML family, and comes + with a heavily optimising ahead-of-time compiler that presently + generates GPU code via CUDA and OpenCL, although the language itself + is hardware-agnostic. + . + For more information, see the website at https://futhark-lang.org category: Language homepage: https://futhark-lang.org bug-reports: https://github.com/diku-dk/futhark/issues @@ -29,6 +36,7 @@ rts/c/opencl.h rts/c/panic.h rts/c/timing.h + rts/c/tuning.h rts/c/values.h rts/csharp/exceptions.cs rts/csharp/functions.cs @@ -43,6 +51,7 @@ rts/python/opencl.py rts/python/panic.py rts/python/scalar.py + rts/python/tuning.py rts/python/values.py source-repository head @@ -58,7 +67,6 @@ , 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 @@ -67,7 +75,6 @@ , 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 @@ -87,15 +94,14 @@ , 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 + , utf8-string >=1 , vector >=0.12 , vector-binary-instances >=0.2.2.0 , versions >=3.3.1 @@ -104,9 +110,6 @@ build-tools: alex , happy - if !impl(ghc >= 8.0) - build-depends: - semigroups ==0.18.* exposed-modules: Futhark.Actions Futhark.Analysis.AlgSimplify @@ -174,6 +177,7 @@ Futhark.CodeGen.ImpGen.CUDA Futhark.CodeGen.ImpGen.Kernels Futhark.CodeGen.ImpGen.Kernels.Base + Futhark.CodeGen.ImpGen.Kernels.SegGenRed Futhark.CodeGen.ImpGen.Kernels.SegRed Futhark.CodeGen.ImpGen.Kernels.ToOpenCL Futhark.CodeGen.ImpGen.Kernels.Transpose @@ -209,34 +213,6 @@ 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 @@ -257,7 +233,6 @@ Futhark.Pass.ExtractKernels.ISRWIM Futhark.Pass.ExtractKernels.Kernelise Futhark.Pass.ExtractKernels.Segmented - Futhark.Pass.ExtractKernels.Split Futhark.Pass.FirstOrderTransform Futhark.Pass.KernelBabysitting Futhark.Pass.ResolveAssertions @@ -288,7 +263,6 @@ 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 @@ -341,757 +315,13 @@ main-is: src/futhark.hs ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -threaded -rtsopts "-with-rtsopts=-N -qg" build-depends: - aeson - , 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 + base , 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 - , 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.* + , text other-modules: Paths_futhark default-language: Haskell2010 -executable futhark-bench - main-is: src/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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/wrapper.hs - other-modules: - Paths_futhark - ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists - build-depends: - aeson - , 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 - , 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 - , 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 - build-depends: - aeson - , 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 - , 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 - , 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 - test-suite unit type: exitcode-stdio-1.0 main-is: futhark_tests.hs @@ -1099,62 +329,17 @@ unittests ghc-options: -Wall -Wcompat -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists build-depends: - HUnit - , QuickCheck >=2.8 - , aeson - , 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 + QuickCheck >=2.8 + , base + , containers , 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 - , random - , raw-strings-qq >=1.1 - , regex-tdfa >=1.2 - , srcloc >=0.4 + , megaparsec + , mtl + , parser-combinators , tasty , tasty-hunit , tasty-quickcheck - , 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.* + , text other-modules: Futhark.Analysis.ScalExpTests Futhark.Optimise.AlgSimplifyTests @@ -1164,6 +349,9 @@ Futhark.Representation.AST.AttributesTests Futhark.Representation.AST.Syntax.CoreTests Futhark.Representation.AST.SyntaxTests + Futhark.Representation.ExplicitMemory.IndexFunction.Alg + Futhark.Representation.ExplicitMemory.IndexFunctionTests + Futhark.Representation.ExplicitMemory.IndexFunctionWrapper Futhark.Representation.PrimitiveTests Language.Futhark.CoreTests Language.Futhark.SyntaxTests diff --git a/futlib/prelude.fut b/futlib/prelude.fut --- a/futlib/prelude.fut +++ b/futlib/prelude.fut @@ -16,6 +16,13 @@ -- | Create integer from double-precision float. let t64 (x: f64): i32 = i32.f64 x +-- | Semantically just identity, but serves as an optimisation +-- inhibitor. The compiler will treat this function as a black box. +-- You can use this to work around optimisation deficiencies (or +-- bugs), although it should hopefully rarely be necessary. +let opaque 't (x: t): t = + intrinsics.opaque x + -- | Semantically just identity, but in `futharki` the argument value -- will be printed. let trace 't (x: t): t = diff --git a/futlib/soacs.fut b/futlib/soacs.fut --- a/futlib/soacs.fut +++ b/futlib/soacs.fut @@ -87,10 +87,10 @@ 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. +-- | 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 +-- commutative, so plain `reduce`@term will work just as well. -- -- **Work:** *O(n)* -- diff --git a/rts/c/cuda.h b/rts/c/cuda.h --- a/rts/c/cuda.h +++ b/rts/c/cuda.h @@ -242,15 +242,31 @@ return x[chosen].arch_str; } -static char *cuda_nvrtc_build(struct cuda_context *ctx, const char *src) +static char *cuda_nvrtc_build(struct cuda_context *ctx, const char *src, + const char *extra_opts[]) { nvrtcProgram prog; NVRTC_SUCCEED(nvrtcCreateProgram(&prog, src, "futhark-cuda", 0, NULL, NULL)); + int arch_set = 0, num_extra_opts; - size_t n_opts, i = 0, i_dyn, n_opts_alloc = 20 + ctx->cfg.num_sizes; + // nvrtc cannot handle multiple -arch options. Hence, if one of the + // extra_opts is -arch, we have to be careful not to do our usual + // automatic generation. + for (num_extra_opts = 0; extra_opts[num_extra_opts] != NULL; num_extra_opts++) { + if (strstr(extra_opts[num_extra_opts], "-arch") + == extra_opts[num_extra_opts] || + strstr(extra_opts[num_extra_opts], "--gpu-architecture") + == extra_opts[num_extra_opts]) { + arch_set = 1; + } + } + + size_t n_opts, i = 0, i_dyn, n_opts_alloc = 20 + num_extra_opts + ctx->cfg.num_sizes; const char **opts = malloc(n_opts_alloc * sizeof(const char *)); - opts[i++] = "-arch"; - opts[i++] = cuda_nvrtc_get_arch(ctx->dev); + if (!arch_set) { + opts[i++] = "-arch"; + opts[i++] = cuda_nvrtc_get_arch(ctx->dev); + } opts[i++] = "-default-device"; if (ctx->cfg.debugging) { opts[i++] = "-G"; @@ -265,6 +281,13 @@ } opts[i++] = msgprintf("-DLOCKSTEP_WIDTH=%zu", ctx->lockstep_width); opts[i++] = msgprintf("-DMAX_THREADS_PER_BLOCK=%zu", ctx->max_block_size); + + // It is crucial that the extra_opts are last, so that the free() + // logic below does not cause problems. + for (int j = 0; extra_opts[j] != NULL; j++) { + opts[i++] = extra_opts[j]; + } + n_opts = i; if (ctx->cfg.debugging) { @@ -288,7 +311,7 @@ NVRTC_SUCCEED(res); } - for (i = i_dyn; i < n_opts; i++) { free((char *)opts[i]); } + for (i = i_dyn; i < n_opts-num_extra_opts; i++) { free((char *)opts[i]); } free(opts); char *ptx; @@ -356,7 +379,7 @@ if (*size_value == 0) { *size_value = default_value; } else if (max_value > 0 && *size_value > max_value) { - fprintf(stderr, "Note: Device limits %zu to %zu (down from %zu)\n", + fprintf(stderr, "Note: Device limits %s to %zu (down from %zu)\n", size_name, max_value, *size_value); *size_value = max_value; } @@ -394,16 +417,17 @@ } static void cuda_module_setup(struct cuda_context *ctx, - const char *src_fragments[]) + const char *src_fragments[], + const char *extra_opts[]) { char *ptx = NULL, *src = NULL; if (ctx->cfg.load_ptx_from == NULL && ctx->cfg.load_program_from == NULL) { src = concat_fragments(src_fragments); - ptx = cuda_nvrtc_build(ctx, src); + ptx = cuda_nvrtc_build(ctx, src, extra_opts); } else if (ctx->cfg.load_ptx_from == NULL) { load_string_from_file(ctx->cfg.load_program_from, &src, NULL); - ptx = cuda_nvrtc_build(ctx, src); + ptx = cuda_nvrtc_build(ctx, src, extra_opts); } else { if (ctx->cfg.load_program_from != NULL) { fprintf(stderr, @@ -432,7 +456,7 @@ } } -void cuda_setup(struct cuda_context *ctx, const char *src_fragments[]) +void cuda_setup(struct cuda_context *ctx, const char *src_fragments[], const char *extra_opts[]) { CUDA_SUCCEED(cuInit(0)); @@ -450,7 +474,7 @@ ctx->lockstep_width = device_query(ctx->dev, WARP_SIZE); cuda_size_setup(ctx); - cuda_module_setup(ctx, src_fragments); + cuda_module_setup(ctx, src_fragments, extra_opts); } CUresult cuda_free_all(struct cuda_context *ctx); diff --git a/rts/c/opencl.h b/rts/c/opencl.h --- a/rts/c/opencl.h +++ b/rts/c/opencl.h @@ -70,8 +70,7 @@ cfg->default_group_size = 0; cfg->default_num_groups = 0; cfg->default_tile_size = 0; - - cfg->default_threshold = 32*1024; + cfg->default_threshold = 0; cfg->default_group_size_changed = 0; cfg->default_tile_size_changed = 0; @@ -397,31 +396,29 @@ } 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); + cl_int clBuildProgram_error = 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); + if (clBuildProgram_error != CL_SUCCESS && + clBuildProgram_error != CL_BUILD_PROGRAM_FAILURE) { + OPENCL_SUCCEED_FATAL(clBuildProgram_error); } 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); + OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, + device, + CL_PROGRAM_BUILD_STATUS, + sizeof(cl_build_status), + &build_status, + NULL)); 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); + OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, NULL, &ret_val_size)); build_log = malloc(ret_val_size+1); - clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, ret_val_size, build_log, NULL); - assert(ret_val == 0); + OPENCL_SUCCEED_FATAL(clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, ret_val_size, build_log, NULL)); // The spec technically does not say whether the build log is zero-terminated, so let's be careful. build_log[ret_val_size] = '\0'; @@ -445,7 +442,8 @@ static cl_program setup_opencl_with_command_queue(struct opencl_context *ctx, cl_command_queue queue, const char *srcs[], - int required_types) { + int required_types, + const char *extra_build_opts[]) { int error; ctx->queue = queue; @@ -478,7 +476,7 @@ 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", + panic(1, "Program uses double-precision floats, but this is not supported on the chosen device: %s\n", device_option.device_name); } } @@ -578,7 +576,7 @@ } cl_program prog; - error = 0; + error = CL_SUCCESS; const char* src_ptr[] = {fut_opencl_src}; if (ctx->cfg.dump_program_to != NULL) { @@ -590,12 +588,18 @@ if (ctx->cfg.load_binary_from == NULL) { prog = clCreateProgramWithSource(ctx->ctx, 1, src_ptr, &src_size, &error); - assert(error == 0); + OPENCL_SUCCEED_FATAL(error); 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; } + + for (int i = 0; extra_build_opts[i] != NULL; i++) { + compile_opts_size += strlen(extra_build_opts[i] + 1); + } + char *compile_opts = malloc(compile_opts_size); int w = snprintf(compile_opts, compile_opts_size, @@ -609,6 +613,11 @@ (int)ctx->cfg.size_values[i]); } + for (int i = 0; extra_build_opts[i] != NULL; i++) { + w += snprintf(compile_opts+w, compile_opts_size-w, + "%s ", extra_build_opts[i]); + } + OPENCL_SUCCEED_FATAL(build_opencl_program(prog, device_option.device, compile_opts)); free(compile_opts); @@ -650,7 +659,8 @@ static cl_program setup_opencl(struct opencl_context *ctx, const char *srcs[], - int required_types) { + int required_types, + const char *extra_build_opts[]) { ctx->lockstep_width = 0; // Real value set later. @@ -669,15 +679,15 @@ 0 }; - cl_int error; - - ctx->ctx = clCreateContext(properties, 1, &device_option.device, NULL, NULL, &error); - assert(error == 0); + cl_int clCreateContext_error; + ctx->ctx = clCreateContext(properties, 1, &device_option.device, NULL, NULL, &clCreateContext_error); + OPENCL_SUCCEED_FATAL(clCreateContext_error); - cl_command_queue queue = clCreateCommandQueue(ctx->ctx, device_option.device, 0, &error); - assert(error == 0); + cl_int clCreateCommandQueue_error; + cl_command_queue queue = clCreateCommandQueue(ctx->ctx, device_option.device, 0, &clCreateCommandQueue_error); + OPENCL_SUCCEED_FATAL(clCreateCommandQueue_error); - return setup_opencl_with_command_queue(ctx, queue, srcs, required_types); + return setup_opencl_with_command_queue(ctx, queue, srcs, required_types, extra_build_opts); } // Allocate memory from driver. The problem is that OpenCL may perform diff --git a/rts/c/tuning.h b/rts/c/tuning.h new file mode 100644 --- /dev/null +++ b/rts/c/tuning.h @@ -0,0 +1,40 @@ +// Start of tuning.h + +static char* load_tuning_file(const char *fname, + void *cfg, + int (*set_size)(void*, const char*, size_t)) { + const int max_line_len = 1024; + char* line = malloc(max_line_len); + + FILE *f = fopen(fname, "r"); + + if (f == NULL) { + snprintf(line, max_line_len, "Cannot open file: %s", strerror(errno)); + return line; + } + + int lineno = 0; + while (fgets(line, max_line_len, f) != NULL) { + lineno++; + char *eql = strstr(line, "="); + if (eql) { + *eql = 0; + int value = atoi(eql+1); + if (set_size(cfg, line, value) != 0) { + strncpy(eql+1, line, max_line_len-strlen(line)-1); + snprintf(line, max_line_len, "Unknown name '%s' on line %d.", eql+1, lineno); + return line; + } + } else { + snprintf(line, max_line_len, "Invalid line %d (must be of form 'name=int').", + lineno); + return line; + } + } + + free(line); + + return NULL; +} + +// End of tuning.h diff --git a/rts/csharp/opencl.cs b/rts/csharp/opencl.cs --- a/rts/csharp/opencl.cs +++ b/rts/csharp/opencl.cs @@ -374,10 +374,13 @@ cfg.DumpProgramTo = null; cfg.LoadProgramFrom = null; - cfg.DefaultGroupSize = 256; - cfg.DefaultNumGroups = 128; - cfg.DefaultTileSize = 32; - cfg.DefaultThreshold = 32*1024; + // The following are dummy sizes that mean the concrete defaults + // will be set during initialisation via hardware-inspection-based + // heuristics. + cfg.DefaultGroupSize = 0; + cfg.DefaultNumGroups = 0; + cfg.DefaultTileSize = 0; + cfg.DefaultThreshold = 0; cfg.NumSizes = num_sizes; cfg.SizeNames = size_names; @@ -774,6 +777,9 @@ int MaxTileSize = (int) Math.Sqrt(MaxGroupSize); + // Make sure this function is defined. + PostOpenCLSetup(ref ctx, ref device_option); + 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); @@ -798,16 +804,16 @@ string size_name = ctx.OpenCL.Cfg.SizeNames[i]; int max_value, default_value; max_value = default_value = 0; - if (size_class == "group_size") { + if (size_class.StartsWith("group_size")) { max_value = MaxGroupSize; default_value = ctx.OpenCL.Cfg.DefaultGroupSize; - } else if (size_class == "num_groups") { + } else if (size_class.StartsWith("num_groups")) { max_value = MaxGroupSize; // Futhark assumes this constraint. default_value = ctx.OpenCL.Cfg.DefaultNumGroups; - } else if (size_class == "tile_size"){ + } else if (size_class.StartsWith("tile_size")){ max_value = (int) Math.Sqrt(MaxGroupSize); default_value = ctx.OpenCL.Cfg.DefaultTileSize; - } else if (size_class == "threshold") { + } else if (size_class.StartsWith("threshold")) { max_value = 0; // No limit. default_value = ctx.OpenCL.Cfg.DefaultThreshold; } else { @@ -835,9 +841,6 @@ 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); @@ -916,8 +919,31 @@ var name = name_and_value[0]; var value = Convert.ToInt32(name_and_value[1]); - if (!FutharkContextConfigSetSize(ref config, name, value)) + if (name == "default_num_groups") { + config.OpenCL.DefaultNumGroups = value; + } + else if (name == "default_group_size") { + config.OpenCL.DefaultGroupSize = value; + } + else if (name == "default_tile_size") { + config.OpenCL.DefaultTileSize = value; + } + else if (name == "default_threshold") { + config.OpenCL.DefaultThreshold = value; + } + else if (!FutharkContextConfigSetSize(ref config, name, value)) { panic(1, "Unknown size: {0}", name); } +} + +private void FutharkConfigLoadTuning(ref FutharkContextConfig config, string fname) +{ + StreamReader file = new StreamReader(fname); + String line; + while((line = file.ReadLine()) != null) + { + FutharkConfigSetSize(ref config, line); + } + file.Close(); } diff --git a/rts/python/opencl.py b/rts/python/opencl.py --- a/rts/python/opencl.py +++ b/rts/python/opencl.py @@ -55,6 +55,10 @@ device_matches += 1 raise Exception('No OpenCL platform and device matching constraints found.') +def size_assignment(s): + name, value = s.split('=') + return (name, int(value)) + def check_types(self, required_types): if 'f64' in required_types: if self.device.get_info(cl.device_info.PREFERRED_VECTOR_WIDTH_DOUBLE) == 0: @@ -107,6 +111,22 @@ self.max_num_groups = 0 self.free_list = {} + if 'default_group_size' in sizes: + default_group_size = sizes['default_group_size'] + del sizes['default_group_size'] + + if 'default_num_groups' in sizes: + default_num_groups = sizes['default_num_groups'] + del sizes['default_num_groups'] + + if 'default_tile_size' in sizes: + default_tile_size = sizes['default_tile_size'] + del sizes['default_tile_size'] + + if 'default_threshold' in sizes: + default_threshold = sizes['default_threshold'] + del sizes['default_threshold'] + default_group_size_set = default_group_size != None default_tile_size_set = default_tile_size != None default_sizes = apply_size_heuristics(self, size_heuristics, @@ -137,7 +157,7 @@ if k in all_sizes: all_sizes[k]['value'] = v else: - raise Exception('Unknown size: {}'.format(k)) + raise Exception('Unknown size: {}\nKnown sizes: {}'.format(k, ' '.join(all_sizes.keys()))) self.sizes = {} for (k,v) in all_sizes.items(): diff --git a/rts/python/tuning.py b/rts/python/tuning.py new file mode 100644 --- /dev/null +++ b/rts/python/tuning.py @@ -0,0 +1,11 @@ +### start of tuning.py +### +### Reading the .tuning file. + +def read_tuning_file(kvs, f): + for line in f.read().splitlines(): + size, value = line.split('=') + kvs[size] = int(value) + return kvs + +### end of tuning.py diff --git a/src/Futhark/CLI/Bench.hs b/src/Futhark/CLI/Bench.hs --- a/src/Futhark/CLI/Bench.hs +++ b/src/Futhark/CLI/Bench.hs @@ -45,11 +45,12 @@ , optExcludeCase :: [String] , optIgnoreFiles :: [Regex] , optEntryPoint :: Maybe String + , optTuning :: Maybe String } initialBenchOptions :: BenchOptions initialBenchOptions = BenchOptions "c" "futhark" "" 10 [] Nothing (-1) False - ["nobench", "disable"] [] Nothing + ["nobench", "disable"] [] Nothing (Just "tuning") -- | The name we use for compiled programs. binaryName :: FilePath -> FilePath @@ -158,9 +159,13 @@ runBenchmark :: BenchOptions -> (FilePath, [InputOutputs]) -> IO [BenchResult] runBenchmark opts (program, cases) = mapM forInputOutputs $ filter relevant cases where forInputOutputs (InputOutputs entry_name runs) = do - putStr $ "Results for " ++ program' ++ ":\n" + (tuning_opts, tuning_desc) <- determineTuning (optTuning opts) program + + putStr $ "Results for " ++ program' ++ tuning_desc ++ ":\n" + let opts' = opts { optExtraOptions = + optExtraOptions opts ++ tuning_opts } BenchResult program' . catMaybes <$> - mapM (runBenchmarkCase opts program entry_name pad_to) runs + mapM (runBenchmarkCase opts' program entry_name pad_to) runs where program' = if entry_name == "main" then program else program ++ ":" ++ T.unpack entry_name @@ -368,6 +373,13 @@ config { optEntryPoint = Just s }) "NAME") "Only run this entry point." + , Option [] ["tuning"] + (ReqArg (\s -> Right $ \config -> config { optTuning = Just s }) + "EXTENSION") + "Look for tuning files with this extension (defaults: .tuning)." + , Option [] ["no-tuning"] + (NoArg $ Right $ \config -> config { optTuning = Nothing }) + "Do not load tuning files." ] where max_timeout :: Int max_timeout = maxBound `div` 1000000 diff --git a/src/Futhark/CLI/Datacmp.hs b/src/Futhark/CLI/Datacmp.hs --- a/src/Futhark/CLI/Datacmp.hs +++ b/src/Futhark/CLI/Datacmp.hs @@ -3,7 +3,6 @@ import qualified Data.ByteString.Lazy.Char8 as BS import System.Exit -import System.IO import Futhark.Test.Values import Futhark.Util.Options @@ -21,9 +20,8 @@ (Just vs_a, Just vs_b) -> case compareValues vs_a vs_b of [] -> return () - es -> do - mapM_ (hPrint stderr) es - exitWith $ ExitFailure 2 + es -> do mapM_ print es + exitWith $ ExitFailure 2 f _ _ = Nothing diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs --- a/src/Futhark/CLI/Dev.hs +++ b/src/Futhark/CLI/Dev.hs @@ -66,6 +66,9 @@ -- ^ Nothing is distinct from a empty pipeline - -- it means we don't even run the internaliser. , futharkAction :: UntypedAction + , futharkPrintAST :: Bool + -- ^ If true, prints programs as raw ASTs instead + -- of their prettyprinted form. } @@ -121,7 +124,8 @@ representation PolyAction{} = "" newConfig :: Config -newConfig = Config newFutharkConfig (Pipeline []) $ PolyAction printAction printAction printAction +newConfig = Config newFutharkConfig (Pipeline []) + (PolyAction printAction printAction printAction) False changeFutharkConfig :: (FutharkConfig -> FutharkConfig) -> Config -> Config @@ -251,7 +255,7 @@ , Option "t" ["type-check"] (NoArg $ Right $ \opts -> opts { futharkPipeline = TypeCheck }) - "Type-check the program and print errors on standard error." + "Print on standard output the type-checked program." , Option [] ["pretty-print"] (NoArg $ Right $ \opts -> @@ -284,6 +288,9 @@ , Option [] ["defunctionalise"] (NoArg $ Right $ \opts -> opts { futharkPipeline = Defunctionalise }) "Defunctionalise the program." + , Option [] ["ast"] + (NoArg $ Right $ \opts -> opts { futharkPrintAST = True }) + "Output ASTs instead of prettyprinted programs." , typedPassOption soacsProg Kernels firstOrderTransform "f" , soacsPassOption fuseSOACs "o" , soacsPassOption inlineAndRemoveDeadFunctions [] @@ -336,27 +343,32 @@ 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 + Right prog | futharkPrintAST config -> print prog + | otherwise -> putStrLn $ pretty prog + TypeCheck -> do + (_, imports, _) <- readProgram file + liftIO $ forM_ (map snd imports) $ \fm -> + putStrLn $ if futharkPrintAST config + then show $ fileProg fm + else pretty $ fileProg fm Defunctorise -> do (_, imports, src) <- readProgram file - liftIO $ mapM_ (putStrLn . pretty) $ + liftIO $ mapM_ (putStrLn . if futharkPrintAST config then show else pretty) $ evalState (Defunctorise.transformProg imports) src Monomorphise -> do (_, imports, src) <- readProgram file - liftIO $ mapM_ (putStrLn . pretty) $ flip evalState src $ + liftIO $ mapM_ (putStrLn . if futharkPrintAST config then show else pretty) $ + flip evalState src $ Defunctorise.transformProg imports >>= Monomorphise.transformProg Defunctionalise -> do (_, imports, src) <- readProgram file - liftIO $ mapM_ (putStrLn . pretty) $ flip evalState src $ + liftIO $ mapM_ (putStrLn . if futharkPrintAST config then show else pretty) $ + flip evalState src $ Defunctorise.transformProg imports >>= Monomorphise.transformProg >>= Defunctionalise.transformProg diff --git a/src/Futhark/CLI/Misc.hs b/src/Futhark/CLI/Misc.hs --- a/src/Futhark/CLI/Misc.hs +++ b/src/Futhark/CLI/Misc.hs @@ -2,9 +2,11 @@ -- Various small subcommands that are too simple to deserve their own file. module Futhark.CLI.Misc ( mainCheck + , mainImports ) where +import Data.List (isPrefixOf) import Control.Monad.State import System.IO import System.Exit @@ -29,3 +31,14 @@ _ -> Nothing where check file = do (warnings, _, _) <- readProgram file liftIO $ hPutStr stderr $ show warnings + +mainImports :: String -> [String] -> IO () +mainImports = mainWithOptions () [] "program" $ \args () -> + case args of + [file] -> Just $ runFutharkM' $ findImports file + _ -> Nothing + where findImports file = do + (_, prog_imports, _) <- readProgram file + liftIO $ putStr $ unlines $ map (++ ".fut") + $ filter (\f -> not ("futlib/" `isPrefixOf` f)) + $ map fst prog_imports diff --git a/src/Futhark/CLI/Pkg.hs b/src/Futhark/CLI/Pkg.hs --- a/src/Futhark/CLI/Pkg.hs +++ b/src/Futhark/CLI/Pkg.hs @@ -283,7 +283,7 @@ Nothing -> liftIO $ T.putStrLn $ "Added new required package " <> p <> " " <> prettySemVer v <> "." putPkgManifest m' - liftIO $ T.putStrLn "Remember to run 'futhark-pkg sync'." + liftIO $ T.putStrLn "Remember to run 'futhark pkg sync'." doRemove :: String -> [String] -> IO () doRemove = cmdMain "PKGPATH" $ \args cfg -> @@ -326,6 +326,9 @@ m <- getPkgManifest rs <- traverse (mapM (traverse upgrade)) $ manifestRequire m putPkgManifest m { manifestRequire = rs } + if rs == manifestRequire m + then liftIO $ T.putStrLn "Nothing to upgrade." + else liftIO $ T.putStrLn "Remember to run 'futhark pkg sync'." _ -> Nothing where upgrade req = do v <- lookupNewestRev $ requiredPkg req diff --git a/src/Futhark/CLI/REPL.hs b/src/Futhark/CLI/REPL.hs --- a/src/Futhark/CLI/REPL.hs +++ b/src/Futhark/CLI/REPL.hs @@ -91,7 +91,7 @@ confirmQuit :: Haskeline.InputT IO Bool confirmQuit = do - c <- Haskeline.getInputChar "Quit futharki? (y/n) " + c <- Haskeline.getInputChar "Quit REPL? (y/n) " case c of Nothing -> return True -- EOF Just 'y' -> return True diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs --- a/src/Futhark/CLI/Test.hs +++ b/src/Futhark/CLI/Test.hs @@ -169,9 +169,13 @@ context ("Compiling with --backend=" <> T.pack backend) $ do compileTestProgram extra_options futhark backend program warnings mapM_ (testMetrics progs program) structures - unless (mode == Compile) $ + unless (mode == Compile) $ do + (tuning_opts, _) <- + liftIO $ determineTuning (configTuning progs) program + let progs' = progs { configExtraOptions = + tuning_opts ++ configExtraOptions progs } context "Running compiled program" $ - accErrors_ $ map (runCompiledEntry program progs) ios + accErrors_ $ map (runCompiledEntry program progs') ios unless (mode == Compile || mode == Compiled) $ context "Interpreting" $ accErrors_ $ map (runInterpretedEntry futhark program) ios @@ -380,8 +384,10 @@ _ <- 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" + let report | isTTY = reportTable + | otherwise = reportText + clear | isTTY = clearFromCursorToScreenEnd + |otherwise = putStr "\n" numTestCases tc = case testAction $ testCaseTest tc of @@ -442,9 +448,8 @@ -- Removes "Now testing" output. when isTTY $ cursorUpLine 1 >> clearLine - let excluded_str = if null excluded - then "" - else " (" ++ show (length excluded) ++ " program(s) excluded).\n" + let excluded_str | null excluded = "" + | otherwise = " (" ++ show (length excluded) ++ " program(s) excluded).\n" putStr excluded_str exitWith $ case testStatusFail ts of 0 -> ExitSuccess _ -> ExitFailure 1 @@ -473,6 +478,7 @@ , configRunner = "" , configExtraOptions = [] , configExtraCompilerOptions = [] + , configTuning = Just "tuning" } , configLineOutput = False } @@ -482,6 +488,7 @@ , configFuthark :: FilePath , configRunner :: FilePath , configExtraCompilerOptions :: [String] + , configTuning :: Maybe String , configExtraOptions :: [String] -- ^ Extra options passed to the programs being run. } diff --git a/src/Futhark/CodeGen/Backends/CCUDA.hs b/src/Futhark/CodeGen/Backends/CCUDA.hs --- a/src/Futhark/CodeGen/Backends/CCUDA.hs +++ b/src/Futhark/CodeGen/Backends/CCUDA.hs @@ -51,24 +51,31 @@ cliOptions :: [Option] cliOptions = [ Option { optionLongName = "dump-cuda" , optionShortName = Nothing - , optionArgument = RequiredArgument - , optionAction = [C.cstm|futhark_context_config_dump_program_to(cfg, optarg);|] + , optionArgument = RequiredArgument "FILE" + , optionAction = [C.cstm|{futhark_context_config_dump_program_to(cfg, optarg); + entry_point = NULL;}|] } , Option { optionLongName = "load-cuda" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "FILE" , optionAction = [C.cstm|futhark_context_config_load_program_from(cfg, optarg);|] } , Option { optionLongName = "dump-ptx" , optionShortName = Nothing - , optionArgument = RequiredArgument - , optionAction = [C.cstm|futhark_context_config_dump_ptx_to(cfg, optarg);|] + , optionArgument = RequiredArgument "FILE" + , optionAction = [C.cstm|{futhark_context_config_dump_ptx_to(cfg, optarg); + entry_point = NULL;}|] } , Option { optionLongName = "load-ptx" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "FILE" , optionAction = [C.cstm|futhark_context_config_load_ptx_from(cfg, optarg);|] } + , Option { optionLongName = "nvrtc-option" + , optionShortName = Nothing + , optionArgument = RequiredArgument "OPT" + , optionAction = [C.cstm|futhark_context_config_add_nvrtc_option(cfg, optarg);|] + } , Option { optionLongName = "print-sizes" , optionShortName = Nothing , optionArgument = NoArgument @@ -81,6 +88,17 @@ exit(0); }|] } + , Option { optionLongName = "tuning" + , optionShortName = Nothing + , optionArgument = RequiredArgument "FILE" + , optionAction = [C.cstm|{ + char *fname = optarg; + char *ret = load_tuning_file(optarg, cfg, (int(*)(void*, const char*, size_t)) + futhark_context_config_set_size); + if (ret != NULL) { + panic(1, "When loading tuning from '%s': %s\n", optarg, ret); + }}|] + } ] writeCUDAScalar :: GC.WriteScalar OpenCL () @@ -138,12 +156,17 @@ ++ "' from '" ++ show srcSpace ++ "'." staticCUDAArray :: GC.StaticArray OpenCL () -staticCUDAArray name "device" t vals = do +staticCUDAArray name "device" t vs = do let ct = GC.primTypeToCType t - vals' = [[C.cinit|$exp:v|] | v <- map GC.compilePrimValue vals] - num_elems = length vals name_realtype <- newVName $ baseString name ++ "_realtype" - GC.libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:num_elems] = {$inits:vals'};|] + num_elems <- case vs of + ArrayValues vs' -> do + let vs'' = [[C.cinit|$exp:v|] | v <- map GC.compilePrimValue vs'] + GC.libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs'')] = {$inits:vs''};|] + return $ length vs'' + ArrayZeros n -> do + GC.libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|] + return n -- Fake a memory block. GC.contextField (pretty name) [C.cty|struct memblock_device|] Nothing -- During startup, copy the data to where we need it. @@ -274,4 +297,4 @@ printSizes = intercalate [[C.cstm|fprintf(stderr, ", ");|]] . map printSize printSize e = - [[C.cstm|fprintf(stderr, "%zu", $exp:e);|]] + [[C.cstm|fprintf(stderr, "%d", $exp:e);|]] diff --git a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs @@ -77,6 +77,8 @@ ([C.cedecl|struct $id:s;|], [C.cedecl|struct $id:s { struct cuda_config cu_cfg; size_t sizes[$int:num_sizes]; + int num_nvrtc_opts; + const char **nvrtc_opts; };|]) let size_value_inits = map (\i -> [C.cstm|cfg->sizes[$int:i] = 0;|]) @@ -89,6 +91,9 @@ return NULL; } + cfg->num_nvrtc_opts = 0; + cfg->nvrtc_opts = malloc(sizeof(const char*)); + cfg->nvrtc_opts[0] = NULL; $stms:size_value_inits cuda_config_init(&cfg->cu_cfg, $int:num_sizes, size_names, size_vars, @@ -99,9 +104,19 @@ 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->nvrtc_opts); free(cfg); }|]) + GC.publicDef_ "context_config_add_nvrtc_option" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) { + cfg->nvrtc_opts[cfg->num_nvrtc_opts] = opt; + cfg->num_nvrtc_opts++; + cfg->nvrtc_opts = realloc(cfg->nvrtc_opts, (cfg->num_nvrtc_opts+1) * sizeof(const char*)); + cfg->nvrtc_opts[cfg->num_nvrtc_opts] = NULL; + }|]) + 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) { @@ -222,7 +237,7 @@ create_lock(&ctx->lock); $stms:init_fields - cuda_setup(&ctx->cuda, cuda_program); + cuda_setup(&ctx->cuda, cuda_program, cfg->nvrtc_opts); $stms:(map (loadKernelByName) kernel_names) $stms:final_inits diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs --- a/src/Futhark/CodeGen/Backends/COpenCL.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL.hs @@ -45,6 +45,7 @@ } include_opencl_h = unlines ["#define CL_USE_DEPRECATED_OPENCL_1_2_APIS", "#ifdef __APPLE__", + "#define CL_SILENCE_DEPRECATION", "#include ", "#else", "#include ", @@ -53,54 +54,61 @@ cliOptions :: [Option] cliOptions = [ Option { optionLongName = "platform" , optionShortName = Just 'p' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "NAME" , optionAction = [C.cstm|futhark_context_config_set_platform(cfg, optarg);|] } , Option { optionLongName = "device" , optionShortName = Just 'd' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "NAME" , optionAction = [C.cstm|futhark_context_config_set_device(cfg, optarg);|] } , Option { optionLongName = "default-group-size" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "INT" , optionAction = [C.cstm|futhark_context_config_set_default_group_size(cfg, atoi(optarg));|] } , Option { optionLongName = "default-num-groups" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "INT" , optionAction = [C.cstm|futhark_context_config_set_default_num_groups(cfg, atoi(optarg));|] } , Option { optionLongName = "default-tile-size" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "INT" , optionAction = [C.cstm|futhark_context_config_set_default_tile_size(cfg, atoi(optarg));|] } , Option { optionLongName = "default-threshold" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "INT" , 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);|] + , optionArgument = RequiredArgument "FILE" + , optionAction = [C.cstm|{futhark_context_config_dump_program_to(cfg, optarg); + entry_point = NULL;}|] } , Option { optionLongName = "load-opencl" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "FILE" , optionAction = [C.cstm|futhark_context_config_load_program_from(cfg, optarg);|] } , Option { optionLongName = "dump-opencl-binary" , optionShortName = Nothing - , optionArgument = RequiredArgument - , optionAction = [C.cstm|futhark_context_config_dump_binary_to(cfg, optarg);|] + , optionArgument = RequiredArgument "FILE" + , optionAction = [C.cstm|{futhark_context_config_dump_binary_to(cfg, optarg); + entry_point = NULL;}|] } , Option { optionLongName = "load-opencl-binary" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "FILE" , optionAction = [C.cstm|futhark_context_config_load_binary_from(cfg, optarg);|] } + , Option { optionLongName = "build-option" + , optionShortName = Nothing + , optionArgument = RequiredArgument "OPT" + , optionAction = [C.cstm|futhark_context_config_add_build_option(cfg, optarg);|] + } , Option { optionLongName = "print-sizes" , optionShortName = Nothing , optionArgument = NoArgument @@ -115,7 +123,7 @@ } , Option { optionLongName = "size" , optionShortName = Nothing - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "NAME=INT" , optionAction = [C.cstm|{ char *name = optarg; char *equals = strstr(optarg, "="); @@ -130,6 +138,17 @@ panic(1, "Invalid argument for size option: %s\n", optarg); }}|] } + , Option { optionLongName = "tuning" + , optionShortName = Nothing + , optionArgument = RequiredArgument "FILE" + , optionAction = [C.cstm|{ + char *fname = optarg; + char *ret = load_tuning_file(optarg, cfg, (int(*)(void*, const char*, size_t)) + futhark_context_config_set_size); + if (ret != NULL) { + panic(1, "When loading tuning from '%s': %s\n", optarg, ret); + }}|] + } ] -- We detect the special case of writing a constant and turn it into a @@ -240,10 +259,15 @@ 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'};|] + num_elems <- case vs of + ArrayValues vs' -> do + let vs'' = [[C.cinit|$exp:v|] | v <- map GC.compilePrimValue vs'] + GC.libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs'')] = {$inits:vs''};|] + return $ length vs'' + ArrayZeros n -> do + GC.libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|] + return n -- Fake a memory block. GC.contextField (pretty name) [C.cty|struct memblock_device|] Nothing -- During startup, copy the data to where we need it. diff --git a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs @@ -61,6 +61,8 @@ ([C.cedecl|struct $id:s;|], [C.cedecl|struct $id:s { struct opencl_config opencl; size_t sizes[$int:num_sizes]; + int num_build_opts; + const char **build_opts; };|]) let size_value_inits = map (\i -> [C.cstm|cfg->sizes[$int:i] = 0;|]) [0..M.size sizes-1] @@ -72,6 +74,9 @@ return NULL; } + cfg->num_build_opts = 0; + cfg->build_opts = malloc(sizeof(const char*)); + cfg->build_opts[0] = NULL; $stms:size_value_inits opencl_config_init(&cfg->opencl, $int:num_sizes, size_names, size_vars, @@ -82,9 +87,19 @@ 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->build_opts); free(cfg); }|]) + GC.publicDef_ "context_config_add_build_option" GC.InitDecl $ \s -> + ([C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, const char *opt) { + cfg->build_opts[cfg->num_build_opts] = opt; + cfg->num_build_opts++; + cfg->build_opts = realloc(cfg->build_opts, (cfg->num_build_opts+1) * sizeof(const char*)); + cfg->build_opts[cfg->num_build_opts] = NULL; + }|]) + 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) { @@ -170,6 +185,27 @@ return 0; } } + + if (strcmp(size_name, "default_group_size") == 0) { + cfg->opencl.default_group_size = size_value; + return 0; + } + + if (strcmp(size_name, "default_num_groups") == 0) { + cfg->opencl.default_num_groups = size_value; + return 0; + } + + if (strcmp(size_name, "default_threshold") == 0) { + cfg->opencl.default_threshold = size_value; + return 0; + } + + if (strcmp(size_name, "default_tile_size") == 0) { + cfg->opencl.default_tile_size = size_value; + return 0; + } + return 1; }|]) @@ -231,7 +267,7 @@ $stms:set_required_types init_context_early(cfg, ctx); - typename cl_program prog = setup_opencl(&ctx->opencl, opencl_program, required_types); + typename cl_program prog = setup_opencl(&ctx->opencl, opencl_program, required_types, cfg->build_opts); init_context_late(cfg, ctx, prog); return ctx; }|]) @@ -248,7 +284,7 @@ $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); + typename cl_program prog = setup_opencl_with_command_queue(&ctx->opencl, queue, opencl_program, required_types, cfg->build_opts); init_context_late(cfg, ctx, prog); return ctx; }|]) @@ -403,6 +439,7 @@ NumGroups -> [C.cexp|ctx->cfg.default_num_groups|] GroupSize -> [C.cexp|ctx->cfg.default_group_size|] TileSize -> [C.cexp|ctx->cfg.default_tile_size|] + Threshold -> [C.cexp|ctx->cfg.default_threshold|] get_size = case what of HeuristicConst x -> diff --git a/src/Futhark/CodeGen/Backends/CSOpenCL.hs b/src/Futhark/CodeGen/Backends/CSOpenCL.hs --- a/src/Futhark/CodeGen/Backends/CSOpenCL.hs +++ b/src/Futhark/CodeGen/Backends/CSOpenCL.hs @@ -337,11 +337,17 @@ -- 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) + CS.staticMemDecl $ AssignTyped (Composite $ ArrayT t') (Var tmp_arr) $ Just $ + case vs of Imp.ArrayValues vs' -> + CreateArray (CS.compilePrimTypeToAST t) $ Right $ map CS.compilePrimValue vs' + Imp.ArrayZeros n -> + CreateArray (CS.compilePrimTypeToAST t) $ Left n -- Create memory block on the device. ptr <- newVName' "ptr" - let size = Integer $ genericLength vs * Imp.primByteSize t + let num_elems = case vs of Imp.ArrayValues vs' -> length vs' + Imp.ArrayZeros n -> n + size = Integer $ toInteger num_elems * Imp.primByteSize t CS.staticMemAlloc $ Reassign (Var name') (CS.simpleCall "EmptyMemblock" [Var "Ctx.EMPTY_MEM_HANDLE"]) errcode <- CS.compileName <$> newVName "errCode" @@ -375,7 +381,7 @@ let createTuple = "createTuple_"++ pretty bt' return $ CS.simpleCall createTuple [ memblockFromMem mem, Var "Ctx.OpenCL.Queue", nbytes - , CreateArray (Primitive $ CSInt Int64T) dims'] + , CreateArray (Primitive $ CSInt Int64T) $ Right dims'] where dims' = map CS.compileDim dims packArrayOutput _ sid _ _ _ = diff --git a/src/Futhark/CodeGen/Backends/CSOpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/CSOpenCL/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/CSOpenCL/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/CSOpenCL/Boilerplate.hs @@ -271,6 +271,7 @@ NumGroups -> Var "Ctx.OpenCL.Cfg.DefaultNumGroups" GroupSize -> Var "Ctx.OpenCL.Cfg.DefaultGroupSize" TileSize -> Var "Ctx.OpenCL.Cfg.DefaultTileSize" + Threshold -> Var "Ctx.OpenCL.Cfg.DefaultThreshold" get_size = case what of HeuristicConst x -> diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -161,7 +161,7 @@ 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 () +type StaticArray op s = VName -> SpaceId -> PrimType -> ArrayContents -> CompilerM op s () -- | Copy from one memory block to another. type Copy op s = C.Exp -> C.Exp -> Space -> @@ -1221,12 +1221,12 @@ benchmarkOptions = [ Option { optionLongName = "write-runtime-to" , optionShortName = Just 't' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "FILE" , optionAction = set_runtime_file } , Option { optionLongName = "runs" , optionShortName = Just 'r' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "INT" , optionAction = set_num_runs } , Option { optionLongName = "debugging" @@ -1241,8 +1241,8 @@ } , Option { optionLongName = "entry-point" , optionShortName = Just 'e' - , optionArgument = RequiredArgument - , optionAction = [C.cstm|entry_point = optarg;|] + , optionArgument = RequiredArgument "NAME" + , optionAction = [C.cstm|if (entry_point != NULL) entry_point = optarg;|] } , Option { optionLongName = "binary-output" , optionShortName = Just 'b' @@ -1356,8 +1356,11 @@ static typename FILE *runtime_file; static int perform_warmup = 0; static int num_runs = 1; +// If the entry point is NULL, the program will terminate after doing initialisation and such. static const char *entry_point = "main"; +$esc:tuning_h + $func:option_parser $edecls:cli_entry_point_decls @@ -1390,31 +1393,33 @@ 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 != 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); + 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; } - return 1; - } - entry_point_fun(ctx); + entry_point_fun(ctx); - if (runtime_file != NULL) { - fclose(runtime_file); - } + if (runtime_file != NULL) { + fclose(runtime_file); + } - futhark_debugging_report(ctx); + futhark_debugging_report(ctx); + } futhark_context_free(ctx); futhark_context_config_free(cfg); @@ -1489,6 +1494,7 @@ values_h = $(embedStringFile "rts/c/values.h") timing_h = $(embedStringFile "rts/c/timing.h") lock_h = $(embedStringFile "rts/c/lock.h") + tuning_h = $(embedStringFile "rts/c/tuning.h") compileFun :: (Name, Function op) -> CompilerM op s (C.Definition, C.Func) compileFun (fname, Function _ outputs inputs body _ _) = do @@ -1805,10 +1811,14 @@ 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'};|] + let ct = primTypeToCType t + case vs of + ArrayValues vs' -> do + let vs'' = [[C.cinit|$exp:(compilePrimValue v)|] | v <- vs'] + libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs')] = {$inits:vs''};|] + ArrayZeros n -> + libDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|] -- Fake a memory block. contextField (pretty name) [C.cty|struct memblock|] $ diff --git a/src/Futhark/CodeGen/Backends/GenericC/Options.hs b/src/Futhark/CodeGen/Backends/GenericC/Options.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Options.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Options.hs @@ -28,7 +28,8 @@ -- | Whether an option accepts an argument. data OptionArgument = NoArgument - | RequiredArgument + | RequiredArgument String + -- ^ The 'String' becomes part of the help text. | OptionalArgument -- | Generate an option parser as a function of the given name, that @@ -52,7 +53,8 @@ panic(-1, "Missing argument for option %s\n", argv[optind-1]); } if ($id:chosen_option == '?') { - panic(-1, "Unknown option %s\n", argv[optind-1]); + fprintf(stderr, "Usage: %s: %s\n", fut_progname, $string:option_descriptions); + panic(1, "Unknown option: %s\n", argv[optind-1]); } } return optind; @@ -62,15 +64,28 @@ option_string = ':' : optionString options option_applications = optionApplications chosen_option options option_fields = optionFields options + option_descriptions = unwords $ map describeOption options +describeOption :: Option -> String +describeOption opt = + concat [ "[" + , maybe "" (\c -> "-" ++ [c] ++ "/") $ optionShortName opt + , "--" ++ optionLongName opt + , case optionArgument opt of + NoArgument -> "" + RequiredArgument what -> " " ++ what + OptionalArgument -> " [ARG]" + , "]" + ] + 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" + NoArgument -> "no_argument" + RequiredArgument _ -> "required_argument" + OptionalArgument -> "optional_argument" optionApplications :: String -> [Option] -> [C.Stm] optionApplications chosen_option = zipWith check [(1::Int)..] @@ -86,6 +101,6 @@ short <- optionShortName option return $ short : case optionArgument option of - NoArgument -> "" - RequiredArgument -> ":" - OptionalArgument -> "::" + NoArgument -> "" + RequiredArgument _ -> ":" + OptionalArgument -> "::" diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp.hs b/src/Futhark/CodeGen/Backends/GenericCSharp.hs --- a/src/Futhark/CodeGen/Backends/GenericCSharp.hs +++ b/src/Futhark/CodeGen/Backends/GenericCSharp.hs @@ -114,7 +114,7 @@ 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 () +type StaticArray op s = VName -> Imp.SpaceId -> PrimType -> Imp.ArrayContents -> CompilerM op s () -- | Construct the C# array being returned from an entry point. type EntryOutput op s = VName -> Imp.SpaceId -> @@ -381,6 +381,11 @@ , optionAction = [ Reassign (Var "EntryPoint") $ Var "optarg" ] } + , Option { optionLongName = "tuning" + , optionShortName = Nothing + , optionArgument = RequiredArgument + , optionAction = [Escape "FutharkConfigLoadTuning(ref Cfg, optarg);"] + } ] -- | The class generated by the code generator must have a @@ -644,7 +649,7 @@ 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] + return $ simpleCall createTuple [src, CreateArray (Primitive $ CSInt Int64T) $ Right $ map compileDim dims] entryPointOutput (Imp.TransparentValue (Imp.ArrayValue mem _ (Imp.Space sid) bt ept dims)) = do unRefMem mem (Imp.Space sid) @@ -804,7 +809,7 @@ 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) + let size = callMethod (CreateArray (Primitive $ CSInt Int32T) $ Right $ map compileDim $ outer:shape) "Aggregate" [ Integer 1 , Lambda (Tuple [Var "acc", Var "val"]) [Exp $ BinOp "*" (Var "acc") (Var "val")] @@ -1257,13 +1262,15 @@ stms [Assign (Var $ "init_"++name') $ simpleCall "unwrapArray" [ - CreateArray (compilePrimTypeToAST t) (map compilePrimValue vs) + case vs of Imp.ArrayValues vs' -> + CreateArray (compilePrimTypeToAST t) $ Right $ map compilePrimValue vs' + Imp.ArrayZeros n -> + CreateArray (compilePrimTypeToAST t) $ Left n , 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 <*> diff --git a/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs b/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs --- a/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs +++ b/src/Futhark/CodeGen/Backends/GenericCSharp/AST.hs @@ -143,7 +143,7 @@ | Call CSExp [CSArg] | CallMethod CSExp CSExp [CSArg] | CreateObject CSExp [CSArg] - | CreateArray CSType [CSExp] + | CreateArray CSType (Either Int [CSExp]) | CreateSystemTuple [CSExp] | AllocArray CSType CSExp | Cast CSType CSExp @@ -182,7 +182,8 @@ 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 (CreateArray t (Left n)) = text "new" <+> ppr t <> brackets (ppr n) + ppr (CreateArray t (Right 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? diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs --- a/src/Futhark/CodeGen/Backends/GenericPython.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython.hs @@ -91,7 +91,7 @@ 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 () +type StaticArray op s = VName -> Imp.SpaceId -> PrimType -> Imp.ArrayContents -> CompilerM op s () -- | Construct the Python array being returned from an entry point. type EntryOutput op s = VName -> Imp.SpaceId -> @@ -249,7 +249,7 @@ standardOptions = [ Option { optionLongName = "write-runtime-to" , optionShortName = Just 't' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "str" , optionAction = [ If (Var "runtime_file") @@ -260,7 +260,7 @@ }, Option { optionLongName = "runs" , optionShortName = Just 'r' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "str" , optionAction = [ Assign (Var "num_runs") $ Var "optarg" , Assign (Var "do_warmup_run") $ Bool True @@ -268,7 +268,7 @@ }, Option { optionLongName = "entry-point" , optionShortName = Just 'e' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "str" , optionAction = [ Assign (Var "entry_point") $ Var "optarg" ] }, @@ -277,6 +277,11 @@ , optionShortName = Just 'b' , optionArgument = NoArgument , optionAction = [Pass] + }, + Option { optionLongName = "tuning" + , optionShortName = Nothing + , optionArgument = RequiredArgument "open" + , optionAction = [Exp $ simpleCall "read_tuning_file" [Var "sizes", Var "optarg"]] } ] @@ -312,7 +317,9 @@ Just _ -> "" return $ maybe_shebang ++ pretty (PyProg $ imports ++ - [Import "argparse" Nothing] ++ + [ Import "argparse" Nothing + , Assign (Var "sizes") $ Dict [] + ] ++ defines ++ [Escape pyUtility] ++ prog') @@ -879,9 +886,15 @@ -- 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") arr_name) $ case vs of + Imp.ArrayValues vs' -> + Call (Var "np.array") + [Arg $ List $ map compilePrimValue vs', + ArgKeyword "dtype" $ Var $ compilePrimToNp t] + Imp.ArrayZeros n -> + Call (Var "np.zeros") + [Arg $ Integer $ fromIntegral n, + ArgKeyword "dtype" $ Var $ compilePrimToNp t] atInit $ Assign (Field (Var "self") name') $ simpleCall "unwrapArray" [Field (Var "self") arr_name] diff --git a/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs b/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs --- a/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython/Definitions.hs @@ -4,6 +4,7 @@ , pyUtility , pyValues , pyPanic + , pyTuning ) where import Data.FileEmbed @@ -19,3 +20,6 @@ pyPanic :: String pyPanic = $(embedStringFile "rts/python/panic.py") + +pyTuning :: String +pyTuning = $(embedStringFile "rts/python/tuning.py") diff --git a/src/Futhark/CodeGen/Backends/GenericPython/Options.hs b/src/Futhark/CodeGen/Backends/GenericPython/Options.hs --- a/src/Futhark/CodeGen/Backends/GenericPython/Options.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython/Options.hs @@ -24,7 +24,7 @@ -- | Whether an option accepts an argument. data OptionArgument = NoArgument - | RequiredArgument + | RequiredArgument String | OptionalArgument -- | Generate option parsing code that accepts the given command line options. Will read from @sys.argv@. @@ -43,13 +43,15 @@ map executeOption options where parseOption option = Exp $ Call (Var "parser.add_argument") $ - map (Arg . String) name_args ++ argument_args + map (Arg . String) name_args ++ + argument_args where name_args = maybe id ((:) . ('-':) . (:[])) (optionShortName option) ["--" ++ optionLongName option] argument_args = case optionArgument option of - RequiredArgument -> + RequiredArgument t -> [ArgKeyword "action" (String "append"), - ArgKeyword "default" $ List []] + ArgKeyword "default" $ List [], + ArgKeyword "type" $ Var t] NoArgument -> [ArgKeyword "action" (String "append_const"), diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs --- a/src/Futhark/CodeGen/Backends/PyOpenCL.hs +++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs @@ -34,10 +34,15 @@ [Assign (Var "synchronous") $ Bool False, Assign (Var "preferred_platform") None, Assign (Var "preferred_device") None, + Assign (Var "default_threshold") None, + Assign (Var "default_group_size") None, + Assign (Var "default_num_groups") None, + Assign (Var "default_tile_size") None, Assign (Var "fut_opencl_src") $ RawStringLiteral $ opencl_prelude ++ opencl_code, Escape pyValues, Escape pyFunctions, - Escape pyPanic] + Escape pyPanic, + Escape pyTuning] let imports = [Import "sys" Nothing, Import "numpy" $ Just "np", Import "ctypes" $ Just "ct", @@ -50,23 +55,59 @@ , "interactive=False" , "platform_pref=preferred_platform" , "device_pref=preferred_device" - , "default_group_size=None" - , "default_num_groups=None" - , "default_tile_size=None" - , "sizes={}"] + , "default_group_size=default_group_size" + , "default_num_groups=default_num_groups" + , "default_tile_size=default_tile_size" + , "default_threshold=default_threshold" + , "sizes=sizes"] [Escape $ openClInit types assign sizes] options = [ Option { optionLongName = "platform" , optionShortName = Just 'p' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "str" , optionAction = [ Assign (Var "preferred_platform") $ Var "optarg" ] } , Option { optionLongName = "device" , optionShortName = Just 'd' - , optionArgument = RequiredArgument + , optionArgument = RequiredArgument "str" , optionAction = [ Assign (Var "preferred_device") $ Var "optarg" ] - }] + } + , Option { optionLongName = "default-threshold" + , optionShortName = Nothing + , optionArgument = RequiredArgument "int" + , optionAction = + [ Assign (Var "default_threshold") $ Var "optarg" ] + } + , Option { optionLongName = "default-group-size" + , optionShortName = Nothing + , optionArgument = RequiredArgument "int" + , optionAction = + [ Assign (Var "default_group_size") $ Var "optarg" ] + } + , Option { optionLongName = "default-num-groups" + , optionShortName = Nothing + , optionArgument = RequiredArgument "int" + , optionAction = + [ Assign (Var "default_num_groups") $ Var "optarg" ] + } + , Option { optionLongName = "default-tile-size" + , optionShortName = Nothing + , optionArgument = RequiredArgument "int" + , optionAction = + [ Assign (Var "default_tile_size") $ Var "optarg" ] + } + , Option { optionLongName = "size" + , optionShortName = Nothing + , optionArgument = RequiredArgument "size_assignment" + , optionAction = + [Assign (Index (Var "sizes") + (IdxExp (Index (Var "optarg") + (IdxExp (Integer 0))))) + (Index (Var "optarg") (IdxExp (Integer 1))) + ] + } + ] Right <$> Py.compileProg module_name constructor imports defines operations () [Exp $ Py.simpleCall "self.queue.finish" []] options prog' @@ -216,14 +257,22 @@ 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] + Py.stm $ Assign (Var name') $ case vs of + Imp.ArrayValues vs' -> + Call (Var "np.array") + [Arg $ List $ map Py.compilePrimValue vs', + ArgKeyword "dtype" $ Var $ Py.compilePrimToNp t] + Imp.ArrayZeros n -> + Call (Var "np.zeros") + [Arg $ Integer $ fromIntegral n, + ArgKeyword "dtype" $ Var $ Py.compilePrimToNp t] + let num_elems = case vs of Imp.ArrayValues vs' -> length vs' + Imp.ArrayZeros n -> n + -- Create memory block on the device. static_mem <- newVName "static_mem" - let size = Integer $ fromIntegral (length vs) * Imp.primByteSize t + let size = Integer $ toInteger num_elems * Imp.primByteSize t allocateOpenCLBuffer static_mem size "device" -- Copy Numpy array to the device memory block. diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs @@ -35,6 +35,7 @@ default_group_size=default_group_size, default_num_groups=default_num_groups, default_tile_size=default_tile_size, + default_threshold=default_threshold, size_heuristics=size_heuristics, required_types=$types', user_sizes=sizes, @@ -68,6 +69,7 @@ NumGroups -> String "num_groups" GroupSize -> String "group_size" TileSize -> String "tile_size" + Threshold -> String "threshold" what' = case what of HeuristicConst x -> Integer $ toInteger x diff --git a/src/Futhark/CodeGen/Backends/SequentialPython.hs b/src/Futhark/CodeGen/Backends/SequentialPython.hs --- a/src/Futhark/CodeGen/Backends/SequentialPython.hs +++ b/src/Futhark/CodeGen/Backends/SequentialPython.hs @@ -27,7 +27,7 @@ Import "numpy" $ Just "np", Import "ctypes" $ Just "ct", Import "time" Nothing] - defines = [Escape pyValues, Escape pyFunctions, Escape pyPanic] + defines = [Escape pyValues, Escape pyFunctions, Escape pyPanic, Escape pyTuning] operations :: GenericPython.Operations Imp.Sequential () operations = GenericPython.defaultOperations { GenericPython.opsCompiler = const $ return () diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs --- a/src/Futhark/CodeGen/ImpCode.hs +++ b/src/Futhark/CodeGen/ImpCode.hs @@ -34,6 +34,7 @@ , index , ErrorMsg(..) , ErrorMsgPart(..) + , ArrayContents(..) -- * Typed enumerations , Count (..) @@ -137,13 +138,22 @@ -- | Type alias for namespace control. type Function = FunctionT +-- | The contents of a statically declared constant array. Such +-- arrays are always unidimensional, and reshaped if necessary in the +-- code that uses them. +data ArrayContents = ArrayValues [PrimValue] + -- ^ Precisely these values. + | ArrayZeros Int + -- ^ This many zeroes. + deriving (Show) + 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] + | DeclareArray VName Space PrimType ArrayContents -- ^ Create an array containing the given values. The -- lifetime of the array will be the entire application. -- This is mostly used for constant arrays, but also for @@ -295,6 +305,10 @@ ppr (ConstSize x) = ppr x ppr (VarSize v) = ppr v +instance Pretty ArrayContents where + ppr (ArrayValues vs) = braces (commasep $ map ppr vs) + ppr (ArrayZeros n) = braces (text "0") <+> text "*" <+> ppr n + instance Pretty op => Pretty (Code op) where ppr (Op op) = ppr op ppr Skip = text "skip" @@ -313,7 +327,7 @@ 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) + equals <+> ppr vs ppr (Allocate name e space) = ppr name <+> text "<-" <+> text "malloc" <> parens (ppr e) <> ppr space ppr (Free name space) = diff --git a/src/Futhark/CodeGen/ImpCode/Kernels.hs b/src/Futhark/CodeGen/ImpCode/Kernels.hs --- a/src/Futhark/CodeGen/ImpCode/Kernels.hs +++ b/src/Futhark/CodeGen/ImpCode/Kernels.hs @@ -103,11 +103,11 @@ instance Pretty KernelUse where ppr (ScalarUse name t) = - text "scalar_copy" <> parens (commasep [ppr name, ppr t]) + oneLine $ text "scalar_copy" <> parens (commasep [ppr name, ppr t]) ppr (MemoryUse name) = - text "mem_copy" <> parens (commasep [ppr name]) + oneLine $ text "mem_copy" <> parens (commasep [ppr name]) ppr (ConstUse name e) = - text "const" <> parens (commasep [ppr name, ppr e]) + oneLine $ text "const" <> parens (commasep [ppr name, ppr e]) instance Pretty HostOp where ppr (GetSize dest key size_class) = diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs --- a/src/Futhark/CodeGen/ImpGen.hs +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -78,9 +78,9 @@ , sComment , sIf, sWhen, sUnless , sOp - , sAlloc + , sDeclareMem, sAlloc, sAlloc_ , sArray, sAllocArray, sStaticArray - , sWrite + , sWrite, sUpdate , (<--) ) where @@ -675,11 +675,8 @@ defCompileBasicOp _ Index{} = return () -defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) = do - MemLocation mem shape ixfun <- entryArrayLocation <$> lookupArray (patElemName pe) - let memdest = sliceArray (MemLocation mem shape ixfun) $ - map (fmap (compileSubExpOfType int32)) slice - copyDWIMDest (ArrayDestination $ Just memdest) [] se [] +defCompileBasicOp (Pattern _ [pe]) (Update _ slice se) = + sUpdate (patElemName pe) (map (fmap (compileSubExpOfType int32)) slice) se defCompileBasicOp (Pattern _ [pe]) (Replicate (Shape ds) se) = do ds' <- mapM compileSubExp ds @@ -736,7 +733,7 @@ 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 + emit $ Imp.DeclareArray static_array dest_space t $ Imp.ArrayValues 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 @@ -1237,19 +1234,33 @@ sOp :: op -> ImpM lore op () sOp = emit . Imp.Op -sAlloc :: String -> Count Bytes -> Space -> ImpM lore op VName -sAlloc name size space = do +dSize :: Imp.Count u -> ImpM lore op Imp.Size +dSize size = + case Imp.innerExp size of + Imp.LeafExp (Imp.ScalarVar size') _ -> return $ Imp.VarSize size' + Imp.ValueExp (IntValue (Int64Value v)) -> return $ Imp.ConstSize v + _ -> do size_var <- dPrim "local_buf_size" int32 + size_var <-- Imp.innerExp size + return $ Imp.VarSize size_var + +sDeclareMem :: String -> Count Bytes -> Space -> ImpM lore op (VName, Imp.MemSize) +sDeclareMem name size space = do + name' <- newVName name - size' <- case Imp.innerExp size of - Imp.LeafExp (Imp.ScalarVar size') _ -> return $ Imp.VarSize size' - Imp.ValueExp (IntValue (Int64Value v)) -> return $ Imp.ConstSize v - _ -> do size_var <- dPrim "local_buf_size" int32 - size_var <-- Imp.innerExp size - return $ Imp.VarSize size_var + size' <- dSize size emit $ Imp.DeclareMem name' space - fake <- asks $ elem space . envFakeMemory - unless fake $ emit $ Imp.Allocate name' size space addVar name' $ MemVar Nothing $ MemEntry size' space + return (name', size') + +sAlloc_ :: VName -> Imp.MemSize -> Space -> ImpM lore op () +sAlloc_ name' size' space = do + fake <- asks $ elem space . envFakeMemory + unless fake $ emit $ Imp.Allocate name' (Imp.memSizeToExp size') space + +sAlloc :: String -> Count Bytes -> Space -> ImpM lore op VName +sAlloc name size space = do + (name', size') <- sDeclareMem name size space + sAlloc_ name' size' space return name' sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM lore op VName @@ -1268,21 +1279,28 @@ ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape -- | Uses linear/iota index function. -sStaticArray :: String -> Space -> PrimType -> [PrimValue] -> ImpM lore op VName +sStaticArray :: String -> Space -> PrimType -> Imp.ArrayContents -> ImpM lore op VName sStaticArray name space pt vs = do - let shape = Shape [constant $ length vs] - size = Imp.ConstSize $ fromIntegral (length vs) * primByteSize pt + let num_elems = case vs of Imp.ArrayValues vs' -> genericLength vs' + Imp.ArrayZeros n -> fromIntegral n + shape = Shape [constant num_elems] + mem_size = Imp.ConstSize $ num_elems * primByteSize pt mem <- newVName $ name ++ "_mem" emit $ Imp.DeclareArray mem space pt vs - addVar mem $ MemVar Nothing $ MemEntry size space - sArray name pt shape $ - ArrayIn mem $ IxFun.iota $ map (primExpFromSubExp int32) $ shapeDims shape + addVar mem $ MemVar Nothing $ MemEntry mem_size space + sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems] sWrite :: VName -> [Imp.Exp] -> PrimExp Imp.ExpLeaf -> ImpM lore op () sWrite arr is v = do (mem, space, offset) <- fullyIndexArray arr is vol <- asks envVolatility emit $ Imp.Write mem offset (primExpType v) space vol v + +sUpdate :: VName -> Slice Imp.Exp -> SubExp -> ImpM lore op () +sUpdate arr slice v = do + MemLocation mem shape ixfun <- entryArrayLocation <$> lookupArray arr + let memdest = sliceArray (MemLocation mem shape ixfun) slice + copyDWIMDest (ArrayDestination $ Just memdest) [] v [] -- | ASsignment. (<--) :: VName -> Imp.Exp -> ImpM lore op () diff --git a/src/Futhark/CodeGen/ImpGen/Kernels.hs b/src/Futhark/CodeGen/ImpGen/Kernels.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels.hs @@ -23,9 +23,9 @@ import qualified Futhark.CodeGen.ImpGen as ImpGen import Futhark.CodeGen.ImpGen.Kernels.Base import Futhark.CodeGen.ImpGen.Kernels.SegRed +import Futhark.CodeGen.ImpGen.Kernels.SegGenRed import Futhark.CodeGen.ImpGen (sFor, sWhen, - sOp, - dPrim, dPrim_, dPrimV) + sOp) import Futhark.CodeGen.ImpGen.Kernels.Transpose import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun import Futhark.CodeGen.SetDefaultSpace @@ -108,14 +108,12 @@ show (baseTag $ kernelGlobalThreadIdVar constants) } --- First handle the simple case of a non-segmented reduction. Our --- strategy is the conventional approach of generating two kernels: --- one where each group is given a chunk of the total input and --- produces a partial result per group, and then a final kernel that --- combines the per-group partial results. kernelCompiler pat (SegRed space comm red_op nes _ body) = compileSegRed pat space comm red_op nes body +kernelCompiler pat (SegGenRed space ops _ body) = + compileSegGenRed pat space ops body + kernelCompiler pat e = compilerBugS $ "ImpGen.kernelCompiler: Invalid pattern\n " ++ pretty pat ++ "\nfor expression\n " ++ pretty e @@ -127,33 +125,11 @@ n' <- ImpGen.compileSubExp n x' <- ImpGen.compileSubExp x s' <- ImpGen.compileSubExp s - destloc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray (patElemName pe) - (constants, set_constants) <- simpleKernelConstants n' "iota" - sKernel constants "iota" $ do - set_constants - let gtid = kernelGlobalThreadId constants - sWhen (kernelThreadActive constants) $ do - (destmem, destspace, destidx) <- - ImpGen.fullyIndexArray' destloc [gtid] (IntType et) - - ImpGen.emit $ - Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ - Imp.ConvOpExp (SExt Int32 et) gtid * s' + x' - -expCompiler (Pattern _ [pe]) (BasicOp (Replicate (Shape ds) se)) = do - t <- subExpType se - - dims <- mapM ImpGen.compileSubExp $ ds ++ arrayDims t - (constants, set_constants) <- - simpleKernelConstants (product dims) "replicate" - - let is' = unflattenIndex dims $ kernelGlobalThreadId constants + sIota (patElemName pe) n' x' s' et - sKernel constants "replicate" $ do - set_constants - sWhen (kernelThreadActive constants) $ - ImpGen.copyDWIM (patElemName pe) is' se $ drop (length ds) is' +expCompiler (Pattern _ [pe]) (BasicOp (Replicate shape se)) = + sReplicate (patElemName pe) shape se -- Allocation in the "local" space is just a placeholder. expCompiler _ (Op (Alloc _ (Space "local"))) = @@ -195,29 +171,7 @@ srcmem (bytes srcoffset) srcspace $ (n * row_size) `Imp.withElemType` bt - | otherwise = do - - -- 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 - kernel_size = Imp.innerExp n * product (drop 1 shape) - - (constants, set_constants) <- simpleKernelConstants kernel_size "copy" - - sKernel constants "copy" $ do - set_constants - - let gtid = kernelGlobalThreadId constants - dest_is = unflattenIndex shape_se gtid - src_is = dest_is - - (_, destspace, destidx) <- ImpGen.fullyIndexArray' destloc dest_is bt - (_, srcspace, srcidx) <- ImpGen.fullyIndexArray' srcloc src_is bt - - sWhen (gtid .<. kernel_size) $ ImpGen.emit $ - Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ - Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile + | otherwise = sCopy bt destloc srcloc n mapTransposeForType :: PrimType -> ImpGen.ImpM ExplicitMemory Imp.HostOp Name mapTransposeForType bt = do @@ -372,39 +326,6 @@ let (mapped, notmapped) = splitAt r1 shape (pretrans, posttrans) = f $ splitAt r2 notmapped in (product mapped, product pretrans, product posttrans) - -simpleKernelConstants :: Imp.Exp -> String - -> CallKernelGen (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ()) -simpleKernelConstants kernel_size desc = do - thread_gtid <- newVName $ desc ++ "_gtid" - thread_ltid <- newVName $ desc ++ "_ltid" - group_id <- newVName $ desc ++ "_gid" - (group_size, num_groups) <- computeMapKernelGroups kernel_size - let set_constants = do - dPrim_ thread_gtid int32 - dPrim_ thread_ltid int32 - dPrim_ group_id int32 - sOp (Imp.GetGlobalId thread_gtid 0) - sOp (Imp.GetLocalId thread_ltid 0) - sOp (Imp.GetGroupId group_id 0) - - return (KernelConstants - (Imp.var thread_gtid int32) (Imp.var thread_ltid int32) (Imp.var group_id int32) - thread_gtid thread_ltid group_id - group_size num_groups (group_size*num_groups) 0 - [] (Imp.var thread_gtid int32 .<. kernel_size) mempty, - - set_constants) - -computeMapKernelGroups :: Imp.Exp -> CallKernelGen (Imp.Exp, Imp.Exp) -computeMapKernelGroups kernel_size = do - group_size <- dPrim "group_size" int32 - fname <- asks ImpGen.envFunction - let group_size_var = Imp.var group_size int32 - group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty group_size - sOp $ Imp.GetSize group_size group_size_key Imp.SizeGroup - num_groups <- dPrimV "num_groups" $ kernel_size `quotRoundingUp` Imp.ConvOpExp (SExt Int32 Int32) group_size_var - return (Imp.var group_size int32, Imp.var num_groups int32) compileKernelBody :: Pattern InKernel -> KernelConstants diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs @@ -19,6 +19,14 @@ , groupScan , isActive , sKernel + , sReplicate + , sIota + , sCopy + + , atomicUpdate + , atomicUpdateLocking + , Locking(..) + , AtomicUpdate ) where @@ -30,7 +38,7 @@ import qualified Data.Set as S import Data.List -import Prelude hiding (quot) +import Prelude hiding (quot, rem) import Futhark.Error import Futhark.MonadFreshNames @@ -45,7 +53,7 @@ dPrim, dPrim_, dPrimV) import Futhark.Tools (partitionChunkedKernelLambdaParameters) import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem, IntegralExp) -import Futhark.Util (splitAt3, maybeNth) +import Futhark.Util (splitAt3, maybeNth, takeLast) type CallKernelGen = ImpGen.ImpM ExplicitMemory Imp.HostOp type InKernelGen = ImpGen.ImpM InKernel Imp.KernelOp @@ -271,12 +279,17 @@ -- Check if bucket is in-bounds bucket' <- mapM ImpGen.compileSubExp bucket w' <- mapM ImpGen.compileSubExp w - sWhen (indexInBounds bucket' w') $ - atomicUpdate arrs bucket op values locking + num_locks <- ImpGen.compileSubExpOfType int32 . arraySize 0 <$> lookupType locks + let locking = Locking locks 0 1 0 $ (`rem` num_locks) . sum + values_params = takeLast (length values) $ lambdaParams op + + sWhen (indexInBounds bucket' w') $ do + forM_ (zip values_params values) $ \(p, v) -> + ImpGen.copyDWIM (paramName p) [] v [] + atomicUpdate arrs bucket' op locking where indexInBounds inds bounds = foldl1 (.&&.) $ zipWith checkBound inds bounds where checkBound ind bound = 0 .<=. ind .&&. ind .<. bound - locking = Locking locks 0 1 0 compileKernelExp _ dest e = compilerBugS $ unlines ["Invalid target", " " ++ show dest, @@ -290,35 +303,57 @@ ImpGen.emit =<< ImpGen.subImpM_ (inKernelOperations constants') m -- | Locking strategy used for an atomic update. -data Locking = Locking { lockingArray :: VName -- ^ Array containing the lock. - , lockingIsUnlocked :: Imp.Exp -- ^ Value for us to consider the lock free. - , lockingToLock :: Imp.Exp -- ^ What to write when we lock it. - , lockingToUnlock :: Imp.Exp -- ^ What to write when we unlock it. - } +data Locking = + Locking { lockingArray :: VName + -- ^ Array containing the lock. + , lockingIsUnlocked :: Imp.Exp + -- ^ Value for us to consider the lock free. + , lockingToLock :: Imp.Exp + -- ^ What to write when we lock it. + , lockingToUnlock :: Imp.Exp + -- ^ What to write when we unlock it. + , lockingMapping :: [Imp.Exp] -> Imp.Exp + -- ^ A transformation from the logical lock index to the + -- physical position in the array. This can also be used + -- to make the lock array smaller. + } +-- | A function for generating code for an atomic update. Assumes +-- that the bucket is in-bounds. +type AtomicUpdate lore = + [VName] -> [Imp.Exp] -> ImpGen.ImpM lore Imp.KernelOp () + atomicUpdate :: ExplicitMemorish lore => - [VName] -> [SubExp] -> Lambda lore -> [SubExp] -> Locking + [VName] -> [Imp.Exp] -> Lambda lore -> Locking -> ImpGen.ImpM lore Imp.KernelOp () +atomicUpdate arrs bucket lam locking = + case atomicUpdateLocking lam of + Left f -> f arrs bucket + Right f -> f locking arrs bucket -atomicUpdate arrs bucket lam values _ +-- | 'atomicUpdate', but where it is explicitly visible whether a +-- locking strategy is necessary. +atomicUpdateLocking :: ExplicitMemorish lore => + Lambda lore + -> Either (AtomicUpdate lore) (Locking -> AtomicUpdate lore) + +atomicUpdateLocking lam | Just ops_and_ts <- splitOp lam, - all ((==32) . primBitSize . snd) ops_and_ts = + all (\(_, t, _) -> primBitSize t == 32) ops_and_ts = Left $ \arrs bucket -> -- If the operator is a vectorised binary operator on 32-bit values, -- we can use a particularly efficient implementation. If the -- operator has an atomic implementation we use that, otherwise it -- is still a binary operator which can be implemented by atomic -- compare-and-swap if 32 bits. - forM_ (zip3 arrs ops_and_ts values) $ \(a, (op, t), val) -> do + forM_ (zip arrs ops_and_ts) $ \(a, (op, t, val)) -> do -- Common variables. old <- dPrim "old" t - bucket' <- mapM ImpGen.compileSubExp bucket - (arr', _a_space, bucket_offset) <- ImpGen.fullyIndexArray a bucket' + (arr', _a_space, bucket_offset) <- ImpGen.fullyIndexArray a bucket - val' <- ImpGen.compileSubExp val case opHasAtomicSupport old arr' bucket_offset op of - Just f -> sOp $ f val' + Just f -> sOp $ f val Nothing -> do -- Code generation target: @@ -331,7 +366,7 @@ -- } while(assumed != old); assumed <- dPrim "assumed" t run_loop <- dPrimV "run_loop" 1 - ImpGen.copyDWIM old [] (Var a) bucket' + ImpGen.copyDWIM old [] (Var a) bucket -- Critical section x <- dPrim "x" t @@ -344,7 +379,7 @@ _ -> (id, id) sWhile (Imp.var run_loop int32) $ do assumed <-- Imp.var old t - x <-- val' + x <-- val y <-- Imp.var assumed t x <-- Imp.BinOpExp op (Imp.var x t) (Imp.var y t) old_bits <- dPrim "old_bits" int32 @@ -359,69 +394,69 @@ let atomic f = Imp.Atomic . f old arr' bucket' atomic <$> Imp.atomicBinOp bop -atomicUpdate arrs bucket op values locking = do +atomicUpdateLocking op = Right $ \locking arrs bucket -> do old <- dPrim "old" int32 continue <- dPrimV "continue" true - -- Check if bucket is in-bounds - bucket' <- mapM ImpGen.compileSubExp bucket - -- Correctly index into locks. (locks', _locks_space, locks_offset) <- - ImpGen.fullyIndexArray (lockingArray locking) bucket' - - -- Preparing parameters - let (acc_params, arr_params) = - splitAt (length values) $ lambdaParams op + ImpGen.fullyIndexArray (lockingArray locking) [lockingMapping locking bucket] -- Critical section let try_acquire_lock = sOp $ Imp.Atomic $ Imp.AtomicCmpXchg old locks' locks_offset (lockingIsUnlocked locking) (lockingToLock locking) lock_acquired = Imp.var old int32 .==. lockingIsUnlocked locking - release_lock = ImpGen.everythingVolatile $ - ImpGen.sWrite (lockingArray locking) bucket' $ lockingToUnlock locking + -- Even the releasing is done with an atomic rather than a + -- simple write, for memory coherency reasons. + release_lock = + sOp $ Imp.Atomic $ + Imp.AtomicCmpXchg old locks' locks_offset (lockingToLock locking) (lockingToUnlock locking) break_loop = continue <-- false - -- We copy the current value and the new value to the parameters. - -- It is important that the right-hand-side is bound first for the - -- (rare) case when we are dealing with arrays. - let bind_acc_params = + -- Preparing parameters. It is assumed that the caller has already + -- filled the arr_params. We copy the current value to the + -- accumulator parameters. + -- + -- Note the use of 'everythingVolatile' when reading and writing the + -- buckets. This was necessary to ensure correct execution on a + -- newer NVIDIA GPU (RTX 2080). The 'volatile' modifiers likely + -- make the writes pass through the (SM-local) L1 cache, which is + -- necessary here, because we are really doing device-wide + -- synchronisation without atomics (naughty!). + let (acc_params, _arr_params) = splitAt (length arrs) $ lambdaParams op + bind_acc_params = + ImpGen.everythingVolatile $ ImpGen.sComment "bind lhs" $ forM_ (zip acc_params arrs) $ \(acc_p, arr) -> - ImpGen.copyDWIM (paramName acc_p) [] (Var arr) bucket' - - let bind_arr_params = - ImpGen.sComment "bind rhs" $ - forM_ (zip arr_params values) $ \(arr_p, val) -> - ImpGen.copyDWIM (paramName arr_p) [] val [] + ImpGen.copyDWIM (paramName acc_p) [] (Var arr) bucket let op_body = ImpGen.sComment "execute operation" $ ImpGen.compileBody' acc_params $ lambdaBody op - do_gen_reduce = ImpGen.sComment "update global result" $ - zipWithM_ (writeArray bucket') arrs $ map (Var . paramName) acc_params + do_gen_reduce = + ImpGen.everythingVolatile $ + ImpGen.sComment "update global result" $ + zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params -- While-loop: Try to insert your value sWhile (Imp.var continue Bool) $ do try_acquire_lock sWhen lock_acquired $ do - ImpGen.dLParams $ lambdaParams op - bind_arr_params + ImpGen.dLParams acc_params bind_acc_params op_body do_gen_reduce + sOp Imp.MemFence release_lock break_loop sOp Imp.MemFence - where writeArray bucket' arr val = - ImpGen.copyDWIM arr bucket' val [] + where writeArray bucket arr val = ImpGen.copyDWIM arr bucket val [] -- | Horizontally fission a lambda that models a binary operator. -splitOp :: Attributes lore => Lambda lore -> Maybe [(BinOp, PrimType)] +splitOp :: Attributes lore => Lambda lore -> Maybe [(BinOp, PrimType, Imp.Exp)] splitOp lam = mapM splitStm $ bodyResult $ lambdaBody lam where n = length $ lambdaReturnType lam - splitStm :: SubExp -> Maybe (BinOp, PrimType) splitStm (Var res) = do Let (Pattern [] [pe]) _ (BasicOp (BinOp op (Var x) (Var y))) <- find (([res]==) . patternNames . stmPattern) $ @@ -432,7 +467,7 @@ guard $ paramName xp == x guard $ paramName yp == y Prim t <- Just $ patElemType pe - return (op, t) + return (op, t, Imp.var (paramName yp) t) splitStm _ = Nothing computeKernelUses :: FreeIn a => @@ -494,10 +529,10 @@ hasExp (ImpGen.ScalarVar e _) = e hasExp (ImpGen.MemVar e _) = e --- | Only some constant expressions quality as *static* expressions, +-- | Only some constant expressions qualify 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 daring an array in C. +-- when declaring an array in C. isStaticExp :: Imp.KernelConstExp -> Bool isStaticExp LeafExp{} = True isStaticExp ValueExp{} = True @@ -853,6 +888,8 @@ sComment "restore correct values for first block" $ sWhen is_first_block write_final_result + sOp Imp.LocalBarrier + inBlockScan :: Maybe (Imp.Exp -> Imp.Exp -> Imp.Exp) -> Imp.Exp -> Imp.Exp @@ -944,6 +981,39 @@ collapse l@((g,_):_) = (g, map snd l) +computeMapKernelGroups :: Imp.Exp -> CallKernelGen (Imp.Exp, Imp.Exp) +computeMapKernelGroups kernel_size = do + group_size <- dPrim "group_size" int32 + fname <- asks ImpGen.envFunction + let group_size_var = Imp.var group_size int32 + group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty group_size + sOp $ Imp.GetSize group_size group_size_key Imp.SizeGroup + num_groups <- dPrimV "num_groups" $ kernel_size `quotRoundingUp` Imp.ConvOpExp (SExt Int32 Int32) group_size_var + return (Imp.var group_size int32, Imp.var num_groups int32) + +simpleKernelConstants :: Imp.Exp -> String + -> CallKernelGen (KernelConstants, ImpGen.ImpM InKernel Imp.KernelOp ()) +simpleKernelConstants kernel_size desc = do + thread_gtid <- newVName $ desc ++ "_gtid" + thread_ltid <- newVName $ desc ++ "_ltid" + group_id <- newVName $ desc ++ "_gid" + (group_size, num_groups) <- computeMapKernelGroups kernel_size + let set_constants = do + dPrim_ thread_gtid int32 + dPrim_ thread_ltid int32 + dPrim_ group_id int32 + sOp (Imp.GetGlobalId thread_gtid 0) + sOp (Imp.GetLocalId thread_ltid 0) + sOp (Imp.GetGroupId group_id 0) + + return (KernelConstants + (Imp.var thread_gtid int32) (Imp.var thread_ltid int32) (Imp.var group_id int32) + thread_gtid thread_ltid group_id + group_size num_groups (group_size*num_groups) 0 + [] (Imp.var thread_gtid int32 .<. kernel_size) mempty, + + set_constants) + sKernel :: KernelConstants -> String -> ImpGen.ImpM InKernel Imp.KernelOp a -> CallKernelGen () sKernel constants name m = do body <- makeAllMemoryGlobal $ @@ -958,3 +1028,69 @@ , Imp.kernelName = nameFromString $ name ++ "_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) } + +-- | Perform a Replicate with a kernel. +sReplicate :: VName -> Shape -> SubExp + -> CallKernelGen () +sReplicate arr (Shape ds) se = do + t <- subExpType se + + dims <- mapM ImpGen.compileSubExp $ ds ++ arrayDims t + (constants, set_constants) <- + simpleKernelConstants (product dims) "replicate" + + let is' = unflattenIndex dims $ kernelGlobalThreadId constants + + sKernel constants "replicate" $ do + set_constants + sWhen (kernelThreadActive constants) $ + ImpGen.copyDWIM arr is' se $ drop (length ds) is' + +-- | Perform an Iota with a kernel. +sIota :: VName -> Imp.Exp -> Imp.Exp -> Imp.Exp -> IntType + -> CallKernelGen () +sIota arr n x s et = do + destloc <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray arr + (constants, set_constants) <- simpleKernelConstants n "iota" + + sKernel constants "iota" $ do + set_constants + let gtid = kernelGlobalThreadId constants + sWhen (kernelThreadActive constants) $ do + (destmem, destspace, destidx) <- + ImpGen.fullyIndexArray' destloc [gtid] (IntType et) + + ImpGen.emit $ + Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ + Imp.ConvOpExp (SExt Int32 et) gtid * s + x + +sCopy :: PrimType + -> ImpGen.MemLocation + -> ImpGen.MemLocation + -> Imp.Count Imp.Elements + -> CallKernelGen () +sCopy bt + destloc@(ImpGen.MemLocation destmem _ _) + srcloc@(ImpGen.MemLocation srcmem srcshape _) + n = do + -- 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 + kernel_size = Imp.innerExp n * product (drop 1 shape) + + (constants, set_constants) <- simpleKernelConstants kernel_size "copy" + + sKernel constants "copy" $ do + set_constants + + let gtid = kernelGlobalThreadId constants + dest_is = unflattenIndex shape_se gtid + src_is = dest_is + + (_, destspace, destidx) <- ImpGen.fullyIndexArray' destloc dest_is bt + (_, srcspace, srcidx) <- ImpGen.fullyIndexArray' srcloc src_is bt + + sWhen (gtid .<. kernel_size) $ ImpGen.emit $ + Imp.Write destmem destidx bt destspace Imp.Nonvolatile $ + Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegGenRed.hs @@ -0,0 +1,275 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleContexts #-} +-- | Our compilation strategy for 'SegGenRed' is based around avoiding +-- bin conflicts. We do this by splitting the input into chunks, and +-- for each chunk computing a single subhistogram. Then we combine +-- the subhistograms using an ordinary segmented reduction ('SegRed'). +-- +-- There are some branches around to efficiently handle the case where +-- we use only a single subhistogram (because it's large), so that we +-- respect the asymptotics, and do not copy the destination array. +module Futhark.CodeGen.ImpGen.Kernels.SegGenRed + ( compileSegGenRed ) + where + +import Control.Monad.Except +import Data.Maybe +import Data.List + +import Prelude hiding (quot, rem) + +import Futhark.MonadFreshNames +import Futhark.Representation.ExplicitMemory +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFun +import Futhark.Pass.ExplicitAllocations() +import qualified Futhark.CodeGen.ImpCode.Kernels as Imp +import qualified Futhark.CodeGen.ImpGen as ImpGen +import Futhark.CodeGen.ImpGen ((<--), + sFor, sComment, sIf, sWhen, sArray, + dPrim_, dPrimV) +import Futhark.CodeGen.ImpGen.Kernels.SegRed (compileSegRed') +import Futhark.CodeGen.ImpGen.Kernels.Base +import Futhark.Util.IntegralExp (quotRoundingUp, quot, rem) +import Futhark.Util (chunks, mapAccumLM, splitFromEnd, takeLast) +import Futhark.Construct (fullSliceNum) + +prepareAtomicUpdate :: Maybe Locking -> [VName] -> Lambda InKernel + -> CallKernelGen (Maybe Locking, + [Imp.Exp] -> ImpGen.ImpM InKernel Imp.KernelOp ()) +prepareAtomicUpdate l dests lam = + -- We need a separate lock array if the opterators are not all of a + -- particularly simple form that permits pure atomic operations. + case (l, atomicUpdateLocking lam) of + (_, Left f) -> return (l, f dests) + (Just l', Right f) -> return (l, f l' dests) + (Nothing, Right f) -> do + -- The number of locks used here is too low, but since we are + -- currently forced to inline a huge list, I'm keeping it down + -- for now. Some quick experiments suggested that it has little + -- impact anyway (maybe the locking case is just too slow). + -- + -- A fun solution would also be to use a simple hashing + -- algorithm to ensure good distribution of locks. + let num_locks = 10000 + locks <- + ImpGen.sStaticArray "genred_locks" (Space "device") int32 $ + Imp.ArrayZeros num_locks + let l' = Locking locks 0 1 0 ((`rem` fromIntegral num_locks) . sum) + return (Just l', f l' dests) + +prepareIntermediateArrays :: [SubExp] -> Imp.Exp -> [GenReduceOp InKernel] + -> CallKernelGen + [(VName, + [VName], + [Imp.Exp] -> ImpGen.ImpM InKernel Imp.KernelOp ())] +prepareIntermediateArrays segment_dims num_threads = fmap snd . mapAccumLM onOp Nothing + where + onOp l op = do + -- 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) + -- + -- Careful to avoid division by zero when genReduceWidth==0. + num_histos <- dPrimV "num_histos" $ num_threads `quotRoundingUp` + BinOpExp (SMax Int32) 1 (ImpGen.compileSubExpOfType int32 (genReduceWidth op)) + + ImpGen.emit $ Imp.DebugPrint "num_histograms" int32 $ Imp.var num_histos int32 + + -- Initialise sub-histograms. + -- + -- 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. + + dests <- forM (zip (genReduceDest op) (genReduceNeutral op)) $ \(dest, ne) -> do + dest_t <- lookupType dest + dest_mem <- ImpGen.entryArrayLocation <$> ImpGen.lookupArray dest + let num_elems = foldl' (*) (Imp.var num_histos int32) $ + map (ImpGen.compileSubExpOfType int32) $ + arrayDims dest_t + let size = Imp.elements num_elems `Imp.withElemType` int32 + + (sub_mem, size') <- + ImpGen.sDeclareMem "subhistogram_mem" size $ Space "device" + + let num_segments = length segment_dims + sub_shape = Shape (segment_dims++[Var num_histos]) <> + stripDims num_segments (arrayShape dest_t) + sub_membind = ArrayIn sub_mem $ IxFun.iota $ + map (primExpFromSubExp int32) $ shapeDims sub_shape + subhisto <- sArray "genred_dest" (elemType dest_t) sub_shape sub_membind + + let unitHistoCase = + ImpGen.emit $ + Imp.SetMem sub_mem (ImpGen.memLocationName dest_mem) $ + Space "device" + + multiHistoCase = do + ImpGen.sAlloc_ sub_mem size' $ Space "device" + sReplicate subhisto (Shape $ segment_dims ++ [Var num_histos, genReduceWidth op]) ne + subhisto_t <- lookupType subhisto + let slice = fullSliceNum (map (ImpGen.compileSubExpOfType int32) $ arrayDims subhisto_t) $ + map (unitSlice 0 . ImpGen.compileSubExpOfType int32) segment_dims ++ + [DimFix 0] + ImpGen.sUpdate subhisto slice $ Var dest + + sIf (Imp.var num_histos int32 .==. 1) unitHistoCase multiHistoCase + + return subhisto + + (l', do_op) <- prepareAtomicUpdate l dests $ genReduceOp op + + return (l', (num_histos, dests, do_op)) + +genRedKernel :: [PatElem ExplicitMemory] + -> KernelSpace + -> [GenReduceOp InKernel] + -> Body InKernel + -> CallKernelGen [(VName, [VName])] +genRedKernel map_pes space ops body = do + (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () + let constants = base_constants { kernelThreadActive = true } + (space_is, space_sizes) = unzip $ spaceDimensions space + i32_to_i64 = ConvOpExp (SExt Int32 Int64) + space_sizes_64 = map (i32_to_i64 . ImpGen.compileSubExpOfType int32) space_sizes + total_w_64 = product space_sizes_64 + + histograms <- prepareIntermediateArrays (init space_sizes) (kernelNumThreads constants) ops + + elems_per_thread_64 <- dPrimV "elems_per_thread_64" $ + total_w_64 `quotRoundingUp` + ConvOpExp (SExt Int32 Int64) (kernelNumThreads constants) + + sKernel constants "seggenred" $ allThreads constants $ do + init_constants + + i <- newVName "i" + + -- Compute subhistogram index for each thread, per histogram. + subhisto_inds <- forM histograms $ \(num_histograms, _, _) -> + dPrimV "subhisto_ind" $ + kernelGlobalThreadId constants `quot` + (kernelNumThreads constants `quotRoundingUp` Imp.var num_histograms int32) + + sFor i Int64 (Imp.var elems_per_thread_64 int64) $ 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 <- dPrimV "offset" $ + (i32_to_i64 (kernelGroupId constants) * + (Imp.var elems_per_thread_64 int64 * + i32_to_i64 (kernelGroupSize constants))) + + (Imp.var i int64 * i32_to_i64 (kernelGroupSize constants)) + + j <- dPrimV "j" $ Imp.var offset int64 + i32_to_i64 (kernelLocalThreadId constants) + + -- Construct segment indices. + let setIndex v e = do dPrim_ v int32 + v <-- e + zipWithM_ setIndex space_is $ + map (ConvOpExp (SExt Int64 Int32)) . unflattenIndex space_sizes_64 $ Imp.var j int64 + + -- We execute the bucket function once and update each histogram serially. + -- We apply the bucket function if j=offset+ltid is less than + -- num_elements. This also involves writing to the mapout + -- arrays. + let input_in_bounds = Imp.var j int32 .<. total_w_64 + + sWhen input_in_bounds $ ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do + let (red_res, map_res) = splitFromEnd (length map_pes) $ bodyResult body + + sComment "save map-out results" $ + forM_ (zip map_pes map_res) $ \(pe, se) -> + ImpGen.copyDWIM (patElemName pe) + (map ((`Imp.var` int32) . fst) $ kernelDimensions constants) se [] + + let (buckets, vs) = splitAt (length ops) red_res + perOp = chunks $ map (length . genReduceDest) ops + + sComment "perform atomic updates" $ + forM_ (zip5 ops histograms buckets (perOp vs) subhisto_inds) $ + \(GenReduceOp dest_w _ _ shape lam, + (_, _, do_op), bucket, vs', subhisto_ind) -> do + + let bucket' = ImpGen.compileSubExpOfType int32 bucket + dest_w' = ImpGen.compileSubExpOfType int32 dest_w + bucket_in_bounds = 0 .<=. bucket' .&&. bucket' .<. dest_w' + bucket_is = map (`Imp.var` int32) (init space_is) ++ + [Imp.var subhisto_ind int32, bucket'] + vs_params = takeLast (length vs') $ lambdaParams lam + + sWhen bucket_in_bounds $ do + ImpGen.dLParams vs_params + vectorLoops [] (shapeDims shape) $ \is -> do + forM_ (zip vs_params vs') $ \(p, v) -> + ImpGen.copyDWIM (paramName p) [] v is + do_op (bucket_is ++ is) + + let histogramInfo (num_histos, dests, _) = (num_histos, dests) + return $ map histogramInfo histograms + +vectorLoops :: [Imp.Exp] -> [SubExp] + -> ([Imp.Exp] -> ImpGen.ImpM lore op ()) + -> ImpGen.ImpM lore op () +vectorLoops is [] f = f $ reverse is +vectorLoops is (d:ds) f = do + i <- newVName "i" + d' <- ImpGen.compileSubExp d + ImpGen.sFor i Int32 d' $ vectorLoops (Imp.var i int32:is) ds f + +compileSegGenRed :: Pattern ExplicitMemory + -> KernelSpace + -> [GenReduceOp InKernel] + -> Body InKernel + -> CallKernelGen () +compileSegGenRed (Pattern _ pes) genred_space ops body = do + let num_red_res = length ops + sum (map (length . genReduceNeutral) ops) + (all_red_pes, map_pes) = splitAt num_red_res pes + + infos <- genRedKernel map_pes genred_space ops body + let pes_per_op = chunks (map (length . genReduceDest) ops) all_red_pes + + forM_ (zip3 infos pes_per_op ops) $ \((num_histos, subhistos), red_pes, op) -> do + let unitHistoCase = + -- This is OK because the memory blocks are at least as + -- large as the ones we are supposed to use for the result. + forM_ (zip red_pes subhistos) $ \(pe, subhisto) -> do + pe_mem <- ImpGen.memLocationName . ImpGen.entryArrayLocation <$> + ImpGen.lookupArray (patElemName pe) + subhisto_mem <- ImpGen.memLocationName . ImpGen.entryArrayLocation <$> + ImpGen.lookupArray subhisto + ImpGen.emit $ Imp.SetMem pe_mem subhisto_mem $ Space "device" + + sIf (Imp.var num_histos int32 .==. 1) unitHistoCase $ do + -- For the segmented reduction, we keep the segment dimensions + -- unchanged. To this, we add two dimensions: one over the number + -- of buckets, and one over the number of subhistograms. This + -- inner dimension is the one that is collapsed in the reduction. + let segment_dims = init $ spaceDimensions genred_space + num_buckets = genReduceWidth op + + bucket_id <- newVName "bucket_id" + subhistogram_id <- newVName "subhistogram_id" + vector_ids <- mapM (const $ newVName "vector_id") $ + shapeDims $ genReduceShape op + gtid <- newVName $ baseString $ spaceGlobalId genred_space + let lam = genReduceOp op + segred_space = + genred_space + { spaceStructure = + FlatThreadSpace $ + segment_dims ++ + [(bucket_id, num_buckets)] ++ + zip vector_ids (shapeDims $ genReduceShape op) ++ + [(subhistogram_id, Var num_histos)] + , spaceGlobalId = gtid + } + + compileSegRed' (Pattern [] red_pes) segred_space + Commutative lam (genReduceNeutral op) $ \red_dests _ -> + forM_ (zip red_dests subhistos) $ \((d, is), subhisto) -> + ImpGen.copyDWIM d is (Var subhisto) $ map (`Imp.var` int32) $ + map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegRed.hs @@ -43,12 +43,12 @@ -- have any element to read, which becomes highly inefficient. module Futhark.CodeGen.ImpGen.Kernels.SegRed ( compileSegRed + , compileSegRed' ) where import Control.Monad.Except import Data.Maybe -import qualified Data.Set as S import Data.List import Prelude hiding (quot, rem) @@ -83,14 +83,29 @@ i <- newVName "i" sFor i Int32 iterations $ m $ group_id + Imp.var i int32 * kernelNumGroups constants --- Compile 'SegRed' instance to host-level code with calls to various --- kernels. +-- | Compile 'SegRed' instance to host-level code with calls to +-- various kernels. compileSegRed :: Pattern ExplicitMemory -> KernelSpace -> Commutativity -> Lambda InKernel -> [SubExp] -> Body InKernel -> CallKernelGen () -compileSegRed pat space comm red_op nes body +compileSegRed pat space comm red_op nes body = + compileSegRed' pat space comm red_op nes $ \red_dests map_dests -> + ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do + let (red_res, map_res) = splitAt (length nes) $ bodyResult body + sComment "save results to be reduced" $ + forM_ (zip red_dests red_res) $ \((d,is), se) -> ImpGen.copyDWIM d is se [] + sComment "save map-out results" $ + forM_ (zip map_dests map_res) $ \((d,is), se) -> ImpGen.copyDWIM d is se [] + +-- | Like 'compileSegRed', but where the body is a monadic action. +compileSegRed' :: Pattern ExplicitMemory + -> KernelSpace + -> Commutativity -> Lambda InKernel -> [SubExp] + -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ()) + -> CallKernelGen () +compileSegRed' pat space comm red_op nes body | [(_, Constant (IntValue (Int32Value 1))), _] <- spaceDimensions space = nonsegmentedReduction pat space comm red_op nes body | otherwise = do @@ -98,52 +113,14 @@ ImpGen.compileSubExp $ last $ map snd $ spaceDimensions space group_size <- ImpGen.compileSubExp $ spaceGroupSize space let use_small_segments = segment_size * 2 .<. group_size - sIf (segment_size .==. 1) - (unitSegmentsReduction pat space nes body) $ - sIf use_small_segments + sIf use_small_segments (smallSegmentsReduction pat space red_op nes body) (largeSegmentsReduction pat space comm red_op nes body) --- Handle degenerate case where segments are of size 1, meaning --- that it is really just a 'map' in disguise. -unitSegmentsReduction :: Pattern ExplicitMemory - -> KernelSpace - -> [SubExp] - -> Body InKernel - -> CallKernelGen () -unitSegmentsReduction (Pattern _ segred_pes) space nes body = do - (constants, init_constants) <- kernelInitialisationSetSpace space $ return () - - let (gtids, dims) = unzip $ spaceDimensions space - (redout_pes, mapout_pes) = splitAt (length nes) segred_pes - - dims' <- mapM ImpGen.compileSubExp dims - - let num_segments = product $ init dims' - required_groups = num_segments `quotRoundingUp` kernelGroupSize constants - - ImpGen.emit $ Imp.DebugPrint "num_segments" int32 num_segments - ImpGen.emit $ Imp.DebugPrint "required_groups" int32 required_groups - - sKernel constants "segred_mapseg" $ do - init_constants - virtualiseGroups constants required_groups $ \group_id -> do - setSpaceIndices (group_id * kernelGroupSize constants + kernelLocalThreadId constants) space - ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ - sWhen (kernelThreadActive constants) $ do - let (redout_ses, mapout_ses) = splitAt (length nes) $ bodyResult body - forM_ (zip redout_pes redout_ses) $ \(pe, se) -> - ImpGen.copyDWIM (patElemName pe) - (map (`Imp.var` int32) (init gtids)) se [] - - forM_ (zip mapout_pes mapout_ses) $ \(pe, se) -> - ImpGen.copyDWIM (patElemName pe) - (map (`Imp.var` int32) gtids) se [] - nonsegmentedReduction :: Pattern ExplicitMemory -> KernelSpace -> Commutativity -> Lambda InKernel -> [SubExp] - -> Body InKernel + -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ()) -> CallKernelGen () nonsegmentedReduction segred_pat space comm red_op nes body = do (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () @@ -166,7 +143,7 @@ counter <- ImpGen.sStaticArray "counter" (Space "device") int32 $ - replicate 1 $ IntValue $ Int32Value 0 + Imp.ArrayValues $ replicate 1 $ IntValue $ Int32Value 0 group_res_arrs <- forM (lambdaReturnType red_op) $ \t -> do let pt = elemType t @@ -198,14 +175,10 @@ (kernelNumGroups constants) group_result_params red_acc_params red_op_renamed nes 1 counter sync_arr group_res_arrs red_arrs -hasMemoryAccesses :: Body InKernel -> ImpGen.ImpM InKernel Imp.KernelOp Bool -hasMemoryAccesses body = or <$> mapM isArray (S.toList $ freeInBody body) - where isArray = fmap (not . primType) . lookupType - smallSegmentsReduction :: Pattern ExplicitMemory -> KernelSpace -> Lambda InKernel -> [SubExp] - -> Body InKernel + -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ()) -> CallKernelGen () smallSegmentsReduction (Pattern _ segred_pes) space red_op nes body = do (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () @@ -215,8 +188,12 @@ dims' <- mapM ImpGen.compileSubExp dims let segment_size = last dims' + -- Careful to avoid division by zero now. + segment_size_nonzero_v <- dPrimV "segment_size_nonzero" $ + BinOpExp (SMax Int32) 1 segment_size + let segment_size_nonzero = Imp.var segment_size_nonzero_v int32 num_segments = product $ init dims' - segments_per_group = kernelGroupSize constants `quot` segment_size + segments_per_group = kernelGroupSize constants `quot` segment_size_nonzero required_groups = num_segments `quotRoundingUp` segments_per_group let red_op_params = lambdaParams red_op @@ -248,44 +225,37 @@ -- the segment ID, and are computed from the group id. The inner -- is computed from the local thread id, and may be out-of-bounds. let ltid = kernelLocalThreadId constants - segment_index = (ltid `quot` segment_size) + (group_id' * segments_per_group) + segment_index = (ltid `quot` segment_size_nonzero) + (group_id' * segments_per_group) index_within_segment = ltid `rem` segment_size zipWithM_ (<--) (init gtids) $ unflattenIndex (init dims') segment_index last gtids <-- index_within_segment let toLocalMemory ses = - forM_ (zip red_arrs ses) $ \(arr, se) -> do - se_t <- subExpType se - when (primType se_t) $ - ImpGen.copyDWIM arr [ltid] se [] + forM_ (zip red_arrs ses) $ \(arr, se) -> + ImpGen.copyDWIM arr [ltid] se [] in_bounds = - ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do - let (red_res, map_res) = splitAt (length nes) $ bodyResult body - - sComment "save results to be reduced" $ - toLocalMemory red_res - - sComment "save map-out results" $ - forM_ (zip (drop (length nes) segred_pes) map_res) $ \(pe, se) -> - ImpGen.copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids) se [] + body (zip red_arrs $ repeat [ltid]) + (zip (map patElemName $ drop (length nes) segred_pes) $ + repeat $ map (`Imp.var` int32) gtids) sComment "apply map function if in bounds" $ - sIf (isActive (init $ zip gtids dims) .&&. + sIf (segment_size .>. 0 .&&. + isActive (init $ zip gtids dims) .&&. ltid .<. segment_size * segments_per_group) in_bounds (toLocalMemory nes) sOp Imp.LocalBarrier index_i <- newVName "index_i" index_j <- newVName "index_j" - let crossesSegment from to = - (to-from) .>. (to `rem` segment_size) + let crossesSegment from to = (to-from) .>. (to `rem` segment_size) red_op' = red_op { lambdaParams = Param index_i (MemPrim int32) : Param index_j (MemPrim int32) : lambdaParams red_op } - sComment "perform segmented scan to imitate reduction" $ + sWhen (segment_size .>. 0) $ + sComment "perform segmented scan to imitate reduction" $ groupScan constants (Just crossesSegment) (segment_size*segments_per_group) red_op' red_arrs sOp Imp.LocalBarrier @@ -298,12 +268,16 @@ let flat_segment_index = group_id' * segments_per_group + ltid gtids' = unflattenIndex (init dims') flat_segment_index ImpGen.copyDWIM (patElemName pe) gtids' - (Var arr) [(ltid+1) * segment_size - 1] + (Var arr) [(ltid+1) * segment_size_nonzero - 1] + -- Finally another barrier, because we will be writing to the + -- local memory array first thing in the next iteration. + sOp Imp.LocalBarrier + largeSegmentsReduction :: Pattern ExplicitMemory -> KernelSpace -> Commutativity -> Lambda InKernel -> [SubExp] - -> Body InKernel + -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ()) -> CallKernelGen () largeSegmentsReduction segred_pat space comm red_op nes body = do (base_constants, init_constants) <- kernelInitialisationSetSpace space $ return () @@ -368,7 +342,7 @@ let num_counters = 1024 counter <- ImpGen.sStaticArray "counter" (Space "device") int32 $ - replicate num_counters $ IntValue $ Int32Value 0 + Imp.ArrayZeros num_counters sync_arr <- ImpGen.sAllocArray "sync_arr" Bool (Shape [intConst Int32 1]) $ Space "local" @@ -408,11 +382,12 @@ sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment +-- Careful to avoid division by zero here. groupsPerSegmentAndElementsPerThread :: Imp.Exp -> Imp.Exp -> Imp.Exp -> Imp.Exp -> (Imp.Exp, Imp.Count Imp.Elements) groupsPerSegmentAndElementsPerThread segment_size num_segments num_groups_hint group_size = let groups_per_segment = - num_groups_hint `quotRoundingUp` num_segments + num_groups_hint `quotRoundingUp` BinOpExp (SMax Int32) 1 num_segments elements_per_thread = segment_size `quotRoundingUp` (group_size * groups_per_segment) in (groups_per_segment, Imp.elements elements_per_thread) @@ -427,7 +402,7 @@ -> LambdaT InKernel -> [SubExp] -> [VName] - -> Body InKernel + -> ([(VName, [Imp.Exp])] -> [(VName, [Imp.Exp])] -> InKernelGen ()) -> InKernelGen ([LParam InKernel], Lambda InKernel) reductionStageOne constants (Pattern _ segred_pes) num_elements global_tid elems_per_thread threads_per_segment comm red_op nes red_arrs body = do @@ -436,13 +411,11 @@ (gtids, _dims) = unzip $ kernelDimensions constants gtid = last gtids local_tid = kernelLocalThreadId constants - index_in_segment = global_tid `quot` kernelGroupSize constants -- Figure out how many elements this thread should process. chunk_size <- dPrim "chunk_size" int32 let ordering = case comm of Commutative -> SplitStrided $ Var threads_per_segment Noncommutative -> SplitContiguous - accesses_memory <- hasMemoryAccesses body computeThreadChunkSize ordering global_tid elems_per_thread num_elements chunk_size ImpGen.dScope Nothing $ scopeOfLParams $ lambdaParams red_op @@ -462,6 +435,8 @@ groupReduce constants (kernelGroupSize constants) red_op_renamed red_arrs + sOp Imp.LocalBarrier + i <- newVName "i" -- If this is a non-commutative reduction, each thread must run the -- loop the same number of iterations, because we will be performing @@ -478,31 +453,24 @@ Commutative -> global_tid + Imp.var threads_per_segment int32 * Imp.var i int32 - Noncommutative | accesses_memory -> - local_tid + - (index_in_segment * Imp.innerExp elems_per_thread + Imp.var i int32) * - kernelGroupSize constants Noncommutative -> - Imp.var i int32 + - global_tid * Imp.innerExp elems_per_thread - - check_bounds $ sComment "apply map function" $ - ImpGen.compileStms mempty (stmsToList $ bodyStms body) $ do - let (red_res, map_res) = splitAt (length nes) $ bodyResult body + let index_in_segment = global_tid `quot` kernelGroupSize constants + in local_tid + + (index_in_segment * Imp.innerExp elems_per_thread + Imp.var i int32) * + kernelGroupSize constants - sComment "save results to be reduced" $ - forM_ (zip red_next_params red_res) $ \(p, se) -> - ImpGen.copyDWIM (paramName p) [] se [] + let red_dests = zip (map paramName red_next_params) $ repeat [] + map_dests = zip (map patElemName $ drop (length nes) segred_pes) $ + repeat $ map (`Imp.var` int32) gtids - sComment "save map-out results" $ - forM_ (zip (drop (length nes) segred_pes) map_res) $ \(pe, se) -> - ImpGen.copyDWIM (patElemName pe) (map (`Imp.var` int32) gtids) se [] + check_bounds $ sComment "apply map function" $ do + body red_dests map_dests - sComment "apply reduction operator" $ - ImpGen.compileBody' red_acc_params $ lambdaBody red_op + sComment "apply reduction operator" $ + ImpGen.compileBody' red_acc_params $ lambdaBody red_op case comm of - Noncommutative | accesses_memory -> do + Noncommutative -> do doTheReduction sComment "first thread takes carry-out; others neutral element" $ do let carry_out = @@ -512,17 +480,12 @@ forM_ (zip red_acc_params nes) $ \(p, ne) -> ImpGen.copyDWIM (paramName p) [] ne [] sIf (local_tid .==. 0) carry_out reset_to_neutral - _ -> - return () - - group_result_params <- case comm of - Noncommutative | accesses_memory -> - return red_acc_params - - _ -> do - doTheReduction + _ -> return () - return $ lambdaParams red_op_renamed + group_result_params <- + case comm of Noncommutative -> return red_acc_params + _ -> do doTheReduction + return $ lambdaParams red_op_renamed return (group_result_params, red_op_renamed) @@ -590,6 +553,8 @@ load_group_result load_neutral_element when (primType $ paramType p) $ ImpGen.copyDWIM arr [local_tid] (Var $ paramName p) [] + + sOp Imp.LocalBarrier sComment "reduce the per-group results" $ do groupReduce constants group_size red_op_renamed red_arrs diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs @@ -29,7 +29,7 @@ import qualified Futhark.CodeGen.ImpCode.OpenCL as ImpOpenCL import Futhark.MonadFreshNames import Futhark.Util (zEncodeString) -import Futhark.Util.Pretty (pretty, prettyOneLine) +import Futhark.Util.Pretty (pretty) kernelsToCUDA, kernelsToOpenCL :: ImpKernels.Program -> Either InternalError ImpOpenCL.Program @@ -65,17 +65,15 @@ type UsedFunctions = [(String,C.Func)] -- The ordering is important! -data OpenClRequirements = - OpenClRequirements { kernelUsedTypes :: S.Set PrimType - , _kernelConstants :: [(VName, KernelConstExp)] - } +newtype OpenClRequirements = + OpenClRequirements { kernelUsedTypes :: S.Set PrimType } instance Semigroup OpenClRequirements where - OpenClRequirements ts1 consts1 <> OpenClRequirements ts2 consts2 = - OpenClRequirements (ts1 <> ts2) (consts1 <> consts2) + OpenClRequirements ts1 <> OpenClRequirements ts2 = + OpenClRequirements (ts1 <> ts2) instance Monoid OpenClRequirements where - mempty = OpenClRequirements mempty mempty + mempty = OpenClRequirements mempty data ToOpenCL = ToOpenCL { clExtraFuns :: M.Map Name ImpOpenCL.Function , clKernels :: M.Map KernelName C.Func @@ -140,16 +138,17 @@ params = perm_params ++ catMaybes local_memory_params ++ use_params + const_defs = mapMaybe constDef $ kernelUses kernel + tell mempty { clExtraFuns = mempty , clKernels = M.singleton name [C.cfun|__kernel void $id:name ($params:params) { + $items:const_defs $items:block_dim_init $items:local_memory_init $items:kernel_body }|] - , clRequirements = OpenClRequirements - (typesInKernel kernel) - (mapMaybe useAsConst $ kernelUses kernel) + , clRequirements = OpenClRequirements (typesInKernel kernel) } return $ LaunchKernel name (kernelArgs kernel) num_groups group_size @@ -186,9 +185,11 @@ useAsParam ConstUse{} = Nothing -useAsConst :: KernelUse -> Maybe (VName, KernelConstExp) -useAsConst (ConstUse v e) = Just (v,e) -useAsConst _ = Nothing +constDef :: KernelUse -> Maybe C.BlockItem +constDef (ConstUse v e) = Just [C.citem|const $ty:t $id:v = $exp:e';|] + where t = GenericC.primTypeToCType $ primExpType e + e' = compilePrimExp e +constDef _ = Nothing openClCode :: [C.Func] -> String openClCode kernels = @@ -198,9 +199,13 @@ kernel_func <- kernels ] genOpenClPrelude :: OpenClRequirements -> [C.Definition] -genOpenClPrelude (OpenClRequirements ts consts) = +genOpenClPrelude (OpenClRequirements ts) = -- 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:("#ifdef cl_clang_storage_class_specifiers")|] + , [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_clang_storage_class_specifiers : enable")|] + , [C.cedecl|$esc:("#endif")|] + , [C.cedecl|$esc:("#pragma OPENCL EXTENSION cl_khr_byte_addressable_store : enable")|]] + ++ [[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. @@ -222,10 +227,24 @@ typedef ulong uint64_t; $esc:("#define ALIGNED_LOCAL_MEMORY(m,size) __local unsigned char m[size] __attribute__ ((align))") + +// NVIDIAs OpenCL does not create device-wide memory fences (see #734), so we +// use inline assembly if we detect we are on an NVIDIA GPU. +$esc:("#ifdef cl_nv_pragma_unroll") +static inline void mem_fence_global() { + asm("membar.gl;"); +} +$esc:("#else") +static inline void mem_fence_global() { + mem_fence(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE); +} +$esc:("#endif") +static inline void mem_fence_local() { + mem_fence(CLK_LOCAL_MEM_FENCE); +} |] ++ cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ - (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else []) ++ - [ [C.cedecl|$esc:def|] | def <- map constToDefine consts ] + (if uses_float64 then cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps else []) where uses_float64 = FloatType Float64 `S.member` ts @@ -254,11 +273,10 @@ }|] | t <- types] genCUDAPrelude :: OpenClRequirements -> [C.Definition] -genCUDAPrelude (OpenClRequirements _ consts) = - cudafy ++ cudaAtomicOps ++ defs ++ ops +genCUDAPrelude (OpenClRequirements _) = + cudafy ++ cudaAtomicOps ++ ops where ops = cIntOps ++ cFloat32Ops ++ cFloat32Funs ++ cFloat64Ops ++ cFloat64Funs ++ cFloatConvOps - defs = [ [C.cedecl|$esc:def|] | def <- map constToDefine consts ] cudafy = [CUDAC.cunit| typedef char int8_t; typedef short int16_t; @@ -349,25 +367,17 @@ { __syncthreads(); } -static inline void mem_fence(int x) -{ - if (x == CLK_LOCAL_MEM_FENCE) { - __threadfence_block(); - } else { - __threadfence(); - } +static inline void mem_fence_local() { + __threadfence_block(); } +static inline void mem_fence_global() { + __threadfence(); +} $esc:("#define NAN (0.0/0.0)") $esc:("#define INFINITY (1.0/0.0)") extern volatile __shared__ char shared_mem[]; |] -constToDefine :: (VName, KernelConstExp) -> String -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) = @@ -413,7 +423,7 @@ kernelOps GlobalBarrier = GenericC.stm [C.cstm|barrier(CLK_GLOBAL_MEM_FENCE);|] kernelOps MemFence = - GenericC.stm [C.cstm|mem_fence(CLK_GLOBAL_MEM_FENCE);|] + GenericC.stm [C.cstm|mem_fence_global();|] kernelOps (Atomic aop) = atomicOps aop atomicOps (AtomicAdd old arr ind val) = do diff --git a/src/Futhark/CodeGen/OpenCL/Kernels.hs b/src/Futhark/CodeGen/OpenCL/Kernels.hs --- a/src/Futhark/CodeGen/OpenCL/Kernels.hs +++ b/src/Futhark/CodeGen/OpenCL/Kernels.hs @@ -28,7 +28,7 @@ | HeuristicDeviceInfo String -- | A size that can be assigned a default. -data WhichSize = LockstepWidth | NumGroups | GroupSize | TileSize +data WhichSize = LockstepWidth | NumGroups | GroupSize | TileSize | Threshold -- | A heuristic for setting the default value for something. data SizeHeuristic = @@ -47,9 +47,11 @@ , SizeHeuristic "" DeviceGPU NumGroups $ HeuristicConst 256 , SizeHeuristic "" DeviceGPU GroupSize $ HeuristicConst 256 , SizeHeuristic "" DeviceGPU TileSize $ HeuristicConst 32 + , SizeHeuristic "" DeviceGPU Threshold $ HeuristicConst $ 32*1024 , SizeHeuristic "" DeviceCPU LockstepWidth $ HeuristicConst 1 , SizeHeuristic "" DeviceCPU NumGroups $ HeuristicDeviceInfo "MAX_COMPUTE_UNITS" , SizeHeuristic "" DeviceCPU GroupSize $ HeuristicConst 32 , SizeHeuristic "" DeviceCPU TileSize $ HeuristicConst 4 + , SizeHeuristic "" DeviceCPU Threshold $ HeuristicDeviceInfo "MAX_COMPUTE_UNITS" ] diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -30,6 +30,7 @@ , eSliceArray , eSplitArray + , eOutOfBounds , eWriteArray , asIntZ, asIntS @@ -285,16 +286,13 @@ 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 +-- | Are these indexes out-of-bounds for the array? +eOutOfBounds :: MonadBinder m => + VName -> [m (Exp (Lore m))] -> m (Exp (Lore m)) +eOutOfBounds arr is = 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)) @@ -302,11 +300,19 @@ BasicOp $ CmpOp (CmpSle Int32) w i letSubExp "outside_bounds_dim" $ BasicOp $ BinOp LogOr less_than_zero greater_than_size + foldBinOp LogOr (constant False) =<< zipWithM checkDim ws is' - outside_bounds <- - letSubExp "outside_bounds" =<< - foldBinOp LogOr (constant False) =<< - zipWithM checkDim ws is' +-- | 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 + is' <- mapM (letSubExp "write_i") =<< sequence is + v' <- letSubExp "write_v" =<< v + + outside_bounds <- letSubExp "outside_bounds" =<< eOutOfBounds arr is outside_bounds_branch <- insertStmsM $ resultBodyM [Var arr] diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs --- a/src/Futhark/Doc/Generator.hs +++ b/src/Futhark/Doc/Generator.hs @@ -715,10 +715,9 @@ describeDec _ ImportDec{} = Nothing valBindWhat :: ValBind -> IndexWhat -valBindWhat vb = - if null (valBindParams vb) && orderZero (unInfo (valBindRetType vb)) - then IndexValue - else IndexFunction +valBindWhat vb | null (valBindParams vb), + orderZero (unInfo (valBindRetType vb)) = IndexValue + | otherwise = IndexFunction describeSpecs :: [Spec] -> DocM Html describeSpecs specs = diff --git a/src/Futhark/Internalise.hs b/src/Futhark/Internalise.hs --- a/src/Futhark/Internalise.hs +++ b/src/Futhark/Internalise.hs @@ -459,11 +459,11 @@ mergeinit_ts' <- mapM subExpType mergeinit' - let ctxinit = argShapes - (map I.paramName shapepat) - (map I.paramType mergepat') - mergeinit_ts' - ctxmerge = zip shapepat ctxinit + ctxinit <- argShapes + (map I.paramName shapepat) + (map I.paramType mergepat') + mergeinit_ts' + let ctxmerge = zip shapepat ctxinit valmerge = zip mergepat' mergeinit' merge = ctxmerge ++ valmerge dropCond = case form of E.While{} -> drop 1 @@ -487,10 +487,10 @@ 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 + shapeargs <- argShapes + (map I.paramName shapepat) + (map I.paramType mergepat') + sets return (resultBody $ shapeargs ++ ses, (form', shapepat, @@ -536,10 +536,10 @@ -- 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 + 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) -> @@ -555,10 +555,10 @@ 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 + shapeargs <- argShapes + (map I.paramName shapepat) + (map I.paramType mergepat') + sets -- Careful not to clobber anything. loop_end_cond_body <- renameBody <=< insertStmsM $ do @@ -618,9 +618,6 @@ 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 @@ -636,33 +633,6 @@ 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' @@ -1420,48 +1390,9 @@ ressize <- foldM sumdims outer_size =<< mapM (fmap (arraysSize 0) . mapM lookupType) [ys] - let conc xarr yarr = do - -- All dimensions except the outermost must match. An - -- empty array matches anything. - xt <- lookupType xarr - yt <- lookupType yarr - let matches n m = - letSubExp "match" $ - I.BasicOp $ I.CmpOp (I.CmpEq I.int32) n m - - emptyRow arr_t = - letSubExp "empty_row" =<< - foldBinOp I.LogOr (constant False) =<< - mapM (matches (intConst Int32 0)) (arrayDims $ rowType arr_t) - - all_match <- letSubExp "all_match" =<< - foldBinOp I.LogAnd (constant True) =<< - zipWithM matches - (arrayDims (rowType xt)) (arrayDims (rowType yt)) - xarr_empty <- emptyRow xt - yarr_empty <- emptyRow yt - either_empty <- letSubExp "either_empty" $ - I.BasicOp $ I.BinOp I.LogOr xarr_empty yarr_empty - matchcs <- assertingOne $ letExp "concat_ok" =<< - eAssert (pure $ I.BasicOp $ I.BinOp I.LogOr either_empty all_match) - "row sizes do not match when concatenating" loc - - let updims (j, xd, yd) - | j == 0 = - return (xd, yd) - | otherwise = do - d <- letSubExp "dim" $ I.BasicOp $ I.BinOp (SMax Int32) xd yd - return (d, d) - - (xdims, ydims) <- unzip <$> - mapM updims (zip3 [(0::Int)..] (I.arrayDims xt) (I.arrayDims yt)) - - xarr' <- certifying matchcs $ letExp "concat_x_reshaped" $ - shapeCoerce xdims xarr - yarr' <- certifying matchcs $ letExp "concat_y_reshaped" $ - shapeCoerce ydims yarr - return $ I.BasicOp $ I.Concat 0 xarr' [yarr'] ressize - letSubExps desc =<< zipWithM conc xs ys + let conc xarr yarr = + I.BasicOp $ I.Concat 0 xarr [yarr] ressize + letSubExps desc $ zipWith conc xs ys handle [TupLit [offset, e] _] "rotate" = Just $ \desc -> do offset' <- internaliseExp1 "rotation_offset" offset @@ -1635,8 +1566,8 @@ (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 ++ + shapeargs <- argShapes shapes value_paramts argts + let 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 ++ diff --git a/src/Futhark/Internalise/AccurateSizes.hs b/src/Futhark/Internalise/AccurateSizes.hs --- a/src/Futhark/Internalise/AccurateSizes.hs +++ b/src/Futhark/Internalise/AccurateSizes.hs @@ -15,6 +15,7 @@ import Control.Monad import Data.Loc import qualified Data.Map.Strict as M +import qualified Data.Set as S import Futhark.Construct import Futhark.Representation.AST @@ -26,7 +27,7 @@ runBodyBinder $ do ses <- bodyBind body sets <- mapM subExpType ses - return $ resultBody $ argShapes shapenames ts sets + resultBody <$> argShapes shapenames ts sets annotateArrayShape :: ArrayShape shape => TypeBase shape u -> [Int] -> TypeBase Shape u @@ -34,13 +35,25 @@ t `setArrayShape` Shape (take (arrayRank t) $ map (intConst Int32 . toInteger) $ newshape ++ repeat 0) -argShapes :: [VName] -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> [SubExp] +-- Some trickery is needed here to predict sensible values for +-- dimensions that are used exclusively as the inner dimension of an +-- array. The issue is that the dimension may be inside an empty +-- array. In this case, the dimension inside the empty array should +-- not count, as it will be zero. The solution we use is to take the +-- maximum of such sizes; this will effectively disregard the zeroes. +argShapes :: MonadBinder m => + [VName] -> [TypeBase Shape u0] -> [TypeBase Shape u1] -> m [SubExp] argShapes shapes valts valargts = - map addShape shapes + mapM addShape shapes where mapping = shapeMapping valts valargts - addShape name - | Just se <- M.lookup name mapping = se - | otherwise = intConst Int32 0 + outer_dims = map (arraySize 0) valts + addShape name = + case M.lookup name mapping of + Just s | x:xs <- S.toList s -> + if Var name `elem` outer_dims + then return x + else letSubExp "size" =<< foldBinOp (SMax Int32) x xs + _ -> return $ intConst Int32 0 ensureResultShape :: MonadBinder m => (m Certificates -> m Certificates) diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs --- a/src/Futhark/Internalise/Defunctionalise.hs +++ b/src/Futhark/Internalise/Defunctionalise.hs @@ -363,15 +363,6 @@ 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) @@ -497,67 +488,42 @@ 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 of 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 e0_t $ typeOf e0' + -- Lift lambda to top-level function definition. We put in + -- a lot of effort to try to infer the uniqueness attributes + -- of the lifted function, but this is ultimately all a sham + -- and a hack. There is some piece we're missing. + let params = [closure_pat, pat'] + params_for_rettype = params ++ svParams sv1 ++ svParams sv2 + svParams (LambdaSV _ sv_pat _ _ _) = [sv_pat] + svParams _ = [] + rettype = buildRetType closure_env params_for_rettype e0_t $ 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' + -- 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) + 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 + apply_e = Apply e1' e2' d (Info $ foldFunType argtypes' rettype + `setAliases` aliases ret) loc in return (apply_e, sv) -- Propagate the 'IntrinsicsSV' until we reach the outermost application, @@ -681,7 +647,7 @@ comb (Record fs_annot) (Record fs_got) = Record $ M.intersectionWith comb fs_annot fs_got comb Arrow{} t = vacuousShapeAnnotations $ descend t - comb got _ = fromStruct got + comb got et = descend $ fromStruct got `setUniqueness` uniqueness et `setAliases` aliases et descend t@Array{} | any (problematic . aliasVar) (aliases t) = t `setUniqueness` Nonunique @@ -865,8 +831,6 @@ 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 VConstr0{} -> mempty diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs --- a/src/Futhark/Internalise/Monomorphise.hs +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE FlexibleContexts #-} -- | This monomorphization module converts a well-typed, polymorphic, -- module-free Futhark program into an equivalent monomorphic program. -- @@ -29,15 +30,19 @@ import Control.Monad.RWS import Control.Monad.State +import Control.Monad.Writer +import Data.Bitraversable +import Data.Bifunctor import Data.Loc import qualified Data.Map.Strict as M +import qualified Data.Set as S 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.Semantic (TypeBinding(..)) import Language.Futhark.TypeChecker.Types -- | The monomorphization monad reads 'PolyBinding's and writes 'ValBinding's. @@ -135,11 +140,24 @@ (Nothing, Nothing) -> return fname -- A polymorphic function. (Nothing, Just funbind) -> do - (fname', funbind') <- monomorphizeBinding funbind t + (fname', funbind') <- monomorphizeBinding False funbind t tell $ Seq.singleton (fname, funbind') addLifted fname t fname' return fname' +-- | This carries out record replacements in the alias information of a type. +transformType :: TypeBase dim Aliasing -> MonoM (TypeBase dim Aliasing) +transformType t = do + rrs <- asks envRecordReplacements + let replace (AliasBound v) | Just d <- M.lookup v rrs = + S.fromList $ map (AliasBound . fst) $ M.elems d + replace x = S.singleton x + -- As an attempt at an optimisation, only transform the aliases if + -- they refer to a variable we have record-replaced. + return $ if any ((`M.member` rrs) . aliasVar) $ aliases t + then bimap id (mconcat . map replace . S.toList) t + else t + -- | Monomorphization of expressions. transformExp :: Exp -> MonoM Exp transformExp e@Literal{} = return e @@ -159,9 +177,10 @@ RecordLit <$> mapM transformField fs <*> pure loc where transformField (RecordFieldExplicit name e loc') = RecordFieldExplicit name <$> transformExp e <*> pure loc' - transformField (RecordFieldImplicit v t _) = + transformField (RecordFieldImplicit v t _) = do + t' <- traverse transformType t transformField $ RecordFieldExplicit (baseName v) - (Var (qualName v) (vacuousShapeAnnotations <$> t) loc) loc + (Var (qualName v) (vacuousShapeAnnotations <$> t') loc) loc transformExp (ArrayLit es tp loc) = ArrayLit <$> mapM transformExp es <*> pure tp <*> pure loc @@ -176,13 +195,15 @@ 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 + let toField (f, (f_v, f_t)) = do + f_t' <- transformType f_t + let f_v' = Var (qualName f_v) (Info $ vacuousShapeAnnotations f_t') loc + return $ RecordFieldExplicit f f_v' loc + RecordLit <$> mapM toField (M.toList fs) <*> pure loc Nothing -> do fname' <- transformFName fname (toStructural t) - return $ Var (QualName qs fname') (Info t) loc + t' <- transformType t + return $ Var (QualName qs fname') (Info t') loc transformExp (Ascript e tp loc) = Ascript <$> transformExp e <*> pure tp <*> pure loc @@ -213,7 +234,8 @@ e1' <- transformExp e1 e2' <- transformExp e2 e3' <- transformExp e3 - return $ If e1' e2' e3' tp loc + tp' <- traverse transformType tp + 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 @@ -254,7 +276,8 @@ _ -> do e1' <- transformExp e1 e2' <- transformExp e2 - return $ Apply e1' e2' d tp loc + tp' <- traverse transformType tp + return $ Apply e1' e2' d tp' loc where intrinsic s (QualName _ v) = baseTag v <= maxIntrinsicTag && baseName v == nameFromString s @@ -364,14 +387,6 @@ <*> 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 @@ -380,7 +395,7 @@ transformExp e@VConstr0{} = return e transformExp (Match e cs t loc) = - Match <$> transformExp e <*> mapM transformCase cs <*> pure t <*> pure loc + Match <$> transformExp e <*> mapM transformCase cs <*> traverse transformType t <*> pure loc transformCase :: Case -> MonoM Case transformCase (CasePat p e loc) = do @@ -447,7 +462,7 @@ 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 + (,) <$> newVName (nameToString f) <*> transformType ft return (RecordPattern (zip (map fst fs') (zipWith3 Id fs_ks (map Info fs_ts) $ repeat loc)) loc, @@ -463,7 +478,9 @@ expandRecordPattern (PatternParens pat loc) = do (pat', rr) <- expandRecordPattern pat return (PatternParens pat' loc, rr) -expandRecordPattern (Wildcard t loc) = return (Wildcard t loc, mempty) +expandRecordPattern (Wildcard t loc) = do + t' <- traverse transformType t + return (Wildcard t' loc, mempty) expandRecordPattern (PatternAscription pat td loc) = do (pat', rr) <- expandRecordPattern pat return (PatternAscription pat' td loc, rr) @@ -472,24 +489,27 @@ -- | 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 rr (name, tparams, params, retdecl, rettype, body, loc)) t = +monomorphizeBinding :: Bool -> PolyBinding -> TypeBase () () -> MonoM (VName, ValBind) +monomorphizeBinding entry (PolyBinding rr (name, tparams, params, retdecl, rettype, body, loc)) t = replaceRecordReplacements rr $ do t' <- removeTypeVariablesInType t let bind_t = foldFunType (map (toStructural . patternType) params) $ toStructural rettype - substs = M.map Subst $ typeSubsts bind_t t' - rettype' = applySubst (`M.lookup` substs) rettype - params' = map (substPattern $ applySubst (`M.lookup` substs)) params + (substs, t_shape_params) <- typeSubstsM loc bind_t t' + let substs' = M.map Subst substs + rettype' = substTypesAny (`M.lookup` substs') rettype + substPatternType = + substTypesAny (fmap (fmap fromStruct) . (`M.lookup` substs')) + params' = map (substPattern entry substPatternType) params (params'', rrs) <- unzip <$> mapM expandRecordPattern params' mapM_ noticeDims $ rettype : map patternStructType params'' - body' <- updateExpTypes (`M.lookup` substs) body + body' <- updateExpTypes (`M.lookup` M.map (fmap toStructural) 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'') + return (name', toValBinding t_shape_params name' params'' rettype' body'') where shape_params = filter (not . isTypeParam) tparams @@ -503,45 +523,69 @@ , mapOnPatternType = pure . applySubst substs } - toValBinding name' params'' rettype' body'' = + toValBinding t_shape_params name' params'' rettype' body'' = ValBind { valBindEntryPoint = False , valBindName = name' , valBindRetDecl = retdecl , valBindRetType = Info rettype' - , valBindTypeParams = shape_params + , valBindTypeParams = shape_params ++ t_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 Enum{} Enum{} = mempty -typeSubsts t1 t2 = error $ unlines ["typeSubsts: mismatched types:", pretty t1, pretty t2] +-- Careful not to introduce size parameters for non-positive positions +-- (i.e. function parameters). +typeSubstsM :: MonadFreshNames m => + SrcLoc -> TypeBase () () -> TypeBase () () + -> m (M.Map VName StructType, [TypeParam]) +typeSubstsM loc orig_t1 orig_t2 = + let (t1_pts, t1_rt) = unfoldFunType orig_t1 + (t2_pts, t2_rt) = unfoldFunType orig_t2 + m = do zipWithM_ (sub True) t1_pts t2_pts + sub False t1_rt t2_rt + in runWriterT $ execStateT m mempty + where sub pos (TypeVar _ _ v _) t = addSubst pos v t + sub pos (Record fields1) (Record fields2) = + zipWithM_ (sub pos) + (map snd $ sortFields fields1) (map snd $ sortFields fields2) + sub _ Prim{} Prim{} = return () + sub _ Enum{} Enum{} = return () + sub pos t1@Array{} t2@Array{} + | Just t1' <- peelArray (arrayRank t1) t1, + Just t2' <- peelArray (arrayRank t1) t2 = + sub pos t1' t2' + sub pos (Arrow _ _ t1a t1b) (Arrow _ _ t2a t2b) = do + sub False t1a t2a + sub pos t1b t2b + + sub _ t1 t2 = error $ unlines ["typeSubstsM: mismatched types:", pretty t1, pretty t2] + + addSubst pos (TypeName _ v) t = do + exists <- gets $ M.member v + unless exists $ do + t' <- if pos + then bitraverse onDim pure t + else pure $ vacuousShapeAnnotations t + modify $ M.insert v t' + + onDim () = do d <- lift $ lift $ newVName "d" + tell [TypeParamDim d loc] + return $ NamedDim $ qualName d + -- | 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 +substPattern :: Bool -> (PatternType -> PatternType) -> Pattern -> Pattern +substPattern entry f pat = case pat of + TuplePattern pats loc -> TuplePattern (map (substPattern entry 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 + where substField (n, p) = (n, substPattern entry f p) + PatternParens p loc -> PatternParens (substPattern entry 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 + PatternAscription p td loc | entry -> PatternAscription (substPattern False f p) td loc + | otherwise -> substPattern False f p PatternLit e (Info tp) loc -> PatternLit e (Info $ f tp) loc toPolyBinding :: ValBind -> PolyBinding @@ -549,27 +593,26 @@ PolyBinding mempty (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 +removeTypeVariables :: Bool -> ValBind -> MonoM ValBind +removeTypeVariables entry valbind@(ValBind _ _ _ (Info rettype) _ pats body _ _) = do subs <- asks $ M.map TypeSub . envTypeBindings - let substPatternType = fromStruct . substituteTypes subs . toStruct - mapper = ASTMapper { + let mapper = ASTMapper { mapOnExp = astMap mapper , mapOnName = pure , mapOnQualName = pure , mapOnType = pure . removeShapeAnnotations . substituteTypes subs . vacuousShapeAnnotations - , mapOnCompType = pure . fromStruct . removeShapeAnnotations . + , mapOnCompType = pure . removeShapeAnnotations . substituteTypes subs . - vacuousShapeAnnotations . toStruct + vacuousShapeAnnotations , mapOnStructType = pure . substituteTypes subs - , mapOnPatternType = pure . substPatternType + , mapOnPatternType = pure . substituteTypes subs } body' <- astMap mapper body return valbind { valBindRetType = Info $ substituteTypes subs rettype - , valBindParams = map (substPattern substPatternType) pats + , valBindParams = map (substPattern entry $ substituteTypes subs) pats , valBindBody = body' } @@ -580,12 +623,12 @@ transformValBind :: ValBind -> MonoM Env transformValBind valbind = do - valbind' <- toPolyBinding <$> removeTypeVariables valbind + valbind' <- toPolyBinding <$> removeTypeVariables (valBindEntryPoint valbind) valbind when (valBindEntryPoint valbind) $ do t <- removeTypeVariablesInType $ removeShapeAnnotations $ foldFunType (map patternStructType (valBindParams valbind)) $ unInfo $ valBindRetType valbind - (name, valbind'') <- monomorphizeBinding valbind' t + (name, valbind'') <- monomorphizeBinding True valbind' t tell $ Seq.singleton (name, valbind'' { valBindEntryPoint = True}) addLifted (valBindName valbind) t name return mempty { envPolyBindings = M.singleton (valBindName valbind) valbind' } diff --git a/src/Futhark/Optimise/CSE.hs b/src/Futhark/Optimise/CSE.hs --- a/src/Futhark/Optimise/CSE.hs +++ b/src/Futhark/Optimise/CSE.hs @@ -47,8 +47,9 @@ import qualified Futhark.Representation.ExplicitMemory as ExplicitMemory import Futhark.Transform.Substitute import Futhark.Pass +import Futhark.Util (takeLast) --- | Perform CSE on every functioon in a program. +-- | Perform CSE on every function in a program. performCSE :: (Attributes lore, CanBeAliased (Op lore), CSEInOp (OpWithAliases (Op lore))) => Bool -> Pass lore lore @@ -61,22 +62,27 @@ Bool -> FunDef lore -> FunDef lore cseInFunDef cse_arrays fundec = fundec { funDefBody = - runReader (cseInBody $ funDefBody fundec) $ newCSEState cse_arrays + runReader (cseInBody ds $ funDefBody fundec) $ newCSEState cse_arrays } + where ds = map (diet . declExtTypeOf) $ funDefRetType fundec 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 + [Diet] -> Body lore -> CSEM lore (Body lore) +cseInBody ds (Body bodyattr bnds res) = + cseInStms (res_cons <> consumedInStms bnds res) (stmsToList bnds) $ do CSEState (_, nsubsts) _ <- ask return $ Body bodyattr mempty $ substituteNames nsubsts res + where res_cons = mconcat $ zipWith consumeResult ds $ + takeLast (length ds) res + consumeResult Consume se = freeIn se + consumeResult _ _ = mempty cseInLambda :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => Lambda lore -> CSEM lore (Lambda lore) cseInLambda lam = do - body' <- cseInBody $ lambdaBody lam + body' <- cseInBody (map (const Observe) $ lambdaReturnType lam) $ lambdaBody lam return lam { lambdaBody = body' } cseInStms :: (Attributes lore, Aliased lore, CSEInOp (Op lore)) => @@ -90,12 +96,17 @@ bnd'' <- mapM nestedCSE bnd' return $ Body bodyattr (stmsFromList bnd''<>bnds') es where nestedCSE bnd' = do - e <- mapExpM cse $ stmExp bnd' + let ds = map patElemDiet $ patternValueElements $ stmPattern bnd' + e <- mapExpM (cse ds) $ stmExp bnd' return bnd' { stmExp = e } - cse = identityMapper { mapOnBody = const cseInBody - , mapOnOp = cseInOp - } + cse ds = identityMapper { mapOnBody = const $ cseInBody ds + , mapOnOp = cseInOp + } + + patElemDiet pe | patElemName pe `S.member` consumed = Consume + | otherwise = Observe + cseInStm :: Attributes lore => Names -> Stm lore -> ([Stm lore] -> CSEM lore a) @@ -169,29 +180,34 @@ instance (Attributes lore, Aliased lore, CSEInOp (Op lore)) => CSEInOp (Kernel.Kernel lore) where cseInOp = subCSE . Kernel.mapKernelM - (Kernel.KernelMapper return cseInLambda cseInBody + (Kernel.KernelMapper return cseInLambda + (\b -> cseInBody (map (const Observe) $ bodyResult b) b) 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 [] + Body _ bnds' _ <- cseInBody (map (const Observe) res) $ 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 + subCSE $ KernelExp.Combine cspace ts active <$> + cseInBody (map (const Observe) ts) body cseInOp (KernelExp.GroupReduce w lam input) = - subCSE $ KernelExp.GroupReduce w <$> cseInLambda lam <*> pure 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 + 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 + body' <- cseInBody (map (const Observe) $ KernelExp.groupStreamAccParams lam) $ + KernelExp.groupStreamLambdaBody lam return lam { KernelExp.groupStreamLambdaBody = body' } diff --git a/src/Futhark/Optimise/InliningDeadFun.hs b/src/Futhark/Optimise/InliningDeadFun.hs --- a/src/Futhark/Optimise/InliningDeadFun.hs +++ b/src/Futhark/Optimise/InliningDeadFun.hs @@ -63,13 +63,14 @@ 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))) + inline (Let pat aux (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) + res_stms = map (certify $ stmAuxCerts aux) $ + zipWith reshapeIfNecessary (patternIdents pat) $ + bodyResult $ funDefBody fun in param_stms ++ body_stms ++ res_stms inline stm = [inlineInStm inlcallees stm] diff --git a/src/Futhark/Optimise/MemoryBlockMerging.hs b/src/Futhark/Optimise/MemoryBlockMerging.hs deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging.hs +++ /dev/null @@ -1,28 +0,0 @@ --- | 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/ActualVariables.hs +++ /dev/null @@ -1,358 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/AllExpVars.hs +++ /dev/null @@ -1,96 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/AuxiliaryInfo.hs +++ /dev/null @@ -1,63 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing.hs +++ /dev/null @@ -1,31 +0,0 @@ --- | 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/AllocationMovingUp.hs +++ /dev/null @@ -1,94 +0,0 @@ --- | 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Core.hs +++ /dev/null @@ -1,624 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/Exps.hs +++ /dev/null @@ -1,70 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition2.hs +++ /dev/null @@ -1,110 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition3.hs +++ /dev/null @@ -1,136 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Coalescing/SafetyCondition5.hs +++ /dev/null @@ -1,120 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/CrudeMovingUp.hs +++ /dev/null @@ -1,263 +0,0 @@ --- | 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Existentials.hs +++ /dev/null @@ -1,80 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Liveness/FirstUse.hs +++ /dev/null @@ -1,198 +0,0 @@ -{-# 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 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Liveness/Interference.hs +++ /dev/null @@ -1,520 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Liveness/LastUse.hs +++ /dev/null @@ -1,281 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/MemoryAliases.hs +++ /dev/null @@ -1,162 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/MemoryUpdater.hs +++ /dev/null @@ -1,418 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Miscellaneous.hs +++ /dev/null @@ -1,263 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/PrimExps.hs +++ /dev/null @@ -1,105 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Reuse.hs +++ /dev/null @@ -1,30 +0,0 @@ --- | 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeMovingUp.hs +++ /dev/null @@ -1,32 +0,0 @@ --- | 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizeUses.hs +++ /dev/null @@ -1,127 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/AllocationSizes.hs +++ /dev/null @@ -1,132 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Reuse/Core.hs +++ /dev/null @@ -1,747 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/Types.hs +++ /dev/null @@ -1,89 +0,0 @@ -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 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 Semigroup Log where - Log a <> Log b = Log $ M.unionWith (++) a b - -instance Monoid Log where - mempty = Log M.empty diff --git a/src/Futhark/Optimise/MemoryBlockMerging/VariableAliases.hs b/src/Futhark/Optimise/MemoryBlockMerging/VariableAliases.hs deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/VariableAliases.hs +++ /dev/null @@ -1,82 +0,0 @@ -{-# 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 deleted file mode 100644 --- a/src/Futhark/Optimise/MemoryBlockMerging/VariableMemory.hs +++ /dev/null @@ -1,99 +0,0 @@ -{-# 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/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs --- a/src/Futhark/Pass/ExpandAllocations.hs +++ b/src/Futhark/Pass/ExpandAllocations.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE TypeFamilies, FlexibleContexts #-} +{-# LANGUAGE TypeFamilies, FlexibleContexts, GeneralizedNewtypeDeriving #-} -- | Expand allocations inside of maps when possible. module Futhark.Pass.ExpandAllocations ( expandAllocations ) @@ -31,6 +31,7 @@ import Futhark.Util.IntegralExp import Futhark.Util (mapAccumLM) + expandAllocations :: Pass ExplicitMemory ExplicitMemory expandAllocations = Pass "expand allocations" "Expand allocations" $ @@ -63,6 +64,12 @@ where transform = identityMapper { mapOnBody = \scope -> localScope scope . transformBody } +nameInfoConv :: NameInfo ExplicitMemory -> NameInfo InKernel +nameInfoConv (LetInfo mem_info) = LetInfo mem_info +nameInfoConv (FParamInfo mem_info) = FParamInfo mem_info +nameInfoConv (LParamInfo mem_info) = LParamInfo mem_info +nameInfoConv (IndexInfo it) = IndexInfo it + transformExp :: Exp ExplicitMemory -> ExpandM (Stms ExplicitMemory, Exp ExplicitMemory) transformExp (Op (Inner (Kernel desc kspace ts kbody))) = do @@ -74,9 +81,10 @@ (alloc_stms, alloc_offsets) <- memoryRequirements kspace (kernelBodyStms kbody) variant_allocs invariant_allocs - kbody'' <- either compilerLimitationS pure $ - offsetMemoryInKernelBody alloc_offsets - kbody' + scope <- askScope + let scope' = scopeOfKernelSpace kspace <> M.map nameInfoConv scope + kbody'' <- either compilerLimitationS pure $ + runOffsetM scope' alloc_offsets $ offsetMemoryInKernelBody kbody' return (alloc_stms, Op $ Inner $ Kernel desc kspace ts kbody'') @@ -96,9 +104,11 @@ (alloc_stms, alloc_offsets) <- memoryRequirements kspace (bodyStms kbody) variant_allocs invariant_allocs - either compilerLimitationS pure $ do - kbody'' <- offsetMemoryInBody alloc_offsets kbody' - red_op'' <- offsetMemoryInLambda alloc_offsets red_op' + scope <- askScope + let scope' = scopeOfKernelSpace kspace <> M.map nameInfoConv scope + either compilerLimitationS pure $ runOffsetM scope' alloc_offsets $ do + kbody'' <- offsetMemoryInBody kbody' + red_op'' <- localScope (scopeOf red_op') $ offsetMemoryInLambda red_op' return (alloc_stms, Op $ Inner $ SegRed kspace comm red_op'' nes ts kbody'') @@ -108,9 +118,53 @@ M.keys (scopeOfKernelSpace kspace <> scopeOf (bodyStms kbody)) +transformExp (Op (Inner (SegGenRed kspace ops ts kbody))) = do + let (kbody', kbody_allocs) = extractBodyAllocations kbody + (ops', ops_allocs) = unzip $ map extractGenRedOpAllocations ops + variantAlloc (Var v) = v `S.member` bound_in_kernel + variantAlloc _ = False + allocs = kbody_allocs <> mconcat ops_allocs + (variant_allocs, invariant_allocs) = M.partition (variantAlloc . fst) allocs + + allocsForBody variant_allocs invariant_allocs kspace kbody' $ \alloc_stms kbody'' -> do + ops'' <- mapM offsetMemoryInGenRedOp ops' + + return (alloc_stms, + Op $ Inner $ SegGenRed kspace ops'' ts kbody'') + + where bound_in_kernel = + S.fromList $ map fst (spaceDimensions kspace) ++ + M.keys (scopeOfKernelSpace kspace <> + scopeOf (bodyStms kbody)) + + extractGenRedOpAllocations op = + let (lam, allocs) = extractLambdaAllocations $ genReduceOp op + in (op { genReduceOp = lam }, allocs) + + offsetMemoryInGenRedOp op = do + lam <- localScope (scopeOf (genReduceOp op)) $ + offsetMemoryInLambda $ genReduceOp op + return op { genReduceOp = lam } + transformExp e = return (mempty, e) +allocsForBody :: M.Map VName (SubExp, Space) + -> M.Map VName (SubExp, Space) + -> KernelSpace + -> Body InKernel + -> (Stms ExplicitMemory -> Body InKernel -> OffsetM b) + -> ExpandM b +allocsForBody variant_allocs invariant_allocs kspace kbody' m = do + (alloc_stms, alloc_offsets) <- + memoryRequirements kspace (bodyStms kbody') variant_allocs invariant_allocs + + scope <- askScope + let scope' = scopeOfKernelSpace kspace <> M.map nameInfoConv scope + either compilerLimitationS pure $ runOffsetM scope' alloc_offsets $ do + kbody'' <- offsetMemoryInBody kbody' + m alloc_stms kbody'' + memoryRequirements :: KernelSpace -> Stms InKernel -> M.Map VName (SubExp, Space) @@ -280,100 +334,147 @@ 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 +newtype OffsetM a = OffsetM (ReaderT (Scope InKernel) + (ReaderT RebaseMap (Either String)) a) + deriving (Applicative, Functor, Monad, + HasScope InKernel, LocalScope InKernel, + MonadError String) -offsetMemoryInKernelBody :: RebaseMap -> KernelBody InKernel - -> Either String (KernelBody InKernel) -offsetMemoryInKernelBody initial_offsets kbody = do - stms' <- snd <$> mapAccumLM offsetMemoryInStm initial_offsets +runOffsetM :: Scope InKernel -> RebaseMap -> OffsetM a -> Either String a +runOffsetM scope offsets (OffsetM m) = + runReaderT (runReaderT m scope) offsets + +askRebaseMap :: OffsetM RebaseMap +askRebaseMap = OffsetM $ lift ask + +lookupNewBase :: VName -> ([PrimExp VName], PrimType) -> OffsetM (Maybe IxFun) +lookupNewBase name x = do + offsets <- askRebaseMap + return $ ($ x) <$> M.lookup name offsets + +offsetMemoryInKernelBody :: KernelBody InKernel -> OffsetM (KernelBody InKernel) +offsetMemoryInKernelBody kbody = do + scope <- askScope + stms' <- stmsFromList . snd <$> + mapAccumLM (\scope' -> localScope scope' . offsetMemoryInStm) scope (stmsToList $ kernelBodyStms kbody) - return kbody { kernelBodyStms = stmsFromList stms' } + return kbody { kernelBodyStms = stms' } -offsetMemoryInBody :: RebaseMap -> Body InKernel -> Either String (Body InKernel) -offsetMemoryInBody offsets (Body attr stms res) = do - stms' <- stmsFromList . snd <$> mapAccumLM offsetMemoryInStm offsets (stmsToList stms) +offsetMemoryInBody :: Body InKernel -> OffsetM (Body InKernel) +offsetMemoryInBody (Body attr stms res) = do + scope <- askScope + stms' <- stmsFromList . snd <$> + mapAccumLM (\scope' -> localScope scope' . offsetMemoryInStm) scope + (stmsToList stms) return $ Body attr stms' res -offsetMemoryInLambda :: RebaseMap -> Lambda InKernel -> Either String (Lambda InKernel) -offsetMemoryInLambda offset lam = do - body <- offsetMemoryInBody offset $ lambdaBody lam - return $ lam { lambdaBody = body } +offsetMemoryInStm :: Stm InKernel -> OffsetM (Scope InKernel, Stm InKernel) +offsetMemoryInStm (Let pat attr e) = do + pat' <- offsetMemoryInPattern pat + e' <- localScope (scopeOfPattern pat') $ offsetMemoryInExp e + scope <- askScope + -- Try to recompute the index function. Fall back to creating rebase + -- operations with the RebaseMap. + rts <- runReaderT (expReturns e') scope + let pat'' = Pattern (patternContextElements pat') + (zipWith pick (patternValueElements pat') rts) + stm = Let pat'' attr e' + let scope' = scopeOf stm <> scope + return (scope', stm) + where pick :: PatElemT (MemInfo SubExp NoUniqueness MemBind) -> + ExpReturns -> PatElemT (MemInfo SubExp NoUniqueness MemBind) + pick (PatElem name (MemArray pt s u _ret)) + (MemArray _ _ _ (Just (ReturnsInBlock m extixfun))) + | Just ixfun <- instantiateIxFun extixfun = + PatElem name (MemArray pt s u (ArrayIn m ixfun)) + pick p _ = p -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') + instantiateIxFun :: ExtIxFun -> Maybe IxFun + instantiateIxFun = traverse (traverse inst) + where inst Ext{} = Nothing + inst (Free x) = return x -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 +offsetMemoryInPattern :: Pattern InKernel -> OffsetM (Pattern InKernel) +offsetMemoryInPattern (Pattern ctx vals) = do + mapM_ inspectCtx ctx + Pattern ctx <$> mapM inspectVal vals + where inspectVal patElem = do + new_attr <- offsetMemoryInMemBound $ patElemAttr patElem + return patElem { patElemAttr = new_attr } + inspectCtx 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 + | otherwise = return () -offsetMemoryInParam :: RebaseMap -> Param (MemBound u) -> Param (MemBound u) -offsetMemoryInParam offsets fparam = - fparam { paramAttr = offsetMemoryInMemBound offsets $ paramAttr fparam } +offsetMemoryInParam :: Param (MemBound u) -> OffsetM (Param (MemBound u)) +offsetMemoryInParam fparam = do + fparam' <- offsetMemoryInMemBound $ paramAttr fparam + return fparam { 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 +offsetMemoryInMemBound :: MemBound u -> OffsetM (MemBound u) +offsetMemoryInMemBound summary@(MemArray pt shape u (ArrayIn mem ixfun)) = do + new_base <- lookupNewBase mem (IxFun.base ixfun, pt) + return $ fromMaybe summary $ do + new_base' <- new_base + return $ MemArray pt shape u $ ArrayIn mem $ IxFun.rebase new_base' ixfun +offsetMemoryInMemBound summary = return 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 +offsetMemoryInBodyReturns :: BodyReturns -> OffsetM BodyReturns +offsetMemoryInBodyReturns br@(MemArray pt shape u (ReturnsInBlock mem ixfun)) + | Just ixfun' <- isStaticIxFun ixfun = do + new_base <- lookupNewBase mem (IxFun.base ixfun', pt) + return $ fromMaybe br $ do + new_base' <- new_base + return $ + MemArray pt shape u $ ReturnsInBlock mem $ + IxFun.rebase (fmap (fmap Free) new_base') ixfun +offsetMemoryInBodyReturns br = return 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 +offsetMemoryInLambda :: Lambda InKernel -> OffsetM (Lambda InKernel) +offsetMemoryInLambda lam = do + body <- offsetMemoryInBody $ lambdaBody lam + return $ lam { lambdaBody = body } + +offsetMemoryInExp :: Exp InKernel -> OffsetM (Exp InKernel) +offsetMemoryInExp (DoLoop ctx val form body) = do + let (ctxparams, ctxinit) = unzip ctx + (valparams, valinit) = unzip val + ctxparams' <- mapM offsetMemoryInParam ctxparams + valparams' <- mapM offsetMemoryInParam valparams + body' <- localScope (scopeOfFParams ctxparams' <> scopeOfFParams valparams' <> scopeOf form) (offsetMemoryInBody body) + return $ DoLoop (zip ctxparams' ctxinit) (zip valparams' valinit) form body' +offsetMemoryInExp (Op (Inner (GroupStream w max_chunk lam accs arrs))) = do + lam_accs <- mapM offsetMemoryInParam $ groupStreamAccParams lam + lam_arrs <- mapM offsetMemoryInParam $ groupStreamArrParams lam + let lam' = lam { groupStreamAccParams = lam_accs + , groupStreamArrParams = lam_arrs } - return $ Op $ Inner $ GroupStream w max_chunk lam' accs arrs -offsetMemoryInExp offsets (Op (Inner (GroupReduce w lam input))) = do - body <- offsetMemoryInBody offsets $ lambdaBody lam + body <- localScope (scopeOf lam') $ offsetMemoryInBody $ groupStreamLambdaBody lam + let lam'' = lam' { groupStreamLambdaBody = body } + return $ Op $ Inner $ GroupStream w max_chunk lam'' accs arrs +offsetMemoryInExp (Op (Inner (GroupReduce w lam input))) = do + body <- localScope (scopeOf lam) $ offsetMemoryInBody $ 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 +offsetMemoryInExp (Op (Inner (GroupScan w lam input))) = do + body <- localScope (scopeOf lam) $ offsetMemoryInBody $ lambdaBody lam + let lam' = lam { lambdaBody = body } + return $ Op $ Inner $ GroupScan w lam' input +offsetMemoryInExp (Op (Inner (GroupGenReduce w dests lam nes vals locks))) = do + lam_params <- mapM offsetMemoryInParam $ lambdaParams lam + let lam' = lam { lambdaParams = lam_params } + body <- localScope (scopeOf lam') $ offsetMemoryInBody $ lambdaBody lam + let lam'' = lam' { lambdaBody = body } + return $ Op $ Inner $ GroupGenReduce w dests lam'' nes vals locks +offsetMemoryInExp (Op (Inner (Combine cspace ts active body))) = + Op . Inner . Combine cspace ts active <$> offsetMemoryInBody body +offsetMemoryInExp e = mapExpM recurse e where recurse = identityMapper - { mapOnBody = const $ offsetMemoryInBody offsets - , mapOnBranchType = return . offsetMemoryInBodyReturns offsets + { mapOnBody = \bscope -> localScope bscope . offsetMemoryInBody + , mapOnBranchType = offsetMemoryInBodyReturns } ---- Slicing allocation sizes out of a kernel. diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -545,6 +545,13 @@ red_op' <- allocInSegRedLambda (spaceGlobalId space) (spaceNumThreads space) red_op return $ Inner $ SegRed space comm red_op' nes ts body' +handleKernel (SegGenRed space ops ts body) = do + body' <- subInKernel $ localScope (scopeOfKernelSpace space) $ allocInBodyNoDirect body + ops' <- forM ops $ \op -> do + lam <- allocInSegRedLambda (spaceGlobalId space) (spaceNumThreads space) $ genReduceOp op + return op { genReduceOp = lam } + return $ Inner $ SegGenRed space ops' ts body' + subInKernel :: AllocM InInKernel OutInKernel a -> AllocM fromlore2 ExplicitMemory a subInKernel = subAllocM handleKernelExp True diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs --- a/src/Futhark/Pass/ExtractKernels.hs +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -171,6 +171,7 @@ import Data.List import Futhark.Representation.SOACS +import qualified Futhark.Representation.SOACS.SOAC as SOAC import Futhark.Representation.SOACS.Simplify (simplifyStms, simpleSOACS) import qualified Futhark.Representation.Kernels as Out import Futhark.Representation.Kernels.Kernel @@ -462,9 +463,9 @@ addStms bnds letBind_ pat $ Op kernel -transformStm path (Let orig_pat (StmAux cs _) (Op (GenReduce w ops bucket_fun imgs))) = do +transformStm _ (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 + genReduceKernel orig_pat [] [] cs w ops bfun' imgs transformStm _ bnd = runBinder_ $ FOT.transformStmRecursively bnd @@ -1281,7 +1282,7 @@ -> [Int] -> Certificates -> SubExp - -> [GenReduceOp SOACS] + -> [SOAC.GenReduceOp SOACS] -> InKernelLambda -> [VName] -> KernelM KernelsStms @@ -1292,129 +1293,48 @@ (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 + ops' <- forM ops $ \(SOAC.GenReduceOp num_bins dests nes op) -> + SOAC.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' liftDistribM $ (nest_stms<>) <$> inScopeOf nest_stms - (genReduceKernel path (kernelNestLoops $ removeArraysFromNest all_dests nest) - orig_pat ispace inputs cs genred_w ops' lam arrs) + (genReduceKernel 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 :: KernelPath -> [LoopNesting] - -> Pattern -> [(VName, SubExp)] -> [KernelInput] - -> Certificates -> SubExp -> [GenReduceOp SOACS] +genReduceKernel :: Pattern -> [(VName, SubExp)] -> [KernelInput] + -> Certificates -> SubExp -> [SOAC.GenReduceOp SOACS] -> InKernelLambda -> [VName] -> DistribM 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 +genReduceKernel orig_pat ispace inputs cs genred_w ops lam arrs = do + ops' <- forM ops $ \(SOAC.GenReduceOp num_bins dests nes op) -> do + let (shape, op') = isVectorMap op + Out.GenReduceOp num_bins dests nes shape <$> Kernelise.transformLambda op' - let isDest = flip elem $ concatMap genReduceDest ops' + let isDest = flip elem $ concatMap Out.genReduceDest ops' inputs' = filter (not . isDest . kernelInputArray) inputs - (histos, k_stms) <- blockedGenReduce genred_w ispace inputs' ops' lam arrs - - let histos' = chunks (map (length . genReduceDest) ops') histos - pes = chunks (map (length . genReduceDest) ops') $ patternElements orig_pat - - (fmap (certify cs) k_stms<>) . mconcat <$> - inScopeOf k_stms (mapM combineIntermediateResults (zip3 pes ops histos')) - - where depth = length nests - - mkBodies num_histos pes num_bins nes op histos = runBinder $ do - 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 - - return (body_with_reshape, collapse_body) - - combineIntermediateResults (pes, GenReduceOp num_bins _ nes op, histos) = do - num_histos <- arraysSize depth <$> mapM lookupType histos - - ((body_with_reshape, collapse_body), aux_stms) <- mkBodies num_histos pes num_bins nes op histos - - segmented_reduce_stms <- - inScopeOf aux_stms $ transformStms path $ stmsToList $ bodyStms collapse_body - - let body_with_segred = mkBody segmented_reduce_stms $ - bodyResult collapse_body - - runBinder_ $ do - addStms aux_stms - - -- 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 - - letBindNames (map patElemName pes) $ - If num_histos_is_one body_with_reshape body_with_segred $ - IfAttr (staticShapes $ map patElemType pes) IfNormal + k_stms <- blockedGenReduce orig_pat genred_w ispace inputs' ops' lam arrs -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 - } + return $ certify cs <$> k_stms -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 +isVectorMap :: Lambda -> (Shape, Lambda) +isVectorMap lam + | [Let (Pattern [] pes) _ (Op (Screma w form arrs))] <- + stmsToList $ bodyStms $ lambdaBody lam, + bodyResult (lambdaBody lam) == map (Var . patElemName) pes, + Just map_lam <- isMapSOAC form, + arrs == map paramName (lambdaParams lam) = + let (shape, lam') = isVectorMap map_lam + in (Shape [w] <> shape, lam') + | otherwise = (mempty, lam) segmentedScanomapKernel :: KernelNest -> [Int] diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs --- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs @@ -42,7 +42,7 @@ 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 qualified Futhark.Representation.SOACS.SOAC as SOAC import Futhark.Util import Futhark.Util.IntegralExp @@ -345,7 +345,7 @@ -> [(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 <- SOAC.composeLambda SOAC.nilFn reduce_lam map_lam fold_lam' <- chunkLambda pat nes fold_lam let arr_idents = drop (length nes) $ patternIdents pat @@ -358,152 +358,34 @@ ispace nes (arrs ++ map_out_arrs) blockedGenReduce :: (MonadFreshNames m, HasScope Kernels m) => - SubExp + Pattern Kernels + -> 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 + -> m (Stms Kernels) +blockedGenReduce pat arr_w ispace inps ops lam arrs = runBinder_ $ do + let (_, segment_sizes) = unzip ispace 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) <- + (_, KernelSize num_groups group_size _ _ 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 + gtid <- newVName "gtid" + kspace <- newKernelSpace (num_groups, group_size, num_threads) $ + FlatThreadSpace $ ispace ++ [(gtid, arr_w)] - result <- letTupExp "result" $ DoLoop [] merge form loop_body - return $ map KernelInPlaceReturn result + body <- runBodyBinder $ localScope (scopeOfKernelSpace kspace) $ do + mapM_ (addStm <=< readKernelInput) inps + forM_ (zip (lambdaParams lam) arrs) $ \(p, arr) -> do + arr_t <- lookupType arr + letBindNames_ [paramName p] $ + BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var gtid] + return $ lambdaBody lam - let kbody = KernelBody () kstms kres - letTupExp "histograms" $ Op $ Kernel (KernelDebugHints "gen_reduce" []) kspace dest_ts kbody + letBind_ pat $ Op $ SegGenRed kspace ops (lambdaReturnType lam) body blockedMap :: (MonadFreshNames m, HasScope Kernels m) => Pattern Kernels -> SubExp @@ -653,8 +535,8 @@ -- First stage scan kernel. scanKernel1 :: (MonadBinder m, Lore m ~ Kernels) => SubExp -> KernelSize - -> Scan InKernel - -> Reduce InKernel + -> SOAC.Scan InKernel + -> SOAC.Reduce InKernel -> Lambda InKernel -> [VName] -> m (Kernel InKernel) scanKernel1 w scan_sizes (scan_lam, scan_nes) (_comm, red_lam, red_nes) foldlam arrs = do @@ -902,13 +784,13 @@ -- need them. blockedScan :: (MonadBinder m, Lore m ~ Kernels) => Pattern Kernels -> SubExp - -> Scan InKernel - -> Reduce InKernel + -> SOAC.Scan InKernel + -> SOAC.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 + foldlam <- SOAC.composeLambda scan_lam red_lam map_lam (_, first_scan_size) <- blockedKernelSize =<< asIntS Int64 w my_index <- newVName "my_index" diff --git a/src/Futhark/Pass/ExtractKernels/Segmented.hs b/src/Futhark/Pass/ExtractKernels/Segmented.hs --- a/src/Futhark/Pass/ExtractKernels/Segmented.hs +++ b/src/Futhark/Pass/ExtractKernels/Segmented.hs @@ -13,7 +13,7 @@ import Futhark.Representation.Kernels import Futhark.Representation.SOACS.SOAC (nilFn) import Futhark.MonadFreshNames -import Futhark.Tools hiding (true, false) +import Futhark.Tools hiding (false) import Futhark.Pass.ExtractKernels.BlockedKernel addFlagToLambda :: (MonadBinder m, Lore m ~ Kernels) => diff --git a/src/Futhark/Pass/ExtractKernels/Split.hs b/src/Futhark/Pass/ExtractKernels/Split.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/Split.hs +++ /dev/null @@ -1,41 +0,0 @@ --- | Functionality for identifying chunks of interesting parallelism --- inside of a map nesting. -module Futhark.Pass.ExtractKernels.Split - ( splitMap) 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 - -splitMap :: (MonadFreshNames m) => - Scope SOACS -> a -> m [a] -splitMap scope loop = return [loop] diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -19,7 +19,6 @@ 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 @@ -63,49 +62,6 @@ 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 >>> @@ -138,9 +94,7 @@ , simplifyExplicitMemory , doubleBuffer , simplifyExplicitMemory - ] >>> - memoryBlockMergingCoalescingMaybeCPU >>> - memoryBlockMergingReuseMaybeCPU + ] gpuPipeline :: Pipeline SOACS ExplicitMemory gpuPipeline = @@ -153,6 +107,4 @@ , simplifyExplicitMemory , expandAllocations , simplifyExplicitMemory - ] >>> - memoryBlockMergingCoalescingMaybeGPU >>> - memoryBlockMergingReuseMaybeGPU + ] diff --git a/src/Futhark/Representation/AST/Attributes/Types.hs b/src/Futhark/Representation/AST/Attributes/Types.hs --- a/src/Futhark/Representation/AST/Attributes/Types.hs +++ b/src/Futhark/Representation/AST/Attributes/Types.hs @@ -53,7 +53,6 @@ , generaliseExtTypes , existentialiseExtTypes , shapeMapping - , shapeMapping' , shapeExtMapping -- * Abbreviations @@ -73,7 +72,7 @@ import Control.Monad.State import Data.Maybe import Data.Monoid ((<>)) -import Data.List (elemIndex) +import Data.List (elemIndex, foldl') import qualified Data.Set as S import qualified Data.Map.Strict as M @@ -431,32 +430,34 @@ -- 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@. +-- to a set of the corresponding dimensions in @ts2@ (because they may +-- not fit exactly). -- -- 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 :: [TypeBase Shape u0] -> [TypeBase Shape u1] -> M.Map VName (S.Set 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 +shapeMapping' :: Ord a => [TypeBase Shape u] -> [[a]] -> M.Map VName (S.Set a) +shapeMapping' = dimMapping arrayDims id match (M.unionWith (<>)) where match Constant{} _ = M.empty - match (Var v) dim = M.singleton v dim + match (Var v) dim = M.singleton v $ S.singleton 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 +shapeExtMapping = dimMapping arrayExtDims arrayDims match mappend where match Free{} _ = mempty match (Ext i) dim = M.singleton i dim dimMapping :: Monoid res => (t1 -> [dim1]) -> (t2 -> [dim2]) -> (dim1 -> dim2 -> res) + -> (res -> res -> res) -> [t1] -> [t2] -> res -dimMapping getDims1 getDims2 f ts1 ts2 = - mconcat $ concat $ zipWith (zipWith f) (map getDims1 ts1) (map getDims2 ts2) +dimMapping getDims1 getDims2 f comb ts1 ts2 = + foldl' comb mempty $ concat $ zipWith (zipWith f) (map getDims1 ts1) (map getDims2 ts2) int8 :: PrimType int8 = IntType Int8 diff --git a/src/Futhark/Representation/Aliases.hs b/src/Futhark/Representation/Aliases.hs --- a/src/Futhark/Representation/Aliases.hs +++ b/src/Futhark/Representation/Aliases.hs @@ -56,7 +56,6 @@ 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. diff --git a/src/Futhark/Representation/ExplicitMemory.hs b/src/Futhark/Representation/ExplicitMemory.hs --- a/src/Futhark/Representation/ExplicitMemory.hs +++ b/src/Futhark/Representation/ExplicitMemory.hs @@ -554,7 +554,7 @@ TC.bad $ TC.TypeError $ "Array " ++ pretty v ++ " returned by function, but has nontrivial index function " ++ - pretty ixfun ++ " " ++ show ixfun + pretty ixfun matchBranchReturnType :: ExplicitMemorish lore => [BodyReturns] @@ -644,7 +644,7 @@ let x_ixfun' = substConstsInExtIndFun x_ixfun y_ixfun' = existentialiseIxFun0 y_ixfun unless (x_ixfun' == y_ixfun') $ - throwError $ unwords ["Index function unification fails1!", + throwError $ unwords ["Index function unification failed (ReturnsInBlock)", "\nixfun of body result: ", pretty y_ixfun', "\nixfun of return type: ", pretty x_ixfun', "\nand context elements: ", pretty ctx_res] @@ -654,7 +654,7 @@ let x_ixfun' = substConstsInExtIndFun x_ixfun y_ixfun' = existentialiseIxFun0 y_ixfun unless (x_ixfun' == y_ixfun') $ - throwError $ unwords ["Index function unification fails2!", + throwError $ unwords ["Index function unification failed (ReturnsNewBlock)", "\nixfun of body result: ", pretty y_ixfun', "\nixfun of return type: ", pretty x_ixfun', "\nand context elements: ", pretty ctx_res] @@ -1017,6 +1017,8 @@ extendedScope (varReturns arr) (castScope $ scopeOf $ kernelBodyStms body) correct _ ret = return ret + opReturns (Inner (SegGenRed _ ops _ _)) = + concat <$> mapM (mapM varReturns . genReduceDest) ops opReturns k = extReturns <$> opType k diff --git a/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs --- a/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs +++ b/src/Futhark/Representation/ExplicitMemory/IndexFunction.hs @@ -1,447 +1,789 @@ --- | 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 +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} +-- | This module contains a representation for the index function based on +-- linear-memory accessor descriptors; see Zhu, Hoeflinger and David work. +module Futhark.Representation.ExplicitMemory.IndexFunction + ( IxFun(..) + , index + , iota + , offsetIndex + , strideIndex + , permute + , rotate + , reshape + , slice + , rebase + , repeat + , isContiguous + , shape + , rank + , linearWithOffset + , rearrangeWithOffset + , isDirect + , isLinear + , substituteInIxFun + ) + where + +import Prelude hiding (mod, repeat) +import Data.List hiding (repeat) +import qualified Data.List.NonEmpty as NE +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Function (on) +import Data.Maybe (isJust) +import Control.Monad.Identity +import Control.Monad.Writer +import qualified Data.Map.Strict as M + +import Futhark.Transform.Substitute +import Futhark.Transform.Rename +import Futhark.Representation.AST.Syntax + (ShapeChange, DimChange(..), DimIndex(..), Slice, unitSlice, dimFix) +import Futhark.Representation.AST.Attributes +import Futhark.Util.IntegralExp +import Futhark.Util.Pretty +import Futhark.Analysis.PrimExp.Convert + + +-- | LMAD's representation consists of a permutation, a general offset, and, for +-- each dimension a stride, rotate factor, number of elements, permutation, and +-- monotonicity and unit-stride info for each dimension. Note that the +-- permutation is not strictly necessary in that the permutation can be +-- performed directly on LMAD dimensions, but then it is difficult to extract +-- the permutation back from an LMAD. +-- +-- LMAD algebra is closed under composition w.r.t. operators such as permute, +-- repeat, index and slice. However, other operations, such as reshape, cannot +-- always be represented inside the LMAD algebra. +-- +-- It follows that the general representation of an index function is a list of +-- LMADS, in which each following LMAD in the list implicitly corresponds to an +-- irregular reshaping operation. +-- +-- However, we expect that the common case is when the index function is one +-- LMAD -- we call this the 'nice' representation. +-- +-- Finally, the list of LMADs is tupled with the shape of the original array, +-- and with contiguous info, i.e., if we instantiate all the points of the +-- current index function, do we get a contiguous memory interval? +-- +-- By definition, the LMAD denotes the set of points (simplified): +-- +-- \{ 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 offset dims) = + braces $ semisep [ text "offset: " <> oneLine (ppr offset) + , text "strides: " <> p ldStride + , text "rotates: " <> p ldRotate + , text "shape: " <> p ldShape + , text "permutation: " <> p ldPerm + , text "monotonicity: " <> p ldMon + ] + where p f = oneLine $ brackets $ commasep $ map (ppr . f) dims + +instance Pretty num => Pretty (IxFun num) where + ppr (IxFun lmads oshp cg) = + braces $ semisep [ text "base: " <> brackets (commasep $ map ppr oshp) + , text "contiguous: " <> text (show cg) + , text "LMADs: " <> brackets (commasep $ NE.toList $ NE.map ppr lmads) + ] + + +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 offset dims) = + LMAD <$> f offset <*> traverse f' dims + where f' (LMADDim s r n p m) = + LMADDim <$> f s <*> f r <*> f n <*> pure p <*> pure m + +instance Traversable IxFun where + traverse f (IxFun lmads oshp cg) = + IxFun <$> traverse (traverse f) lmads <*> traverse f oshp <*> pure cg + +(++@) :: [a] -> NonEmpty a -> NonEmpty a +es ++@ (ne :| nes) = case es of + e : es' -> e :| es' ++ [ne] ++ nes + [] -> ne :| nes + +(@++@) :: NonEmpty a -> NonEmpty a -> NonEmpty a +(x :| xs) @++@ (y :| ys) = x :| xs ++ [y] ++ ys + +invertMonotonicity :: Monotonicity -> Monotonicity +invertMonotonicity Inc = Dec +invertMonotonicity Dec = Inc +invertMonotonicity Unknown = Unknown + +lmadPermutation :: LMAD num -> Permutation +lmadPermutation = map ldPerm . lmadDims + +setLMADPermutation :: Permutation -> LMAD num -> LMAD num +setLMADPermutation perm lmad = + lmad { lmadDims = zipWith (\dim p -> dim { ldPerm = p }) (lmadDims lmad) perm } + +setLMADShape :: Shape num -> LMAD num -> LMAD num +setLMADShape shp lmad = lmad { lmadDims = zipWith (\dim s -> dim { ldShape = s }) (lmadDims lmad) shp } + +-- | Substitute a name with a PrimExp in an LMAD. +substituteInLMAD :: Ord a => M.Map a (PrimExp a) -> LMAD (PrimExp a) + -> LMAD (PrimExp a) +substituteInLMAD tab (LMAD offset dims) = + let offset' = substituteInPrimExp tab offset + dims' = map (\(LMADDim s r n p m) -> + LMADDim + (substituteInPrimExp tab s) + (substituteInPrimExp tab r) + (substituteInPrimExp tab n) + p m) + dims + in LMAD offset' dims' + +-- | Substitute 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 (IxFun lmads oshp cg) = + IxFun (NE.map (substituteInLMAD tab) lmads) + (map (substituteInPrimExp tab) oshp) + cg + +-- | Is this is a row-major array? +isDirect :: (Eq num, IntegralExp num) => IxFun num -> Bool +isDirect ixfun@(IxFun (LMAD offset dims :| []) oshp True) = + let strides_expected = reverse $ scanl (*) 1 (reverse (tail oshp)) + in hasContiguousPerm ixfun && + length oshp == length dims && + offset == 0 && + all (\(LMADDim s r n p _, m, d, se) -> + s == se && r == 0 && n == d && p == m) + (zip4 dims [0..length dims - 1] oshp strides_expected) +isDirect _ = False + +-- | Does the index function have contiguous memory support? +isContiguous :: (Eq num, IntegralExp num) => IxFun num -> Bool +isContiguous ixfun = ixfunContig ixfun && hasContiguousPerm ixfun + +-- | Does the index function have an ascending permutation? +hasContiguousPerm :: IxFun num -> Bool +hasContiguousPerm (IxFun (lmad :| []) _ _) = + let perm = lmadPermutation lmad + in perm == sort perm +hasContiguousPerm _ = False + +-- | Shape of an index function. +shape :: (Eq num, IntegralExp num) => IxFun num -> Shape num +shape (IxFun (lmad :| _) _ _) = lmadShape lmad + +-- | Shape of an LMAD. +lmadShape :: (Eq num, IntegralExp num) => LMAD num -> Shape num +lmadShape lmad = permuteInv (lmadPermutation lmad) $ lmadShapeBase lmad + +-- | Shape of an LMAD, ignoring permutations. +lmadShapeBase :: (Eq num, IntegralExp num) => LMAD num -> Shape num +lmadShapeBase = map ldShape . lmadDims + +-- | Compute 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 = indexFromLMADs . ixfunLMADs + where indexFromLMADs :: (IntegralExp num, Eq num) => + NonEmpty (LMAD num) -> Indices num -> num -> num + indexFromLMADs (lmad :| []) inds elm_size = indexLMAD lmad inds elm_size + indexFromLMADs (lmad1 :| lmad2 : lmads) inds elm_size = + let i_flat = indexLMAD lmad1 inds 1 + new_inds = unflattenIndex (permuteFwd (lmadPermutation lmad2) $ lmadShapeBase lmad2) i_flat + in indexFromLMADs (lmad2 :| lmads) new_inds elm_size + + -- | Compute the flat index of an LMAD. + indexLMAD :: (IntegralExp num, Eq num) => + LMAD num -> Indices num -> num -> num + indexLMAD lmad@(LMAD off dims) inds elm_size = + let prod = sum $ zipWith flatOneDim + (map (\(LMADDim s r n _ _) -> (s, r, n)) dims) + (permuteInv (lmadPermutation lmad) inds) + ind = off + prod + in if elm_size == 1 then ind else ind * elm_size + +-- | iota. +iota :: IntegralExp num => Shape num -> IxFun num +iota ns = + let rs = replicate (length ns) 0 + in IxFun (makeRotIota Inc 0 (zip rs ns) :| []) ns True + +-- | Permute dimensions. +permute :: IntegralExp num => + IxFun num -> Permutation -> IxFun num +permute (IxFun (lmad :| lmads) oshp cg) perm_new = + let perm_cur = lmadPermutation lmad + perm = map (perm_cur !!) perm_new + in IxFun (setLMADPermutation perm lmad :| lmads) oshp cg + +-- | Repeat dimensions. +repeat :: (Eq num, IntegralExp num) => + IxFun num -> [Shape num] -> Shape num -> IxFun num +repeat (IxFun (lmad@(LMAD off dims) :| lmads) oshp _) shps shp = + let perm = lmadPermutation 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] + + dims' = concatMap (\(shp_k, srnp) -> + map fakeDim shp_k ++ [srnp] + ) $ zip shps' dims + lmad' = setLMADPermutation perm'' $ LMAD off (dims' ++ map fakeDim shp) + in IxFun (lmad' :| lmads) oshp False -- XXX: Can we be less conservative? + where fakeDim x = LMADDim 0 0 x 0 Unknown + +-- | Rotate an index function. +rotate :: (Eq num, IntegralExp num) => + IxFun num -> Indices num -> IxFun num +rotate (IxFun (lmad@(LMAD off dims) :| lmads) oshp cg) offs = + let dims' = zipWith (\(LMADDim s r n p f) o -> + if s == 0 then LMADDim 0 0 n p Unknown + else LMADDim s (r + o) n p f + ) dims (permuteInv (lmadPermutation lmad) offs) + in IxFun (LMAD off dims' :| lmads) oshp cg + +-- | Handle the case where a slice can stay within a single LMAD. +sliceOneLMAD :: (Eq num, IntegralExp num) => + IxFun num -> Slice num -> Maybe (IxFun num) +sliceOneLMAD (IxFun (lmad@(LMAD _ ldims) :| lmads) oshp cg) is = do + let perm = lmadPermutation lmad + is' = permuteInv perm is + cg' = cg && slicePreservesContiguous lmad is' + guard $ harmlessRotation lmad is' + let lmad' = foldl sliceOne (LMAD (lmadOffset lmad) []) $ zip is' ldims + -- need to remove the fixed dims from the permutation + perm' = updatePerm perm $ map fst $ filter (isJust . dimFix . snd) $ + zip [0..length is' - 1] is' + + return $ IxFun (setLMADPermutation perm' lmad' :| lmads) oshp cg' + where 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] + + harmlessRotation' :: (Eq num, IntegralExp num) => + LMADDim num -> DimIndex num -> Bool + harmlessRotation' _ (DimFix _) = True + harmlessRotation' (LMADDim 0 _ _ _ _) _ = True + harmlessRotation' (LMADDim _ 0 _ _ _) _ = True + harmlessRotation' (LMADDim _ _ n _ _) dslc + | dslc == DimSlice (n - 1) n (-1) || + dslc == unitSlice 0 n = True + harmlessRotation' _ _ = False + + harmlessRotation :: (Eq num, IntegralExp num) => + LMAD num -> Slice num -> Bool + harmlessRotation (LMAD _ dims) iss = + and $ zipWith harmlessRotation' dims iss + + -- XXX: TODO: what happens to r on a negative-stride slice; is there + -- such a case? + sliceOne :: (Eq num, IntegralExp num) => + LMAD num -> (DimIndex num, LMADDim num) -> LMAD num + sliceOne (LMAD off dims) (DimFix i, LMADDim s r n _ _) = + LMAD (off + flatOneDim (s, r, n) i) dims + sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ _ p _) = + LMAD off (dims ++ [LMADDim 0 0 ne p Unknown]) + sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ _ n _ _)) + | dmind == unitSlice 0 n = LMAD off (dims ++ [dim]) + sliceOne (LMAD off dims) (dmind, LMADDim s r n p m) + | dmind == DimSlice (n - 1) n (-1) = + let r' = if r == 0 then 0 else n - r + off' = off + flatOneDim (s, 0, n) (n - 1) + in LMAD off' (dims ++ [LMADDim (s * (-1)) r' n p (invertMonotonicity m)]) + sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s r n p _) = + LMAD (off + flatOneDim (s, r, n) b) (dims ++ [LMADDim 0 0 ne p Unknown]) + sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s 0 _ p m) = + let m' = case sgn ss of + Just 1 -> m + Just (-1) -> invertMonotonicity m + _ -> Unknown + in LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) 0 ns p m']) + sliceOne _ _ = error "slice: reached impossible case" + + slicePreservesContiguous :: (Eq num, IntegralExp num) => + LMAD num -> Slice num -> Bool + slicePreservesContiguous (LMAD _ dims) slc = + -- remove from the slice the LMAD dimensions that 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 (dims', slc') = unzip $ + filter ((/= 0) . ldStride . fst) $ + zip dims $ 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 unrotated or full. + -- 3. the rest of inner sliced dims are full. + (_, success) = + foldl (\(found, res) (slcdim, LMADDim _ 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' dims' + in success + + normIndex :: (Eq num, IntegralExp num) => + DimIndex num -> DimIndex num + normIndex (DimSlice b 1 _) = DimFix b + normIndex (DimSlice b _ 0) = DimFix b + normIndex d = d + +-- | Slice an index function. +slice :: (Eq num, IntegralExp num) => + IxFun num -> Slice num -> IxFun num +slice _ [] = error "slice: empty slice" +slice ixfun@(IxFun (lmad@(LMAD _ _) :| lmads) oshp cg) dim_slices + -- Avoid identity slicing. + | dim_slices == map (unitSlice 0) (shape ixfun) = ixfun + | Just ixfun' <- sliceOneLMAD ixfun dim_slices = ixfun' + | otherwise = + case sliceOneLMAD (iota (lmadShape lmad)) dim_slices of + Just (IxFun (lmad' :| []) _ cg') -> + IxFun (lmad' :| lmad : lmads) oshp (cg && cg') + _ -> error "slice: reached impossible case" + +-- | Handle the simple case where all reshape dimensions are coercions. +reshapeCoercion :: (Eq num, IntegralExp num) => + IxFun num -> ShapeChange num -> Maybe (IxFun num) +reshapeCoercion (IxFun (lmad@(LMAD off dims) :| lmads) _ cg) newshape = do + let perm = lmadPermutation lmad + (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape + let hd_len = length head_coercions + num_coercions = hd_len + length tail_coercions + dims' = permuteFwd perm dims + mid_dims = take (length dims - num_coercions) $ drop hd_len dims' + num_rshps = length reshapes + guard (num_rshps == 0 || (num_rshps == 1 && length mid_dims == 1)) + let dims'' = map snd $ sortBy (compare `on` fst) $ + zipWith (\ld n -> (ldPerm ld, ld { ldShape = n })) + dims' (newDims newshape) + lmad' = LMAD off dims'' + return $ IxFun (lmad' :| lmads) (newDims newshape) cg + +-- | Handle the case where a reshape operation can stay inside a single LMAD. +-- +-- There are four conditions that all must hold for the result of a reshape +-- operation to remain in 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 these conditions do 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. +reshapeOneLMAD :: (Eq num, IntegralExp num) => + IxFun num -> ShapeChange num -> Maybe (IxFun num) +reshapeOneLMAD ixfun@(IxFun (lmad@(LMAD off dims) :| lmads) _ cg) newshape = do + let perm = lmadPermutation lmad + (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape + let hd_len = length head_coercions + num_coercions = hd_len + length tail_coercions + dims_perm = permuteFwd perm dims + mid_dims = take (length dims - num_coercions) $ drop hd_len dims_perm + -- Ignore rotates, as we only care about not having rotates in the + -- dimensions that aren't coercions (@mid_dims@), which we check + -- separately. + mon = ixfunMonotonicityRots True ixfun + + guard $ + -- checking conditions (2) and (3) + all (\ (LMADDim s r _ _ _) -> s /= 0 && r == 0) mid_dims && + -- checking condition (1) + consecutive hd_len (map ldPerm mid_dims) && + -- checking condition (4) + hasContiguousPerm ixfun && cg && (mon == Inc || mon == Dec) + + -- make new permutation + let rsh_len = length reshapes + diff = length newshape - length dims + 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_dims not affected + else let p = ldPerm (dims !! ind) + in if p < hd_len + then p + else p + diff + ) iota_shape + -- split the dimensions + (support_inds, repeat_inds) = + foldl (\(sup, rpt) (i, shpdim, ip) -> + case (i < hd_len, i >= hd_len + rsh_len, shpdim) of + (True, _, DimCoercion n) -> + case dims_perm !! i of + (LMADDim 0 _ _ _ _) -> ( sup, (ip, n) : rpt ) + (LMADDim _ r _ _ _) -> ( (ip, (r, n)) : sup, rpt ) + (_, True, DimCoercion n) -> + case dims_perm !! (i-diff) of + (LMADDim 0 _ _ _ _) -> ( sup, (ip, n) : rpt ) + (LMADDim _ 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 $ sortBy (compare `on` fst) support_inds + (rpt_inds, repeats) = unzip repeat_inds + LMAD off' dims_sup = makeRotIota mon off support + repeats' = map (\n -> LMADDim 0 0 n 0 Unknown) repeats + dims' = map snd $ sortBy (compare `on` fst) + $ zip sup_inds dims_sup ++ zip rpt_inds repeats' + lmad' = LMAD off' dims' + return $ IxFun (setLMADPermutation perm' lmad' :| lmads) (newDims newshape) cg + where consecutive _ [] = True + consecutive i [p]= i == p + consecutive i ps = and $ zipWith (==) ps [i, i+1..] + +splitCoercions :: (Eq num, IntegralExp num) => + ShapeChange num -> Maybe (ShapeChange num, ShapeChange num, ShapeChange num) +splitCoercions newshape' = do + let (head_coercions, newshape'') = span isCoercion newshape' + (reshapes, tail_coercions) = break isCoercion newshape'' + guard (all isCoercion tail_coercions) + return (head_coercions, reshapes, tail_coercions) + where isCoercion DimCoercion{} = True + isCoercion _ = False + +-- | Reshape an index function. +reshape :: (Eq num, IntegralExp num) => + IxFun num -> ShapeChange num -> IxFun num +reshape ixfun new_shape + | Just ixfun' <- reshapeCoercion ixfun new_shape = ixfun' + | Just ixfun' <- reshapeOneLMAD ixfun new_shape = ixfun' +reshape (IxFun (lmad0 :| lmad0s) oshp cg) new_shape = + case iota (newDims new_shape) of + IxFun (lmad :| []) _ _ -> IxFun (lmad :| lmad0 : lmad0s) oshp cg + _ -> error "reshape: reached impossible case" + +rank :: IntegralExp num => + IxFun num -> Int +rank (IxFun (LMAD _ sss :| _) _ _) = length sss + +-- | Handle the case where a rebase operation can stay within m + n - 1 LMADs, +-- where m is the number of LMADs in the index function, and n is the number of +-- LMADs in the new base. If both index function have only on LMAD, this means +-- that we stay within the single-LMAD domain. +-- +-- We can often stay in that domain if the original ixfun is essentially a +-- slice, e.g. `x[i, (k1,m,s1), (k2,n,s2)] = orig`. +-- +-- XXX: TODO: handle repetitions in both lmads. +-- +-- 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 ixfun0 such that: +-- +-- new-base == rebase ixfun0 ixfun, or equivalently: +-- new-base == ixfun o ixfun0 +-- +-- because then I can go bottom up and compose with ixfun0 all the index +-- functions corresponding to the memory block associated with ixfun. +rebaseNice :: (Eq num, IntegralExp num) => + IxFun num -> IxFun num -> Maybe (IxFun num) +rebaseNice + new_base@(IxFun (lmad_base :| lmads_base) _ cg_base) + ixfun@(IxFun lmads shp cg) = do + let (lmad_full :| lmads') = NE.reverse lmads + ((outer_shapes, inner_shape), lmad) = shaveoffRepeats lmad_full + dims = lmadDims lmad + perm = lmadPermutation lmad + perm_base = lmadPermutation lmad_base + + guard $ + -- Core rebase condition. + base ixfun == shape new_base + -- Conservative safety conditions: ixfun is contiguous and has known + -- monotonicity for all dimensions. + && cg && all ((/= Unknown) . ldMon) dims + -- XXX: We should be able to handle some basic cases where both index + -- functions have non-trivial permutations. + && (hasContiguousPerm ixfun || hasContiguousPerm new_base) + -- We need the permutations to be of the same size if we want to compose + -- them. They don't have to be of the same size if the ixfun has a trivial + -- permutation. Supporting this latter case allows us to rebase when ixfun + -- has been created by slicing with fixed dimensions. + && (length perm == length perm_base || hasContiguousPerm ixfun) + -- To not have to worry about ixfun having non-1 strides, we also check that + -- it is a row-major array (modulo permutation, which is handled + -- separately). Accept a non-full innermost dimension. XXX: Maybe this can + -- be less conservative? + && and (zipWith3 (\sn ld inner -> sn == ldShape ld || (inner && ldStride ld == 1)) + shp dims (replicate (length dims - 1) False ++ [True])) + + -- Compose permutations, reverse strides and adjust offset if necessary. + let perm_base' = if hasContiguousPerm ixfun + then perm_base + else map (perm !!) perm_base + lmad_base' = setLMADPermutation perm_base' lmad_base + dims_base = lmadDims lmad_base' + n_fewer_dims = length dims_base - length dims + (dims_base', offs_contrib) = unzip $ + zipWith (\(LMADDim s1 r1 n1 p1 _) (LMADDim _ r2 _ _ m2) -> + let (s', off') | m2 == Inc = (s1, 0) + | otherwise = (s1 * (-1), s1 * (n1 - 1)) + r' | m2 == Inc = if r2 == 0 then r1 else r1 + r2 + | r1 == 0 = r2 + | r2 == 0 = n1 - r1 + | otherwise = n1 - r1 + r2 + in (LMADDim s' r' n1 (p1 - n_fewer_dims) Inc, off')) + -- If @dims@ is morally a slice, it might have fewer dimensions than + -- @dims_base@. Drop extraneous outer dimensions. + (drop n_fewer_dims dims_base) dims + off_base = lmadOffset lmad_base' + sum offs_contrib + lmad_base'' + | lmadOffset lmad == 0 = LMAD off_base dims_base' + | otherwise = + -- If the innermost dimension of the ixfun was not full (but still + -- had a stride of 1), add its offset relative to the new base. + setLMADShape (lmadShape lmad) + (LMAD (off_base + ldStride (last dims_base) * lmadOffset lmad) + dims_base') + new_base' = IxFun (lmad_base'' :| lmads_base) shp cg_base + IxFun lmads_base' _ _ = if all null outer_shapes && null inner_shape + then new_base' + else repeat new_base' outer_shapes inner_shape + lmads'' = lmads' ++@ lmads_base' + return $ IxFun lmads'' shp (cg && cg_base) + where shaveoffRepeats :: (Eq num, IntegralExp num) => + LMAD num -> (([Shape num], Shape num), LMAD num) + shaveoffRepeats lmad = + -- 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. + let perm = lmadPermutation lmad + dims = lmadDims lmad + -- compute the Repeat: + resacc= foldl (\acc (LMADDim s _ n _ _) -> + case acc of + rpt : acc0 -> + if s == 0 then (n : rpt) : acc0 + else [] : (rpt : acc0) + _ -> error "shaveoffRepeats: empty accumulator" + ) [[]] $ reverse $ permuteFwd perm dims + last_shape = last resacc + shapes = take (length resacc - 1) resacc + -- update permutation and lmad: + howManyRepLT k = + foldl (\i (LMADDim s _ _ p _) -> + if s == 0 && p < k then i + 1 else i + ) 0 dims + dims' = foldl (\acc (LMADDim s r n p info) -> + if s == 0 then acc + else let p' = p - howManyRepLT p + in LMADDim s r n p' info : acc + ) [] $ reverse dims + lmad' = LMAD (lmadOffset lmad) dims' + in ((shapes, last_shape), lmad') + +-- | Rebase an index function on top of a new base. +rebase :: (Eq num, IntegralExp num) => + IxFun num -> IxFun num -> IxFun num +rebase new_base@(IxFun lmads_base shp_base cg_base) ixfun@(IxFun lmads shp cg) + | Just ixfun' <- rebaseNice new_base ixfun = ixfun' + -- In the general case just concatenate LMADs since this refers to index + -- function composition, which is always safe. + | otherwise = + let (lmads_base', shp_base') = + if base ixfun == shape new_base + then (lmads_base, shp_base) + else let IxFun lmads' shp_base'' _ = reshape new_base $ map DimCoercion shp + in (lmads', shp_base'') + in IxFun (lmads @++@ lmads_base') shp_base' (cg && cg_base) + +ixfunMonotonicity :: (Eq num, IntegralExp num) => IxFun num -> Monotonicity +ixfunMonotonicity = ixfunMonotonicityRots False + +-- | Offset index. 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" + +-- | Stride index. 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@(IxFun (lmad :| []) _ cg) elem_size + | hasContiguousPerm ixfun && cg && ixfunMonotonicity ixfun == Inc = + Just $ lmadOffset lmad * elem_size +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 (lmad :| []) oshp cg) elem_size = do + -- Note that @cg@ describes whether the index function is + -- contiguous, *ignoring permutations*. This function requires that + -- functionality. + let perm = lmadPermutation lmad + perm_contig = [0..length perm-1] + offset <- linearWithOffset + (IxFun (setLMADPermutation perm_contig lmad :| []) oshp cg) elem_size + return (offset, zip perm (permuteFwd perm (lmadShapeBase lmad))) +rearrangeWithOffset _ _ = Nothing + +isLinear :: (Eq num, IntegralExp num) => IxFun num -> Bool +isLinear = (== Just 0) . flip linearWithOffset 1 + +permuteFwd :: Permutation -> [a] -> [a] +permuteFwd ps elems = map (elems !!) ps + +permuteInv :: Permutation -> [a] -> [a] +permuteInv ps elems = map snd $ sortBy (compare `on` fst) $ zip ps elems + +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 + +-- | Generalised iota with user-specified offset and strides. +makeRotIota :: IntegralExp num => + Monotonicity -> num -> [(num, num)] -> LMAD num +makeRotIota mon off support + | mon == Inc || mon == Dec = + let rk = length support + (rs, ns) = unzip support + ss0 = reverse $ take rk $ scanl (*) 1 $ reverse ns + ss = if mon == Inc + then ss0 + else map (* (-1)) ss0 + ps = map fromIntegral [0..rk-1] + fi = replicate rk mon + in LMAD off $ zipWith5 LMADDim ss rs ns ps fi + | otherwise = error "makeRotIota: requires Inc or Dec" + +-- | Check monotonicity of an index function. +ixfunMonotonicityRots :: (Eq num, IntegralExp num) => + Bool -> IxFun num -> Monotonicity +ixfunMonotonicityRots ignore_rots (IxFun (lmad :| lmads) _ _) = + let mon0 = lmadMonotonicityRots lmad + in if all ((== mon0) . lmadMonotonicityRots) lmads + then mon0 + else Unknown + where lmadMonotonicityRots :: (Eq num, IntegralExp num) => + LMAD num -> Monotonicity + lmadMonotonicityRots (LMAD _ dims) + | all (isMonDim Inc) dims = Inc + | all (isMonDim Dec) dims = Dec + | otherwise = Unknown + + isMonDim :: (Eq num, IntegralExp num) => + Monotonicity -> LMADDim num -> Bool + isMonDim mon (LMADDim s r _ _ ldmon) = + s == 0 || ((ignore_rots || r == 0) && mon == ldmon) diff --git a/src/Futhark/Representation/ExplicitMemory/Lmad.hs b/src/Futhark/Representation/ExplicitMemory/Lmad.hs deleted file mode 100644 --- a/src/Futhark/Representation/ExplicitMemory/Lmad.hs +++ /dev/null @@ -1,761 +0,0 @@ -{-# 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/Kernels/Kernel.hs b/src/Futhark/Representation/Kernels/Kernel.hs --- a/src/Futhark/Representation/Kernels/Kernel.hs +++ b/src/Futhark/Representation/Kernels/Kernel.hs @@ -11,6 +11,7 @@ ( Kernel(..) , kernelType , KernelDebugHints(..) + , GenReduceOp(..) , KernelBody(..) , KernelSpace(..) , spaceDimensions @@ -35,6 +36,7 @@ ) where +import Control.Arrow (first) import Control.Monad.Writer hiding (mapM_) import Control.Monad.Identity hiding (mapM_) import qualified Data.Set as S @@ -79,6 +81,20 @@ } deriving (Eq, Show, Ord) +data GenReduceOp lore = + GenReduceOp { genReduceWidth :: SubExp + , genReduceDest :: [VName] + , genReduceNeutral :: [SubExp] + , genReduceShape :: Shape + -- ^ In case this operator is semantically a + -- vectorised operator (corresponding to a perfect map + -- nest in the SOACS representation), these are the + -- logical "dimensions". This is used to generate + -- more efficient code. + , genReduceOp :: LambdaT lore + } + deriving (Eq, Ord, Show) + data Kernel lore = GetSize Name SizeClass -- ^ Produce some runtime-configurable size. | GetSizeMax SizeClass -- ^ The maximum size of some class. @@ -88,6 +104,7 @@ | SegRed KernelSpace Commutativity (Lambda lore) [SubExp] [Type] (Body lore) -- ^ The KernelSpace must always have at least two dimensions, -- implying that the result of a SegRed is always an array. + | SegGenRed KernelSpace [GenReduceOp lore] [Type] (Body lore) deriving (Eq, Show, Ord) data KernelSpace = KernelSpace { spaceGlobalId :: VName @@ -200,6 +217,18 @@ <*> mapM (mapOnKernelSubExp tv) nes <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts <*> mapOnKernelBody tv lam +mapKernelM tv (SegGenRed space ops ts body) = + SegGenRed + <$> mapOnKernelSpace tv space + <*> mapM onGenRedOp ops + <*> mapM (mapOnType $ mapOnKernelSubExp tv) ts + <*> mapOnKernelBody tv body + where onGenRedOp (GenReduceOp w arrs nes shape op) = + GenReduceOp <$> mapOnKernelSubExp tv w + <*> mapM (mapOnKernelVName tv) arrs + <*> mapM (mapOnKernelSubExp tv) nes + <*> (Shape <$> mapM (mapOnKernelSubExp tv) (shapeDims shape)) + <*> mapOnKernelLambda tv op mapKernelM tv (Kernel desc space ts kernel_body) = Kernel <$> mapOnKernelDebugHints desc <*> mapOnKernelSpace tv space <*> @@ -421,6 +450,13 @@ dims = map snd $ spaceDimensions space outer_dims = init dims +kernelType (SegGenRed space ops _ _) = do + op <- ops + let shape = Shape (segment_dims <> [genReduceWidth op]) <> genReduceShape op + map (`arrayOfShape` shape) (lambdaReturnType $ genReduceOp op) + where dims = map snd $ spaceDimensions space + segment_dims = init dims + kernelType GetSize{} = [Prim int32] kernelType GetSizeMax{} = [Prim int32] kernelType CmpSizeLe{} = [Prim Bool] @@ -442,6 +478,9 @@ mconcat (map consumedByReturn (kernelBodyResult kbody)) where consumedByReturn (WriteReturn _ a _) = S.singleton a consumedByReturn _ = mempty + consumedInOp (SegGenRed _ ops _ body) = + S.fromList (concatMap genReduceDest ops) <> + consumedInBody body consumedInOp _ = mempty aliasAnalyseKernelBody :: (Attributes lore, @@ -567,6 +606,9 @@ mconcat $ map UT.consumedUsage $ S.toList $ consumedInKernelBody kbody usageInOp (SegRed _ _ _ _ _ body) = mconcat $ map UT.consumedUsage $ S.toList $ consumedInBody body + usageInOp (SegGenRed _ ops _ body) = + mconcat $ map UT.consumedUsage $ S.toList (consumedInBody body) <> + concatMap genReduceDest ops usageInOp GetSize{} = mempty usageInOp GetSizeMax{} = mempty usageInOp CmpSizeLe{} = mempty @@ -598,6 +640,44 @@ TC.checkLambdaBody ts body +typeCheckKernel (SegGenRed space ops ts body) = do + checkSpace space + mapM_ TC.checkType ts + + TC.binding (scopeOfKernelSpace space) $ do + forM_ ops $ \(GenReduceOp dest_w dests nes shape op) -> do + TC.require [Prim int32] dest_w + nes' <- mapM TC.checkArg nes + mapM_ (TC.require [Prim int32]) $ shapeDims shape + + -- Operator type must match the type of neutral elements. + let stripVecDims = stripArray $ shapeRank shape + TC.checkLambda op $ map (TC.noArgAliases .first stripVecDims) $ nes' ++ nes' + let nes_t = map TC.argType nes' + unless (nes_t == map (`arrayOfShape` shape) (lambdaReturnType op)) $ + TC.bad $ TC.TypeError $ "SegGenRed operator has return type " ++ + prettyTuple (lambdaReturnType op) ++ " but neutral element has type " ++ + prettyTuple nes_t + + -- Arrays must have proper type. + let dest_shape = Shape $ segment_dims <> [dest_w] + forM_ (zip nes_t dests) $ \(t, dest) -> do + TC.requireI [t `arrayOfShape` dest_shape] dest + TC.consume =<< TC.lookupAliases dest + + TC.checkLambdaBody ts body + + -- 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 == ts) $ + TC.bad $ TC.TypeError $ "SegGenRed body has return type " ++ + prettyTuple ts ++ " but should have type " ++ + prettyTuple bucket_ret_t + + where segment_dims = init $ map snd $ spaceDimensions space + typeCheckKernel (Kernel _ space kts kbody) = do checkSpace space mapM_ TC.checkType kts @@ -664,6 +744,9 @@ kernelBodyMetrics = mapM_ bindingMetrics . kernelBodyStms opMetrics (SegRed _ _ red_op _ _ body) = inside "SegRed" $ lambdaMetrics red_op >> bodyMetrics body + opMetrics (SegGenRed _ ops _ body) = + inside "SegGenRed" $ do mapM_ (lambdaMetrics . genReduceOp) ops + bodyMetrics body opMetrics GetSize{} = seen "GetSize" opMetrics GetSizeMax{} = seen "GetSizeMax" opMetrics CmpSizeLe{} = seen "CmpSizeLe" @@ -691,6 +774,18 @@ PP.nestedBlock "{" "}" (ppr body) where name = case comm of Commutative -> "segred_comm" Noncommutative -> "segred" + + ppr (SegGenRed space ops ts body) = + text "seggenred" <> + PP.parens (PP.braces (mconcat $ intersperse (PP.comma <> PP.line) $ map ppOp ops)) + PP.align (ppr space) <+> PP.colon <+> ppTuple' ts <+> + PP.nestedBlock "{" "}" (ppr body) + where ppOp (GenReduceOp w dests nes shape op) = + ppr w <> PP.comma + PP.braces (PP.commasep $ map ppr dests) <> PP.comma + PP.braces (PP.commasep $ map ppr nes) <> PP.comma + ppr shape <> PP.comma + ppr op instance Pretty KernelSpace where ppr (KernelSpace f_gtid f_ltid gid num_threads num_groups group_size structure) = diff --git a/src/Futhark/Representation/Kernels/Simplify.hs b/src/Futhark/Representation/Kernels/Simplify.hs --- a/src/Futhark/Representation/Kernels/Simplify.hs +++ b/src/Futhark/Representation/Kernels/Simplify.hs @@ -92,8 +92,64 @@ Engine.simplifyLambda red_op $ replicate (length nes * 2) Nothing red_op_hoisted' <- mapM processHoistedStm red_op_hoisted + (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env ts body + + return (SegRed space' comm red_op' nes' ts' body', + red_op_hoisted' <> body_hoisted) + + where scope_vtable = ST.fromScope scope + scope = scopeOfKernelSpace space + +simplifyKernelOp mk_ops env (SegGenRed space ops ts body) = do + outer_vtable <- Engine.askVtable + + space' <- Engine.simplify space + ts' <- mapM Engine.simplify ts + + (ops', ops_hoisted) <- fmap unzip $ forM ops $ + \(GenReduceOp w arrs nes dims lam) -> do + w' <- Engine.simplify w + arrs' <- Engine.simplify arrs + nes' <- Engine.simplify nes + dims' <- Engine.simplify dims + (lam', op_hoisted) <- + Engine.subSimpleM (mk_ops space) env outer_vtable $ + Engine.localVtable (<>scope_vtable) $ + Engine.simplifyLambda lam $ + replicate (length nes * 2) Nothing + return (GenReduceOp w' arrs' nes' dims' lam', + op_hoisted) + + red_op_hoisted' <- mapM processHoistedStm $ mconcat ops_hoisted + + (body', body_hoisted) <- hoistFromBody space' (mk_ops space') env ts body + + return (SegGenRed space' ops' ts' body', + red_op_hoisted' <> body_hoisted) + + where scope_vtable = ST.fromScope scope + scope = scopeOfKernelSpace space + +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) + +hoistFromBody :: (Engine.SimplifiableLore lore, + SameScope lore outerlore, + BodyAttr outerlore ~ (), BodyAttr lore ~ (), + ExpAttr lore ~ ExpAttr outerlore, + RetType lore ~ RetType outerlore, + BranchType lore ~ BranchType outerlore) => + KernelSpace -> Simplify.SimpleOps lore -> Engine.Env lore + -> [Type] -> Body lore + -> Engine.SimpleM outerlore (Body (Wise lore), Stms (Wise outerlore)) +hoistFromBody kspace ops env ts body = do + outer_vtable <- Engine.askVtable + ((body_stms, body_res), body_hoisted) <- - Engine.subSimpleM (mk_ops space) env outer_vtable $ do + Engine.subSimpleM ops env outer_vtable $ do par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers Engine.localVtable (<>scope_vtable) $ Engine.blockIf (Engine.hasFree bound_here @@ -101,21 +157,15 @@ `Engine.orIf` par_blocker `Engine.orIf` Engine.isConsumed) $ Engine.simplifyBody (replicate (length ts) Observe) body + body_hoisted' <- mapM processHoistedStm body_hoisted - return (SegRed space' comm red_op' nes' ts' $ - mkWiseBody () body_stms body_res, - red_op_hoisted' <> body_hoisted') + return (mkWiseBody () body_stms body_res, + body_hoisted') where scope_vtable = ST.fromScope scope - scope = scopeOfKernelSpace space + scope = scopeOfKernelSpace kspace 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, diff --git a/src/Futhark/Test.hs b/src/Futhark/Test.hs --- a/src/Futhark/Test.hs +++ b/src/Futhark/Test.hs @@ -17,6 +17,7 @@ , compileProgram , runProgram , ensureReferenceOutput + , determineTuning , Mismatch , ProgramTest (..) @@ -184,6 +185,10 @@ lexstr :: T.Text -> Parser () lexstr = void . try . lexeme . string +-- | Like 'lexstr', but does not consume trailing linebreaks. +lexstr' :: T.Text -> Parser () +lexstr' = void . try . lexeme' . string + braces :: Parser a -> Parser a braces p = lexstr "{" *> p <* lexstr "}" @@ -210,7 +215,7 @@ tagConstituent c = isAlphaNum c || c == '_' || c == '-' parseAction :: Parser TestAction -parseAction = CompileTimeFailure <$> (lexstr "error:" *> parseExpectedError) <|> +parseAction = CompileTimeFailure <$> (lexstr' "error:" *> parseExpectedError) <|> (RunCases <$> parseInputOutputs <*> many parseExpectedStructure <*> many parseWarning) @@ -232,7 +237,7 @@ return s parseRunCases :: Parser [TestRun] -parseRunCases = parseRunCases' (0::Int) +parseRunCases = parseRunCases' (1::Int) where parseRunCases' i = (:) <$> parseRunCase i <*> parseRunCases' (i+1) <|> pure [] parseRunCase i = do @@ -509,12 +514,12 @@ genValues :: [GenValue] -> IO SBS.ByteString genValues gens = do - (code, stdout, stderr) <- readProcessWithExitCode "futhark-dataset" args mempty + (code, stdout, stderr) <- readProcessWithExitCode "futhark" ("dataset":args) mempty case code of ExitSuccess -> return stdout ExitFailure e -> - fail $ "futhark-dataset failed with exit code " ++ show e ++ " and stderr:\n" ++ + fail $ "'futhark dataset' failed with exit code " ++ show e ++ " and stderr:\n" ++ map (chr . fromIntegral) (SBS.unpack stderr) where args = "-b" : concatMap argForGen gens argForGen g = ["-g", genValueType g] @@ -613,8 +618,10 @@ fail $ "Reference dataset generation failed with exit code " ++ show e ++ " and stderr:\n" ++ map (chr . fromIntegral) (SBS.unpack stderr) - ExitSuccess -> - SBS.writeFile (file (entry, tr)) stdout + ExitSuccess -> do + let f = file (entry, tr) + liftIO $ createDirectoryIfMissing True $ takeDirectory f + SBS.writeFile f stdout where file (entry, tr) = takeDirectory prog testRunReferenceOutput prog entry tr @@ -625,3 +632,15 @@ liftIO . fmap not . doesFileExist . file $ (entry, tr) | otherwise = return False + +-- | Determine the --tuning options to pass to the program. The first +-- argument is the extension of the tuning file, or 'Nothing' if none +-- should be used. +determineTuning :: MonadIO m => Maybe FilePath -> FilePath -> m ([String], String) +determineTuning Nothing _ = return ([], mempty) +determineTuning (Just ext) program = do + exists <- liftIO $ doesFileExist (program <.> ext) + if exists + then return (["--tuning", program <.> ext], + " (using " <> takeFileName (program <.> ext) <> ")") + else return ([], mempty) diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs --- a/src/Futhark/Transform/FirstOrderTransform.hs +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -25,6 +25,7 @@ import Control.Monad.State import qualified Data.Map.Strict as M import qualified Data.Set as S +import Data.List (zip4) import qualified Futhark.Representation.AST as AST import Futhark.Representation.SOACS @@ -221,22 +222,34 @@ 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] + hists_out'' <- forM (zip4 hists_out' ops inds vals) $ \(hist, op, idx, val) -> do + -- Check whether the indexes are in-bound. If they are not, we + -- return the histograms unchanged. + let outside_bounds_branch = insertStmsM $ resultBodyM $ map Var hist + oob = case hist of [] -> eSubExp $ constant True + arr:_ -> eOutOfBounds arr [eSubExp 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 + letTupExp "new_histo" <=< + eIf oob outside_bounds_branch $ do + -- Read values from histogram. + h_val <- forM hist $ \arr -> do + arr_t <- lookupType arr + letSubExp "read_hist" $ BasicOp $ Index arr $ fullSlice arr_t [DimFix idx] - -- 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) + -- Apply operator. + h_val' <- bindLambda (genReduceOp op) $ + map (BasicOp . SubExp) $ h_val ++ val - return $ resultBody $ map Var $ concat ress + -- Write values back to histograms. + hist' <- forM (zip hist h_val') $ \(arr, v) -> do + arr_t <- lookupType arr + letInPlace "hist_out" arr (fullSlice arr_t [DimFix idx]) $ + BasicOp $ SubExp v + + return $ resultBody $ map Var hist' + + return $ resultBody $ map Var $ concat hists_out'' + -- Wrap up the above into a for-loop. letBind_ pat $ DoLoop [] merge (ForLoop iter Int32 len []) loopBody diff --git a/src/Futhark/Version.hs b/src/Futhark/Version.hs --- a/src/Futhark/Version.hs +++ b/src/Futhark/Version.hs @@ -20,11 +20,16 @@ -- | The version of Futhark that we are using, as a 'String' versionString :: String -versionString = showVersion version ++ "\n" ++ gitversion +versionString = showVersion version ++ + if used_hash /= "UNKNOWN" + then "\n" ++ gitversion + else "" where + used_hash = take 7 $(gitHash) + gitversion = concat ["git: " , branch - , take 7 $(gitHash) + , used_hash , " (", $(gitCommitDate), ")" , dirty ] diff --git a/src/Language/Futhark/Attributes.hs b/src/Language/Futhark/Attributes.hs --- a/src/Language/Futhark/Attributes.hs +++ b/src/Language/Futhark/Attributes.hs @@ -456,9 +456,6 @@ 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 diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -205,24 +205,30 @@ data Env = Env { envTerm :: M.Map VName TermBinding , envType :: M.Map VName T.TypeBinding + , envShapes :: M.Map VName Shape + -- ^ A mapping from type parameters to the shapes of + -- the value to which they were initially bound. } instance Monoid Env where - mempty = Env mempty mempty + mempty = Env mempty mempty mempty instance Semigroup Env where - Env vm1 tm1 <> Env vm2 tm2 = Env (vm1 <> vm2) (tm1 <> tm2) + Env vm1 tm1 sm1 <> Env vm2 tm2 sm2 = + Env (vm1 <> vm2) (tm1 <> tm2) (sm1 <> sm2) newtype InterpreterError = InterpreterError String valEnv :: M.Map VName (Maybe T.BoundV, Value) -> Env valEnv m = Env { envTerm = M.map (uncurry TermValue) m , envType = mempty + , envShapes = mempty } modEnv :: M.Map VName Module -> Env modEnv m = Env { envTerm = M.map TermModule m , envType = mempty + , envShapes = mempty } instance Show InterpreterError where @@ -287,90 +293,92 @@ 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 -> Pattern -> Value -> EvalM Env matchPattern env p v = do - m <- runMaybeT $ patternMatch env mempty p v + m <- runMaybeT $ patternMatch env p v case m of Nothing -> error $ "matchPattern: missing case for " ++ pretty p ++ " and " ++ pretty v - Just binds -> return binds + Just env' -> return env' -patternMatch :: Env -> M.Map VName (Maybe T.BoundV, Value) - -> Pattern -> Value - -> MaybeT EvalM (M.Map VName (Maybe T.BoundV, Value)) -patternMatch _ m (Id v (Info t) _) val = - lift $ pure $ M.insert v (Just $ T.BoundV [] $ toStruct t, val) m -patternMatch _ m Wildcard{} _ = - lift $ pure m -patternMatch env m (TuplePattern ps _) (ValueRecord vs) +patternMatch :: Env -> Pattern -> Value -> MaybeT EvalM Env +patternMatch env (Id v (Info t) _) val = + lift $ pure $ valEnv (M.singleton v (Just $ T.BoundV [] $ toStruct t, val)) <> env +patternMatch env Wildcard{} _ = + lift $ pure env +patternMatch env (TuplePattern ps _) (ValueRecord vs) | length ps == length vs' = - foldM (\m' (p,v) -> patternMatch env m' p v) m $ + foldM (\env' (p,v) -> patternMatch env' p v) env $ zip ps (map snd $ sortFields vs) where vs' = sortFields vs -patternMatch env m (RecordPattern ps _) (ValueRecord vs) +patternMatch env (RecordPattern ps _) (ValueRecord vs) | length ps == length vs' = - foldM (\m' (p,v) -> patternMatch env m' p v) m $ + foldM (\env' (p,v) -> patternMatch env' p v) env $ zip (map snd $ sortFields $ M.fromList ps) (map snd $ sortFields vs) where vs' = sortFields vs -patternMatch env m (PatternParens p _) v = patternMatch env m p v -patternMatch env m (PatternAscription p td loc) v = do - t <- lift $ evalType env $ unInfo $ expandedType td - case matchValueToType env m t v of +patternMatch env (PatternParens p _) v = patternMatch env p v +patternMatch env (PatternAscription p td loc) v = do + let t = evalType env $ unInfo $ expandedType td + case matchValueToType env t v of Left err -> lift $ bad loc env err - Right m' -> patternMatch env m' p v -patternMatch env m (PatternLit e _ _) v = do + Right env' -> patternMatch env' p v +patternMatch env (PatternLit e _ _) v = do v' <- lift $ eval env e if v == v' - then pure m + then pure env else mzero -patternMatch _ _ _ _ = mzero +patternMatch _ _ _ = mzero -- | 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) +matchValueToType :: Env -> StructType -> Value - -> Either String (M.Map VName (Maybe T.BoundV, Value)) + -> Either String Env -matchValueToType env m t@(Array _ _ _ (ShapeDecl ds@(d:_))) val@(ValueArray arr) = +matchValueToType env t@(TypeVar _ _ tn []) val + | Just shape <- M.lookup (typeLeaf tn) $ envShapes env, + shape /= valueShape val = + Left $ "Value passed for type parameter `" <> prettyName (typeLeaf tn) <> + "` does not match shape " <> pretty shape <> + " of previously observed value." + | Nothing <- M.lookup (typeLeaf tn) $ envShapes env = + matchValueToType (tnenv <> env) t val + where tnenv = Env mempty mempty $ M.singleton (typeLeaf tn) (valueShape val) + +matchValueToType env t@(Array _ _ _ (ShapeDecl ds@(d:_))) val@(ValueArray arr) = case d of NamedDim v | Just x <- look v -> if x == arr_n - then continue m + then continue env else emptyOrWrong $ "`" <> 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 + continue $ + valEnv (M.singleton (qualLeaf v) + (Just $ T.BoundV [] $ Prim $ Signed Int32, + ValuePrim $ SignedValue $ Int32Value arr_n)) + <> env + AnyDim -> continue env ConstDim x - | fromIntegral x == arr_n -> continue m + | fromIntegral x == arr_n -> continue env | otherwise -> emptyOrWrong $ 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 - [] -> - -- We have to ensure that remaining unbound shape - -- parameters become zeroes. - return $ m' <> mconcat (map namedAreZero ds) - v:_ -> - matchValueToType env m' (stripArray 1 t) v + continue env' = case elems arr of + [] -> return env' + v:_ -> matchValueToType env' (stripArray 1 t) v -- Empty arrays always match if nothing else does. emptyOrWrong x | any zeroDim ds, emptyShape (valueShape val) = - Right $ m <> mconcat (map namedAreZero ds) + Right env | otherwise = wrong x wrong x = Left $ "Size annotation " <> x <> @@ -380,20 +388,17 @@ 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 $ +matchValueToType env (Record fs) (ValueRecord arr) = + foldM (\env' (t, v) -> matchValueToType env' t v) env $ M.intersectionWith (,) fs arr -matchValueToType _ m _ _ = return m +matchValueToType env _ _ = return env +bindToZero :: [VName] -> Env +bindToZero = valEnv . M.fromList . map f + where f v = (v, (Just $ T.BoundV [] $ Prim $ Signed Int32, + ValuePrim $ SignedValue $ Int32Value 0)) + data Indexing = IndexingFix Int32 | IndexingSlice (Maybe Int32) (Maybe Int32) (Maybe Int32) @@ -442,7 +447,8 @@ -- | 'signum', but with 0 as 1. signum' :: (Eq p, Num p) => p -> p -signum' x = if x == 0 then 1 else signum x +signum' 0 = 1 +signum' x = signum x indexArray :: [Indexing] -> Value -> Maybe Value indexArray (IndexingFix i:is) (ValueArray arr) @@ -497,41 +503,41 @@ -- | 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 -> StructType -> StructType +evalType _ (Prim pt) = Prim pt +evalType env (Record fs) = Record $ fmap (evalType env) fs evalType env (Arrow () p t1 t2) = - Arrow () p <$> evalType env t1 <*> evalType env t2 -evalType env t@(Array _ u _ shape) = do + Arrow () p (evalType env t1) (evalType env t2) +evalType env t@(Array _ u _ shape) = 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 + et' = evalType env et + shape' = fmap evalDim shape + in 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 + ConstDim $ fromIntegral x + evalDim d = 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 + Just (T.TypeAbbr _ ps t') -> + let (substs, types) = mconcat $ zipWith matchPtoA ps args + 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 + in if null ps then bimap onDim id t' + else evalType (Env mempty types mempty <> env) $ bimap onDim id t' + Nothing -> t + where matchPtoA (TypeParamDim p _) (TypeArgDim (NamedDim qv) _) = - return (M.singleton p $ NamedDim qv, mempty) + (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 -evalType _ (Enum cs) = return $ Enum cs + (M.singleton p $ ConstDim k, mempty) + matchPtoA (TypeParamType l p _) (TypeArgType t' _) = + let t'' = evalType env t' + in (mempty, M.singleton p $ T.TypeAbbr l [] t'') + matchPtoA _ _ = mempty +evalType _ (Enum cs) = Enum cs eval :: Env -> Exp -> EvalM Value @@ -580,16 +586,16 @@ eval env (Ascript e td loc) = do v <- eval env e - t <- evalType env $ unInfo $ expandedType td - case matchValueToType env mempty t v of + let t = evalType env $ unInfo $ expandedType td + case matchValueToType env t v of Right _ -> return v Left err -> bad loc env $ "Value `" <> pretty v <> "` cannot match shape of type `" <> pretty (declaredType td) <> "` (`" <> pretty t <> "`): " ++ err eval env (LetPat _ p e body _) = do v <- eval env e - p_env <- valEnv <$> matchPattern env p v - eval (p_env <> env) body + env' <- matchPattern env p v + eval 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 @@ -684,25 +690,31 @@ -- 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 +eval env (Lambda tparams [] body _ (Info (_, t)) loc) = do + -- All remaining size parameters that have not yet been assigned a + -- value (because they were inner dimensions of empty arrays) are + -- now assigned a zero. + let unbound_dims = bindToZero $ map typeParamName $ filter isDimParam tparams + v <- eval (env <> unbound_dims) body case (t, v) of (Arrow _ _ _ rt, ValueFun f) -> return $ ValueFun $ \arg -> do r <- f arg - rt' <- evalType env rt - match rt' r + match (evalType env rt) r _ -> match t v where match vt v = - case matchValueToType env mempty vt v of + case matchValueToType env vt v of Right _ -> return v Left err -> bad loc env $ "Value `" <> pretty v <> "` cannot match type `" <> pretty vt <> "`: " ++ err + isDimParam TypeParamDim{} = True + isDimParam _ = False + 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 + env' <- matchPattern env p v + eval env' $ Lambda tparams ps body mrd (Info (als, ret)) loc eval env (OpSection qv _ _) = evalTermVar env qv @@ -733,7 +745,7 @@ foldM (forInLoop in_pat) init_v in_vs While cond -> whileLoop cond init_v - where withLoopParams v = (<>env) . valEnv <$> matchPattern env pat v + where withLoopParams = matchPattern env pat inc = (`P.doAdd` Int64Value 1) zero = (`P.doMul` Int64Value 0) @@ -755,8 +767,8 @@ forInLoop in_pat v in_v = do env' <- withLoopParams v - pat_env <- matchPattern env' in_pat in_v - eval (valEnv pat_env <> env') body + env'' <- matchPattern env' in_pat in_v + eval env'' body eval env (Project f e _ _) = do v <- eval env e @@ -789,8 +801,8 @@ evalCase :: Value -> Env -> CaseBase Info VName -> EvalM (Maybe Value) evalCase v env (CasePat p cExp _) = runMaybeT $ do - pEnv <- valEnv <$> patternMatch env mempty p v - lift $ eval (pEnv <> env) cExp + env' <- patternMatch env p v + lift $ eval env' cExp substituteInModule :: M.Map VName VName -> Module -> Module substituteInModule substs = onModule @@ -801,8 +813,8 @@ 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 (Module (Env terms types _)) = + Module $ Env (replaceM onTerm terms) (replaceM onType types) mempty onModule (ModuleFun f) = ModuleFun $ \m -> onModule <$> f (substituteInModule rev_substs m) onTerm (TermValue t v) = TermValue t v @@ -829,10 +841,11 @@ Just m -> return $ Module m evalModExp env (ModDecs ds _) = do - Env terms types <- foldM evalDec env ds + 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) + mempty evalModExp env (ModVar qv _) = evalModuleVar env qv @@ -857,8 +870,8 @@ 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 + let t' = evalType env t + 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 @@ -872,8 +885,8 @@ 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' + let abbr = T.TypeAbbr Lifted ps $ + evalType env $ unInfo $ expandedType 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 @@ -892,8 +905,8 @@ initialCtx :: Ctx initialCtx = Ctx (Env (M.insert (VName (nameFromString "intrinsics") 0) - (TermModule (Module $ Env terms types)) terms) - types) + (TermModule (Module $ Env terms types mempty)) terms) + types mempty) mempty where terms = M.mapMaybeWithKey (const . def . baseString) intrinsics diff --git a/src/Language/Futhark/Parser/Parser.y b/src/Language/Futhark/Parser/Parser.y --- a/src/Language/Futhark/Parser/Parser.y +++ b/src/Language/Futhark/Parser/Parser.y @@ -25,6 +25,7 @@ import Control.Arrow import Data.Array import qualified Data.Text as T +import Codec.Binary.UTF8.String (encode) import Data.Char (ord) import Data.Maybe (fromMaybe, fromJust) import Data.Loc hiding (L) -- Lexer has replacements. @@ -598,7 +599,7 @@ | 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 } + ArrayLit (map (flip Literal loc . UnsignedValue . Int8Value . fromIntegral) $ encode s) NoInfo loc } | '(' Exp ')' FieldAccesses { foldl (\x (y, _) -> Project y x NoInfo (srclocOf x)) (Parens $2 (srcspan $1 $3)) @@ -763,7 +764,7 @@ | intlit { let L loc (INTLIT x) = $1 in (IntLit x NoInfo loc, loc) } | floatlit { let L loc (FLOATLIT x) = $1 in (FloatLit x NoInfo loc, loc) } | stringlit { let L loc (STRINGLIT s) = $1 in - (ArrayLit (map (flip Literal loc . SignedValue . Int32Value . fromIntegral . ord) s) NoInfo loc, loc) } + (ArrayLit (map (flip Literal loc . UnsignedValue . Int8Value . fromIntegral) $ encode s) NoInfo loc, loc) } | VConstr0 { (VConstr0 (fst $1) NoInfo (snd $1), snd $1) } LoopForm :: { LoopFormBase NoInfo Name } @@ -874,7 +875,7 @@ StringValue :: { Value } StringValue : stringlit { let L pos (STRINGLIT s) = $1 in - ArrayValue (arrayFromList $ map (PrimValue . SignedValue . Int32Value . fromIntegral . ord) s) $ Prim $ Signed Int32 } + ArrayValue (arrayFromList $ map (PrimValue . UnsignedValue . Int8Value . fromIntegral) $ encode s) $ Prim $ Signed Int32 } BoolValue :: { Value } BoolValue : true { PrimValue $ BoolValue True } @@ -898,15 +899,13 @@ : 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 } + case qn of + QualName ["f32"] "inf" -> return (Float32Value (1/0), loc) + QualName ["f32"] "nan" -> return (Float32Value (0/0), loc) + QualName ["f64"] "inf" -> return (Float64Value (1/0), loc) + QualName ["f64"] "nan" -> return (Float64Value (0/0), loc) + _ -> parseErrorAt (snd $1) Nothing + } | floatlit { let L loc (FLOATLIT num) = $1 in (Float64Value num, loc) } ArrayValue :: { Value } diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs --- a/src/Language/Futhark/Pretty.hs +++ b/src/Language/Futhark/Pretty.hs @@ -237,7 +237,7 @@ then equals indent 2 (ppr e) else equals <+> align (ppr e)) (case body of LetPat{} -> ppr body - _ -> text "in" <+> ppr body) + _ -> text "in" <+> align (ppr body)) where linebreak = case e of Map{} -> True Reduce{} -> True @@ -297,9 +297,6 @@ 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 _ _) = diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs --- a/src/Language/Futhark/Syntax.hs +++ b/src/Language/Futhark/Syntax.hs @@ -711,14 +711,6 @@ -- 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 @@ -768,8 +760,6 @@ 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 diff --git a/src/Language/Futhark/Traversals.hs b/src/Language/Futhark/Traversals.hs --- a/src/Language/Futhark/Traversals.hs +++ b/src/Language/Futhark/Traversals.hs @@ -125,11 +125,6 @@ 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) = diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs --- a/src/Language/Futhark/TypeChecker.hs +++ b/src/Language/Futhark/TypeChecker.hs @@ -168,6 +168,12 @@ mempty { envTypeTable = M.singleton v $ TypeAbbr l [] $ TypeVar () Nonunique (typeName v) [] } +-- In this function, after the recursion, we add the Env of the +-- current Spec *after* the one that is returned from the recursive +-- call. This implements the behaviour that specs later in a module +-- type can override those earlier (it rarely matters, but it affects +-- the specific structure of substitutions in case some module type is +-- redundantly imported multiple times). checkSpecs :: [SpecBase NoInfo Name] -> TypeM (TySet, Env, [SpecBase Info VName]) checkSpecs [] = return (mempty, mempty, []) @@ -195,7 +201,7 @@ (tenv, tdec') <- checkTypeBind tdec (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs return (abstypes, - tenv <> env, + env <> tenv, TypeAbbrSpec tdec' : specs') checkSpecs (TypeSpec l name ps doc loc : specs) = @@ -211,7 +217,7 @@ } (abstypes, env, specs') <- localEnv tenv $ checkSpecs specs return (M.insert (qualName name') l abstypes, - tenv <> env, + env <> tenv, TypeSpec l name' ps' doc loc : specs') checkSpecs (ModSpec name sig doc loc : specs) = @@ -223,7 +229,7 @@ } (abstypes, env, specs') <- localEnv senv $ checkSpecs specs return (M.mapKeys (qualify name') (mtyAbs mty) <> abstypes, - senv <> env, + env <> senv, ModSpec name' sig' doc loc : specs') checkSpecs (IncludeSpec e loc : specs) = do @@ -232,8 +238,8 @@ mapM_ (warnIfShadowing . fmap baseName) $ M.keys e_abs (abstypes, env, specs') <- localEnv e_env $ checkSpecs specs - return (e_abs <> abstypes, - e_env <> env, + return (abstypes <> e_abs, + env <> e_env, IncludeSpec e' loc : specs') where warnIfShadowing qn = (lookupType loc qn >> warnAbout qn) diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs --- a/src/Language/Futhark/TypeChecker/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Monad.hs @@ -6,6 +6,7 @@ , runTypeM , askEnv , askRootEnv + , askImportName , localTmpEnv , checkQualNameWithEnv , bindSpaced @@ -154,6 +155,10 @@ askEnv = asks contextEnv askRootEnv = asks contextRootEnv +-- | The name of the current file/import. +askImportName :: TypeM ImportName +askImportName = asks contextImportName + localTmpEnv :: Env -> TypeM a -> TypeM a localTmpEnv env = local $ \ctx -> ctx { contextEnv = env <> contextEnv ctx } @@ -311,7 +316,7 @@ -- 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 +qualifyTypeVars outer_env except ref_qs = runIdentity . astMap mapper where mapper = ASTMapper { mapOnExp = pure , mapOnName = pure , mapOnQualName = pure . qual @@ -321,9 +326,16 @@ , 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 + | name `elem` except || reachable orig_qs name outer_env = + QualName orig_qs name + | otherwise = + prependAsNecessary [] ref_qs $ QualName orig_qs name + + prependAsNecessary qs rem_qs (QualName orig_qs name) + | reachable (qs++orig_qs) name outer_env = QualName (qs++orig_qs) name + | otherwise = case rem_qs of + q:rem_qs' -> prependAsNecessary (qs++[q]) rem_qs' (QualName orig_qs name) + [] -> QualName (qs++orig_qs) name reachable [] name env = isJust $ find matches $ M.elems (envTypeTable env) diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -28,6 +28,7 @@ import Prelude hiding (mod) import Language.Futhark +import Language.Futhark.Semantic (includeToString) import Language.Futhark.Traversals import Language.Futhark.TypeChecker.Monad hiding (BoundV, checkQualNameWithEnv) import Language.Futhark.TypeChecker.Types hiding (checkTypeDecl) @@ -347,6 +348,9 @@ checkIntrinsic :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) checkIntrinsic space qn@(QualName _ name) loc | Just v <- M.lookup (space, name) intrinsicsNameMap = do + me <- liftTypeM askImportName + unless ("/futlib" `isPrefixOf` includeToString me) $ + warn loc "Using intrinsic functions directly can easily crash the compiler or result in wrong code generation." scope <- ask return (scope, v) | otherwise = @@ -420,14 +424,14 @@ 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." + "Unique return value of function " ++ quote (pretty fname) ++ + " is aliased to " ++ quote (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." + "A unique tuple element of return value of `" ++ + quote (pretty fname) ++ "` is aliased to some other tuple component." --- Basic checking @@ -955,7 +959,7 @@ return $ LetFun name' (tparams', params', maybe_retdecl', Info rettype, e') body' loc -checkExp (LetWith dest src idxes ve body pos) = do +checkExp (LetWith dest src idxes ve body loc) = do (t, _) <- newArrayType (srclocOf src) "src" $ length idxes let elemt = stripArray (length $ filter isFix idxes) t sequentially (checkIdent src) $ \src' _ -> do @@ -965,18 +969,26 @@ void $ unifies t src'' unless (unique $ unInfo $ identType src') $ - typeError pos $ "Source " ++ quote (pretty (identName src)) ++ - " has type " ++ pretty (unInfo $ identType src') ++ ", which is not unique" + typeError loc $ "Source " ++ quote (pretty (identName src)) ++ + " has type " ++ pretty (unInfo $ identType src') ++ ", which is not unique." + vtable <- asks scopeVtable + forM_ (aliases $ unInfo $ identType src') $ \v -> + case aliasVar v `M.lookup` vtable of + Just (BoundV Local _ v_t) + | not $ unique v_t -> + typeError loc $ "Source " ++ quote (pretty (identName src)) ++ + " aliases " ++ quote (prettyName (aliasVar v)) ++ ", which is not consumable." + _ -> return () idxes' <- mapM checkDimIndex idxes sequentially (unifies elemt =<< checkExp ve) $ \ve' _ -> do ve_t <- expType ve' when (AliasBound (identName src') `S.member` aliases ve_t) $ - badLetWithValue pos + badLetWithValue loc bindingIdent dest (unInfo (identType src') `setAliases` S.empty) $ \dest' -> do body' <- consuming src' $ checkExp body - return $ LetWith dest' src' idxes' ve' body' pos + return $ LetWith dest' src' idxes' ve' body' loc where isFix DimFix{} = True isFix _ = False @@ -1020,44 +1032,6 @@ 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 - - e_ts <- mapM expType $ e':es' - ts <- forM (zip (e':es') e_ts) $ \(arr_e, arr_e_t) -> - case typeToRecordArrayElem =<< 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 (mconcat $ map aliases e_ts) u - (ArrayRecordElem $ M.fromList $ zip tupleFieldNames ts) - (rank (1+i)) - 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 _ u (ArrayRecordElem fs) shape - | 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 mempty u et' shape - RecordArrayArrayElem et' et_shape -> - Array mempty u et' (shape <> et_shape) - checkExp (Unsafe e loc) = Unsafe <$> checkExp e <*> pure loc @@ -1080,14 +1054,15 @@ 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, inferReturnUniqueness params' body_t) + body_t <- expType body' + let (maybe_retdecl'', rettype) = + case maybe_retdecl' of + Just retdecl'@(TypeDecl _ (Info st)) -> (Just retdecl', st) + Nothing -> (Nothing, inferReturnUniqueness params' body_t) - closure' <- lexicalClosure params' closure + checkGlobalAliases params' body_t loc + closure' <- lexicalClosure params' closure return $ Lambda tparams' params' body' maybe_retdecl'' (Info (closure', rettype)) loc checkExp (OpSection op _ loc) = do @@ -1253,19 +1228,25 @@ | unique pat_v_t, v:_ <- S.toList $ S.map aliasVar (aliases t) `S.intersection` bound_outside = lift $ typeError loc $ "Loop return value corresponding to merge parameter " ++ - prettyName pat_v ++ " aliases " ++ prettyName v ++ "." + quote (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." + quote (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) + quote (prettyName pat_v) ++ " aliases previously returned value." if unique pat_v_t then put (cons<>aliases t, obs) else put (cons, obs<>aliases t) + checkMergeReturn (PatternParens p _) t = + checkMergeReturn p t + checkMergeReturn (PatternAscription p _ _) t = + checkMergeReturn p t + checkMergeReturn (RecordPattern pfs _) (Record tfs) = + sequence_ $ M.elems $ M.intersectionWith checkMergeReturn (M.fromList pfs) tfs checkMergeReturn (TuplePattern pats _) t | Just ts <- isTupleRecord t = zipWithM_ checkMergeReturn pats ts checkMergeReturn _ _ = @@ -1648,21 +1629,8 @@ bindSpaced [(Term, fname)] $ do fname' <- checkName Term fname loc - vtable <- asks scopeVtable - let isLocal v = case v `M.lookup` vtable of - Just (BoundV Local _ _) -> True - _ -> False - let als = filter (not . isLocal) $ S.toList $ - boundArrayAliases body_t `S.difference` - S.map identName (mconcat (map patIdentSet params')) - case als of - v:_ | not $ null params -> - typeError loc $ - unlines [ "Function result aliases the free variable " <> - quote (prettyName v) <> "." - , "Use " ++ quote "copy" ++ " to break the aliasing."] - _ -> - return (fname', tparams'', params'', maybe_retdecl'', rettype, body') + checkGlobalAliases params'' body_t loc + return (fname', tparams'', params'', maybe_retdecl'', rettype, body') where -- | Check that unique return values do not alias a -- non-consumed parameter. @@ -1695,6 +1663,25 @@ concat $ M.elems $ M.intersectionWith returnAliasing ets1 ets2 returnAliasing expected got = [(uniqueness expected, S.map aliasVar $ aliases got)] + +checkGlobalAliases :: [Pattern] -> CompType -> SrcLoc -> TermTypeM () +checkGlobalAliases params body_t loc = do + vtable <- asks scopeVtable + let isLocal v = case v `M.lookup` vtable of + Just (BoundV Local _ _) -> True + _ -> False + let als = filter (not . isLocal) $ S.toList $ + boundArrayAliases body_t `S.difference` + S.map identName (mconcat (map patIdentSet params)) + case als of + v:_ | not $ null params -> + typeError loc $ + unlines [ "Function result aliases the free variable " <> + quote (prettyName v) <> "." + , "Use " ++ quote "copy" ++ " to break the aliasing."] + _ -> + return () + inferReturnUniqueness :: [Pattern] -> CompType -> StructType inferReturnUniqueness params t = diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs --- a/src/Language/Futhark/TypeChecker/Types.hs +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -375,18 +375,18 @@ type TypeSubs = M.Map VName TypeSub -substituteTypes :: TypeSubs -> StructType -> StructType +substituteTypes :: Monoid als => TypeSubs -> TypeBase (DimDecl VName) als -> TypeBase (DimDecl VName) als substituteTypes substs ot = case ot of Array als u at shape -> maybe nope (`addAliases` (<>als)) $ arrayOf (substituteTypesInArrayElem at) (substituteInShape shape) u Prim t -> Prim t - TypeVar () u v targs + TypeVar als 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 + applyType ps (t `setAliases` mempty) (map substituteInTypeArg targs) + `setUniqueness` u `addAliases` (<>als) + | otherwise -> TypeVar als u v $ map substituteInTypeArg targs Record ts -> Record $ fmap (substituteTypes substs) ts Arrow als v t1 t2 -> @@ -399,9 +399,9 @@ substituteTypesInArrayElem (ArrayPolyElem v targs) | Just (TypeSub (TypeAbbr _ ps t)) <- M.lookup (qualLeaf (qualNameFromTypeName v)) substs = - applyType ps t (map substituteInTypeArg targs) + applyType ps (t `setAliases` mempty) (map substituteInTypeArg targs) | otherwise = - TypeVar () Nonunique v (map substituteInTypeArg targs) + TypeVar mempty Nonunique v (map substituteInTypeArg targs) substituteTypesInArrayElem (ArrayRecordElem ts) = Record ts' where ts' = fmap (substituteTypes substs . recordArrayElemToType) ts @@ -424,7 +424,8 @@ substituteTypesInBoundV substs (BoundV tps t) = BoundV tps (substituteTypes substs t) -applyType :: [TypeParam] -> StructType -> [StructTypeArg] -> StructType +applyType :: Monoid als => + [TypeParam] -> TypeBase (DimDecl VName) als -> [StructTypeArg] -> TypeBase (DimDecl VName) als applyType ps t args = substituteTypes substs t where substs = M.fromList $ zipWith mkSubst ps args diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -29,11 +29,9 @@ 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 Language.Futhark.TypeChecker.Monad hiding (BoundV) +import Language.Futhark.TypeChecker.Types import Futhark.Util.Pretty (Pretty) -- | Mapping from fresh type variables, instantiated from the type diff --git a/src/futhark.hs b/src/futhark.hs --- a/src/futhark.hs +++ b/src/futhark.hs @@ -63,6 +63,7 @@ , ("pkg", (Pkg.main, "Manage local packages.")) , ("check", (Misc.mainCheck, "Type check a program.")) + , ("imports", (Misc.mainImports, "Print all non-library imported Futhark files to standard out and exit.")) ] msg :: String diff --git a/src/futharki.hs b/src/futharki.hs deleted file mode 100644 --- a/src/futharki.hs +++ /dev/null @@ -1,19 +0,0 @@ -module Main (main) where - -import System.Environment -import System.Process -import System.IO -import System.Exit - -main :: IO () -main = do - prog <- getProgName - args <- getArgs - let suffix = case args of - [] -> "repl" - _ -> "run" - hPutStrLn stderr $ - "'" ++ prog ++ "' is deprecated. Use '" ++ - unwords ["futhark", suffix] ++ "' instead." - (_, _, _, h) <- createProcess $ proc "futhark" $ suffix:args - exitWith =<< waitForProcess h diff --git a/src/wrapper.hs b/src/wrapper.hs deleted file mode 100644 --- a/src/wrapper.hs +++ /dev/null @@ -1,29 +0,0 @@ --- Wrapper program that translates @futhark-foo@ to @futhark foo@; --- using whichever @futhark@ binary is in the user's search path. - -module Main (main) where - -import Data.Maybe - -import System.Environment -import System.Process -import System.IO -import System.Exit - -nameChanges :: [(String, String)] -nameChanges = [ ("py", "python") - , ("cs", "csharp") - ] - -main :: IO () -main = do - prog <- getProgName - args <- getArgs - let suffix = drop 1 $ dropWhile (/='-') prog - suffix' = fromMaybe suffix $ lookup suffix nameChanges - - hPutStrLn stderr $ - prog ++ ": this command is deprecated. Use '" ++ - unwords ["futhark", suffix'] ++ "' instead." - (_, _, _, h) <- createProcess $ proc "futhark" $ suffix':args - exitWith =<< waitForProcess h diff --git a/unittests/Futhark/Analysis/ScalExpTests.hs b/unittests/Futhark/Analysis/ScalExpTests.hs --- a/unittests/Futhark/Analysis/ScalExpTests.hs +++ b/unittests/Futhark/Analysis/ScalExpTests.hs @@ -3,7 +3,6 @@ module Futhark.Analysis.ScalExpTests ( tests , parseScalExp - , parseScalExp' ) where @@ -24,11 +23,8 @@ 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 +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 diff --git a/unittests/Futhark/Optimise/AlgSimplifyTests.hs b/unittests/Futhark/Optimise/AlgSimplifyTests.hs --- a/unittests/Futhark/Optimise/AlgSimplifyTests.hs +++ b/unittests/Futhark/Optimise/AlgSimplifyTests.hs @@ -9,7 +9,7 @@ import Futhark.Representation.AST import Futhark.Analysis.ScalExp -import Futhark.Analysis.ScalExpTests (parseScalExp') +import Futhark.Analysis.ScalExpTests (parseScalExp) import Futhark.Analysis.AlgSimplify tests :: TestTree @@ -33,7 +33,7 @@ ] where vars = declareVars [("x", int32)] simplify'' e = simplify' vars e [] - scalExp = parseScalExp' vars + scalExp = parseScalExp vars cfoldTest input expected = testCase ("constant-fold " ++ input) $ @@ -85,11 +85,11 @@ (lookupVarName name varinfo, (i, fixBound lower, fixBound upper))) fixBound "" = Nothing - fixBound s = Just $ parseScalExp' varinfo s + fixBound s = Just $ parseScalExp varinfo s simplify' :: VarInfo -> String -> RangesRep' -> ScalExp simplify' varinfo s r = simplify e r' - where e = parseScalExp' varinfo s + where e = parseScalExp varinfo s r' = instantiateRanges varinfo r mkSuffConds' :: VarInfo -> String -> RangesRep' -> [[ScalExp]] @@ -97,5 +97,5 @@ case mkSuffConds e r' of Left _ -> [[e]] Right sc -> sc - where e = simplify (parseScalExp' varinfo s) r' + where e = simplify (parseScalExp varinfo s) r' r' = instantiateRanges varinfo r diff --git a/unittests/Futhark/Representation/AST/AttributesTests.hs b/unittests/Futhark/Representation/AST/AttributesTests.hs --- a/unittests/Futhark/Representation/AST/AttributesTests.hs +++ b/unittests/Futhark/Representation/AST/AttributesTests.hs @@ -7,7 +7,6 @@ import Test.Tasty -import Futhark.Representation.AST.SyntaxTests () import qualified Futhark.Representation.AST.Attributes.ReshapeTests import qualified Futhark.Representation.AST.Attributes.RearrangeTests diff --git a/unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs b/unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/ExplicitMemory/IndexFunction/Alg.hs @@ -0,0 +1,178 @@ +-- | A simple index operation representation. Every operation corresponds to a +-- constructor. +module Futhark.Representation.ExplicitMemory.IndexFunction.Alg + ( IxFun(..) + , iota + , offsetIndex + , strideIndex + , permute + , rotate + , reshape + , slice + , rebase + , repeat + , shape + , rank + , index + ) +where + +import Data.Monoid ((<>)) +import Data.List hiding (repeat) + +import Prelude hiding (repeat, mod) + +import Futhark.Representation.AST.Syntax + (ShapeChange, DimChange(..), Slice, sliceDims, DimIndex(..), unitSlice) +import Futhark.Representation.AST.Attributes +import Futhark.Representation.AST.Pretty () +import Futhark.Util.IntegralExp +import Futhark.Util.Pretty + +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) + | OffsetIndex (IxFun num) num + | StrideIndex (IxFun num) num + | Rebase (IxFun num) (IxFun 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])) + ppr (OffsetIndex fun i) = + ppr fun <> text "->offset_index" <> parens (ppr i) + ppr (StrideIndex fun s) = + ppr fun <> text "->stride_index" <> parens (ppr s) + ppr (Rebase new_base fun) = + text "rebase(" <> ppr new_base <> text ", " <> ppr fun <> text ")" + + +iota :: Shape num -> IxFun num +iota = Direct + +offsetIndex :: IxFun num -> num -> IxFun num +offsetIndex = OffsetIndex + +strideIndex :: IxFun num -> num -> IxFun num +strideIndex = StrideIndex + +permute :: IxFun num -> Permutation -> IxFun num +permute = Permute + +rotate :: IxFun num -> Indices num -> IxFun num +rotate = Rotate + +repeat :: IxFun num -> [Shape num] -> Shape num -> IxFun num +repeat = Repeat + +slice :: IxFun num -> Slice num -> IxFun num +slice = Index + +rebase :: IxFun num -> IxFun num -> IxFun num +rebase = Rebase + +reshape :: IxFun num -> ShapeChange num -> IxFun num +reshape = Reshape + +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] +shape (OffsetIndex ixfun _) = + shape ixfun +shape (StrideIndex ixfun _) = + shape ixfun +shape (Rebase _ ixfun) = + shape ixfun + +rank :: IntegralExp num => + IxFun num -> Int +rank = length . shape + + +index :: (IntegralExp num, Eq 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 +index (OffsetIndex fun i) is element_size = + case shape fun of + d : ds -> + index (Index fun (DimSlice i (d-i) 1 : map (unitSlice 0) ds)) is element_size + [] -> error "index: OffsetIndex: underlying index function has rank zero" +index (StrideIndex fun s) is element_size = + case shape fun of + d : ds -> + index (Index fun (DimSlice 0 d s : map (unitSlice 0) ds)) is element_size + [] -> error "index: StrideIndex: underlying index function has rank zero" +index (Rebase new_base fun) is element_size = + let fun' = case fun of + Direct old_shape -> + if old_shape == shape new_base + then new_base + else reshape new_base $ map DimCoercion old_shape + Permute ixfun perm -> + permute (rebase new_base ixfun) perm + Rotate ixfun offsets -> + rotate (rebase new_base ixfun) offsets + Index ixfun iis -> + slice (rebase new_base ixfun) iis + Reshape ixfun new_shape -> + reshape (rebase new_base ixfun) new_shape + Repeat ixfun outer_shapes inner_shape -> + repeat (rebase new_base ixfun) outer_shapes inner_shape + StrideIndex ixfun i -> + strideIndex (rebase new_base ixfun) i + OffsetIndex ixfun s -> + offsetIndex (rebase new_base ixfun) s + r@Rebase{} -> + r + in index fun' is element_size diff --git a/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs b/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionTests.hs @@ -0,0 +1,371 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +module Futhark.Representation.ExplicitMemory.IndexFunctionTests + ( tests + ) +where + +import Prelude hiding (span, repeat) +import qualified Prelude as P +import qualified Data.List as DL + +import Test.Tasty +import Test.Tasty.HUnit + +import Futhark.Representation.AST.Syntax +import Futhark.Representation.AST.Syntax.Core() +import qualified Futhark.Util.Pretty as PR +import qualified Futhark.Util.IntegralExp as IE +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as IxFunLMAD +import qualified Futhark.Representation.ExplicitMemory.IndexFunction.Alg as IxFunAlg + +import qualified Futhark.Representation.ExplicitMemory.IndexFunctionWrapper as IxFunWrap +import Futhark.Representation.ExplicitMemory.IndexFunctionWrapper + + +instance IE.IntegralExp Int where + quot = P.quot + rem = P.rem + div = P.div + mod = P.mod + sgn = Just . P.signum + + fromInt8 = fromInteger . toInteger + fromInt16 = fromInteger . toInteger + fromInt32 = fromInteger . toInteger + fromInt64 = fromInteger . toInteger + +allPoints :: [Int] -> [[Int]] +allPoints dims = + let total = product dims + strides = drop 1 $ DL.reverse $ scanl (*) 1 $ DL.reverse dims + in map (unflatInd strides) [0..total-1] + where unflatInd :: [Int] -> Int -> [Int] + unflatInd strides x = fst $ + foldl (\(res, acc) span -> + (res ++ [acc `P.div` span], acc `P.mod` span)) + ([], x) strides + +compareIxFuns :: IxFunLMAD.IxFun Int -> IxFunAlg.IxFun Int -> Assertion +compareIxFuns ixfunLMAD ixfunAlg = + let lmadShape = IxFunLMAD.shape ixfunLMAD + algShape = IxFunAlg.shape ixfunAlg + points = allPoints lmadShape + resLMAD = map (\is -> IxFunLMAD.index ixfunLMAD is 4) points + resAlg = map (\is -> IxFunAlg.index ixfunAlg is 4) points + errorMessage = "lmad ixfun: " ++ PR.pretty ixfunLMAD ++ "\n" ++ + "alg ixfun: " ++ PR.pretty ixfunAlg ++ "\n" ++ + "lmad shape: " ++ show lmadShape ++ "\n" ++ + "alg shape: " ++ show algShape ++ "\n" ++ + "lmad points length: " ++ show (length resLMAD) ++ "\n" ++ + "alg points length: " ++ show (length resAlg) ++ "\n" ++ + "lmad points: " ++ show resLMAD ++ "\n" ++ + "alg points: " ++ show resAlg + in (lmadShape == algShape && resLMAD == resAlg) @? errorMessage + +compareOps :: IxFunWrap.IxFun Int -> Assertion +compareOps (ixfunLMAD, ixfunAlg) = compareIxFuns ixfunLMAD ixfunAlg + +-- XXX: Clean this up. +n :: Int +n = 19 +slice3 :: [DimIndex Int] +slice3 = [ DimSlice 2 (n `P.div` 3) 3 + , DimFix (n `P.div` 2) + , DimSlice 1 (n `P.div` 2) 2 + ] + + +-- Actual tests. +tests :: TestTree +tests = testGroup "IndexFunctionTests" + $ concat + [ test_iota + , test_slice_iota + , test_reshape_slice_iota1 + , test_permute_slice_iota + , test_repeat_slice_iota + , test_rotate_rotate_permute_slice_iota + , test_slice_rotate_permute_slice_iota1 + , test_slice_rotate_permute_slice_iota2 + , test_slice_rotate_permute_slice_iota3 + , test_permute_rotate_slice_permute_slice_iota + , test_reshape_rotate_iota + , test_reshape_permute_iota + , test_reshape_slice_iota2 + , test_reshape_slice_iota3 + , test_complex1 + , test_complex2 + , test_complex3 + , test_rebase1 + , test_rebase2 + , test_rebase3 + , test_rebase4_5 + , test_rebase6 + ] + +singleton :: TestTree -> [TestTree] +singleton = (: []) + +test_iota :: [TestTree] +test_iota = singleton $ testCase "iota" $ compareOps $ + iota [n] + +test_slice_iota :: [TestTree] +test_slice_iota = singleton $ testCase "slice . iota" $ compareOps $ + slice (iota [n, n, n]) slice3 + +test_reshape_slice_iota1 :: [TestTree] +test_reshape_slice_iota1 = singleton $ testCase "reshape . slice . iota 1" $ compareOps $ + reshape (slice (iota [n, n, n]) slice3) + [DimNew (n `P.div` 2), DimNew (n `P.div` 3)] + +test_permute_slice_iota :: [TestTree] +test_permute_slice_iota = singleton $ testCase "permute . slice . iota" $ compareOps $ + permute (slice (iota [n, n, n]) slice3) [1, 0] + +test_repeat_slice_iota :: [TestTree] +test_repeat_slice_iota = singleton $ testCase "repeat . slice . iota" $ compareOps $ + repeat (slice (iota [n, n, n]) slice3) [[2, 3], [3, 2]] [4, 4] + +test_rotate_rotate_permute_slice_iota :: [TestTree] +test_rotate_rotate_permute_slice_iota = + singleton $ testCase "rotate . rotate . permute . slice . iota" $ compareOps $ + let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] + in rotate (rotate ixfun [2, 1]) [1, 2] + +test_slice_rotate_permute_slice_iota1 :: [TestTree] +test_slice_rotate_permute_slice_iota1 = + singleton $ testCase "slice . rotate . permute . slice . iota 1" $ compareOps $ + let slice2 = [ DimSlice 0 n 1 + , DimSlice 1 (n `P.div` 2) 2 + , DimSlice 0 n 1 + ] + slice13 = [ DimSlice 2 (n `P.div` 3) 3 + , DimSlice 0 (n `P.div` 2) 1 + , DimSlice 1 (n `P.div` 2) 2 + ] + ixfun = permute (slice (iota [n, n, n]) slice2) [2, 1, 0] + ixfun' = slice (rotate ixfun [3, 1, 2]) slice13 + in ixfun' + +test_slice_rotate_permute_slice_iota2 :: [TestTree] +test_slice_rotate_permute_slice_iota2 = + singleton $ testCase "slice . rotate . permute . slice . iota 2" $ compareOps $ + let slice2 = [ DimSlice 0 (n `P.div` 2) 1 + , DimFix (n `P.div` 2) + , DimSlice 0 (n `P.div` 3) 1 + ] + slice13 = [ DimSlice 2 (n `P.div` 3) 3 + , DimSlice 0 n 1 + , DimSlice 1 (n `P.div` 2) 2 + ] + ixfun = permute (slice (iota [n, n, n]) slice13) [2, 1, 0] + ixfun' = slice (rotate ixfun [3, 1, 2]) slice2 + in ixfun' + +test_slice_rotate_permute_slice_iota3 :: [TestTree] +test_slice_rotate_permute_slice_iota3 = + singleton $ testCase "slice . rotate . permute . slice . iota 3" $ compareOps $ + -- full-slice of (-1) stride + let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] + ixfun' = rotate ixfun [2, 1] + + (n1, m1) = case IxFunLMAD.shape (fst ixfun') of + [a, b] -> (a, b) + _ -> error "expecting 2 dimensions at this point!" + negslice = [DimSlice 0 n1 1, DimSlice (m1 - 1) m1 (-1)] + ixfun'' = rotate (slice ixfun' negslice) [1,2] + in ixfun'' + +test_permute_rotate_slice_permute_slice_iota :: [TestTree] +test_permute_rotate_slice_permute_slice_iota = + singleton $ testCase "permute . rotate . slice . permute . slice . iota" $ compareOps $ + -- contiguousness + let slice33 = [ DimFix (n `P.div` 2) + , DimSlice (n - 1) (n `P.div` 3) (-1) + , DimSlice 0 n 1 + ] + ixfun = permute (slice (iota [n, n, n]) slice33) [1, 0] + m = n `P.div` 3 + slice1 = [DimSlice (n - 1) n (-1), DimSlice 2 (m - 2) 1] + ixfun' = permute (rotate (slice ixfun slice1) [1, 2]) [1, 0] + in ixfun' + +test_reshape_rotate_iota :: [TestTree] +test_reshape_rotate_iota = + -- negative reshape test + singleton $ testCase "reshape . rotate . iota" $ compareOps $ + let newdims = [DimNew (n * n), DimCoercion n] + in reshape (rotate (iota [n, n, n]) [1, 0, 0]) newdims + +test_reshape_permute_iota :: [TestTree] +test_reshape_permute_iota = + -- negative reshape test + singleton $ testCase "reshape . permute . iota" $ compareOps $ + let newdims = [DimNew (n * n), DimCoercion n] + in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims + +test_reshape_slice_iota2 :: [TestTree] +test_reshape_slice_iota2 = + -- negative reshape test + singleton $ testCase "reshape . slice . iota 2" $ compareOps $ + let newdims = [DimNew (n*n), DimCoercion n] + slc = [ DimFix (n `P.div` 2) + , DimSlice (n-1) n (-1) + , DimSlice 0 n 1 + , DimSlice (n-1) n (-1) + ] + in reshape (slice (iota [n, n, n, n]) slc) newdims + +test_reshape_slice_iota3 :: [TestTree] +test_reshape_slice_iota3 = + -- negative reshape test + singleton $ testCase "reshape . slice . iota 3" $ compareOps $ + let newdims = [DimNew (n*n), DimCoercion n] + slc = [ DimFix (n `P.div` 2) + , DimSlice 0 n 1 + , DimSlice 0 (n `P.div` 2) 1 + , DimSlice 0 n 1 + ] + in reshape (slice (iota [n, n, n, n]) slc) newdims + +test_complex1 :: [TestTree] +test_complex1 = + singleton $ testCase "reshape . permute . rotate . slice . permute . slice . iota 1" $ compareOps $ + let newdims = [ DimCoercion n + , DimCoercion n + , DimNew n + , DimCoercion ((n `P.div` 3) - 2) + ] + slice33 = [ DimSlice (n-1) (n `P.div` 3) (-1) + , DimSlice (n-1) n (-1) + , DimSlice (n-1) n (-1) + , DimSlice 0 n 1 + ] + ixfun = permute (slice (iota [n, n, n, n, n]) slice33) [3, 1, 2, 0] + m = n `P.div` 3 + slice1 = [DimSlice 0 n 1, DimSlice (n-1) n (-1), DimSlice (n-1) n (-1), DimSlice 1 (m-2) (-1)] + ixfun' = reshape (rotate (slice ixfun slice1) [1, 2, 3, 4]) newdims + in ixfun' + +test_complex2 :: [TestTree] +test_complex2 = + singleton $ testCase "reshape . permute . rotate . slice . permute . slice . iota 2" $ compareOps $ + let newdims = [ DimCoercion n + , DimNew (n*n) + , DimCoercion ((n `P.div` 3) - 2)] + slc2 = [ DimFix (n `P.div` 2) + , DimSlice (n-1) (n `P.div` 3) (-1) + , DimSlice (n-1) n (-1) + , DimSlice (n-1) n (-1) + , DimSlice 0 n 1 + ] + ixfun = permute (slice (iota [n, n, n, n, n]) slc2) [3, 1, 2, 0] + m = n `P.div` 3 + slice1 = [DimSlice 0 n 1, DimSlice (n-1) n (-1), DimSlice (n-1) n (-1), DimSlice 1 (m-2) (-1)] + ixfun' = reshape (rotate (slice ixfun slice1) [1, 0, 0, 2]) newdims + in ixfun' + +test_complex3 :: [TestTree] +test_complex3 = + singleton $ testCase "reshape . permute . rotate . slice . permute . slice . iota 3" $ compareOps $ + let newdims = [ DimCoercion 1 + , DimCoercion n + , DimNew (n*n) + , DimCoercion 2 + , DimCoercion ((n `P.div` 3) - 2) + ] + slc3 = [ DimFix (n `P.div` 2) + , DimSlice (n-1) (n `P.div` 3) (-1) + , DimSlice (n-1) n (-1) + , DimSlice (n-1) n (-1) + , DimSlice 0 n 1 + ] + ixfun = permute (slice (iota [n, n, n, n, n]) slc3) [3, 1, 2, 0] + m = n `P.div` 3 + slice1 = [DimSlice 0 n 1, DimSlice (n-1) n (-1), DimSlice (n-1) n (-1), DimSlice 1 (m-2) (-1)] + repeats = [[1],[],[],[2]] + ixfun' = reshape (repeat (rotate (slice ixfun slice1) [1, 0, 0, 2]) repeats []) newdims + in ixfun' + +test_rebase1 :: [TestTree] +test_rebase1 = + singleton $ testCase "rebase 1" $ compareOps $ + let slice_base = [ DimFix (n `P.div` 2) + , DimSlice 2 (n-2) 1 + , DimSlice 3 (n-3) 1 + ] + ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] + ixfn_orig = rotate (permute (iota [n-3, n-2]) [1, 0]) [1, 2] + ixfn_rebase = rebase ixfn_base ixfn_orig + in ixfn_rebase + +test_rebase2 :: [TestTree] +test_rebase2 = + singleton $ testCase "rebase 2" $ compareOps $ + let slice_base = [ DimFix (n `P.div` 2) + , DimSlice (n-1) (n-2) (-1) + , DimSlice (n-1) (n-3) (-1) + ] + slice_orig = [ DimSlice (n-4) (n-3) (-1) + , DimSlice (n-3) (n-2) (-1) + ] + ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] + ixfn_orig = rotate (permute (slice (iota [n-3, n-2]) slice_orig) [1, 0]) [1, 2] + ixfn_rebase = rebase ixfn_base ixfn_orig + in ixfn_rebase + +test_rebase3 :: [TestTree] +test_rebase3 = + singleton $ testCase "rebase full orig but not monotonic" $ compareOps $ + let n2 = (n-2) `P.div` 3 + n3 = (n-3) `P.div` 2 + slice_base = [ DimFix (n `P.div` 2) + , DimSlice (n-1) n2 (-3) + , DimSlice (n-1) n3 (-2) + ] + slice_orig = [ DimSlice (n3-1) n3 (-1) + , DimSlice (n2-1) n2 (-1) + ] + ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] + ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] + ixfn_rebase = rebase ixfn_base ixfn_orig + in ixfn_rebase + +test_rebase4_5 :: [TestTree] +test_rebase4_5 = + let n2 = (n-2) `P.div` 3 + n3 = (n-3) `P.div` 2 + slice_base = [ DimFix (n `P.div` 2) + , DimSlice (n-1) n2 (-3) + , DimSlice 3 n3 2 + ] + slice_orig = [ DimSlice (n3-1) n3 (-1) + , DimSlice 0 n2 1 + ] + ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] + ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] + in [ testCase "rebase mixed monotonicities" $ compareOps $ + rebase ixfn_base ixfn_orig + + , testCase "rebase repetitions and mixed monotonicities 1" $ compareOps $ + let ixfn_orig' = repeat ixfn_orig [[2, 2], [3, 3]] [2, 3] + in rebase ixfn_base ixfn_orig' + ] + +test_rebase6 :: [TestTree] +test_rebase6 = + singleton $ testCase "rebase repetitions and mixed monotonicities 2" $ compareOps $ + let n2 = (n-2) `P.div` 3 + n3 = (n-3) `P.div` 2 + slice_base = [ DimFix (n `P.div` 2) + , DimSlice (n-1) n2 (-3) + ] + slice_orig = [ DimSlice (n3-1) n3 (-1) + , DimSlice 0 n2 1 + ] + ixfn_base = permute (repeat (rotate (slice (iota [n, n]) slice_base) [1]) [[], []] [n3]) [1, 0] + ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] + ixfn_orig' = repeat ixfn_orig [[2, 2],[3, 3]] [2, 3] + ixfn_rebase = rebase ixfn_base ixfn_orig' + in ixfn_rebase diff --git a/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionWrapper.hs b/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionWrapper.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/Representation/ExplicitMemory/IndexFunctionWrapper.hs @@ -0,0 +1,65 @@ +-- | Perform index function operations in both algebraic and LMAD +-- representations. +module Futhark.Representation.ExplicitMemory.IndexFunctionWrapper + ( IxFun + , iota + , offsetIndex + , strideIndex + , permute + , rotate + , reshape + , slice + , rebase + , repeat + ) +where + +import Prelude hiding (repeat) + +import Futhark.Util.IntegralExp +import Futhark.Representation.AST.Syntax (ShapeChange, Slice) +import qualified Futhark.Representation.ExplicitMemory.IndexFunction as I +import qualified Futhark.Representation.ExplicitMemory.IndexFunction.Alg as IA + + +type Shape num = [num] +type Indices num = [num] +type Permutation = [Int] + +type IxFun num = (I.IxFun num, IA.IxFun num) + +iota :: IntegralExp num => + Shape num -> IxFun num +iota x = (I.iota x, IA.iota x) + +offsetIndex :: (Eq num, IntegralExp num) => + IxFun num -> num -> IxFun num +offsetIndex (l, a) x = (I.offsetIndex l x, IA.offsetIndex a x) + +strideIndex :: (Eq num, IntegralExp num) => + IxFun num -> num -> IxFun num +strideIndex (l, a) x = (I.strideIndex l x, IA.strideIndex a x) + +permute :: IntegralExp num => + IxFun num -> Permutation -> IxFun num +permute (l, a) x = (I.permute l x, IA.permute a x) + +rotate :: (Eq num, IntegralExp num) => + IxFun num -> Indices num -> IxFun num +rotate (l, a) x = (I.rotate l x, IA.rotate a x) + +repeat :: (Eq num, IntegralExp num) => + IxFun num -> [Shape num] -> Shape num -> IxFun num +repeat (l, a) x y = (I.repeat l x y, IA.repeat a x y) + +reshape :: (Eq num, IntegralExp num) => + IxFun num -> ShapeChange num -> IxFun num +reshape (l, a) x = (I.reshape l x, IA.reshape a x) + +slice :: (Eq num, IntegralExp num) => + IxFun num -> Slice num -> IxFun num +slice (l, a) x = (I.slice l x, IA.slice a x) + +rebase :: (Eq num, IntegralExp num) => + IxFun num -> IxFun num -> IxFun num +rebase (l, a) (l1, a1) = (I.rebase l l1, IA.rebase a a1) diff --git a/unittests/futhark_tests.hs b/unittests/futhark_tests.hs --- a/unittests/futhark_tests.hs +++ b/unittests/futhark_tests.hs @@ -3,8 +3,11 @@ import qualified Language.Futhark.SyntaxTests import qualified Futhark.Representation.AST.Syntax.CoreTests import qualified Futhark.Representation.AST.AttributesTests +import qualified Futhark.Representation.ExplicitMemory.IndexFunctionTests import qualified Futhark.Optimise.AlgSimplifyTests import qualified Futhark.Pkg.SolveTests +import qualified Futhark.Representation.PrimitiveTests +import qualified Futhark.Analysis.ScalExpTests import Test.Tasty @@ -16,6 +19,9 @@ , Futhark.Optimise.AlgSimplifyTests.tests , Futhark.Representation.AST.Syntax.CoreTests.tests , Futhark.Pkg.SolveTests.tests + , Futhark.Representation.ExplicitMemory.IndexFunctionTests.tests + , Futhark.Representation.PrimitiveTests.tests + , Futhark.Analysis.ScalExpTests.tests ] main :: IO ()