futhark 0.25.32 → 0.25.33

raw patch · 155 files changed

+1865/−2753 lines, 155 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +34/−0
• docs/conf.py +13/−4
• docs/installation.rst +2/−2
• docs/language-reference.rst +3/−15
• docs/man/futhark-literate.rst +4/−0
• docs/man/futhark-script.rst +2/−2
• docs/performance.rst +6/−0
• futhark.cabal +3/−2
• src-testing/Futhark/Analysis/DataDependenciesTests.hs +44/−0
• src-testing/Futhark/IR/GPUTests.hs +1/−1
• src-testing/Futhark/IR/MCTests.hs +1/−1
• src-testing/Futhark/IR/SOACSTests.hs +17/−0
• src-testing/futhark_tests.hs +2/−0
• src/Futhark/AD/Derivatives.hs +1/−1
• src/Futhark/AD/Fwd.hs +4/−30
• src/Futhark/AD/Rev.hs +71/−24
• src/Futhark/AD/Rev/Hist.hs +25/−36
• src/Futhark/AD/Rev/Loop.hs +5/−5
• src/Futhark/AD/Rev/Map.hs +2/−2
• src/Futhark/AD/Rev/Monad.hs +20/−11
• src/Futhark/AD/Rev/Reduce.hs +4/−4
• src/Futhark/AD/Rev/SOAC.hs +0/−12
• src/Futhark/AD/Rev/Scan.hs +17/−17
• src/Futhark/AD/Rev/Scatter.hs +0/−187
• src/Futhark/Actions.hs +11/−3
• src/Futhark/Analysis/AccessPattern.hs +1/−1
• src/Futhark/Analysis/Alias.hs +3/−3
• src/Futhark/Analysis/HORep/MapNest.hs +4/−6
• src/Futhark/Analysis/HORep/SOAC.hs +0/−20
• src/Futhark/Analysis/Interference.hs +4/−4
• src/Futhark/Analysis/LastUse.hs +1/−1
• src/Futhark/Analysis/MemAlias.hs +4/−4
• src/Futhark/Analysis/Metrics.hs +1/−1
• src/Futhark/Analysis/PrimExp/Table.hs +2/−2
• src/Futhark/Builder.hs +10/−9
• src/Futhark/Builder/Class.hs +4/−4
• src/Futhark/CLI/Dev.hs +11/−9
• src/Futhark/CLI/Fmt.hs +0/−1
• src/Futhark/CLI/Literate.hs +26/−18
• src/Futhark/CLI/Script.hs +20/−14
• src/Futhark/CodeGen/Backends/GenericC/Code.hs +2/−2
• src/Futhark/CodeGen/Backends/GenericC/Fun.hs +2/−2
• src/Futhark/CodeGen/Backends/GenericC/Monad.hs +1/−1
• src/Futhark/CodeGen/Backends/GenericC/Types.hs +28/−12
• src/Futhark/CodeGen/Backends/MulticoreC.hs +10/−10
• src/Futhark/CodeGen/Backends/MulticoreISPC.hs +20/−13
• src/Futhark/CodeGen/Backends/SimpleRep.hs +9/−2
• src/Futhark/CodeGen/ImpCode/Multicore.hs +2/−2
• src/Futhark/CodeGen/ImpGen.hs +23/−23
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +15/−22
• src/Futhark/CodeGen/ImpGen/GPU/Block.hs +11/−13
• src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs +9/−7
• src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs +5/−5
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +7/−7
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs +2/−2
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +2/−2
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +1/−1
• src/Futhark/CodeGen/ImpGen/Multicore.hs +2/−2
• src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +8/−10
• src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs +6/−6
• src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs +1/−9
• src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs +3/−3
• src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs +5/−5
• src/Futhark/Construct.hs +15/−15
• src/Futhark/Doc/Generator.hs +4/−0
• src/Futhark/Fmt/Printer.hs +2/−2
• src/Futhark/FreshNames.hs +1/−1
• src/Futhark/IR/Aliases.hs +13/−9
• src/Futhark/IR/GPU/Simplify.hs +1/−1
• src/Futhark/IR/GPUMem.hs +1/−1
• src/Futhark/IR/Mem/Simplify.hs +1/−1
• src/Futhark/IR/Parse.hs +34/−32
• src/Futhark/IR/Prop.hs +12/−0
• src/Futhark/IR/Prop/Aliases.hs +1/−1
• src/Futhark/IR/Prop/Names.hs +6/−4
• src/Futhark/IR/Prop/TypeOf.hs +2/−2
• src/Futhark/IR/Rephrase.hs +1/−1
• src/Futhark/IR/SOACS.hs +0/−2
• src/Futhark/IR/SOACS/SOAC.hs +0/−203
• src/Futhark/IR/SOACS/Simplify.hs +44/−77
• src/Futhark/IR/SegOp.hs +80/−233
• src/Futhark/IR/Syntax.hs +15/−10
• src/Futhark/IR/Syntax/Core.hs +7/−0
• src/Futhark/IR/TypeCheck.hs +2/−0
• src/Futhark/Internalise/AccurateSizes.hs +4/−4
• src/Futhark/Internalise/Bindings.hs +1/−1
• src/Futhark/Internalise/Defunctionalise.hs +9/−9
• src/Futhark/Internalise/Entry.hs +6/−2
• src/Futhark/Internalise/Exps.hs +72/−130
• src/Futhark/Internalise/FullNormalise.hs +20/−19
• src/Futhark/Internalise/LiftLambdas.hs +2/−2
• src/Futhark/Internalise/Monad.hs +1/−1
• src/Futhark/Internalise/Monomorphise.hs +6/−6
• src/Futhark/Internalise/ReplaceRecords.hs +1/−1
• src/Futhark/Internalise/TypesValues.hs +5/−2
• src/Futhark/MonadFreshNames.hs +8/−24
• src/Futhark/Optimise/ArrayLayout/Transform.hs +11/−15
• src/Futhark/Optimise/ArrayShortCircuiting.hs +8/−8
• src/Futhark/Optimise/ArrayShortCircuiting/ArrayCoalescing.hs +11/−17
• src/Futhark/Optimise/BlkRegTiling.hs +56/−60
• src/Futhark/Optimise/CSE.hs +2/−2
• src/Futhark/Optimise/DoubleBuffer.hs +6/−6
• src/Futhark/Optimise/Fusion.hs +85/−41
• src/Futhark/Optimise/Fusion/GraphRep.hs +24/−23
• src/Futhark/Optimise/Fusion/RulesWithAccs.hs +1/−310
• src/Futhark/Optimise/Fusion/TryFusion.hs +2/−54
• src/Futhark/Optimise/GenRedOpt.hs +24/−16
• src/Futhark/Optimise/HistAccs.hs +7/−7
• src/Futhark/Optimise/MemoryBlockMerging.hs +20/−20
• src/Futhark/Optimise/MergeGPUBodies.hs +3/−3
• src/Futhark/Optimise/ReduceDeviceSyncs.hs +8/−11
• src/Futhark/Optimise/Simplify/Rep.hs +4/−4
• src/Futhark/Optimise/Simplify/Rules.hs +4/−0
• src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +5/−5
• src/Futhark/Optimise/Simplify/Rules/Index.hs +2/−2
• src/Futhark/Optimise/Sink.hs +2/−2
• src/Futhark/Optimise/TileLoops.hs +9/−9
• src/Futhark/Optimise/TileLoops/Shared.hs +36/−40
• src/Futhark/Optimise/Unstream.hs +3/−14
• src/Futhark/Pass/ExpandAllocations.hs +14/−14
• src/Futhark/Pass/ExplicitAllocations.hs +8/−8
• src/Futhark/Pass/ExplicitAllocations/GPU.hs +4/−4
• src/Futhark/Pass/ExplicitAllocations/SegOp.hs +3/−4
• src/Futhark/Pass/ExtractKernels.hs +6/−81
• src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +5/−5
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +47/−131
• src/Futhark/Pass/ExtractKernels/Distribution.hs +2/−2
• src/Futhark/Pass/ExtractKernels/ISRWIM.hs +1/−1
• src/Futhark/Pass/ExtractKernels/Interchange.hs +3/−3
• src/Futhark/Pass/ExtractKernels/Intrablock.hs +4/−41
• src/Futhark/Pass/ExtractKernels/StreamKernel.hs +3/−3
• src/Futhark/Pass/ExtractKernels/ToGPU.hs +5/−5
• src/Futhark/Pass/ExtractMulticore.hs +1/−15
• src/Futhark/Pass/LiftAllocations.hs +8/−8
• src/Futhark/Pass/LowerAllocations.hs +8/−8
• src/Futhark/Pipeline.hs +7/−1
• src/Futhark/Script.hs +164/−98
• src/Futhark/Test.hs +3/−3
• src/Futhark/Tools.hs +63/−2
• src/Futhark/Transform/FirstOrderTransform.hs +28/−42
• src/Futhark/Transform/Rename.hs +4/−4
• src/Futhark/Transform/Substitute.hs +1/−1
• src/Futhark/Util.hs +19/−0
• src/Language/Futhark/Core.hs +0/−5
• src/Language/Futhark/FreeVars.hs +1/−1
• src/Language/Futhark/Interpreter.hs +164/−161
• src/Language/Futhark/Parser/Parser.y +6/−5
• src/Language/Futhark/Pretty.hs +1/−1
• src/Language/Futhark/Prop.hs +1/−1
• src/Language/Futhark/Query.hs +1/−1
• src/Language/Futhark/Syntax.hs +1/−1
• src/Language/Futhark/TypeChecker/Consumption.hs +1/−1
• src/Language/Futhark/TypeChecker/Names.hs +3/−3
• src/Language/Futhark/TypeChecker/Terms.hs +7/−6
• src/Language/Futhark/TypeChecker/Terms/Monad.hs +1/−1

@@ -5,6 +5,40 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.25.33]++### Added++* Futhark now implements the cachedir specification, such that e.g. the `data`+  directories created by `futhark test` and `futhark bench` contain a+  `CACHEDIR.TAG` file.++* C types corresponding to records now have nicer names.++* ``futhark script`` now has a ``$restore`` procedure, and the ``$store``+  procedure has been augmented with support for opaque types.++* The expression guarded by an `assert` may now be any expression, and need no+  longer be an atom.++### Fixed++* `futhark fmt --check` no longer prints the program on failure.++* Use of unsigned types in entry points could cause invalid C to be generated.+  (#2306)++* `futhark script` now frees values before terminating.++* Correct source locations when warning about unused local functions.++* Unpacking a unary sum type directly in a parameter or `let`-binding was+  defective. (#2314)++* The derivative of `x**1` for `x==0` would be NaN.++* `futhark fmt` now prints multi-line `assert` in a less horrible way.+ ## [0.25.32]  ### Added

@@ -62,7 +62,7 @@  # General information about the project. project = "Futhark"-copyright = "2013-2020, DIKU, University of Copenhagen"+copyright = "2013-2025, DIKU, University of Copenhagen"  # The version info for the project you're documenting, acts as replacement for # |version| and |release|, also used in various other places throughout the@@ -287,13 +287,22 @@  # The theme to use for HTML and HTML Help pages.  See the documentation for # a list of builtin themes.-html_theme = "futhark"-html_theme_path = ["_theme"]+html_theme = "classic"  # Theme options are theme-specific and customize the look and feel of a theme # further.  For a list of options available for each theme, see the # documentation.-html_theme_options: Dict[Any, Any] = {}+html_theme_options: Dict[str, str] = {+    "sidebarbgcolor": "#5f021f",+    "footerbgcolor": "#5f021f",+    "relbarbgcolor": "#5f021f",+    "bgcolor": "#fff9e5",+    "linkcolor": "#5f021f",+    "headbgcolor": "#eeeeee",+    "sidebartextcolor": "white",+    "sidebarlinkcolor": "white",+    "codebgcolor": "#eeeeee",+}  # Add any paths that contain custom themes here, relative to this directory. # html_theme_path = []

@@ -96,8 +96,8 @@ **Linux (x86_64)**   `futhark-nightly-linux-x86_64.tar.xz <https://futhark-lang.org/releases/futhark-nightly-linux-x86_64.tar.xz>`_ -**macOS (x86_64)**-  `futhark-nightly-macos-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-macos-x86_64.tar.xz>`_+**macOS (ARM64)**+  `futhark-nightly-macos-arm64.zip <https://futhark-lang.org/releases/futhark-nightly-macos-arm64.tar.xz>`_  **Windows (x86_64)**   `futhark-nightly-windows-x86_64.zip <https://futhark-lang.org/releases/futhark-nightly-windows-x86_64.zip>`_

@@ -503,7 +503,7 @@       : | "loop" `pat` ["=" `exp`] `loopform` "do" `exp`       : | "#[" `attr` "]" `exp`       : | "unsafe" `exp`-      : | "assert" `atom` `atom`+      : | "assert" `atom` `exp`       : | `exp` "with" `slice` "=" `exp`       : | `exp` "with" `fieldid` ("." `fieldid`)* "=" `exp`       : | "match" `exp` ("case" `pat` "->" `exp`)+@@ -590,8 +590,8 @@   ``t``.  To pass a single array-typed parameter, enclose it in   parens. -* The bodies of ``let``, ``if``, and ``loop`` extend as far to the-  right as possible.+* The bodies of ``let``, ``if``, and ``loop`` extend as far to the right as+  possible, as does the expression guarded by an ``assert``.  * The following table describes the precedence and associativity of   infix operators in both expressions and type expressions.  All@@ -933,18 +933,6 @@ Apply the given attribute to the expression.  Attributes are an ad-hoc and optional mechanism for providing extra information, directives, or hints to the compiler.  See :ref:`attributes` for more information.--``unsafe e``-............--Elide safety checks and assertions (such as bounds checking) that-occur during execution of ``e``.  This is useful if the compiler is-otherwise unable to avoid bounds checks (e.g. when using indirect-indexes), but you really do not want them there.  Make very sure that-the code is correct; eliding such checks can lead to memory-corruption.--This construct is deprecated.  Use the ``#[unsafe]`` attribute instead.  .. _assert: 

@@ -278,6 +278,10 @@   soundfile. Most common audio-formats are supported, including mp3, ogg, wav,   flac and opus. +* ``$restore "type" "file"`` loads a serialised value of type ``type`` from+  ``file``. The usual caveats apply regarding the stability of the value+  serialisation format.+ FutharkScript supports a form of automatic uncurrying. If a function taking *n* parameters is applied to a single argument that is an *n*-element tuple, the function is applied to the elements of the

@@ -87,8 +87,8 @@ ADDITIONAL BUILTINS =================== -* ``$store "file" v`` store the value *v* (which must be a primitive-  or an array) as a binary value in the given file.+* ``$store "file" v`` store the value *v* as a binary value in the given file.+  It can be restored again with ``$restore``.  BUGS ====

@@ -257,6 +257,12 @@ which is not great.  Take caution when you use sum types with large arrays in their payloads. +Unary sum types+!!!!!!!!!!!!!!!++As an optimisation, the ``i8`` tag is elided when a sum type has only a single+constructor. This means you can always use unary sum types with zero overhead.+ Functions ~~~~~~~~~ 

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.25.32+version:        0.25.33 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -123,7 +123,6 @@       Futhark.AD.Rev.Monad       Futhark.AD.Rev.Reduce       Futhark.AD.Rev.Scan-      Futhark.AD.Rev.Scatter       Futhark.AD.Rev.SOAC       Futhark.Analysis.AccessPattern       Futhark.Analysis.AlgSimplify@@ -512,6 +511,7 @@   exposed-modules:       Futhark.AD.DerivativesTests       Futhark.Analysis.AlgSimplifyTests+      Futhark.Analysis.DataDependenciesTests       Futhark.Analysis.PrimExp.TableTests       Futhark.BenchTests       Futhark.IR.GPUTests@@ -523,6 +523,7 @@       Futhark.IR.Prop.RearrangeTests       Futhark.IR.Prop.ReshapeTests       Futhark.IR.PropTests+      Futhark.IR.SOACSTests       Futhark.IR.Syntax.CoreTests       Futhark.IR.SyntaxTests       Futhark.Internalise.TypesValuesTests

@@ -0,0 +1,44 @@+module Futhark.Analysis.DataDependenciesTests (tests) where++import Data.Map qualified as M+import Futhark.Analysis.DataDependencies+import Futhark.IR.SOACS+import Futhark.IR.SOACSTests ()+import Test.Tasty+import Test.Tasty.HUnit++names :: [VName] -> Names+names = namesFromList++tests :: TestTree+tests =+  testGroup+    "DataDependenciesTests"+    [ testGroup+        "lambdaDependencies"+        [ testCase "if" $+            lambdaDependencies+              mempty+              ("\\{x_0: i32} : {i32} -> {x_0}" :: Lambda SOACS)+              [names ["y_1"]]+              @?= [names ["y_1"]],+          testCase "flip" $+            lambdaDependencies+              mempty+              ("\\{x_0: i32, x_1: bool} : {bool, i32} -> {x_1, x_0}" :: Lambda SOACS)+              [names ["y_2"], names ["y_3"]]+              @?= [names ["y_3"], names ["y_2"]],+          testCase "add" $+            lambdaDependencies+              mempty+              ("\\{x_0: i32, x_1: i32} : {i32} -> let {x_2: i32} = add32(x_1, x_0) in {x_2}" :: Lambda SOACS)+              [names ["y_2"], names ["y_3"]]+              @?= [names ["y_3", "y_2"]],+          testCase "outer" $+            lambdaDependencies+              (M.fromList [("x_1", names ["x_4"])])+              ("\\{x_0: i32} : {i32} -> let {x_2: i32} = add32(x_1, x_0) in {x_2}" :: Lambda SOACS)+              [names ["y_2"]]+              @?= [names ["y_2", "x_1"]]+        ]+    ]

@@ -17,4 +17,4 @@   fromString = parseString "Body GPU" parseBodyGPU  instance IsString (Prog GPU) where-  fromString = parseString "Prog GPU" parseGPU+  fromString = snd . parseString "Prog GPU" parseGPU

@@ -17,4 +17,4 @@   fromString = parseString "Body MC" parseBodyMC  instance IsString (Prog MC) where-  fromString = parseString "Prog MC" parseMC+  fromString = snd . parseString "Prog MC" parseMC

@@ -0,0 +1,17 @@+{-# OPTIONS_GHC -fno-warn-orphans #-}++module Futhark.IR.SOACSTests () where++import Data.String+import Futhark.IR.Parse+import Futhark.IR.SOACS+import Futhark.IR.SyntaxTests (parseString)++instance IsString (Lambda SOACS) where+  fromString = parseString "Lambda" parseLambdaSOACS++instance IsString (Body SOACS) where+  fromString = parseString "Body" parseBodySOACS++instance IsString (Stm SOACS) where+  fromString = parseString "Stm" parseStmSOACS

@@ -2,6 +2,7 @@  import Futhark.AD.DerivativesTests qualified import Futhark.Analysis.AlgSimplifyTests qualified+import Futhark.Analysis.DataDependenciesTests qualified import Futhark.BenchTests qualified import Futhark.IR.Mem.IntervalTests qualified import Futhark.IR.Mem.IxFunTests qualified@@ -35,6 +36,7 @@       Language.Futhark.PrimitiveTests.tests,       Futhark.Optimise.MemoryBlockMerging.GreedyColoringTests.tests,       Futhark.Analysis.AlgSimplifyTests.tests,+      Futhark.Analysis.DataDependenciesTests.tests,       Language.Futhark.TypeCheckerTests.tests,       Language.Futhark.SemanticTests.tests,       Futhark.Optimise.ArrayLayoutTests.tests

@@ -135,7 +135,7 @@   floatBinOp derivs derivs derivs ft a b   where     derivs x y =-      ( y * (x ** (y - 1)),+      ( condExp (x .==. 0 .&&. y .<. 1) 0 (y * (x ** (y - 1))),         (x ** y) * condExp (x .<=. 0) 0 (log x)       ) pdBinOp (FMax ft) a b =

@@ -81,7 +81,7 @@         (RState mempty vn)  tanVName :: VName -> ADM VName-tanVName v = newVName (baseString v <> "_tan")+tanVName v = newVName (baseName v <> "_tan")  insertTan :: VName -> VName -> ADM () insertTan v v' =@@ -160,7 +160,7 @@       Just v_tan -> pure v_tan       Nothing -> do         t <- lookupType v-        letExp (baseString v <> "_implicit_tan") $ zeroExp t+        letExp (baseName v <> "_implicit_tan") $ zeroExp t   bundleTan v = do     t <- lookupType v     if isAcc t@@ -330,32 +330,6 @@             histNeutral = interleave nes nes_tan,             histOp = op'           }-fwdSOAC (Pat pes) aux (Scatter w ivs as lam) = do-  as_tan <- mapM (\(s, n, a) -> do a_tan <- tangent a; pure (s, n, a_tan)) as-  pes_tan <- mapM newTan pes-  ivs' <- bundleTangents ivs-  let (as_ws, as_ns, _as_vs) = unzip3 as-      n_indices = sum $ zipWith (*) as_ns $ map length as_ws-  lam' <- fwdScatterLambda n_indices lam-  let s = Let (Pat (pes ++ pes_tan)) aux $ Op $ Scatter w ivs' (as ++ as_tan) lam'-  addStm s-  where-    fwdScatterLambda :: Int -> Lambda SOACS -> ADM (Lambda SOACS)-    fwdScatterLambda n_indices (Lambda params ret body) = do-      params' <- bundleNewList params-      ret_tan <- mapM tangent $ drop n_indices ret-      body' <- fwdBodyScatter n_indices body-      let indices = concat $ replicate 2 $ take n_indices ret-          ret' = indices ++ drop n_indices ret ++ ret_tan-      pure $ Lambda params' ret' body'-    fwdBodyScatter :: Int -> Body SOACS -> ADM (Body SOACS)-    fwdBodyScatter n_indices (Body _ stms res) = do-      (res_tan, stms') <- collectStms $ do-        mapM_ fwdStm stms-        mapM tangent $ drop n_indices res-      let indices = concat $ replicate 2 $ take n_indices res-          res' = indices ++ drop n_indices res ++ res_tan-      pure $ mkBody stms' res' fwdSOAC _ _ JVP {} =   error "fwdSOAC: nested JVP not allowed." fwdSOAC _ _ VJP {} =@@ -394,8 +368,8 @@                       _ -> error $ "fwdStm.convertTo: " ++ prettyString (f, tt, e_t)                 where                   e_t = primExpType e-          zipWithM_ (letBindNames . pure) (patNames pat_tan)-            =<< mapM toExp (zipWith (~*~) (map (convertTo ret) arg_tans) derivs)+          letBindNames (patNames pat_tan)+            =<< toExp (foldl1 (~+~) $ zipWith (~*~) (map (convertTo ret) arg_tans) derivs) fwdStm (Let pat aux (Match ses cases defbody (MatchDec ret ifsort))) = do   cases' <- slocal' $ mapM (traverse fwdBody) cases   defbody' <- slocal' $ fwdBody defbody

@@ -9,6 +9,7 @@ module Futhark.AD.Rev (revVJP) where  import Control.Monad+import Control.Monad.Identity import Data.List ((\\)) import Data.List.NonEmpty (NonEmpty (..)) import Data.Map qualified as M@@ -22,7 +23,7 @@ import Futhark.Tools import Futhark.Transform.Rename import Futhark.Transform.Substitute-import Futhark.Util (takeLast)+import Futhark.Util (chunks, takeLast)  patName :: Pat Type -> ADM VName patName (Pat [pe]) = pure $ patElemName pe@@ -33,8 +34,7 @@   v_t <- lookupType v   case v_t of     Array {} ->-      letExp (baseString v <> "_copy") . BasicOp $-        Replicate mempty (Var v)+      letExp (baseName v <> "_copy") . BasicOp $ Replicate mempty (Var v)     _ -> pure v  -- The vast majority of BasicOps require no special treatment in the@@ -144,7 +144,7 @@         ne <- letSubExp "zero" $ zeroExp x_t         n <- letSubExp "rep_size" =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ns         pat_adj_flat <--          letExp (baseString pat_adj <> "_flat") . BasicOp $+          letExp (baseName pat_adj <> "_flat") . BasicOp $             Reshape pat_adj (reshapeAll (Shape ns) (Shape $ n : arrayDims x_t))         reduce <- reduceSOAC [Reduce Commutative lam [ne]]         updateSubExpAdj x@@ -159,7 +159,7 @@               let w = arraySize 0 v_t                   slice = DimSlice start w (intConst Int64 1)               pat_adj_slice <--                letExp (baseString pat_adj <> "_slice") $+                letExp (baseName pat_adj <> "_slice") $                   BasicOp $                     Index pat_adj (sliceAt v_t d [slice])               start' <- letSubExp "start" $ BasicOp $ BinOp (Add Int64 OverflowUndef) start w@@ -203,7 +203,8 @@       pat_adjs <- mapM lookupAdjVal (patNames pat)       returnSweepCode $ do         forM_ (zip pat_adjs vs) $ \(adj, v) -> do-          adj_i <- letExp "updateacc_val_adj" $ BasicOp $ Index adj $ Slice $ map DimFix is+          adj_t <- lookupType adj+          adj_i <- letExp "updateacc_val_adj" $ BasicOp $ Index adj $ fullSlice adj_t $ map DimFix is           updateSubExpAdj v adj_i  vjpOps :: VjpOps@@ -213,6 +214,27 @@       vjpStm = diffStm     } +-- | Transform updates on accumulators matching the given certificates into+-- updates that write provided zero values.+zeroOutUpdates :: [(VName, [SubExp])] -> Lambda SOACS -> Lambda SOACS+zeroOutUpdates certs_to_zeroes lam = lam {lambdaBody = onBody $ lambdaBody lam}+  where+    onExp = runIdentity . mapExpM mapper+      where+        mapper =+          (identityMapper :: (Monad m) => Mapper SOACS SOACS m)+            { mapOnOp = traverseSOACStms (\_ stms -> pure $ onStms stms),+              mapOnBody = \_ body -> pure $ onBody body+            }+    onStms = fmap onStm+    onStm (Let (Pat [pe]) aux (BasicOp (UpdateAcc safety acc is _)))+      | Acc c _ _ _ <- patElemType pe,+        Just zero <- lookup c certs_to_zeroes =+          Let (Pat [pe]) aux (BasicOp (UpdateAcc safety acc is zero))+    onStm (Let pat aux e) = Let pat aux $ onExp e++    onBody body = body {bodyStms = onStms $ bodyStms body}+ diffStm :: Stm SOACS -> ADM () -> ADM () diffStm (Let pat aux (BasicOp e)) m =   diffBasicOp pat aux e m@@ -285,16 +307,20 @@     free_accs <- filterM (fmap isAcc . lookupType) free_vars     let free_vars' = free_vars \\ free_accs     lam'' <- diffLambda' adjs free_vars' lam'-    inputs' <- mapM renameInputLambda inputs-    free_adjs <- letTupExp "with_acc_contrib" $ WithAcc inputs' lam''+    (inputs_zeroes, inputs') <-+      unzip <$> zipWithM renameInputLambda (chunks lengths adjs) inputs+    let certs = map paramName $ take (length inputs) $ lambdaParams lam''+    free_adjs <- letTupExp "with_acc_contrib" $ WithAcc inputs' $ zeroOutUpdates (zip certs inputs_zeroes) lam''     zipWithM_ insAdj (arrs <> free_vars') free_adjs   where+    lengths = map (\(_, as, _) -> length as) inputs     arrs = concatMap (\(_, as, _) -> as) inputs-    renameInputLambda (shape, as, Just (f, nes)) = do-      f' <- renameLambda f-      pure (shape, as, Just (f', nes))-    renameInputLambda input = pure input-    diffLambda' res_adjs get_adjs_for (Lambda params ts body) =+    renameInputLambda as_adj (shape, as, _) = do+      nes_ts <- mapM (fmap (stripArray (shapeRank shape)) . lookupType) as+      zeroes <- mapM (zeroArray mempty) nes_ts+      as' <- mapM adjVal as_adj+      pure (map Var zeroes, (shape, as', Nothing))+    diffLambda' res_adjs get_adjs_for (Lambda params ts body) = do       localScope (scopeOfLParams params) $ do         Body () stms res <- diffBody res_adjs get_adjs_for body         let body' = Body () stms $ take (length inputs) res <> takeLast (length get_adjs_for) res@@ -358,10 +384,16 @@ -- some assumptions and lay down a basic design. -- -- First, we assume that any WithAccs that occur in the program are--- the result of previous invocations of VJP.  This means we can rely--- on the operator having a constant adjoint (it's some kind of--- addition).+-- come from one of these sources: --+-- - A previous instance of VJP, which means we can rely on the operator having+--   a constant adjoint (it's addition as appropriate to the type).+--+-- - A scatter, meaning there is no operator.+--+-- (These can actually be distinguished by the presence of an operator, although+-- we do not currently bother.)+-- -- Second, the adjoint of an accumulator is an array of the same type -- as the underlying array.  For example, the adjoint type of the -- primal type 'acc(c, [n], {f64})' is '[n]f64'.  In principle the@@ -369,23 +401,38 @@ -- '[]f64', '[]f32'.  Our current design assumes that adjoints are -- single variables.  This is fixable. --+-- In the return sweep, when inserting the with_acc, we still compute the+-- "original" accumulator result, but modified such that its initial value is+-- the adjoint of the result of the accumulator. We also modify the update_accs+-- of these accumulators to be with zero values. This means that the array that+-- is produced will be equal to the adjoint of the result, except for those+-- places that have been updated, where it will be zero. This is intuitively+-- sensible - values that have been overwritten (and so do not contribute to the+-- result) should obviously have zero sensitivity.+-- -- # Adjoint of UpdateAcc -----   Consider primal code+-- Consider primal code -- --     update_acc(acc, i, v) -----   Interpreted as an imperative statement, this means+-- Interpreted as an imperative statement, this means -- --     acc[i] ⊕= v -----   for some '⊕'.  Normally all the compiler knows of '⊕' is that it---   is associative and commutative, but because we assume that all---   accumulators are the result of previous AD transformations, we---   can assume that '⊕' actually behaves like addition - that is, has---   unit partial derivatives.  So the return sweep is+-- for some '⊕'.  Normally all the compiler knows of '⊕' is that it+-- is associative and commutative, but because we assume that all+-- accumulators are the result of previous AD transformations, we+-- can assume that '⊕' actually behaves like addition - that is, has+-- unit partial derivatives.  So the return sweep is -----     v += acc_adj[i]+--     v_adj += acc_adj[i]+--+-- Further, we modify the primal code so that it becomes+--+--     update_acc(acc, i, 0)+--+-- for some appropriate notion of zero. -- -- # Adjoint of Map --

@@ -56,7 +56,7 @@     $ elseIf t cs bs elseIf _ _ _ = error "In elseIf, Hist.hs: input not supported" -bindSubExpRes :: (MonadBuilder m) => String -> [SubExpRes] -> m [VName]+bindSubExpRes :: (MonadBuilder m) => Name -> [SubExpRes] -> m [VName] bindSubExpRes s =   traverse     ( \(SubExpRes cs se) -> do@@ -130,20 +130,9 @@    letExp "is'" $ Op $ Screma n (pure is) $ mapSOAC is'_lam -multiScatter :: SubExp -> [VName] -> VName -> [VName] -> ADM [VName]-multiScatter n dst is vs = do-  tps <- traverse lookupType vs-  par_i <- newParam "i" $ Prim int64-  scatter_params <- traverse (newParam "scatter_param" . rowType) tps-  scatter_lam <--    mkLambda (par_i : scatter_params) $-      fmap subExpsRes . mapM (letSubExp "scatter_map_res") =<< do-        p1 <- replicateM (length scatter_params) $ eParam par_i-        p2 <- traverse eParam scatter_params-        pure $ p1 <> p2--  let spec = zipWith (\t -> (,,) (Shape $ pure $ arraySize 0 t) 1) tps dst-  letTupExp "scatter_res" . Op $ Scatter n (is : vs) spec scatter_lam+multiScatter :: [VName] -> VName -> [VName] -> ADM [VName]+multiScatter dst is vs =+  doScatter "scatter_res" 1 dst (is : vs) $ pure . map (Var . paramName)  multiIndex :: [VName] -> [DimIndex SubExp] -> ADM [VName] multiIndex vs s = do@@ -189,7 +178,7 @@   let dst_dims = arrayDims dst_type    dst_cpy <--    letExp (baseString dst <> "_copy") . BasicOp $+    letExp (baseName dst <> "_copy") . BasicOp $       Replicate mempty (Var dst)    acc_v_p <- newParam "acc_v" $ Prim t@@ -243,7 +232,7 @@    let hist_op = HistOp (Shape [w]) rf [dst_cpy, dst_minus_ones] [ne, if nr_dims == 1 then intConst Int64 (-1) else ne_minus_ones] hist_lam   f' <- mkIdentityLambda [Prim int64, rowType vs_type, rowType $ Array int64 (Shape vs_dims) NoUniqueness]-  x_inds <- newVName (baseString x <> "_inds")+  x_inds <- newVName (baseName x <> "_inds")   auxing aux $     letBindNames [x, x_inds] $       Op $@@ -253,8 +242,8 @@    x_bar <- lookupAdjVal x -  x_ind_dst <- newParam (baseString x <> "_ind_param") $ Prim int64-  x_bar_dst <- newParam (baseString x <> "_bar_param") $ Prim t+  x_ind_dst <- newParam (baseName x <> "_ind_param") $ Prim int64+  x_bar_dst <- newParam (baseName x <> "_bar_param") $ Prim t   dst_lam_inner <-     mkLambda [x_ind_dst, x_bar_dst] $       fmap varsRes . letTupExp "dst_bar"@@ -265,7 +254,7 @@   dst_lam <- nestedmap inner_dims [int64, vs_elm_type] dst_lam_inner    dst_bar <--    letExp (baseString dst <> "_bar") . Op $+    letExp (baseName dst <> "_bar") . Op $       Screma w [x_inds, x_bar] (mapSOAC dst_lam)    updateAdj dst dst_bar@@ -275,8 +264,8 @@   inds' <- traverse (letExp "inds" . BasicOp . Replicate (Shape [w]) . Var) =<< mk_indices inner_dims []   let inds = x_inds : inds' -  par_x_ind_vs <- replicateM nr_dims $ newParam (baseString x <> "_ind_param") $ Prim int64-  par_x_bar_vs <- newParam (baseString x <> "_bar_param") $ Prim t+  par_x_ind_vs <- replicateM nr_dims $ newParam (baseName x <> "_ind_param") $ Prim int64+  par_x_bar_vs <- newParam (baseName x <> "_bar_param") $ Prim t   vs_lam_inner <-     mkLambda (par_x_bar_vs : par_x_ind_vs) $       fmap varsRes . letTupExp "res"@@ -286,7 +275,7 @@           ( eBody $               pure $ do                 vs_bar_i <--                  letSubExp (baseString vs_bar <> "_el") . BasicOp $+                  letSubExp (baseName vs_bar <> "_el") . BasicOp $                     Index vs_bar . Slice $                       fmap (DimFix . Var . paramName) par_x_ind_vs                 eBinOp (getBinOpPlus t) (eParam par_x_bar_vs) (eSubExp vs_bar_i)@@ -294,7 +283,7 @@   vs_lam <- nestedmap inner_dims (vs_elm_type : replicate nr_dims int64) vs_lam_inner    vs_bar_p <--    letExp (baseString vs <> "_partial") . Op $+    letExp (baseName vs <> "_partial") . Op $       Screma w (x_bar : inds) (mapSOAC vs_lam)    q <-@@ -309,10 +298,10 @@       letExp "flat" . BasicOp . Reshape v $         reshapeAll (arrayShape v_t) (Shape [q]) -  f'' <- mkIdentityLambda $ replicate nr_dims (Prim int64) ++ [Prim t]   vs_bar' <--    letExp (baseString vs <> "_bar") . Op $-      Scatter q scatter_inps [(Shape vs_dims, 1, vs_bar)] f''+    fmap head $+      doScatter (baseName vs <> "_bar") nr_dims [vs_bar] scatter_inps $+        pure . map (Var . paramName)   insAdj vs vs_bar'   where     mk_indices :: [SubExp] -> [SubExp] -> ADM [VName]@@ -432,7 +421,7 @@    lam_mul'' <- renameLambda lam_mul'   dst_bar_res <- eLambda lam_mul'' $ map (eSubExp . Var) [h_part, x_bar]-  dst_bar <- bindSubExpRes (baseString dst <> "_bar") dst_bar_res+  dst_bar <- bindSubExpRes (baseName dst <> "_bar") dst_bar_res   updateAdj dst $ head dst_bar    lam_mul''' <- renameLambda lam_mul'@@ -478,7 +467,7 @@           (eBody $ pure $ pure $ zeroExp $ rowType dst_type)    vs_bar <--    letExp (baseString vs <> "_bar") $ Op $ Screma n [is, vs] $ mapSOAC lam_vsbar+    letExp (baseName vs <> "_bar") $ Op $ Screma n [is, vs] $ mapSOAC lam_vsbar    updateAdj vs vs_bar @@ -502,7 +491,7 @@   let t = paramDec $ head $ lambdaParams add    dst_cpy <--    letExp (baseString dst <> "_copy") . BasicOp $+    letExp (baseName dst <> "_copy") . BasicOp $       Replicate mempty (Var dst)    f <- mkIdentityLambda [Prim int64, t]@@ -516,7 +505,7 @@   updateAdj dst x_bar    x_type <- lookupType x-  i_param <- newParam (baseString vs <> "_i") $ Prim int64+  i_param <- newParam (baseName vs <> "_i") $ Prim int64   let i = paramName i_param   lam_vsbar <-     mkLambda [i_param] $@@ -526,7 +515,7 @@           (eBody $ pure $ pure $ BasicOp $ Index x_bar $ fullSlice x_type [DimFix $ Var i])           (eBody $ pure $ eSubExp ne) -  vs_bar <- letExp (baseString vs <> "_bar") $ Op $ Screma n [is] $ mapSOAC lam_vsbar+  vs_bar <- letExp (baseName vs <> "_bar") $ Op $ Screma n [is] $ mapSOAC lam_vsbar   updateAdj vs vs_bar  -- Special case for vectorised combining operator. Rewrite@@ -684,7 +673,7 @@   nis <- letExp "nis" $ Op $ Screma n (bins : offsets) $ mapSOAC map_lam    scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) tps-  multiScatter n scatter_dst nis xs+  multiScatter scatter_dst nis xs   where     iConst c = eSubExp $ intConst Int64 c @@ -700,7 +689,7 @@   iters <- letSubExp "iters" =<< toExp (untyped (pe64 logw + 1) ~/~ untyped (pe64 (intConst Int64 2)))    types <- traverse lookupType xs-  params <- zipWithM (\x -> newParam (baseString x) . flip toDecl Nonunique) xs types+  params <- zipWithM (\x -> newParam (baseName x) . flip toDecl Nonunique) xs types   i <- newVName "i"   loopbody <- buildBody_ . localScope (scopeOfFParams params) $     fmap varsRes $ do@@ -790,7 +779,7 @@   nes <- traverse (letExp "new_dst" . BasicOp . Replicate (Shape $ pure $ head w)) ne    h_map <- mkIdentityLambda $ Prim int64 : map rowType as_type-  h_part <- traverse (newVName . flip (<>) "_h_part" . baseString) xs+  h_part <- traverse (newVName . flip (<>) "_h_part" . baseName) xs   auxing aux . letBindNames h_part . Op $     Hist n as [HistOp (Shape w) rf nes ne lam0] h_map @@ -935,7 +924,7 @@       =<< eLambda f_adj (map (eSubExp . Var) $ ls_arr <> sas <> rs_arr)    scatter_dst <- traverse (\t -> letExp "scatter_dst" $ BasicOp $ Scratch (elemType t) (arrayDims t)) as_type-  as_bar <- multiScatter n scatter_dst siota sas_bar+  as_bar <- multiScatter scatter_dst siota sas_bar    zipWithM_ updateAdj (tail as) as_bar   where

@@ -204,7 +204,7 @@      empty_saved_array <-       forM loop_params_to_copy $ \p ->-        letSubExp (baseString (paramName p) <> "_empty_saved")+        letSubExp (baseName (paramName p) <> "_empty_saved")           =<< eBlank (arrayOf (paramDec p) (Shape [bound64]) NoUniqueness)      (body', (saved_pats, saved_params)) <- buildBody $@@ -217,15 +217,15 @@             forM loop_params_to_copy $ \p -> do               let v = paramName p                   t = paramDec p-              saved_param_v <- newVName $ baseString v <> "_saved"-              saved_pat_v <- newVName $ baseString v <> "_saved"+              saved_param_v <- newVName $ baseName v <> "_saved"+              saved_pat_v <- newVName $ baseName v <> "_saved"               setLoopTape v saved_pat_v               let saved_param = Param mempty saved_param_v $ arrayOf t (Shape [bound64]) Unique                   saved_pat = PatElem saved_pat_v $ arrayOf t (Shape [bound64]) NoUniqueness               saved_update <-                 localScope (scopeOfFParams [saved_param])                   $ letInPlace-                    (baseString v <> "_saved_update")+                    (baseName v <> "_saved_update")                     saved_param_v                     (fullSlice (fromDecl $ paramDec saved_param) [DimFix i_i64])                   $ substituteNames copy_substs@@ -262,7 +262,7 @@      (i_rev, i_stms) <- collectStms $       localScope (scopeOfLoopForm form) $ do-        letExp (baseString i <> "_rev") $+        letExp (baseName i <> "_rev") $           BasicOp $             BinOp (Sub it OverflowWrap) bound_minus_one (Var i) 

@@ -10,7 +10,7 @@ import Futhark.Analysis.PrimExp.Convert import Futhark.Builder import Futhark.IR.SOACS-import Futhark.Tools+import Futhark.Tools hiding (withAcc) import Futhark.Transform.Rename import Futhark.Util (splitAt3) @@ -238,6 +238,6 @@           zero <- letSubExp "zero" $ zeroExp t           reduce <- reduceSOAC [Reduce Commutative lam [zero]]           contrib_sum <--            letExp (baseString v <> "_contrib_sum") . Op $+            letExp (baseName v <> "_contrib_sum") . Op $               Screma w [contribs] reduce           void $ updateAdj v contrib_sum

@@ -40,6 +40,7 @@     tabNest,     oneExp,     zeroExp,+    zeroArray,     unitAdjOfType,     addLambda,     --@@ -126,6 +127,9 @@   substituteNames m (AdjVal (Var v)) = AdjVal $ Var $ substituteNames m v   substituteNames _ adj = adj +-- | Create an array of the given shape and element type consisting of zeroes.+-- The shape may be empty, meaning this function can (despite its name) also+-- create non-arrays. zeroArray :: (MonadBuilder m) => Shape -> Type -> m VName zeroArray shape t   | shapeRank shape == 0 =@@ -241,7 +245,7 @@ insAdj v = setAdj v . AdjVal . Var  adjVName :: VName -> ADM VName-adjVName v = newVName (baseString v <> "_adj")+adjVName v = newVName (baseName v <> "_adj")  -- | Create copies of all arrays consumed in the given statement, and -- return statements which include copies of the consumed arrays.@@ -256,14 +260,17 @@             case v_t of               Array {} -> do                 v' <--                  letExp (baseString v <> "_ad_copy") . BasicOp $+                  letExp (baseName v <> "_ad_copy") . BasicOp $                     Replicate mempty (Var v)                 addSubstitution v' v                 pure [(v, v')]               _ -> pure mempty-       in M.fromList . mconcat-            <$> mapM onConsumed (namesToList $ consumedInStms $ fst (Alias.analyseStms mempty (oneStm stm))) +          consumed =+            namesToList . consumedInStms $+              fst (Alias.analyseStms mempty (oneStm stm))+       in M.fromList . mconcat <$> mapM onConsumed consumed+ copyConsumedArrsInBody :: [VName] -> Body SOACS -> ADM Substitutions copyConsumedArrsInBody dontCopy b =   mconcat <$> mapM onConsumed (filter (`notElem` dontCopy) $ namesToList $ consumedInBody (Alias.analyseBody mempty b))@@ -275,7 +282,7 @@         Array {} ->           M.singleton v             <$> letExp-              (baseString v <> "_ad_copy")+              (baseName v <> "_ad_copy")               (BasicOp $ Replicate mempty (Var v))         _ -> pure mempty @@ -324,7 +331,7 @@           Iota w (intConst Int64 0) (intConst Int64 1) Int64       iparam <- newParam "i" $ Prim int64       params <- forM vs $ \v ->-        newParam (baseString v <> "_p") . rowType =<< lookupType v+        newParam (baseName v <> "_p") . rowType =<< lookupType v       ((ret, res), stms) <- collectStms . localScope (scopeOfLParams (iparam : params)) $ do         res <- tabNest' (paramName iparam : is) (n - 1) (map paramName params) f         ret <- mapM lookupType res@@ -374,6 +381,8 @@       v_t <- lookupType v       case v_t of         Acc _ shape [Prim t] _ -> pure $ AdjZero shape t+        Acc _ shape [t] _ -> pure $ AdjZero (shape <> arrayShape t) (elemType t)+        Acc {} -> error $ "lookupAdj: Non-singleton accumulator adjoint: " <> prettyString v_t         _ -> pure $ AdjZero (arrayShape v_t) (elemType v_t)     Just v_adj -> pure v_adj @@ -413,11 +422,11 @@           _ -> do             let stms s = do                   v_adj_i <--                    letExp (baseString v_adj <> "_i") . BasicOp $+                    letExp (baseName v_adj <> "_i") . BasicOp $                       Index v_adj $                         fullSlice v_adj_t [DimFix i]                   se_update <- letSubExp "updated_adj_i" =<< addExp se_v v_adj_i-                  letExp (baseString v_adj) . BasicOp $+                  letExp (baseName v_adj) . BasicOp $                     Update s v_adj (fullSlice v_adj_t [DimFix i]) se_update             case check of               CheckBounds _ -> stms Safe@@ -442,7 +451,7 @@                 UpdateAcc safety v_adj' (map Var is) [Var d']           insAdj v v_adj'         _ -> do-          v_adj' <- letExp (baseString v <> "_adj") =<< addExp v_adj d+          v_adj' <- letExp (baseName v <> "_adj") =<< addExp v_adj d           insAdj v v_adj'  updateAdjSliceWithSafety :: Slice SubExp -> VName -> VName -> Safety -> ADM ()@@ -465,14 +474,14 @@             traverse (toSubExp "index") $               fixSlice (fmap pe64 slice) $                 map le64 is-          letTupExp (baseString v_adj') . BasicOp $+          letTupExp (baseName v_adj') . BasicOp $             UpdateAcc safety v_adj' slice' [Var d']       pure v_adj'     _ -> do       v_adjslice <-         if primType t           then pure v_adj-          else letExp (baseString v ++ "_slice") $ BasicOp $ Index v_adj slice+          else letExp (baseName v <> "_slice") $ BasicOp $ Index v_adj slice       letInPlace "updated_adj" v_adj slice =<< addExp v_adjslice d   insAdj v v_adj' 

@@ -25,11 +25,11 @@       BinOp (Sub Int64 OverflowUndef) w (intConst Int64 1)   let stride = intConst Int64 (-1)       slice = fullSlice arr_t [DimSlice start w stride]-  letExp (baseString arr <> "_rev") $ BasicOp $ Index arr slice+  letExp (baseName arr <> "_rev") $ BasicOp $ Index arr slice  scanExc ::   (MonadBuilder m, Rep m ~ SOACS) =>-  String ->+  Name ->   Scan SOACS ->   [VName] ->   m [VName]@@ -78,7 +78,7 @@   | Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,     isAdd op = do       adj_rep <--        letExp (baseString adj <> "_rep") $+        letExp (baseName adj <> "_rep") $           BasicOp $             Replicate (Shape [w]) $               Var adj@@ -183,7 +183,7 @@       BasicOp $         Iota w (intConst Int64 0) (intConst Int64 1) Int64   form <- reduceSOAC [Reduce Commutative red_lam [ne, intConst Int64 (-1)]]-  x_ind <- newVName (baseString x <> "_ind")+  x_ind <- newVName (baseName x <> "_ind")   auxing aux $ letBindNames [x, x_ind] $ Op $ Screma w [as, red_iota] form    m

@@ -8,7 +8,6 @@ import Futhark.AD.Rev.Monad import Futhark.AD.Rev.Reduce import Futhark.AD.Rev.Scan-import Futhark.AD.Rev.Scatter import Futhark.Analysis.PrimExp.Convert import Futhark.Builder import Futhark.IR.SOACS@@ -137,17 +136,6 @@       (mapstm, scanstm) <-         scanomapToMapAndScan pat (w, scans, map_lam, as)       vjpStm ops mapstm $ vjpStm ops scanstm m---- Differentiating Scatter-vjpSOAC ops pat aux (Scatter w ass written_info lam) m-  | isIdentityLambda lam =-      vjpScatter ops pat aux (w, ass, lam, written_info) m-  | otherwise = do-      map_idents <- mapM (\t -> newIdent "map_res" (arrayOfRow t w)) $ lambdaReturnType lam-      let map_stm = mkLet map_idents $ Op $ Screma w ass $ mapSOAC lam-      lam_id <- mkIdentityLambda $ lambdaReturnType lam-      let scatter_stm = Let pat aux $ Op $ Scatter w (map identName map_idents) written_info lam_id-      vjpStm ops map_stm $ vjpStm ops scatter_stm m  -- Differentiating Histograms vjpSOAC ops pat aux (Hist n as histops f) m

@@ -78,7 +78,7 @@         ( buildBody_ $ do             j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))             y_s <- forM ys_adj $ \y_ ->-              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ [eSubExp j]+              letSubExp (baseName y_ <> "_j") =<< eIndex y_ [eSubExp j]             let zso = orderArgs s y_s             let ido = orderArgs s $ caseJac k sc idmat             pure $ subExpsRes $ concat $ zipWith (++) zso $ fmap concat ido@@ -87,7 +87,7 @@             j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))             j1 <- letSubExp "j1" =<< toExp (pe64 w - le64 i)             y_s <- forM ys_adj $ \y_ ->-              letSubExp (baseString y_ ++ "_j") =<< eIndex y_ [eSubExp j]+              letSubExp (baseName y_ <> "_j") =<< eIndex y_ [eSubExp j]              let args =                   map (`eIndex` [eSubExp j]) ys ++ map (`eIndex` [eSubExp j1]) xs@@ -158,7 +158,7 @@    par_i <- newParam "i" $ Prim int64   let i = paramName par_i-  par_x <- zipWithM (\x -> newParam (baseString x ++ "_par_x")) xs eltps+  par_x <- zipWithM (\x -> newParam (baseName x <> "_par_x")) xs eltps    map_lam <-     mkLambda (par_i : par_x) $ do@@ -166,7 +166,7 @@        dj <-         forM ds $ \dd ->-          letExp (baseString dd ++ "_dj") =<< eIndex dd [eSubExp j]+          letExp (baseName dd <> "_dj") =<< eIndex dd [eSubExp j]        fmap varsRes . letTupExp "scan_contribs"         =<< eIf@@ -177,7 +177,7 @@                im1 <- letSubExp "im1" =<< toExp (le64 i - 1)               ys_im1 <- forM ys $ \y ->-                letSubExp (baseString y <> "_im1") =<< eIndex y [eSubExp im1]+                letSubExp (baseName y <> "_im1") =<< eIndex y [eSubExp im1]                let args = map eSubExp $ ys_im1 ++ map (Var . paramName) par_x               eLambda lam args@@ -269,8 +269,8 @@   as_types <- mapM lookupType as   let arg_type_row = map rowType as_types -  par_a1 <- zipWithM (\x -> newParam (baseString x <> "_par_a1")) as arg_type_row-  par_a2 <- zipWithM (\x -> newParam (baseString x <> "_par_a2")) as arg_type_row+  par_a1 <- zipWithM (\x -> newParam (baseName x <> "_par_a1")) as arg_type_row+  par_a2 <- zipWithM (\x -> newParam (baseName x <> "_par_a2")) as arg_type_row   -- Just the original operator but with par_a1 and par_a2 swapped.   rev_op <- mkLambda (par_a1 <> par_a2) $ do     op <- renameLambda $ scanLambda scan@@ -303,12 +303,12 @@ mkPPADOpLifted ops as scan = do   as_types <- mapM lookupType as   let arg_type_row = map rowType as_types-  par_x1 <- zipWithM (\x -> newParam (baseString x ++ "_par_x1")) as arg_type_row-  par_x2_unused <- zipWithM (\x -> newParam (baseString x ++ "_par_x2_unused")) as arg_type_row-  par_a1 <- zipWithM (\x -> newParam (baseString x ++ "_par_a1")) as arg_type_row-  par_a2 <- zipWithM (\x -> newParam (baseString x ++ "_par_a2")) as arg_type_row-  par_y1_h <- zipWithM (\x -> newParam (baseString x ++ "_par_y1_h")) as arg_type_row-  par_y2_h <- zipWithM (\x -> newParam (baseString x ++ "_par_y2_h")) as arg_type_row+  par_x1 <- zipWithM (\x -> newParam (baseName x <> "_par_x1")) as arg_type_row+  par_x2_unused <- zipWithM (\x -> newParam (baseName x <> "_par_x2_unused")) as arg_type_row+  par_a1 <- zipWithM (\x -> newParam (baseName x <> "_par_a1")) as arg_type_row+  par_a2 <- zipWithM (\x -> newParam (baseName x <> "_par_a2")) as arg_type_row+  par_y1_h <- zipWithM (\x -> newParam (baseName x <> "_par_y1_h")) as arg_type_row+  par_y2_h <- zipWithM (\x -> newParam (baseName x <> "_par_y2_h")) as arg_type_row    add_lams <- mapM addLambda arg_type_row @@ -373,9 +373,9 @@ finalMapPPAD ops as scan = do   as_types <- mapM lookupType as   let arg_type_row = map rowType as_types-  par_y_right <- zipWithM (\x -> newParam (baseString x ++ "_par_y_right")) as arg_type_row-  par_a <- zipWithM (\x -> newParam (baseString x ++ "_par_a")) as arg_type_row-  par_r_adj <- zipWithM (\x -> newParam (baseString x ++ "_par_r_adj")) as arg_type_row+  par_y_right <- zipWithM (\x -> newParam (baseName x <> "_par_y_right")) as arg_type_row+  par_a <- zipWithM (\x -> newParam (baseName x <> "_par_a")) as arg_type_row+  par_r_adj <- zipWithM (\x -> newParam (baseName x <> "_par_r_adj")) as arg_type_row    mkLambda (par_y_right ++ par_a ++ par_r_adj) $ do     op_bar_2 <- mkScanAdjointLam ops (scanLambda scan) WrtSecond (Var . paramName <$> par_r_adj)@@ -446,7 +446,7 @@      transp_as <-       forM as $ \a ->-        letExp (baseString a ++ "_transp") $ BasicOp $ Rearrange a rear+        letExp (baseName a <> "_transp") $ BasicOp $ Rearrange a rear      ts <- traverse lookupType transp_as     let n = arraysSize 0 ts

@@ -1,187 +0,0 @@-{-# LANGUAGE TypeFamilies #-}--module Futhark.AD.Rev.Scatter (vjpScatter) where--import Control.Monad-import Futhark.AD.Rev.Monad-import Futhark.Analysis.PrimExp.Convert-import Futhark.Builder-import Futhark.IR.SOACS-import Futhark.Tools-import Futhark.Util (chunk)--withinBounds :: [(SubExp, VName)] -> TPrimExp Bool VName-withinBounds [] = TPrimExp $ ValueExp (BoolValue True)-withinBounds [(q, i)] = (le64 i .<. pe64 q) .&&. (pe64 (intConst Int64 (-1)) .<. le64 i)-withinBounds (qi : qis) = withinBounds [qi] .&&. withinBounds qis---- Generates a potential tower-of-maps lambda body for an indexing operation.--- Assuming parameters:---   `arr`   the array that is indexed---   `[(w_1, i_1), (w_2, i_2), ..., (w_k, i_k)]` outer lambda formal parameters and their bounds---   `[n_1,n_2,...]ptp` the type of the index expression `arr[i_1,i_2,...,i_k]`--- Generates something like:--- (\ i_1 i_2 ->---    map (\j_1 -> ... if (i_1 >= 0 && i_1 < w_1) &&---                        (i_2 >= 0 && i_2 < w_2) && ...---                     then arr[i_1, i_2, ... j_1, ...]---                     else 0---        ) (iota n_1)--- )--- The idea is that you do not want to put under the `if` something---     that is an array because it would not flatten well!-genIdxLamBody :: VName -> [(SubExp, Param Type)] -> Type -> ADM (Body SOACS)-genIdxLamBody as wpis = genRecLamBody as wpis []-  where-    genRecLamBody :: VName -> [(SubExp, Param Type)] -> [Param Type] -> Type -> ADM (Body SOACS)-    genRecLamBody arr w_pis nest_pis (Array t (Shape []) _) =-      genRecLamBody arr w_pis nest_pis (Prim t)-    genRecLamBody arr w_pis nest_pis (Array t (Shape (s : ss)) _) = do-      new_ip <- newParam "i" (Prim int64)-      let t' = Prim t `arrayOfShape` Shape ss-      inner_lam <--        mkLambda [new_ip] $-          bodyBind =<< genRecLamBody arr w_pis (nest_pis ++ [new_ip]) t'-      let (_, orig_pis) = unzip w_pis-      buildBody_ . localScope (scopeOfLParams (orig_pis ++ nest_pis)) $ do-        iota_v <- letExp "iota" $ BasicOp $ Iota s (intConst Int64 0) (intConst Int64 1) Int64-        r <- letSubExp (baseString arr ++ "_elem") $ Op $ Screma s [iota_v] (mapSOAC inner_lam)-        pure [subExpRes r]-    genRecLamBody arr w_pis nest_pis (Prim ptp) = do-      let (ws, orig_pis) = unzip w_pis-      let inds = map paramName (orig_pis ++ nest_pis)-      localScope (scopeOfLParams (orig_pis ++ nest_pis)) $-        eBody-          [ eIf-              (toExp $ withinBounds $ zip ws $ map paramName orig_pis)-              ( do-                  r <- letSubExp "r" $ BasicOp $ Index arr $ Slice $ map (DimFix . Var) inds-                  resultBodyM [r]-              )-              (resultBodyM [Constant $ blankPrimValue ptp])-          ]-    genRecLamBody _ _ _ _ = error "In Rev.hs, helper function genRecLamBody, unreachable case reached!"------- Original:---   let ys = scatter xs is vs--- Assumes no duplicate indices in `is`--- Forward Sweep:---   let xs_save = gather xs is---   let ys = scatter xs is vs--- Return Sweep:---   let vs_ctrbs = gather is ys_adj---   let vs_adj \overline{+}= vs_ctrbs -- by map or generalized reduction---   let xs_adj = scatter ys_adj is \overline{0}---   let xs = scatter ys is xs_save-vjpScatter1 ::-  PatElem Type ->-  StmAux () ->-  (SubExp, [VName], (ShapeBase SubExp, Int, VName)) ->-  ADM () ->-  ADM ()-vjpScatter1 pys aux (w, ass, (shp, num_vals, xs)) m = do-  let rank = length $ shapeDims shp-      (all_inds, val_as) = splitAt (rank * num_vals) ass-      inds_as = chunk rank all_inds-  xs_t <- lookupType xs-  let val_t = stripArray (shapeRank shp) xs_t-  -- computing xs_save-  xs_saves <- mkGather inds_as xs xs_t-  -- performing the scatter-  id_lam <--    mkIdentityLambda $-      replicate (shapeRank shp) (Prim int64) ++ replicate (shapeRank shp) val_t-  addStm $ Let (Pat [pys]) aux $ Op $ Scatter w ass [(shp, num_vals, xs)] id_lam-  m-  let ys = patElemName pys-  -- XXX: Since our restoration of xs will consume ys, we have to-  -- make a copy of ys in the chance that it is actually the result-  -- of the program.  In that case the asymptotics will not be-  -- (locally) preserved, but since ys must necessarily have been-  -- constructed somewhere close, they are probably globally OK.-  ys_copy <--    letExp (baseString ys <> "_copy") . BasicOp $-      Replicate mempty (Var ys)-  returnSweepCode $ do-    ys_adj <- lookupAdjVal ys-    -- computing vs_ctrbs and updating vs_adj-    vs_ctrbs <- mkGather inds_as ys_adj xs_t-    zipWithM_ updateAdj val_as vs_ctrbs -- use Slice?-    -- creating xs_adj-    zeros <--      replicateM (length val_as) . letExp "zeros" $-        zeroExp $-          xs_t `setOuterSize` w-    let f_tps = replicate (rank * num_vals) (Prim int64) ++ replicate num_vals val_t-    f <- mkIdentityLambda f_tps-    xs_adj <--      letExp (baseString xs ++ "_adj") . Op $-        Scatter w (all_inds ++ zeros) [(shp, num_vals, ys_adj)] f-    insAdj xs xs_adj -- reusing the ys_adj for xs_adj!-    f' <- mkIdentityLambda f_tps-    xs_rc <--      auxing aux . letExp (baseString xs <> "_rc") . Op $-        Scatter w (all_inds ++ xs_saves) [(shp, num_vals, ys)] f'-    addSubstitution xs xs_rc-    addSubstitution ys ys_copy-  where-    -- Creates a potential map-nest that indexes in full the array,-    --   and applies the condition of indices within bounds at the-    --   deepest level in the nest so that everything can be parallel.-    mkGather :: [[VName]] -> VName -> Type -> ADM [VName]-    mkGather inds_as arr arr_t = do-      ips <- forM inds_as $ \idxs ->-        mapM (\idx -> newParam (baseString idx ++ "_elem") (Prim int64)) idxs--      gather_lam <- mkLambda (concat ips) . fmap mconcat . forM ips $ \idxs -> do-        let q = length idxs-            (ws, eltp) = (take q $ arrayDims arr_t, stripArray q arr_t)-        bodyBind =<< genIdxLamBody arr (zip ws idxs) eltp-      let soac = Screma w (concat inds_as) (mapSOAC gather_lam)-      letTupExp (baseString arr ++ "_gather") $ Op soac--vjpScatter ::-  VjpOps ->-  Pat Type ->-  StmAux () ->-  (SubExp, [VName], Lambda SOACS, [(Shape, Int, VName)]) ->-  ADM () ->-  ADM ()-vjpScatter ops (Pat pes) aux (w, ass, lam, written_info) m-  | isIdentityLambda lam,-    [(shp, num_vals, xs)] <- written_info,-    [pys] <- pes =-      vjpScatter1 pys aux (w, ass, (shp, num_vals, xs)) m-  | isIdentityLambda lam = do-      let sind = splitInd written_info-          (inds, vals) = splitAt sind ass-      lst_stms <- chunkScatterInps (inds, vals) (zip pes written_info)-      diffScatters (stmsFromList lst_stms)-  | otherwise =-      error "vjpScatter: cannot handle"-  where-    splitInd [] = 0-    splitInd ((shp, num_res, _) : rest) =-      num_res * length (shapeDims shp) + splitInd rest-    chunkScatterInps (acc_inds, acc_vals) [] =-      case (acc_inds, acc_vals) of-        ([], []) -> pure []-        _ -> error "chunkScatterInps: cannot handle"-    chunkScatterInps-      (acc_inds, acc_vals)-      ((pe, info@(shp, num_vals, _)) : rest) = do-        let num_inds = num_vals * length (shapeDims shp)-            (curr_inds, other_inds) = splitAt num_inds acc_inds-            (curr_vals, other_vals) = splitAt num_vals acc_vals-        vtps <- mapM lookupType curr_vals-        f <- mkIdentityLambda (replicate num_inds (Prim int64) ++ vtps)-        let stm =-              Let (Pat [pe]) aux . Op $-                Scatter w (curr_inds ++ curr_vals) [info] f-        stms_rest <- chunkScatterInps (other_inds, other_vals) rest-        pure $ stm : stms_rest-    diffScatters all_stms-      | Just (stm, stms) <- stmsHead all_stms =-          vjpStm ops stm $ diffScatters stms-      | otherwise = m

@@ -28,6 +28,7 @@  import Control.Monad import Control.Monad.IO.Class+import Control.Monad.State (get) import Data.Bifunctor import Data.List (intercalate) import Data.Map qualified as M@@ -62,19 +63,26 @@ import Futhark.IR.SeqMem (SeqMem) import Futhark.Optimise.Fusion.GraphRep qualified import Futhark.Util (runProgramWithExitCode, unixEnvironment)+import Futhark.Util.Pretty (Doc, pretty, putDocLn, (</>)) import Futhark.Version (versionString) import System.Directory import System.Exit-import System.FilePath+import System.FilePath hiding ((</>)) import System.Info qualified +-- | Convert the program and compiler state to a textual form.+progDoc :: (PrettyRep rep) => Prog rep -> FutharkM (Doc a)+progDoc prog = do+  state <- get+  pure $ pretty state </> mempty </> pretty prog+ -- | Print the result to stdout. printAction :: (ASTRep rep) => Action rep printAction =   Action     { actionName = "Prettyprint",       actionDescription = "Prettyprint the resulting internal representation on standard output.",-      actionProcedure = liftIO . putStrLn . prettyString+      actionProcedure = liftIO . putDocLn <=< progDoc     }  -- | Print the result to stdout, alias annotations.@@ -83,7 +91,7 @@   Action     { actionName = "Prettyprint",       actionDescription = "Prettyprint the resulting internal representation on standard output.",-      actionProcedure = liftIO . putStrLn . prettyString . aliasAnalysis+      actionProcedure = liftIO . putDocLn . pretty . aliasAnalysis     }  -- | Print last use information to stdout.

@@ -554,7 +554,7 @@           . unSegSpace           $ segSpace op    in -- Analyse statements in the SegOp body-      analyseStms segspace_context (SegOpName . segOpType op) pats . stmsToList . kernelBodyStms $ segBody op+      analyseStms segspace_context (SegOpName . segOpType op) pats . stmsToList . bodyStms $ segBody op  analyseSizeOp :: SizeOp -> Context rep -> [VName] -> (Context rep, IndexTable rep) analyseSizeOp op ctx pats =

@@ -49,10 +49,10 @@  -- | Perform alias analysis on Body. analyseBody ::-  (AliasableRep rep) =>+  (AliasableRep rep, IsResult res) =>   AliasTable ->-  Body rep ->-  Body (Aliases rep)+  GBody rep res ->+  GBody (Aliases rep) res analyseBody atable (Body rep stms result) =   let (stms', _atable') = analyseStms atable stms    in mkAliasedBody rep stms' result

@@ -151,7 +151,7 @@                     Nothing               boundUsedInBody =                 mapMaybe isBound $ namesToList $ freeIn lam-          newParams <- mapM (newIdent' (++ "_wasfree")) boundUsedInBody+          newParams <- mapM (newIdent' (<> "_wasfree")) boundUsedInBody           let subst =                 M.fromList $                   zip (map identName boundUsedInBody) (map identName newParams)@@ -206,10 +206,10 @@     inspect (_, SOAC.Input ts v _)       | Just (p, pInp) <- find (isParam v) ourInps = do           let pInp' = SOAC.transformRows ts pInp-          p' <- newNameFromString $ baseString p+          p' <- newName p           pure (p', pInp')     inspect (param, SOAC.Input ts a t) = do-      param' <- newNameFromString (baseString param ++ "_rep")+      param' <- newVName (baseName param <> "_rep")       pure (param', SOAC.Input (ts SOAC.|> SOAC.Replicate mempty (Shape [w])) a t)  -- | Reshape a map nest. It is assumed that any validity tests have@@ -231,9 +231,7 @@       | shapeRank nest_shape == 0 =           pure $ MapNest w map_lam nests inps'       | otherwise = do-          nest_params <--            mapM (newVName . baseString . paramName) $-              lambdaParams map_lam+          nest_params <- mapM (newName . paramName) $ lambdaParams map_lam           res <-             replicateM               (length $ lambdaReturnType map_lam)

@@ -86,7 +86,6 @@ import Futhark.IR qualified as Futhark import Futhark.IR.SOACS.SOAC   ( HistOp (..),-    ScatterSpec,     ScremaForm (..),     scremaType,   )@@ -399,7 +398,6 @@ -- | A definite representation of a SOAC expression. data SOAC rep   = Stream SubExp [Input] [SubExp] (Lambda rep)-  | Scatter SubExp [Input] (ScatterSpec VName) (Lambda rep)   | Screma SubExp [Input] (ScremaForm rep)   | Hist SubExp [Input] [HistOp rep] (Lambda rep)   deriving (Eq, Show)@@ -407,7 +405,6 @@ -- | Returns the inputs used in a SOAC. inputs :: SOAC rep -> [Input] inputs (Stream _ arrs _ _) = arrs-inputs (Scatter _ arrs _lam _spec) = arrs inputs (Screma _ arrs _) = arrs inputs (Hist _ inps _ _) = inps @@ -415,8 +412,6 @@ setInputs :: [Input] -> SOAC rep -> SOAC rep setInputs arrs (Stream w _ nes lam) =   Stream (newWidth arrs w) arrs nes lam-setInputs arrs (Scatter w _ lam spec) =-  Scatter (newWidth arrs w) arrs lam spec setInputs arrs (Screma w _ form) =   Screma w arrs form setInputs inps (Hist w _ ops lam) =@@ -429,7 +424,6 @@ -- | The lambda used in a given SOAC. lambda :: SOAC rep -> Lambda rep lambda (Stream _ _ _ lam) = lam-lambda (Scatter _len _ivs _spec lam) = lam lambda (Screma _ _ (ScremaForm lam _ _)) = lam lambda (Hist _ _ _ lam) = lam @@ -437,8 +431,6 @@ setLambda :: Lambda rep -> SOAC rep -> SOAC rep setLambda lam (Stream w arrs nes _) =   Stream w arrs nes lam-setLambda lam (Scatter len arrs spec _lam) =-  Scatter len arrs spec lam setLambda lam (Screma w arrs (ScremaForm _ scan red)) =   Screma w arrs (ScremaForm lam scan red) setLambda lam (Hist w ops inps _) =@@ -453,12 +445,6 @@           | t <- drop (length nes) (lambdaReturnType lam)         ]    in accrtps ++ arrtps-typeOf (Scatter _w _ivs dests lam) =-  zipWith arrayOfShape val_ts ws-  where-    indexes = sum $ zipWith (*) ns $ map length ws-    val_ts = drop indexes $ lambdaReturnType lam-    (ws, ns, _) = unzip3 dests typeOf (Screma w _ form) =   scremaType w form typeOf (Hist _ _ ops _) = do@@ -469,7 +455,6 @@ -- inputs _after_ input-transforms have been carried out. width :: SOAC rep -> SubExp width (Stream w _ _ _) = w-width (Scatter len _lam _ivs _as) = len width (Screma w _ _) = w width (Hist w _ _ _) = w @@ -484,8 +469,6 @@ toSOAC :: (MonadBuilder m) => SOAC (Rep m) -> m (Futhark.SOAC (Rep m)) toSOAC (Stream w inps nes lam) =   Futhark.Stream w <$> inputsToSubExps inps <*> pure nes <*> pure lam-toSOAC (Scatter w ivs dests lam) =-  Futhark.Scatter w <$> inputsToSubExps ivs <*> pure dests <*> pure lam toSOAC (Screma w arrs form) =   Futhark.Screma w <$> inputsToSubExps arrs <*> pure form toSOAC (Hist w arrs ops lam) =@@ -507,8 +490,6 @@   m (Either NotSOAC (SOAC rep)) fromExp (Op (Futhark.Stream w as nes lam)) =   Right <$> (Stream w <$> traverse varInput as <*> pure nes <*> pure lam)-fromExp (Op (Futhark.Scatter w arrs spec lam)) =-  Right <$> (Scatter w <$> traverse varInput arrs <*> pure spec <*> pure lam) fromExp (Op (Futhark.Screma w arrs form)) =   Right <$> (Screma w <$> traverse varInput arrs <*> pure form) fromExp (Op (Futhark.Hist w arrs ops lam)) =@@ -712,4 +693,3 @@   pretty (Screma w arrs form) = Futhark.ppScrema w arrs form   pretty (Hist len imgs ops bucket_fun) = Futhark.ppHist len imgs ops bucket_fun   pretty (Stream w arrs nes lam) = Futhark.ppStream w arrs nes lam-  pretty (Scatter w arrs dests lam) = Futhark.ppScatter w arrs dests lam

@@ -128,7 +128,7 @@   InUse ->   KernelBody GPUMem ->   m (InUse, LastUsed, Graph VName)-analyseKernelBody lumap inuse body = analyseStms lumap inuse $ kernelBodyStms body+analyseKernelBody lumap inuse body = analyseStms lumap inuse $ bodyStms body  analyseBody ::   (LocalScope GPUMem m) =>@@ -274,11 +274,11 @@     memSizesExp :: (LocalScope GPUMem m) => Exp GPUMem -> m (Map VName Int)     memSizesExp (Op (Inner (SegOp segop))) =       let body = segBody segop-       in inScopeOf (kernelBodyStms body)+       in inScopeOf (bodyStms body)             $ fmap mconcat               <$> mapM memSizesStm             $ stmsToList-            $ kernelBodyStms body+            $ bodyStms body     memSizesExp (Match _ cases defbody _) = do       mconcat <$> mapM (memSizes . bodyStms) (defbody : map caseBody cases)     memSizesExp (Loop _ _ body) =@@ -295,7 +295,7 @@       M.singleton name sp     getSpacesStm (Let _ _ (Op (Alloc _ _))) = error "impossible"     getSpacesStm (Let _ _ (Op (Inner (SegOp segop)))) =-      foldMap getSpacesStm $ kernelBodyStms $ segBody segop+      foldMap getSpacesStm $ bodyStms $ segBody segop     getSpacesStm (Let _ _ (Match _ cases defbody _)) =       foldMap (foldMap getSpacesStm . bodyStms) $ defbody : map caseBody cases     getSpacesStm (Let _ _ (Loop _ _ body)) =

@@ -165,7 +165,7 @@   --      (i) an updated last-use table,   --     (ii) an updated set of used names (including the binding).   LastUseM rep (LUTabFun, Names)-lastUseKernelBody bdy@(KernelBody _ stms result) (lutab, used_nms) =+lastUseKernelBody bdy@(Body _ stms result) (lutab, used_nms) =   inScopeOf stms $ do     -- perform analysis bottom-up in bindings: results are known to be used,     -- hence they are added to the used_nms set.

@@ -67,13 +67,13 @@  analyzeHostOp :: MemAliases -> HostOp NoOp GPUMem -> MemAliasesM (HostOp NoOp GPUMem) MemAliases analyzeHostOp m (SegOp (SegMap _ _ _ kbody)) =-  analyzeStms (kernelBodyStms kbody) m+  analyzeStms (bodyStms kbody) m analyzeHostOp m (SegOp (SegRed _ _ _ kbody _)) =-  analyzeStms (kernelBodyStms kbody) m+  analyzeStms (bodyStms kbody) m analyzeHostOp m (SegOp (SegScan _ _ _ kbody _)) =-  analyzeStms (kernelBodyStms kbody) m+  analyzeStms (bodyStms kbody) m analyzeHostOp m (SegOp (SegHist _ _ _ kbody _)) =-  analyzeStms (kernelBodyStms kbody) m+  analyzeStms (bodyStms kbody) m analyzeHostOp m SizeOp {} = pure m analyzeHostOp m GPUBody {} = pure m analyzeHostOp m (OtherOp NoOp) = pure m

@@ -90,7 +90,7 @@ funDefMetrics = bodyMetrics . funDefBody  -- | Compute metrics for this body.-bodyMetrics :: (OpMetrics (Op rep)) => Body rep -> MetricsM ()+bodyMetrics :: (OpMetrics (Op rep)) => GBody rep res -> MetricsM () bodyMetrics = mapM_ stmMetrics . bodyStms  -- | Compute metrics for this statement.

@@ -70,9 +70,9 @@   Scope rep ->   KernelBody rep ->   State PrimExpTable ()-kernelToBodyPrimExps scope kbody = mapM_ (stmToPrimExps scope') (kernelBodyStms kbody)+kernelToBodyPrimExps scope kbody = mapM_ (stmToPrimExps scope') (bodyStms kbody)   where-    scope' = scope <> scopeOf (kernelBodyStms kbody)+    scope' = scope <> scopeOf (bodyStms kbody)  -- | Adds a statement to the PrimExpTable. If it can't be resolved as a `PrimExp`, -- it will be added as `Nothing`.

@@ -45,10 +45,10 @@     Exp rep ->     m (ExpDec rep)   mkBodyB ::-    (MonadBuilder m, Rep m ~ rep) =>+    (MonadBuilder m, Rep m ~ rep, IsResult res) =>     Stms rep ->-    Result ->-    m (Body rep)+    [res] ->+    m (GBody rep res)   mkLetNamesB ::     (MonadBuilder m, Rep m ~ rep) =>     [VName] ->@@ -63,10 +63,10 @@   mkExpDecB pat e = pure $ mkExpDec pat e    default mkBodyB ::-    (MonadBuilder m, Buildable rep) =>+    (MonadBuilder m, Buildable rep, IsResult res) =>     Stms rep ->-    Result ->-    m (Body rep)+    [res] ->+    m (GBody rep res)   mkBodyB stms res = pure $ mkBody stms res    default mkLetNamesB ::@@ -201,10 +201,11 @@   ( Buildable rep,     MonadFreshNames m,     HasScope somerep m,-    SameScope somerep rep+    SameScope somerep rep,+    IsResult res   ) =>-  Builder rep Result ->-  m (Body rep)+  Builder rep [res] ->+  m (GBody rep res) runBodyBuilder =   fmap (uncurry $ flip insertStms) . runBuilder . fmap (mkBody mempty) 

@@ -43,7 +43,7 @@   where   mkExpPat :: [Ident] -> Exp rep -> Pat (LetDec rep)   mkExpDec :: Pat (LetDec rep) -> Exp rep -> ExpDec rep-  mkBody :: Stms rep -> Result -> Body rep+  mkBody :: (IsResult res) => Stms rep -> [res] -> GBody rep res   mkLetNames ::     (MonadFreshNames m, HasScope rep m) =>     [VName] ->@@ -72,7 +72,7 @@   where   type Rep m :: Data.Kind.Type   mkExpDecM :: Pat (LetDec (Rep m)) -> Exp (Rep m) -> m (ExpDec (Rep m))-  mkBodyM :: Stms (Rep m) -> Result -> m (Body (Rep m))+  mkBodyM :: (IsResult res) => Stms (Rep m) -> [res] -> m (GBody (Rep m) res)   mkLetNamesM :: [VName] -> Exp (Rep m) -> m (Stm (Rep m))    -- | Add a statement to the 'Stms' under construction.@@ -170,9 +170,9 @@   pure res  -- | Add several bindings at the outermost level of a t'Body'.-insertStms :: (Buildable rep) => Stms rep -> Body rep -> Body rep+insertStms :: (Buildable rep, IsResult res) => Stms rep -> GBody rep res -> GBody rep res insertStms stms1 (Body _ stms2 res) = mkBody (stms1 <> stms2) res  -- | Add a single binding at the outermost level of a t'Body'.-insertStm :: (Buildable rep) => Stm rep -> Body rep -> Body rep+insertStm :: (Buildable rep, IsResult res) => Stm rep -> GBody rep res -> GBody rep res insertStm = insertStms . oneStm

@@ -32,6 +32,7 @@ import Futhark.Internalise.LiftLambdas as LiftLambdas import Futhark.Internalise.Monomorphise as Monomorphise import Futhark.Internalise.ReplaceRecords as ReplaceRecords+import Futhark.MonadFreshNames import Futhark.Optimise.ArrayLayout import Futhark.Optimise.ArrayShortCircuiting qualified as ArrayShortCircuiting import Futhark.Optimise.CSE@@ -829,13 +830,14 @@         Pipeline {} -> do           let (base, ext) = splitExtension file -              readCore parse construct = do+              readIR parse construct = do                 logMsg $ "Reading " <> file <> "..."                 input <- liftIO $ T.readFile file                 logMsg ("Parsing..." :: T.Text)                 case parse file input of                   Left err -> externalErrorS $ T.unpack err-                  Right prog -> do+                  Right (src, prog) -> do+                    modifyNameSource $ const ((), src)                     logMsg ("Typechecking..." :: T.Text)                     case checkProg $ Alias.aliasAnalysis prog of                       Left err -> externalErrorS $ show err@@ -847,13 +849,13 @@                       prog <- runPipelineOnProgram (futharkConfig config) id file                       runPolyPasses config base (SOACS prog)                   ),-                  (".fut_soacs", readCore parseSOACS SOACS),-                  (".fut_seq", readCore parseSeq Seq),-                  (".fut_seq_mem", readCore parseSeqMem SeqMem),-                  (".fut_gpu", readCore parseGPU GPU),-                  (".fut_gpu_mem", readCore parseGPUMem GPUMem),-                  (".fut_mc", readCore parseMC MC),-                  (".fut_mc_mem", readCore parseMCMem MCMem)+                  (".fut_soacs", readIR parseSOACS SOACS),+                  (".fut_seq", readIR parseSeq Seq),+                  (".fut_seq_mem", readIR parseSeqMem SeqMem),+                  (".fut_gpu", readIR parseGPU GPU),+                  (".fut_gpu_mem", readIR parseGPUMem GPUMem),+                  (".fut_mc", readIR parseMC MC),+                  (".fut_mc_mem", readIR parseMCMem MCMem)                 ]           case lookup ext handlers of             Just handler -> handler

@@ -40,7 +40,6 @@         if cfgCheck cfg           then unless (docText doc == file_s) $ do             T.hPutStrLn stderr $ T.pack file <> ": not formatted correctly."-            T.hPutStr stderr $ docText doc             exitFailure           else withFile file WriteMode $ \h -> hPutDoc h doc   where

@@ -40,6 +40,7 @@ import Futhark.Test.Values import Futhark.Util   ( directoryContents,+    ensureCacheDirectory,     fancyTerminal,     hashText,     nubOrd,@@ -52,7 +53,6 @@ import Futhark.Util.ProgressBar import System.Directory   ( copyFile,-    createDirectoryIfMissing,     doesFileExist,     getCurrentDirectory,     removePathForcibly,@@ -637,28 +637,36 @@     _ -> throwError "$loadImg failed to detect the number of channels in the audio input"  literateBuiltin :: EvalBuiltin ScriptM-literateBuiltin "loadimg" vs =-  case vs of-    [ValueAtom v]-      | Just path <- getValue v -> do+literateBuiltin server "loadimg" vs+  | [v] <- vs = do+      v' <- getHaskellValue server v+      case v' of+        Just path -> do           let path' = map (chr . fromIntegral) (path :: [Word8])-          loadImage path'-    _ ->+          valToExpValue <$> loadImage path'+        _ -> bad+  | otherwise = bad+  where+    bad =       throwError $         "$loadimg does not accept arguments of types: "-          <> T.intercalate ", " (map (prettyText . fmap valueType) vs)-literateBuiltin "loadaudio" vs =-  case vs of-    [ValueAtom v]-      | Just path <- getValue v -> do+          <> T.intercalate ", " (map (prettyText . fmap scriptValueType) vs)+literateBuiltin server "loadaudio" vs+  | [v] <- vs = do+      v' <- getHaskellValue server v+      case v' of+        Just path -> do           let path' = map (chr . fromIntegral) (path :: [Word8])-          loadAudio path'-    _ ->+          valToExpValue <$> loadAudio path'+        _ -> bad+  | otherwise = bad+  where+    bad =       throwError $         "$loadaudio does not accept arguments of types: "-          <> T.intercalate ", " (map (prettyText . fmap valueType) vs)-literateBuiltin f vs =-  scriptBuiltin "." f vs+          <> T.intercalate ", " (map (prettyText . fmap scriptValueType) vs)+literateBuiltin server f vs =+  scriptBuiltin "." server f vs  -- | Some of these only make sense for @futhark literate@, but enough -- are also sensible for @futhark script@ that we can share them.@@ -703,7 +711,7 @@   let fname_base = fromMaybe (T.unpack (envHash env) <> "-" <> template) fname_desired       fname = envImgDir env </> fname_base   exists <- liftIO $ doesFileExist fname-  liftIO $ createDirectoryIfMissing True $ envImgDir env+  liftIO $ ensureCacheDirectory $ envImgDir env   when (exists && scriptVerbose (envOpts env) > 0) $     liftIO . T.hPutStrLn stderr $       "Using existing file: " <> T.pack fname

@@ -2,7 +2,7 @@ module Futhark.CLI.Script (main) where  import Control.Monad.Except-import Control.Monad.IO.Class (MonadIO, liftIO)+import Control.Monad.IO.Class (MonadIO) import Data.Binary qualified as Bin import Data.ByteString.Lazy.Char8 qualified as BS import Data.Char (chr)@@ -15,7 +15,7 @@     scriptCommandLineOptions,   ) import Futhark.Script-import Futhark.Test.Values (Compound (..), getValue, valueType)+import Futhark.Test.Values (Compound (..)) import Futhark.Util.Options import Futhark.Util.Pretty (prettyText) import System.Exit@@ -86,17 +86,23 @@  -- A few extra procedures that are not handled by scriptBuiltin. extScriptBuiltin :: (MonadError T.Text m, MonadIO m) => EvalBuiltin m-extScriptBuiltin "store" [ValueAtom fv, ValueAtom vv]-  | Just path <- getValue fv = do-      let path' = map (chr . fromIntegral) (path :: [Bin.Word8])-      liftIO $ BS.writeFile path' $ Bin.encode vv-      pure $ ValueTuple []-extScriptBuiltin "store" vs =-  throwError $-    "$store does not accept arguments of types: "-      <> T.intercalate ", " (map (prettyText . fmap valueType) vs)-extScriptBuiltin f vs =-  scriptBuiltin "." f vs+extScriptBuiltin server "store" vs+  | [fv, v@(ValueAtom (SValue _ _))] <- vs = do+      fv' <- getHaskellValue server fv+      case fv' of+        Just path -> do+          let path' = map (chr . fromIntegral) (path :: [Bin.Word8])+          storeExpValue server path' v+          pure $ ValueTuple []+        _ -> bad+  | otherwise = bad+  where+    bad =+      throwError $+        "$store does not accept arguments of types: "+          <> T.intercalate ", " (map (prettyText . fmap scriptValueType) vs)+extScriptBuiltin server f vs =+  scriptBuiltin "." server f vs  -- | Run @futhark script@. main :: String -> [String] -> IO ()@@ -114,7 +120,7 @@             vs <- mapM (evalExp extScriptBuiltin s) scripts'             case reverse vs of               [] -> pure Nothing-              v : _ -> Just <$> getExpValue s v+              v : _ -> Just <$> getExpValue s v <* freeValue s v         case r of           Left e -> do             T.hPutStrLn stderr e

@@ -212,7 +212,7 @@   wrap <- memNeedsWrapping v   if wrap     then do-      v' <- newVName $ baseString v <> "_struct"+      v' <- newVName $ baseName v <> "_struct"       item [C.citem|$ty:(fatMemType DefaultSpace) $id:v' = {.references = NULL, .mem = $exp:v};|]       pure (v', [C.cstms|$id:v = $id:v'.mem;|])     else pure (v, mempty)@@ -360,7 +360,7 @@   let ct = primTypeToCType t   decl [C.cdecl|$tyquals:(volQuals vol) $ty:ct $id:name;|] compileCode (DeclareArray name t vs) = do-  name_realtype <- newVName $ baseString name ++ "_realtype"+  name_realtype <- newVName $ baseName name <> "_realtype"   let ct = primTypeToCType t   case vs of     ArrayValues vs' -> do

@@ -47,11 +47,11 @@ compileOutput :: Param -> CompilerM op s (C.Param, C.Exp) compileOutput (ScalarParam name bt) = do   let ctp = primTypeToCType bt-  p_name <- newVName $ "out_" ++ baseString name+  p_name <- newVName $ "out_" <> baseName name   pure ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|]) compileOutput (MemParam name space) = do   ty <- memToCType name space-  p_name <- newVName $ baseString name ++ "_p"+  p_name <- newVName $ baseName name <> "_p"   pure ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|])  compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func)

@@ -658,7 +658,7 @@   let cached = M.keys $ M.filter (== DefaultSpace) lexical    cached' <- forM cached $ \mem -> do-    size <- newVName $ prettyString mem <> "_cached_size"+    size <- newVName $ nameFromText (prettyText mem) <> "_cached_size"     pure (mem, size)    let lexMem env =

@@ -312,7 +312,7 @@           then [C.citems|v->$id:(tupleField j) = obj->$id:(tupleField i);|]           else             [C.citems|v->$id:(tupleField j) = malloc(sizeof(*v->$id:(tupleField j)));-                      *v->$id:(tupleField j) = *obj->$id:(tupleField i);+                      memcpy(v->$id:(tupleField j), obj->$id:(tupleField i), sizeof(*obj->$id:(tupleField i)));                       (void)(*(v->$id:(tupleField j)->mem.references))++;|]   pure     ( [C.cty|$ty:ct *|],@@ -360,8 +360,12 @@   | r == 0 =       [C.cstm|v->$id:(tupleField i) = $exp:e;|]   | otherwise =+      -- We use a memcpy instead of a straight assignment because the types may+      -- not actually be exactly the same - this is because we ignore array+      -- signedness when representing the payload of opaque types. However, the+      -- array types will have the same layout, so we can copy like this.       [C.cstm|{v->$id:(tupleField i) = malloc(sizeof(*$exp:e));-               *v->$id:(tupleField i) = *$exp:e;+               memcpy(v->$id:(tupleField i), $exp:e, sizeof(*$exp:e));                (void)(*(v->$id:(tupleField i)->mem.references))++;}|]  recordNewSetFields ::@@ -394,7 +398,7 @@               ( param_name,                 [C.cparam|const $ty:ct* $id:param_name|],                 [C.citem|{v->$id:(tupleField offset) = malloc(sizeof($ty:ct));-                          *v->$id:(tupleField offset) = *$id:param_name;+                          memcpy(v->$id:(tupleField offset), $id:param_name, sizeof($ty:ct));                           (void)(*(v->$id:(tupleField offset)->mem.references))++;}|]               )             )@@ -637,6 +641,8 @@    zipWithM onVariant [0 :: Int ..] variants   where+    unary = length variants == 1+     constructFunction ops ctx_ty opaque_ty i fname payload = do       (params, new_stms) <- unzip <$> zipWithM constructPayload [0 ..] payload @@ -644,6 +650,9 @@       set_unused_stms <-         mapM setUnused $ filter ((`notElem` used) . fst) (zip [0 ..] vds) +      let set_variant =+            [[C.cstm|v->$id:(tupleField 0) = $int:i;|] | not unary]+       headerDecl         (OpaqueDecl desc)         [C.cedecl|int $id:fname($ty:ctx_ty *ctx,@@ -656,7 +665,7 @@                                 $params:params) {                     (void)ctx;                     $ty:opaque_ty* v = malloc(sizeof($ty:opaque_ty));-                    v->$id:(tupleField 0) = $int:i;+                    $stms:set_variant                     { $items:(criticalSection ops new_stms) }                     // Set other fields                     { $items:set_unused_stms }@@ -707,6 +716,8 @@      destructFunction ops ctx_ty opaque_ty i fname payload = do       (params, destruct_stms) <- unzip <$> zipWithM (destructPayload ops) [0 ..] payload+      let check_stms =+            [[C.cstm|assert(obj->$id:(tupleField 0) == $int:i);|] | not unary]       headerDecl         (OpaqueDecl desc)         [C.cedecl|int $id:fname($ty:ctx_ty *ctx,@@ -718,7 +729,7 @@                                 $params:params,                                 const $ty:opaque_ty *obj) {                     (void)ctx;-                    assert(obj->$id:(tupleField 0) == $int:i);+                    $stms:check_stms                     $stms:destruct_stms                     return FUTHARK_SUCCESS;                   }|]@@ -734,20 +745,25 @@                   }|]         ) -sumVariantFunction :: Name -> CompilerM op s Manifest.CFuncName-sumVariantFunction desc = do+sumVariantFunction :: Int -> Name -> CompilerM op s Manifest.CFuncName+sumVariantFunction num_cs desc = do   opaque_ty <- opaqueToCType desc   ctx_ty <- contextType   variant <- publicName $ "variant_" <> opaqueName desc   headerDecl     (OpaqueDecl desc)     [C.cedecl|int $id:variant($ty:ctx_ty *ctx, const $ty:opaque_ty* v);|]-  -- This depends on the assumption that the first value always-  -- encodes the variant.+  -- This depends on the assumption that the first value always encodes the+  -- variant, which means we need to treat the unary case specially.+  let e =+        if num_cs == 1+          then [C.cexp|0|]+          else+            [C.cexp|v->$id:(tupleField 0)|]   libDecl     [C.cedecl|int $id:variant($ty:ctx_ty *ctx, const $ty:opaque_ty* v) {-                (void)ctx;-                return v->$id:(tupleField 0);+                (void)ctx; (void)v;+                return $exp:e;               }|]   pure variant @@ -764,7 +780,7 @@   Just . Manifest.OpaqueSum     <$> ( Manifest.SumOps             <$> sumVariants desc cs vds-            <*> sumVariantFunction desc+            <*> sumVariantFunction (length cs) desc         ) opaqueExtraOps _ types desc (OpaqueRecord fs) vds =   Just . Manifest.OpaqueRecord

@@ -234,12 +234,12 @@   GC.contextField (functionTiming name) [C.cty|typename int64_t|] $ Just [C.cexp|0|]   GC.contextField (functionIterations name) [C.cty|typename int64_t|] $ Just [C.cexp|0|] -multicoreName :: String -> GC.CompilerM op s Name+multicoreName :: Name -> GC.CompilerM op s Name multicoreName s = do-  s' <- newVName ("futhark_mc_" ++ s)-  pure $ nameFromString $ baseString s' ++ "_" ++ show (baseTag s')+  s' <- newVName $ "futhark_mc_" <> s+  pure $ baseName s' <> "_" <> nameFromString (show (baseTag s')) -type DefSpecifier s = String -> (Name -> GC.CompilerM Multicore s C.Definition) -> GC.CompilerM Multicore s Name+type DefSpecifier s = Name -> (Name -> GC.CompilerM Multicore s C.Definition) -> GC.CompilerM Multicore s Name  multicoreDef :: DefSpecifier s multicoreDef s f = do@@ -250,7 +250,7 @@ generateParLoopFn ::   (C.ToIdent a) =>   M.Map VName Space ->-  String ->+  Name ->   MCCode ->   a ->   [(VName, (C.Type, ValueType))] ->@@ -285,7 +285,7 @@  prepareTaskStruct ::   DefSpecifier s ->-  String ->+  Name ->   [VName] ->   [(C.Type, ValueType)] ->   [VName] ->@@ -332,7 +332,7 @@   fstruct <-     prepareTaskStruct multicoreDef "task" free_args free_ctypes retval_args retval_ctypes -  fpar_task <- generateParLoopFn lexical (name ++ "_task") seq_code fstruct free retval+  fpar_task <- generateParLoopFn lexical (name <> "_task") seq_code fstruct free retval   addTimingFields fpar_task    let ftask_name = fstruct <> "_task"@@ -353,7 +353,7 @@   case par_task of     Just (ParallelTask nested_code) -> do       let lexical_nested = lexicalMemoryUsageMC TraverseKernels $ Function Nothing [] params nested_code-      fnpar_task <- generateParLoopFn lexical_nested (name ++ "_nested_task") nested_code fstruct free retval+      fnpar_task <- generateParLoopFn lexical_nested (name <> "_nested_task") nested_code fstruct free retval       GC.stm [C.cstm|$id:ftask_name.nested_fn = $id:fnpar_task;|]     Nothing ->       GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|]@@ -374,9 +374,9 @@   let lexical = lexicalMemoryUsageMC TraverseKernels $ Function Nothing [] free body    fstruct <--    prepareTaskStruct multicoreDef (s' ++ "_parloop_struct") free_args free_ctypes mempty mempty+    prepareTaskStruct multicoreDef (s' <> "_parloop_struct") free_args free_ctypes mempty mempty -  ftask <- multicoreDef (s' ++ "_parloop") $ \s -> do+  ftask <- multicoreDef (s' <> "_parloop") $ \s -> do     fbody <- benchmarkCode s <=< GC.inNewFunction $       GC.cachingMemory lexical $ \decl_cached free_cached -> GC.collect $ do         GC.items [C.citems|$decls:(compileGetStructVals fstruct free_args free_ctypes)|]

@@ -271,12 +271,12 @@       pure ([C.cparam|$tyquals:vari $ty:ty * $tyquals:vari $id:name|], [C.cexp|*$id:name|])      compileOutputsExtern vari (ScalarParam name bt) = do-      p_name <- newVName $ "out_" ++ baseString name+      p_name <- newVName $ "out_" <> baseName name       let ctp = GC.primTypeToCType bt       pure ([C.cparam|$tyquals:vari $ty:ctp * $tyquals:vari $id:p_name|], [C.cexp|$id:p_name|])     compileOutputsExtern vari (MemParam name space) = do       ty <- GC.memToCType name space-      p_name <- newVName $ baseString name ++ "_p"+      p_name <- newVName $ baseName name <> "_p"       pure ([C.cparam|$tyquals:vari $ty:ty * $tyquals:vari $id:p_name|], [C.cexp|$id:p_name|])      compileInputsUniform (ScalarParam name bt) = do@@ -291,14 +291,14 @@       pure (params, args)      compileOutputsUniform (ScalarParam name bt) = do-      p_name <- newVName $ "out_" ++ baseString name+      p_name <- newVName $ "out_" <> baseName name       let ctp = GC.primTypeToCType bt           params = [C.cparam|$tyqual:uniform $ty:ctp *$tyqual:uniform $id:p_name|]           args = [C.cexp|$id:p_name|]       pure (params, args)     compileOutputsUniform (MemParam name space) = do       ty <- GC.memToCType name space-      p_name <- newVName $ baseString name ++ "_p"+      p_name <- newVName $ baseName name <> "_p"       let params = [C.cparam|$tyqual:uniform $ty:ty $id:p_name|]           args = [C.cexp|&$id:p_name|]       pure (params, args)@@ -311,7 +311,7 @@       pure (params, args, pre_body)     compileInputsVarying (MemParam name space) = do       typ <- GC.memToCType name space-      newvn <- newVName $ "aos_" <> baseString name+      newvn <- newVName $ "aos_" <> baseName name       let params = [C.cparam|$ty:typ $id:name|]           args = [C.cexp|&$id:(newvn)[i]|]           pre_body =@@ -320,9 +320,9 @@       pure (params, args, pre_body)      compileOutputsVarying (ScalarParam name bt) = do-      p_name <- newVName $ "out_" ++ baseString name-      deref_name <- newVName $ "aos_" ++ baseString name-      vari_p_name <- newVName $ "convert_" ++ baseString name+      p_name <- newVName $ "out_" <> baseName name+      deref_name <- newVName $ "aos_" <> baseName name+      vari_p_name <- newVName $ "convert_" <> baseName name       let ctp = GC.primTypeToCType bt           pre_body =             [C.citems|$tyqual:varying $ty:ctp $id:vari_p_name = *$id:p_name;@@ -334,7 +334,7 @@       pure (params, args, pre_body, post_body)     compileOutputsVarying (MemParam name space) = do       typ <- GC.memToCType name space-      newvn <- newVName $ "aos_" <> baseString name+      newvn <- newVName $ "aos_" <> baseName name       let params = [C.cparam|$ty:typ $id:name|]           args = [C.cexp|&$id:(newvn)[i]|]           pre_body =@@ -481,7 +481,7 @@   quals <- getVariabilityQuals name   GC.decl [C.cdecl|$tyquals:quals $ty:ct $id:name;|] compileCode (DeclareArray name t vs) = do-  name_realtype <- newVName $ baseString name ++ "_realtype"+  name_realtype <- newVName $ baseName name <> "_realtype"   let ct = GC.primTypeToCType t   case vs of     ArrayValues vs' -> do@@ -752,7 +752,7 @@   fstruct <-     MC.prepareTaskStruct sharedDef "task" free_args free_ctypes retval_args retval_ctypes -  fpar_task <- MC.generateParLoopFn lexical (name ++ "_task") seq_code fstruct free retval+  fpar_task <- MC.generateParLoopFn lexical (name <> "_task") seq_code fstruct free retval   MC.addTimingFields fpar_task    let ftask_name = fstruct <> "_task"@@ -775,7 +775,14 @@     case par_task of       Just (ParallelTask nested_code) -> do         let lexical_nested = lexicalMemoryUsageMC OpaqueKernels $ Function Nothing [] params nested_code-        fnpar_task <- MC.generateParLoopFn lexical_nested (name ++ "_nested_task") nested_code fstruct free retval+        fnpar_task <-+          MC.generateParLoopFn+            lexical_nested+            (name <> "_nested_task")+            nested_code+            fstruct+            free+            retval         GC.stm [C.cstm|$id:ftask_name.nested_fn = $id:fnpar_task;|]       Nothing ->         GC.stm [C.cstm|$id:ftask_name.nested_fn=NULL;|]@@ -944,7 +951,7 @@   let cached = M.keys $ M.filter (== DefaultSpace) lexical    cached' <- forM cached $ \mem -> do-    size <- newVName $ prettyString mem <> "_cached_size"+    size <- newVName $ nameFromText (prettyText mem) <> "_cached_size"     pure (mem, size)    let lexMem env =

@@ -146,14 +146,17 @@   either (Left . errorBundlePretty) Right . parse (p <* eof) "type name"   where     p :: Parsec Void T.Text T.Text-    p = choice [pArr, pTup, pQual]+    p = choice [pArr, pTup, pRec, pQual]     pArr = do       dims <- some "[]"       (("arr" <> showText (length dims) <> "d_") <>) <$> p     pTup = between "(" ")" $ do       ts <- p `sepBy` pComma       pure $ "tup" <> showText (length ts) <> "_" <> T.intercalate "_" ts-    pAtom = T.pack <$> some (satisfy (`notElem` ("[]{}(),." :: String)))+    pRec = between "{" "}" $ do+      fs <- pField `sepBy` pComma+      pure $ "rec__" <> T.intercalate "__" fs+    pAtom = T.pack <$> some (satisfy (`notElem` ("[]{}(),.:" :: String)))     pComma = void $ "," <* space     -- Rewrite 'x.y' to 'x_y'.     pQual = do@@ -162,6 +165,10 @@         [ "." >> ((x <> "_") <>) <$> pAtom,           pure x         ]+    pField = do+      f <- pAtom <* ":"+      t <- p+      pure $ f <> "_" <> t  -- | The name of exposed opaque types. opaqueName :: Name -> T.Text

@@ -22,8 +22,8 @@  -- | A multicore operation. data Multicore-  = SegOp String [Param] ParallelTask (Maybe ParallelTask) [Param] SchedulerInfo-  | ParLoop String MCCode [Param]+  = SegOp Name [Param] ParallelTask (Maybe ParallelTask) [Param] SchedulerInfo+  | ParLoop Name MCCode [Param]   | -- | A kernel of ISPC code, or a scoped block in regular C.     ISPCKernel MCCode [Param]   | -- | A foreach loop in ISPC, or a regular for loop in C.

@@ -717,7 +717,7 @@   -- variables mirroring the merge parameters, and then copy this   -- buffer to the merge parameters.  This is efficient, because the   -- operations are all scalar operations.-  tmpnames <- mapM (newVName . (++ "_tmp") . baseString . paramName) mergeparams+  tmpnames <- mapM (newVName . (<> "_tmp") . baseName . paramName) mergeparams   compileStms (freeIn ses) stms $ do     copy_to_merge_params <- forM (zip3 mergeparams tmpnames ses) $ \(p, tmp, SubExpRes _ se) ->       case typeOf p of@@ -1076,7 +1076,7 @@ dLParams :: (Mem rep inner) => [LParam rep] -> ImpM rep r op () dLParams = dScope Nothing . scopeOfLParams -dPrimVol :: String -> PrimType -> Imp.TExp t -> ImpM rep r op (TV t)+dPrimVol :: Name -> PrimType -> Imp.TExp t -> ImpM rep r op (TV t) dPrimVol name t e = do   name' <- newVName name   emit $ Imp.DeclareScalar name' Imp.Volatile t@@ -1092,7 +1092,7 @@ -- | Create variable of some provided dynamic type. You'll need this -- when you are compiling program code of Haskell-level unknown type. -- For other things, use other functions.-dPrimS :: String -> PrimType -> ImpM rep r op VName+dPrimS :: Name -> PrimType -> ImpM rep r op VName dPrimS name t = do   name' <- newVName name   dPrim_ name' t@@ -1100,11 +1100,11 @@  -- | Create 'TV' of some provided dynamic type. No guarantee that the -- dynamic type matches the inferred type.-dPrimSV :: String -> PrimType -> ImpM rep r op (TV t)+dPrimSV :: Name -> PrimType -> ImpM rep r op (TV t) dPrimSV name t = TV <$> dPrimS name t <*> pure t  -- | Create 'TV' of some fixed type.-dPrim :: (MkTV t) => String -> ImpM rep r op (TV t)+dPrim :: (MkTV t) => Name -> ImpM rep r op (TV t) dPrim name = do   name' <- newVName name   let tv = mkTV name'@@ -1118,7 +1118,7 @@   where     t = primExpType $ untyped e -dPrimV :: String -> Imp.TExp t -> ImpM rep r op (TV t)+dPrimV :: Name -> Imp.TExp t -> ImpM rep r op (TV t) dPrimV name e = do   name' <- dPrimS name pt   let tv = TV name' pt@@ -1127,7 +1127,7 @@   where     pt = primExpType $ untyped e -dPrimVE :: String -> Imp.TExp t -> ImpM rep r op (Imp.TExp t)+dPrimVE :: Name -> Imp.TExp t -> ImpM rep r op (Imp.TExp t) dPrimVE name e = do   name' <- dPrimS name pt   let tv = TV name' pt@@ -1309,16 +1309,16 @@ askFunction = asks envFunction  -- | Generate a 'VName', prefixed with 'askFunction' if it exists.-newVNameForFun :: String -> ImpM rep r op VName+newVNameForFun :: Name -> ImpM rep r op VName newVNameForFun s = do-  fname <- fmap nameToString <$> askFunction-  newVName $ maybe "" (++ ".") fname ++ s+  fname <- askFunction+  newVName $ maybe "" (<> ".") fname <> s  -- | Generate a 'Name', prefixed with 'askFunction' if it exists.-nameForFun :: String -> ImpM rep r op Name+nameForFun :: Name -> ImpM rep r op Name nameForFun s = do   fname <- askFunction-  pure $ maybe "" (<> ".") fname <> nameFromString s+  pure $ maybe "" (<> ".") fname <> s  askEnv :: ImpM rep r op r askEnv = asks envEnv@@ -1749,7 +1749,7 @@   body' <- collect body   emit $ Imp.For i bound body' -sFor :: String -> Imp.TExp t -> (Imp.TExp t -> ImpM rep r op ()) -> ImpM rep r op ()+sFor :: Name -> Imp.TExp t -> (Imp.TExp t -> ImpM rep r op ()) -> ImpM rep r op () sFor i bound body = do   i' <- newVName i   sFor' i' (untyped bound) $@@ -1793,7 +1793,7 @@ sOp :: op -> ImpM rep r op () sOp = emit . Imp.Op -sDeclareMem :: String -> Space -> ImpM rep r op VName+sDeclareMem :: Name -> Space -> ImpM rep r op VName sDeclareMem name space = do   name' <- newVName name   emit $ Imp.DeclareMem name' space@@ -1807,20 +1807,20 @@     Nothing -> emit $ Imp.Allocate name' size' space     Just allocator' -> allocator' name' size' -sAlloc :: String -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM rep r op VName+sAlloc :: Name -> Count Bytes (Imp.TExp Int64) -> Space -> ImpM rep r op VName sAlloc name size space = do   name' <- sDeclareMem name space   sAlloc_ name' size space   pure name' -sArray :: String -> PrimType -> ShapeBase SubExp -> VName -> LMAD -> ImpM rep r op VName+sArray :: Name -> PrimType -> ShapeBase SubExp -> VName -> LMAD -> ImpM rep r op VName sArray name bt shape mem lmad = do   name' <- newVName name   dArray name' bt shape mem lmad   pure name'  -- | Declare an array in row-major order in the given memory block.-sArrayInMem :: String -> PrimType -> ShapeBase SubExp -> VName -> ImpM rep r op VName+sArrayInMem :: Name -> PrimType -> ShapeBase SubExp -> VName -> ImpM rep r op VName sArrayInMem name pt shape mem =   sArray name pt shape mem $     LMAD.iota 0 $@@ -1828,28 +1828,28 @@         shapeDims shape  -- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified.-sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM rep r op VName+sAllocArrayPerm :: Name -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM rep r op VName sAllocArrayPerm name pt shape space perm = do   let permuted_dims = rearrangeShape perm $ shapeDims shape-  mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space+  mem <- sAlloc (name <> "_mem") (typeSize (Array pt shape NoUniqueness)) space   let iota_lmad = LMAD.iota 0 $ map (isInt64 . primExpFromSubExp int64) permuted_dims   sArray name pt shape mem $     LMAD.permute iota_lmad $       rearrangeInverse perm  -- | Uses linear/iota index function.-sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName+sAllocArray :: Name -> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName sAllocArray name pt shape space =   sAllocArrayPerm name pt shape space [0 .. shapeRank shape - 1]  -- | Uses linear/iota index function.-sStaticArray :: String -> PrimType -> Imp.ArrayContents -> ImpM rep r op VName+sStaticArray :: Name -> PrimType -> Imp.ArrayContents -> ImpM rep r op VName sStaticArray name pt vs = do   let num_elems = case vs of         Imp.ArrayValues vs' -> length vs'         Imp.ArrayZeros n -> fromIntegral n       shape = Shape [intConst Int64 $ toInteger num_elems]-  mem <- newVNameForFun $ name ++ "_mem"+  mem <- newVNameForFun $ name <> "_mem"   emit $ Imp.DeclareArray mem pt vs   addVar mem $ MemVar Nothing $ MemEntry DefaultSpace   sArray name pt shape mem $ LMAD.iota 0 [fromIntegral num_elems]@@ -1977,7 +1977,7 @@ -- | Like 'dIndexSpace', but invent some new names for the indexes -- based on the given template. dIndexSpace' ::-  String ->+  Name ->   [Imp.TExp Int64] ->   Imp.TExp Int64 ->   ImpM rep r op [Imp.TExp Int64]

@@ -117,7 +117,7 @@  keyWithEntryPoint :: Maybe Name -> Name -> Name keyWithEntryPoint fname key =-  nameFromString $ maybe "" ((++ ".") . nameToString) fname ++ nameToString key+  maybe "" (<> ".") fname <> key  allocLocal :: AllocCompiler GPUMem r Imp.KernelOp allocLocal mem size =@@ -180,7 +180,7 @@  -- | Generate a constant device array of 32-bit integer zeroes with -- the given number of elements.  Initialised with a replicate.-genZeroes :: String -> Int -> CallKernelGen VName+genZeroes :: Name -> Int -> CallKernelGen VName genZeroes desc n = genConstants $ do   counters_mem <- sAlloc (desc <> "_mem") (4 * fromIntegral n) (Space "device")   let shape = Shape [intConst Int64 (fromIntegral n)]@@ -285,11 +285,11 @@ kernelConstToExp = traverse f   where     f (Imp.SizeMaxConst c) = do-      v <- dPrimS (prettyString c) int64+      v <- dPrimS (nameFromText $ prettyText c) int64       sOp $ Imp.GetSizeMax v c       pure v     f (Imp.SizeConst k c) = do-      v <- dPrimS (nameToString k) int64+      v <- dPrimS (nameFromText $ prettyText k) int64       sOp $ Imp.GetSize v k c       pure v @@ -1058,7 +1058,7 @@  simpleKernelConstants ::   Imp.TExp Int64 ->-  String ->+  Name ->   CallKernelGen     ( (Imp.TExp Int64 -> InKernelGen ()) -> InKernelGen (),       KernelConstants@@ -1070,9 +1070,9 @@   -- GPU will possibly need.  Feel free to come back and laugh at me   -- in the future.   let max_num_tblocks = 1024 * 1024-  thread_gtid <- newVName $ desc ++ "_gtid"-  thread_ltid <- newVName $ desc ++ "_ltid"-  tblock_id <- newVName $ desc ++ "_gid"+  thread_gtid <- newVName $ desc <> "_gtid"+  thread_ltid <- newVName $ desc <> "_ltid"+  tblock_id <- newVName $ desc <> "_gid"   inner_tblock_size <- newVName "tblock_size"   (virt_num_tblocks, num_tblocks, tblock_size) <-     simpleKernelBlocks max_num_tblocks kernel_size@@ -1175,11 +1175,11 @@  -- | Retrieve a size of the given size class and put it in a variable -- with the given name.-getSize :: String -> SizeClass -> CallKernelGen (TV Int64)+getSize :: Name -> SizeClass -> CallKernelGen (TV Int64) getSize desc size_class = do   v <- dPrim desc   fname <- askFunction-  let v_key = keyWithEntryPoint fname $ nameFromString $ prettyString $ tvVar v+  let v_key = keyWithEntryPoint fname $ nameFromText $ prettyText $ tvVar v   sOp $ Imp.GetSize (tvVar v) v_key size_class   pure v @@ -1205,7 +1205,7 @@ sKernel ::   Operations GPUMem KernelEnv Imp.KernelOp ->   (KernelConstants -> Imp.TExp Int64) ->-  String ->+  Name ->   VName ->   KernelAttrs ->   InKernelGen () ->@@ -1213,14 +1213,14 @@ sKernel ops flatf name v attrs f = do   (constants, set_constants) <-     kernelInitialisationSimple (kAttrNumBlocks attrs) (kAttrBlockSize attrs)-  name' <- nameForFun $ name ++ "_" ++ show (baseTag v)+  name' <- nameForFun $ name <> "_" <> nameFromString (show (baseTag v))   sKernelOp attrs constants ops name' $ do     set_constants     dPrimV_ v $ flatf constants     f  sKernelThread ::-  String ->+  Name ->   VName ->   KernelAttrs ->   InKernelGen () ->@@ -1308,8 +1308,8 @@   fname <- askFunction   let name =         keyWithEntryPoint fname $-          nameFromString $-            "replicate_" ++ show (baseTag $ tvVar $ kernelGlobalThreadIdVar constants)+          "replicate_"+            <> nameFromString (show (baseTag $ tvVar $ kernelGlobalThreadIdVar constants))    sKernelFailureTolerant True threadOperations constants name $     virtualise $ \gtid -> do@@ -1470,12 +1470,5 @@ compileThreadResult space pe (Returns _ _ what) = do   let is = map (Imp.le64 . fst) $ unSegSpace space   copyDWIMFix (patElemName pe) is what []-compileThreadResult _ pe (WriteReturns _ arr dests) = do-  arr_t <- lookupType arr-  let rws' = map pe64 $ arrayDims arr_t-  forM_ dests $ \(slice, e) -> do-    let slice' = fmap pe64 slice-        write = inBounds slice' rws'-    sWhen write $ copyDWIM (patElemName pe) (unSlice slice') e [] compileThreadResult _ _ TileReturns {} =   compilerBugS "compileThreadResult: TileReturns unhandled."

@@ -39,7 +39,7 @@ flattenArray k flat arr = do   ArrayEntry arr_loc pt <- lookupArray arr   let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc)-  sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $+  sArray (baseName arr <> "_flat") pt flat_shape (memLocName arr_loc) $     fromMaybe (error "flattenArray") $       LMAD.reshape (memLocLMAD arr_loc) (map pe64 $ shapeDims flat_shape) @@ -52,7 +52,7 @@           (map Imp.pe64 (arrayDims arr_t))           [DimSlice start (tvExp size) 1]   sArray-    (baseString arr ++ "_chunk")+    (baseName arr <> "_chunk")     (elemType arr_t)     (arrayShape arr_t `setOuterDim` Var (tvVar size))     mem@@ -361,9 +361,9 @@   compileFlatId space    blockCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms body) $+    compileStms mempty (bodyStms body) $       zipWithM_ (compileThreadResult space) (patElems pat) $-        kernelBodyResult body+        bodyResult body   sOp $ Imp.ErrorSync Imp.FenceLocal compileBlockOp pat (Inner (SegOp (SegScan lvl space _ body scans))) = do   compileFlatId space@@ -372,8 +372,8 @@       dims' = map pe64 dims    blockCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms body) $-      forM_ (zip (patNames pat) $ kernelBodyResult body) $ \(dest, res) ->+    compileStms mempty (bodyStms body) $+      forM_ (zip (patNames pat) $ bodyResult body) $ \(dest, res) ->         copyDWIMFix           dest           (map Imp.le64 ltids)@@ -421,9 +421,9 @@    tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops   blockCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms body) $ do+    compileStms mempty (bodyStms body) $ do       let (red_res, map_res) =-            splitAt (segBinOpResults ops) $ kernelBodyResult body+            splitAt (segBinOpResults ops) $ bodyResult body       forM_ (zip tmp_arrs red_res) $ \(dest, res) ->         copyDWIMFix dest (map Imp.le64 ltids) (kernelResultSubExp res) []       zipWithM_ (compileThreadResult space) map_pes map_res@@ -551,8 +551,8 @@   sOp $ Imp.Barrier Imp.FenceLocal    blockCoverSegSpace (segVirt lvl) space $-    compileStms mempty (kernelBodyStms kbody) $ do-      let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody+    compileStms mempty (bodyStms kbody) $ do+      let (red_res, map_res) = splitAt num_red_res $ bodyResult kbody           (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res       zipWithM_ (compileThreadResult space) map_pes map_res @@ -601,7 +601,7 @@  -- | Create a kernel with GPU operations at the block level. sKernelBlock ::-  String ->+  Name ->   VName ->   KernelAttrs ->   InKernelGen () ->@@ -698,8 +698,6 @@     -- block.  TODO: also do this if the array is in global memory     -- (but this is a bit more tricky, synchronisation-wise).       copyDWIMFix (patElemName pe) gids what []-compileBlockResult _ _ WriteReturns {} =-  compilerLimitationS "compileBlockResult: WriteReturns not handled yet."  -- | The sizes of nested iteration spaces in the kernel. type SegOpSizes = S.Set [SubExp]

@@ -41,11 +41,13 @@ import Data.List qualified as L import Data.Map qualified as M import Data.Maybe+import Futhark.Analysis.Alias qualified as Alias import Futhark.CodeGen.ImpCode.GPU qualified as Imp import Futhark.CodeGen.ImpGen import Futhark.CodeGen.ImpGen.GPU.Base import Futhark.CodeGen.ImpGen.GPU.SegRed (compileSegRed') import Futhark.Construct (fullSliceNum)+import Futhark.IR.Aliases (consumedInBody) import Futhark.IR.GPUMem import Futhark.IR.Mem.LMAD qualified as LMAD import Futhark.Pass.ExplicitAllocations ()@@ -103,14 +105,14 @@     dest_mem <- entryArrayLoc <$> lookupArray dest      subhistos_mem <--      sDeclareMem (baseString dest ++ "_subhistos_mem") (Space "device")+      sDeclareMem (baseName dest <> "_subhistos_mem") (Space "device")      let subhistos_shape =           Shape (map snd segment_dims ++ [tvSize num_subhistos])             <> stripDims num_segments (arrayShape dest_t)     subhistos <-       sArray-        (baseString dest ++ "_subhistos")+        (baseName dest <> "_subhistos")         (elemType dest_t)         subhistos_shape         subhistos_mem@@ -197,7 +199,7 @@  bodyPassage :: KernelBody GPUMem -> Passage bodyPassage kbody-  | mempty == consumedInKernelBody (aliasAnalyseKernelBody mempty kbody) =+  | mempty == consumedInBody (Alias.analyseBody mempty kbody) =       MayBeMultiPass   | otherwise =       MustBeSinglePass@@ -418,8 +420,8 @@       let input_in_bounds = offset .<. total_w_64        sWhen input_in_bounds $-        compileStms mempty (kernelBodyStms kbody) $ do-          let (red_res, map_res) = splitFromEnd (length map_pes) $ kernelBodyResult kbody+        compileStms mempty (bodyStms kbody) $ do+          let (red_res, map_res) = splitFromEnd (length map_pes) $ bodyResult kbody            sComment "save map-out results" $             forM_ (zip map_pes map_res) $ \(pe, res) ->@@ -720,11 +722,11 @@           -- serially.  This also involves writing to the mapout arrays if           -- this is the first chunk. -          compileStms mempty (kernelBodyStms kbody) $ do+          compileStms mempty (bodyStms kbody) $ do             let (red_res, map_res) =                   splitFromEnd (length map_pes) $                     map kernelResultSubExp $-                      kernelBodyResult kbody+                      bodyResult kbody              sWhen (chk_i .==. 0) $               sComment "save map-out results" $

@@ -45,20 +45,20 @@           dIndexSpace (zip is dims') global_tid            sWhen (isActive $ unSegSpace space) $-            compileStms mempty (kernelBodyStms kbody) $+            compileStms mempty (bodyStms kbody) $               zipWithM_ (compileThreadResult space) (patElems pat) $-                kernelBodyResult kbody+                bodyResult kbody     SegBlock {} -> do-      pc <- precomputeConstants tblock_size' $ kernelBodyStms kbody+      pc <- precomputeConstants tblock_size' $ bodyStms kbody       virt_num_tblocks <- dPrimVE "virt_num_tblocks" $ sExt32 $ product dims'       sKernelBlock "segmap_intrablock" (segFlat space) attrs $ do         precomputedConstants pc $           virtualiseBlocks (segVirt lvl) virt_num_tblocks $ \tblock_id -> do             dIndexSpace (zip is dims') $ sExt64 tblock_id -            compileStms mempty (kernelBodyStms kbody) $+            compileStms mempty (bodyStms kbody) $               zipWithM_ (compileBlockResult space) (patElems pat) $-                kernelBodyResult kbody+                bodyResult kbody     SegThreadInBlock {} ->       error "compileSegMap: SegThreadInBlock"   emit $ Imp.DebugPrint "" Nothing

@@ -113,8 +113,8 @@    compileSegRed' pat grid space segbinops $ \red_cont ->     sComment "apply map function" $-      compileStms mempty (kernelBodyStms map_kbody) $ do-        let (red_res, map_res) = splitAt (segBinOpResults segbinops) $ kernelBodyResult map_kbody+      compileStms mempty (bodyStms map_kbody) $ do+        let (red_res, map_res) = splitAt (segBinOpResults segbinops) $ bodyResult map_kbody          let mapout_arrs = drop (segBinOpResults segbinops) $ patElems pat         unless (null mapout_arrs) $@@ -222,7 +222,7 @@       lmem <- sAlloc "local_mem" lmem_total_size (Space "shared")       let arrInLMem ptype name len_se offset =             sArray-              (name ++ "_" ++ prettyString ptype)+              (name <> "_" <> nameFromString (prettyString ptype))               ptype               (Shape [len_se])               lmem@@ -236,7 +236,7 @@               <$> arrInLMem ptype "coll_copy_arr" block_worksize 0               <*> arrInLMem ptype "block_red_arr" tblock_size offset               <*> sAllocArray-                ("chunk_" ++ prettyString ptype)+                ("chunk_" <> nameFromText (prettyText ptype))                 ptype                 (Shape [chunk])                 (ScalarSpace [chunk] ptype)@@ -262,7 +262,7 @@       MemArray pt shape _ (ArrayIn mem ixfun) -> do         let shape' = Shape [tblock_size] <> shape         let shape_E = map pe64 $ shapeDims shape'-        sArray ("red_arr_" ++ prettyString pt) pt shape' mem $+        sArray ("red_arr_" <> nameFromText (prettyText pt)) pt shape' mem $           -- This 'segmented' thing here is a hack, related to #2227.           -- There absolutely must be some unifying principle we are           -- missing.@@ -272,7 +272,7 @@       _ -> do         let pt = elemType $ paramType p             shape = Shape [tblock_size]-        sAllocArray ("red_arr_" ++ prettyString pt) pt shape $ Space "shared"+        sAllocArray ("red_arr_" <> nameFromText (prettyText pt)) pt shape $ Space "shared"  -- | Arrays for storing block results. --@@ -647,7 +647,7 @@     mkAcc p block_res_arr       | Prim t <- paramType p,         shapeRank (segBinOpShape op) == 0 = do-          block_res_acc <- dPrimS (baseString (paramName p) <> "_block_res_acc") t+          block_res_acc <- dPrimS (baseName (paramName p) <> "_block_res_acc") t           pure (block_res_acc, [])       -- if this is a non-primitive reduction, the global mem result array will       -- double as accumulator.

@@ -352,9 +352,9 @@           dIndexSpace (zip gtids dims') virt_tid           -- Perform the map           let in_bounds =-                compileStms mempty (kernelBodyStms map_kbody) $ do+                compileStms mempty (bodyStms map_kbody) $ do                   let (all_scan_res, map_res) =-                        splitAt (segBinOpResults [scan_op]) $ kernelBodyResult map_kbody+                        splitAt (segBinOpResults [scan_op]) $ bodyResult map_kbody                    -- Write map results to their global memory destinations                   forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(dest, src) ->

@@ -196,9 +196,9 @@           in_bounds =             foldl1 (.&&.) $ zipWith (.<.) (map Imp.le64 gtids) dims' -          when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do+          when_in_bounds = compileStms mempty (bodyStms kbody) $ do             let (all_scan_res, map_res) =-                  splitAt (segBinOpResults scans) $ kernelBodyResult kbody+                  splitAt (segBinOpResults scans) $ bodyResult kbody                 per_scan_res =                   segBinOpChunks scans all_scan_res 

@@ -637,7 +637,7 @@     kernelOps (MemFence FenceGlobal) =       GC.stm [C.cstm|mem_fence_global();|]     kernelOps (SharedAlloc name size) = do-      name' <- newVName $ prettyString name ++ "_backing"+      name' <- newVName $ nameFromText $ prettyText name <> "_backing"       GC.modifyUserState $ \s ->         s {kernelSharedMemory = (name', size) : kernelSharedMemory s}       GC.stm [C.cstm|$id:name = (__local unsigned char*) $id:name';|]

@@ -32,7 +32,7 @@     free_params <- freeParams body     emit $ Imp.Op $ Imp.ParLoop "copy" body free_params   free_params <- freeParams seq_code-  s <- prettyString <$> newVName "copy"+  s <- nameFromText . prettyText <$> newVName "copy"   iterations <- dPrimVE "iterations" $ product $ map pe64 srcshape   let scheduling = Imp.SchedulerInfo (untyped iterations) Imp.Static   emit . Imp.Op $@@ -152,7 +152,7 @@       pure $ Just $ Imp.ParallelTask par_code     Nothing -> pure Nothing -  s <- segOpString op+  s <- segOpName op   let seq_task = Imp.ParallelTask seq_code   free_params <- filter (`notElem` retvals) <$> freeParams (par_task, seq_task)   let code =

@@ -18,7 +18,7 @@     getLoopBounds,     getIterationDomain,     getReturnParams,-    segOpString,+    segOpName,     ChunkLoopVectorization (..),     generateChunkLoop,     generateUniformizeLoop,@@ -58,11 +58,11 @@  type MulticoreGen = ImpM MCMem HostEnv Imp.Multicore -segOpString :: SegOp () MCMem -> MulticoreGen String-segOpString SegMap {} = pure "segmap"-segOpString SegRed {} = pure "segred"-segOpString SegScan {} = pure "segscan"-segOpString SegHist {} = pure "seghist"+segOpName :: SegOp () MCMem -> MulticoreGen Name+segOpName SegMap {} = pure "segmap"+segOpName SegRed {} = pure "segred"+segOpName SegScan {} = pure "segscan"+segOpName SegHist {} = pure "seghist"  arrParam :: VName -> MulticoreGen Imp.Param arrParam arr = do@@ -134,8 +134,6 @@ compileThreadResult space pe (Returns _ _ what) = do   let is = map (Imp.le64 . fst) $ unSegSpace space   copyDWIMFix (patElemName pe) is what []-compileThreadResult _ _ WriteReturns {} =-  compilerBugS "compileThreadResult: WriteReturns unhandled." compileThreadResult _ _ TileReturns {} =   compilerBugS "compileThreadResult: TileReturns unhandled." compileThreadResult _ _ RegTileReturns {} =@@ -219,7 +217,7 @@ -- The action is called with the (symbolic) index of the current -- iteration. generateChunkLoop ::-  String ->+  Name ->   ChunkLoopVectorization ->   (Imp.TExp Int64 -> MulticoreGen ()) ->   MulticoreGen ()@@ -302,7 +300,7 @@   body' <- collect body   emit $ Imp.Op $ Imp.ForEach i (Imp.ValueExp $ blankPrimValue $ Imp.IntType Imp.Int64) bound body' -sForVectorized :: String -> Imp.TExp t -> (Imp.TExp t -> MulticoreGen ()) -> MulticoreGen ()+sForVectorized :: Name -> Imp.TExp t -> (Imp.TExp t -> MulticoreGen ()) -> MulticoreGen () sForVectorized i bound body = do   i' <- newVName i   sForVectorized' i' (untyped bound) $

@@ -120,9 +120,9 @@     sOp $ Imp.GetTaskId (segFlat space)     generateChunkLoop "SegHist" Scalar $ \flat_idx -> do       zipWithM_ dPrimV_ is $ unflattenIndex ns_64 flat_idx-      compileStms mempty (kernelBodyStms kbody) $ do+      compileStms mempty (bodyStms kbody) $ do         let (red_res, map_res) =-              splitFromEnd (length map_pes) $ kernelBodyResult kbody+              splitFromEnd (length map_pes) $ bodyResult kbody             red_res_split = splitHistResults histops $ map kernelResultSubExp red_res          let pes_per_op = chunks (map (length . histDest) histops) all_red_pes@@ -230,11 +230,11 @@     inISPC $       generateChunkLoop "SegRed" Vectorized $ \i -> do         zipWithM_ dPrimV_ is $ unflattenIndex ns_64 i-        compileStms mempty (kernelBodyStms kbody) $ do+        compileStms mempty (bodyStms kbody) $ do           let (red_res, map_res) =                 splitFromEnd (length map_pes) $                   map kernelResultSubExp $-                    kernelBodyResult kbody+                    bodyResult kbody            sComment "save map-out results" $             forM_ (zip map_pes map_res) $ \(pe, res) ->@@ -360,11 +360,11 @@       zipWithM_ dPrimV_ (init is) $ unflattenIndex (init ns_64) idx       dPrimV_ (last is) i -      compileStms mempty (kernelBodyStms kbody) $ do+      compileStms mempty (bodyStms kbody) $ do         let (red_res, map_res) =               splitFromEnd (length map_pes) $                 map kernelResultSubExp $-                  kernelBodyResult kbody+                  bodyResult kbody         forM_ (zip3 per_red_pes histops (splitHistResults histops red_res)) $           \(red_pes, HistOp dest_shape _ _ _ shape lam, (bucket, vs')) -> do             let (is_params, vs_params) = splitAt (length vs') $ lambdaParams lam

@@ -18,14 +18,6 @@   MulticoreGen () writeResult is pe (Returns _ _ se) =   copyDWIMFix (patElemName pe) (map Imp.le64 is) se []-writeResult _ pe (WriteReturns _ arr idx_vals) = do-  arr_t <- lookupType arr-  let (iss, vs) = unzip idx_vals-      rws' = map pe64 $ arrayDims arr_t-  forM_ (zip iss vs) $ \(slice, v) -> do-    let slice' = fmap pe64 slice-    sWhen (inBounds slice' rws') $-      copyDWIM (patElemName pe) (unSlice slice') v [] writeResult _ _ res =   error $ "writeResult: cannot handle " ++ prettyString res @@ -34,7 +26,7 @@   SegSpace ->   KernelBody MCMem ->   MulticoreGen Imp.MCCode-compileSegMapBody pat space (KernelBody _ kstms kres) = collect $ do+compileSegMapBody pat space (Body _ kstms kres) = collect $ do   let (is, ns) = unzip $ unSegSpace space       ns' = map pe64 ns   dPrim_ (segFlat space) int64

@@ -25,8 +25,8 @@   MulticoreGen Imp.MCCode compileSegRed pat space reds kbody nsubtasks =   compileSegRed' pat space reds nsubtasks $ \red_cont ->-    compileStms mempty (kernelBodyStms kbody) $ do-      let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult kbody+    compileStms mempty (bodyStms kbody) $ do+      let (red_res, map_res) = splitAt (segBinOpResults reds) $ bodyResult kbody        sComment "save map-out results" $ do         let map_arrs = drop (segBinOpResults reds) $ patElems pat@@ -82,7 +82,7 @@  -- | Arrays for storing group results shared between threads groupResultArrays ::-  String ->+  Name ->   SubExp ->   [SegBinOp MCMem] ->   MulticoreGen [[VName]]

@@ -38,7 +38,7 @@ lamBody = lambdaBody . segBinOpLambda  -- Arrays for storing worker results.-carryArrays :: String -> TV Int32 -> [SegBinOp MCMem] -> MulticoreGen [[VName]]+carryArrays :: Name -> TV Int32 -> [SegBinOp MCMem] -> MulticoreGen [[VName]] carryArrays s nsubtasks segops =   forM segops $ \(SegBinOp _ lam _ shape) ->     forM (lambdaReturnType lam) $ \t -> do@@ -184,12 +184,12 @@   ImpM MCMem HostEnv Imp.Multicore () genScanLoop typ pat space kbody scan_ops local_accs i = do   let (all_scan_res, map_res) =-        splitAt (segBinOpResults scan_ops) $ kernelBodyResult kbody+        splitAt (segBinOpResults scan_ops) $ bodyResult kbody   let (is, ns) = unzip $ unSegSpace space       ns' = map pe64 ns    zipWithM_ dPrimV_ is $ unflattenIndex ns' i-  compileStms mempty (kernelBodyStms kbody) $ do+  compileStms mempty (bodyStms kbody) $ do     let map_arrs = drop (segBinOpResults scan_ops) $ patElems pat     sComment "write mapped values results to memory" $       zipWithM_ (compileThreadResult space) map_arrs map_res@@ -460,8 +460,8 @@       sFor "i" inner_bound $ \i -> do         zipWithM_ dPrimV_ (init is) $ unflattenIndex (init ns_64) segment_i         dPrimV_ (last is) i-        compileStms mempty (kernelBodyStms kbody) $ do-          let (scan_res, map_res) = splitAt (length $ segBinOpNeutral scan_op) $ kernelBodyResult kbody+        compileStms mempty (bodyStms kbody) $ do+          let (scan_res, map_res) = splitAt (length $ segBinOpNeutral scan_op) $ bodyResult kbody           sComment "write to-scan values to parameters" $             forM_ (zip scan_y_params scan_res) $ \(p, se) ->               copyDWIMFix (paramName p) [] (kernelResultSubExp se) []

@@ -135,7 +135,7 @@ -- 'letTupExp'. letSubExp ::   (MonadBuilder m) =>-  String ->+  Name ->   Exp (Rep m) ->   m SubExp letSubExp _ (BasicOp (SubExp se)) = pure se@@ -144,7 +144,7 @@ -- | Like 'letSubExp', but returns a name rather than a t'SubExp'. letExp ::   (MonadBuilder m) =>-  String ->+  Name ->   Exp (Rep m) ->   m VName letExp _ (BasicOp (SubExp (Var v))) =@@ -162,19 +162,19 @@ -- updated array is returned. letInPlace ::   (MonadBuilder m) =>-  String ->+  Name ->   VName ->   Slice SubExp ->   Exp (Rep m) ->   m VName letInPlace desc src slice e = do-  tmp <- letSubExp (desc ++ "_tmp") e+  tmp <- letSubExp (desc <> "_tmp") e   letExp desc $ BasicOp $ Update Unsafe src slice tmp  -- | Like 'letExp', but the expression may return multiple values. letTupExp ::   (MonadBuilder m) =>-  String ->+  Name ->   Exp (Rep m) ->   m [VName] letTupExp _ (BasicOp (SubExp (Var v))) =@@ -188,7 +188,7 @@ -- | Like 'letTupExp', but returns t'SubExp's instead of 'VName's. letTupExp' ::   (MonadBuilder m) =>-  String ->+  Name ->   Exp (Rep m) ->   m [SubExp] letTupExp' _ (BasicOp (SubExp se)) = pure [se]@@ -465,8 +465,8 @@     _ -> error "asInt: wrong type"   where     s = case e of-      Var v -> baseString v-      _ -> "to_" ++ prettyString to_it+      Var v -> baseName v+      _ -> nameFromText $ "to_" <> prettyText to_it  -- | Apply a binary operator to several subexpressions.  A left-fold. foldBinOp ::@@ -605,9 +605,9 @@ -- value, then return the body constructed from the 'Result' and any -- statements added during the action, along the auxiliary value. buildBody ::-  (MonadBuilder m) =>-  m (Result, a) ->-  m (Body (Rep m), a)+  (MonadBuilder m, IsResult res) =>+  m ([res], a) ->+  m (GBody (Rep m) res, a) buildBody m = do   ((res, v), stms) <- collectStms m   body <- mkBodyM stms res@@ -615,9 +615,9 @@  -- | As 'buildBody', but there is no auxiliary value. buildBody_ ::-  (MonadBuilder m) =>-  m Result ->-  m (Body (Rep m))+  (MonadBuilder m, IsResult res) =>+  m [res] ->+  m (GBody (Rep m) res) buildBody_ m = fst <$> buildBody ((,()) <$> m)  -- | Change that result where evaluation of the body would stop.  Also@@ -711,5 +711,5 @@   toExp = pure . BasicOp . SubExp . Var  -- | A convenient composition of 'letSubExp' and 'toExp'.-toSubExp :: (MonadBuilder m, ToExp a) => String -> a -> m SubExp+toSubExp :: (MonadBuilder m, ToExp a) => Name -> a -> m SubExp toSubExp s e = letSubExp s =<< toExp e

@@ -68,6 +68,10 @@ -- name. type FileMap = M.Map VName (FilePath, String) +-- | Return the base 'Name' converted to a string.+baseString :: VName -> String+baseString = nameToString . baseName+ vnameToFileMap :: Imports -> FileMap vnameToFileMap = mconcat . map forFile   where

@@ -273,7 +273,7 @@   fmt e@Update {} = fmtUpdate e   fmt e@RecordUpdate {} = fmtUpdate e   fmt (Assert e1 e2 _ loc) =-    addComments loc $ "assert" <+> fmt e1 <+> fmt e2+    addComments loc $ "assert" <+> fmt e1 </> fmt e2   fmt (Lambda params body rettype _ loc) =     addComments loc $       "\\"@@ -363,7 +363,7 @@       sub         | null sizes = fmt pat         | otherwise = sizes' <+> fmt pat-  fmt (LetFun fname (tparams, params, retdecl, _, e) body loc) =+  fmt (LetFun (fname, _) (tparams, params, retdecl, _, e) body loc) =     addComments loc $       lineIndent         e

@@ -1,6 +1,6 @@ -- | This module provides facilities for generating unique names. module Futhark.FreshNames-  ( VNameSource,+  ( VNameSource (..),     blankNameSource,     newNameSource,     newName,

@@ -269,11 +269,15 @@ -- | Augment a body decoration with aliasing information provided by -- the statements and result of that body. mkAliasedBody ::-  (ASTRep rep, AliasedOp (OpC rep), ASTConstraints (OpC rep (Aliases rep))) =>+  ( ASTRep rep,+    AliasedOp (OpC rep),+    ASTConstraints (OpC rep (Aliases rep)),+    IsResult res+  ) =>   BodyDec rep ->   Stms (Aliases rep) ->-  Result ->-  Body (Aliases rep)+  [res] ->+  GBody (Aliases rep) res mkAliasedBody dec stms res =   Body (mkBodyAliasing stms res, dec) stms res @@ -301,9 +305,9 @@ -- in scope outside of it.  Note that this does *not* include aliases -- of results that are not bound in the statements! mkBodyAliasing ::-  (Aliased rep) =>+  (Aliased rep, IsResult res) =>   Stms rep ->-  Result ->+  [res] ->   BodyAliasing mkBodyAliasing stms res =   -- We need to remove the names that are bound in stms from the alias@@ -318,14 +322,14 @@ -- | The aliases of the result and everything consumed in the given -- statements. mkStmsAliases ::-  (Aliased rep) =>+  (Aliased rep, IsResult res) =>   Stms rep ->-  Result ->+  [res] ->   ([Names], Names) mkStmsAliases stms res = delve mempty $ stmsToList stms   where     delve (aliasmap, consumed) [] =-      ( map (aliasClosure aliasmap . subExpAliases . resSubExp) res,+      ( map (aliasClosure aliasmap . resAliases) res,         consumed       )     delve (aliasmap, consumed) (stm : stms') =@@ -344,7 +348,7 @@  -- | The variables consumed in these statements. consumedInStms :: (Aliased rep) => Stms rep -> Names-consumedInStms = snd . flip mkStmsAliases []+consumedInStms = snd . flip mkStmsAliases ([] :: [()])  -- | A helper function for computing the aliases of a sequence of -- statements.  You'd use this while recursing down the statements

@@ -140,7 +140,7 @@     Just (slice_y_bef, DimFix j, []) <- focusNth (length slice_y - 1) slice_y = Simplify $ do       let slice_x' = Slice $ slice_x_bef ++ [DimSlice i (intConst Int64 1) (intConst Int64 1)]           slice_y' = Slice $ slice_y_bef ++ [DimSlice j (intConst Int64 1) (intConst Int64 1)]-      v' <- letExp (baseString v ++ "_slice") $ BasicOp $ Index arr_y slice_y'+      v' <- letExp (baseName v <> "_slice") $ BasicOp $ Index arr_y slice_y'       certifying cs_y . auxing aux $         letBind pat $           BasicOp $

@@ -95,7 +95,7 @@     -- importantly, past the versioning If, but see also #1569).     usage (SegOp (SegMap _ _ _ kbody)) = localAllocs kbody     usage _ = mempty-    localAllocs = foldMap stmSharedAlloc . kernelBodyStms+    localAllocs = foldMap stmSharedAlloc . bodyStms     stmSharedAlloc = expSharedAlloc . stmExp     expSharedAlloc (Op (Alloc (Var v) _)) =       UT.sizeUsage v

@@ -126,7 +126,7 @@     Pat [PatElem _ (_, MemArray _ _ _ (ArrayIn mem _))] <- pat =       Simplify $ do         ~(MemArray pt shape u (ArrayIn _ v1_lmad)) <- lookupMemInfo v1-        v0' <- newVName (baseString v1 <> "_manifest")+        v0' <- newVName (baseName v1 <> "_manifest")         let manifest_pat =               Pat [PatElem v0' $ MemArray pt shape u $ ArrayIn mem v1_lmad]             stm = mkWiseStm manifest_pat mempty $ BasicOp $ Manifest v0 perm

@@ -9,15 +9,20 @@     parseSeq,     parseSeqMem, -    -- * Fragments+    -- * Representation-agnostic fragments     parseType,     parseDeclExtType,     parseDeclType,     parseVName,     parseSubExp,     parseSubExpRes,++    -- * Representation-specific fragments+    parseLambdaSOACS,+    parseBodySOACS,     parseBodyGPU,     parseBodyMC,+    parseStmSOACS,     parseStmGPU,     parseStmMC,   )@@ -33,6 +38,7 @@ import Data.Text qualified as T import Data.Void import Futhark.Analysis.PrimExp.Parse+import Futhark.FreshNames (VNameSource (..)) import Futhark.IR import Futhark.IR.GPU (GPU) import Futhark.IR.GPU.Op qualified as GPU@@ -722,6 +728,9 @@ pOpaqueTypes :: Parser OpaqueTypes pOpaqueTypes = keyword "types" $> OpaqueTypes <*> braces (many pOpaqueType) +pVNameSource :: Parser VNameSource+pVNameSource = keyword "name_source" *> (VNameSource <$> braces pInt)+ pProg :: PR rep -> Parser (Prog rep) pProg pr =   Prog@@ -731,6 +740,9 @@   where     noTypes = OpaqueTypes mempty +pStateAndProg :: PR rep -> Parser (VNameSource, Prog rep)+pStateAndProg pr = (,) <$> (pVNameSource <|> pure (VNameSource 0)) <*> pProg pr+ pSOAC :: PR rep -> Parser (SOAC.SOAC rep) pSOAC pr =   choice@@ -740,7 +752,6 @@       keyword "screma" *> pScrema pScremaForm,       keyword "vjp" *> pVJP,       keyword "jvp" *> pJVP,-      pScatter,       pHist,       pStream     ]@@ -774,20 +785,6 @@         <*> pure []     pMapForm =       SOAC.ScremaForm <$> pLambda pr <*> pure mempty <*> pure mempty-    pScatter =-      keyword "scatter"-        *> parens-          ( SOAC.Scatter-              <$> pSubExp-              <* pComma-              <*> braces (pVName `sepBy` pComma)-              <* pComma-              <*> many (pDest <* pComma)-              <*> pLambda pr-          )-      where-        pDest =-          parens $ (,,) <$> pShape <* pComma <*> pInt <* pComma <*> pVName     pHist =       keyword "hist"         *> parens@@ -909,11 +906,6 @@           ]         <*> pure cs         <*> pSubExp,-      try $-        SegOp.WriteReturns cs-          <$> pVName-          <* keyword "with"-          <*> parens (pWrite `sepBy` pComma),       try "tile"         *> parens (SegOp.TileReturns cs <$> (pTile `sepBy` pComma))         <*> pVName,@@ -929,11 +921,10 @@         blk_tile <- pSubExp <* pAsterisk         reg_tile <- pSubExp         pure (dim, blk_tile, reg_tile)-    pWrite = (,) <$> pSlice <* pEqual <*> pSubExp  pKernelBody :: PR rep -> Parser (SegOp.KernelBody rep) pKernelBody pr =-  SegOp.KernelBody (pBodyDec pr)+  Body (pBodyDec pr)     <$> pStms pr     <* keyword "return"     <*> braces (pKernelResult `sepBy` pComma)@@ -1168,28 +1159,28 @@   either (Left . T.pack . errorBundlePretty) Right $     parse (whitespace *> p <* eof) fname s -parseRep :: PR rep -> FilePath -> T.Text -> Either T.Text (Prog rep)-parseRep = parseFull . pProg+parseRep :: PR rep -> FilePath -> T.Text -> Either T.Text (VNameSource, Prog rep)+parseRep = parseFull . pStateAndProg -parseSOACS :: FilePath -> T.Text -> Either T.Text (Prog SOACS)+parseSOACS :: FilePath -> T.Text -> Either T.Text (VNameSource, Prog SOACS) parseSOACS = parseRep prSOACS -parseSeq :: FilePath -> T.Text -> Either T.Text (Prog Seq)+parseSeq :: FilePath -> T.Text -> Either T.Text (VNameSource, Prog Seq) parseSeq = parseRep prSeq -parseSeqMem :: FilePath -> T.Text -> Either T.Text (Prog SeqMem)+parseSeqMem :: FilePath -> T.Text -> Either T.Text (VNameSource, Prog SeqMem) parseSeqMem = parseRep prSeqMem -parseGPU :: FilePath -> T.Text -> Either T.Text (Prog GPU)+parseGPU :: FilePath -> T.Text -> Either T.Text (VNameSource, Prog GPU) parseGPU = parseRep prGPU -parseGPUMem :: FilePath -> T.Text -> Either T.Text (Prog GPUMem)+parseGPUMem :: FilePath -> T.Text -> Either T.Text (VNameSource, Prog GPUMem) parseGPUMem = parseRep prGPUMem -parseMC :: FilePath -> T.Text -> Either T.Text (Prog MC)+parseMC :: FilePath -> T.Text -> Either T.Text (VNameSource, Prog MC) parseMC = parseRep prMC -parseMCMem :: FilePath -> T.Text -> Either T.Text (Prog MCMem)+parseMCMem :: FilePath -> T.Text -> Either T.Text (VNameSource, Prog MCMem) parseMCMem = parseRep prMCMem  --- Fragment parsers@@ -1211,6 +1202,17 @@  parseSubExpRes :: FilePath -> T.Text -> Either T.Text SubExpRes parseSubExpRes = parseFull pSubExpRes++-- Rep-specific fragment parsers++parseLambdaSOACS :: FilePath -> T.Text -> Either T.Text (Lambda SOACS)+parseLambdaSOACS = parseFull $ pLambda prSOACS++parseBodySOACS :: FilePath -> T.Text -> Either T.Text (Body SOACS)+parseBodySOACS = parseFull $ pBody prSOACS++parseStmSOACS :: FilePath -> T.Text -> Either T.Text (Stm SOACS)+parseStmSOACS = parseFull $ pStm prSOACS  parseBodyGPU :: FilePath -> T.Text -> Either T.Text (Body GPU) parseBodyGPU = parseFull $ pBody prGPU

@@ -35,6 +35,7 @@     ASTConstraints,     IsOp (..),     ASTRep (..),+    IsResult (..),   ) where @@ -204,6 +205,17 @@   safeOp NoOp = True   cheapOp NoOp = True   opDependencies NoOp = []++-- | Something that can be returned from a 'GBody'.+class (FreeIn res) => IsResult res where+  -- | The names that may or may not contribute to the aliases of this result.+  resAliases :: res -> Names++instance IsResult SubExpRes where+  resAliases = freeIn . resSubExp++instance IsResult () where+  resAliases () = mempty  -- | Representation-specific attributes; also means the rep supports -- some basic facilities.

@@ -45,7 +45,7 @@   bodyAliases :: Body rep -> [Names]    -- | The variables consumed in the body.-  consumedInBody :: Body rep -> Names+  consumedInBody :: GBody rep res -> Names  vnameAliases :: VName -> Names vnameAliases = oneName

@@ -176,10 +176,11 @@     FreeDec (BodyDec rep),     FreeIn (RetType rep),     FreeIn (BranchType rep),-    FreeDec (ExpDec rep)+    FreeDec (ExpDec rep),+    FreeIn res   ) =>   Stms rep ->-  Result ->+  res ->   FV freeInStmsAndRes stms res =   fvBind (boundByStms stms) $ foldMap freeIn' stms <> freeIn' res@@ -264,9 +265,10 @@     FreeIn (LetDec rep),     FreeIn (RetType rep),     FreeIn (BranchType rep),-    FreeIn (Op rep)+    FreeIn (Op rep),+    FreeIn res   ) =>-  FreeIn (Body rep)+  FreeIn (GBody rep res)   where   freeIn' (Body dec stms res) =     precomputed dec $ freeIn' dec <> freeInStmsAndRes stms res

@@ -82,14 +82,14 @@ basicOpType (Index ident slice) =   result <$> lookupType ident   where-    result t = [Prim (elemType t) `arrayOfShape` shape]+    result t = [t `setArrayShape` shape]     shape = Shape $ sliceDims slice basicOpType (Update _ src _ _) =   pure <$> lookupType src basicOpType (FlatIndex ident slice) =   result <$> lookupType ident   where-    result t = [Prim (elemType t) `arrayOfShape` shape]+    result t = [t `setArrayShape` shape]     shape = Shape $ flatSliceDims slice basicOpType (FlatUpdate src _ _) =   pure <$> lookupType src

@@ -81,7 +81,7 @@   Param attrs name <$> rephraser from  -- | Rephrase a body.-rephraseBody :: (Monad m) => Rephraser m from to -> Body from -> m (Body to)+rephraseBody :: (Monad m) => Rephraser m from to -> GBody from res -> m (GBody to res) rephraseBody rephraser (Body rep stms res) =   Body     <$> rephraseBodyDec rephraser rep

@@ -61,8 +61,6 @@         || any (lamUsesAD . redLambda) reds     expUsesAD (Op (Hist _ _ ops lam)) =       lamUsesAD lam || any (lamUsesAD . histOp) ops-    expUsesAD (Op (Scatter _ _ _ lam)) =-      lamUsesAD lam     expUsesAD (Match _ cases def_case _) =       any (bodyUsesAD . caseBody) cases || bodyUsesAD def_case     expUsesAD (Loop _ _ body) = bodyUsesAD body

@@ -6,7 +6,6 @@ module Futhark.IR.SOACS.SOAC   ( SOAC (..),     ScremaForm (..),-    ScatterSpec,     HistOp (..),     Scan (..),     scanResults,@@ -35,10 +34,6 @@     ppScrema,     ppHist,     ppStream,-    ppScatter,-    groupScatterResults,-    groupScatterResults',-    splitScatterResults,      -- * Generic traversal     SOACMapper (..),@@ -53,7 +48,6 @@ import Control.Monad.Identity import Control.Monad.State.Strict import Control.Monad.Writer-import Data.Function ((&)) import Data.List (intersperse) import Data.Map.Strict qualified as M import Data.Maybe@@ -75,51 +69,9 @@ import Futhark.Util.Pretty qualified as PP import Prelude hiding (id, (.)) --- | How the results of a scatter operation should be written. Each--- element of the list consists of a @v@ (often a `VName`) specifying--- which array to scatter to, a `Shape` describing the shape of that--- array, and an `Int` describing how many elements should be written--- to that array for each invocation of the scatter.-type ScatterSpec v = [(Shape, Int, v)]- -- | A second-order array combinator (SOAC). data SOAC rep   = Stream SubExp [VName] [SubExp] (Lambda rep)-  | -- | @Scatter <length> <inputs> <lambda> <spec>@-    ---    -- Scatter maps values from a set of input arrays to indices and values of a-    -- set of output arrays. It is able to write multiple values to multiple-    -- outputs each of which may have multiple dimensions.-    ---    -- <inputs> is a list of input arrays, all having size <length>, elements of-    -- which are applied to the <lambda> function. For instance, if there are-    -- two arrays, <lambda> will get two values as input, one from each array.-    ---    -- <spec> specifies the result of the <lambda> and which arrays to-    -- write to.-    ---    -- <lambda> is a function that takes inputs from <inputs> and-    -- returns values according to <spec>. It returns values in the-    -- following manner:-    ---    --     [index_0, index_1, ..., index_n, value_0, value_1, ..., value_m]-    ---    -- For each output in <spec>, <lambda> returns <i> * <j> index-    -- values and <j> output values, where <i> is the number of-    -- dimensions (rank) of the given output, and <j> is the number of-    -- output values written to the given output.-    ---    -- For example, given the following scatter specification:-    ---    --     [([x1, y1, z1], 2, arr1), ([x2, y2], 1, arr2)]-    ---    -- <lambda> will produce 6 (3 * 2) index values and 2 output values for-    -- <arr1>, and 2 (2 * 1) index values and 1 output value for-    -- arr2. Additionally, the results are grouped, so the first 6 index values-    -- will correspond to the first two output values, and so on. For this-    -- example, <lambda> should return a total of 11 values, 8 index values and-    -- 3 output values.  See also 'splitScatterResults'.-    Scatter SubExp [VName] (ScatterSpec VName) (Lambda rep)   | -- | @Hist <length> <input arrays> <dest-arrays-and-ops> <bucket fun>@     --     -- The final lambda produces indexes and values for the 'HistOp's.@@ -325,56 +277,6 @@   guard $ null reds   pure map_lam --- | @splitScatterResults <scatter specification> <results>@------ Splits the results array into indices and values according to the--- specification.------ See 'groupScatterResults' for more information.-splitScatterResults :: [(Shape, Int, array)] -> [a] -> ([a], [a])-splitScatterResults output_spec results =-  let (shapes, ns, _) = unzip3 output_spec-      num_indices = sum $ zipWith (*) ns $ map length shapes-   in splitAt num_indices results---- | @groupScatterResults' <scatter specification> <results>@------ Blocks the index values and result values of <results> according to--- the specification. This is the simpler version of--- @groupScatterResults@, which doesn't return any information about--- shapes or output arrays.------ See 'groupScatterResults' for more information,-groupScatterResults' :: [(Shape, Int, array)] -> [a] -> [([a], a)]-groupScatterResults' output_spec results =-  let (indices, values) = splitScatterResults output_spec results-      (shapes, ns, _) = unzip3 output_spec-      chunk_sizes =-        concat $ zipWith (\shp n -> replicate n $ length shp) shapes ns-   in zip (chunks chunk_sizes indices) values---- | @groupScatterResults <scatter specification> <results>@------ Blocks the index values and result values of <results> according to the--- <scatter specification>.------ This function is used for extracting and grouping the results of a--- scatter. In the SOACS representation, the lambda inside a 'Scatter' returns--- all indices and values as one big list. This function groups each value with--- its corresponding indices (as determined by the t'Shape' of the output array).------ The elements of the resulting list correspond to the shape and name of the--- output parameters, in addition to a list of values written to that output--- parameter, along with the array indices marking where to write them to.------ See 'Scatter' for more information.-groupScatterResults :: ScatterSpec array -> [a] -> [(Shape, array, [([a], a)])]-groupScatterResults output_spec results =-  let (shapes, ns, arrays) = unzip3 output_spec-   in groupScatterResults' output_spec results-        & chunks ns-        & zip3 shapes arrays- -- | Like 'Mapper', but just for 'SOAC's. data SOACMapper frep trep m = SOACMapper   { mapOnSOACSubExp :: SubExp -> m SubExp,@@ -415,19 +317,6 @@     <*> mapM (mapOnSOACVName tv) arrs     <*> mapM (mapOnSOACSubExp tv) accs     <*> mapOnSOACLambda tv lam-mapSOACM tv (Scatter w ivs as lam) =-  Scatter-    <$> mapOnSOACSubExp tv w-    <*> mapM (mapOnSOACVName tv) ivs-    <*> mapM-      ( \(aw, an, a) ->-          (,,)-            <$> mapM (mapOnSOACSubExp tv) aw-            <*> pure an-            <*> mapOnSOACVName tv a-      )-      as-    <*> mapOnSOACLambda tv lam mapSOACM tv (Hist w arrs ops bucket_fun) =   Hist     <$> mapOnSOACSubExp tv w@@ -524,11 +413,6 @@     nms = map paramName $ take (1 + length accs) params     substs = M.fromList $ zip nms (outersize : accs)     Lambda params rtp _ = lam-soacType (Scatter _w _ivs dests lam) =-  zipWith arrayOfShape (map (snd . head) rets) shapes-  where-    (shapes, _, rets) =-      unzip3 $ groupScatterResults dests $ lambdaReturnType lam soacType (Hist _ _ ops _bucket_fun) = do   op <- ops   map (`arrayOfShape` histShape op) (lambdaReturnType $ histOp op)@@ -557,8 +441,6 @@       -- Drop the chunk parameter, which cannot alias anything.       paramsToInput =         zip (map paramName $ drop 1 $ lambdaParams lam) (accs ++ map Var arrs)-  consumedInOp (Scatter _ _ spec _) =-    namesFromList $ map (\(_, _, a) -> a) spec   consumedInOp (Hist _ _ ops _) =     namesFromList $ concatMap histDest ops @@ -576,8 +458,6 @@     VJP args vec (Alias.analyseLambda aliases lam)   addOpAliases aliases (Stream size arr accs lam) =     Stream size arr accs $ Alias.analyseLambda aliases lam-  addOpAliases aliases (Scatter len arrs dests lam) =-    Scatter len arrs dests (Alias.analyseLambda aliases lam)   addOpAliases aliases (Hist w arrs ops bucket_fun) =     Hist       w@@ -623,12 +503,6 @@       concatIndicesToEachValue is vs =         let is_flat = mconcat is          in map (is_flat <>) vs-  opDependencies (Scatter w arrs outputs lam) =-    let deps = lambdaDependencies mempty lam (depsOfArrays w arrs)-     in map flattenBlocks (groupScatterResults outputs deps)-    where-      flattenBlocks (_, arr, ivs) =-        oneName arr <> mconcat (map (mconcat . fst) ivs) <> mconcat (map snd ivs)   opDependencies (JVP args vec lam) =     mconcat $       replicate 2 $@@ -751,62 +625,6 @@   -- arr, so we can later check that aliases of arr are not used inside lam.   let fake_lamarrs' = map asArg lamarrs'   TC.checkLambda lam $ asArg inttp : accargs ++ fake_lamarrs'-typeCheckSOAC (Scatter w arrs as lam) = do-  -- Requirements:-  ---  --   0. @lambdaReturnType@ of @lam@ must be a list-  --      [index types..., value types, ...].-  ---  --   1. The number of index types and value types must be equal to the number-  --      of return values from @lam@.-  ---  --   2. Each index type must have the type i64.-  ---  --   3. Each array in @as@ and the value types must have the same type-  ---  --   4. Each array in @as@ is consumed.  This is not really a check, but more-  --      of a requirement, so that e.g. the source is not hoisted out of a-  --      loop, which will mean it cannot be consumed.-  ---  --   5. Each of arrs must be an array matching a corresponding lambda-  --      parameters.-  ---  -- Code:--  -- First check the input size.-  TC.require [Prim int64] w--  -- 0.-  let (as_ws, as_ns, _as_vs) = unzip3 as-      indexes = sum $ zipWith (*) as_ns $ map length as_ws-      rts = lambdaReturnType lam-      rtsI = take indexes rts-      rtsV = drop indexes rts--  -- 1.-  unless (length rts == sum as_ns + sum (zipWith (*) as_ns $ map length as_ws)) $-    TC.bad $-      TC.TypeError "Scatter: number of index types, value types and array outputs do not match."--  -- 2.-  forM_ rtsI $ \rtI ->-    unless (Prim int64 == rtI) $-      TC.bad $-        TC.TypeError "Scatter: Index return type must be i64."--  forM_ (zip (chunks as_ns rtsV) as) $ \(rtVs, (aw, _, a)) -> do-    -- All lengths must have type i64.-    mapM_ (TC.require [Prim int64]) aw--    -- 3.-    forM_ rtVs $ \rtV -> TC.requireI [arrayOfShape rtV aw] a--    -- 4.-    TC.consume =<< TC.lookupAliases a--  -- 5.-  arrargs <- TC.checkSOACArrayArgs w arrs-  TC.checkLambda lam arrargs typeCheckSOAC (Hist w arrs ops bucket_fun) = do   TC.require [Prim int64] w @@ -897,8 +715,6 @@     JVP args vec <$> rephraseLambda r lam   rephraseInOp r (Stream w arrs acc lam) =     Stream w arrs acc <$> rephraseLambda r lam-  rephraseInOp r (Scatter w arrs dests lam) =-    Scatter w arrs dests <$> rephraseLambda r lam   rephraseInOp r (Hist w arrs ops lam) =     Hist w arrs <$> mapM onOp ops <*> rephraseLambda r lam     where@@ -922,8 +738,6 @@     inside "JVP" $ lambdaMetrics lam   opMetrics (Stream _ _ _ lam) =     inside "Stream" $ lambdaMetrics lam-  opMetrics (Scatter _len _ _ lam) =-    inside "Scatter" $ lambdaMetrics lam   opMetrics (Hist _ _ ops bucket_fun) =     inside "Hist" $ mapM_ (lambdaMetrics . histOp) ops >> lambdaMetrics bucket_fun   opMetrics (Screma _ _ (ScremaForm map_lam scans reds)) =@@ -951,8 +765,6 @@         )   pretty (Stream size arrs acc lam) =     ppStream size arrs acc lam-  pretty (Scatter w arrs dests lam) =-    ppScatter w arrs dests lam   pretty (Hist w arrs ops bucket_fun) =     ppHist w arrs ops bucket_fun   pretty (Screma w arrs (ScremaForm map_lam scans reds))@@ -1011,21 +823,6 @@             </> ppTuple' (map pretty arrs)           <> comma             </> ppTuple' (map pretty acc)-          <> comma-            </> pretty lam-      )---- | Prettyprint the given Scatter.-ppScatter ::-  (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> [(Shape, Int, VName)] -> Lambda rep -> Doc ann-ppScatter w arrs dests lam =-  "scatter"-    <> (parens . align)-      ( pretty w-          <> comma-            </> ppTuple' (map pretty arrs)-          <> comma-            </> commasep (map pretty dests)           <> comma             </> pretty lam       )

@@ -24,9 +24,10 @@ import Control.Monad.Identity import Control.Monad.State import Control.Monad.Writer+import Data.Bifunctor import Data.Either import Data.Foldable-import Data.List (partition, transpose, unzip4, unzip6, zip6)+import Data.List (partition, transpose, unzip4) import Data.List.NonEmpty (NonEmpty (..)) import Data.Map.Strict qualified as M import Data.Maybe@@ -98,12 +99,6 @@   arr' <- mapM Engine.simplify arr   (lam', lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam   pure (Stream outerdim' arr' nes' lam', lam_hoisted)-simplifySOAC (Scatter w ivs as lam) = do-  w' <- Engine.simplify w-  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam-  ivs' <- mapM Engine.simplify ivs-  as' <- mapM Engine.simplify as-  pure (Scatter w' ivs' as' lam', hoisted) simplifySOAC (Hist w imgs ops bfun) = do   w' <- Engine.simplify w   (ops', hoisted) <- fmap unzip $@@ -198,7 +193,6 @@ topDownRules =   [ RuleOp hoistCerts,     RuleOp removeReplicateMapping,-    RuleOp removeReplicateWrite,     RuleOp removeUnusedSOACInput,     RuleOp simplifyClosedFormReduce,     RuleOp simplifyKnownIterationSOAC,@@ -214,7 +208,6 @@ bottomUpRules =   [ RuleOp removeDeadMapping,     RuleOp removeDeadReduction,-    RuleOp removeDeadWrite,     RuleBasicOp removeUnnecessaryCopy,     RuleOp liftIdentityStreaming,     RuleOp mapOpToOp@@ -346,14 +339,6 @@       auxing aux $ letBind pat $ Op $ soacOp $ Screma w arrs' $ mapSOAC fun' removeReplicateMapping _ _ _ _ = Skip --- | Like 'removeReplicateMapping', but for 'Scatter'.-removeReplicateWrite :: TopDownRuleOp (Wise SOACS)-removeReplicateWrite vtable pat aux (Scatter w ivs as lam)-  | Just (stms, lam', ivs') <- removeReplicateInput vtable lam ivs = Simplify $ do-      forM_ stms $ \(vs, cs, e) -> certifying cs $ letBindNames vs e-      auxing aux $ letBind pat $ Op $ Scatter w ivs' as lam'-removeReplicateWrite _ _ _ _ = Skip- removeReplicateInput ::   (Aliased rep) =>   ST.SymbolTable rep ->@@ -402,10 +387,6 @@     Just (used_arrs, map_lam') <- remove map_lam arrs =       Simplify . auxing aux . letBind pat . Op $         soacOp (Screma w used_arrs (ScremaForm map_lam' scan reduce))-  | Just (Scatter w arrs dests map_lam :: SOAC rep) <- asSOAC op,-    Just (used_arrs, map_lam') <- remove map_lam arrs =-      Simplify . auxing aux . letBind pat . Op $-        soacOp (Scatter w used_arrs dests map_lam')   where     used_in_body map_lam = freeIn $ lambdaBody map_lam     usedInput map_lam (param, _) = paramName param `nameIn` used_in_body map_lam@@ -576,64 +557,50 @@     removeDeadReduction' _ _ _ _ = Skip removeDeadReduction _ _ _ _ = Skip --- | If we are writing to an array that is never used, get rid of it.-removeDeadWrite :: BottomUpRuleOp (Wise SOACS)-removeDeadWrite (_, used) pat aux (Scatter w arrs dests fun) =-  let (i_ses, v_ses) = unzip $ groupScatterResults' dests $ bodyResult $ lambdaBody fun-      (i_ts, v_ts) = unzip $ groupScatterResults' dests $ lambdaReturnType fun-      isUsed (bindee, _, _, _, _, _) = (`UT.used` used) $ patElemName bindee-      (pat', i_ses', v_ses', i_ts', v_ts', dests') =-        unzip6 $ filter isUsed $ zip6 (patElems pat) i_ses v_ses i_ts v_ts dests-      fun' =-        fun-          { lambdaBody =-              mkBody (bodyStms (lambdaBody fun)) (concat i_ses' ++ v_ses'),-            lambdaReturnType = concat i_ts' ++ v_ts'-          }-   in if pat /= Pat pat'-        then-          Simplify . auxing aux . letBind (Pat pat') $-            Op (Scatter w arrs dests' fun')-        else Skip-removeDeadWrite _ _ _ _ = Skip---- handles now concatenation of more than two arrays+{-# NOINLINE fuseConcatScatter #-} fuseConcatScatter :: TopDownRuleOp (Wise SOACS)-fuseConcatScatter vtable pat _ (Scatter _ arrs dests fun)-  | Just (ws@(w' : _), xss, css) <- unzip3 <$> mapM isConcat arrs,+fuseConcatScatter vtable pat aux (Screma _ arrs form)+  | Just lam <- isMapSOAC form,+    all isAcc $ lambdaReturnType lam,+    Just (accs, (ws@(w' : _), ps, xss, css)) <-+      second unzip4 . partitionEithers+        <$> mapM isConcatOrAcc (zip (lambdaParams lam) arrs),     xivs <- transpose xss,-    all (w' ==) ws = Simplify $ do+    all (w' ==) ws = Simplify $ auxing aux $ certifying (mconcat css) $ do+      -- r is the amount of arrays being concatenated.       let r = length xivs-      fun2s <- replicateM (r - 1) (renameLambda fun)-      let (fun_is, fun_vs) =-            unzip . map (splitScatterResults dests . bodyResult . lambdaBody) $-              fun : fun2s-          (its, vts) =-            unzip . replicate r . splitScatterResults dests $ lambdaReturnType fun-          new_stmts = mconcat $ map (bodyStms . lambdaBody) (fun : fun2s)-      let fun' =-            Lambda-              { lambdaParams = mconcat $ map lambdaParams (fun : fun2s),-                lambdaBody = mkBody new_stmts $ mix fun_is <> mix fun_vs,-                lambdaReturnType = mix its <> mix vts-              }-      certifying (mconcat css) . letBind pat . Op $-        Scatter w' (concat xivs) (map (incWrites r) dests) fun'+          num_accs = length accs+      lams <- replicateM r (renameLambda lam {lambdaParams = map fst accs <> ps})+      let acc_params = map fst accs+          input_params = concatMap (drop num_accs . lambdaParams) lams+      lam' <-+        mkLambda (acc_params <> input_params) $+          subExpsRes <$> recurse (map (Var . paramName) acc_params) lams+      letBind pat $ Op $ Screma w' (map snd accs <> concat xivs) (mapSOAC lam')   where+    recurse accs [] = pure accs+    recurse accs (lam : lams) = do+      -- We know that the accumulators are the first params and that the rest are+      -- already bound.+      forM_ (zip accs (lambdaParams lam)) $ \(acc, p) ->+        letBindNames [paramName p] $ BasicOp $ SubExp acc+      accs' <- map resSubExp <$> bodyBind (lambdaBody lam)+      recurse accs' lams+     sizeOf :: VName -> Maybe SubExp     sizeOf x = arraySize 0 . typeOf <$> ST.lookup x vtable-    mix = concat . transpose-    incWrites r (w, n, a) = (w, n * r, a) -- ToDO: is it (n*r) or (n+r-1)??-    isConcat v = case ST.lookupExp v vtable of+    isConcatOrAcc (p@(Param _ _ Acc {}), v) =+      pure (Left (p, v))+    isConcatOrAcc (p, v) = case ST.lookupExp v vtable of       Just (BasicOp (Concat 0 (x :| ys) _), cs) -> do         x_w <- sizeOf x         y_ws <- mapM sizeOf ys         guard $ all (x_w ==) y_ws-        pure (x_w, x : ys, cs)+        pure (Right (x_w, p, x : ys, cs))       Just (BasicOp (Reshape arr newshape), cs)         | ReshapeCoerce <- reshapeKind newshape -> do-            (a, b, cs') <- isConcat arr-            pure (a, b, cs <> cs')+            Right (a, _, b, cs') <- isConcatOrAcc (p, arr)+            pure (Right (a, p, b, cs <> cs'))       _ -> Nothing fuseConcatScatter _ _ _ _ = Skip @@ -662,10 +629,10 @@               patElems pat           bindMapParam p a = do             a_t <- lookupType a-            letBindNames [paramName p] $-              BasicOp $-                Index a $-                  fullSlice a_t [DimFix $ constant (0 :: Int64)]+            letBindNames [paramName p] . BasicOp $+              if isAcc a_t+                then SubExp $ Var a+                else Index a $ fullSlice a_t [DimFix $ constant (0 :: Int64)]           bindArrayResult pe (SubExpRes cs se) =             certifying cs . letBindNames [patElemName pe] $               BasicOp $@@ -891,7 +858,7 @@     properArr [] arr = pure arr     properArr js arr = do       arr_t <- lookupType arr-      letExp (baseString arr) $ BasicOp $ Index arr $ fullSlice arr_t $ map DimFix js+      letExp (baseName arr) $ BasicOp $ Index arr $ fullSlice arr_t $ map DimFix js      mapOverArr w (pat, js, ArrayIndexing cs arr slice) = do       arr' <- properArr js arr@@ -900,9 +867,9 @@         if arraySize 0 arr_t == w           then pure arr'           else-            certifying cs . letExp (baseString arr ++ "_prefix") . BasicOp . Index arr' $+            certifying cs . letExp (baseName arr <> "_prefix") . BasicOp . Index arr' $               fullSlice arr_t [DimSlice (intConst Int64 0) w (intConst Int64 1)]-      arr_elem_param <- newParam (baseString arr ++ "_elem") (rowType arr_t)+      arr_elem_param <- newParam (baseName arr <> "_elem") (rowType arr_t)       pure $         Just           ( arr'',@@ -974,7 +941,7 @@           arr_t <- lookupType arr           let whole_dim = DimSlice (intConst Int64 0) (arraySize 0 arr_t) (intConst Int64 1)           arr_transformed <- certifying (arrayOpCerts op) $-            letExp (baseString arr ++ "_transformed") $+            letExp (baseName arr <> "_transformed") $               case op of                 ArrayIndexing _ _ (Slice slice) ->                   BasicOp $ Index arr $ Slice $ whole_dim : slice@@ -987,7 +954,7 @@                 ArrayVar {} ->                   BasicOp $ SubExp $ Var arr           arr_transformed_t <- lookupType arr_transformed-          arr_transformed_row <- newVName $ baseString arr ++ "_transformed_row"+          arr_transformed_row <- newVName $ baseName arr <> "_transformed_row"           pure $             Just               ( arr_transformed,@@ -1065,7 +1032,7 @@       (transformed, map_infos, stms <> oneStm stm)      mkTransformed (t, pe, (arr, cs, f)) = do-      v <- newVName (baseString (patElemName pe) <> "_pretr")+      v <- newVName (baseName (patElemName pe) <> "_pretr")       let bind = letBindNames [patElemName pe] $ f v       pure (SubExpRes cs (Var arr), t, v, bind) moveTransformToOutput _ _ _ _ =

@@ -21,9 +21,7 @@     SegBinOp (..),     segBinOpResults,     segBinOpChunks,-    KernelBody (..),-    aliasAnalyseKernelBody,-    consumedInKernelBody,+    KernelBody,     ResultManifest (..),     KernelResult (..),     kernelResultCerts,@@ -158,17 +156,7 @@ segBinOpChunks = chunks . map (length . segBinOpNeutral)  -- | The body of a 'SegOp'.-data KernelBody rep = KernelBody-  { kernelBodyDec :: BodyDec rep,-    kernelBodyStms :: Stms rep,-    kernelBodyResult :: [KernelResult]-  }--deriving instance (RepTypes rep) => Ord (KernelBody rep)--deriving instance (RepTypes rep) => Show (KernelBody rep)--deriving instance (RepTypes rep) => Eq (KernelBody rep)+type KernelBody rep = GBody rep KernelResult  -- | Metadata about whether there is a subtle point to this -- 'KernelResult'.  This is used to protect things like tiling, which@@ -193,10 +181,6 @@     -- Whether this is a result-per-thread or a     -- result-per-block depends on where the 'SegOp' occurs.     Returns ResultManifest Certs SubExp-  | WriteReturns-      Certs-      VName -- Destination array-      [(Slice SubExp, SubExp)]   | TileReturns       Certs       [(SubExp, SubExp)] -- Total/tile for each dimension@@ -214,49 +198,31 @@       VName -- Tile returned by this thread/block.   deriving (Eq, Show, Ord) +instance IsResult KernelResult where+  resAliases = freeIn . kernelResultSubExp+ -- | Get the certs for this 'KernelResult'. kernelResultCerts :: KernelResult -> Certs kernelResultCerts (Returns _ cs _) = cs-kernelResultCerts (WriteReturns cs _ _) = cs kernelResultCerts (TileReturns cs _ _) = cs kernelResultCerts (RegTileReturns cs _ _) = cs  -- | Get the root t'SubExp' corresponding values for a 'KernelResult'. kernelResultSubExp :: KernelResult -> SubExp kernelResultSubExp (Returns _ _ se) = se-kernelResultSubExp (WriteReturns _ arr _) = Var arr kernelResultSubExp (TileReturns _ _ v) = Var v kernelResultSubExp (RegTileReturns _ _ v) = Var v  instance FreeIn KernelResult where   freeIn' (Returns _ cs what) = freeIn' cs <> freeIn' what-  freeIn' (WriteReturns cs arr res) = freeIn' cs <> freeIn' arr <> freeIn' res   freeIn' (TileReturns cs dims v) =     freeIn' cs <> freeIn' dims <> freeIn' v   freeIn' (RegTileReturns cs dims_n_tiles v) =     freeIn' cs <> freeIn' dims_n_tiles <> freeIn' v -instance (ASTRep rep) => FreeIn (KernelBody rep) where-  freeIn' (KernelBody dec stms res) =-    fvBind bound_in_stms $ freeIn' dec <> freeIn' stms <> freeIn' res-    where-      bound_in_stms = foldMap boundByStm stms--instance (ASTRep rep) => Substitute (KernelBody rep) where-  substituteNames subst (KernelBody dec stms res) =-    KernelBody-      (substituteNames subst dec)-      (substituteNames subst stms)-      (substituteNames subst res)- instance Substitute KernelResult where   substituteNames subst (Returns manifest cs se) =     Returns manifest (substituteNames subst cs) (substituteNames subst se)-  substituteNames subst (WriteReturns cs arr res) =-    WriteReturns-      (substituteNames subst cs)-      (substituteNames subst arr)-      (substituteNames subst res)   substituteNames subst (TileReturns cs dims v) =     TileReturns       (substituteNames subst cs)@@ -268,47 +234,52 @@       (substituteNames subst dims_n_tiles)       (substituteNames subst v) -instance (ASTRep rep) => Rename (KernelBody rep) where-  rename (KernelBody dec stms res) = do-    dec' <- rename dec-    renamingStms stms $ \stms' ->-      KernelBody dec' stms' <$> rename res- instance Rename KernelResult where   rename = substituteRename --- | Perform alias analysis on a 'KernelBody'.-aliasAnalyseKernelBody ::-  (Alias.AliasableRep rep) =>-  AliasTable ->-  KernelBody rep ->-  KernelBody (Aliases rep)-aliasAnalyseKernelBody aliases (KernelBody dec stms res) =-  let Body dec' stms' _ = Alias.analyseBody aliases $ Body dec stms []-   in KernelBody dec' stms' res+checkKernelResult ::+  (TC.Checkable rep) =>+  KernelResult ->+  Type ->+  TC.TypeM rep ()+checkKernelResult (Returns _ cs what) t = do+  TC.checkCerts cs+  TC.require [t] what+checkKernelResult (TileReturns cs dims v) t = do+  TC.checkCerts cs+  forM_ dims $ \(dim, tile) -> do+    TC.require [Prim int64] dim+    TC.require [Prim int64] tile+  vt <- lookupType v+  unless (vt == t `arrayOfShape` Shape (map snd dims)) $+    TC.bad $+      TC.TypeError $+        "Invalid type for TileReturns " <> prettyText v+checkKernelResult (RegTileReturns cs dims_n_tiles arr) t = do+  TC.checkCerts cs+  mapM_ (TC.require [Prim int64]) dims+  mapM_ (TC.require [Prim int64]) blk_tiles+  mapM_ (TC.require [Prim int64]) reg_tiles --- | The variables consumed in the kernel body.-consumedInKernelBody ::-  (Aliased rep) =>-  KernelBody rep ->-  Names-consumedInKernelBody (KernelBody dec stms res) =-  consumedInBody (Body dec stms []) <> mconcat (map consumedByReturn res)+  -- assert that arr is of element type t and shape (rev outer_tiles ++ reg_tiles)+  arr_t <- lookupType arr+  unless (arr_t == expected) $+    TC.bad . TC.TypeError $+      "Invalid type for TileReturns. Expected:\n  "+        <> prettyText expected+        <> ",\ngot:\n  "+        <> prettyText arr_t   where-    consumedByReturn (WriteReturns _ a _) = oneName a-    consumedByReturn _ = mempty+    (dims, blk_tiles, reg_tiles) = unzip3 dims_n_tiles+    expected = t `arrayOfShape` Shape (blk_tiles <> reg_tiles)  checkKernelBody ::   (TC.Checkable rep) =>   [Type] ->   KernelBody (Aliases rep) ->   TC.TypeM rep ()-checkKernelBody ts (KernelBody (_, dec) stms kres) = do+checkKernelBody ts (Body (_, dec) stms kres) = do   TC.checkBodyDec dec-  -- We consume the kernel results (when applicable) before-  -- type-checking the stms, so we will get an error if a statement-  -- uses an array that is written to in a result.-  mapM_ consumeKernelResult kres   TC.checkStms stms $ do     unless (length ts == length kres) $       TC.bad . TC.TypeError $@@ -318,62 +289,9 @@           <> prettyText (length kres)           <> " values."     zipWithM_ checkKernelResult kres ts-  where-    consumeKernelResult (WriteReturns _ arr _) =-      TC.consume =<< TC.lookupAliases arr-    consumeKernelResult _ =-      pure () -    checkKernelResult (Returns _ cs what) t = do-      TC.checkCerts cs-      TC.require [t] what-    checkKernelResult (WriteReturns cs arr res) t = do-      TC.checkCerts cs-      arr_t <- lookupType arr-      unless (arr_t == t) $-        TC.bad . TC.TypeError $-          "WriteReturns result type annotation for "-            <> prettyText arr-            <> " is "-            <> prettyText t-            <> ", but inferred as"-            <> prettyText arr_t-      forM_ res $ \(slice, e) -> do-        TC.checkSlice arr_t slice-        TC.require [t `setArrayShape` sliceShape slice] e-    checkKernelResult (TileReturns cs dims v) t = do-      TC.checkCerts cs-      forM_ dims $ \(dim, tile) -> do-        TC.require [Prim int64] dim-        TC.require [Prim int64] tile-      vt <- lookupType v-      unless (vt == t `arrayOfShape` Shape (map snd dims)) $-        TC.bad $-          TC.TypeError $-            "Invalid type for TileReturns " <> prettyText v-    checkKernelResult (RegTileReturns cs dims_n_tiles arr) t = do-      TC.checkCerts cs-      mapM_ (TC.require [Prim int64]) dims-      mapM_ (TC.require [Prim int64]) blk_tiles-      mapM_ (TC.require [Prim int64]) reg_tiles--      -- assert that arr is of element type t and shape (rev outer_tiles ++ reg_tiles)-      arr_t <- lookupType arr-      unless (arr_t == expected) $-        TC.bad . TC.TypeError $-          "Invalid type for TileReturns. Expected:\n  "-            <> prettyText expected-            <> ",\ngot:\n  "-            <> prettyText arr_t-      where-        (dims, blk_tiles, reg_tiles) = unzip3 dims_n_tiles-        expected = t `arrayOfShape` Shape (blk_tiles <> reg_tiles)--kernelBodyMetrics :: (OpMetrics (Op rep)) => KernelBody rep -> MetricsM ()-kernelBodyMetrics = mapM_ stmMetrics . kernelBodyStms- instance (PrettyRep rep) => Pretty (KernelBody rep) where-  pretty (KernelBody _ stms res) =+  pretty (Body _ stms res) =     PP.stack (map pretty (stmsToList stms))       </> "return"       <+> PP.braces (PP.commastack $ map pretty res)@@ -390,12 +308,6 @@     hsep $ certAnnots cs <> ["returns (private)" <+> pretty what]   pretty (Returns ResultMaySimplify cs what) =     hsep $ certAnnots cs <> ["returns" <+> pretty what]-  pretty (WriteReturns cs arr res) =-    hsep $-      certAnnots cs-        <> [pretty arr </> "with" <+> PP.apply (map ppRes res)]-    where-      ppRes (slice, e) = pretty slice <+> "=" <+> pretty e   pretty (TileReturns cs dims v) =     hsep $ certAnnots cs <> ["tile" <> apply (map onDim dims) <+> pretty v]     where@@ -443,9 +355,7 @@ -- -- The type list is usually the type of the element returned by a -- single thread. The result of the SegOp is then an array of that--- type, with the shape of the 'SegSpace' prepended. One exception is--- for 'WriteReturns', where the type annotation is the /full/ type of--- the result.+-- type, with the shape of the 'SegSpace' prepended. data SegOp lvl rep   = SegMap lvl SegSpace [Type] (KernelBody rep)   | -- | The KernelSpace must always have at least two dimensions,@@ -479,8 +389,6 @@     SegHist _ _ _ body _ -> body  segResultShape :: SegSpace -> Type -> KernelResult -> Type-segResultShape _ t (WriteReturns {}) =-  t segResultShape space t Returns {} =   foldr (flip arrayOfRow) t $ segSpaceDims space segResultShape _ t (TileReturns _ dims _) =@@ -491,13 +399,13 @@ -- | The return type of a 'SegOp'. segOpType :: SegOp lvl rep -> [Type] segOpType (SegMap _ space ts kbody) =-  zipWith (segResultShape space) ts $ kernelBodyResult kbody+  zipWith (segResultShape space) ts $ bodyResult kbody segOpType (SegRed _ space ts kbody reds) =   red_ts     ++ zipWith       (segResultShape space)       map_ts-      (drop (length red_ts) $ kernelBodyResult kbody)+      (drop (length red_ts) $ bodyResult kbody)   where     map_ts = drop (length red_ts) ts     segment_dims = init $ segSpaceDims space@@ -510,7 +418,7 @@     ++ zipWith       (segResultShape space)       map_ts-      (drop (length scan_ts) $ kernelBodyResult kbody)+      (drop (length scan_ts) $ bodyResult kbody)   where     map_ts = drop (length scan_ts) ts     scan_ts = do@@ -532,13 +440,13 @@   opAliases = map (const mempty) . segOpType    consumedInOp (SegMap _ _ _ kbody) =-    consumedInKernelBody kbody+    consumedInBody kbody   consumedInOp (SegRed _ _ _ kbody _) =-    consumedInKernelBody kbody+    consumedInBody kbody   consumedInOp (SegScan _ _ _ kbody _) =-    consumedInKernelBody kbody+    consumedInBody kbody   consumedInOp (SegHist _ _ _ kbody ops) =-    namesFromList (concatMap histDest ops) <> consumedInKernelBody kbody+    namesFromList (concatMap histDest ops) <> consumedInBody kbody  -- | Type check a 'SegOp', given a checker for its level. typeCheckSegOp ::@@ -760,28 +668,20 @@ rephraseBinOp r (SegBinOp comm lam nes shape) =   SegBinOp comm <$> rephraseLambda r lam <*> pure nes <*> pure shape -rephraseKernelBody ::-  (Monad f) =>-  Rephraser f from rep ->-  KernelBody from ->-  f (KernelBody rep)-rephraseKernelBody r (KernelBody dec stms res) =-  KernelBody <$> rephraseBodyDec r dec <*> traverse (rephraseStm r) stms <*> pure res- instance RephraseOp (SegOp lvl) where   rephraseInOp r (SegMap lvl space ts body) =-    SegMap lvl space ts <$> rephraseKernelBody r body+    SegMap lvl space ts <$> rephraseBody r body   rephraseInOp r (SegRed lvl space ts body reds) =     SegRed lvl space ts-      <$> rephraseKernelBody r body+      <$> rephraseBody r body       <*> mapM (rephraseBinOp r) reds   rephraseInOp r (SegScan lvl space ts body scans) =     SegScan lvl space ts-      <$> rephraseKernelBody r body+      <$> rephraseBody r body       <*> mapM (rephraseBinOp r) scans   rephraseInOp r (SegHist lvl space ts body hists) =     SegHist lvl space ts-      <$> rephraseKernelBody r body+      <$> rephraseBody r body       <*> mapM onOp hists     where       onOp (HistOp w rf arrs nes shape op) =@@ -798,8 +698,8 @@         { mapOnSegOpLambda = traverseLambdaStms f',           mapOnSegOpBody = onBody         }-    onBody (KernelBody dec stms res) =-      KernelBody dec <$> f seg_scope stms <*> pure res+    onBody (Body dec stms res) =+      Body dec <$> f seg_scope stms <*> pure res  instance   (ASTRep rep, Substitute lvl) =>@@ -840,19 +740,19 @@  instance (OpMetrics (Op rep)) => OpMetrics (SegOp lvl rep) where   opMetrics (SegMap _ _ _ body) =-    inside "SegMap" $ kernelBodyMetrics body+    inside "SegMap" $ bodyMetrics body   opMetrics (SegRed _ _ _ body reds) =     inside "SegRed" $ do       mapM_ (inside "SegBinOp" . lambdaMetrics . segBinOpLambda) reds-      kernelBodyMetrics body+      bodyMetrics body   opMetrics (SegScan _ _ _ body scans) =     inside "SegScan" $ do       mapM_ (inside "SegBinOp" . lambdaMetrics . segBinOpLambda) scans-      kernelBodyMetrics body+      bodyMetrics body   opMetrics (SegHist _ _ _ body ops) =     inside "SegHist" $ do       mapM_ (lambdaMetrics . histOp) ops-      kernelBodyMetrics body+      bodyMetrics body  instance Pretty SegSpace where   pretty (SegSpace phys dims) =@@ -929,31 +829,22 @@         SegOpMapper           pure           (pure . Alias.analyseLambda aliases)-          (pure . aliasAnalyseKernelBody aliases)+          (pure . Alias.analyseBody aliases)           pure           pure -informKernelBody :: (Informing rep) => KernelBody rep -> KernelBody (Wise rep)-informKernelBody (KernelBody dec stms res) =-  mkWiseKernelBody dec (informStms stms) res- instance CanBeWise (SegOp lvl) where   addOpWisdom = runIdentity . mapSegOpM add     where       add =-        SegOpMapper-          pure-          (pure . informLambda)-          (pure . informKernelBody)-          pure-          pure+        SegOpMapper pure (pure . informLambda) (pure . informBody) pure pure  instance (ASTRep rep) => ST.IndexOp (SegOp lvl rep) where   indexOp vtable k (SegMap _ space _ kbody) is = do-    Returns ResultMaySimplify _ se <- maybeNth k $ kernelBodyResult kbody+    Returns ResultMaySimplify _ se <- maybeNth k $ bodyResult kbody     guard $ length gtids <= length is     let idx_table = M.fromList $ zip gtids $ map (ST.Indexed mempty . untyped) is-        idx_table' = foldl' expandIndexedTable idx_table $ kernelBodyStms kbody+        idx_table' = foldl' expandIndexedTable idx_table $ bodyStms kbody     case se of       Var v -> M.lookup v idx_table'       _ -> Nothing@@ -1006,11 +897,6 @@ instance Engine.Simplifiable KernelResult where   simplify (Returns manifest cs what) =     Returns manifest <$> Engine.simplify cs <*> Engine.simplify what-  simplify (WriteReturns cs a res) =-    WriteReturns-      <$> Engine.simplify cs-      <*> Engine.simplify a-      <*> Engine.simplify res   simplify (TileReturns cs dims what) =     TileReturns <$> Engine.simplify cs <*> Engine.simplify dims <*> Engine.simplify what   simplify (RegTileReturns cs dims_n_tiles what) =@@ -1019,35 +905,12 @@       <*> Engine.simplify dims_n_tiles       <*> Engine.simplify what -mkWiseKernelBody ::-  (Informing rep) =>-  BodyDec rep ->-  Stms (Wise rep) ->-  [KernelResult] ->-  KernelBody (Wise rep)-mkWiseKernelBody dec stms res =-  let Body dec' _ _ = mkWiseBody dec stms $ subExpsRes res_vs-   in KernelBody dec' stms res-  where-    res_vs = map kernelResultSubExp res--mkKernelBodyM ::-  (MonadBuilder m) =>-  Stms (Rep m) ->-  [KernelResult] ->-  m (KernelBody (Rep m))-mkKernelBodyM stms kres = do-  Body dec' _ _ <- mkBodyM stms $ subExpsRes res_ses-  pure $ KernelBody dec' stms kres-  where-    res_ses = map kernelResultSubExp kres- simplifyKernelBody ::   (Engine.SimplifiableRep rep, BodyDec rep ~ ()) =>   SegSpace ->   KernelBody (Wise rep) ->   Engine.SimpleM rep (KernelBody (Wise rep), Stms (Wise rep))-simplifyKernelBody space (KernelBody _ stms res) = do+simplifyKernelBody space (Body _ stms res) = do   par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers    let blocker =@@ -1060,8 +923,7 @@    -- Ensure we do not try to use anything that is consumed in the result.   (body_res, body_stms, hoisted) <--    Engine.localVtable (flip (foldl' (flip ST.consume)) (foldMap consumedInResult res))-      . Engine.localVtable (<> scope_vtable)+    Engine.localVtable (<> scope_vtable)       . Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True})       . Engine.enterLoop       $ Engine.blockIf blocker stms@@ -1071,16 +933,11 @@             mapM Engine.simplify res         pure (res', UT.usages $ freeIn res') -  pure (mkWiseKernelBody () body_stms body_res, hoisted)+  pure (mkWiseBody () body_stms body_res, hoisted)   where     scope_vtable = segSpaceSymbolTable space     bound_here = namesFromList $ M.keys $ scopeOfSegSpace space -    consumedInResult (WriteReturns _ arr _) =-      [arr]-    consumedInResult _ =-      []- simplifyLambda ::   (Engine.SimplifiableRep rep) =>   Names ->@@ -1221,14 +1078,14 @@   Rule rep -- If a SegOp produces something invariant to the SegOp, turn it -- into a replicate.-topDownSegOp vtable (Pat kpes) dec (SegMap lvl space ts (KernelBody _ kstms kres)) = Simplify $ do+topDownSegOp vtable (Pat kpes) dec (SegMap lvl space ts (Body _ kstms kres)) = Simplify $ do   (ts', kpes', kres') <-     unzip3 <$> filterM checkForInvarianceResult (zip3 ts kpes kres)    -- Check if we did anything at all.   when (kres == kres') cannotSimplify -  kbody <- mkKernelBodyM kstms kres'+  kbody <- mkBodyM kstms kres'   addStm $ Let (Pat kpes') dec $ Op $ segOp $ SegMap lvl space ts' kbody   where     isInvariant Constant {} = True@@ -1260,12 +1117,12 @@           (red_pes', red_ts', red_res') = unzip3 $ concat aux           pes' = red_pes' ++ map_pes           ts' = red_ts' ++ map_ts-          kbody' = kbody {kernelBodyResult = red_res' ++ map_res}+          kbody' = kbody {bodyResult = red_res' ++ map_res}       letBind (Pat pes') $ Op $ segOp $ SegRed lvl space ts' kbody' ops'   where     (red_pes, map_pes) = splitAt (segBinOpResults ops) pes     (red_ts, map_ts) = splitAt (segBinOpResults ops) ts-    (red_res, map_res) = splitAt (segBinOpResults ops) $ kernelBodyResult kbody+    (red_res, map_res) = splitAt (segBinOpResults ops) $ bodyResult kbody      sameShape (op1, _) (op2, _) =       segBinOpShape op1 == segBinOpShape op2@@ -1335,11 +1192,11 @@   -- results, so we only deal with map results for now.   | (_, kpes', kts', kres') <- unzip4 $ filter keep $ zip4 [0 ..] kpes kts kres,     kpes' /= kpes = Simplify $ do-      kbody' <- localScope (scopeOfSegSpace space) $ mkKernelBodyM kstms kres'+      kbody' <- localScope (scopeOfSegSpace space) $ mkBodyM kstms kres'       addStm $ Let (Pat kpes') dec $ Op $ segOp $ mk_segop kts' kbody'   where     space = segSpace segop-    (kts, KernelBody _ kstms kres, num_nonmap_results, mk_segop) =+    (kts, Body _ kstms kres, num_nonmap_results, mk_segop) =       segOpGuts segop      keep (i, pe, _, _) =@@ -1358,11 +1215,11 @@   when (kpes' == kpes) cannotSimplify    kbody' <--    localScope (scopeOfSegSpace space) $ mkKernelBodyM kstms' kres'+    localScope (scopeOfSegSpace space) $ mkBodyM kstms' kres'    addStm $ Let (Pat kpes') dec $ Op $ segOp $ mk_segop kts' kbody'   where-    (kts, KernelBody _ kstms kres, num_nonmap_results, mk_segop) =+    (kts, Body _ kstms kres, num_nonmap_results, mk_segop) =       segOpGuts segop     free_in_kstms = foldMap freeIn kstms     space = segSpace segop@@ -1393,7 +1250,7 @@                 letBindNames [patElemName kpe'] . BasicOp . Index arr $                   Slice $                     outer_slice <> remaining_slice-          precopy <- newVName $ baseString (patElemName kpe) <> "_precopy"+          precopy <- newVName $ baseName (patElemName kpe) <> "_precopy"           index kpe {patElemName = precopy}           letBindNames [patElemName kpe] $ BasicOp $ Replicate mempty $ Var precopy           pure@@ -1421,26 +1278,16 @@  --- Memory -kernelBodyReturns ::-  (Mem rep inner, HasScope rep m, Monad m) =>-  KernelBody somerep ->-  [ExpReturns] ->-  m [ExpReturns]-kernelBodyReturns = zipWithM correct . kernelBodyResult-  where-    correct (WriteReturns _ arr _) _ = varReturns arr-    correct _ ret = pure ret- -- | Like 'segOpType', but for memory representations. segOpReturns ::   (Mem rep inner, Monad m, HasScope rep m) =>   SegOp lvl rep ->   m [ExpReturns]-segOpReturns k@(SegMap _ _ _ kbody) =-  kernelBodyReturns kbody . extReturns =<< opType k-segOpReturns k@(SegRed _ _ _ kbody _) =-  kernelBodyReturns kbody . extReturns =<< opType k-segOpReturns k@(SegScan _ _ _ kbody _) =-  kernelBodyReturns kbody . extReturns =<< opType k+segOpReturns k@(SegMap {}) =+  extReturns <$> opType k+segOpReturns k@(SegRed {}) =+  extReturns <$> opType k+segOpReturns k@(SegScan {}) =+  extReturns <$> opType k segOpReturns (SegHist _ _ _ _ ops) =   concat <$> mapM (mapM varReturns . histDest) ops

@@ -125,7 +125,8 @@     Stms,     SubExpRes (..),     Result,-    Body (..),+    GBody (..),+    Body,     BasicOp (..),     UnOp (..),     BinOp (..),@@ -300,20 +301,24 @@ -- | The result of a body is a sequence of subexpressions. type Result = [SubExpRes] --- | A body consists of a sequence of statements, terminating in a--- list of result values.-data Body rep = Body+-- | A generalised body consists of a sequence of statements, terminated in some+-- kind of result.+data GBody rep res = Body   { bodyDec :: BodyDec rep,     bodyStms :: Stms rep,-    bodyResult :: Result+    bodyResult :: [res]   } -deriving instance (RepTypes rep) => Ord (Body rep)+-- | A body consists of a sequence of statements, terminating in a+-- list of result values.+type Body rep = GBody rep SubExpRes -deriving instance (RepTypes rep) => Show (Body rep)+deriving instance (RepTypes rep, Ord res) => Ord (GBody rep res) -deriving instance (RepTypes rep) => Eq (Body rep)+deriving instance (RepTypes rep, Show res) => Show (GBody rep res) +deriving instance (RepTypes rep, Eq res) => Eq (GBody rep res)+ -- | Apart from being Opaque, what else is going on here? data OpaqueOp   = -- | No special operation.@@ -395,7 +400,7 @@     Index VName (Slice SubExp)   | -- | An in-place update of the given array at the given position.     -- Consumes the array.  If 'Safe', perform a run-time bounds check-    -- and ignore the write if out of bounds (like @Scatter@).+    -- and ignore the write if out of bounds (scatter-like).     Update Safety VName (Slice SubExp) SubExp   | FlatIndex VName (FlatSlice SubExp)   | FlatUpdate VName (FlatSlice SubExp) VName@@ -433,7 +438,7 @@   | -- | Update an accumulator at the given index with the given     -- value. Consumes the accumulator and produces a new one. If     -- 'Safe', perform a run-time bounds check and ignore the write if-    -- out of bounds (like @Scatter@).+    -- out of bounds (scatter-like).     UpdateAcc Safety VName [SubExp] [SubExp]   deriving (Eq, Ord, Show) 

@@ -185,9 +185,13 @@   -- | Check whether one shape if a subset of another shape.   subShapeOf :: a -> a -> Bool +  -- | Prepend the dimensions of a 'Shape'.+  prependShape :: Shape -> a -> a+ instance ArrayShape (ShapeBase SubExp) where   shapeRank (Shape l) = length l   subShapeOf = (==)+  prependShape = (<>)  instance ArrayShape (ShapeBase ExtSize) where   shapeRank (Shape l) = length l@@ -210,6 +214,8 @@             put $ M.insert y x extmap             pure True +  prependShape shape = (fmap Free shape <>)+ instance Semigroup Rank where   Rank x <> Rank y = Rank $ x + y @@ -219,6 +225,7 @@ instance ArrayShape Rank where   shapeRank (Rank x) = x   subShapeOf = (==)+  prependShape shape (Rank x) = Rank $ shapeRank shape + x  -- | The memory space of a block.  If 'DefaultSpace', this is the "default" -- space, whatever that is.  The exact meaning of the 'SpaceId'

@@ -964,6 +964,8 @@         <> prettyText (length is)         <> " provided." +  mapM_ (require [Prim int64]) is+   zipWithM_ require (map pure ts) ses   consume =<< lookupAliases acc 

@@ -85,7 +85,7 @@ ensureExtShape ::   ErrorMsg SubExp ->   ExtType ->-  String ->+  Name ->   SubExp ->   InternaliseM SubExp ensureExtShape msg t name orig@@ -97,7 +97,7 @@ ensureShape ::   ErrorMsg SubExp ->   Type ->-  String ->+  Name ->   SubExp ->   InternaliseM SubExp ensureShape msg = ensureExtShape msg . staticShapes1@@ -119,12 +119,12 @@     ensureArgShape t (Var v)       | arrayRank t < 1 = pure $ Var v       | otherwise =-          ensureShape msg t (baseString v) $ Var v+          ensureShape msg t (baseName v) $ Var v  ensureShapeVar ::   ErrorMsg SubExp ->   ExtType ->-  String ->+  Name ->   VName ->   InternaliseM VName ensureShapeVar msg t name v

@@ -147,7 +147,7 @@       let name = E.identName bindee       case internalisedTypeSize $ E.unInfo $ E.identType bindee of         1 -> pure [name]-        n -> replicateM n $ newVName $ baseString name+        n -> replicateM n $ newName name  bindingFlatPat ::   (Show t) =>

@@ -765,7 +765,7 @@       let t' = second (const d) t       x <- case p of         Named x | x `notElem` prev -> pure x-        _ -> newNameFromString "eta_p"+        _ -> newVName "eta_p"       pure         ( x : prev,           ( Id x (Info t') mempty,@@ -874,8 +874,8 @@           (params, body, ret) <- etaExpand (RetType [] $ toRes Nonunique t) e           defuncFun [] params body ret mempty       | otherwise -> do-          fname <- newVName $ "dyn_" <> baseString (qualLeaf qn)-          let (pats, e0, sv') = liftDynFun (prettyString qn) sv num_args+          fname <- newVName $ "dyn_" <> baseName (qualLeaf qn)+          let (pats, e0, sv') = liftDynFun (nameFromText (prettyText qn)) sv num_args               (argtypes', rettype') = dynamicFunType sv' argtypes               dims' = mempty @@ -902,15 +902,15 @@ -- Embed some information about the original function -- into the name of the lifted function, to make the -- result slightly more human-readable.-liftedName :: Int -> Exp -> String+liftedName :: Int -> Exp -> Name liftedName i (Var f _ _) =-  "defunc_" ++ show i ++ "_" ++ baseString (qualLeaf f)+  "defunc_" <> nameFromString (show i) <> "_" <> baseName (qualLeaf f) liftedName i (AppExp (Apply f _ _) _) =   liftedName (i + 1) f liftedName _ _ = "defunc"  defuncApplyArg ::-  (String, SrcLoc) ->+  (Name, SrcLoc) ->   (Exp, StaticVal) ->   ((Maybe VName, Exp), [ParamType]) ->   DefM (Exp, StaticVal)@@ -950,7 +950,7 @@   let bound_sizes = S.fromList (dims <> more_dims) <> globals   params' <- instAnySizes params -  fname <- newNameFromString fname_s+  fname <- newVName fname_s   liftValDec     fname     lifted_rettype@@ -1037,7 +1037,7 @@ -- dimensions, a list of parameters, a function body, and the -- appropriate static value for applying the function at the given -- depth of partial application.-liftDynFun :: String -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal)+liftDynFun :: Name -> StaticVal -> Int -> ([Pat ParamType], Exp, StaticVal) liftDynFun _ (DynamicFun (e, sv) _) 0 = ([], e, sv) liftDynFun s (DynamicFun clsr@(_, LambdaSV pat _ _ _) sv) d   | d > 0 =@@ -1045,7 +1045,7 @@        in (pat : pats, e', DynamicFun clsr sv') liftDynFun s sv d =   error $-    s+    nameToString s       ++ " Tried to lift a StaticVal "       ++ take 100 (show sv)       ++ ", but expected a dynamic function.\n"

@@ -179,7 +179,10 @@       let ns = map (internalisedTypeSize . E.entryType) ets           is' = chunks ns is       ets' <- map snd <$> zipWithM (entryPointType types) ets (map (map (ts !!)) is')-      pure $ zip ets' $ map (map (+ 1)) is' -- Adjust for tag.+      pure . zip ets' $+        if length cs == 1+          then is'+          else map (map (+ 1)) is' -- Adjust for tag.  entryPointTypeName :: I.EntryPointType -> Name entryPointTypeName (I.TypeOpaque v) = v@@ -223,8 +226,9 @@           let (_, places) = internaliseSumTypeRep cs               cs' = sumConstrs types cs $ E.entryAscribed t               cs'' = zip (map fst cs') (zip (map snd cs') (map snd places))+              ts' = if length cs == 1 then ts else drop 1 ts           addType desc . I.OpaqueSum (map valueType ts)-            =<< opaqueSum types cs'' (drop 1 ts)+            =<< opaqueSum types cs'' ts'         E.Array _ shape (E.Record fs)           | not $ null fs -> do               let fs' = recordFields types fs $ E.entryAscribed t

@@ -61,7 +61,7 @@             ]      (body', rettype') <- buildBody $ do-      body_res <- internaliseExp (baseString fname <> "_res") body+      body_res <- internaliseExp (baseName fname <> "_res") body       (rettype', retals) <-         first zeroExts . unzip . internaliseReturnType (map (fmap paramDeclType) all_params) rettype           <$> mapM subExpType body_res@@ -156,7 +156,7 @@   where     zeroExts ts = generaliseExtTypes ts ts -internaliseBody :: String -> E.Exp -> InternaliseM (Body SOACS)+internaliseBody :: Name -> E.Exp -> InternaliseM (Body SOACS) internaliseBody desc e =   buildBody_ $ subExpsRes <$> internaliseExp (desc <> "_res") e @@ -169,7 +169,7 @@  -- | Only returns those pattern names that are not used in the pattern -- itself (the "non-existential" part, you could say).-letValExp :: String -> I.Exp SOACS -> InternaliseM [VName]+letValExp :: Name -> I.Exp SOACS -> InternaliseM [VName] letValExp name e = do   e_t <- expExtType e   names <- replicateM (length e_t) $ newVName name@@ -177,11 +177,11 @@   let ctx = shapeContext e_t   pure $ map fst $ filter ((`S.notMember` ctx) . snd) $ zip names [0 ..] -letValExp' :: String -> I.Exp SOACS -> InternaliseM [SubExp]+letValExp' :: Name -> I.Exp SOACS -> InternaliseM [SubExp] letValExp' _ (BasicOp (SubExp se)) = pure [se] letValExp' name ses = map I.Var <$> letValExp name ses -internaliseAppExp :: String -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp]+internaliseAppExp :: Name -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp] internaliseAppExp desc _ (E.Index e idxs _) = do   vs <- internaliseExpToVars "indexed" e   dims <- case vs of@@ -357,22 +357,22 @@     (FunctionName qfname, args) -> do       -- Argument evaluation is outermost-in so that any existential sizes       -- created by function applications can be brought into scope.-      let fname = nameFromString $ prettyString $ baseName $ qualLeaf qfname-          arg_desc = nameToString fname ++ "_arg"+      let fname = baseName $ qualLeaf qfname+          arg_desc = fname <> "_arg"        -- Some functions are magical (overloaded) and we handle that here.       case () of         ()           -- Short-circuiting operators are magical.           | baseTag (qualLeaf qfname) <= maxIntrinsicTag,-            baseString (qualLeaf qfname) == "&&",+            baseName (qualLeaf qfname) == "&&",             [(x, _), (y, _)] <- args ->               internaliseExp desc $                 E.AppExp                   (E.If x y (E.Literal (E.BoolValue False) mempty) mempty)                   (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])           | baseTag (qualLeaf qfname) <= maxIntrinsicTag,-            baseString (qualLeaf qfname) == "||",+            baseName (qualLeaf qfname) == "||",             [(x, _), (y, _)] <- args ->               internaliseExp desc $                 E.AppExp@@ -399,7 +399,7 @@               funcall desc qfname args' internaliseAppExp desc _ (E.LetPat sizes pat e body _) =   internalisePat desc sizes pat e $ internaliseExp desc body-internaliseAppExp _ _ (E.LetFun ofname _ _ _) =+internaliseAppExp _ _ (E.LetFun (ofname, _) _ _ _) =   error $ "Unexpected LetFun " ++ prettyString ofname internaliseAppExp desc _ (E.Loop sparams mergepat loopinit form loopbody _) = do   ses <- internaliseExp "loop_init" $ loopInitExp loopinit@@ -572,7 +572,7 @@       (E.LetPat [] pat e body loc)       (Info (AppRes (E.typeOf body) mempty)) internaliseAppExp desc _ (E.Match e orig_cs _) = do-  ses <- internaliseExp (desc ++ "_scrutinee") e+  ses <- internaliseExp (desc <> "_scrutinee") e   cs <- mapM (onCase ses) orig_cs   case NE.uncons cs of     (I.Case _ body, Nothing) ->@@ -594,7 +594,7 @@ internaliseAppExp _ _ e@E.BinOp {} =   error $ "internaliseAppExp: Unexpected BinOp " ++ prettyString e -internaliseExp :: String -> E.Exp -> InternaliseM [I.SubExp]+internaliseExp :: Name -> E.Exp -> InternaliseM [I.SubExp] internaliseExp desc (E.Parens e _) =   internaliseExp desc e internaliseExp desc (E.Hole (Info t) loc) = do@@ -837,8 +837,11 @@   ts' <- instantiateShapes noExt $ map fromDecl ts    case lookupWithIndex c constr_map of-    Just (i, js) ->-      (intConst Int8 (toInteger i) :) <$> clauses 0 ts' (zip js es')+    Just (i, js)+      | length fs == 1 ->+          clauses 0 ts' (zip js es')+      | otherwise ->+          (intConst Int8 (toInteger i) :) <$> clauses 0 ts' (zip js es')     Nothing ->       error "internaliseExp Constr: missing constructor"   where@@ -902,7 +905,7 @@ internaliseExp _ e@E.IndexSection {} =   error $ "internaliseExp: Unexpected index section at " ++ locStr (srclocOf e) -internaliseArg :: String -> (E.Exp, Maybe VName) -> InternaliseM [SubExp]+internaliseArg :: Name -> (E.Exp, Maybe VName) -> InternaliseM [SubExp] internaliseArg desc (arg, argdim) = do   exists <- askScope   case argdim of@@ -937,7 +940,10 @@   where     compares (E.PatLit l (Info t) _) (se : ses) =       pure ([Just $ internalisePatLit l t], [se], ses)-    compares (E.PatConstr c (Info (E.Scalar (E.Sum fs))) pats _) (_ : ses) = do+    compares (E.PatConstr c (Info (E.Scalar (E.Sum fs))) pats _) sum_ses = do+      -- Handle the unary sum type special case.+      let unary = length fs == 1+          ses = if unary then sum_ses else tail sum_ses       (payload_ts, m) <- internaliseSumType $ M.map (map toStruct) fs       case lookupWithIndex c m of         Just (tag, payload_is) -> do@@ -949,8 +955,8 @@                   Just j -> cmps !! j                   Nothing -> Nothing           pure-            ( Just (I.IntValue $ intValue Int8 $ toInteger tag)-                : zipWith missingCmps [0 ..] payload_ses,+            ( [Just (I.IntValue $ intValue Int8 $ toInteger tag) | not unary]+                <> zipWith missingCmps [0 ..] payload_ses,               pertinent,               ses'             )@@ -991,7 +997,7 @@         )  internalisePat ::-  String ->+  Name ->   [E.SizeBinder VName] ->   E.Pat StructType ->   E.Exp ->@@ -1002,7 +1008,7 @@   internalisePat' sizes p ses m   where     desc' = case E.patIdents p of-      [v] -> baseString $ E.identName v+      [v] -> baseName $ E.identName v       _ -> desc  internalisePat' ::@@ -1179,20 +1185,20 @@     one = constant (1 :: Int64)  internaliseScanOrReduce ::-  String ->-  String ->+  Name ->+  Name ->   (SubExp -> I.Lambda SOACS -> [SubExp] -> [VName] -> InternaliseM (SOAC SOACS)) ->   (E.Exp, E.Exp, E.Exp) ->   InternaliseM [SubExp] internaliseScanOrReduce desc what f (lam, ne, arr) = do-  arrs <- internaliseExpToVars (what ++ "_arr") arr-  nes <- internaliseExp (what ++ "_ne") ne+  arrs <- internaliseExpToVars (what <> "_arr") arr+  nes <- internaliseExp (what <> "_ne") ne   nes' <- forM (zip nes arrs) $ \(ne', arr') -> do     rowtype <- I.stripArray 1 <$> lookupType arr'     ensureShape       "Row shape of input array does not match shape of neutral element"       rowtype-      (what ++ "_ne_right_shape")+      (what <> "_ne_right_shape")       ne'   nests <- mapM I.subExpType nes'   arrts <- mapM lookupType arrs@@ -1202,7 +1208,7 @@  internaliseHist ::   Int ->-  String ->+  Name ->   E.Exp ->   E.Exp ->   E.Exp ->@@ -1251,7 +1257,7 @@     I.Hist w_img (buckets' ++ img') [HistOp shape_hist rf' hist' ne_shp op'] lam'  internaliseStreamAcc ::-  String ->+  Name ->   E.Exp ->   Maybe (E.Exp, E.Exp) ->   E.Exp ->@@ -1292,7 +1298,7 @@     letTupExp desc $       WithAcc [(I.Shape [destw], dest', op')] withacc_lam -internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp+internaliseExp1 :: Name -> E.Exp -> InternaliseM I.SubExp internaliseExp1 desc e = do   vs <- internaliseExp desc e   case vs of@@ -1301,14 +1307,14 @@  -- | Promote to dimension type as appropriate for the original type. -- Also return original type.-internaliseSizeExp :: String -> E.Exp -> InternaliseM (I.SubExp, IntType)+internaliseSizeExp :: Name -> E.Exp -> InternaliseM (I.SubExp, IntType) internaliseSizeExp s e = do   e' <- internaliseExp1 s e   case E.typeOf e of     E.Scalar (E.Prim (E.Signed it)) -> (,it) <$> asIntS Int64 e'     _ -> error "internaliseSizeExp: bad type" -internaliseExpToVars :: String -> E.Exp -> InternaliseM [I.VName]+internaliseExpToVars :: Name -> E.Exp -> InternaliseM [I.VName] internaliseExpToVars desc e =   mapM asIdent =<< internaliseExp desc e   where@@ -1316,7 +1322,7 @@     asIdent se = letExp desc $ I.BasicOp $ I.SubExp se  internaliseOperation ::-  String ->+  Name ->   E.Exp ->   (I.VName -> InternaliseM I.BasicOp) ->   InternaliseM [I.SubExp]@@ -1351,7 +1357,7 @@   certifying c m  internaliseBinOp ::-  String ->+  Name ->   E.BinOp ->   I.SubExp ->   I.SubExp ->@@ -1438,7 +1444,7 @@ internaliseBinOp desc E.Equal x y t _ =   simpleCmpOp desc (I.CmpEq $ internalisePrimType t) x y internaliseBinOp desc E.NotEqual x y t _ = do-  eq <- letSubExp (desc ++ "true") $ I.BasicOp $ I.CmpOp (I.CmpEq $ internalisePrimType t) x y+  eq <- letSubExp (desc <> "true") $ I.BasicOp $ I.CmpOp (I.CmpEq $ internalisePrimType t) x y   fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp (I.Neg I.Bool) eq internaliseBinOp desc E.Less x y (E.Signed t) _ =   simpleCmpOp desc (I.CmpSlt t) x y@@ -1484,7 +1490,7 @@       ++ prettyString t2  simpleBinOp ::-  String ->+  Name ->   I.BinOp ->   I.SubExp ->   I.SubExp ->@@ -1493,7 +1499,7 @@   letTupExp' desc $ I.BasicOp $ I.BinOp bop x y  simpleCmpOp ::-  String ->+  Name ->   I.CmpOp ->   I.SubExp ->   I.SubExp ->@@ -1548,11 +1554,11 @@ -- | Overloaded operators are treated here. isOverloadedFunction ::   E.QualName VName ->-  String ->+  Name ->   Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp]) isOverloadedFunction qname desc = do   guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag-  handle $ baseString $ qualLeaf qname+  handle $ baseName $ qualLeaf qname   where     -- Handle equality and inequality specially, to treat the case of     -- arrays.@@ -1608,7 +1614,10 @@               letSubExp "arrays_equal"                 =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False])     handle name-      | Just bop <- find ((name ==) . prettyString) [minBound .. maxBound :: E.BinOp] =+      | Just bop <-+          find+            ((name ==) . nameFromText . prettyText)+            [minBound .. maxBound :: E.BinOp] =           Just $ \[(x_t, [x']), (y_t, [y'])] ->             case (x_t, y_t) of               (E.Scalar (E.Prim t1), E.Scalar (E.Prim t2)) ->@@ -1616,13 +1625,21 @@               _ -> error "Futhark.Internalise.internaliseExp: non-primitive type in BinOp."     handle _ = Nothing +scatterF :: Int -> E.Exp -> E.Exp -> E.Exp -> Name -> InternaliseM [SubExp]+scatterF rank dest is v desc = do+  is' <- internaliseExpToVars "write_arg_i" is+  v' <- internaliseExpToVars "write_arg_v" v+  dest' <- internaliseExpToVars "write_arg_a" dest+  map I.Var+    <$> doScatter desc rank dest' (is' <> v') (pure . map (I.Var . I.paramName))+ -- | Handle intrinsic functions.  These are only allowed to be called -- in the prelude, and their internalisation may involve inspecting -- the AST. isIntrinsicFunction ::   E.QualName VName ->   [E.Exp] ->-  Maybe (String -> InternaliseM [SubExp])+  Maybe (Name -> InternaliseM [SubExp]) isIntrinsicFunction qname args = do   guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag   let handlers =@@ -1633,7 +1650,7 @@           handleAD,           handleRest         ]-  msum [h args $ baseString $ qualLeaf qname | h <- handlers]+  msum [h args $ baseName $ qualLeaf qname | h <- handlers]   where     handleSign [x] "sign_i8" = Just $ toSigned I.Int8 x     handleSign [x] "sign_i16" = Just $ toSigned I.Int16 x@@ -1646,20 +1663,20 @@     handleSign _ _ = Nothing      handleOps [x] s-      | Just unop <- find ((== s) . prettyString) allUnOps = Just $ \desc -> do+      | Just unop <- find ((== s) . nameFromText . prettyText) allUnOps = Just $ \desc -> do           x' <- internaliseExp1 "x" x           fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x'     handleOps [TupLit [x, y] _] s-      | Just bop <- find ((== s) . prettyString) allBinOps = Just $ \desc -> do+      | Just bop <- find ((== s) . nameFromText . prettyText) allBinOps = Just $ \desc -> do           x' <- internaliseExp1 "x" x           y' <- internaliseExp1 "y" y           fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y'-      | Just cmp <- find ((== s) . prettyString) allCmpOps = Just $ \desc -> do+      | Just cmp <- find ((== s) . nameFromText . prettyText) allCmpOps = Just $ \desc -> do           x' <- internaliseExp1 "x" x           y' <- internaliseExp1 "y" y           fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y'     handleOps [x] s-      | Just conv <- find ((== s) . prettyString) allConvOps = Just $ \desc -> do+      | Just conv <- find ((== s) . nameFromText . prettyText) allConvOps = Just $ \desc -> do           x' <- internaliseExp1 "x" x           fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x'     handleOps _ _ = Nothing@@ -1806,8 +1823,8 @@     handleRest [x, y] "zip" = Just $ \desc ->       mapM (letSubExp "zip_copy" . BasicOp . Replicate mempty . I.Var)         =<< ( (++)-                <$> internaliseExpToVars (desc ++ "_zip_x") x-                <*> internaliseExpToVars (desc ++ "_zip_y") y+                <$> internaliseExpToVars (desc <> "_zip_x") x+                <*> internaliseExpToVars (desc <> "_zip_y") y             )     handleRest [x] "unzip" = Just $ \desc ->       mapM (letSubExp desc . BasicOp . Replicate mempty . I.Var)@@ -1860,67 +1877,7 @@           letTupExp' desc $ I.BasicOp $ I.ConvOp (I.FPToUI float_from int_to) e'         _ -> error "Futhark.Internalise.internaliseExp: non-numeric type in ToUnsigned" -    scatterF dim a si v desc = do-      si' <- internaliseExpToVars "write_arg_i" si-      svs <- internaliseExpToVars "write_arg_v" v-      sas <- internaliseExpToVars "write_arg_a" a--      si_w <- I.arraysSize 0 <$> mapM lookupType si'-      sv_ts <- mapM lookupType svs--      svs' <- forM (zip svs sv_ts) $ \(sv, sv_t) -> do-        let sv_shape = I.arrayShape sv_t-            sv_w = arraySize 0 sv_t--        -- Generate an assertion and reshapes to ensure that sv and si' are the same-        -- size.-        cmp <--          letSubExp "write_cmp" $-            I.BasicOp $-              I.CmpOp (I.CmpEq I.int64) si_w sv_w-        c <--          assert-            "write_cert"-            cmp-            "length of index and value array does not match"-        certifying c $-          letExp (baseString sv ++ "_write_sv") $-            shapeCoerce (I.shapeDims (reshapeOuter (I.Shape [si_w]) 1 sv_shape)) sv--      indexType <- fmap rowType <$> mapM lookupType si'-      indexName <- mapM (\_ -> newVName "write_index") indexType-      valueNames <- replicateM (length sv_ts) $ newVName "write_value"--      sa_ts <- mapM lookupType sas-      let bodyTypes = concat (replicate (length sv_ts) indexType) ++ map (I.stripArray dim) sa_ts-          paramTypes = indexType <> map rowType sv_ts-          bodyNames = indexName <> valueNames-          bodyParams = zipWith (I.Param mempty) bodyNames paramTypes--      -- This body is boring right now, as every input is exactly the output.-      -- But it can get funky later on if fused with something else.-      body <- localScope (scopeOfLParams bodyParams) . buildBody_ $ do-        let outs = concat (replicate (length valueNames) indexName) ++ valueNames-        results <- forM outs $ \name ->-          letSubExp "write_res" $ I.BasicOp $ I.SubExp $ I.Var name-        ensureResultShape-          "scatter value has wrong size"-          bodyTypes-          (subExpsRes results)--      let lam =-            I.Lambda-              { I.lambdaParams = bodyParams,-                I.lambdaReturnType = bodyTypes,-                I.lambdaBody = body-              }-          sivs = si' <> svs'--      let sa_ws = map (I.Shape . take dim . arrayDims) sa_ts-          spec = zip3 sa_ws (repeat 1) sas-      letTupExp' desc $ I.Op $ I.Scatter si_w sivs spec lam--flatIndexHelper :: String -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp]+flatIndexHelper :: Name -> E.Exp -> E.Exp -> [(E.Exp, E.Exp)] -> InternaliseM [SubExp] flatIndexHelper desc arr offset slices = do   arrs <- internaliseExpToVars "arr" arr   offset' <- internaliseExp1 "offset" offset@@ -1968,7 +1925,7 @@     forM arrs $ \arr' ->       letSubExp desc $ I.BasicOp $ I.FlatIndex arr' slice -flatUpdateHelper :: String -> E.Exp -> E.Exp -> [E.Exp] -> E.Exp -> InternaliseM [SubExp]+flatUpdateHelper :: Name -> E.Exp -> E.Exp -> [E.Exp] -> E.Exp -> InternaliseM [SubExp] flatUpdateHelper desc arr1 offset slices arr2 = do   arrs1 <- internaliseExpToVars "arr" arr1   offset' <- internaliseExp1 "offset" offset@@ -2019,7 +1976,7 @@       letSubExp desc $ I.BasicOp $ I.FlatUpdate arr1' slice arr2'  funcall ::-  String ->+  Name ->   QualName VName ->   [SubExp] ->   InternaliseM [SubExp]@@ -2140,32 +2097,17 @@         Scratch (I.elemType arr_t) (w : drop 1 (I.arrayDims arr_t))    -- Now write into the result.-  write_lam <- do-    c_param <- newParam "c" (I.Prim int64)-    offset_params <- replicateM k $ newParam "offset" (I.Prim int64)-    value_params <- mapM (newParam "v" . I.rowType) arr_ts-    (offset, offset_stms) <--      collectStms $+  results <-+    doScatter "partition_res" 1 blanks (classes : all_offsets ++ arrs) $ \params -> do+      let ([c_param], offset_params, value_params) =+            splitAt3 1 (length all_offsets) params+      offset <-         mkOffsetLambdaBody           (map I.Var sizes)           (I.Var $ I.paramName c_param)           0           offset_params-    pure-      I.Lambda-        { I.lambdaParams = c_param : offset_params ++ value_params,-          I.lambdaReturnType =-            replicate (length arr_ts) (I.Prim int64)-              ++ map I.rowType arr_ts,-          I.lambdaBody =-            mkBody offset_stms $-              replicate (length arr_ts) (subExpRes offset)-                ++ I.varsRes (map I.paramName value_params)-        }-  let spec = zip3 (repeat $ I.Shape [w]) (repeat 1) blanks-  results <--    letTupExp "partition_res" . I.Op $-      I.Scatter w (classes : all_offsets ++ arrs) spec write_lam+      pure $ offset : map (I.Var . I.paramName) value_params   sizes' <-     letSubExp "partition_sizes" $       I.BasicOp $

@@ -27,6 +27,7 @@ import Data.Map qualified as M import Data.Text qualified as T import Futhark.MonadFreshNames+import Futhark.Util (showText) import Language.Futhark import Language.Futhark.Traversals import Language.Futhark.TypeChecker.Types@@ -62,9 +63,9 @@ --   they have to be in the same list to conserve their order -- Direct interaction with the inside state should be done with caution, that's why their -- no instance of `MonadState`.-newtype OrderingM a = OrderingM (StateT NormState (Reader String) a)+newtype OrderingM a = OrderingM (StateT NormState (Reader Name) a)   deriving-    (Functor, Applicative, Monad, MonadReader String, MonadState NormState)+    (Functor, Applicative, Monad, MonadReader Name, MonadState NormState)  instance MonadFreshNames OrderingM where   getNameSource = OrderingM $ gets snd@@ -97,7 +98,7 @@         then ((a, binds), src)         else error "not all bind modifiers were freed" -naming :: String -> OrderingM a -> OrderingM a+naming :: Name -> OrderingM a -> OrderingM a naming s = local (const s)  -- | From now, we say an expression is "final" if it's going to be stored in a let-bind@@ -107,7 +108,7 @@ nameExp :: Bool -> Exp -> OrderingM Exp nameExp True e = pure e nameExp False e = do-  name <- newNameFromString =<< ask -- "e<{" ++ prettyString e ++ "}>"+  name <- newVName =<< ask -- "e<{" ++ prettyString e ++ "}>"   let ty = typeOf e       loc = srclocOf e       pat = Id name (Info ty) loc@@ -116,18 +117,18 @@  -- An evocative name to use when naming subexpressions of the -- expression bound to this pattern.-patRepName :: Pat t -> String+patRepName :: Pat t -> Name patRepName (PatAscription p _ _) = patRepName p-patRepName (Id v _ _) = baseString v+patRepName (Id v _ _) = baseName v patRepName _ = "tmp" -expRepName :: Exp -> String-expRepName (Var v _ _) = prettyString v-expRepName e = "d<{" ++ prettyString (bareExp e) ++ "}>"+expRepName :: Exp -> Name+expRepName (Var v _ _) = nameFromText $ prettyText v+expRepName e = "d<{" <> nameFromText (prettyText (bareExp e)) <> "}>"  -- An evocative name to use when naming arguments to an application.-argRepName :: Exp -> Int -> String-argRepName e i = expRepName e <> "_arg" <> show i+argRepName :: Exp -> Int -> Name+argRepName e i = expRepName e <> "_arg" <> nameFromText (showText i)  -- Modify an expression as describe in module introduction, -- introducing the let-bindings in the state.@@ -214,7 +215,7 @@   pure $ Var qn ty loc getOrdering final (OpSectionLeft op ty e (Info (xp, _, xext), Info (yp, yty)) (Info (RetType dims ret), Info exts) loc) = do   x <- getOrdering False e-  yn <- newNameFromString "y"+  yn <- newVName "y"   let y = Var (qualName yn) (Info $ toStruct yty) mempty       ret' = applySubst (pSubst x y) ret       body =@@ -227,7 +228,7 @@       | Named p <- yp, p == vn = Just $ ExpSubst y       | otherwise = Nothing getOrdering final (OpSectionRight op ty e (Info (xp, xty), Info (yp, _, yext)) (Info (RetType dims ret)) loc) = do-  xn <- newNameFromString "x"+  xn <- newVName "x"   y <- getOrdering False e   let x = Var (qualName xn) (Info $ toStruct xty) mempty       ret' = applySubst (pSubst x y) ret@@ -239,7 +240,7 @@       | Named p <- yp, p == vn = Just $ ExpSubst y       | otherwise = Nothing getOrdering final (ProjectSection names (Info ty) loc) = do-  xn <- newNameFromString "x"+  xn <- newVName "x"   let (xty, RetType dims ret) = case ty of         Scalar (Arrow _ _ d xty' ret') -> (toParam d xty', ret')         _ -> error $ "not a function type for project section: " ++ prettyString ty@@ -260,7 +261,7 @@               ++ prettyString field getOrdering final (IndexSection slice (Info ty) loc) = do   slice' <- astMap mapper slice-  xn <- newNameFromString "x"+  xn <- newVName "x"   let (xty, RetType dims ret) = case ty of         Scalar (Arrow _ _ d xty' ret') -> (toParam d xty', ret')         _ -> error $ "not a function type for index section: " ++ prettyString ty@@ -290,7 +291,7 @@   expr' <- naming (patRepName pat) $ getOrdering True expr   addBind $ PatBind sizes pat expr'   getOrdering final body-getOrdering final (AppExp (LetFun vn (tparams, params, mrettype, rettype, body) e _) _) = do+getOrdering final (AppExp (LetFun (vn, _) (tparams, params, mrettype, rettype, body) e _) _) = do   body' <- transformBody body   addBind $ FunBind vn (tparams, params, mrettype, rettype, body')   getOrdering final e@@ -318,8 +319,8 @@       er' <- naming "and_rhs" $ transformBody er       pure $ AppExp (If el' er' (Literal (BoolValue False) mempty) loc) (Info resT)     (False, False) -> do-      el' <- naming (prettyString op <> "_lhs") $ getOrdering False el-      er' <- naming (prettyString op <> "_rhs") $ getOrdering False er+      el' <- naming (nameFromText (prettyText op) <> "_lhs") $ getOrdering False el+      er' <- naming (nameFromText (prettyText op) <> "_rhs") $ getOrdering False er       pure $ mkApply (Var op opT oloc) [(elp, el'), (erp, er')] resT   nameExp final expr'   where@@ -363,7 +364,7 @@     f body (PatBind sizes p expr) =       AppExp (LetPat sizes p expr body mempty) appRes     f body (FunBind vn infos) =-      AppExp (LetFun vn infos body mempty) appRes+      AppExp (LetFun (vn, mempty) infos body mempty) appRes  transformValBind :: (MonadFreshNames m) => ValBind -> m ValBind transformValBind valbind = do

@@ -152,9 +152,9 @@ transformPat = traverse transformType  transformExp :: Exp -> LiftM Exp-transformExp (AppExp (LetFun fname (tparams, params, _, Info ret, funbody) body _) _) = do+transformExp (AppExp (LetFun (fname, _) (tparams, params, _, Info ret, funbody) body _) _) = do   funbody' <- bindingParams (map typeParamName tparams) params $ transformExp funbody-  fname' <- newVName $ "lifted_" ++ baseString fname+  fname' <- newVName $ "lifted_" <> baseName fname   lifted_call <- liftFunction fname' tparams params ret funbody'   replacing fname lifted_call $ transformExp body transformExp e@(Lambda params body _ (Info ret) _) = do

@@ -191,7 +191,7 @@ -- | Construct an 'Assert' statement, but taking attributes into account. Always -- use this function, and never construct 'Assert' directly in the internaliser! assert ::-  String ->+  Name ->   SubExp ->   ErrorMsg SubExp ->   InternaliseM Certs

@@ -343,8 +343,8 @@ lookupLifted :: VName -> MonoType -> MonoM (Maybe (VName, InferSizeArgs)) lookupLifted fname t = lookup (fname, t) <$> getLifts -sizeVarName :: Exp -> String-sizeVarName e = "d<{" <> prettyString (bareExp e) <> "}>"+sizeVarName :: Exp -> Name+sizeVarName e = "d<{" <> nameFromText (prettyText (bareExp e)) <> "}>"  -- | Creates a new expression replacement if needed, this always produces normalised sizes. -- (e.g. single variable or constant)@@ -360,7 +360,7 @@         (Just vn, _) -> pure $ sizeFromName (qualName vn) (srclocOf e)         (Nothing, Just vn) -> pure $ sizeFromName (qualName vn) (srclocOf e)         (Nothing, Nothing) -> do-          vn <- newNameFromString $ sizeVarName e+          vn <- newVName $ sizeVarName e           modify ((e', vn) :)           pure $ sizeFromName (qualName vn) (srclocOf e)   where@@ -571,7 +571,7 @@      makeVarParam arg = do       let argtype = typeOf arg-      x <- newNameFromString "binop_p"+      x <- newVName "binop_p"       pure         ( Var (qualName x) (Info argtype) mempty,           Id x (Info argtype) mempty@@ -590,7 +590,7 @@   if S.null implicitDims     then pure $ AppExp (Match e' cs' loc) (Info res')     else do-      tmpVar <- newNameFromString "matched_variable"+      tmpVar <- newVName "matched_variable"       pure $         AppExp           ( LetPat@@ -752,7 +752,7 @@     Lambda (p1 ++ p2) body Nothing (Info rettype'') loc   where     patAndVar argtype = do-      x <- newNameFromString "x"+      x <- newVName "x"       pure         ( x,           Id x (Info argtype) mempty,

@@ -100,7 +100,7 @@   let fs' = M.toList fs   (fs_ks, fs_ts) <- fmap unzip $     forM fs' $ \(f, ft) ->-      (,) <$> newVName (nameToString f) <*> pure ft+      (,) <$> newVName f <*> pure ft   pure     ( RecordPat         (zip (map (L noLoc . fst) fs') (zipWith3 Id fs_ks (map Info fs_ts) $ repeat loc))

@@ -257,7 +257,7 @@           concat <$> mapM (internaliseTypeM exts . snd) (E.sortFields ets)     E.Scalar (E.TypeVar u tn [E.TypeArgType arr_t])       | baseTag (E.qualLeaf tn) <= E.maxIntrinsicTag,-        baseString (E.qualLeaf tn) == "acc" -> do+        baseName (E.qualLeaf tn) == "acc" -> do           ts <-             foldMap (toList . fmap (fromDecl . onAccType))               <$> internaliseTypeM exts (E.toRes Nonunique arr_t)@@ -274,7 +274,10 @@       (ts, _) <-         internaliseConstructors           <$> traverse (fmap concat . mapM (internaliseTypeM exts)) cs-      pure $ Pure (I.Prim (I.IntType I.Int8)) : ts+      pure $+        if length cs == 1+          then ts+          else Pure (I.Prim (I.IntType I.Int8)) : ts   where     internaliseShape = mapM (internaliseDim exts) . E.shapeDims     array [Free ts] = Free ts

@@ -10,7 +10,6 @@   ( MonadFreshNames (..),     modifyNameSource,     newName,-    newNameFromString,     newVName,     newIdent,     newIdent',@@ -80,49 +79,34 @@ newName :: (MonadFreshNames m) => VName -> m VName newName = modifyNameSource . flip FreshNames.newName --- | As @newName@, but takes a 'String' for the name template.-newNameFromString :: (MonadFreshNames m) => String -> m VName-newNameFromString s = newName $ VName (nameFromString s) 0- -- | Produce a fresh 'VName', using the given base name as a template.-newID :: (MonadFreshNames m) => Name -> m VName-newID s = newName $ VName s 0---- | Produce a fresh 'VName', using the given base name as a template.-newVName :: (MonadFreshNames m) => String -> m VName-newVName = newID . nameFromString+newVName :: (MonadFreshNames m) => Name -> m VName+newVName s = newName $ VName s 0  -- | Produce a fresh 'Ident', using the given name as a template. newIdent ::-  (MonadFreshNames m) =>-  String ->-  Type ->-  m Ident-newIdent s t = do-  s' <- newID $ nameFromString s-  pure $ Ident s' t+  (MonadFreshNames m) => Name -> Type -> m Ident+newIdent s t = Ident <$> newVName s <*> pure t  -- | Produce a fresh 'Ident', using the given 'Ident' as a template, -- but possibly modifying the name. newIdent' ::   (MonadFreshNames m) =>-  (String -> String) ->+  (Name -> Name) ->   Ident ->   m Ident newIdent' f ident =   newIdent-    (f $ nameToString $ baseName $ identName ident)+    (f $ baseName $ identName ident)     (identType ident)  -- | Produce a fresh 'Param', using the given name as a template. newParam ::   (MonadFreshNames m) =>-  String ->+  Name ->   dec ->   m (Param dec)-newParam s t = do-  s' <- newID $ nameFromString s-  pure $ Param mempty s' t+newParam s t = Param mempty <$> newVName s <*> pure t  -- Utility instance defintions for MTL classes.  This requires -- UndecidableInstances, but saves on typing elsewhere.

@@ -124,8 +124,8 @@   ExpMap rep ->   KernelBody rep ->   TransformM rep (KernelBody rep)-transformSegGroupKernelBody perm_table expmap (KernelBody b stms res) =-  KernelBody b <$> transformStms perm_table expmap stms <*> pure res+transformSegGroupKernelBody perm_table expmap (Body b stms res) =+  Body b <$> transformStms perm_table expmap stms <*> pure res  -- | Transform the statements in the body of a SegThread kernel. transformSegThreadKernelBody ::@@ -152,15 +152,13 @@   VName ->   KernelBody rep ->   TransformM rep (KernelBody rep)-transformKernelBody perm_table expmap p seg_name (KernelBody b stms res) = do+transformKernelBody perm_table expmap p seg_name (Body b stms res) = do   stms' <- transformStms perm_table expmap stms-  evalStateT-    ( traverseKernelBodyArrayIndexes-        seg_name-        (ensureTransformedAccess perm_table p)-        (KernelBody b stms' res)-    )-    mempty+  flip evalStateT mempty $+    traverseKernelBodyArrayIndexes+      seg_name+      (ensureTransformedAccess perm_table p)+      (Body b stms' res)  traverseKernelBodyArrayIndexes ::   forall m rep.@@ -169,10 +167,8 @@   ArrayIndexTransform m ->   KernelBody rep ->   m (KernelBody rep)-traverseKernelBodyArrayIndexes seg_name coalesce (KernelBody b kstms kres) =-  KernelBody b . stmsFromList-    <$> mapM onStm (stmsToList kstms)-    <*> pure kres+traverseKernelBodyArrayIndexes seg_name coalesce (Body b kstms kres) =+  Body b <$> traverse onStm kstms <*> pure kres   where     onLambda lam =       (\body' -> lam {lambdaBody = body'})@@ -244,7 +240,7 @@       pure $ Just (arr', slice)      manifest perm array =-      auxing aux . letExp (baseString array ++ "_coalesced") $+      auxing aux . letExp (baseName array <> "_coalesced") $         BasicOp (Manifest array perm)  lookupPermutation :: LayoutTable -> VName -> IndexExprName -> VName -> Maybe Permutation

@@ -164,17 +164,17 @@   SegOp lvl rep ->   UpdateM (inner rep) (SegOp lvl rep) replaceInSegOp (SegMap lvl sp tps body) = do-  stms <- updateStms $ kernelBodyStms body-  pure $ SegMap lvl sp tps $ body {kernelBodyStms = stms}+  stms <- updateStms $ bodyStms body+  pure $ SegMap lvl sp tps $ body {bodyStms = stms} replaceInSegOp (SegRed lvl sp tps body binops) = do-  stms <- updateStms $ kernelBodyStms body-  pure $ SegRed lvl sp tps (body {kernelBodyStms = stms}) binops+  stms <- updateStms $ bodyStms body+  pure $ SegRed lvl sp tps (body {bodyStms = stms}) binops replaceInSegOp (SegScan lvl sp tps body binops) = do-  stms <- updateStms $ kernelBodyStms body-  pure $ SegScan lvl sp tps (body {kernelBodyStms = stms}) binops+  stms <- updateStms $ bodyStms body+  pure $ SegScan lvl sp tps (body {bodyStms = stms}) binops replaceInSegOp (SegHist lvl sp tps body hist_ops) = do-  stms <- updateStms $ kernelBodyStms body-  pure $ SegHist lvl sp tps (body {kernelBodyStms = stms}) hist_ops+  stms <- updateStms $ bodyStms body+  pure $ SegHist lvl sp tps (body {bodyStms = stms}) hist_ops  replaceInHostOp :: HostOp NoOp GPUMem -> UpdateM (HostOp NoOp GPUMem) (HostOp NoOp GPUMem) replaceInHostOp (SegOp op) = SegOp <$> replaceInSegOp op

@@ -268,8 +268,8 @@ shortCircuitGPUMem lutab pat certs (Inner (GPU.SegOp op)) td_env bu_env =   shortCircuitSegOp isSegThread lutab pat certs op td_env bu_env shortCircuitGPUMem lutab pat certs (Inner (GPU.GPUBody _ body)) td_env bu_env = do-  fresh1 <- newNameFromString "gpubody"-  fresh2 <- newNameFromString "gpubody"+  fresh1 <- newVName "gpubody"+  fresh2 <- newVName "gpubody"   shortCircuitSegOpHelper     0     isSegThread@@ -336,7 +336,7 @@  bodyToKernelBody :: Body (Aliases GPUMem) -> KernelBody (Aliases GPUMem) bodyToKernelBody (Body dec stms res) =-  KernelBody dec stms $ map (\(SubExpRes cert subexps) -> Returns ResultNoSimplify cert subexps) res+  Body dec stms $ map (\(SubExpRes cert subexps) -> Returns ResultNoSimplify cert subexps) res  -- | A helper for all the different kinds of 'SegOp'. --@@ -371,7 +371,7 @@   -- that correspond to reductions.   let ps_space_and_res =         zip3 ps0 (replicate num_reds (dropLastSegSpace space0) <> repeat space0) $-          kernelBodyResult kernel_body+          bodyResult kernel_body   -- Create coalescing relations between pattern elements and kernel body   -- results   let (actv0, inhibit0) =@@ -390,7 +390,7 @@   let actv0' = M.map (\etry -> etry {memrefs = mempty}) $ actv0 <> actv_return   -- Process kernel body statements   bu_env' <--    mkCoalsTabStms lutab (kernelBodyStms kernel_body) td_env $+    mkCoalsTabStms lutab (bodyStms kernel_body) td_env $       bu_env {activeCoals = actv0', inhibit = inhibit_return}    let actv_coals_after =@@ -522,11 +522,11 @@   (CoalsTab, InhibitTab) makeSegMapCoals lvlOK lvl td_env kernel_body pat_certs (active, inhb) (PatElem pat_name (_, MemArray _ _ _ (ArrayIn pat_mem pat_ixf)), space, Returns _ _ (Var return_name))   | Just (MemBlock tp return_shp return_mem _) <--      getScopeMemInfo return_name $ scope td_env <> scopeOf (kernelBodyStms kernel_body),+      getScopeMemInfo return_name $ scope td_env <> scopeOf (bodyStms kernel_body),     lvlOK lvl,     MemMem pat_space <- runReader (lookupMemInfo pat_mem) $ removeScopeAliases $ scope td_env,     MemMem return_space <--      scope td_env <> scopeOf (kernelBodyStms kernel_body) <> scopeOfSegSpace space+      scope td_env <> scopeOf (bodyStms kernel_body) <> scopeOfSegSpace space         & removeScopeAliases         & runReader (lookupMemInfo return_mem),     pat_space == return_space =@@ -586,10 +586,6 @@         & map (DimFix . TPrimExp . flip LeafExp (IntType Int64) . fst)         & Slice     resultSlice ixf = LMAD.slice ixf $ fullSlice (LMAD.shape ixf) thread_slice-makeSegMapCoals _ _ td_env _ _ x (_, _, WriteReturns _ return_name _) =-  case getScopeMemInfo return_name $ scope td_env of-    Just (MemBlock _ _ return_mem _) -> markFailedCoal x return_mem-    Nothing -> error "Should not happen?" makeSegMapCoals _ _ td_env _ _ x (_, _, result) =   freeIn result     & namesToList@@ -1467,11 +1463,9 @@   scopetab   (Pat [PatElem dst (_, MemArray dst_pt _ _ (ArrayIn dst_mem dst_ixf))])   certs-  (SegMap _ space _ kernel_body@KernelBody {kernelBodyResult = [Returns {}]})+  (SegMap _ space _ kernel_body@Body {bodyResult = [Returns {}]})     | (src, MemBlock src_pt shp src_mem src_ixf) : _ <--        mapMaybe getPotentialMapShortCircuit $-          stmsToList $-            kernelBodyStms kernel_body =+        mapMaybe getPotentialMapShortCircuit $ stmsToList $ bodyStms kernel_body =         Just [(MapCoal, id, dst, dst_mem, dst_ixf, src, src_mem, src_ixf, src_pt, shp, certs)]     where       iterators = map fst $ unSegSpace space@@ -1744,10 +1738,10 @@   concatMapM     ( computeScalarTable $         scope_table-          <> scopeOf (kernelBodyStms $ segBody segop)+          <> scopeOf (bodyStms $ segBody segop)           <> scopeOfSegSpace (segSpace segop)     )-    (stmsToList $ kernelBodyStms $ segBody segop)+    (stmsToList $ bodyStms $ segBody segop)  computeScalarTableGPUMem ::   ComputeScalarTable GPUMem (GPU.HostOp NoOp (Aliases GPUMem))

@@ -59,7 +59,7 @@ isInnerCoal _ _ _ =   error "kkLoopBody.isInnerCoal: not an error, but I would like to know why!" -scratch :: (MonadBuilder m) => String -> PrimType -> [SubExp] -> m VName+scratch :: (MonadBuilder m) => Name -> PrimType -> [SubExp] -> m VName scratch se_name t shape = letExp se_name $ BasicOp $ Scratch t shape  -- | Main helper function for Register-and-Block Tiling@@ -190,15 +190,15 @@       copyGlb2ShMem is_B kk (gtid, ii, ptp_X_el, parlen_X, inp_X, load_X, x_loc_init') = do         let (t_par, r_par, tseq_div_tpar) = (tx, rx, tk_div_tx)             is_inner_coal = isInnerCoal env inp_X load_X-            str_A = baseString inp_X+            str_A = baseName inp_X         x_loc <--          segScatter2D (str_A ++ "_glb2loc") x_loc_init' [r_par, tseq_div_tpar] (t_par, t_par) $+          segScatter2D (str_A <> "_glb2loc") x_loc_init' [r_par, tseq_div_tpar] (t_par, t_par) $             scatterFun is_inner_coal         pure (x_loc, indexLocMem is_inner_coal str_A x_loc)         where           indexLocMem ::             Bool ->-            String ->+            Name ->             VName ->             VName ->             VName ->@@ -208,13 +208,13 @@             let (r_par, t_seq, tr_par) = (rx, tk, tx_rx)             let pad_term = if is_B then pe64 se1 else pe64 se0             x_loc_ind_32 <--              letExp (str_A ++ "_loc_ind_64")+              letExp (str_A <> "_loc_ind_64")                 =<< toExp                   ( if is_inner_coal -- ToDo: check this is correct + turn to i32                       then le64 k + (le64 ltid_yx * pe64 r_par + le64 ij) * (pe64 t_seq + pad_term)                       else le64 ij + le64 ltid_yx * pe64 r_par + le64 k * pe64 tr_par                   )-            index (str_A ++ "_loc_elem") x_loc [x_loc_ind_32]+            index (str_A <> "_loc_elem") x_loc [x_loc_ind_32]           --           scatterFun ::             Bool ->@@ -222,14 +222,14 @@             (VName, VName) ->             Builder GPU (SubExp, SubExp)           scatterFun is_inner_coal [i0, k0] (thd_y, thd_x) = do-            let str_A = baseString inp_X+            let str_A = baseName inp_X                 t_seq = tk             (i, k, epx_loc_fi) <- mk_ik is_B is_inner_coal (thd_y, thd_x) (i0, k0)             letBindNames [gtid] =<< toExp (le64 ii + le64 i)-            a_seqdim_idx <- letExp (str_A ++ "_seqdim_idx") =<< toExp (le64 kk + le64 k)+            a_seqdim_idx <- letExp (str_A <> "_seqdim_idx") =<< toExp (le64 kk + le64 k)              a_elem <--              letSubExp (str_A ++ "_elem")+              letSubExp (str_A <> "_elem")                 =<< eIf                   ( toExp $                       le64 gtid@@ -246,7 +246,7 @@                   (eBody [eBlank $ Prim ptp_X_el])              a_loc_ind <--              letSubExp (str_A ++ "_loc_ind")+              letSubExp (str_A <> "_loc_ind")                 =<< eIf                   (toExp $ le64 k .<. pe64 t_seq)                   (eBody [toExp epx_loc_fi])@@ -267,7 +267,7 @@  mmBlkRegTilingAcc :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) mmBlkRegTilingAcc env (Let pat aux (Op (SegOp (SegMap SegThread {} seg_space ts old_kbody))))-  | KernelBody () kstms [Returns ResultMaySimplify cs (Var res_nm)] <- old_kbody,+  | Body () kstms [Returns ResultMaySimplify cs (Var res_nm)] <- old_kbody,     cs == mempty,     -- check kernel has one result of primitive type     [res_tp] <- ts,@@ -389,7 +389,7 @@         let grid = KernelGrid (Count grid_size) (Count tblock_size)             level' = SegBlock SegNoVirt (Just grid)             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_t, gridDim_t), (gid_y, gridDim_y), (gid_x, gridDim_x)])-            kbody' = KernelBody () stms_seggroup ret_seggroup+            kbody' = Body () stms_seggroup ret_seggroup         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'       pure $ Just (host_stms, new_kernel)   where@@ -462,7 +462,7 @@  mmBlkRegTilingNrm :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) mmBlkRegTilingNrm env (Let pat aux (Op (SegOp (SegMap SegThread {} seg_space ts old_kbody))))-  | KernelBody () kstms [Returns ResultMaySimplify cs (Var res_nm)] <- old_kbody,+  | Body () kstms [Returns ResultMaySimplify cs (Var res_nm)] <- old_kbody,     cs == mempty,     -- check kernel has one result of primitive type     [res_tp] <- ts,@@ -562,7 +562,7 @@         let grid = KernelGrid (Count grid_size) (Count tblock_size)             level' = SegBlock SegNoVirt (Just grid)             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_y, gridDim_y), (gid_x, gridDim_x)])-            kbody' = KernelBody () stms_seggroup ret_seggroup+            kbody' = Body () stms_seggroup ret_seggroup         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'       pure $ Just (host_stms, new_kernel)   where@@ -718,8 +718,8 @@       SubExp     ) mkTileMemSizes height_A _width_B common_dim is_B_not_transp = do-  tk_name <- nameFromString . prettyString <$> newVName "Tk"-  ty_name <- nameFromString . prettyString <$> newVName "Ty"+  tk_name <- nameFromText . prettyText <$> newVName "Tk"+  ty_name <- nameFromText . prettyText <$> newVName "Ty"   ry_name <- nameFromString . prettyString <$> newVName "Ry"    -- until we change the copying to use lmads we need to@@ -892,7 +892,7 @@       Just (ss Seq.|> stm', tab)   | otherwise = Nothing -getParTiles :: (String, String) -> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp)+getParTiles :: (Name, Name) -> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp) getParTiles (t_str, r_str) (t_name, r_name) len_dim =   case len_dim of     Constant (IntValue (Int64Value 8)) ->@@ -906,7 +906,7 @@       r <- letSubExp r_str $ Op $ SizeOp $ GetSize r_name SizeRegTile       pure (t, r) -getSeqTile :: String -> Name -> SubExp -> SubExp -> SubExp -> Builder GPU SubExp+getSeqTile :: Name -> Name -> SubExp -> SubExp -> SubExp -> Builder GPU SubExp getSeqTile tk_str tk_name len_dim tx ty =   case (tx, ty) of     (Constant (IntValue (Int64Value v_x)), Constant (IntValue (Int64Value v_y))) ->@@ -992,7 +992,7 @@ -- mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread{} seg_space ts old_kbody)))) doRegTiling3D :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) doRegTiling3D (Let pat aux (Op (SegOp old_kernel)))-  | SegMap SegThread {} space kertp (KernelBody () kstms kres) <- old_kernel,+  | SegMap SegThread {} space kertp (Body () kstms kres) <- old_kernel,     -- build the variance table, that records, for     -- each variable name, the variables it depends on     initial_variance <- M.map mempty $ scopeOfSegSpace space,@@ -1100,7 +1100,7 @@           -- scratch the shared-memory arrays corresponding to the arrays that are           --   input to the redomap and are invariant to the outermost parallel dimension.           loc_arr_nms <- forM (M.toList tab_out) $ \(nm, (ptp, _)) ->-            scratch (baseString nm ++ "_loc") ptp [rz]+            scratch (baseName nm <> "_loc") ptp [rz]            prologue_res_list <-             forLoop' common_dim (reg_arr_nms ++ loc_arr_nms) $@@ -1112,44 +1112,40 @@                 loc_arr_nms' <-                   forLoop' count_shmem loc_arr_merge_nms $ \tt loc_arr_merge2_nms -> do                     loc_arr_merge2_nms' <--                      forM (zip loc_arr_merge2_nms (M.toList tab_out)) $ \(loc_Y_nm, (glb_Y_nm, (ptp_Y, load_Y))) -> do-                        ltid_flat <- newVName "ltid_flat"-                        ltid <- newVName "ltid"-                        let segspace = SegSpace ltid_flat [(ltid, tblock_size)]-                        ((res_v, res_i), stms) <- runBuilder $ do-                          offs <- letExp "offs" =<< toExp (pe64 tblock_size * le64 tt)-                          loc_ind <- letExp "loc_ind" =<< toExp (le64 ltid + le64 offs)-                          letBindNames [gtid_z] =<< toExp (le64 ii + le64 loc_ind)-                          let glb_ind = gtid_z-                          y_elm <--                            letSubExp "y_elem"-                              =<< eIf-                                (toExp $ le64 glb_ind .<. pe64 d_M)-                                ( do-                                    addStm load_Y-                                    res <- index "Y_elem" glb_Y_nm [q]-                                    resultBodyM [Var res]-                                )-                                (eBody [eBlank $ Prim ptp_Y])-                          y_ind <--                            letSubExp "y_loc_ind"-                              =<< eIf-                                (toExp $ le64 loc_ind .<. pe64 rz)-                                (toExp loc_ind >>= letTupExp' "loc_fi" >>= resultBodyM)-                                (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])-                          -- y_tp  <- subExpType y_elm-                          pure (y_elm, y_ind)+                      forM (zip loc_arr_merge2_nms (M.toList tab_out)) $ \(loc_Y_nm, (glb_Y_nm, (ptp_Y, load_Y))) ->+                        letExp "Y_glb2loc" <=< withAcc [loc_Y_nm] 1 $ \ ~[acc] -> do+                          ltid_flat <- newVName "ltid_flat"+                          ltid <- newVName "ltid"+                          let segspace = SegSpace ltid_flat [(ltid, tblock_size)]+                          body <- runBodyBuilder $ do+                            offs <- letExp "offs" =<< toExp (pe64 tblock_size * le64 tt)+                            loc_ind <- letExp "loc_ind" =<< toExp (le64 ltid + le64 offs)+                            letBindNames [gtid_z] =<< toExp (le64 ii + le64 loc_ind)+                            let glb_ind = gtid_z+                            y_elm <-+                              letSubExp "y_elem"+                                =<< eIf+                                  (toExp $ le64 glb_ind .<. pe64 d_M)+                                  ( do+                                      addStm load_Y+                                      res <- index "Y_elem" glb_Y_nm [q]+                                      resultBodyM [Var res]+                                  )+                                  (eBody [eBlank $ Prim ptp_Y])+                            y_ind <-+                              letSubExp "y_loc_ind"+                                =<< eIf+                                  (toExp $ le64 loc_ind .<. pe64 rz)+                                  (toExp loc_ind >>= letTupExp' "loc_fi" >>= resultBodyM)+                                  (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])+                            acc' <- letExp (baseName acc) $ BasicOp $ UpdateAcc Safe acc [y_ind] [y_elm]+                            pure [Returns ResultMaySimplify mempty $ Var acc'] -                        let ret = WriteReturns mempty loc_Y_nm [(Slice [DimFix res_i], res_v)]-                        let body = KernelBody () stms [ret]-                        loc_Y_nm_t <- lookupType loc_Y_nm+                          acc_t <- lookupType acc -                        res_nms <--                          letTupExp "Y_glb2loc" <=< renameExp $+                          letTupExp' "Y_glb2loc" <=< renameExp $                             Op . SegOp $-                              SegMap segthd_lvl segspace [loc_Y_nm_t] body-                        let res_nm : _ = res_nms-                        pure res_nm+                              SegMap segthd_lvl segspace [acc_t] body                     resultBodyM $ map Var loc_arr_merge2_nms'                  redomap_res <-@@ -1166,7 +1162,7 @@                               inp_scals_invar_outer <-                                 forM (M.toList tab_inn) $ \(inp_arr_nm, load_stm) -> do                                   addStm load_stm-                                  index (baseString inp_arr_nm) inp_arr_nm [q]+                                  index (baseName inp_arr_nm) inp_arr_nm [q]                               -- build the loop of count R whose body is semantically the redomap code                               reg_arr_merge_nms' <-                                 forLoop' rz reg_arr_merge_nms_slc $ \i reg_arr_mm_nms -> do@@ -1195,7 +1191,7 @@                                           map_res_scals <- eLambda map_lam' (map (eSubExp . Var) map_inp_scals)                                           red_res <- eLambda red_lam' (map eSubExp (map Var cs ++ map resSubExp map_res_scals))                                           css <- forM (zip reg_arr_mm_nms red_res) $ \(reg_arr_nm, c) ->-                                            update (baseString reg_arr_nm) reg_arr_nm [i] (resSubExp c)+                                            update (baseName reg_arr_nm) reg_arr_nm [i] (resSubExp c)                                           resultBodyM $ map Var css                                       )                                       (resultBodyM $ map Var reg_arr_mm_nms)@@ -1280,7 +1276,7 @@         let grid = KernelGrid (Count grid_size) (Count tblock_size)             level' = SegBlock SegNoVirt (Just grid)             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_z, gridDim_z), (gid_y, gridDim_y), (gid_x, gridDim_x)])-            kbody' = KernelBody () stms_seggroup ret_seggroup+            kbody' = Body () stms_seggroup ret_seggroup          pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' kertp kbody'       -- END (new_kernel, host_stms) <- runBuilder $ do@@ -1289,7 +1285,7 @@     getResNm (Returns ResultMaySimplify _ (Var res_nm)) = Just res_nm     getResNm _ = Nothing -    limitTile :: String -> SubExp -> SubExp -> Builder GPU SubExp+    limitTile :: Name -> SubExp -> SubExp -> Builder GPU SubExp     limitTile t_str t d_K = letSubExp t_str $ BasicOp $ BinOp (SMin Int64) t d_K     insertTranspose ::       VarianceTable ->@@ -1306,7 +1302,7 @@             i : _ -> do               arr_tp <- lookupType arr_nm               let perm = [i + 1 .. arrayRank arr_tp - 1] ++ [0 .. i]-              let arr_tr_str = baseString arr_nm ++ "_transp"+              let arr_tr_str = baseName arr_nm <> "_transp"               arr_tr_nm <- letExp arr_tr_str $ BasicOp $ Manifest arr_nm perm               let e_ind' = BasicOp $ Index arr_tr_nm slc               let stm' = Let patt yy e_ind'

@@ -318,9 +318,9 @@   (Aliased rep, CSEInOp (Op rep)) =>   GPU.KernelBody rep ->   CSEM rep (GPU.KernelBody rep)-cseInKernelBody (GPU.KernelBody bodydec stms res) = do+cseInKernelBody (GPU.Body bodydec stms res) = do   Body _ stms' _ <- cseInBody (map (const Observe) res) $ Body bodydec stms []-  pure $ GPU.KernelBody bodydec stms' res+  pure $ GPU.Body bodydec stms' res  instance (CSEInOp (op rep)) => CSEInOp (Memory.MemOp op rep) where   cseInOp o@Memory.Alloc {} = pure o

@@ -180,8 +180,8 @@   KernelBody rep ->   DoubleBufferM rep (KernelBody rep) optimiseKernelBody kbody = do-  stms' <- optimiseStms $ stmsToList $ kernelBodyStms kbody-  pure $ kbody {kernelBodyStms = stms'}+  stms' <- optimiseStms $ stmsToList $ bodyStms kbody+  pure $ kbody {bodyStms = stms'}  optimiseLambda ::   (ASTRep rep) =>@@ -268,21 +268,21 @@           num_bytes <-             letSubExp "num_bytes" =<< toExp (primByteSize pt * (1 + LMAD.range arr_lmad))           arr_mem_in <--            letExp (baseString arg_v <> "_in") $ Op $ Alloc num_bytes space+            letExp (baseName arg_v <> "_in") $ Op $ Alloc num_bytes space           addStm arr_mem_out_alloc            -- Construct additional pattern element and parameter for           -- the memory block that is not used afterwards.           pe_unused <-             PatElem-              <$> newVName (baseString (patElemName pe) <> "_unused")+              <$> newVName (baseName (patElemName pe) <> "_unused")               <*> pure (MemMem space)           param_out <--            newParam (baseString (paramName param) <> "_out") (MemMem space)+            newParam (baseName (paramName param) <> "_out") (MemMem space)            -- Copy the initial array value to the input memory, with           -- the same index function as the result.-          arr_v_copy <- newVName $ baseString arr_v <> "_db_copy"+          arr_v_copy <- newVName $ baseName arr_v <> "_db_copy"           let arr_initial_info =                 MemArray pt shape NoUniqueness $ ArrayIn arr_mem_in arr_lmad               arr_initial_pe =

@@ -211,7 +211,7 @@     mkLetNames [name'] $ BasicOp $ Replicate mempty $ Var name   pure (stmsFromList copies, M.fromList $ zip vs vs')   where-    makeNewName name = newVName $ baseString name <> "_copy"+    makeNewName name = newVName $ baseName name <> "_copy"  okToFuseProducer :: H.SOAC SOACS -> FusionM Bool okToFuseProducer (H.Screma _ _ form) = do@@ -220,7 +220,12 @@ okToFuseProducer _ = pure True  -- First node is producer, second is consumer.-vFuseNodeT :: [EdgeT] -> [VName] -> (NodeT, [EdgeT], [EdgeT]) -> (NodeT, [EdgeT]) -> FusionM (Maybe NodeT)+vFuseNodeT ::+  [EdgeT] ->+  [VName] ->+  (NodeT, [EdgeT], [EdgeT]) ->+  (NodeT, [EdgeT]) ->+  FusionM (Maybe NodeT) vFuseNodeT _ infusible (s1, _, e1s) (MatchNode stm2 dfused, _)   | isRealNode s1,     null infusible =@@ -250,8 +255,15 @@         preserveEdge e = isDep e         preserve = namesFromList $ map getName $ filter preserveEdge i1s     ok <- okToFuseProducer soac1+    -- It is not safe to fuse if any accumulators are updated by both, as the+    -- semantics require that any updates done by the consumer take precedence+    -- over those in the producer. This is implemented with a manual check here+    -- for convenience, but it could be argued that this should really be a Fake+    -- edge in the graph.+    let isProducedAcc (H.Input _ v Acc {}) = v `elem` patNames pats1+        isProducedAcc _ = False     r <--      if ok && ots1 == mempty+      if ok && ots1 == mempty && not (any isProducedAcc (H.inputs soac2))         then TF.attemptFusion TF.Vertical preserve (patNames pats1) soac1 ker         else pure Nothing     case r of@@ -302,35 +314,36 @@ --    for as long as it is returned by the withAcc. If `infusible` is empty --    then the extranous result will be simplified away. vFuseNodeT-  _edges+  edges   _infusible-  (SoacNode ots1 pat1 soac@(H.Screma _w _form _s_inps) aux1, is1, _os1)+  (SoacNode ots1 pat1 soac@(H.Screma _w _form _s_inps) aux1, _is1, os1)   (StmNode (Let pat2 aux2 (WithAcc w_inps lam0)), _os2)     | ots1 == mempty,+      not $ any isFake edges,       wacc_cons_nms <- namesFromList $ concatMap (\(_, nms, _) -> nms) w_inps,       soac_prod_nms <- map patElemName $ patElems pat1,-      soac_indep_nms <- map getName is1,-      all (`notNameIn` wacc_cons_nms) (soac_indep_nms ++ soac_prod_nms) =-        do-          lam <- fst <$> doFusionInLambda lam0-          bdy' <--            runBodyBuilder $ inScopeOf lam $ do-              soac' <- H.toExp soac-              addStm $ Let pat1 aux1 soac'-              lam_res <- bodyBind $ lambdaBody lam-              let pat1_res = map (SubExpRes (Certs []) . Var) soac_prod_nms-              pure $ lam_res ++ pat1_res-          let lam_ret_tp = lambdaReturnType lam ++ map patElemType (patElems pat1)-              pat = Pat $ patElems pat2 ++ patElems pat1-          lam' <- renameLambda $ lam {lambdaBody = bdy', lambdaReturnType = lam_ret_tp}-          -- see if bringing the map inside the scatter has actually benefitted fusion-          (lam'', success) <- doFusionInLambda lam'-          if not success-            then pure Nothing-            else do-              -- `aux1` already appear in the moved SOAC stm; is there-              -- any need to add it to the enclosing withAcc stm as well?-              fusedSomething $ StmNode $ Let pat aux2 $ WithAcc w_inps lam''+      soac_indep_nms <- map getName os1,+      all (`notNameIn` wacc_cons_nms) (soac_indep_nms ++ soac_prod_nms) = do+        lam <- fst <$> doFusionInLambda lam0+        bdy' <-+          runBodyBuilder $ inScopeOf lam $ do+            soac' <- H.toExp soac+            addStm $ Let pat1 aux1 soac'+            lam_res <- bodyBind $ lambdaBody lam+            let pat1_res = map (SubExpRes (Certs []) . Var) soac_prod_nms+            pure $ lam_res ++ pat1_res+        let lam_ret_tp = lambdaReturnType lam ++ map patElemType (patElems pat1)+            pat = Pat $ patElems pat2 ++ patElems pat1+        lam' <- renameLambda $ lam {lambdaBody = bdy', lambdaReturnType = lam_ret_tp}+        -- see if bringing the map inside the scatter has actually benefitted fusion+        (lam'', success) <- doFusionInLambda lam'+        if not success+          then pure Nothing+          else do+            -- `aux1` already appear in the moved SOAC stm; is there+            -- any need to add it to the enclosing withAcc stm as well?+            fusedSomething $ StmNode $ Let pat aux2 $ WithAcc w_inps lam''+ -- -- The reverse of the case above, i.e., fusing a screma at the back of an --   WithAcc such as to (hopefully) enable more fusion there.@@ -379,11 +392,12 @@             fusedSomething $ StmNode $ Let pat aux1 $ WithAcc w_inps wlam'' -- the case of fusing two withaccs vFuseNodeT-  _edges+  edges   infusible   (StmNode (Let pat1 aux1 (WithAcc w1_inps lam1)), is1, _os1)   (StmNode (Let pat2 aux2 (WithAcc w2_inps lam2)), _os2)-    | wacc2_cons_nms <- namesFromList $ concatMap (\(_, nms, _) -> nms) w2_inps,+    | not $ any isFake edges,+      wacc2_cons_nms <- namesFromList $ concatMap (\(_, nms, _) -> nms) w2_inps,       wacc1_indep_nms <- map getName is1,       all (`notNameIn` wacc2_cons_nms) wacc1_indep_nms = do         -- the other safety checks are done inside `tryFuseWithAccs`@@ -429,6 +443,31 @@                 zipWith PatElem (TF.fsOutNames ker') (H.typeOf (TF.fsSOAC ker'))           fusedSomething $ SoacNode mempty (Pat pats2') (TF.fsSOAC ker') (aux1 <> aux2)         Nothing -> pure Nothing+hFuseNodeT+  (StmNode (Let pat1 aux1 (WithAcc w1_inps lam1)))+  (StmNode (Let pat2 aux2 (WithAcc w2_inps lam2))) = do+    -- The only tricky thing here is that we have to put all the+    -- accumulator-based results first.+    let num_inputs1 = length w1_inps+        num_inputs2 = length w2_inps+        num_arrs1 = sum $ map (\(_, as, _) -> length as) w1_inps+        num_arrs2 = sum $ map (\(_, as, _) -> length as) w2_inps+        w3_inps = w1_inps <> w2_inps+        reorder f n a m b =+          let (a_xs, a_ys) = splitAt n $ f a+              (b_xs, b_ys) = splitAt m $ f b+           in a_xs <> b_xs <> a_ys <> b_ys+        lam3 =+          Lambda+            (reorder lambdaParams num_inputs1 lam1 num_inputs2 lam2)+            (reorder lambdaReturnType num_inputs1 lam1 num_inputs2 lam2)+            $ mkBody+              (bodyStms (lambdaBody lam1) <> bodyStms (lambdaBody lam2))+              (reorder (bodyResult . lambdaBody) num_inputs1 lam1 num_inputs2 lam2)+    fusedSomething $+      StmNode $+        Let (Pat $ reorder patElems num_arrs1 pat1 num_arrs2 pat2) (aux1 <> aux2) $+          WithAcc w3_inps lam3 hFuseNodeT _ _ = pure Nothing  removeOutputsExcept :: [VName] -> NodeT -> NodeT@@ -468,13 +507,8 @@ tryFuseNodeInGraph node_to_fuse dg@DepGraph {dgGraph = g}   | not (G.gelem (nodeFromLNode node_to_fuse) g) = pure dg -- \^ Node might have been fused away since.-tryFuseNodeInGraph node_to_fuse dg@DepGraph {dgGraph = g} = do-  spec_rule_res <- SF.ruleMFScat node_to_fuse dg-  -- \^ specialized fusion rules such as the one-  --   enabling map-flatten-scatter fusion-  case spec_rule_res of-    Just dg' -> pure dg'-    Nothing -> applyAugs (map (vTryFuseNodesInGraph node_to_fuse_id) fuses_with) dg+tryFuseNodeInGraph node_to_fuse dg@DepGraph {dgGraph = g} =+  applyAugs (map (vTryFuseNodesInGraph node_to_fuse_id) fuses_with) dg   where     node_to_fuse_id = nodeFromLNode node_to_fuse     relevant (n, InfDep _) = isWithAccNodeId n dg@@ -488,13 +522,13 @@     relevant (_, ResNode {}) = False     relevant _ = True --- | For each pair of SOAC nodes that share an input, attempt to fuse--- them horizontally.+-- | For each pair of SOAC nodes that share an input, or any WithAcc nodes,+-- attempt to fuse them horizontally. doHorizontalFusion :: DepGraphAug FusionM-doHorizontalFusion dg = applyAugs pairs dg+doHorizontalFusion dg = applyAugs (soac_pairs <> withacc_pairs) dg   where-    pairs :: [DepGraphAug FusionM]-    pairs = do+    soac_pairs, withacc_pairs :: [DepGraphAug FusionM]+    soac_pairs = do       (x, SoacNode _ _ soac_x _) <- G.labNodes $ dgGraph dg       (y, SoacNode _ _ soac_y _) <- G.labNodes $ dgGraph dg       guard $ x < y@@ -503,6 +537,16 @@         any           ((`elem` map H.inputArray (H.inputs soac_x)) . H.inputArray)           (H.inputs soac_y)+      pure $ \dg' -> do+        -- Nodes might have been fused away by now.+        if G.gelem x (dgGraph dg') && G.gelem y (dgGraph dg')+          then hTryFuseNodesInGraph x y dg'+          else pure dg'++    withacc_pairs = do+      (x, StmNode (Let _ _ (WithAcc {}))) <- G.labNodes $ dgGraph dg+      (y, StmNode (Let _ _ (WithAcc {}))) <- G.labNodes $ dgGraph dg+      guard $ x < y       pure $ \dg' -> do         -- Nodes might have been fused away by now.         if G.gelem x (dgGraph dg') && G.gelem y (dgGraph dg')

@@ -36,6 +36,7 @@     isCons,     isDep,     isInf,+    isFake,      -- * Construction     mkDepGraph,@@ -74,7 +75,7 @@   | Cons VName   | Fake VName   | Res VName-  deriving (Eq, Ord)+  deriving (Show, Eq, Ord)  -- | Information associated with a node in the graph. data NodeT@@ -93,24 +94,7 @@     FreeNode VName   | MatchNode (Stm SOACS) [(NodeT, [EdgeT])]   | DoNode (Stm SOACS) [(NodeT, [EdgeT])]-  deriving (Eq)--instance Show EdgeT where-  show (Dep vName) = "Dep " <> prettyString vName-  show (InfDep vName) = "iDep " <> prettyString vName-  show (Cons _) = "Cons"-  show (Fake _) = "Fake"-  show (Res _) = "Res"-  show (Alias _) = "Alias"--instance Show NodeT where-  show (StmNode (Let pat _ _)) = L.intercalate ", " $ map prettyString $ patNames pat-  show (SoacNode _ pat _ _) = prettyString pat-  show (TransNode _ tr _) = prettyString tr-  show (ResNode name) = prettyString $ "Res: " ++ prettyString name-  show (FreeNode name) = prettyString $ "Input: " ++ prettyString name-  show (MatchNode stm _) = "Match: " ++ L.intercalate ", " (map prettyString $ stmNames stm)-  show (DoNode stm _) = "Do: " ++ L.intercalate ", " (map prettyString $ stmNames stm)+  deriving (Show, Eq)  -- | The name that this edge depends on. getName :: EdgeT -> VName@@ -132,8 +116,23 @@  -- | Prettyprint dependency graph. pprg :: DepGraph -> String-pprg = G.showDot . G.fglToDotString . G.nemap show show . dgGraph+pprg = G.showDot . G.fglToDotString . G.nemap onNode onEdge . dgGraph+  where+    onEdge (Dep vName) = "Dep " <> prettyString vName+    onEdge (InfDep vName) = "iDep " <> prettyString vName+    onEdge (Cons _) = "Cons"+    onEdge (Fake _) = "Fake"+    onEdge (Res _) = "Res"+    onEdge (Alias _) = "Alias" +    onNode (StmNode (Let pat _ _)) = L.intercalate ", " $ map prettyString $ patNames pat+    onNode (SoacNode _ pat _ _) = prettyString pat+    onNode (TransNode _ tr _) = prettyString tr+    onNode (ResNode name) = prettyString $ "Res: " ++ prettyString name+    onNode (FreeNode name) = prettyString $ "Input: " ++ prettyString name+    onNode (MatchNode stm _) = "Match: " ++ L.intercalate ", " (map prettyString $ stmNames stm)+    onNode (DoNode stm _) = "Do: " ++ L.intercalate ", " (map prettyString $ stmNames stm)+ -- | A pair of a 'G.Node' and the node label. type DepNode = G.LNode NodeT @@ -392,9 +391,7 @@ expInputs (Op soac) = case soac of   Futhark.Screma w is form -> inputs is <> freeClassifications (w, form)   Futhark.Hist w is ops lam -> inputs is <> freeClassifications (w, ops, lam)-  Futhark.Scatter w is lam iws -> inputs is <> freeClassifications (w, lam, iws)-  Futhark.Stream w is nes lam ->-    inputs is <> freeClassifications (w, nes, lam)+  Futhark.Stream w is nes lam -> inputs is <> freeClassifications (w, nes, lam)   Futhark.JVP {} -> freeClassifications soac   Futhark.VJP {} -> freeClassifications soac   where@@ -430,6 +427,10 @@   InfDep _ -> True   Fake _ -> True -- this is infusible to avoid simultaneous cons/dep edges   _ -> False++isFake :: EdgeT -> Bool+isFake (Fake _) = True+isFake _ = False  -- | Is this a 'Cons' edge? isCons :: EdgeT -> Bool

@@ -26,325 +26,16 @@ --        introduced in the code, it is more important that --        they can be transformed by various optimizations passes. module Futhark.Optimise.Fusion.RulesWithAccs-  ( ruleMFScat,-    tryFuseWithAccs,+  ( tryFuseWithAccs,   ) where  import Control.Monad-import Data.Graph.Inductive.Graph qualified as G import Data.Map.Strict qualified as M-import Data.Maybe-import Futhark.Analysis.HORep.SOAC qualified as H import Futhark.Construct import Futhark.IR.SOACS hiding (SOAC (..))-import Futhark.IR.SOACS qualified as F-import Futhark.Optimise.Fusion.GraphRep-import Futhark.Tools import Futhark.Transform.Rename import Futhark.Transform.Substitute--se0 :: SubExp-se0 = intConst Int64 0--se1 :: SubExp-se1 = intConst Int64 1------------------------------------------- I. Map-Flatten-Scatter Fusion --------------------------------------------- helper data structures-type IotaInp = ((VName, LParam SOACS), (SubExp, SubExp, SubExp, IntType))--- ^           ((array-name, lambda param), (len, start, stride, Int64))--type RshpInp = ((VName, LParam SOACS), (Shape, Shape, Type))--- ^           ((array-name, lambda param), (flat-shape, unflat-shape, elem-type))---- | Implements a specialized rule for fusing a pattern---   formed by a map o flatten o scatter, i.e.,---      let (inds,   vals) = map-nest f inps---          (finds, fvals) = (flatten inds, flatten vals)---          let res = scatter res0 finds fvals---   where inds & vals have higher rank than finds & fvals.-ruleMFScat ::-  (HasScope SOACS m, MonadFreshNames m) =>-  DepNode ->-  DepGraph ->-  m (Maybe DepGraph)-ruleMFScat node_to_fuse dg@DepGraph {dgGraph = g}-  | soac_nodeT <- snd node_to_fuse,-    scat_node_id <- nodeFromLNode node_to_fuse,-    SoacNode node_out_trsfs scat_pat scat_soac scat_aux <- soac_nodeT,-    H.nullTransforms node_out_trsfs,-    -- \^ for simplicity we do not allow transforms on scatter's result.-    H.Scatter _len scat_inp scat_out scat_lam <- scat_soac,-    -- \^ get the scatter-    scat_trsfs <- map H.inputTransforms (H.inputs scat_soac),-    -- \^ get the transforms on the input-    any (/= mempty) scat_trsfs,-    scat_ctx <- G.context g scat_node_id,-    (out_deps, _, _, inp_deps) <- scat_ctx,-    cons_deps <- filter (isCons . fst) inp_deps,-    drct_deps <- filter (isDep . fst) inp_deps,-    cons_ctxs <- map (G.context g . snd) cons_deps,-    drct_ctxs <- map (G.context g . snd) drct_deps,-    _cons_nTs <- map getNodeTfromCtx cons_ctxs, -- not used!!-    drct_tups0 <- mapMaybe (pairUp (zip drct_ctxs (map fst drct_deps))) scat_inp,-    length drct_tups0 == length scat_inp,-    -- \^ checks that all direct dependencies are also array-    --   inputs to scatter-    (t1s, t2s) <- unzip drct_tups0,-    drct_tups <- zip t1s $ zip t2s (lambdaParams scat_lam),-    (ctxs_iots, drct_iots) <- unzip $ filter (isIota . snd . fst . snd) drct_tups,-    (ctxs_rshp, drct_rshp) <- unzip $ filter (not . isIota . snd . fst . snd) drct_tups,-    length drct_iots + length drct_rshp == length scat_inp,-    -- \^ direct dependencies are either flatten reshapes or iotas.-    rep_iotas <- mapMaybe getRepIota drct_iots,-    length rep_iotas == length drct_iots,-    rep_rshps_certs <- mapMaybe getRepRshpArr drct_rshp,-    (rep_rshps, certs_rshps) <- unzip rep_rshps_certs,-    -- \^ gather the representations for the iotas and reshapes, that use-    --   the helper types `IotaInp` and `RshpInp`-    not (null rep_rshps),-    -- \^ at least one flatten-reshaped array-    length rep_rshps == length drct_rshp,-    (_, (s1, s2, _)) : _ <- rep_rshps,-    all (\(_, (s1', s2', _)) -> s1 == s1' && s2 == s2') rep_rshps,-    -- \^ Check that all unflatten shape dimensions are the same,-    --   so that we can construct a map nest;--    -- check profitability, which is conservatively defined as all-    -- the reshaped and consumer arrays are used solely by the-    -- scatter AND all reshape dependencies originate in the same-    -- map.-    checkSafeAndProfitable dg scat_node_id ctxs_rshp cons_ctxs = do-      -- generate the withAcc statement-      let cons_patels_outs = zip (patElems scat_pat) scat_out-      wacc_stm <- mkWithAccStm rep_iotas rep_rshps cons_patels_outs scat_aux scat_lam-      let all_cert_rshp = mconcat certs_rshps-          aux = stmAux wacc_stm-          aux' = aux {stmAuxCerts = all_cert_rshp <> stmAuxCerts aux}-          wacc_stm' = wacc_stm {stmAux = aux'}-          -- get the input deps of iotas-          fiot acc (_, _, _, inp_deps_iot) =-            acc <> inp_deps_iot-          deps_of_iotas = foldl fiot mempty ctxs_iots-          ---          iota_nms = namesFromList $ map (fst . fst) rep_iotas-          inp_deps_wo_iotas = filter ((`notNameIn` iota_nms) . getName . fst) inp_deps-          -- generate a new node for the with-acc-stmt and its associated context:-          --   add the inp-deps of iotas but remove the iota themselves from deps.-          new_withacc_nT = StmNode wacc_stm'-          inp_deps' = inp_deps_wo_iotas <> deps_of_iotas-          new_withacc_ctx = (out_deps, scat_node_id, new_withacc_nT, inp_deps')-          -- construct the new WithAcc node/graph; do we need to use `fusedSomething` ??-          new_node = G.node' new_withacc_ctx-          dg' = dg {dgGraph = new_withacc_ctx G.& G.delNodes [new_node] g}-      -- result-      pure $ Just dg'-  where-    ---    getNodeTfromCtx (_, _, nT, _) = nT-    findCtxOf ctxes nm-      | [ctxe] <- filter (\x -> nm == getName (snd x)) ctxes =-          Just ctxe-    findCtxOf _ _ = Nothing-    pairUp :: [(DepContext, EdgeT)] -> H.Input -> Maybe (DepContext, (H.Input, NodeT))-    pairUp ctxes inp@(H.Input _arrtrsfs nm _tp)-      | Just (ctx@(_, _, nT, _), _) <- findCtxOf ctxes nm =-          Just (ctx, (inp, nT))-    pairUp _ _ = Nothing-    ---    isIota :: NodeT -> Bool-    isIota (StmNode (Let _ _ (BasicOp (Iota {})))) = True-    isIota _ = False-    ---    getRepIota :: ((H.Input, NodeT), LParam SOACS) -> Maybe IotaInp-    getRepIota ((H.Input iottrsf arr_nm _arr_tp, nt), farg)-      | mempty == iottrsf,-        StmNode (Let _ _ (BasicOp (Iota n x s Int64))) <- nt =-          Just ((arr_nm, farg), (n, x, s, Int64))-    getRepIota _ = Nothing-    ---    getRepRshpArr :: ((H.Input, NodeT), LParam SOACS) -> Maybe (RshpInp, Certs)-    getRepRshpArr ((H.Input outtrsf arr_nm arr_tp, _nt), farg)-      | rshp_trsfm H.:< other_trsfms <- H.viewf outtrsf,-        H.Reshape aux shp_flat <- rshp_trsfm,-        ReshapeArbitrary <- reshapeKind shp_flat,-        other_trsfms == mempty,-        eltp <- paramDec farg,-        Just shp_flat' <- checkShp eltp $ newShape shp_flat,-        Array _ptp shp_unflat _ <- arr_tp,-        Just shp_unflat' <- checkShp eltp shp_unflat,-        shapeRank shp_flat' == 1,-        shapeRank shp_flat' < shapeRank shp_unflat' =-          Just (((arr_nm, farg), (shp_flat', shp_unflat', eltp)), stmAuxCerts aux)-    getRepRshpArr _ = Nothing-    ---    checkShp (Prim _) shp_arr = Just shp_arr-    checkShp (Array _ptp shp_elm _) shp_arr =-      let dims_elm = shapeDims shp_elm-          dims_arr = shapeDims shp_arr-          (m, n) = (length dims_elm, length dims_arr)-          shp' = Shape $ take (n - m) dims_arr-          dims_com = drop (n - m) dims_arr-       in if all (uncurry (==)) (zip dims_com dims_elm)-            then Just shp'-            else Nothing-    checkShp _ _ = Nothing--- default fails:-ruleMFScat _ _ = pure Nothing--checkSafeAndProfitable :: DepGraph -> G.Node -> [DepContext] -> [DepContext] -> Bool-checkSafeAndProfitable dg scat_node_id ctxs_rshp@(_ : _) ctxs_cons =-  let all_deps = concatMap (\(x, _, _, _) -> x) $ ctxs_rshp ++ ctxs_cons-      prof1 = all (\(_, dep_id) -> dep_id == scat_node_id) all_deps-      -- \^ scatter is the sole target to all consume & unflatten-reshape deps-      (_, map_node_id, map_nT, _) = head ctxs_rshp-      prof2 = all (\(_, nid, _, _) -> nid == map_node_id) ctxs_rshp-      prof3 = isMap map_nT-      -- \^ all reshapes come from the same node, which is a map-      safe = vFusionFeasability dg map_node_id scat_node_id-   in safe && prof1 && prof2 && prof3-  where-    isMap nT-      | SoacNode out_trsfs _pat soac _ <- nT,-        H.Screma _ _ form <- soac,-        ScremaForm _ [] [] <- form =-          H.nullTransforms out_trsfs-    isMap _ = False-checkSafeAndProfitable _ _ _ _ = False---- | produces the withAcc statement that constitutes the translation of---   the scater o flatten o map composition in which the map inputs are---   reshaped in the same way-mkWithAccStm ::-  (HasScope SOACS m, MonadFreshNames m) =>-  [IotaInp] ->-  [RshpInp] ->-  [(PatElem (LetDec SOACS), (Shape, Int, VName))] ->-  StmAux (ExpDec SOACS) ->-  Lambda SOACS ->-  m (Stm SOACS)-mkWithAccStm iota_inps rshp_inps cons_patels_outs scatter_aux scatter_lam-  -- iotas are assumed to operate on Int64 values-  -- ToDo: maybe simplify rshp_inps-  --       check that the unflat shape is the same across reshapes-  --       check that the rank of the unflatten shape is higher than the flatten-  | rshp_inp : _ <- rshp_inps,-    (_, (_, s_unflat, _)) <- rshp_inp,-    (_ : _) <- shapeDims s_unflat = do-      ---      (cert_params, acc_params) <- fmap unzip $-        forM cons_patels_outs $ \(patel, (shp, _, nm)) -> do-          cert_param <- newParam "acc_cert_p" $ Prim Unit-          let arr_tp = patElemType patel-              acc_tp = stripArray (shapeRank shp) arr_tp-          acc_param <--            newParam (baseString nm) $-              Acc (paramName cert_param) shp [acc_tp] NoUniqueness-          pure (cert_param, acc_param)-      let cons_params_outs = zip acc_params $ map snd cons_patels_outs-      acc_bdy <- mkWithAccBdy s_unflat iota_inps rshp_inps cons_params_outs scatter_lam-      let withacc_lam =-            Lambda-              { lambdaParams = cert_params ++ acc_params,-                lambdaReturnType = map paramDec acc_params,-                lambdaBody = acc_bdy-              }-          withacc_inps = map (\(_, (shp, _, nm)) -> (shp, [nm], Nothing)) cons_patels_outs-          withacc_pat = Pat $ map fst cons_patels_outs-          stm =-            Let withacc_pat scatter_aux $-              WithAcc withacc_inps withacc_lam-      pure stm-mkWithAccStm _ _ _ _ _ =-  error "Unreachable case reached!"---- | Wrapper function for constructing the body of the withAcc---   translation of the scatter-mkWithAccBdy ::-  (HasScope SOACS m, MonadFreshNames m) =>-  Shape ->-  [IotaInp] ->-  [RshpInp] ->-  [(LParam SOACS, (Shape, Int, VName))] ->-  Lambda SOACS ->-  m (Body SOACS)-mkWithAccBdy shp iota_inps rshp_inps cons_params_outs scat_lam = do-  let cons_ps = map fst cons_params_outs-      scat_res_info = map snd cons_params_outs-      static_arg = (iota_inps, rshp_inps, scat_res_info, scat_lam)-      mkParam ((nm, _), (_, s, t)) = Param mempty nm (arrayOfShape t s)-      rshp_ps = map mkParam rshp_inps-  mkWithAccBdy' static_arg (shapeDims shp) [] [] rshp_ps cons_ps---- | builds a body that essentially consists of a map-nest with accumulators,---   i.e., one level for each level of the unflatten shape of scatter's reshaped---   input arrays-mkWithAccBdy' ::-  (HasScope SOACS m, MonadFreshNames m) =>-  ([IotaInp], [RshpInp], [(Shape, Int, VName)], Lambda SOACS) ->-  [SubExp] ->-  [SubExp] ->-  [VName] ->-  [LParam SOACS] ->-  [LParam SOACS] ->-  m (Body SOACS)---- | the base case below addapts the scatter's lambda-mkWithAccBdy' static_arg [] dims_rev iot_par_nms rshp_ps cons_ps = do-  let (iota_inps, rshp_inps, scat_res_info, scat_lam) = static_arg-      tp_int = Prim $ IntType Int64-  scope <- askScope-  runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do-    -- handle iota args-    let strides_rev = scanl (*) (pe64 se1) $ map pe64 dims_rev-        strides = tail $ reverse strides_rev-        prods = zipWith (*) (map le64 iot_par_nms) strides-        i_pe = sum prods-    i_norm <- letExp "iota_norm_arg" =<< toExp i_pe-    forM_ iota_inps $ \arg -> do-      let ((_, i_par), (_, b, s, _)) = arg-      i_new <- letExp "tmp" =<< toExp (pe64 b + le64 i_norm * pe64 s)-      letBind (Pat [PatElem (paramName i_par) tp_int]) $ BasicOp $ SubExp $ Var i_new-    -- handle rshp args-    let rshp_lam_args = map (snd . fst) rshp_inps-    forM_ (zip rshp_lam_args rshp_ps) $ \(old_par, new_par) -> do-      let pat = Pat [PatElem (paramName old_par) (paramDec old_par)]-      letBind pat $ BasicOp $ SubExp $ Var $ paramName new_par-    -- add the body of the scatter's lambda-    mapM_ addStm $ bodyStms $ lambdaBody scat_lam-    -- add the withAcc update statements-    let iv_ses = groupScatterResults' scat_res_info $ bodyResult $ lambdaBody scat_lam-    res_nms <--      forM (zip cons_ps iv_ses) $ \(cons_p, (i_ses, v_se)) -> do-        -- i_ses is a list-        let f nm_in i_se =-              letExp (baseString nm_in) $ BasicOp $ UpdateAcc Safe nm_in [resSubExp i_se] [resSubExp v_se]-        foldM f (paramName cons_p) i_ses-    let lam_certs = foldMap resCerts $ bodyResult $ lambdaBody scat_lam-    pure $ map (SubExpRes lam_certs . Var) res_nms--- \| the recursive case builds a call to a map soac.-mkWithAccBdy' static_arg (dim : dims) dims_rev iot_par_nms rshp_ps cons_ps = do-  scope <- askScope-  runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do-    iota_arr <- letExp "iota_arr" $ BasicOp $ Iota dim se0 se1 Int64-    iota_p <- newParam "iota_arg" $ Prim $ IntType Int64-    rshp_ps' <- forM (zip [0 .. length rshp_ps - 1] (map paramDec rshp_ps)) $-      \(i, arr_tp) ->-        newParam ("rshp_arg_" ++ show i) $ stripArray 1 arr_tp-    cons_ps' <- forM (zip [0 .. length cons_ps - 1] (map paramDec cons_ps)) $-      \(i, arr_tp) ->-        newParam ("acc_arg_" ++ show i) arr_tp-    map_lam_bdy <--      mkWithAccBdy' static_arg dims (dim : dims_rev) (iot_par_nms ++ [paramName iota_p]) rshp_ps' cons_ps'-    let map_lam = Lambda (rshp_ps' ++ [iota_p] ++ cons_ps') (map paramDec cons_ps') map_lam_bdy-        map_inps = map paramName rshp_ps ++ [iota_arr] ++ map paramName cons_ps-        map_soac = F.Screma dim map_inps $ ScremaForm map_lam [] []-    res_nms <- letTupExp "acc_res" $ Op map_soac-    pure $ map (subExpRes . Var) res_nms  --------------------------------------------------- --- II. WithAcc-WithAcc Fusion

@@ -238,7 +238,7 @@   guard $ all inputOrUnfus outVars    outPairs <- forM (zip outVars $ map rowType $ SOAC.typeOf soac_p) $ \(outVar, t) -> do-    outVar' <- newVName $ baseString outVar ++ "_elem"+    outVar' <- newVName $ baseName outVar <> "_elem"     pure (outVar, Ident outVar' t)    let mapLikeFusionCheck =@@ -308,29 +308,6 @@                 new_inp                 (ScremaForm res_lam' (scans_p ++ scans_c) (reds_p ++ reds_c)) -    -------------------    -- Scatter fusion ---    --------------------    -- Map-Scatter fusion.-    ---    -- The 'inplace' mechanism for kernels already takes care of-    -- checking that the Scatter is not writing to any array used in-    -- the Map.-    ( SOAC.Scatter _len _ivs dests _lam,-      SOAC.Screma _ _ form,-      _-      )-        | isJust $ isMapSOAC form,-          -- 1. all arrays produced by the map are ONLY used (consumed)-          --    by the scatter, i.e., not used elsewhere.-          all (`notNameIn` unfus_set) outVars,-          -- 2. all arrays produced by the map are input to the scatter.-          mapWriteFusionOK outVars ker -> do-            let (extra_nms, res_lam', new_inp) = mapLikeFusionCheck-            success (fsOutNames ker ++ extra_nms) $-              SOAC.Scatter w new_inp dests res_lam'-     -- Map-Hist fusion.     --     -- The 'inplace' mechanism for kernels already takes care of@@ -379,35 +356,6 @@                 }         success (fsOutNames ker ++ returned_outvars) $           SOAC.Hist w (inp_c_arr <> inp_p_arr) (ops_c <> ops_p) lam'--    -- Scatter-write fusion.-    ( SOAC.Scatter _w_c ivs_c as_c _lam_c,-      SOAC.Scatter _w_p ivs_p as_p _lam_p,-      Horizontal-      ) -> do-        let zipW as_xs xs as_ys ys = xs_indices ++ ys_indices ++ xs_vals ++ ys_vals-              where-                (xs_indices, xs_vals) = splitScatterResults as_xs xs-                (ys_indices, ys_vals) = splitScatterResults as_ys ys-        let (body_p, body_c) = (lambdaBody lam_p, lambdaBody lam_c)-        let body' =-              Body-                { bodyDec = bodyDec body_p, -- body_p and body_c have the same decorations-                  bodyStms = bodyStms body_p <> bodyStms body_c,-                  bodyResult = zipW as_c (bodyResult body_c) as_p (bodyResult body_p)-                }-        let lam' =-              Lambda-                { lambdaParams = lambdaParams lam_c ++ lambdaParams lam_p,-                  lambdaBody = body',-                  lambdaReturnType = zipW as_c (lambdaReturnType lam_c) as_p (lambdaReturnType lam_p)-                }-        success (fsOutNames ker ++ returned_outvars) $-          SOAC.Scatter w (ivs_c ++ ivs_p) (as_c ++ as_p) lam'-    (SOAC.Scatter {}, _, _) ->-      fail "Cannot fuse a scatter with anything else than a scatter or a map"-    (_, SOAC.Scatter {}, _) ->-      fail "Cannot fuse a scatter with anything else than a scatter or a map"     (_, SOAC.Hist {}, _) ->       fail "Cannot fuse a Hist with anything else than a Hist or a Map"     (SOAC.Hist {}, _, _) ->@@ -690,7 +638,7 @@           sliceRes (SubExpRes rcs (Var v)) =             certifying rcs               . fmap subExpRes-              . letSubExp (baseString v <> "_sliced")+              . letSubExp (baseName v <> "_sliced")               $ BasicOp (Index v (Slice inner_ds))           sliceRes r = pure r           inner_changed =

@@ -94,7 +94,7 @@ genRed2Tile2d env kerstm@(Let pat_ker aux (Op (SegOp (SegMap seg_thd seg_space kres_tps old_kbody))))   | SegThread _novirt _ <- seg_thd,     -- novirt == SegNoVirtFull || novirt == SegNoVirt,-    KernelBody () kstms kres <- old_kbody,+    Body () kstms kres <- old_kbody,     Just (css, r_ses) <- allGoodReturns kres,     null css,     -- build the variance table, that records, for@@ -105,8 +105,11 @@     -- some `code1`, followed by one accumulation, followed by some `code2`     -- UpdateAcc VName [SubExp] [SubExp]     (code1, Just accum_stmt, code2) <- matchCodeAccumCode kstms,-    Let pat_accum _aux_acc (BasicOp (UpdateAcc safety acc_nm acc_inds acc_vals)) <- accum_stmt,-    [pat_acc_nm] <- patNames pat_accum,+    Let+      pat_accum@(Pat [PatElem pat_acc_nm acc_tp])+      _aux_acc+      (BasicOp (UpdateAcc safety acc_nm acc_inds acc_vals)) <-+      accum_stmt,     -- check that the `acc_inds` are invariant to at least one     -- parallel kernel dimensions, and return the innermost such one:     Just (invar_gid, gid_ind) <- isInvarToParDim mempty seg_space variance acc_inds,@@ -126,12 +129,18 @@     --   memory accesses: if more than two are re-executed, then we     --   should abort.     cost <- costRedundantExecution variance pat_acc_nm r_ses kstms,-    maxCost cost (Small 2) == Small 2 = do+    maxCost cost (Small 2) == Small 2,+    -- Must hav eaccumulation operator.+    Just ((redop0, neutral), el_tps) <- getAccLambda acc_tp,+    -- HACK: if any of the indexes depend on names that will not be in scope in+    -- the generated kernel, then we do not do the transformation. A better+    -- solution would be to actually put the necessary statements in the kernel.+    not $+      freeIn acc_inds+        `namesIntersect` namesFromList (concatMap (patNames . stmPat) kstms) = do       -- 1. create the first kernel-      acc_tp <- lookupType acc_nm       let inv_dim_len = segSpaceDims seg_space !! gid_ind-          -- 1.1. get the accumulation operator-          ((redop0, neutral), el_tps) = getAccLambda acc_tp+      -- 1.1. get the accumulation operator       redop <- renameLambda redop0       let red =             Reduce@@ -152,7 +161,7 @@         iota <- letExp "iota" $ BasicOp $ Iota inv_dim_len (intConst Int64 0) (intConst Int64 1) Int64         let op_exp = Op (OtherOp (Screma inv_dim_len [iota] (ScremaForm map_lam [] [red])))         res_redmap <- letTupExp "res_mapred" op_exp-        letSubExp (baseString pat_acc_nm ++ "_big_update") $+        letSubExp (baseName pat_acc_nm <> "_big_update") $           BasicOp (UpdateAcc safety acc_nm acc_inds $ map Var res_redmap)        -- 1.3. build the kernel expression and rename it!@@ -160,18 +169,17 @@       let space1 = SegSpace gid_flat_1 gid_dims_new        let level1 = SegThread (SegNoVirtFull (SegSeqDims [])) Nothing -- novirt ?-          kbody1 = KernelBody () ker1_stms [Returns ResultMaySimplify (Certs []) k1_res]+          kbody1 = Body () ker1_stms [Returns ResultMaySimplify (Certs []) k1_res]        -- is it OK here to use the "aux" from the parrent kernel?       ker_exp <- renameExp $ Op (SegOp (SegMap level1 space1 [acc_tp] kbody1))       let ker1 = Let pat_accum aux ker_exp        -- 2 build the second kernel-      let ker2_body = old_kbody {kernelBodyStms = code1 <> code2}+      let ker2_body = old_kbody {bodyStms = code1 <> code2}       ker2_exp <- renameExp $ Op (SegOp (SegMap seg_thd seg_space kres_tps ker2_body))       let ker2 = Let pat_ker aux ker2_exp-      pure $-        Just (code1_tr_host <> oneStm ker1, ker2)+      pure $ Just (code1_tr_host <> oneStm ker1, ker2)   where     isIndVarToParDim _ (Constant _) _ = False     isIndVarToParDim variance (Var acc_ind) par_dim =@@ -196,8 +204,8 @@       case acc_tp of         (Acc tp_id _shp el_tps _) ->           case M.lookup tp_id (fst env) of-            Just lam -> (lam, el_tps)-            _ -> error $ "Lookup in environment failed! " ++ prettyString tp_id ++ " env: " ++ show (fst env)+            Just lam -> Just (lam, el_tps)+            _ -> Nothing         _ -> error "Illegal accumulator type!"     -- is a subexp invariant to a gid of a parallel dimension?     isSeInvar2 variance gid (Var x) =@@ -261,8 +269,8 @@         ii /= length dims - 1,         perm <- [0 .. ii - 1] ++ [ii + 1 .. length dims - 1] ++ [ii] = do           (arr_tr, stms_tr) <- runBuilderT' $ do-            arr' <- letExp (baseString arr ++ "_trsp") $ BasicOp $ Rearrange arr perm-            letExp (baseString arr' ++ "_opaque") $ BasicOp $ Opaque OpaqueNil $ Var arr'+            arr' <- letExp (baseName arr <> "_trsp") $ BasicOp $ Rearrange arr perm+            letExp (baseName arr' <> "_opaque") $ BasicOp $ Opaque OpaqueNil $ Var arr'           let tab' = M.insert arr (perm, arr_tr, stms_tr) tab               slc' = Slice $ map (dims !!) perm               stm' = Let pat aux $ BasicOp $ Index arr_tr slc'

@@ -47,10 +47,10 @@       pure (x, oneStm stm <> stms'')  mkHistBody :: Accs GPU -> KernelBody GPU -> Maybe (KernelBody GPU, WithAccInput GPU, VName)-mkHistBody accs (KernelBody () stms [Returns rm cs (Var v)]) = do+mkHistBody accs (Body () stms [Returns rm cs (Var v)]) = do   ((acc_input, acc, is, vs), stms') <- extractUpdate accs v stms   pure-    ( KernelBody () stms' $ map (Returns rm cs) is ++ map (Returns rm cs) vs,+    ( Body () stms' $ map (Returns rm cs) is <> map (Returns rm cs) vs,       acc_input,       acc     )@@ -75,18 +75,18 @@   (acc', stms) <- localScope (scopeOfSegSpace space) . collectStms $ do     vs <- forM arrs $ \arr -> do       arr_t <- lookupType arr-      letSubExp (baseString arr <> "_elem") $+      letSubExp (baseName arr <> "_elem") $         BasicOp $           Index arr $             fullSlice arr_t $               map (DimFix . Var) gtids-    letExp (baseString acc <> "_upd") $+    letExp (baseName acc <> "_upd") $       BasicOp $         UpdateAcc Safe acc (map Var gtids) vs    acc_t <- lookupType acc   pure . Op . SegOp . SegMap lvl space [acc_t] $-    KernelBody () stms [Returns ResultMaySimplify mempty (Var acc')]+    Body () stms [Returns ResultMaySimplify mempty (Var acc')]  flatKernelBody ::   (MonadBuilder m) =>@@ -106,9 +106,9 @@           unflattenIndex (map pe64 (segSpaceDims space)) (pe64 $ Var gtid)     zipWithM_ letBindNames (map (pure . fst) (unSegSpace space))       =<< mapM toExp new_inds-    addStms $ kernelBodyStms kbody+    addStms $ bodyStms kbody -  pure (space', kbody {kernelBodyStms = kbody_stms})+  pure (space', kbody {bodyStms = kbody_stms})  optimiseStm :: Accs GPU -> Stm GPU -> OptM (Stms GPU) -- TODO: this is very restricted currently, but shows the idea.

@@ -36,13 +36,13 @@  getAllocsSegOp :: SegOp lvl GPUMem -> Allocs getAllocsSegOp (SegMap _ _ _ body) =-  foldMap getAllocsStm (kernelBodyStms body)+  foldMap getAllocsStm (bodyStms body) getAllocsSegOp (SegRed _ _ _ body _) =-  foldMap getAllocsStm (kernelBodyStms body)+  foldMap getAllocsStm (bodyStms body) getAllocsSegOp (SegScan _ _ _ body _) =-  foldMap getAllocsStm (kernelBodyStms body)+  foldMap getAllocsStm (bodyStms body) getAllocsSegOp (SegHist _ _ _ body _) =-  foldMap getAllocsStm (kernelBodyStms body)+  foldMap getAllocsStm (bodyStms body)  setAllocsStm :: Map VName SubExp -> Stm GPUMem -> Stm GPUMem setAllocsStm m stm@(Let (Pat [PatElem name _]) _ (Op (Alloc _ _)))@@ -68,19 +68,19 @@   SegOp lvl GPUMem setAllocsSegOp m (SegMap lvl sp tps body) =   SegMap lvl sp tps $-    body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body}+    body {bodyStms = setAllocsStm m <$> bodyStms body} setAllocsSegOp m (SegRed lvl sp tps body ops) =   SegRed lvl sp tps body' ops   where-    body' = body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body}+    body' = body {bodyStms = setAllocsStm m <$> bodyStms body} setAllocsSegOp m (SegScan lvl sp tps body ops) =   SegScan lvl sp tps body' ops   where-    body' = body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body}+    body' = body {bodyStms = setAllocsStm m <$> bodyStms body} setAllocsSegOp m (SegHist lvl sp tps body ops) =   SegHist lvl sp tps body' ops   where-    body' = body {kernelBodyStms = setAllocsStm m <$> kernelBodyStms body}+    body' = body {bodyStms = setAllocsStm m <$> bodyStms body}  maxSubExp :: (MonadBuilder m) => Set SubExp -> m SubExp maxSubExp = helper . S.toList@@ -106,17 +106,17 @@   (Stms GPUMem -> m (Stms GPUMem)) ->   m (SegOp lvl GPUMem) onKernelBodyStms (SegMap lvl space ts body) f = do-  stms <- f $ kernelBodyStms body-  pure $ SegMap lvl space ts $ body {kernelBodyStms = stms}+  stms <- f $ bodyStms body+  pure $ SegMap lvl space ts $ body {bodyStms = stms} onKernelBodyStms (SegRed lvl space ts body binops) f = do-  stms <- f $ kernelBodyStms body-  pure $ SegRed lvl space ts (body {kernelBodyStms = stms}) binops+  stms <- f $ bodyStms body+  pure $ SegRed lvl space ts (body {bodyStms = stms}) binops onKernelBodyStms (SegScan lvl space ts body binops) f = do-  stms <- f $ kernelBodyStms body-  pure $ SegScan lvl space ts (body {kernelBodyStms = stms}) binops+  stms <- f $ bodyStms body+  pure $ SegScan lvl space ts (body {bodyStms = stms}) binops onKernelBodyStms (SegHist lvl space ts body binops) f = do-  stms <- f $ kernelBodyStms body-  pure $ SegHist lvl space ts (body {kernelBodyStms = stms}) binops+  stms <- f $ bodyStms body+  pure $ SegHist lvl space ts (body {bodyStms = stms}) binops  -- | This is the actual optimiser. Given an interference graph and a @SegOp@, -- replace allocations and references to memory blocks inside with a (hopefully)@@ -150,19 +150,19 @@   pure $ case segop' of     SegMap lvl sp tps body ->       SegMap lvl sp tps $-        body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body}+        body {bodyStms = maxstms <> stms <> bodyStms body}     SegRed lvl sp tps body ops ->       SegRed lvl sp tps body' ops       where-        body' = body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body}+        body' = body {bodyStms = maxstms <> stms <> bodyStms body}     SegScan lvl sp tps body ops ->       SegScan lvl sp tps body' ops       where-        body' = body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body}+        body' = body {bodyStms = maxstms <> stms <> bodyStms body}     SegHist lvl sp tps body ops ->       SegHist lvl sp tps body' ops       where-        body' = body {kernelBodyStms = maxstms <> stms <> kernelBodyStms body}+        body' = body {bodyStms = maxstms <> stms <> bodyStms body}  -- | Helper function that modifies kernels found inside some statements. onKernels ::

@@ -26,7 +26,7 @@ import Futhark.Error import Futhark.IR.Aliases import Futhark.IR.GPU-import Futhark.MonadFreshNames hiding (newName)+import Futhark.MonadFreshNames qualified import Futhark.Pass  -- | An optimization pass that reorders and merges 'GPUBody' statements to@@ -635,5 +635,5 @@   pure name  -- | Produce a fresh name, using the given string as a template.-newName :: String -> RewriteM VName-newName s = lift $ lift (newNameFromString s)+newName :: Name -> RewriteM VName+newName s = lift $ lift $ Futhark.MonadFreshNames.newVName s

@@ -326,18 +326,18 @@ -- that depends on migrated scalars. optimizeHostOp :: HostOp op GPU -> ReduceM (HostOp op GPU) optimizeHostOp (SegOp (SegMap lvl space types kbody)) =-  SegOp . SegMap lvl space types <$> addReadsToKernelBody kbody+  SegOp . SegMap lvl space types <$> addReadsToBody kbody optimizeHostOp (SegOp (SegRed lvl space types kbody ops)) = do   ops' <- mapM addReadsToSegBinOp ops-  kbody' <- addReadsToKernelBody kbody+  kbody' <- addReadsToBody kbody   pure . SegOp $ SegRed lvl space types kbody' ops' optimizeHostOp (SegOp (SegScan lvl space types kbody ops)) = do   ops' <- mapM addReadsToSegBinOp ops-  kbody' <- addReadsToKernelBody kbody+  kbody' <- addReadsToBody kbody   pure . SegOp $ SegScan lvl space types kbody' ops' optimizeHostOp (SegOp (SegHist lvl space types kbody ops)) = do   ops' <- mapM addReadsToHistOp ops-  kbody' <- addReadsToKernelBody kbody+  kbody' <- addReadsToBody kbody   pure . SegOp $ SegHist lvl space types kbody' ops' optimizeHostOp (SizeOp op) =   pure (SizeOp op)@@ -658,16 +658,13 @@   pure (f {lambdaBody = body'})  -- | Rewrite generic body dependencies to scalars that have been migrated.-addReadsToBody :: Body GPU -> ReduceM (Body GPU)+addReadsToBody ::+  (FreeIn res, Substitute res) =>+  GBody GPU res ->+  ReduceM (GBody GPU res) addReadsToBody body = do   (body', prologue) <- addReadsHelper body   pure body' {bodyStms = prologue >< bodyStms body'}---- | Rewrite kernel body dependencies to scalars that have been migrated.-addReadsToKernelBody :: KernelBody GPU -> ReduceM (KernelBody GPU)-addReadsToKernelBody kbody = do-  (kbody', prologue) <- addReadsHelper kbody-  pure kbody' {kernelBodyStms = prologue >< kernelBodyStms kbody'}  -- | Rewrite migrated scalar dependencies within anything. The returned -- statements must be added to the scope of the rewritten construct.

@@ -230,11 +230,11 @@  -- | Produce a body with simplifier information. mkWiseBody ::-  (Informing rep) =>+  (Informing rep, IsResult res) =>   BodyDec rep ->   Stms (Wise rep) ->-  Result ->-  Body (Wise rep)+  [res] ->+  GBody (Wise rep) res mkWiseBody dec stms res =   Body     ( BodyWisdom aliases consumed (AliasDec $ freeIn $ freeInStmsAndRes stms res),@@ -314,7 +314,7 @@ informStms = fmap informStm  -- | Construct a 'Wise' body.-informBody :: (Informing rep) => Body rep -> Body (Wise rep)+informBody :: (Informing rep, IsResult res) => GBody rep res -> GBody (Wise rep) res informBody (Body dec stms res) = mkWiseBody dec (informStms stms) res  -- | Construct a 'Wise' lambda.

@@ -101,6 +101,10 @@     Just result <- fun args' =       Simplify . auxing aux . letBind pat $         BasicOp (SubExp $ Constant result)+  | [(Constant (BoolValue b), _), (x, _), (y, _)] <- args,+    Just _ <- isCondFun $ nameToText fname =+      Simplify . auxing aux . letBind pat $+        BasicOp (SubExp $ if b then x else y)   where     isConst (Constant v) = Just v     isConst _ = Nothing

@@ -191,7 +191,7 @@         Var v | not $ null $ sliceDims is -> do           v_t <- lookupType v           v_reshaped <--            letSubExp (baseString v ++ "_reshaped") . BasicOp $+            letSubExp (baseName v <> "_reshaped") . BasicOp $               Reshape v $                 reshapeAll (arrayShape v_t) (arrayShape dest_t)           letBind pat $ BasicOp $ Replicate mempty v_reshaped@@ -368,13 +368,13 @@              flipRearrangeReshape perm v3_shape            ) of         (Just perm', _) -> Simplify $ do-          v1' <- letExp (baseString v1) $ BasicOp $ Rearrange v0 perm'+          v1' <- letExp (baseName v1) $ BasicOp $ Rearrange v0 perm'           v1_shape' <- arrayShape <$> lookupType v1'           auxing aux . certifying (v1_cs <> v2_cs) . letBind pat $             BasicOp (Reshape v1' (reshapeAll v1_shape' (newShape v3_shape)))         (_, Just (v3_shape', perm')) -> Simplify $ do           v2' <--            auxing aux . certifying (v1_cs <> v2_cs) . letExp (baseString v2) $+            auxing aux . certifying (v1_cs <> v2_cs) . letExp (baseName v2) $               BasicOp (Reshape v1 v3_shape')           letBind pat $ BasicOp (Rearrange v2' perm')         _ ->@@ -387,7 +387,7 @@     ST.available v1 vtable =       Simplify $ do         v1' <--          certifying cs . auxing aux . letExp (baseString v1) . BasicOp $+          certifying cs . auxing aux . letExp (baseName v1) . BasicOp $             Reshape v1 newshape         letBind pat $ BasicOp $ Replicate (Shape []) (Var v1') ruleBasicOp vtable pat aux (Rearrange v2 perm)@@ -395,7 +395,7 @@     ST.available v1 vtable =       Simplify $ do         v1' <--          certifying cs . auxing aux . letExp (baseString v1) . BasicOp $+          certifying cs . auxing aux . letExp (baseName v1) . BasicOp $             Rearrange v1 perm         letBind pat $ BasicOp $ Replicate (Shape []) (Var v1') ruleBasicOp _ _ _ _ =

@@ -101,10 +101,10 @@             else Just $ do               arr_sliced <-                 certifying cs $-                  letExp (baseString arr <> "_sliced") . BasicOp . Index arr . Slice+                  letExp (baseName arr <> "_sliced") . BasicOp . Index arr . Slice                     =<< sequence inds'''               arr_sliced_tr <--                letSubExp (baseString arr_sliced <> "_tr") $+                letSubExp (baseName arr_sliced <> "_tr") $                   BasicOp (Rearrange arr_sliced perm)               pure $ SubExpResult mempty arr_sliced_tr       where

@@ -208,9 +208,9 @@   (Constraints rep) =>   Sinker rep (Op rep) ->   Sinker rep (KernelBody rep)-optimiseKernelBody onOp vtable sinking (KernelBody attr stms res) =+optimiseKernelBody onOp vtable sinking (Body attr stms res) =   let (stms', sunk) = optimiseStms onOp vtable sinking stms $ freeIn res-   in (KernelBody attr stms' res, sunk)+   in (Body attr stms' res, sunk)  optimiseSegOp ::   (Constraints rep) =>

@@ -80,10 +80,10 @@   KernelBody GPU ->   TileM (Stms GPU, (SegLevel, SegSpace, KernelBody GPU)) tileInKernelBody branch_variant initial_variance lvl initial_kspace ts kbody-  | Just kbody_res <- mapM isSimpleResult $ kernelBodyResult kbody = do+  | Just kbody_res <- mapM isSimpleResult $ bodyResult kbody = do       maybe_tiled <-         tileInBody branch_variant initial_variance lvl initial_kspace ts $-          Body () (kernelBodyStms kbody) kbody_res+          Body () (bodyStms kbody) kbody_res       case maybe_tiled of         Just (host_stms, tiling, tiledBody) -> do           (res', stms') <-@@ -92,7 +92,7 @@             ( host_stms,               ( tilingLevel tiling,                 tilingSpace tiling,-                KernelBody () stms' res'+                Body () stms' res'               )             )         Nothing ->@@ -394,7 +394,7 @@             doPrelude tiling privstms precomputed_variant_prestms live_set          mergeparams' <- forM mergeparams $ \(Param attrs pname pt) ->-          Param attrs <$> newVName (baseString pname ++ "_group") <*> pure (tileDim pt)+          Param attrs <$> newVName (baseName pname <> "_group") <*> pure (tileDim pt)          let merge_ts = map paramType mergeparams @@ -577,7 +577,7 @@ -- the kernel. data Tiling = Tiling   { tilingSegMap ::-      String ->+      Name ->       ResultManifest ->       (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result) ->       Builder GPU [VName],@@ -606,7 +606,7 @@   gtids -> kdims -> SubExp -> Builder GPU Tiling  protectOutOfBounds ::-  String ->+  Name ->   PrimExp VName ->   [Type] ->   Builder GPU Result ->@@ -694,7 +694,7 @@       merge <- forM (zip (lambdaParams red_lam) mergeinits) $ \(p, mergeinit) ->         (,)           <$> newParam-            (baseString (paramName p) ++ "_merge")+            (baseName (paramName p) <> "_merge")             (paramType p `arrayOfShape` tile_shape `toDecl` Unique)           <*> pure (Var mergeinit) @@ -740,7 +740,7 @@       then pure arr       else do         let new_shape = Shape $ unit_dims ++ arrayDims arr_t-        letExp (baseString arr) . BasicOp $+        letExp (baseName arr) . BasicOp $           Reshape arr (reshapeAll (arrayShape arr_t) new_shape)   let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims   pure $ TileReturns mempty tile_dims arr'@@ -913,7 +913,7 @@   gid <- newVName "gid"   gid_flat <- newVName "gid_flat" -  tile_size_key <- nameFromString . prettyString <$> newVName "tile_size"+  tile_size_key <- nameFromText . prettyText <$> newVName "tile_size"   tile_size <- letSubExp "tile_size" $ Op $ SizeOp $ GetSize tile_size_key SizeThreadBlock   let tblock_size = tile_size