futhark 0.20.6 → 0.20.7

raw patch · 80 files changed

+1664/−1039 lines, 80 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• docs/conf.py +4/−4
• docs/index.rst +1/−0
• docs/man/futhark-autotune.rst +6/−6
• docs/performance.rst +424/−0
• docs/requirements.txt +1/−1
• futhark.cabal +1/−1
• prelude/prelude.fut +4/−0
• rts/python/scalar.py +2/−2
• rts/python/values.py +1/−1
• src/Futhark/Analysis/PrimExp/Parse.hs +21/−8
• src/Futhark/Analysis/UsageTable.hs +1/−0
• src/Futhark/CodeGen/Backends/GenericC.hs +29/−27
• src/Futhark/CodeGen/Backends/GenericPython.hs +19/−18
• src/Futhark/CodeGen/ImpCode.hs +22/−45
• src/Futhark/CodeGen/ImpCode/GPU.hs +2/−2
• src/Futhark/CodeGen/ImpGen.hs +37/−34
• src/Futhark/CodeGen/ImpGen/GPU.hs +22/−22
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +65/−58
• src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs +8/−11
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +6/−6
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs +10/−8
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +12/−13
• src/Futhark/CodeGen/ImpGen/GPU/ToOpenCL.hs +5/−4
• src/Futhark/CodeGen/ImpGen/GPU/Transpose.hs +120/−98
• src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +1/−1
• src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs +8/−8
• src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs +1/−1
• src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs +5/−5
• src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs +4/−4
• src/Futhark/CodeGen/ImpGen/Transpose.hs +60/−53
• src/Futhark/CodeGen/SetDefaultSpace.hs +18/−42
• src/Futhark/Doc/Generator.hs +2/−2
• src/Futhark/IR/Aliases.hs +1/−1
• src/Futhark/IR/GPU/Op.hs +4/−0
• src/Futhark/IR/GPU/Simplify.hs +2/−2
• src/Futhark/IR/GPUMem.hs +1/−6
• src/Futhark/IR/MC/Op.hs +7/−2
• src/Futhark/IR/Mem.hs +2/−0
• src/Futhark/IR/Mem/Simplify.hs +2/−19
• src/Futhark/IR/Parse.hs +7/−7
• src/Futhark/IR/Primitive.hs +12/−2
• src/Futhark/IR/Prop/Aliases.hs +20/−5
• src/Futhark/IR/SOACS/SOAC.hs +2/−3
• src/Futhark/IR/SOACS/Simplify.hs +9/−14
• src/Futhark/IR/SegOp.hs +29/−15
• src/Futhark/Internalise/Defunctionalise.hs +11/−18
• src/Futhark/Internalise/Defunctorise.hs +6/−14
• src/Futhark/Internalise/Exps.hs +22/−14
• src/Futhark/Internalise/LiftLambdas.hs +2/−4
• src/Futhark/Internalise/Monomorphise.hs +10/−11
• src/Futhark/Optimise/DoubleBuffer.hs +8/−2
• src/Futhark/Optimise/InPlaceLowering.hs +17/−9
• src/Futhark/Optimise/InliningDeadFun.hs +8/−5
• src/Futhark/Optimise/Simplify.hs +32/−29
• src/Futhark/Optimise/Simplify/Engine.hs +242/−230
• src/Futhark/Optimise/Simplify/Rep.hs +51/−0
• src/Futhark/Optimise/TileLoops.hs +7/−7
• src/Futhark/Pass/ExpandAllocations.hs +1/−2
• src/Futhark/Pass/ExplicitAllocations.hs +1/−1
• src/Futhark/Pass/ExtractKernels.hs +4/−6
• src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +12/−9
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +22/−14
• src/Futhark/Pass/ExtractKernels/Distribution.hs +14/−29
• src/Futhark/Pass/ExtractKernels/Interchange.hs +66/−14
• src/Futhark/Pass/ExtractKernels/Intragroup.hs +2/−2
• src/Futhark/Transform/CopyPropagate.hs +1/−1
• src/Futhark/Transform/FirstOrderTransform.hs +3/−5
• src/Futhark/TypeCheck.hs +19/−4
• src/Futhark/Util.hs +9/−2
• src/Language/Futhark/Interpreter.hs +12/−10
• src/Language/Futhark/Parser/Parser.y +0/−2
• src/Language/Futhark/Pretty.hs +1/−1
• src/Language/Futhark/Prop.hs +12/−1
• src/Language/Futhark/Query.hs +1/−1
• src/Language/Futhark/Syntax.hs +3/−1
• src/Language/Futhark/TypeChecker.hs +2/−2
• src/Language/Futhark/TypeChecker/Terms.hs +11/−14
• src/Language/Futhark/TypeChecker/Terms/Monad.hs +0/−10
• src/Language/Futhark/TypeChecker/Types.hs +17/−2
• unittests/Language/Futhark/TypeChecker/TypesTests.hs +15/−2

@@ -35,7 +35,7 @@ # ones. extensions = [     'sphinx.ext.todo',-    'sphinx.ext.mathjax',+    'sphinx.ext.mathjax' ]  # Add any paths that contain templates here, relative to this directory.@@ -83,7 +83,7 @@  # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files.-exclude_patterns = ['_build']+exclude_patterns = ['_build', 'lib']  # The reST default role (used for this markup: `text`) to use for all # documents.@@ -109,8 +109,8 @@     tokens = {         'root': [             (r'(if|then|else|let|loop|in|val|for|do|with|local|open|include|import|type|entry|module|while|module)\b', token.Keyword),-            (r"[a-zA-Z_][a-zA-Z0-9_']*", token.Name),-            (r"-- .*", token.Comment),+            (r"#?[a-zA-Z_][a-zA-Z0-9_']*", token.Name),+            (r"--.*", token.Comment),             (r'.', token.Text)         ]     }

@@ -33,6 +33,7 @@    c-api.rst    js-api.rst    package-management.rst+   performance.rst    error-index.rst    server-protocol.rst    c-porting-guide.rst

@@ -14,14 +14,14 @@ DESCRIPTION =========== -``futhark-autotune`` attemps to find optimal values for threshold+``futhark autotune`` attemps to find optimal values for threshold parameters given representative datasets.  This is done by repeatedly running running the program through :ref:`futhark-bench(1)` with-different values for the threshold parameters.  When-``futhark-autotune`` finishes tuning a program ``foo.fut``, the-results are written to ``foo.fut.tuning``, which will then-automatically be picked up by subsequent uses of-:ref:`futhark-bench(1)` and :ref:`futhark-test(1)`.+different values for the threshold parameters.  When ``futhark+autotune`` finishes tuning a program ``foo.fut``, the results are+written to ``foo.fut.tuning``, which will then automatically be picked+up by subsequent uses of :ref:`futhark-bench(1)` and+:ref:`futhark-test(1)`.   OPTIONS

@@ -0,0 +1,424 @@+.. _performance:++Writing Fast Futhark Programs+=============================++This document contains tips, tricks, and hints for writing efficient+Futhark code.  Ideally you'd need to know nothing more than an+abstract cost model, but sometimes it is useful to have an idea of how+the compiler will transform your program, what values look like in+memory, and what kind of code the compiler will generate for you.+These details are documented below.  Don't be discouraged by the+complexities mentioned here - most Futhark programs are written+without worrying about any of these details, and they still manage to+run with good performance.  This document focuses on corner cases and+pitfalls, which easily makes for depressing reading.++Parallelism+-----------++The Futhark compiler only generates parallel code for explicitly+parallel constructs such as ``map`` and ``reduce``.  A plain ``loop``+will *not* result in parallel code (unless the loop body itself+contains parallel operations).  The most important parallel constructs+are the *second-order array combinators* (SOACs) such as ``map`` and+``reduce``, but functions such as ``copy`` are also parallel.++When describing the asymptotic cost of a Futhark function, it is not+enough to give a traditional big-O measure of the total amount of+work.  Both ``foldl`` and ``reduce`` involve *O(n)* work, where *n* is+the size of the input array, but ``foldl`` is sequential while+``reduce`` is parallel, and this is an important distinction.  To make+this distinction, each function is described by *two* costs: the+*work*, which is the total amount of operations, and the *span*+(sometimes called *depth*) which is intuitively the "longest chain of+sequential dependencies".  We say that ``foldl`` has span *O(n)*,+while ``reduce`` has span *O(log(n))*.  This explains that+``reduce`` is more parallel than ``foldl``.  The documentation for a+Futhark function should mention both its work and span.  `See this+<https://sigkill.dk/writings/par/cost.html>`_ for more details on+parallel cost models and pointers to literature.++Scans and reductions+~~~~~~~~~~~~~~~~~~~~++The ``scan`` and ``reduce`` SOACs are rather inefficient when their+operators are on arrays.  If possible, use tuples instead (see+:ref:`performance-small-arrays`).  The one exception is when the+operator is a ``map2`` or equivalent.  Example::++  reduce (map2 (+)) (replicate n 0) xss++Such "vectorised" operators are typically handled quite efficiently.+Although to be on the safe side, you can rewrite the above by+interchanging the ``reduce`` and ``map``::++  map (reduce (+) 0) (transpose xss)++Avoid reductions over tiny arrays, e.g. ``reduce (+) 0 [x,y,z]``.  In+such cases the compiler will generate complex code to exploit a+negligible amount of parallelism.  Instead, just unroll the loop+manually (``x+y+z``) or perhaps use ``foldl (+) 0 [x,z,y]``, which+produces a sequential loop.++Histograms+~~~~~~~~~~++The ``reduce_by_index`` construct ("generalised histogram") has a+clever and adaptive implementation that handles multiple updates of+the same bin efficiently.  Its main weakness is when computing a very+large histogram (many millions of bins) where only a tiny fraction of+the bins are updated.  This is because the main mechanism for+optimising conflicts is by duplicating the histogram in memory, but+this is not efficient when it is very large.  If you know your program+involves such a situation, it may be better to implement the histogram+operation by sorting and then performing an irregular segmented+reduction.++Particularly with the GPU backends, ``reduce_by_index`` is much faster+when the operator involves a single 32-bit or 64-bit value.  Even if+you really want an 8-bit or 16-bit result, it may be faster to compute+it with a 32-bit or 64-bit type and manually mask off the excess bits.++Nested parallelism+~~~~~~~~~~~~~~~~~~++Futhark allows nested parallelism, understood as a parallel construct+used inside some other parallel construct.  The simplest example is+nested SOACs.  Example::++  map (\xs -> reduce (+) 0 xs) xss++Nested parallelism is allowed and encouraged, but its compilation to+efficient code is rather complicated, depending on the compiler+backend that is used.  The problem is that sometimes exploiting all+levels of parallelism is not optimal, yet how much to exploit depends+on run-time information that is not available to the compiler.++Sequential backends+!!!!!!!!!!!!!!!!!!!++The sequential backends are straightforward: all parallel operations+are compiled into sequential loops.  Due to Futhark's low-overhead+data representation (see below), this is often surprisingly efficient.++Multicore backend+!!!!!!!!!!!!!!!!!++Whenever the multicore backend encounters nested parallelism, it+generates two code versions: one where the nested parallel constructs+are also parallelised (possibly recursively involving further nested+parallelism), and one where they are turned into sequential loops.  At+runtime, based on the amount of work available and self-tuning+heuristics, the scheduler picks the version that it believes best+balances overhead with exploitation of parallelism.++GPU backends+!!!!!!!!!!!!++The GPU backends handle parallelism through an elaborate program+transformation called *incremental flattening*.  The full details are+beyond the scope of this document, but some properties are useful to+know of.  `See this paper+<https://futhark-lang.org/publications/ppopp19.pdf>`_ for more+details.++The main restriction is that the GPU backends can only handle+*regular* nested parallelism, meaning that the sizes of inner parallel+dimensions are invariant to the outer parallel dimensions.  For+example, this expression contains *irregular* nested parallelism::++  map (\i -> reduce (+) 0 (iota i)) is++This is because the size of the nested parallel construct is ``i``,+and ``i`` has a different value for every iteration of the outer+``map``.  The Futhark compiler will currently turn the irregular+constructs (here, the ``reduce``) into a sequential loop.  Depending+on how complicated the irregularity is, it may even refuse to generate+code entirely.++Incremental flattening is based on generating multiple code versions+to cater to different classes of datasets.  At run-time, one of these+versions will be picked for execution by comparing properties of the+input (its size) with a *threshold parameter*.  These threshold+parameters have sensible defaults, but for optimal performance, they+can be tuned with :ref:`futhark-autotune(1)`.++Value Representation+--------------------++The compiler discards all type abstraction when compiling.  Using the+module system to make a type abstract causes no run-time overhead.++Scalars+~~~~~~~++Scalar values (``i32``, ``f64``, ``bool``, etc) are represented as+themselves.  The internal representation does not distinguish signs,+so ``i32`` and ``u32`` have the same representation, and converting+between types that differ only in sign is free.++Tuples+~~~~~~++Tuples are flattened and then represented directly by their individual+components - there are no *tuple objects* at runtime.  A function that+takes an argument of type ``(f64,f64)`` corresponds to a C function+that takes two arguments of type ``double``.  This has one performance+implication: whenever you pass a tuple to a function, the *entire*+tuple is copied (except any embedded arrays, which are always passed+by reference, see below).  Due to the compiler's heavy use of+inlining, this is rarely a problem in practice, but it can be a+concern when using the ``loop`` construct with a large tuple as the+loop variant parameter.++Records+~~~~~~~++Records are turned into tuples by simply sorting their fields and+discarding the labels.  This means there is no overhead to using a+record compared to using a tuple.++Sum Types+~~~~~~~~~++A starting point, a sum type is turned into a tuple containing all the+payload components in order, prefixed with an `i8` tag to identify the+constructor.  For example,++.. code-block:: futhark++   #foo i32 bool | #bar i32++would be represented as a tuple of type++.. code-block:: futhark++   (i8, i32, bool, i32)++where the value++.. code-block:: futhark++   #foo 42 false++is represented as++.. code-block:: futhark++   (1, 42, false, 0)++where ``#foo`` is assigned the tag ``1`` because it is alphabetically+after ``#bar``.++To shrink the tuples, if multiple constructors have payload elements+of the *same* type, the compiler assigns them to the same elements in+the result tuple. The representation of the above sum type is actually+the following:++.. code-block:: futhark++   (i8, i32, bool)++The types must be the *same* for deduplication to take place - despite+`i32` and `f32` being of the same size, they cannot be assigned the+same tuple element.  This means that the type++.. code-block:: futhark++   #foo [n]i32 | #bar [n]i32++is efficiently represented as++.. code-block:: futhark++   (u8, [n]i32)+++.. code-block:: futhark++   #foo [n]i32 | #bar [n]f32++becomes++.. code-block:: futhark++   (u8, [n]i32, [n]f32)++which is not great.  Take caution when you use sum types with large+arrays in their payloads.++Functions+~~~~~~~~~++Higher-order functions are implemented via defunctionalisation.  At+run-time, they are represented by a record containing their lexical+closure.  Since the type system forbids putting functions in arrays,+this is essentially a constant cost, and not worth worrying about.++Arrays+~~~~~~++Arrays are the only Futhark values that are boxed - that is, are+stored on the heap.++The elements of an array are unboxed, stored adjacent to each other in+memory.  There is zero memory overhead except for the minuscule amount+needed to track the shape of the array.++Multidimensional arrays+!!!!!!!!!!!!!!!!!!!!!!!++At the surface language level, Futhark may appear to support "arrays+of arrays", and this is indeed a convenient aspect of its programming+model, but at runtime multi-dimensional arrays are stored in flattened+form.  A value of type ``[x][y]i32`` is laid out in memory simply as+one array containing *x\*y* integers.  This means that constructing an+array ``[x,y,x]`` can be (relatively) expensive if ``x``, ``y``, ``z``+are themselves large arrays, as they must be copied in their entirety.++Since arrays cannot contain other arrays, memory management only has+to be concerned with one level of indirection.  In practice, it means+that Futhark can use straightforward reference counting to keep track+of when to free the memory backing an array, as circular references+are not possible.  Further, since arrays tend to be large and+relatively few in number, the usual performance impact of naive+reference counting is not present.++Arrays of tuples+!!!!!!!!!!!!!!!!++For arrays of tuples, Futhark uses the so-called `structure of arrays+<https://en.wikipedia.org/wiki/AoS_and_SoA>`_ representation.  In+Futhark terms, an array ``[n](a,b,c)`` is at runtime represented as+the tuple ``([n]a,[n]b,[n]c)``.  This means that the final memory+representation always consists of arrays of scalars.++This has some significant implications.  For example, ``zip`` and+``unzip`` are very cheap, as the actual runtime representation is in+always "unzipped", so these functions don't actually have to do+anything.++Since records and sum types are represented as tuples, this also+explains how arrays of these are represented.++Element order+!!!!!!!!!!!!!++The exact in-memory element ordering is up to the compiler, and+depends on how the array is constructed and how it is used.  Absent+any other information, Futhark represents multidimensional arrays in+row-major order.  However, depending on how the array is traversed,+the compiler may insert code to represent it in some other order.  For+particularly tricky programs, an array may even be duplicated in+memory, represented in different ways, to ensure efficient traversal.+This means you should normally *not* worry about how to represent your+arrays to ensure coalesced access on GPUs or similar.  That is the+compiler's job.++Crucial Optimisations+---------------------++Some of the optimisations done by the Futhark compiler are important,+complex, or subtle enough that it may be useful to know how they work,+and how to write code that caters to their quirks.++Fusion+~~~~~~++Futhark performs fusion of SOACs.  For an expression ``map f (map g+A)``, then the compiler will optimise this into a single ``map`` with+the composition of ``f`` and ``g``, which prevents us from storing an+intermediate array in memory.  This is called *vertical fusion* or+*producer-consumer fusion*.  In this case the *producer* is ``map g``+and the *consumer* is ``map f``.++Fusion does not depend on the expressions being adjacent as in this+example, as the optimisation is performed on a data dependency graph+representing the program.  This means that you can decompose your+programs into many small parallel operations without worrying about+the overhead, as the compiler will fuse them together automatically.++Not all producer-consumer relationships between SOACs can be fused.+Generally, ``map`` can always be fused as a producer, but ``scan``,+``reduce``, and similar SOACs can only act as consumers.++*Horizontal fusion* occurs when two SOACs take as input the same+array, but are not themselves in a producer-consumer relationship.+Example::++   (map f xs, map g xs)++Such cases are fused into a single operation that traverses ``xs``+just once.  More than two SOACs can be involved in horizontal fusion,+and they need not be of the same kind (e.g. one could be a ``map`` and+the other a ``reduce``).++Free Operations+---------------++Some operations such as array slicing, ``take``, ``drop``,+``transpose`` and ``reverse`` are "free" in the sense that they merely+return a different view of some underlying array.  In most cases they+have constant cost, no matter the size of the array they operate on.+This is because they are index space transformations that simply+result in different code being generated when the arrays are+eventually used.++However, there are some cases where the compiler is forced to manifest+such a "view" as an actual array in memory, which involves a full+copy.  An incomplete list follows:++* Any array returned by an entry point is converted to row-major+  order.++* An array returned by an ``if`` branch may be copied if its+  representation is substantially different from that of the other+  branch.++* An array returned by a ``loop`` body may be copied if its+  representation is substantially different from that of the initial+  loop values.++* An array is copied whenever it becomes the element of another+  multidimensional array.  This is most obviously the case for array+  literals (``[x,y,z]``), but also for ``map`` expressions where the+  mapped function returns an array.++Note that this notion of "views" is not part of the Futhark type+system - it is merely an implementation detail.  Strictly speaking,+all functions that return an array (e.g. ``reverse``) should be+considered to have a cost proportional to the size of the array, even+if that cost will almost never actually be paid at run-time.  If you+want to be sure no copy takes place, it may be better to explicitly+maintain tuples of indexes into some other array.++.. _performance-small-arrays:++Small Arrays+------------++The compiler assumes arrays are "large", which for example means that+operations across them are worth parallelising.  It also means they+are boxed and heap-allocated, even when the size is a small constant.+This can cause unexpectedly bad performance when using small+constant-size arrays (say, five elements or less).  Consider using+tuples or records instead.  `This post+<https://futhark-lang.org/blog/2019-01-13-giving-programmers-what-they-want.html>`_+contains more information on how and why.  If in doubt, try both and+measure which is faster.++Inlining+--------++The compiler currently inlines all functions at their call site,+unless they have been marked with the ``noinline`` attribute (see+:ref:`attributes`).  This can lead to code explosion, which mostly+results in slow compile times, but can also affect run-time+performance.  In many cases this is currently unavoidable, but+sometimes the program can be rewritten such that instead of calling+the same function in multiple places, it is called in a single place,+in a loop.  E.g. we might rewrite ``f x (f y (f z v))`` as::++  loop acc = v for a in [z,y,x] do+    f a acc

@@ -1,2 +1,2 @@ pyyaml>=4.2b1-sphinx==3.0.3+sphinx>=4.2.0

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.20.6+version:        0.20.7 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to

@@ -16,6 +16,10 @@ -- | Create integer from double-precision float. let t64 (x: f64): i32 = i32.f64 x +-- | Negate a boolean.  `not x` is the same as `!x`.  This function is+-- mostly useful for passing to `map+let not (x: bool): bool = !x+ -- | Semantically just identity, but serves as an optimisation -- inhibitor.  The compiler will treat this function as a black box. -- You can use this to work around optimisation deficiencies (or

@@ -104,10 +104,10 @@   return signed(max(unsigned(x),unsigned(y)))  def fminN(x,y):-  return min(x,y)+  return np.fmin(x,y)  def fmaxN(x,y):-  return max(x,y)+  return np.fmax(x,y)  def powN(x,y):   return x ** y

@@ -621,7 +621,7 @@         out.write("%di32" % v)     elif type(v) == np.int64:         out.write("%di64" % v)-    elif type(v) in [np.bool, np.bool_]:+    elif type(v) in [bool, np.bool_]:         if v:             out.write("true")         else:

@@ -45,14 +45,27 @@ parens = between (lexeme "(") (lexeme ")")  -- | Parse a 'PrimExp' given a leaf parser.-pPrimExp :: Parser (v, PrimType) -> Parser (PrimExp v)-pPrimExp pLeaf =+pPrimExp :: PrimType -> Parser v -> Parser (PrimExp v)+pPrimExp t pLeaf =   choice-    [ uncurry LeafExp <$> pLeaf,+    [ flip LeafExp t <$> pLeaf,       ValueExp <$> pPrimValue,-      BinOpExp <$> pBinOp <*> pPrimExp pLeaf <*> pPrimExp pLeaf,-      CmpOpExp <$> pCmpOp <*> pPrimExp pLeaf <*> pPrimExp pLeaf,-      ConvOpExp <$> pConvOp <*> pPrimExp pLeaf,-      UnOpExp <$> pUnOp <*> pPrimExp pLeaf,-      parens $ pPrimExp pLeaf+      pBinOp >>= binOpExp,+      pCmpOp >>= cmpOpExp,+      pConvOp >>= convOpExp,+      pUnOp >>= unOpExp,+      parens $ pPrimExp t pLeaf     ]+  where+    binOpExp op =+      BinOpExp op+        <$> pPrimExp (binOpType op) pLeaf+        <*> pPrimExp (binOpType op) pLeaf+    cmpOpExp op =+      CmpOpExp op+        <$> pPrimExp (cmpOpType op) pLeaf+        <*> pPrimExp (cmpOpType op) pLeaf+    convOpExp op =+      ConvOpExp op <$> pPrimExp (fst (convOpType op)) pLeaf+    unOpExp op =+      UnOpExp op <$> pPrimExp (unOpType op) pLeaf

@@ -19,6 +19,7 @@     sizeUsage,     sizeUsages,     Usages,+    consumedU,     usageInStm,   ) where

@@ -271,8 +271,9 @@     [dest]       | isBuiltInFunction fname ->         stm [C.cstm|$id:dest = $id:(funName fname)($args:args');|]-    _ ->-      item [C.citem|if ($id:(funName fname)($args:args') != 0) { err = 1; goto cleanup; }|]+    _ -> do+      free_all_mem <- freeAllocatedMem+      item [C.citem|if ($id:(funName fname)($args:args') != 0) { $items:free_all_mem err = 1; goto cleanup; }|]  -- | A set of operations that fail for every operation involving -- non-default memory spaces.  Uses plain pointers and @malloc@ for@@ -1902,7 +1903,7 @@   return $ derefPointer dest i [C.cty|$tyquals:quals' $ty:elemtype*|]  compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName-compileExpToName _ _ (LeafExp (ScalarVar v) _) =+compileExpToName _ _ (LeafExp v _) =   return v compileExpToName desc t e = do   desc' <- newVName desc@@ -1911,29 +1912,7 @@   return desc'  compileExp :: Exp -> CompilerM op s C.Exp-compileExp = compilePrimExp compileLeaf-  where-    compileLeaf (ScalarVar src) =-      return [C.cexp|$id:src|]-    compileLeaf (Index _ _ Unit __ _) =-      pure $ C.toExp UnitValue mempty-    compileLeaf (Index src (Count iexp) restype DefaultSpace vol) = do-      src' <- rawMem src-      fmap (fromStorage restype) $-        derefPointer src'-          <$> compileExp (untyped iexp)-          <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType restype)*|]-    compileLeaf (Index src (Count iexp) restype (Space space) vol) =-      fmap (fromStorage restype) . join $-        asks envReadScalar-          <*> rawMem src-          <*> compileExp (untyped iexp)-          <*> pure (primStorageType restype)-          <*> pure space-          <*> pure vol-    compileLeaf (Index src (Count iexp) _ ScalarSpace {} _) = do-      iexp' <- compileExp $ untyped iexp-      return [C.cexp|$id:src[$exp:iexp']|]+compileExp = compilePrimExp $ \v -> pure [C.cexp|$id:v|]  -- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. compilePrimExp :: Monad m => (v -> m C.Exp) -> PrimExp v -> m C.Exp@@ -2146,6 +2125,29 @@       <*> pure space       <*> pure vol       <*> (toStorage elemtype <$> compileExp elemexp)+compileCode (Read x _ _ Unit __ _) =+  stm [C.cstm|$id:x = $exp:(UnitValue);|]+compileCode (Read x src (Count iexp) restype DefaultSpace vol) = do+  src' <- rawMem src+  e <-+    fmap (fromStorage restype) $+      derefPointer src'+        <$> compileExp (untyped iexp)+        <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType restype)*|]+  stm [C.cstm|$id:x = $exp:e;|]+compileCode (Read x src (Count iexp) restype (Space space) vol) = do+  e <-+    fmap (fromStorage restype) . join $+      asks envReadScalar+        <*> rawMem src+        <*> compileExp (untyped iexp)+        <*> pure (primStorageType restype)+        <*> pure space+        <*> pure vol+  stm [C.cstm|$id:x = $exp:e;|]+compileCode (Read x src (Count iexp) _ ScalarSpace {} _) = do+  iexp' <- compileExp $ untyped iexp+  stm [C.cstm|$id:x = $id:src[$exp:iexp'];|] compileCode (DeclareMem name space) =   declMem name space compileCode (DeclareScalar name vol t) = do@@ -2178,7 +2180,7 @@ -- For assignments of the form 'x = x OP e', we generate C assignment -- operators to make the resulting code slightly nicer.  This has no -- effect on performance.-compileCode (SetScalar dest (BinOpExp op (LeafExp (ScalarVar x) _) y))+compileCode (SetScalar dest (BinOpExp op (LeafExp x _) y))   | dest == x,     Just f <- assignmentOperator op = do     y' <- compileExp y

@@ -1172,24 +1172,7 @@   simpleCall (futharkFun (pretty h)) <$> mapM (compilePrimExp f) args  compileExp :: Imp.Exp -> CompilerM op s PyExp-compileExp = compilePrimExp compileLeaf-  where-    compileLeaf (Imp.ScalarVar vname) =-      compileVar vname-    compileLeaf (Imp.Index _ _ Unit _ _) =-      return $ compilePrimValue UnitValue-    compileLeaf (Imp.Index src (Imp.Count iexp) restype (Imp.Space space) _) =-      join $-        asks envReadScalar-          <*> compileVar src-          <*> compileExp (Imp.untyped iexp)-          <*> pure restype-          <*> pure space-    compileLeaf (Imp.Index src (Imp.Count iexp) bt _ _) = do-      iexp' <- compileExp $ Imp.untyped iexp-      let bt' = compilePrimType bt-      src' <- compileVar src-      return $ fromStorage bt $ simpleCall "indexArray" [src', iexp', Var bt']+compileExp = compilePrimExp compileVar  errorMsgString :: Imp.ErrorMsg Imp.Exp -> CompilerM op s (String, [PyExp]) errorMsgString (Imp.ErrorMsg parts) = do@@ -1351,4 +1334,22 @@   elemexp' <- toStorage elemtype <$> compileExp elemexp   dest' <- compileVar dest   stm $ Exp $ simpleCall "writeScalarArray" [dest', idx', elemexp']+compileCode (Imp.Read x _ _ Unit _ _) =+  stm =<< Assign <$> compileVar x <*> pure (compilePrimValue UnitValue)+compileCode (Imp.Read x src (Imp.Count iexp) restype (Imp.Space space) _) = do+  x' <- compileVar x+  e <-+    join $+      asks envReadScalar+        <*> compileVar src+        <*> compileExp (Imp.untyped iexp)+        <*> pure restype+        <*> pure space+  stm $ Assign x' e+compileCode (Imp.Read x src (Imp.Count iexp) bt _ _) = do+  x' <- compileVar x+  iexp' <- compileExp $ Imp.untyped iexp+  let bt' = compilePrimType bt+  src' <- compileVar src+  stm $ Assign x' $ fromStorage bt $ simpleCall "indexArray" [src', iexp', Var bt'] compileCode Imp.Skip = return ()

@@ -27,15 +27,11 @@     SpaceId,     Code (..),     PrimValue (..),-    ExpLeaf (..),     Exp,     TExp,     Volatility (..),     Arg (..),     var,-    vi32,-    vi64,-    index,     ErrorMsg (..),     ErrorMsgPart (..),     errorMsgArgTypes,@@ -56,6 +52,7 @@     module Language.Futhark.Core,     module Futhark.IR.Primitive,     module Futhark.Analysis.PrimExp,+    module Futhark.Analysis.PrimExp.Convert,     module Futhark.IR.GPU.Sizes,     module Futhark.IR.Prop.Names,   )@@ -66,6 +63,7 @@ import qualified Data.Set as S import Data.Traversable import Futhark.Analysis.PrimExp+import Futhark.Analysis.PrimExp.Convert import Futhark.IR.GPU.Sizes (Count (..)) import Futhark.IR.Pretty () import Futhark.IR.Primitive@@ -244,10 +242,14 @@     -- @mem@ offset by @i@ elements of type @t@.  The     -- 'Space' argument is the memory space of @mem@     -- (technically redundant, but convenient).  Note that-    -- /reading/ is done with an 'Exp' ('Index').+    -- /reading/ is done with an 'Exp' ('Read').     Write VName (Count Elements (TExp Int64)) PrimType Space Volatility Exp   | -- | Set a scalar variable.     SetScalar VName Exp+  | -- | Read a scalar from memory from memory.  The first 'VName' is+    -- the target scalar variable, and the remaining arguments have+    -- the same meaning as with 'Write'.+    Read VName VName (Count Elements (TExp Int64)) PrimType Space Volatility   | -- | Must be in same space.     SetMem VName VName Space   | -- | Function call.  The results are written to the@@ -338,22 +340,12 @@ calledFuncs (Call _ f _) = S.singleton f calledFuncs _ = mempty --- | The leaves of an 'Exp'.-data ExpLeaf-  = -- | A scalar variable.  The type is stored in the-    -- 'LeafExp' constructor itself.-    ScalarVar VName-  | -- | Reading a value from memory.  The arguments have-    -- the same meaning as with 'Write'.-    Index VName (Count Elements (TExp Int64)) PrimType Space Volatility-  deriving (Eq, Show)- -- | A side-effect free expression whose execution will produce a -- single primitive value.-type Exp = PrimExp ExpLeaf+type Exp = PrimExp VName  -- | Like 'Exp', but with a required/known type.-type TExp t = TPrimExp t ExpLeaf+type TExp t = TPrimExp t VName  -- | A function call argument. data Arg@@ -380,21 +372,9 @@ withElemType :: Count Elements (TExp Int64) -> PrimType -> Count Bytes (TExp Int64) withElemType (Count e) t = bytes $ sExt64 e * primByteSize t --- | Turn a 'VName' into a 'Imp.ScalarVar'.+-- | Turn a 'VName' into a 'Exp'. var :: VName -> PrimType -> Exp-var = LeafExp . ScalarVar---- | Turn a 'VName' into a v'Int32' 'Imp.ScalarVar'.-vi32 :: VName -> TExp Int32-vi32 = TPrimExp . flip var (IntType Int32)---- | Turn a 'VName' into a v'Int64' 'Imp.ScalarVar'.-vi64 :: VName -> TExp Int64-vi64 = TPrimExp . flip var (IntType Int64)---- | Concise wrapper for using 'Index'.-index :: VName -> Count Elements (TExp Int64) -> PrimType -> Space -> Volatility -> Exp-index arr i t s vol = LeafExp (Index arr i t s vol) t+var = LeafExp  -- Prettyprinting definitions. @@ -495,6 +475,13 @@       vol' = case vol of         Volatile -> text "volatile "         Nonvolatile -> mempty+  ppr (Read name v is bt space vol) =+    ppr name <+> text "<-"+      <+> ppr v <> langle <> vol' <> ppr bt <> ppr space <> rangle <> brackets (ppr is)+    where+      vol' = case vol of+        Volatile -> text "volatile "+        Nonvolatile -> mempty   ppr (SetScalar name val) =     ppr name <+> text "<-" <+> ppr val   ppr (SetMem dest from space) =@@ -533,16 +520,6 @@   ppr (MemArg m) = ppr m   ppr (ExpArg e) = ppr e -instance Pretty ExpLeaf where-  ppr (ScalarVar v) =-    ppr v-  ppr (Index v is bt space vol) =-    ppr v <> langle <> vol' <> ppr bt <> ppr space <> rangle <> brackets (ppr is)-    where-      vol' = case vol of-        Volatile -> text "volatile "-        Nonvolatile -> mempty- instance Functor Functions where   fmap = fmapDefault @@ -598,6 +575,8 @@     pure $ Copy dest destoffset destspace src srcoffset srcspace size   traverse _ (Write name i bt val space vol) =     pure $ Write name i bt val space vol+  traverse _ (Read x name i bt space vol) =+    pure $ Read x name i bt space vol   traverse _ (SetScalar name val) =     pure $ SetScalar name val   traverse _ (SetMem dest from space) =@@ -669,6 +648,8 @@     freeIn' x <> freeIn' y   freeIn' (Write v i _ _ _ e) =     freeIn' v <> freeIn' i <> freeIn' e+  freeIn' (Read x v i _ _ _) =+    freeIn' x <> freeIn' v <> freeIn' i   freeIn' (SetScalar x y) =     freeIn' x <> freeIn' y   freeIn' (Call dests _ args) =@@ -685,10 +666,6 @@     maybe mempty freeIn' v   freeIn' (TracePrint msg) =     foldMap freeIn' msg--instance FreeIn ExpLeaf where-  freeIn' (Index v e _ _ _) = freeIn' v <> freeIn' e-  freeIn' (ScalarVar v) = freeIn' v  instance FreeIn Arg where   freeIn' (MemArg m) = freeIn' m

@@ -136,9 +136,9 @@         )  -- | When we do a barrier or fence, is it at the local or global--- level?+-- level?  By the 'Ord' instance, global is greater than local. data Fence = FenceLocal | FenceGlobal-  deriving (Show)+  deriving (Show, Eq, Ord)  -- | An operation that occurs within a kernel body. data KernelOp

@@ -194,7 +194,7 @@       opsAllocCompilers = mempty     } --- | When an array is dared, this is where it is stored.+-- | When an array is declared, this is where it is stored. data MemLoc = MemLoc   { memLocName :: VName,     memLocShape :: [Imp.DimSize],@@ -499,14 +499,11 @@   ImpM rep r op (Either Imp.Param ArrayDecl) compileInParam fparam = case paramDec fparam of   MemPrim bt ->-    return $ Left $ Imp.ScalarParam name bt+    pure $ Left $ Imp.ScalarParam name bt   MemMem space ->-    return $ Left $ Imp.MemParam name space+    pure $ Left $ Imp.MemParam name space   MemArray bt shape _ (ArrayIn mem ixfun) ->-    return $-      Right $-        ArrayDecl name bt $-          MemLoc mem (shapeDims shape) $ fmap (fmap Imp.ScalarVar) ixfun+    pure $ Right $ ArrayDecl name bt $ MemLoc mem (shapeDims shape) ixfun   MemAcc {} ->     error "Functions may not have accumulator parameters."   where@@ -787,7 +784,7 @@     ForLoop i _ bound loopvars -> do       let setLoopParam (p, a)             | Prim _ <- paramType p =-              copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.vi64 i]+              copyDWIM (paramName p) [] (Var a) [DimFix $ Imp.le64 i]             | otherwise =               return () @@ -901,7 +898,7 @@   | otherwise = do     ds' <- mapM toExp ds     is <- replicateM (length ds) (newVName "i")-    copy_elem <- collect $ copyDWIM (patElemName pe) (map (DimFix . Imp.vi64) is) se []+    copy_elem <- collect $ copyDWIM (patElemName pe) (map (DimFix . Imp.le64) is) se []     emit $ foldl (.) id (zipWith Imp.For is ds') copy_elem defCompileBasicOp _ Scratch {} =   return ()@@ -1091,7 +1088,7 @@ memBoundToVarEntry e (MemAcc acc ispace ts _) =   AccVar e (acc, ispace, ts) memBoundToVarEntry e (MemArray bt shape _ (ArrayIn mem ixfun)) =-  let location = MemLoc mem (shapeDims shape) $ fmap (fmap Imp.ScalarVar) ixfun+  let location = MemLoc mem (shapeDims shape) ixfun    in ArrayVar         e         ArrayEntry@@ -1134,10 +1131,12 @@   ImpM rep r op () dScope e = mapM_ (uncurry $ dInfo e) . M.toList -dArray :: VName -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM rep r op ()-dArray name bt shape membind =-  addVar name $-    memBoundToVarEntry Nothing $ MemArray bt shape NoUniqueness membind+dArray :: VName -> PrimType -> ShapeBase SubExp -> VName -> IxFun -> ImpM rep r op ()+dArray name pt shape mem ixfun =+  addVar name $ ArrayVar Nothing $ ArrayEntry location pt+  where+    location =+      MemLoc mem (shapeDims shape) ixfun  everythingVolatile :: ImpM rep r op a -> ImpM rep r op a everythingVolatile = local $ \env -> env {envVolatility = Imp.Volatile}@@ -1207,8 +1206,8 @@   toExp' t (Var v) = Imp.var v t  instance ToExp (PrimExp VName) where-  toExp = pure . fmap Imp.ScalarVar-  toExp' _ = fmap Imp.ScalarVar+  toExp = pure+  toExp' _ = id  addVar :: VName -> VarEntry rep -> ImpM rep r op () addVar name entry =@@ -1489,14 +1488,18 @@ copyElementWise bt dest src = do   let bounds = IxFun.shape $ memLocIxFun src   is <- replicateM (length bounds) (newVName "i")-  let ivars = map Imp.vi64 is+  let ivars = map Imp.le64 is   (destmem, destspace, destidx) <- fullyIndexArray' dest ivars   (srcmem, srcspace, srcidx) <- fullyIndexArray' src ivars   vol <- asks envVolatility+  tmp <- newVName "tmp"   emit $     foldl (.) id (zipWith Imp.For is $ map untyped bounds) $-      Imp.Write destmem destidx bt destspace vol $-        Imp.index srcmem srcidx bt srcspace vol+      mconcat+        [ Imp.DeclareScalar tmp vol bt,+          Imp.Read tmp srcmem srcidx bt srcspace vol,+          Imp.Write destmem destidx bt destspace vol $ Imp.var tmp bt+        ]  -- | Copy from here to there; both destination and source may be -- indexeded.@@ -1522,9 +1525,10 @@       (srcmem, srcspace, srcoffset) <-         fullyIndexArray' srclocation srcis       vol <- asks envVolatility-      return $-        Imp.Write targetmem targetoffset bt destspace vol $-          Imp.index srcmem srcoffset bt srcspace vol+      collect $ do+        tmp <- tvVar <$> dPrim "tmp" bt+        emit $ Imp.Read tmp srcmem srcoffset bt srcspace vol+        emit $ Imp.Write targetmem targetoffset bt destspace vol $ Imp.var tmp bt     | otherwise = do       let destslice' = fullSliceNum (map toInt64Exp destshape) destslice           srcslice' = fullSliceNum (map toInt64Exp srcshape) srcslice@@ -1609,7 +1613,7 @@         (mem, space, i) <-           fullyIndexArray' (entryArrayLoc arr) src_is         vol <- asks envVolatility-        emit $ Imp.SetScalar name $ Imp.index mem i bt space vol+        emit $ Imp.Read name mem i bt space vol       | otherwise ->         error $           unwords@@ -1778,18 +1782,17 @@   sAlloc_ name' size space   return name' -sArray :: String -> PrimType -> ShapeBase SubExp -> MemBind -> ImpM rep r op VName-sArray name bt shape membind = do+sArray :: String -> PrimType -> ShapeBase SubExp -> VName -> IxFun -> ImpM rep r op VName+sArray name bt shape mem ixfun = do   name' <- newVName name-  dArray name' bt shape membind+  dArray name' bt shape mem ixfun   return 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 pt shape mem =-  sArray name pt shape $-    ArrayIn mem $-      IxFun.iota $ map (isInt64 . primExpFromSubExp int64) $ shapeDims shape+  sArray name pt shape mem $+    IxFun.iota $ map (isInt64 . primExpFromSubExp int64) $ shapeDims shape  -- | Like 'sAllocArray', but permute the in-memory representation of the indices as specified. sAllocArrayPerm :: String -> PrimType -> ShapeBase SubExp -> Space -> [Int] -> ImpM rep r op VName@@ -1797,8 +1800,8 @@   let permuted_dims = rearrangeShape perm $ shapeDims shape   mem <- sAlloc (name ++ "_mem") (typeSize (Array pt shape NoUniqueness)) space   let iota_ixfun = IxFun.iota $ map (isInt64 . primExpFromSubExp int64) permuted_dims-  sArray name pt shape $-    ArrayIn mem $ IxFun.permute iota_ixfun $ rearrangeInverse perm+  sArray name pt shape mem $+    IxFun.permute iota_ixfun $ rearrangeInverse perm  -- | Uses linear/iota index function. sAllocArray :: String -> PrimType -> ShapeBase SubExp -> Space -> ImpM rep r op VName@@ -1815,7 +1818,7 @@   mem <- newVNameForFun $ name ++ "_mem"   emit $ Imp.DeclareArray mem space pt vs   addVar mem $ MemVar Nothing $ MemEntry space-  sArray name pt shape $ ArrayIn mem $ IxFun.iota [fromIntegral num_elems]+  sArray name pt shape mem $ IxFun.iota [fromIntegral num_elems]  sWrite :: VName -> [Imp.TExp Int64] -> Imp.Exp -> ImpM rep r op () sWrite arr is v = do@@ -1891,7 +1894,7 @@   where     loop ((v, size) : rest) i = do       dPrimV_ v (i `quot` size)-      i' <- dPrimVE "remnant" $ i - Imp.vi64 v * size+      i' <- dPrimVE "remnant" $ i - Imp.le64 v * size       loop rest i'     loop _ _ = pure () @@ -1905,4 +1908,4 @@ dIndexSpace' desc ds j = do   ivs <- replicateM (length ds) (newVName desc)   dIndexSpace (zip ivs ds) j-  pure $ map Imp.vi64 ivs+  pure $ map Imp.le64 ivs

@@ -367,18 +367,18 @@     -- When an input array has either width==1 or height==1, performing a     -- transpose will be the same as performing a copy.     can_use_copy =-      let onearr = Imp.vi32 num_arrays .==. 1-          height_is_one = Imp.vi32 y .==. 1-          width_is_one = Imp.vi32 x .==. 1+      let onearr = Imp.le32 num_arrays .==. 1+          height_is_one = Imp.le32 y .==. 1+          width_is_one = Imp.le32 x .==. 1        in onearr .&&. (width_is_one .||. height_is_one)      transpose_code =       Imp.If input_is_empty mempty $         mconcat           [ Imp.DeclareScalar muly Imp.Nonvolatile (IntType Int32),-            Imp.SetScalar muly $ untyped $ block_dim `quot` Imp.vi32 x,+            Imp.SetScalar muly $ untyped $ block_dim `quot` Imp.le32 x,             Imp.DeclareScalar mulx Imp.Nonvolatile (IntType Int32),-            Imp.SetScalar mulx $ untyped $ block_dim `quot` Imp.vi32 y,+            Imp.SetScalar mulx $ untyped $ block_dim `quot` Imp.le32 y,             Imp.If can_use_copy copy_code $               Imp.If should_use_lowwidth (callTransposeKernel TransposeLowWidth) $                 Imp.If should_use_lowheight (callTransposeKernel TransposeLowHeight) $@@ -387,28 +387,28 @@           ]      input_is_empty =-      Imp.vi32 num_arrays .==. 0 .||. Imp.vi32 x .==. 0 .||. Imp.vi32 y .==. 0+      Imp.le32 num_arrays .==. 0 .||. Imp.le32 x .==. 0 .||. Imp.le32 y .==. 0      should_use_small =-      Imp.vi32 x .<=. (block_dim `quot` 2)-        .&&. Imp.vi32 y .<=. (block_dim `quot` 2)+      Imp.le32 x .<=. (block_dim `quot` 2)+        .&&. Imp.le32 y .<=. (block_dim `quot` 2)      should_use_lowwidth =-      Imp.vi32 x .<=. (block_dim `quot` 2)-        .&&. block_dim .<. Imp.vi32 y+      Imp.le32 x .<=. (block_dim `quot` 2)+        .&&. block_dim .<. Imp.le32 y      should_use_lowheight =-      Imp.vi32 y .<=. (block_dim `quot` 2)-        .&&. block_dim .<. Imp.vi32 x+      Imp.le32 y .<=. (block_dim `quot` 2)+        .&&. block_dim .<. Imp.le32 x      copy_code =-      let num_bytes = sExt64 $ Imp.vi32 x * Imp.vi32 y * primByteSize bt+      let num_bytes = sExt64 $ Imp.le32 x * Imp.le32 y * primByteSize bt        in Imp.Copy             destmem-            (Imp.Count $ sExt64 $ Imp.vi32 destoffset)+            (Imp.Count $ sExt64 $ Imp.le32 destoffset)             space             srcmem-            (Imp.Count $ sExt64 $ Imp.vi32 srcoffset)+            (Imp.Count $ sExt64 $ Imp.le32 srcoffset)             space             (Imp.Count num_bytes) @@ -418,14 +418,14 @@           (mapTransposeName bt)           block_dim_int           ( destmem,-            Imp.vi32 destoffset,+            Imp.le32 destoffset,             srcmem,-            Imp.vi32 srcoffset,-            Imp.vi32 x,-            Imp.vi32 y,-            Imp.vi32 mulx,-            Imp.vi32 muly,-            Imp.vi32 num_arrays,+            Imp.le32 srcoffset,+            Imp.le32 x,+            Imp.le32 y,+            Imp.le32 mulx,+            Imp.le32 muly,+            Imp.le32 num_arrays,             block           )           bt

@@ -36,7 +36,7 @@ where  import Control.Monad.Except-import Data.List (zip4)+import Data.List (foldl', zip4) import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Set as S@@ -350,10 +350,8 @@               locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness            locks_mem <- sAlloc "locks_mem" (typeSize locks_t) $ Space "local"-          dArray locks int32 (arrayShape locks_t) $-            ArrayIn locks_mem $-              IxFun.iota $-                map pe64 $ arrayDims locks_t+          dArray locks int32 (arrayShape locks_t) locks_mem $+            IxFun.iota $ map pe64 $ arrayDims locks_t            sComment "All locks start out unlocked" $             groupCoverSpace [kernelGroupSize constants] $ \is ->@@ -372,6 +370,21 @@       then m       else sWhen (isActive $ unSegSpace space) m +-- Which fence do we need to protect shared access to this memory space?+fenceForSpace :: Space -> Imp.Fence+fenceForSpace (Space "local") = Imp.FenceLocal+fenceForSpace _ = Imp.FenceGlobal++-- If we are touching these arrays, which kind of fence do we need?+fenceForArrays :: [VName] -> InKernelGen Imp.Fence+fenceForArrays = fmap (foldl' max Imp.FenceLocal) . mapM need+  where+    need arr =+      fmap (fenceForSpace . entryMemSpace) . lookupMemory+        . memLocName+        . entryArrayLoc+        =<< lookupArray arr+ compileGroupOp :: OpCompiler GPUMem KernelEnv Imp.KernelOp compileGroupOp pat (Alloc size space) =   kernelAlloc pat size space@@ -380,11 +393,10 @@ compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do   void $ compileGroupSpace lvl space -  whenActive lvl space $-    localOps threadOperations $-      compileStms mempty (kernelBodyStms body) $-        zipWithM_ (compileThreadResult space) (patElems pat) $-          kernelBodyResult body+  whenActive lvl space . localOps threadOperations $+    compileStms mempty (kernelBodyStms body) $+      zipWithM_ (compileThreadResult space) (patElems pat) $+        kernelBodyResult body    sOp $ Imp.ErrorSync Imp.FenceLocal compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do@@ -392,16 +404,17 @@   let (ltids, dims) = unzip $ unSegSpace space       dims' = map toInt64Exp dims -  whenActive lvl space $+  whenActive lvl space . localOps threadOperations $     compileStms mempty (kernelBodyStms body) $       forM_ (zip (patNames pat) $ kernelBodyResult body) $ \(dest, res) ->         copyDWIMFix           dest-          (map Imp.vi64 ltids)+          (map Imp.le64 ltids)           (kernelResultSubExp res)           [] -  sOp $ Imp.ErrorSync Imp.FenceLocal+  fence <- fenceForArrays $ patNames pat+  sOp $ Imp.ErrorSync fence    let segment_size = last dims'       crossesSegment from to =@@ -409,11 +422,10 @@    -- groupScan needs to treat the scan output as a one-dimensional   -- array of scan elements, so we invent some new flattened arrays-  -- here.  XXX: this assumes that the original index function is just-  -- row-major, but does not actually verify it.+  -- here.   dims_flat <- dPrimV "dims_flat" $ product dims'   let flattened pe = do-        MemLoc mem _ _ <-+        MemLoc mem _ ixfun <-           entryArrayLoc <$> lookupArray (patElemName pe)         let pe_t = typeOf pe             arr_dims = Var (tvVar dims_flat) : drop (length dims') (arrayDims pe_t)@@ -421,7 +433,8 @@           (baseString (patElemName pe) ++ "_flat")           (elemType pe_t)           (Shape arr_dims)-          $ ArrayIn mem $ IxFun.iota $ map pe64 arr_dims+          mem+          $ IxFun.reshape ixfun $ map (DimNew . pe64) arr_dims        num_scan_results = sum $ map (length . segBinOpNeutral) scans @@ -445,12 +458,12 @@   tmp_arrs <- mapM mkTempArr $ concatMap (lambdaReturnType . segBinOpLambda) ops   let tmps_for_ops = chunks (map (length . segBinOpNeutral) ops) tmp_arrs -  whenActive lvl space $+  whenActive lvl space . localOps threadOperations $     compileStms mempty (kernelBodyStms body) $ do       let (red_res, map_res) =             splitAt (segBinOpResults ops) $ kernelBodyResult body       forM_ (zip tmp_arrs red_res) $ \(dest, res) ->-        copyDWIMFix dest (map Imp.vi64 ltids) (kernelResultSubExp res) []+        copyDWIMFix dest (map Imp.le64 ltids) (kernelResultSubExp res) []       zipWithM_ (compileThreadResult space) map_pes map_res    sOp $ Imp.ErrorSync Imp.FenceLocal@@ -481,9 +494,8 @@                   Shape $                     Var (tvVar dims_flat) :                     drop (length ltids) (memLocShape arr_loc)-            sArray "red_arr_flat" pt flat_shape $-              ArrayIn (memLocName arr_loc) $-                IxFun.iota $ map pe64 $ shapeDims flat_shape+            sArray "red_arr_flat" pt flat_shape (memLocName arr_loc) $+              IxFun.iota $ map pe64 $ shapeDims flat_shape        let segment_size = last dims'           crossesSegment from to =@@ -524,7 +536,7 @@   -- Ensure that all locks have been initialised.   sOp $ Imp.Barrier Imp.FenceLocal -  whenActive lvl space $+  whenActive lvl space . localOps threadOperations $     compileStms mempty (kernelBodyStms kbody) $ do       let (red_res, map_res) = splitAt num_red_res $ kernelBodyResult kbody           (red_is, red_vs) = splitAt (length ops) $ map kernelResultSubExp red_res@@ -537,7 +549,7 @@           let bin' = toInt64Exp bin               dest_w' = toInt64Exp dest_w               bin_in_bounds = 0 .<=. bin' .&&. bin' .<. dest_w'-              bin_is = map Imp.vi64 (init ltids) ++ [bin']+              bin_is = map Imp.le64 (init ltids) ++ [bin']               vs_params = takeLast (length op_vs) $ lambdaParams lam            sComment "perform atomic updates" $@@ -703,9 +715,7 @@           sComment "update global result" $             zipWithM_ (writeArray bucket) arrs $ map (Var . paramName) acc_params -      fence = case space of-        Space "local" -> sOp $ Imp.MemFence Imp.FenceLocal-        _ -> sOp $ Imp.MemFence Imp.FenceGlobal+      fence = sOp $ Imp.MemFence $ fenceForSpace space    -- While-loop: Try to insert your value   sWhile (tvExp continue) $ do@@ -823,8 +833,7 @@   ImpM rep r op (Maybe Imp.KernelConstExp) isConstExp vtable size = do   fname <- askFunction-  let onLeaf (Imp.ScalarVar name) _ = lookupConstExp name-      onLeaf Imp.Index {} _ = Nothing+  let onLeaf name _ = lookupConstExp name       lookupConstExp name =         constExp =<< hasExp =<< M.lookup name vtable       constExp (Op (Inner (SizeOp (GetSize key _)))) =@@ -887,16 +896,16 @@   inner_group_size <- newVName "group_size"   let constants =         KernelConstants-          (Imp.vi32 global_tid)-          (Imp.vi32 local_tid)-          (Imp.vi32 group_id)+          (Imp.le32 global_tid)+          (Imp.le32 local_tid)+          (Imp.le32 group_id)           global_tid           local_tid           group_id           num_groups           group_size           (sExt32 (group_size * num_groups))-          (Imp.vi32 wave_size)+          (Imp.le32 wave_size)           true           mempty @@ -922,7 +931,7 @@   where     (is, ws) = unzip limit     actives = zipWith active is $ map toInt64Exp ws-    active i = (Imp.vi64 i .<.)+    active i = (Imp.le64 i .<.)  -- | Change every memory block to be in the global address space, -- except those who are in the local memory space.  This only affects@@ -1056,6 +1065,8 @@    ltid_in_bounds <- dPrimVE "ltid_in_bounds" $ ltid .<. w +  fence <- fenceForArrays arrs+   -- The scan works by splitting the group into blocks, which are   -- scanned separately.  Typically, these blocks are smaller than   -- the lockstep width, which enables barrier-free execution inside@@ -1085,7 +1096,7 @@         | array_scan =           sOp $ Imp.Barrier Imp.FenceGlobal         | otherwise =-          sOp $ Imp.Barrier Imp.FenceLocal+          sOp $ Imp.Barrier fence        group_offset = sExt64 (kernelGroupId constants) * kernelGroupSize constants @@ -1311,9 +1322,9 @@    return     ( KernelConstants-        (Imp.vi32 thread_gtid)-        (Imp.vi32 thread_ltid)-        (Imp.vi32 group_id)+        (Imp.le32 thread_gtid)+        (Imp.le32 thread_ltid)+        (Imp.le32 group_id)         thread_gtid         thread_ltid         group_id@@ -1321,7 +1332,7 @@         group_size         (sExt32 (group_size * num_groups))         0-        (Imp.vi64 thread_gtid .<. kernel_size)+        (Imp.le64 thread_gtid .<. kernel_size)         mempty,       set_constants     )@@ -1356,7 +1367,7 @@     sOp $ Imp.Barrier Imp.FenceGlobal virtualiseGroups _ _ m = do   gid <- kernelGroupIdVar . kernelConstants <$> askEnv-  m $ Imp.vi32 gid+  m $ Imp.le32 gid  sKernelThread ::   String ->@@ -1509,10 +1520,8 @@         shape = Shape [Var num_elems]     function fname [] params $ do       arr <--        sArray "arr" bt shape $-          ArrayIn mem $-            IxFun.iota $-              map pe64 $ shapeDims shape+        sArray "arr" bt shape mem $+          IxFun.iota $ map pe64 $ shapeDims shape       sReplicateKernel arr $ Var val    return fname@@ -1600,16 +1609,14 @@             Imp.ScalarParam s $ IntType bt           ]         shape = Shape [Var n]-        n' = Imp.vi64 n+        n' = Imp.le64 n         x' = Imp.var x $ IntType bt         s' = Imp.var s $ IntType bt      function fname [] params $ do       arr <--        sArray "arr" (IntType bt) shape $-          ArrayIn mem $-            IxFun.iota $-              map pe64 $ shapeDims shape+        sArray "arr" (IntType bt) shape mem $+          IxFun.iota $ map pe64 $ shapeDims shape       sIotaKernel arr (sExt64 n') x' s' bt    return fname@@ -1635,7 +1642,7 @@     else sIotaKernel arr n x s et  sCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp-sCopy bt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do+sCopy pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do   -- Note that the shape of the destination and the source are   -- necessarily the same.   let shape = map toInt64Exp srcdims@@ -1658,10 +1665,10 @@     (_, destspace, destidx) <- fullyIndexArray' destloc is     (_, srcspace, srcidx) <- fullyIndexArray' srcloc is -    sWhen (gtid .<. kernel_size) $-      emit $-        Imp.Write destmem destidx bt destspace Imp.Nonvolatile $-          Imp.index srcmem srcidx bt srcspace Imp.Nonvolatile+    sWhen (gtid .<. kernel_size) $ do+      tmp <- tvVar <$> dPrim "tmp" pt+      emit $ Imp.Read tmp srcmem srcidx pt srcspace Imp.Nonvolatile+      emit $ Imp.Write destmem destidx pt destspace Imp.Nonvolatile $ Imp.var tmp pt  compileGroupResult ::   SegSpace ->@@ -1691,7 +1698,7 @@ compileGroupResult space pe (TileReturns _ dims what) = do   let gids = map fst $ unSegSpace space       out_tile_sizes = map (toInt64Exp . snd) dims-      group_is = zipWith (*) (map Imp.vi64 gids) out_tile_sizes+      group_is = zipWith (*) (map Imp.le64 gids) out_tile_sizes   local_is <- localThreadIDs $ map snd dims   is_for_thread <-     mapM (dPrimV "thread_out_index") $@@ -1709,7 +1716,7 @@       reg_tiles' = map toInt64Exp reg_tiles    -- Which group tile is this group responsible for?-  let group_tile_is = map Imp.vi64 gids+  let group_tile_is = map Imp.le64 gids    -- Within the group tile, which register tile is this thread   -- responsible for?@@ -1739,7 +1746,7 @@ compileGroupResult space pe (Returns _ _ what) = do   constants <- kernelConstants <$> askEnv   in_local_memory <- arrayInLocalMemory what-  let gids = map (Imp.vi64 . fst) $ unSegSpace space+  let gids = map (Imp.le64 . fst) $ unSegSpace space    if not in_local_memory     then@@ -1764,7 +1771,7 @@ compileThreadResult _ _ RegTileReturns {} =   compilerLimitationS "compileThreadResult: RegTileReturns not yet handled." compileThreadResult space pe (Returns _ _ what) = do-  let is = map (Imp.vi64 . fst) $ unSegSpace space+  let is = map (Imp.le64 . fst) $ unSegSpace space   copyDWIMFix (patElemName pe) is what [] compileThreadResult _ pe (ConcatReturns _ SplitContiguous _ per_thread_elems what) = do   constants <- kernelConstants <$> askEnv

@@ -107,16 +107,13 @@     let subhistos_shape =           Shape (map snd segment_dims ++ [tvSize num_subhistos])             <> stripDims num_segments (arrayShape dest_t)-        subhistos_membind =-          ArrayIn subhistos_mem $-            IxFun.iota $-              map pe64 $ shapeDims subhistos_shape     subhistos <-       sArray         (baseString dest ++ "_subhistos")         (elemType dest_t)         subhistos_shape-        subhistos_membind+        subhistos_mem+        $ IxFun.iota $ map pe64 $ shapeDims subhistos_shape      return $       SubhistosInfo subhistos $ do@@ -436,7 +433,7 @@             forM_ (zip map_pes map_res) $ \(pe, res) ->               copyDWIMFix                 (patElemName pe)-                (map (Imp.vi64 . fst) $ unSegSpace space)+                (map (Imp.le64 . fst) $ unSegSpace space)                 (kernelResultSubExp res)                 [] @@ -463,7 +460,7 @@                    sWhen bucket_in_bounds $ do                     let bucket_is =-                          map Imp.vi64 (init space_is)+                          map Imp.le64 (init space_is)                             ++ [sExt64 subhisto_ind, bucket']                     dLParams $ lambdaParams lam                     sLoopNest shape $ \is -> do@@ -705,7 +702,7 @@         sComment "initialize histograms in local memory" $           onAllHistograms $ \dest_local dest_global op ne local_subhisto_i global_subhisto_i local_bucket_is global_bucket_is ->             sComment "First subhistogram is initialised from global memory; others with neutral element." $ do-              let global_is = map Imp.vi64 segment_is ++ [0] ++ global_bucket_is+              let global_is = map Imp.le64 segment_is ++ [0] ++ global_bucket_is                   local_is = sExt64 local_subhisto_i : local_bucket_is               sIf                 (global_subhisto_i .==. 0)@@ -735,7 +732,7 @@                 forM_ (zip map_pes map_res) $ \(pe, se) ->                   copyDWIMFix                     (patElemName pe)-                    (map Imp.vi64 space_is)+                    (map Imp.le64 space_is)                     se                     [] @@ -815,7 +812,7 @@                  sComment "Put final bucket value in global memory." $ do                   let global_is =-                        map Imp.vi64 segment_is+                        map Imp.le64 segment_is                           ++ [sExt64 group_id `rem` unCount groups_per_segment]                           ++ global_bucket_is                   forM_ (zip xparams global_dests) $ \(xp, global_dest) ->@@ -1142,7 +1139,7 @@           red_cont $             flip map subhistos $ \subhisto ->               ( Var subhisto,-                map Imp.vi64 $+                map Imp.le64 $                   map fst segment_dims ++ [subhistogram_id, bucket_id] ++ vector_ids               ) 

@@ -138,8 +138,8 @@     case paramDec p of       MemArray pt shape _ (ArrayIn mem _) -> do         let shape' = Shape [num_threads] <> shape-        sArray "red_arr" pt shape' $-          ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape'+        sArray "red_arr" pt shape' mem $+          IxFun.iota $ map pe64 $ shapeDims shape'       _ -> do         let pt = elemType $ paramType p             shape = Shape [group_size]@@ -183,7 +183,7 @@       dims' = map toInt64Exp dims       num_groups' = fmap toInt64Exp num_groups       group_size' = fmap toInt64Exp group_size-      global_tid = Imp.vi64 $ segFlat space+      global_tid = Imp.le64 $ segFlat space       w = last dims'    counter <-@@ -478,7 +478,7 @@                   pes                   group_id                   flat_segment_id-                  (map Imp.vi64 segment_gtids)+                  (map Imp.le64 segment_gtids)                   (sExt64 first_group_for_segment)                   groups_per_segment                   slug@@ -498,7 +498,7 @@               forM_ (zip slugs segred_pes) $ \(slug, pes) ->                 sWhen (local_tid .==. 0) $                   forM_ (zip pes (slugAccs slug)) $ \(v, (acc, acc_is)) ->-                    copyDWIMFix (patElemName v) (map Imp.vi64 segment_gtids) (Var acc) acc_is+                    copyDWIMFix (patElemName v) (map Imp.le64 segment_gtids) (Var acc) acc_is        sIf (groups_per_segment .==. 1) one_group_per_segment multiple_groups_per_segment @@ -635,7 +635,7 @@     gtid       <-- case comm of         Commutative ->-          global_tid + Imp.vi64 threads_per_segment * i+          global_tid + Imp.le64 threads_per_segment * i         Noncommutative ->           let index_in_segment = global_tid `quot` kernelGroupSize constants            in sExt64 local_tid

@@ -85,7 +85,8 @@         "local_prefix_arr"         ty         (Shape [groupSize])-        $ ArrayIn localMem $ IxFun.iotaOffset off' [pe64 groupSize]+        localMem+        $ IxFun.iotaOffset off' [pe64 groupSize]    warpscan <- sArrayInMem "warpscan" int8 (Shape [constant (warpSize :: Int64)]) localMem   warpExchanges <-@@ -95,7 +96,8 @@         "warp_exchange"         ty         (Shape [constant (warpSize :: Int64)])-        $ ArrayIn localMem $ IxFun.iotaOffset off' [warpSize]+        localMem+        $ IxFun.iotaOffset off' [warpSize]    return (sharedId, transposedArrays, prefixArrays, warpscan, warpExchanges) @@ -245,10 +247,10 @@       m :: Num a => a       m = fromIntegral $ max 1 $ min mem_constraint reg_constraint -  emit $ Imp.DebugPrint "SegScan: number of elements processed sequentially per thread is m:" $ Just $ untyped (m :: TPrimExp Int32 Imp.ExpLeaf)-  emit $ Imp.DebugPrint "SegScan: memory constraints is: " $ Just $ untyped (fromIntegral mem_constraint :: TPrimExp Int32 Imp.ExpLeaf)-  emit $ Imp.DebugPrint "SegScan: register constraints is: " $ Just $ untyped (fromIntegral reg_constraint :: TPrimExp Int32 Imp.ExpLeaf)-  emit $ Imp.DebugPrint "SegScan: sumT' is: " $ Just $ untyped (fromIntegral sumT' :: TPrimExp Int32 Imp.ExpLeaf)+  emit $ Imp.DebugPrint "SegScan: number of elements processed sequentially per thread is m:" $ Just $ untyped (m :: Imp.TExp Int32)+  emit $ Imp.DebugPrint "SegScan: memory constraints is: " $ Just $ untyped (fromIntegral mem_constraint :: Imp.TExp Int32)+  emit $ Imp.DebugPrint "SegScan: register constraints is: " $ Just $ untyped (fromIntegral reg_constraint :: Imp.TExp Int32)+  emit $ Imp.DebugPrint "SegScan: sumT' is: " $ Just $ untyped (fromIntegral sumT' :: Imp.TExp Int32)    -- Allocate the shared memory for output component   numThreads <- dPrimV "numThreads" num_threads@@ -324,7 +326,7 @@                  -- Write map results to their global memory destinations                 forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(dest, src) ->-                  copyDWIMFix (patElemName dest) (map Imp.vi64 gtids) (kernelResultSubExp src) []+                  copyDWIMFix (patElemName dest) (map Imp.le64 gtids) (kernelResultSubExp src) []                  -- Write to-scan results to private memory.                 forM_ (zip privateArrays $ map kernelResultSubExp all_scan_res) $ \(dest, src) ->@@ -620,7 +622,7 @@           sWhen (flat_idx .<. n) $ do             copyDWIMFix               dest-              (map Imp.vi64 gtids)+              (map Imp.le64 gtids)               (Var locmem)               [sExt64 $ flat_idx - tvExp blockOff]         sOp localBarrier

@@ -42,9 +42,8 @@             MemArray pt shape _ (ArrayIn mem _) -> do               let shape' = Shape [num_threads] <> shape               arr <--                lift $-                  sArray "scan_arr" pt shape' $-                    ArrayIn mem $ IxFun.iota $ map pe64 $ shapeDims shape'+                lift . sArray "scan_arr" pt shape' mem $+                  IxFun.iota $ map pe64 $ shapeDims shape'               return (arr, [])             _ -> do               let pt = elemType $ paramType p@@ -100,7 +99,7 @@     forM_ (zip pes scan_res) $ \(pe, res) ->       copyDWIMFix         (patElemName pe)-        (map Imp.vi64 gtids)+        (map Imp.le64 gtids)         (kernelResultSubExp res)         []   | otherwise =@@ -196,7 +195,7 @@       let per_scan_pes = segBinOpChunks scans all_pes            in_bounds =-            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'+            foldl1 (.&&.) $ zipWith (.<.) (map Imp.le64 gtids) dims'            when_in_bounds = compileStms mempty (kernelBodyStms kbody) $ do             let (all_scan_res, map_res) =@@ -212,7 +211,7 @@               forM_ (zip (takeLast (length map_res) all_pes) map_res) $ \(pe, se) ->                 copyDWIMFix                   (patElemName pe)-                  (map Imp.vi64 gtids)+                  (map Imp.le64 gtids)                   (kernelResultSubExp se)                   [] @@ -236,7 +235,7 @@                 sIf                   in_bounds                   ( do-                      readToScanValues (map Imp.vi64 gtids ++ vec_is) pes scan+                      readToScanValues (map Imp.le64 gtids ++ vec_is) pes scan                       readCarries j (tvExp chunk_offset) dims' vec_is pes scan                   )                   ( forM_ (zip (yParams scan) (segBinOpNeutral scan)) $ \(p, ne) ->@@ -272,7 +271,7 @@                   forM_ (zip3 rets pes local_arrs) $ \(t, pe, arr) ->                     copyDWIMFix                       (patElemName pe)-                      (map Imp.vi64 gtids ++ vec_is)+                      (map Imp.le64 gtids ++ vec_is)                       (Var arr)                       [localArrayIndex constants t] @@ -356,10 +355,10 @@     forM_ (zip4 scans per_scan_local_arrs per_scan_rets per_scan_pes) $       \(SegBinOp _ scan_op nes vec_shape, local_arrs, rets, pes) ->         sLoopNest vec_shape $ \vec_is -> do-          let glob_is = map Imp.vi64 gtids ++ vec_is+          let glob_is = map Imp.le64 gtids ++ vec_is                in_bounds =-                foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'+                foldl1 (.&&.) $ zipWith (.<.) (map Imp.le64 gtids) dims'                when_in_bounds = forM_ (zip3 rets local_arrs pes) $ \(t, arr, pe) ->                 copyDWIMFix@@ -437,7 +436,7 @@       -- then the carry was updated in stage 2), and we are not crossing       -- a segment boundary.       let in_bounds =-            foldl1 (.&&.) $ zipWith (.<.) (map Imp.vi64 gtids) dims'+            foldl1 (.&&.) $ zipWith (.<.) (map Imp.le64 gtids) dims'           crosses_segment =             fromMaybe false $               crossesSegment@@ -468,14 +467,14 @@                     (paramName p)                     []                     (Var $ patElemName pe)-                    (map Imp.vi64 gtids ++ vec_is)+                    (map Imp.le64 gtids ++ vec_is)                  compileBody' scan_x_params $ lambdaBody scan_op                  forM_ (zip scan_x_params pes) $ \(p, pe) ->                   copyDWIMFix                     (patElemName pe)-                    (map Imp.vi64 gtids ++ vec_is)+                    (map Imp.le64 gtids ++ vec_is)                     (Var $ paramName p)                     [] 

@@ -584,7 +584,7 @@ kernelArgs = mapMaybe useToArg . kernelUses   where     useToArg (MemoryUse mem) = Just $ MemKArg mem-    useToArg (ScalarUse v bt) = Just $ ValueKArg (LeafExp (ScalarVar v) bt) bt+    useToArg (ScalarUse v pt) = Just $ ValueKArg (LeafExp v pt) pt     useToArg ConstUse {} = Nothing  nextErrorLabel :: GC.CompilerM KernelOp KernelState String@@ -668,7 +668,7 @@         pendingError False         GC.stm [C.cstm|$id:label: barrier($exp:(fence f));|]         GC.stm [C.cstm|if (local_failure) { return; }|]-      GC.stm [C.cstm|barrier(CLK_LOCAL_MEM_FENCE);|] -- intentional+      GC.stm [C.cstm|barrier($exp:(fence f));|]       GC.modifyUserState $ \s -> s {kernelHasBarriers = True}       incErrorLabel     kernelOps (Atomic space aop) = atomicOps space aop@@ -853,6 +853,8 @@     typesInExp e1 <> typesInExp e2 <> typesInExp e3 typesInCode (Write _ (Count (TPrimExp e1)) t _ _ e2) =   typesInExp e1 <> S.singleton t <> typesInExp e2+typesInCode (Read _ _ (Count (TPrimExp e1)) t _ _) =+  typesInExp e1 <> S.singleton t typesInCode (SetScalar _ e) = typesInExp e typesInCode SetMem {} = mempty typesInCode (Call _ _ es) = mconcat $ map typesInArg es@@ -876,5 +878,4 @@     (from, to) = convOpType op typesInExp (UnOpExp _ e) = typesInExp e typesInExp (FunExp _ args t) = S.singleton t <> mconcat (map typesInExp args)-typesInExp (LeafExp (Index _ (Count (TPrimExp e)) t _ _) _) = S.singleton t <> typesInExp e-typesInExp (LeafExp ScalarVar {} _) = mempty+typesInExp LeafExp {} = mempty

@@ -45,95 +45,106 @@     TransposeSmall ->       mconcat         [ get_ids,-          dec our_array_offset $ vi32 get_global_id_0 `quot` (height * width) * (height * width),-          dec x_index $ (vi32 get_global_id_0 `rem` (height * width)) `quot` height,-          dec y_index $ vi32 get_global_id_0 `rem` height,+          dec our_array_offset $ le32 get_global_id_0 `quot` (height * width) * (height * width),+          dec x_index $ (le32 get_global_id_0 `rem` (height * width)) `quot` height,+          dec y_index $ le32 get_global_id_0 `rem` height,+          DeclareScalar val Nonvolatile t,           dec odata_offset $-            (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset,+            (basic_odata_offset `quot` primByteSize t) + le32 our_array_offset,           dec idata_offset $-            (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset,-          dec index_in $ vi32 y_index * width + vi32 x_index,-          dec index_out $ vi32 x_index * height + vi32 y_index,+            (basic_idata_offset `quot` primByteSize t) + le32 our_array_offset,+          dec index_in $ le32 y_index * width + le32 x_index,+          dec index_out $ le32 x_index * height + le32 y_index,           when-            (vi32 get_global_id_0 .<. width * height * num_arrays)-            ( Write odata (elements $ sExt64 $ vi32 odata_offset + vi32 index_out) t (Space "global") Nonvolatile $-                index idata (elements $ sExt64 $ vi32 idata_offset + vi32 index_in) t (Space "global") Nonvolatile+            (le32 get_global_id_0 .<. width * height * num_arrays)+            ( mconcat+                [ Read val idata (elements $ sExt64 $ le32 idata_offset + le32 index_in) t (Space "global") Nonvolatile,+                  Write odata (elements $ sExt64 $ le32 odata_offset + le32 index_out) t (Space "global") Nonvolatile (var val t)+                ]             )         ]     TransposeLowWidth ->       mkTranspose $         lowDimBody-          (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_0 `quot` muly))-          ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_1-              + (vi32 get_local_id_0 `rem` muly) * block_dim+          (le32 get_group_id_0 * block_dim + (le32 get_local_id_0 `quot` muly))+          ( le32 get_group_id_1 * block_dim * muly + le32 get_local_id_1+              + (le32 get_local_id_0 `rem` muly) * block_dim           )-          ( vi32 get_group_id_1 * block_dim * muly + vi32 get_local_id_0-              + (vi32 get_local_id_1 `rem` muly) * block_dim+          ( le32 get_group_id_1 * block_dim * muly + le32 get_local_id_0+              + (le32 get_local_id_1 `rem` muly) * block_dim           )-          (vi32 get_group_id_0 * block_dim + (vi32 get_local_id_1 `quot` muly))+          (le32 get_group_id_0 * block_dim + (le32 get_local_id_1 `quot` muly))     TransposeLowHeight ->       mkTranspose $         lowDimBody-          ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_0-              + (vi32 get_local_id_1 `rem` mulx) * block_dim+          ( le32 get_group_id_0 * block_dim * mulx + le32 get_local_id_0+              + (le32 get_local_id_1 `rem` mulx) * block_dim           )-          (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_1 `quot` mulx))-          (vi32 get_group_id_1 * block_dim + (vi32 get_local_id_0 `quot` mulx))-          ( vi32 get_group_id_0 * block_dim * mulx + vi32 get_local_id_1-              + (vi32 get_local_id_0 `rem` mulx) * block_dim+          (le32 get_group_id_1 * block_dim + (le32 get_local_id_1 `quot` mulx))+          (le32 get_group_id_1 * block_dim + (le32 get_local_id_0 `quot` mulx))+          ( le32 get_group_id_0 * block_dim * mulx + le32 get_local_id_1+              + (le32 get_local_id_0 `rem` mulx) * block_dim           )     TransposeNormal ->       mkTranspose $         mconcat-          [ dec x_index $ vi32 get_global_id_0,-            dec y_index $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_1,-            when (vi32 x_index .<. width) $+          [ dec x_index $ le32 get_global_id_0,+            dec y_index $ le32 get_group_id_1 * tile_dim + le32 get_local_id_1,+            DeclareScalar val Nonvolatile t,+            when (le32 x_index .<. width) $               For j (untyped elemsPerThread) $-                let i = vi32 j * (tile_dim `quot` elemsPerThread)+                let i = le32 j * (tile_dim `quot` elemsPerThread)                  in mconcat-                      [ dec index_in $ (vi32 y_index + i) * width + vi32 x_index,-                        when (vi32 y_index + i .<. height) $-                          Write-                            block-                            ( elements $-                                sExt64 $-                                  (vi32 get_local_id_1 + i) * (tile_dim + 1)-                                    + vi32 get_local_id_0-                            )-                            t-                            (Space "local")-                            Nonvolatile-                            $ index-                              idata-                              (elements $ sExt64 $ vi32 idata_offset + vi32 index_in)-                              t-                              (Space "global")-                              Nonvolatile+                      [ dec index_in $ (le32 y_index + i) * width + le32 x_index,+                        when (le32 y_index + i .<. height) $+                          mconcat+                            [ Read+                                val+                                idata+                                (elements $ sExt64 $ le32 idata_offset + le32 index_in)+                                t+                                (Space "global")+                                Nonvolatile,+                              Write+                                block+                                ( elements $+                                    sExt64 $+                                      (le32 get_local_id_1 + i) * (tile_dim + 1)+                                        + le32 get_local_id_0+                                )+                                t+                                (Space "local")+                                Nonvolatile+                                (var val t)+                            ]                       ],             Op $ Barrier FenceLocal,-            SetScalar x_index $ untyped $ vi32 get_group_id_1 * tile_dim + vi32 get_local_id_0,-            SetScalar y_index $ untyped $ vi32 get_group_id_0 * tile_dim + vi32 get_local_id_1,-            when (vi32 x_index .<. height) $+            SetScalar x_index $ untyped $ le32 get_group_id_1 * tile_dim + le32 get_local_id_0,+            SetScalar y_index $ untyped $ le32 get_group_id_0 * tile_dim + le32 get_local_id_1,+            when (le32 x_index .<. height) $               For j (untyped elemsPerThread) $-                let i = vi32 j * (tile_dim `quot` elemsPerThread)+                let i = le32 j * (tile_dim `quot` elemsPerThread)                  in mconcat-                      [ dec index_out $ (vi32 y_index + i) * height + vi32 x_index,-                        when (vi32 y_index + i .<. width) $-                          Write-                            odata-                            (elements $ sExt64 $ vi32 odata_offset + vi32 index_out)-                            t-                            (Space "global")-                            Nonvolatile-                            $ index-                              block-                              ( elements $-                                  sExt64 $-                                    vi32 get_local_id_0 * (tile_dim + 1) + vi32 get_local_id_1 + i-                              )-                              t-                              (Space "local")-                              Nonvolatile+                      [ dec index_out $ (le32 y_index + i) * height + le32 x_index,+                        when (le32 y_index + i .<. width) $+                          mconcat+                            [ Read+                                val+                                block+                                ( elements . sExt64 $+                                    le32 get_local_id_0 * (tile_dim + 1) + le32 get_local_id_1 + i+                                )+                                t+                                (Space "local")+                                Nonvolatile,+                              Write+                                odata+                                (elements $ sExt64 $ le32 odata_offset + le32 index_out)+                                t+                                (Space "global")+                                Nonvolatile+                                (var val t)+                            ]                       ]           ]   where@@ -174,7 +185,8 @@       get_group_id_0,       get_group_id_1,       get_group_id_2,-      j+      j,+      val       ] =         zipWith (flip VName) [30 ..] $           map@@ -192,7 +204,8 @@               "get_group_id_0",               "get_group_id_1",               "get_group_id_2",-              "j"+              "j",+              "val"             ]      get_ids =@@ -214,11 +227,11 @@     mkTranspose body =       mconcat         [ get_ids,-          dec our_array_offset $ vi32 get_group_id_2 * width * height,+          dec our_array_offset $ le32 get_group_id_2 * width * height,           dec odata_offset $-            (basic_odata_offset `quot` primByteSize t) + vi32 our_array_offset,+            (basic_odata_offset `quot` primByteSize t) + le32 our_array_offset,           dec idata_offset $-            (basic_idata_offset `quot` primByteSize t) + vi32 our_array_offset,+            (basic_idata_offset `quot` primByteSize t) + le32 our_array_offset,           body         ] @@ -226,37 +239,46 @@       mconcat         [ dec x_index x_in_index,           dec y_index y_in_index,-          dec index_in $ vi32 y_index * width + vi32 x_index,-          when (vi32 x_index .<. width .&&. vi32 y_index .<. height) $-            Write-              block-              (elements $ sExt64 $ vi32 get_local_id_1 * (block_dim + 1) + vi32 get_local_id_0)-              t-              (Space "local")-              Nonvolatile-              $ index-                idata-                (elements $ sExt64 $ vi32 idata_offset + vi32 index_in)-                t-                (Space "global")-                Nonvolatile,+          DeclareScalar val Nonvolatile t,+          dec index_in $ le32 y_index * width + le32 x_index,+          when (le32 x_index .<. width .&&. le32 y_index .<. height) $+            mconcat+              [ Read+                  val+                  idata+                  (elements $ sExt64 $ le32 idata_offset + le32 index_in)+                  t+                  (Space "global")+                  Nonvolatile,+                Write+                  block+                  (elements $ sExt64 $ le32 get_local_id_1 * (block_dim + 1) + le32 get_local_id_0)+                  t+                  (Space "local")+                  Nonvolatile+                  (var val t)+              ],           Op $ Barrier FenceLocal,           SetScalar x_index $ untyped x_out_index,           SetScalar y_index $ untyped y_out_index,-          dec index_out $ vi32 y_index * height + vi32 x_index,-          when (vi32 x_index .<. height .&&. vi32 y_index .<. width) $-            Write-              odata-              (elements $ sExt64 (vi32 odata_offset + vi32 index_out))-              t-              (Space "global")-              Nonvolatile-              $ index-                block-                (elements $ sExt64 $ vi32 get_local_id_0 * (block_dim + 1) + vi32 get_local_id_1)-                t-                (Space "local")-                Nonvolatile+          dec index_out $ le32 y_index * height + le32 x_index,+          when (le32 x_index .<. height .&&. le32 y_index .<. width) $+            mconcat+              [ Read+                  val+                  block+                  (elements $ sExt64 $ le32 get_local_id_0 * (block_dim + 1) + le32 get_local_id_1)+                  t+                  (Space "local")+                  Nonvolatile,+                Write+                  odata+                  (elements $ sExt64 (le32 odata_offset + le32 index_out))+                  t+                  (Space "global")+                  Nonvolatile+                  (var val t)+              ]         ]  -- | Generate a transpose kernel.  There is special support to handle

@@ -117,7 +117,7 @@   KernelResult ->   MulticoreGen () compileThreadResult space pe (Returns _ _ what) = do-  let is = map (Imp.vi64 . fst) $ unSegSpace space+  let is = map (Imp.le64 . fst) $ unSegSpace space   copyDWIMFix (patElemName pe) is what [] compileThreadResult _ _ ConcatReturns {} =   compilerBugS "compileThreadResult: ConcatReturn unhandled."

@@ -122,11 +122,11 @@            sComment "save map-out results" $             forM_ (zip map_pes map_res) $ \(pe, res) ->-              copyDWIMFix (patElemName pe) (map Imp.vi64 is) (kernelResultSubExp res) []+              copyDWIMFix (patElemName pe) (map Imp.le64 is) (kernelResultSubExp res) []            sComment "perform updates" $             sWhen bucket_in_bounds $ do-              let bucket_is = map Imp.vi64 (init is) ++ [bucket']+              let bucket_is = map Imp.le64 (init is) ++ [bucket']               dLParams $ lambdaParams lam               sLoopNest shape $ \is' -> do                 forM_ (zip vs_params vs') $ \(p, res) ->@@ -186,7 +186,7 @@       let shape = Shape [tvSize num_histos] <> arrayShape t       sAllocArray "subhistogram" (elemType t) shape DefaultSpace -  let tid' = Imp.vi64 $ segFlat space+  let tid' = Imp.le64 $ segFlat space       flat_idx' = tvExp flat_idx    (local_subhistograms, prebody) <- collect' $ do@@ -221,7 +221,7 @@         forM_ (zip map_pes map_res) $ \(pe, res) ->           copyDWIMFix             (patElemName pe)-            (map Imp.vi64 is)+            (map Imp.le64 is)             (kernelResultSubExp res)             [] @@ -274,7 +274,7 @@       red_cont $         flip map hists $ \subhisto ->           ( Var subhisto,-            map Imp.vi64 $+            map Imp.le64 $               map fst segment_dims ++ [subhistogram_id, bucket_id]           ) @@ -343,7 +343,7 @@              sComment "save map-out results" $               forM_ (zip map_pes map_res) $ \(pe, res) ->-                copyDWIMFix (patElemName pe) (map Imp.vi64 is) res []+                copyDWIMFix (patElemName pe) (map Imp.le64 is) res []              sComment "perform updates" $               sWhen bucket_in_bounds $ do@@ -352,10 +352,10 @@                   -- Index                   let buck = toInt64Exp bucket                   forM_ (zip red_pes is_params) $ \(pe, p) ->-                    copyDWIMFix (paramName p) [] (Var $ patElemName pe) (map Imp.vi64 (init is) ++ [buck] ++ vec_is)+                    copyDWIMFix (paramName p) [] (Var $ patElemName pe) (map Imp.le64 (init is) ++ [buck] ++ vec_is)                   -- Value at index                   forM_ (zip vs_params vs') $ \(p, v) ->                     copyDWIMFix (paramName p) [] v vec_is                   compileStms mempty (bodyStms $ lambdaBody lam) $                     forM_ (zip red_pes $ map resSubExp $ bodyResult $ lambdaBody lam) $-                      \(pe, se) -> copyDWIMFix (patElemName pe) (map Imp.vi64 (init is) ++ [buck] ++ vec_is) se []+                      \(pe, se) -> copyDWIMFix (patElemName pe) (map Imp.le64 (init is) ++ [buck] ++ vec_is) se []

@@ -17,7 +17,7 @@   KernelResult ->   MulticoreGen () writeResult is pe (Returns _ _ se) =-  copyDWIMFix (patElemName pe) (map Imp.vi64 is) se []+  copyDWIMFix (patElemName pe) (map Imp.le64 is) se [] writeResult _ pe (WriteReturns _ (Shape rws) _ idx_vals) = do   let (iss, vs) = unzip idx_vals       rws' = map toInt64Exp rws

@@ -156,7 +156,7 @@   postbody <- collect $     forM_ (zip slugs slug_local_accs) $ \(slug, local_accs) ->       forM (zip (slugResArrs slug) local_accs) $ \(acc, local_acc) ->-        copyDWIMFix acc [Imp.vi64 $ segFlat space] (Var local_acc) []+        copyDWIMFix acc [Imp.le64 $ segFlat space] (Var local_acc) []    free_params <- freeParams (prebody <> fbody <> postbody) (segFlat space : [tvVar flat_idx])   let (body_allocs, fbody') = extractAllocations fbody@@ -170,7 +170,7 @@   MulticoreGen () reductionStage2 pat space nsubtasks slugs = do   let per_red_pes = segBinOpChunks (map slugOp slugs) $ patElems pat-      phys_id = Imp.vi64 (segFlat space)+      phys_id = Imp.le64 (segFlat space)   sComment "neutral-initialise the output" $     forM_ (zip (map slugOp slugs) per_red_pes) $ \(red, red_res) ->       forM_ (zip red_res $ segBinOpNeutral red) $ \(pe, ne) ->@@ -235,7 +235,7 @@       forM_ (zip per_red_pes reds) $ \(pes, red) ->         forM_ (zip pes (segBinOpNeutral red)) $ \(pe, ne) ->           sLoopNest (segBinOpShape red) $ \vec_is ->-            copyDWIMFix (patElemName pe) (map Imp.vi64 (init is) ++ vec_is) ne []+            copyDWIMFix (patElemName pe) (map Imp.le64 (init is) ++ vec_is) ne []      sComment "main body" $ do       dScope Nothing $ scopeOfLParams $ concatMap (lambdaParams . segBinOpLambda) reds@@ -252,7 +252,7 @@               sComment "load accum" $ do                 let acc_params = take (length (segBinOpNeutral red)) $ (lambdaParams . segBinOpLambda) red                 forM_ (zip acc_params pes) $ \(p, pe) ->-                  copyDWIMFix (paramName p) [] (Var $ patElemName pe) (map Imp.vi64 (init is) ++ vec_is)+                  copyDWIMFix (paramName p) [] (Var $ patElemName pe) (map Imp.le64 (init is) ++ vec_is)                sComment "load new val" $ do                 let next_params = drop (length (segBinOpNeutral red)) $ (lambdaParams . segBinOpLambda) red@@ -264,4 +264,4 @@                 compileStms mempty (bodyStms lbody) $                   sComment "write back to res" $                     forM_ (zip pes $ map resSubExp $ bodyResult lbody) $-                      \(pe, se') -> copyDWIMFix (patElemName pe) (map Imp.vi64 (init is) ++ vec_is) se' []+                      \(pe, se') -> copyDWIMFix (patElemName pe) (map Imp.le64 (init is) ++ vec_is) se' []

@@ -120,7 +120,7 @@             compileStms mempty (bodyStms $ lamBody scan_op) $               forM_ (zip3 acc pes $ map resSubExp $ bodyResult $ lamBody scan_op) $                 \(acc', pe, se) -> do-                  copyDWIMFix (patElemName pe) (map Imp.vi64 is ++ vec_is) se []+                  copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []                   copyDWIMFix acc' vec_is se []    free_params <- freeParams (prebody <> body) (segFlat space : [tvVar iter])@@ -240,7 +240,7 @@             compileStms mempty (bodyStms $ lamBody scan_op) $               forM_ (zip3 pes (map resSubExp $ bodyResult $ lamBody scan_op) acc) $                 \(pe, se, acc') -> do-                  copyDWIMFix (patElemName pe) (map Imp.vi64 is ++ vec_is) se []+                  copyDWIMFix (patElemName pe) (map Imp.le64 is ++ vec_is) se []                   copyDWIMFix acc' vec_is se []    free_params' <- freeParams (prebody <> body) (segFlat space : [tvVar iter])@@ -298,10 +298,10 @@            sComment "write mapped values results to memory" $             forM_ (zip (drop (length $ segBinOpNeutral scan_op) $ patElems pat) map_res) $ \(pe, se) ->-              copyDWIMFix (patElemName pe) (map Imp.vi64 is) (kernelResultSubExp se) []+              copyDWIMFix (patElemName pe) (map Imp.le64 is) (kernelResultSubExp se) []            sComment "combine with carry and write to memory" $             compileStms mempty (bodyStms $ lambdaBody $ segBinOpLambda scan_op) $               forM_ (zip3 scan_x_params scan_pes $ map resSubExp $ bodyResult $ lambdaBody $ segBinOpLambda scan_op) $ \(p, pe, se) -> do-                copyDWIMFix (patElemName pe) (map Imp.vi64 is) se []+                copyDWIMFix (patElemName pe) (map Imp.le64 is) se []                 copyDWIMFix (paramName p) [] se []

@@ -46,9 +46,9 @@     []     params     ( mconcat-        [ dec r $ vi64 re - vi64 rb,-          dec c $ vi64 ce - vi64 cb,-          If (vi64 num_arrays .==. 1) doTranspose doMapTranspose+        [ dec r $ le64 re - le64 rb,+          dec c $ le64 ce - le64 cb,+          If (le64 num_arrays .==. 1) doTranspose doMapTranspose         ]     )     []@@ -85,7 +85,8 @@       r,       c,       i,-      j+      j,+      val       ] =         zipWith           (VName . nameFromString)@@ -103,38 +104,44 @@             "r",             "c",             "i",-            "j" -- local+            "j", -- local+            "val"           ]           [0 ..]      dec v e = DeclareScalar v Nonvolatile int32 <> SetScalar v (untyped e)      naiveTranspose =-      For j (untyped $ vi64 c) $-        For i (untyped $ vi64 r) $-          let i' = vi64 i + vi64 rb-              j' = vi64 j + vi64 cb-           in Write-                destmem-                (elements $ vi64 destoffset + j' * vi64 n + i')-                pt-                DefaultSpace-                Nonvolatile-                $ index-                  srcmem-                  (elements $ vi64 srcoffset + i' * vi64 m + j')-                  pt-                  DefaultSpace-                  Nonvolatile+      For j (untyped $ le64 c) $+        For i (untyped $ le64 r) $+          let i' = le64 i + le64 rb+              j' = le64 j + le64 cb+           in mconcat+                [ DeclareScalar val Nonvolatile pt,+                  Read+                    val+                    srcmem+                    (elements $ le64 srcoffset + i' * le64 m + j')+                    pt+                    DefaultSpace+                    Nonvolatile,+                  Write+                    destmem+                    (elements $ le64 destoffset + j' * le64 n + i')+                    pt+                    DefaultSpace+                    Nonvolatile+                    (var val pt)+                ]      recArgs (cb', ce', rb', re') =       [ MemArg destmem,-        ExpArg $ untyped $ vi64 destoffset,+        ExpArg $ untyped $ le64 destoffset,         MemArg srcmem,-        ExpArg $ untyped $ vi64 srcoffset,-        ExpArg $ untyped $ vi64 num_arrays,-        ExpArg $ untyped $ vi64 m,-        ExpArg $ untyped $ vi64 n,+        ExpArg $ untyped $ le64 srcoffset,+        ExpArg $ untyped $ le64 num_arrays,+        ExpArg $ untyped $ le64 m,+        ExpArg $ untyped $ le64 n,         ExpArg $ untyped cb',         ExpArg $ untyped ce',         ExpArg $ untyped rb',@@ -145,28 +152,28 @@     doTranspose =       mconcat         [ If-            (vi64 r .<=. cutoff .&&. vi64 c .<=. cutoff)+            (le64 r .<=. cutoff .&&. le64 c .<=. cutoff)             naiveTranspose             $ If-              (vi64 r .>=. vi64 c)+              (le64 r .>=. le64 c)               ( Call                   []                   fname                   ( recArgs-                      ( vi64 cb,-                        vi64 ce,-                        vi64 rb,-                        vi64 rb + (vi64 r `quot` 2)+                      ( le64 cb,+                        le64 ce,+                        le64 rb,+                        le64 rb + (le64 r `quot` 2)                       )                   )                   <> Call                     []                     fname                     ( recArgs-                        ( vi64 cb,-                          vi64 ce,-                          vi64 rb + vi64 r `quot` 2,-                          vi64 re+                        ( le64 cb,+                          le64 ce,+                          le64 rb + le64 r `quot` 2,+                          le64 re                         )                     )               )@@ -174,20 +181,20 @@                   []                   fname                   ( recArgs-                      ( vi64 cb,-                        vi64 cb + (vi64 c `quot` 2),-                        vi64 rb,-                        vi64 re+                      ( le64 cb,+                        le64 cb + (le64 c `quot` 2),+                        le64 rb,+                        le64 re                       )                   )                   <> Call                     []                     fname                     ( recArgs-                        ( vi64 cb + vi64 c `quot` 2,-                          vi64 ce,-                          vi64 rb,-                          vi64 re+                        ( le64 cb + le64 c `quot` 2,+                          le64 ce,+                          le64 rb,+                          le64 re                         )                     )               )@@ -196,19 +203,19 @@     doMapTranspose =       -- In the map-transpose case, we assume that cb==rb==0, ce==m,       -- re==n.-      For i (untyped $ vi64 num_arrays) $+      For i (untyped $ le64 num_arrays) $         Call           []           fname           [ MemArg destmem,-            ExpArg $ untyped $ vi64 destoffset + vi64 i * vi64 m * vi64 n,+            ExpArg $ untyped $ le64 destoffset + le64 i * le64 m * le64 n,             MemArg srcmem,-            ExpArg $ untyped $ vi64 srcoffset + vi64 i * vi64 m * vi64 n,+            ExpArg $ untyped $ le64 srcoffset + le64 i * le64 m * le64 n,             ExpArg $ untyped (1 :: TExp Int64),-            ExpArg $ untyped $ vi64 m,-            ExpArg $ untyped $ vi64 n,-            ExpArg $ untyped $ vi64 cb,-            ExpArg $ untyped $ vi64 ce,-            ExpArg $ untyped $ vi64 rb,-            ExpArg $ untyped $ vi64 re+            ExpArg $ untyped $ le64 m,+            ExpArg $ untyped $ le64 n,+            ExpArg $ untyped $ le64 cb,+            ExpArg $ untyped $ le64 ce,+            ExpArg $ untyped $ le64 rb,+            ExpArg $ untyped $ le64 re           ]

@@ -56,7 +56,7 @@  setCodeSpace :: Space -> Code op -> Code op setCodeSpace space (Allocate v e old_space) =-  Allocate v (fmap (setTExpSpace space) e) $ setSpace space old_space+  Allocate v e $ setSpace space old_space setCodeSpace space (Free v old_space) =   Free v $ setSpace space old_space setCodeSpace space (DeclareMem name old_space) =@@ -64,66 +64,42 @@ setCodeSpace space (DeclareArray name _ t vs) =   DeclareArray name space t vs setCodeSpace space (Copy dest dest_offset dest_space src src_offset src_space n) =-  Copy-    dest-    (fmap (setTExpSpace space) dest_offset)-    dest_space'-    src-    (fmap (setTExpSpace space) src_offset)-    src_space'-    $ fmap (setTExpSpace space) n+  Copy dest dest_offset dest_space' src src_offset src_space' n   where     dest_space' = setSpace space dest_space     src_space' = setSpace space src_space setCodeSpace space (Write dest dest_offset bt dest_space vol e) =-  Write-    dest-    (fmap (setTExpSpace space) dest_offset)-    bt-    (setSpace space dest_space)-    vol-    (setExpSpace space e)+  Write dest dest_offset bt (setSpace space dest_space) vol e+setCodeSpace space (Read x dest dest_offset bt dest_space vol) =+  Read x dest dest_offset bt (setSpace space dest_space) vol setCodeSpace space (c1 :>>: c2) =   setCodeSpace space c1 :>>: setCodeSpace space c2 setCodeSpace space (For i e body) =-  For i (setExpSpace space e) $ setCodeSpace space body+  For i e $ setCodeSpace space body setCodeSpace space (While e body) =-  While (setTExpSpace space e) $ setCodeSpace space body+  While e $ setCodeSpace space body setCodeSpace space (If e c1 c2) =-  If (setTExpSpace space e) (setCodeSpace space c1) (setCodeSpace space c2)+  If e (setCodeSpace space c1) (setCodeSpace space c2) setCodeSpace space (Comment s c) =   Comment s $ setCodeSpace space c setCodeSpace _ Skip =   Skip setCodeSpace _ (DeclareScalar name vol bt) =   DeclareScalar name vol bt-setCodeSpace space (SetScalar name e) =-  SetScalar name $ setExpSpace space e+setCodeSpace _ (SetScalar name e) =+  SetScalar name e setCodeSpace space (SetMem to from old_space) =   SetMem to from $ setSpace space old_space-setCodeSpace space (Call dests fname args) =-  Call dests fname $ map setArgSpace args-  where-    setArgSpace (MemArg m) = MemArg m-    setArgSpace (ExpArg e) = ExpArg $ setExpSpace space e-setCodeSpace space (Assert e msg loc) =-  Assert (setExpSpace space e) (fmap (setExpSpace space) msg) loc-setCodeSpace space (DebugPrint s v) =-  DebugPrint s $ fmap (setExpSpace space) v-setCodeSpace space (TracePrint msg) =-  TracePrint $ fmap (setExpSpace space) msg+setCodeSpace _ (Call dests fname args) =+  Call dests fname args+setCodeSpace _ (Assert e msg loc) =+  Assert e msg loc+setCodeSpace _ (DebugPrint s v) =+  DebugPrint s v+setCodeSpace _ (TracePrint msg) =+  TracePrint msg setCodeSpace _ (Op op) =   Op op--setExpSpace :: Space -> Exp -> Exp-setExpSpace space = fmap setLeafSpace-  where-    setLeafSpace (Index mem i bt DefaultSpace vol) =-      Index mem i bt space vol-    setLeafSpace e = e--setTExpSpace :: Space -> TExp t -> TExp t-setTExpSpace space = TPrimExp . setExpSpace space . untyped  setSpace :: Space -> Space -> Space setSpace space DefaultSpace = space

@@ -396,7 +396,7 @@   return $ specRow lhs (mhs <> " : ") rhs  valBindHtml :: Html -> ValBind -> DocM (Html, Html, Html)-valBindHtml name (ValBind _ _ retdecl (Info (rettype, _)) tparams params _ _ _ _) = do+valBindHtml name (ValBind _ _ retdecl (Info rettype) tparams params _ _ _ _) = do   let tparams' = mconcat $ map ((" " <>) . typeParamHtml) tparams       noLink' =         noLink $@@ -836,7 +836,7 @@ valBindWhat :: ValBind -> IndexWhat valBindWhat vb   | null (valBindParams vb),-    RetType _ t <- fst $ unInfo $ valBindRetType vb,+    RetType _ t <- unInfo $ valBindRetType vb,     orderZero t =     IndexValue   | otherwise =

@@ -169,7 +169,7 @@     als' ->       Just $         PP.oneLine $-          PP.text "-- Result of " <> PP.ppr name <> PP.text " aliases "+          PP.text "-- Result for " <> PP.ppr name <> PP.text " aliases "             <> PP.commasep (map PP.ppr als')  removeAliases :: CanBeAliased (Op rep) => Rephraser Identity (Aliases rep) rep

@@ -310,6 +310,10 @@   removeOpWisdom (OtherOp op) = OtherOp $ removeOpWisdom op   removeOpWisdom (SizeOp op) = SizeOp op +  addOpWisdom (SegOp op) = SegOp $ addOpWisdom op+  addOpWisdom (OtherOp op) = OtherOp $ addOpWisdom op+  addOpWisdom (SizeOp op) = SizeOp op+ instance (ASTRep rep, ST.IndexOp op) => ST.IndexOp (HostOp rep op) where   indexOp vtable k (SegOp op) is = ST.indexOp vtable k op is   indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is

@@ -47,8 +47,8 @@     BodyDec rep ~ ()   ) =>   Simplify.SimplifyOp rep op ->-  HostOp rep op ->-  Engine.SimpleM rep (HostOp (Wise rep) (OpWithWisdom op), Stms (Wise rep))+  HostOp (Wise rep) op ->+  Engine.SimpleM rep (HostOp (Wise rep) op, Stms (Wise rep)) simplifyKernelOp f (OtherOp op) = do   (op', stms) <- f op   return (OtherOp op', stms)

@@ -89,12 +89,7 @@ simplifyProg = simplifyProgGeneric simpleGPUMem  simplifyStms ::-  (HasScope GPUMem m, MonadFreshNames m) =>-  Stms GPUMem ->-  m-    ( Engine.SymbolTable (Engine.Wise GPUMem),-      Stms GPUMem-    )+  (HasScope GPUMem m, MonadFreshNames m) => Stms GPUMem -> m (Stms GPUMem) simplifyStms = simplifyStmsGeneric simpleGPUMem  simpleGPUMem :: Engine.SimpleOps GPUMem

@@ -119,6 +119,11 @@   removeOpWisdom (OtherOp op) =     OtherOp $ removeOpWisdom op +  addOpWisdom (ParOp par_op op) =+    ParOp (addOpWisdom <$> par_op) (addOpWisdom op)+  addOpWisdom (OtherOp op) =+    OtherOp $ addOpWisdom op+ instance (ASTRep rep, ST.IndexOp op) => ST.IndexOp (MCOp rep op) where   indexOp vtable k (ParOp _ op) is = ST.indexOp vtable k op is   indexOp vtable k (OtherOp op) is = ST.indexOp vtable k op is@@ -152,8 +157,8 @@     BodyDec rep ~ ()   ) =>   Simplify.SimplifyOp rep op ->-  MCOp rep op ->-  Engine.SimpleM rep (MCOp (Wise rep) (OpWithWisdom op), Stms (Wise rep))+  MCOp (Wise rep) op ->+  Engine.SimpleM rep (MCOp (Wise rep) op, Stms (Wise rep)) simplifyMCOp f (OtherOp op) = do   (op', stms) <- f op   return (OtherOp op', stms)

@@ -242,6 +242,8 @@   type OpWithWisdom (MemOp inner) = MemOp (OpWithWisdom inner)   removeOpWisdom (Alloc size space) = Alloc size space   removeOpWisdom (Inner k) = Inner $ removeOpWisdom k+  addOpWisdom (Alloc size space) = Alloc size space+  addOpWisdom (Inner k) = Inner $ addOpWisdom k  instance ST.IndexOp inner => ST.IndexOp (MemOp inner) where   indexOp vtable k (Inner op) is = ST.indexOp vtable k op is

@@ -31,7 +31,7 @@ simpleGeneric ::   (SimplifyMemory rep inner) =>   (OpWithWisdom inner -> UT.UsageTable) ->-  Simplify.SimplifyOp rep inner ->+  Simplify.SimplifyOp rep (OpWithWisdom inner) ->   Simplify.SimpleOps rep simpleGeneric = simplifiable @@ -64,7 +64,7 @@   ) =>   Simplify.SimpleOps rep ->   Stms rep ->-  m (ST.SymbolTable (Wise rep), Stms rep)+  m (Stms rep) simplifyStmsGeneric ops stms = do   scope <- askScope   Simplify.simplifyStms@@ -109,7 +109,6 @@   standardRules     <> ruleBook       [ RuleBasicOp copyCopyToCopy,-        RuleBasicOp removeIdentityCopy,         RuleIf unExistentialiseMemory,         RuleOp decertifySafeAlloc       ]@@ -206,22 +205,6 @@         letExp "rearrange_v0" $ BasicOp $ Rearrange perm v2     letBind pat $ BasicOp $ Copy v0' copyCopyToCopy _ _ _ _ = Skip---- | If the destination of a copy is the same as the source, just--- remove it.-removeIdentityCopy ::-  ( BuilderOps rep,-    LetDec rep ~ (VarWisdom, MemBound u)-  ) =>-  TopDownRuleBasicOp rep-removeIdentityCopy vtable pat@(Pat [pe]) _ (Copy v)-  | (_, MemArray _ _ _ (ArrayIn dest_mem dest_ixfun)) <- patElemDec pe,-    Just (_, MemArray _ _ _ (ArrayIn src_mem src_ixfun)) <--      ST.entryLetBoundDec =<< ST.lookup v vtable,-    dest_mem == src_mem,-    dest_ixfun == src_ixfun =-    Simplify $ letBind pat $ BasicOp $ SubExp $ Var v-removeIdentityCopy _ _ _ _ = Skip  -- If an allocation is statically known to be safe, then we can remove -- the certificates on it.  This can help hoist things that would

@@ -52,7 +52,7 @@   lexeme . fmap nameFromString $     (:) <$> satisfy leading <*> many (satisfy constituent)   where-    leading c = isAlpha c || c == '_'+    leading c = isAlpha c || c `elem` ("_+-*/%=!<>|&^." :: String)  pVName :: Parser VName pVName = lexeme $ do@@ -876,17 +876,17 @@       mon <- pLab "monotonicity" $ brackets (pMon `sepBy` pComma)       pure $ IxFun.LMAD offset $ zipWith5 IxFun.LMADDim strides rotates shape perm mon -pPrimExpLeaf :: Parser (VName, PrimType)-pPrimExpLeaf = (,int64) <$> pVName+pPrimExpLeaf :: Parser VName+pPrimExpLeaf = pVName -pExtPrimExpLeaf :: Parser (Ext VName, PrimType)-pExtPrimExpLeaf = (,int64) <$> pExt pVName+pExtPrimExpLeaf :: Parser (Ext VName)+pExtPrimExpLeaf = pExt pVName  pIxFun :: Parser IxFun-pIxFun = pIxFunBase $ isInt64 <$> pPrimExp pPrimExpLeaf+pIxFun = pIxFunBase $ isInt64 <$> pPrimExp int64 pPrimExpLeaf  pExtIxFun :: Parser ExtIxFun-pExtIxFun = pIxFunBase $ isInt64 <$> pPrimExp pExtPrimExpLeaf+pExtIxFun = pIxFunBase $ isInt64 <$> pPrimExp int64 pExtPrimExpLeaf  pMemInfo :: Parser d -> Parser u -> Parser ret -> Parser (MemInfo d u ret) pMemInfo pd pu pret =

@@ -695,10 +695,20 @@ doBinOp SRem {} = doRiskyIntBinOp doSRem doBinOp SMin {} = doIntBinOp doSMin doBinOp UMin {} = doIntBinOp doUMin-doBinOp FMin {} = doFloatBinOp min min min+doBinOp FMin {} = doFloatBinOp fmin fmin fmin+  where+    fmin x y+      | isNaN x = y+      | isNaN y = x+      | otherwise = min x y doBinOp SMax {} = doIntBinOp doSMax doBinOp UMax {} = doIntBinOp doUMax-doBinOp FMax {} = doFloatBinOp max max max+doBinOp FMax {} = doFloatBinOp fmax fmax fmax+  where+    fmax x y+      | isNaN x = y+      | isNaN y = x+      | otherwise = max x y doBinOp Shl {} = doIntBinOp doShl doBinOp LShr {} = doIntBinOp doLShr doBinOp AShr {} = doIntBinOp doAShr

@@ -31,8 +31,9 @@   ) where -import Control.Arrow (first)+import Data.Bifunctor (first, second) import qualified Data.Kind+import Data.List (find) import qualified Data.Map as M import Futhark.IR.Prop (IsOp, NameInfo (..), Scope) import Futhark.IR.Prop.Names@@ -101,11 +102,25 @@         (bodyAliases tb, consumedInBody tb)         (bodyAliases fb, consumedInBody fb) expAliases (BasicOp op) = basicOpAliases op-expAliases (DoLoop merge _ loopbody) =-  map (`namesSubtract` merge_names) aliases+expAliases (DoLoop merge _ loopbody) = do+  (p, als) <-+    transitive . zip params $ zipWith mappend arg_aliases (bodyAliases loopbody)+  let als' = als `namesSubtract` param_names+  if unique $ paramDeclType p+    then pure mempty+    else pure als'   where-    aliases = bodyAliases loopbody-    merge_names = namesFromList $ map (paramName . fst) merge+    arg_aliases = map (subExpAliases . snd) merge+    params = map fst merge+    param_names = namesFromList $ map paramName params+    transitive merge_and_als =+      let merge_and_als' = map (second expand) merge_and_als+       in if merge_and_als' == merge_and_als+            then merge_and_als+            else transitive merge_and_als'+      where+        look v = maybe mempty snd $ find ((== v) . paramName . fst) merge_and_als+        expand als = als <> foldMap look (namesToList als) expAliases (Apply _ args t _) =   funcallAliases args $ map declExtTypeOf t expAliases (WithAcc inputs lam) =

@@ -607,9 +607,8 @@ instance (ASTRep rep, CanBeWise (Op rep)) => CanBeWise (SOAC rep) where   type OpWithWisdom (SOAC rep) = SOAC (Wise rep) -  removeOpWisdom = runIdentity . mapSOACM remove-    where-      remove = SOACMapper return (return . removeLambdaWisdom) return+  removeOpWisdom = runIdentity . mapSOACM (SOACMapper pure (pure . removeLambdaWisdom) pure)+  addOpWisdom = runIdentity . mapSOACM (SOACMapper pure (pure . informLambda) pure)  instance RepTypes rep => ST.IndexOp (SOAC rep) where   indexOp vtable k soac [i] = do

@@ -76,32 +76,27 @@ simplifyStms ::   (HasScope SOACS m, MonadFreshNames m) =>   Stms SOACS ->-  m (ST.SymbolTable (Wise SOACS), Stms SOACS)+  m (Stms SOACS) simplifyStms stms = do   scope <- askScope-  Simplify.simplifyStms-    simpleSOACS-    soacRules-    Engine.noExtraHoistBlockers-    scope-    stms+  Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers scope stms  simplifyConsts ::   MonadFreshNames m =>   Stms SOACS ->-  m (ST.SymbolTable (Wise SOACS), Stms SOACS)+  m (Stms SOACS) simplifyConsts =   Simplify.simplifyStms simpleSOACS soacRules Engine.noExtraHoistBlockers mempty  simplifySOAC ::   Simplify.SimplifiableRep rep =>-  Simplify.SimplifyOp rep (SOAC rep)+  Simplify.SimplifyOp rep (SOAC (Wise rep)) simplifySOAC (Stream outerdim arr form nes lam) = do   outerdim' <- Engine.simplify outerdim   (form', form_hoisted) <- simplifyStreamForm form   nes' <- mapM Engine.simplify nes   arr' <- mapM Engine.simplify arr-  (lam', lam_hoisted) <- Engine.simplifyLambda lam+  (lam', lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam   return     ( Stream outerdim' arr' form' nes' lam',       form_hoisted <> lam_hoisted@@ -114,7 +109,7 @@       return (Sequential, mempty) simplifySOAC (Scatter w ivs lam as) = do   w' <- Engine.simplify w-  (lam', hoisted) <- Engine.simplifyLambda lam+  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam   ivs' <- mapM Engine.simplify ivs   as' <- mapM Engine.simplify as   return (Scatter w' ivs' lam' as', hoisted)@@ -126,10 +121,10 @@       rf' <- Engine.simplify rf       dests' <- Engine.simplify dests       nes' <- mapM Engine.simplify nes-      (op', hoisted) <- Engine.simplifyLambda op+      (op', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda op       return (HistOp dests_w' rf' dests' nes' op', hoisted)   imgs' <- mapM Engine.simplify imgs-  (bfun', bfun_hoisted) <- Engine.simplifyLambda bfun+  (bfun', bfun_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda bfun   return (Hist w' imgs' ops' bfun', mconcat hoisted <> bfun_hoisted) simplifySOAC (Screma w arrs (ScremaForm scans reds map_lam)) = do   (scans', scans_hoisted) <- fmap unzip $@@ -144,7 +139,7 @@       nes' <- Engine.simplify nes       return (Reduce comm lam' nes', hoisted) -  (map_lam', map_lam_hoisted) <- Engine.simplifyLambda map_lam+  (map_lam', map_lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda map_lam    (,)     <$> ( Screma <$> Engine.simplify w

@@ -996,6 +996,10 @@           return           return +informKernelBody :: Informing rep => KernelBody rep -> KernelBody (Wise rep)+informKernelBody (KernelBody dec stms res) =+  mkWiseKernelBody dec (informStms stms) res+ instance   (CanBeWise (Op rep), ASTRep rep, ASTConstraints lvl) =>   CanBeWise (SegOp lvl rep)@@ -1012,6 +1016,16 @@           return           return +  addOpWisdom = runIdentity . mapSegOpM add+    where+      add =+        SegOpMapper+          return+          (return . informLambda)+          (return . informKernelBody)+          return+          return+ instance ASTRep rep => ST.IndexOp (SegOp lvl rep) where   indexOp vtable k (SegMap _ space _ kbody) is = do     Returns ResultMaySimplify _ se <- maybeNth k $ kernelBodyResult kbody@@ -1124,28 +1138,28 @@ simplifyKernelBody ::   (Engine.SimplifiableRep rep, BodyDec rep ~ ()) =>   SegSpace ->-  KernelBody rep ->+  KernelBody (Wise rep) ->   Engine.SimpleM rep (KernelBody (Wise rep), Stms (Wise rep)) simplifyKernelBody space (KernelBody _ stms res) = do   par_blocker <- Engine.asksEngineEnv $ Engine.blockHoistPar . Engine.envHoistBlockers +  let blocker =+        Engine.hasFree bound_here+          `Engine.orIf` Engine.isOp+          `Engine.orIf` par_blocker+          `Engine.orIf` Engine.isConsumed+   -- Ensure we do not try to use anything that is consumed in the result.-  ((body_stms, body_res), hoisted) <-+  (body_res, body_stms, hoisted) <-     Engine.localVtable (flip (foldl' (flip ST.consume)) (foldMap consumedInResult res))       . Engine.localVtable (<> scope_vtable)       . Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True})       . Engine.enterLoop-      $ Engine.blockIf-        ( Engine.hasFree bound_here-            `Engine.orIf` Engine.isOp-            `Engine.orIf` par_blocker-            `Engine.orIf` Engine.isConsumed-        )-        $ Engine.simplifyStms stms $ do-          res' <--            Engine.localVtable (ST.hideCertified $ namesFromList $ M.keys $ scopeOf stms) $-              mapM Engine.simplify res-          return ((res', UT.usages $ freeIn res'), mempty)+      $ Engine.blockIf blocker stms $ do+        res' <-+          Engine.localVtable (ST.hideCertified $ namesFromList $ M.keys $ scopeOf stms) $+            mapM Engine.simplify res+        pure (res', UT.usages $ freeIn res')    return (mkWiseKernelBody () body_stms body_res, hoisted)   where@@ -1165,7 +1179,7 @@  simplifySegBinOp ::   Engine.SimplifiableRep rep =>-  SegBinOp rep ->+  SegBinOp (Wise rep) ->   Engine.SimpleM rep (SegBinOp (Wise rep), Stms (Wise rep)) simplifySegBinOp (SegBinOp comm lam nes shape) = do   (lam', hoisted) <-@@ -1181,7 +1195,7 @@     BodyDec rep ~ (),     Engine.Simplifiable lvl   ) =>-  SegOp lvl rep ->+  SegOp lvl (Wise rep) ->   Engine.SimpleM rep (SegOp lvl (Wise rep), Stms (Wise rep)) simplifySegOp (SegMap lvl space ts kbody) = do   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)

@@ -1049,16 +1049,12 @@     mkExt _ = Nothing     rettype_st = RetType (mapMaybe mkExt (S.toList (typeDimNames ret)) ++ ret_dims) ret -    (valbind_t, valbind_ext) =-      case pats of-        [] -> (RetType [] $ retType rettype_st, retDims rettype_st)-        _ -> (rettype_st, [])     dec =       ValBind         { valBindEntryPoint = Nothing,           valBindName = fname,           valBindRetDecl = Nothing,-          valBindRetType = Info (valbind_t, valbind_ext),+          valBindRetType = Info rettype_st,           valBindTypeParams = dims',           valBindParams = pats,           valBindBody = body,@@ -1249,7 +1245,7 @@ -- boolean is true if the function is a 'DynamicFun'. defuncValBind :: ValBind -> DefM (ValBind, Env) -- Eta-expand entry points with a functional return type.-defuncValBind (ValBind entry name _ (Info (RetType _ rettype, retext)) tparams params body _ attrs loc)+defuncValBind (ValBind entry name _ (Info (RetType _ rettype)) tparams params body _ attrs loc)   | Scalar Arrow {} <- rettype = do     (body_pats, body', rettype') <- etaExpand (fromStruct rettype) body     defuncValBind $@@ -1257,14 +1253,14 @@         entry         name         Nothing-        (Info (rettype', retext))+        (Info rettype')         tparams         (params <> body_pats)         body'         Nothing         attrs         loc-defuncValBind valbind@(ValBind _ name retdecl (Info (RetType ret_dims rettype, retext)) tparams params body _ _ _) = do+defuncValBind valbind@(ValBind _ name retdecl (Info (RetType ret_dims rettype)) tparams params body _ _ _) = do   when (any isTypeParam tparams) $     error $       prettyName name ++ " has type parameters, "@@ -1279,20 +1275,17 @@         -- applications of lifted functions, we don't properly update         -- the types in the return type annotation.         combineTypeShapes rettype $ first (anyDimIfNotBound bound_sizes) $ toStruct $ typeOf body'+      ret_dims' = filter (`S.member` typeDimNames rettype') ret_dims   (missing_dims, params'') <- sizesForAll bound_sizes params' -  return+  pure     ( valbind         { valBindRetDecl = retdecl,           valBindRetType =-            Info-              ( if null params'-                  then-                    ( RetType [] $ rettype' `setUniqueness` Nonunique,-                      retext-                    )-                  else (RetType ret_dims rettype', retext)-              ),+            Info $+              if null params'+                then RetType ret_dims' $ rettype' `setUniqueness` Nonunique+                else RetType ret_dims' rettype',           valBindTypeParams =             map (`TypeParamDim` mempty) $ tparams' ++ missing_dims,           valBindParams = params'',@@ -1314,7 +1307,7 @@ defuncVals (valbind : ds) = do   (valbind', env) <- defuncValBind valbind   addValBind valbind'-  let globals = valBindName valbind' : snd (unInfo (valBindRetType valbind'))+  let globals = valBindBound valbind'   localEnv env $ areGlobal globals $ defuncVals ds  {-# NOINLINE transformProg #-}

@@ -190,16 +190,14 @@     ModFun f_abs f_closure f_p f_body ->       bindingAbs (f_abs <> S.fromList (unInfo (modParamAbs f_p)))         . extendAbsTypes b_substs-        . extendScope f_closure+        . localScope (const f_closure) -- Start afresh.         . generating         $ do           outer_substs <- scopeSubsts <$> askScope           abs <- asks envAbs           let forward (k, v) = (lookupSubst k outer_substs, v)               p_substs' = M.fromList $ map forward $ M.toList p_substs-              keep k _ =-                k `M.member` p_substs'-                  || k `S.member` abs+              keep k _ = k `M.member` p_substs' || k `S.member` abs               abs_substs =                 M.filterWithKey keep $                   M.map (`lookupSubst` scopeSubsts (modScope arg_mod)) p_substs'@@ -215,7 +213,7 @@             $ do               substs <- scopeSubsts <$> askScope               x <- evalModExp f_body-              return $+              pure $                 addSubsts abs abs_substs $                   -- The next one is dubious, but is necessary to                   -- propagate substitutions from the argument (see@@ -276,14 +274,14 @@ transformExp = transformNames  transformValBind :: ValBind -> TransformM ()-transformValBind (ValBind entry name tdecl (Info (RetType dims t, retext)) tparams params e doc attrs loc) = do+transformValBind (ValBind entry name tdecl (Info (RetType dims t)) tparams params e doc attrs loc) = do   name' <- transformName name   tdecl' <- traverse transformTypeExp tdecl   t' <- transformStructType t   e' <- transformExp e   tparams' <- traverse transformNames tparams   params' <- traverse transformNames params-  emit $ ValDec $ ValBind entry name' tdecl' (Info (RetType dims t', retext)) tparams' params' e' doc attrs loc+  emit $ ValDec $ ValBind entry name' tdecl' (Info (RetType dims t')) tparams' params' e' doc attrs loc  transformTypeBind :: TypeBind -> TransformM () transformTypeBind (TypeBind name l tparams te (Info (RetType dims t)) doc loc) = do@@ -306,13 +304,7 @@         (maybeAscript (srclocOf mb) (modSignature mb) $ modExp mb)         $ modParams mb   mname <- transformName $ modName mb-  abs <- asks envAbs-  -- Copy substitutions involving abstract types out, because they are-  -- always resolved at the outermost level.-  let abs_substs =-        M.filterWithKey (const . flip S.member abs) $-          scopeSubsts $ modScope mod-  return $ Scope abs_substs $ M.singleton mname mod+  return $ Scope (scopeSubsts $ modScope mod) $ M.singleton mname mod  transformDecs :: [Dec] -> TransformM Scope transformDecs ds =

@@ -39,7 +39,7 @@ internaliseFunName = nameFromString . pretty  internaliseValBind :: E.ValBind -> InternaliseM ()-internaliseValBind fb@(E.ValBind entry fname retdecl (Info (rettype, _)) tparams params body _ attrs loc) = do+internaliseValBind fb@(E.ValBind entry fname retdecl (Info rettype) tparams params body _ attrs loc) = do   localConstsScope . bindingFParams tparams params $ \shapeparams params' -> do     let shapenames = map I.paramName shapeparams @@ -94,7 +94,7 @@  generateEntryPoint :: E.EntryPoint -> E.ValBind -> InternaliseM () generateEntryPoint (E.EntryPoint e_params e_rettype) vb = localConstsScope $ do-  let (E.ValBind _ ofname _ (Info (rettype, _)) tparams params _ _ attrs loc) = vb+  let (E.ValBind _ ofname _ (Info rettype) tparams params _ _ attrs loc) = vb   bindingFParams tparams params $ \shapeparams params' -> do     entry_rettype <- internaliseEntryReturnType rettype     let entry' = entryPoint (baseName ofname) (zip e_params params') (e_rettype, entry_rettype)@@ -419,19 +419,27 @@    ctxinit <- argShapes (map I.paramName shapepat) mergepat' mergeinit_ts' -  let ctxmerge = zip shapepat ctxinit-      valmerge = zip mergepat' mergeinit'-      dropCond = case form of+  -- Ensure that the initial loop values match the shapes of the loop+  -- parameters.  XXX: Ideally they should already match (by the+  -- source language type rules), but some of our transformations+  -- (esp. defunctionalisation) strips out some size information.  For+  -- a type-correct source program, these reshapes should simplify+  -- away.+  let args = ctxinit ++ mergeinit'+  args' <-+    ensureArgShapes+      "initial loop values have right shape"+      loc+      (map I.paramName shapepat)+      (map paramType $ shapepat ++ mergepat')+      args++  let dropCond = case form of         E.While {} -> drop 1         _ -> id -  -- Ensure that the result of the loop matches the shapes of the-  -- merge parameters.  XXX: Ideally they should already match (by-  -- the source language type rules), but some of our-  -- transformations (esp. defunctionalisation) strips out some size-  -- information.  For a type-correct source program, these reshapes-  -- should simplify away.-  let merge = ctxmerge ++ valmerge+  -- As above, ensure that the result has the right shape.+  let merge = zip (shapepat ++ mergepat') args'       merge_ts = map (I.paramType . fst) merge   loopbody'' <-     localScope (scopeOfFParams $ map fst merge) . inScopeOf form' . buildBody_ $@@ -439,7 +447,7 @@         . ensureArgShapes           "shape of loop result does not match shapes in loop parameter"           loc-          (map (I.paramName . fst) ctxmerge)+          (map (I.paramName . fst) merge)           merge_ts         . map resSubExp         =<< bodyBind loopbody'@@ -448,7 +456,7 @@   map I.Var . dropCond     <$> attributing       attrs-      (letValExp desc (I.DoLoop (ctxmerge <> valmerge) form' loopbody''))+      (letValExp desc (I.DoLoop merge form' loopbody''))   where     sparams' = map (`TypeParamDim` mempty) sparams 

@@ -51,10 +51,8 @@ addValBind vb = modify $ \s ->   s     { stateValBinds = vb : stateValBinds s,-      stateGlobal = foldl' (flip S.insert) (stateGlobal s) names+      stateGlobal = foldl' (flip S.insert) (stateGlobal s) (valBindBound vb)     }-  where-    names = valBindName vb : snd (unInfo (valBindRetType vb))  replacing :: VName -> Exp -> LiftM a -> LiftM a replacing v e = local $ \env ->@@ -111,7 +109,7 @@         valBindTypeParams = tparams,         valBindParams = free_params ++ params,         valBindRetDecl = Nothing,-        valBindRetType = Info (RetType dims ret, mempty),+        valBindRetType = Info (RetType dims ret),         valBindBody = funbody,         valBindDoc = Nothing,         valBindAttrs = mempty,

@@ -65,7 +65,6 @@         [TypeParam],         [Pat],         StructRetType,-        [VName],         Exp,         [AttrInfo VName],         SrcLoc@@ -300,7 +299,7 @@     -- filter those that are monomorphic versions of the current let-bound     -- function and insert them at this point, and propagate the rest.     rr <- asks envRecordReplacements-    let funbind = PolyBinding rr (fname, tparams, params, ret, [], body, mempty, loc)+    let funbind = PolyBinding rr (fname, tparams, params, ret, body, mempty, loc)     pass $ do       (e', bs) <- listen $ extendEnv fname funbind $ transformExp e       -- Do not remember this one for next time we monomorphise this@@ -625,7 +624,7 @@ -- monomorphic functions with the given expression at the bottom. unfoldLetFuns :: [ValBind] -> Exp -> Exp unfoldLetFuns [] e = e-unfoldLetFuns (ValBind _ fname _ (Info (rettype, _)) dim_params params body _ _ loc : rest) e =+unfoldLetFuns (ValBind _ fname _ (Info rettype) dim_params params body _ _ loc : rest) e =   AppExp (LetFun fname (dim_params, params, Nothing, Info rettype, body) e' loc) (Info $ AppRes e_t mempty)   where     e' = unfoldLetFuns rest e@@ -731,7 +730,7 @@   PolyBinding ->   MonoType ->   MonoM (VName, InferSizeArgs, ValBind)-monomorphiseBinding entry (PolyBinding rr (name, tparams, params, rettype, retext, body, attrs, loc)) inst_t =+monomorphiseBinding entry (PolyBinding rr (name, tparams, params, rettype, body, attrs, loc)) inst_t =   replaceRecordReplacements rr $ do     let bind_t = foldFunType (map patternStructType params) rettype     (substs, t_shape_params) <- typeSubstsM loc (noSizes bind_t) $ noNamedParams inst_t@@ -762,14 +761,14 @@               name'               (shape_params_explicit ++ shape_params_implicit)               params''-              (rettype', retext)+              rettype'               body''           else             toValBinding               name'               shape_params_implicit               (map shapeParam shape_params_explicit ++ params'')-              (rettype', retext)+              rettype'               body''       )   where@@ -879,13 +878,13 @@   PatConstr n (Info tp) ps loc -> PatConstr n (Info $ f tp) ps loc  toPolyBinding :: ValBind -> PolyBinding-toPolyBinding (ValBind _ name _ (Info (rettype, retext)) tparams params body _ attrs loc) =-  PolyBinding mempty (name, tparams, params, rettype, retext, body, attrs, loc)+toPolyBinding (ValBind _ name _ (Info rettype) tparams params body _ attrs loc) =+  PolyBinding mempty (name, tparams, params, rettype, body, attrs, loc)  -- Remove all type variables and type abbreviations from a value binding. removeTypeVariables :: Bool -> ValBind -> MonoM ValBind removeTypeVariables entry valbind = do-  let (ValBind _ _ _ (Info (RetType dims rettype, retext)) _ pats body _ _ _) = valbind+  let (ValBind _ _ _ (Info (RetType dims rettype)) _ pats body _ _ _) = valbind   subs <- asks $ M.map substFromAbbr . envTypeBindings   let mapper =         ASTMapper@@ -904,7 +903,7 @@    return     valbind-      { valBindRetType = Info (applySubst (`M.lookup` subs) $ RetType dims rettype, retext),+      { valBindRetType = Info (applySubst (`M.lookup` subs) $ RetType dims rettype),         valBindParams = map (substPat entry $ applySubst (`M.lookup` subs)) pats,         valBindBody = body'       }@@ -925,7 +924,7 @@       removeTypeVariablesInType $         foldFunType           (map patternStructType (valBindParams valbind))-          $ fst $ unInfo $ valBindRetType valbind+          $ unInfo $ valBindRetType valbind     (name, infer, valbind'') <- monomorphiseBinding True valbind' $ monoType t     tell $ Seq.singleton (name, valbind'' {valBindEntryPoint = valBindEntryPoint valbind})     addLifted (valBindName valbind) (monoType t) (name, infer)

@@ -315,9 +315,15 @@             arg_copy <- newVName (baseString arg <> "_dbcopy")             letBind (Pat [PatElem arg_copy $ MemArray pt shape u $ ArrayIn mem' arg_ixfun]) $               BasicOp $ Copy arg-            pure (param, Var arg_copy)-          _ -> pure (param, Var arg)+            -- We need to make this parameter unique to avoid invalid+            -- hoisting (see #1533), because we are invalidating the+            -- underlying memory.+            pure (fmap mkUnique param, Var arg_copy)+          _ -> pure (fmap mkUnique param, Var arg)     maybeCopyInitial _ (param, arg) = pure (param, arg)++    mkUnique (MemArray bt shape _ ret) = MemArray bt shape Unique ret+    mkUnique x = x  optimiseLoopByCopying :: Constraints rep inner => OptimiseLoop rep optimiseLoopByCopying pat merge body = do

@@ -59,6 +59,8 @@ --    (8) The result of the loop may not alias the merge parameter --    @r1@. --+--    (9) @y@ or its aliases may not be used after the loop.+-- -- FIXME: the implementation is not finished yet.  Specifically, not -- all of the above conditions are checked. module Futhark.Optimise.InPlaceLowering@@ -70,6 +72,7 @@  import Control.Monad.RWS import qualified Data.Map.Strict as M+import Data.Ord (comparing) import Futhark.Analysis.Alias import Futhark.Builder import Futhark.IR.Aliases@@ -78,6 +81,7 @@ import Futhark.IR.Seq (Seq) import Futhark.Optimise.InPlaceLowering.LowerIntoStm import Futhark.Pass+import Futhark.Util (nubByOrd)  -- | Apply the in-place lowering optimisation to the given program. inPlaceLoweringGPU :: Pass GPU GPU@@ -137,14 +141,19 @@   [Stm (Aliases rep)] ->   ForwardingM rep () ->   ForwardingM rep [Stm (Aliases rep)]-optimiseStms [] m = m >> return []+optimiseStms [] m = m >> pure [] optimiseStms (stm : stms) m = do   (stms', bup) <- tapBottomUp $ bindingStm stm $ optimiseStms stms m   stm' <- optimiseInStm stm-  case filter ((`elem` boundHere) . updateValue) $ forwardThese bup of+  -- XXX: unfortunate that we cannot handle duplicate update values.+  -- Would be good to improve this.  See inplacelowering6.fut.+  case nubByOrd (comparing updateValue)+    . filter (not . (`nameIn` bottomUpSeen bup) . updateSource) -- (9)+    . filter ((`elem` boundHere) . updateValue)+    $ forwardThese bup of     [] -> do       checkIfForwardableUpdate stm'-      return $ stm' : stms'+      pure $ stm' : stms'     updates -> do       lower <- asks topLowerUpdate       scope <- askScope@@ -160,13 +169,11 @@       case lower scope stm' updates of         Just lowering -> do           new_stms <- lowering-          new_stms' <--            optimiseStms new_stms $-              tell bup {forwardThese = []}-          return $ new_stms' ++ filter notUpdated stms'+          new_stms' <- optimiseStms new_stms $ tell bup {forwardThese = []}+          pure $ new_stms' ++ filter notUpdated stms'         Nothing -> do           checkIfForwardableUpdate stm'-          return $ stm' : stms'+          pure $ stm' : stms'   where     boundHere = patNames $ stmPat stm @@ -174,7 +181,8 @@       | Pat [PatElem v dec] <- pat,         BasicOp (Update Unsafe src slice (Var ve)) <- e =         maybeForward ve v dec cs src slice-    checkIfForwardableUpdate _ = return ()+    checkIfForwardableUpdate stm' =+      mapM_ seenVar $ namesToList $ freeIn $ stmExp stm'  optimiseInStm :: Constraints rep => Stm (Aliases rep) -> ForwardingM rep (Stm (Aliases rep)) optimiseInStm (Let pat dec e) =

@@ -26,7 +26,7 @@     simplifyFun,   ) import Futhark.Optimise.CSE-import Futhark.Optimise.Simplify.Rep (addScopeWisdom)+import Futhark.Optimise.Simplify.Rep (addScopeWisdom, informStms) import Futhark.Pass import Futhark.Transform.CopyPropagate   ( copyPropagateInFun,@@ -81,9 +81,11 @@                 fdmap $ filter ((`S.member` to_inline_now) . funDefName) dont_inline_in           (vtable', consts') <-             if any (`calledByConsts` cg) to_inline_now-              then-                simplifyConsts . performCSEOnStms True-                  =<< inlineInStms inlinemap consts+              then do+                consts' <-+                  simplifyConsts . performCSEOnStms True+                    =<< inlineInStms inlinemap consts+                pure (ST.insertStms (informStms consts') mempty, consts')               else pure (vtable, consts)            let simplifyFun' fd@@ -94,7 +96,8 @@                 | otherwise =                   copyPropagateInFun simpleSOACS vtable' fd -              onFun = simplifyFun' <=< inlineInFunDef inlinemap+              onFun fd =+                simplifyFun' <=< inlineInFunDef inlinemap $ fd            to_inline_in' <- parMapM onFun to_inline_in 

@@ -1,7 +1,6 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE Strict #-}-{-# LANGUAGE TupleSections #-}  module Futhark.Optimise.Simplify   ( simplifyProg,@@ -21,7 +20,6 @@   ) where -import Data.Bifunctor (second) import qualified Futhark.Analysis.SymbolTable as ST import qualified Futhark.Analysis.UsageTable as UT import Futhark.IR@@ -43,24 +41,22 @@   PassM (Prog rep) simplifyProg simpl rules blockers (Prog consts funs) = do   (consts_vtable, consts') <--    simplifyConsts-      (UT.usages $ foldMap freeIn funs)-      (mempty, consts)+    simplifyConsts (UT.usages $ foldMap freeIn funs) (mempty, informStms consts)    -- We deepen the vtable so it will look like the constants are in an   -- "outer loop"; this communicates useful information to some   -- simplification rules (e.g. seee issue #1302).-  funs' <- parPass (simplifyFun' $ ST.deepen consts_vtable) funs+  funs' <- parPass (simplifyFun' (ST.deepen consts_vtable) . informFunDef) funs   let funs_uses = UT.usages $ foldMap freeIn funs'    (_, consts'') <- simplifyConsts funs_uses (mempty, consts') -  return $ Prog consts'' funs'+  pure $ Prog (fmap removeStmWisdom consts'') (fmap removeFunDefWisdom funs')   where     simplifyFun' consts_vtable =       simplifySomething         (Engine.localVtable (consts_vtable <>) . Engine.simplifyFun)-        removeFunDefWisdom+        id         simpl         rules         blockers@@ -69,18 +65,15 @@     simplifyConsts uses =       simplifySomething         (onConsts uses . snd)-        (second (removeStmWisdom <$>))+        id         simpl         rules         blockers         mempty      onConsts uses consts' = do-      (_, consts'') <--        Engine.simplifyStms consts' (pure ((), mempty))-      (consts''', _) <--        Engine.hoistStms rules (Engine.isFalse False) mempty uses consts''-      return (ST.insertStms consts''' mempty, consts''')+      consts'' <- Engine.simplifyStmsWithUsage uses consts'+      pure (ST.insertStms consts'' mempty, consts'')  -- | Run a simplification operation to convergence. simplifySomething ::@@ -111,7 +104,16 @@   ST.SymbolTable (Wise rep) ->   FunDef rep ->   m (FunDef rep)-simplifyFun = simplifySomething Engine.simplifyFun removeFunDefWisdom+simplifyFun simpl rules blockers vtable fd =+  removeFunDefWisdom+    <$> simplifySomething+      Engine.simplifyFun+      id+      simpl+      rules+      blockers+      vtable+      (informFunDef fd)  -- | Simplify just a single t'Lambda'. simplifyLambda ::@@ -126,14 +128,15 @@   m (Lambda rep) simplifyLambda simpl rules blockers orig_lam = do   vtable <- ST.fromScope . addScopeWisdom <$> askScope-  simplifySomething-    Engine.simplifyLambdaNoHoisting-    removeLambdaWisdom-    simpl-    rules-    blockers-    vtable-    orig_lam+  removeLambdaWisdom+    <$> simplifySomething+      Engine.simplifyLambdaNoHoisting+      id+      simpl+      rules+      blockers+      vtable+      (informLambda orig_lam)  -- | Simplify a sequence of 'Stm's. simplifyStms ::@@ -143,14 +146,14 @@   Engine.HoistBlockers rep ->   Scope rep ->   Stms rep ->-  m (ST.SymbolTable (Wise rep), Stms rep)+  m (Stms rep) simplifyStms simpl rules blockers scope =-  simplifySomething f g simpl rules blockers vtable . (mempty,)+  fmap (fmap removeStmWisdom)+    . simplifySomething f id simpl rules blockers vtable+    . informStms   where     vtable = ST.fromScope $ addScopeWisdom scope-    f (_, stms) =-      Engine.simplifyStms stms ((,mempty) <$> Engine.askVtable)-    g = second $ fmap removeStmWisdom+    f stms = Engine.simplifyStms stms  loopUntilConvergence ::   (MonadFreshNames m, Engine.SimplifiableRep rep) =>@@ -162,4 +165,4 @@   m a loopUntilConvergence env simpl f g x = do   (x', changed) <- modifyNameSource $ Engine.runSimpleM (f x) simpl env-  if changed then loopUntilConvergence env simpl f g (g x') else return $ g x'+  if changed then loopUntilConvergence env simpl f g (g x') else pure $ g x'

@@ -50,13 +50,13 @@     -- * Building blocks     SimplifiableRep,     Simplifiable (..),-    simplifyStms,     simplifyFun,+    simplifyStms,+    simplifyStmsWithUsage,     simplifyLambda,     simplifyLambdaNoHoisting,     bindLParams,     simplifyBody,-    SimplifiedBody,     ST.SymbolTable,     hoistStms,     blockIf,@@ -69,6 +69,7 @@ import Control.Monad.State.Strict import Data.Either import Data.List (find, foldl', mapAccumL)+import qualified Data.Map as M import Data.Maybe import qualified Futhark.Analysis.SymbolTable as ST import qualified Futhark.Analysis.UsageTable as UT@@ -113,6 +114,8 @@  type Protect m = SubExp -> Pat (Rep m) -> Op (Rep m) -> Maybe (m ()) +type SimplifyOp rep op = op -> SimpleM rep (op, Stms (Wise rep))+ data SimpleOps rep = SimpleOps   { mkExpDecS ::       ST.SymbolTable (Wise rep) ->@@ -129,14 +132,12 @@     -- actually be used.     protectHoistedOpS :: Protect (Builder (Wise rep)),     opUsageS :: Op (Wise rep) -> UT.UsageTable,-    simplifyOpS :: SimplifyOp rep (Op rep)+    simplifyOpS :: SimplifyOp rep (Op (Wise rep))   } -type SimplifyOp rep op = op -> SimpleM rep (OpWithWisdom op, Stms (Wise rep))- bindableSimpleOps ::   (SimplifiableRep rep, Buildable rep) =>-  SimplifyOp rep (Op rep) ->+  SimplifyOp rep (Op (Wise rep)) ->   SimpleOps rep bindableSimpleOps =   SimpleOps mkExpDecS' mkBodyS' protectHoistedOpS' (const mempty)@@ -264,21 +265,21 @@   -- | Which side of the branch are we   -- protecting here?   Bool ->-  SimpleM rep (a, Stms (Wise rep)) ->-  SimpleM rep (a, Stms (Wise rep))+  SimpleM rep (Stms (Wise rep), a) ->+  SimpleM rep (Stms (Wise rep), a) protectIfHoisted cond side m = do-  (x, stms) <- m+  (hoisted, x) <- m   ops <- asks $ protectHoistedOpS . fst-  runBuilder $ do-    if not $ all (safeExp . stmExp) stms+  hoisted' <- runBuilder_ $ do+    if not $ all (safeExp . stmExp) hoisted       then do         cond' <-           if side             then return cond             else letSubExp "cond_neg" $ BasicOp $ UnOp Not cond-        mapM_ (protectIf ops unsafeOrCostly cond') stms-      else addStms stms-    return x+        mapM_ (protectIf ops unsafeOrCostly cond') hoisted+      else addStms hoisted+  pure (hoisted', x)   where     unsafeOrCostly e = not (safeExp e) || not (cheapExp e) @@ -289,18 +290,18 @@   SimplifiableRep rep =>   [(FParam (Wise rep), SubExp)] ->   LoopForm (Wise rep) ->-  SimpleM rep (a, Stms (Wise rep)) ->-  SimpleM rep (a, Stms (Wise rep))+  SimpleM rep (a, b, Stms (Wise rep)) ->+  SimpleM rep (a, b, Stms (Wise rep)) protectLoopHoisted merge form m = do-  (x, stms) <- m+  (x, y, stms) <- m   ops <- asks $ protectHoistedOpS . fst-  runBuilder $ do+  stms' <- runBuilder_ $ do     if not $ all (safeExp . stmExp) stms       then do         is_nonempty <- checkIfNonEmpty         mapM_ (protectIf ops (not . safeExp) is_nonempty) stms       else addStms stms-    return x+  pure (x, y, stms')   where     checkIfNonEmpty =       case form of@@ -387,67 +388,115 @@ notWorthHoisting _ _ (Let pat _ e) =   not (safeExp e) && any ((> 0) . arrayRank) (patTypes pat) +-- Top-down simplify a statement (including copy propagation into the+-- pattern and such).  Does not recurse into any sub-Bodies or Ops.+nonrecSimplifyStm ::+  SimplifiableRep rep =>+  Stm (Wise rep) ->+  SimpleM rep (Stm (Wise rep))+nonrecSimplifyStm (Let pat (StmAux cs attrs (_, dec)) e) = do+  cs' <- simplify cs+  (pat', pat_cs) <- collectCerts $ simplifyPat $ removePatWisdom pat+  let aux' = StmAux (cs' <> pat_cs) attrs dec+  mkWiseLetStm pat' aux' <$> simplifyExpBase e++-- Bottom-up simplify a statement.  Recurses into sub-Bodies and Ops.+-- Does not copy-propagate into the pattern and similar, as it is+-- assumed 'nonrecSimplifyStm' has already touched it (and worst case,+-- it'll get it on the next round of the overall fixpoint iteration.)+recSimplifyStm ::+  SimplifiableRep rep =>+  Stm (Wise rep) ->+  UT.UsageTable ->+  SimpleM rep (Stms (Wise rep), Stm (Wise rep))+recSimplifyStm (Let pat (StmAux cs attrs (_, dec)) e) usage = do+  ((e', e_hoisted), e_cs) <- collectCerts $ simplifyExp usage pat e+  let aux' = StmAux (cs <> e_cs) attrs dec+  pure (e_hoisted, mkWiseLetStm (removePatWisdom pat) aux' e')+ hoistStms ::   SimplifiableRep rep =>   RuleBook (Wise rep) ->   BlockPred (Wise rep) ->-  ST.SymbolTable (Wise rep) ->-  UT.UsageTable ->   Stms (Wise rep) ->-  SimpleM-    rep-    ( Stms (Wise rep),-      Stms (Wise rep)-    )-hoistStms rules block vtable uses orig_stms = do-  (blocked, hoisted) <- simplifyStmsBottomUp vtable uses orig_stms+  SimpleM rep (a, UT.UsageTable) ->+  SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))+hoistStms rules block orig_stms final = do+  (a, blocked, hoisted) <- simplifyStmsBottomUp orig_stms   unless (null hoisted) changed-  return (stmsFromList blocked, stmsFromList hoisted)+  pure (a, stmsFromList blocked, stmsFromList hoisted)   where-    simplifyStmsBottomUp vtable' uses' stms = do-      (_, stms') <- simplifyStmsBottomUp' vtable' uses' stms-      -- We need to do a final pass to ensure that nothing is-      -- hoisted past something that it depends on.-      let (blocked, hoisted) = partitionEithers $ blockUnhoistedDeps stms'-      return (blocked, hoisted)--    simplifyStmsBottomUp' vtable' uses' stms = do+    simplifyStmsBottomUp stms = do       opUsage <- asks $ opUsageS . fst       let usageInStm stm =             UT.usageInStm stm               <> case stmExp stm of                 Op op -> opUsage op                 _ -> mempty-      foldM (hoistable usageInStm) (uses', []) $ reverse $ zip (stmsToList stms) vtables-      where-        vtables = scanl (flip ST.insertStm) vtable' $ stmsToList stms+      (x, _, stms') <- hoistableStms usageInStm stms+      -- We need to do a final pass to ensure that nothing is+      -- hoisted past something that it depends on.+      let (blocked, hoisted) = partitionEithers $ blockUnhoistedDeps stms'+      pure (x, blocked, hoisted) -    hoistable usageInStm (uses', stms) (stm, vtable')-      | not $ any (`UT.isUsedDirectly` uses') $ provides stm =-        -- Dead statement.-        return (uses', stms)-      | otherwise = do-        res <--          localVtable (const vtable') $-            bottomUpSimplifyStm rules (vtable', uses') stm-        case res of-          Nothing -- Nothing to optimise - see if hoistable.-            | block vtable' uses' stm ->-              return-                ( expandUsage usageInStm vtable' uses' stm-                    `UT.without` provides stm,-                  Left stm : stms-                )-            | otherwise ->-              return-                ( expandUsage usageInStm vtable' uses' stm,-                  Right stm : stms-                )-          Just optimstms -> do-            changed-            (uses'', stms') <- simplifyStmsBottomUp' vtable' uses' optimstms-            return (uses'', stms' <> stms)+    descend usageInStm stms m =+      case stmsHead stms of+        Nothing -> m+        Just (stms_h, stms_t) -> localVtable (ST.insertStm stms_h) $ do+          (x, usage, stms_t') <- descend usageInStm stms_t m+          process usageInStm stms_h stms_t' usage x +    process usageInStm stm stms usage x = do+      vtable <- askVtable+      res <- bottomUpSimplifyStm rules (vtable, usage) stm+      case res of+        Nothing -- Nothing to optimise - see if hoistable.+          | block vtable usage stm ->+            -- No, not hoistable.+            pure+              ( x,+                expandUsage usageInStm vtable usage stm+                  `UT.without` provides stm,+                Left stm : stms+              )+          | otherwise ->+            -- Yes, hoistable.+            pure+              ( x,+                expandUsage usageInStm vtable usage stm,+                Right stm : stms+              )+        Just optimstms -> do+          changed+          descend usageInStm optimstms $ pure (x, usage, stms)++    hoistableStms usageInStm stms =+      case stmsHead stms of+        Nothing -> do+          (x, usage) <- final+          pure (x, usage, mempty)+        Just (stms_h, stms_t) -> do+          stms_h' <- nonrecSimplifyStm stms_h++          vtable <- askVtable+          simplified <- topDownSimplifyStm rules vtable stms_h'++          case simplified of+            Just newstms -> do+              changed+              hoistableStms usageInStm (newstms <> stms_t)+            Nothing -> do+              (x, usage, stms_t') <-+                localVtable (ST.insertStm stms_h') $+                  hoistableStms usageInStm stms_t+              if not $ any (`UT.isUsedDirectly` usage) $ provides stms_h'+                then -- Dead statement.+                  pure (x, usage, stms_t')+                else do+                  (stms_h_stms, stms_h'') <- recSimplifyStm stms_h' usage+                  descend usageInStm stms_h_stms $+                    process usageInStm stms_h'' stms_t' usage x+ blockUnhoistedDeps ::   ASTRep rep =>   [Either (Stm rep) (Stm rep)] ->@@ -522,19 +571,15 @@     addStms stms     pure res -type SimplifiedBody rep a = ((a, UT.UsageTable), Stms (Wise rep))- blockIf ::   SimplifiableRep rep =>   BlockPred (Wise rep) ->-  SimpleM rep (SimplifiedBody rep a) ->-  SimpleM rep ((Stms (Wise rep), a), Stms (Wise rep))-blockIf block m = do-  ((x, usages), stms) <- m-  vtable <- askVtable+  Stms (Wise rep) ->+  SimpleM rep (a, UT.UsageTable) ->+  SimpleM rep (a, Stms (Wise rep), Stms (Wise rep))+blockIf block stms m = do   rules <- asksEngineEnv envRules-  (blocked, hoisted) <- hoistStms rules block vtable usages stms-  return ((blocked, x), hoisted)+  hoistStms rules block stms m  hasFree :: ASTRep rep => Names -> BlockPred rep hasFree ks _ _ need = ks `namesIntersect` freeIn need@@ -579,22 +624,24 @@  hoistCommon ::   SimplifiableRep rep =>+  UT.UsageTable ->+  [UT.Usages] ->   SubExp ->   IfSort ->-  SimplifiedBody rep Result ->-  SimplifiedBody rep Result ->+  Body (Wise rep) ->+  Body (Wise rep) ->   SimpleM     rep     ( Body (Wise rep),       Body (Wise rep),       Stms (Wise rep)     )-hoistCommon cond ifsort ((res1, usages1), stms1) ((res2, usages2), stms2) = do+hoistCommon res_usage res_usages cond ifsort body1 body2 = do   is_alloc_fun <- asksEngineEnv $ isAllocation . envHoistBlockers   branch_blocker <- asksEngineEnv $ blockHoistBranch . envHoistBlockers   vtable <- askVtable   let -- We are unwilling to hoist things that are unsafe or costly,-+      -- except if they are invariant to the most enclosing loop,       -- because in that case they will also be hoisted past that       -- loop.       --@@ -616,8 +663,8 @@       -- possible.       isNotHoistableBnd _ _ (Let _ _ (BasicOp ArrayLit {})) = False       isNotHoistableBnd _ _ (Let _ _ (BasicOp SubExp {})) = False-      isNotHoistableBnd _ usages (Let pat _ _)-        | any (`UT.isSize` usages) $ patNames pat =+      isNotHoistableBnd _ usage (Let pat _ _)+        | any (`UT.isSize` usage) $ patNames pat =           False       isNotHoistableBnd _ _ stm         | is_alloc_fun stm = False@@ -631,39 +678,55 @@           `orIf` isConsuming           `orIf` isNotHoistableBnd -  rules <- asksEngineEnv envRules-  (body1_stms', safe1) <-+  (hoisted1, body1') <-     protectIfHoisted cond True $-      hoistStms rules block vtable usages1 stms1-  (body2_stms', safe2) <-+      simplifyBody block res_usage res_usages body1+  (hoisted2, body2') <-     protectIfHoisted cond False $-      hoistStms rules block vtable usages2 stms2-  let hoistable = safe1 <> safe2-  body1' <- constructBody body1_stms' res1-  body2' <- constructBody body2_stms' res2-  return (body1', body2', hoistable)+      simplifyBody block res_usage res_usages body2+  return (body1', body2', hoisted1 <> hoisted2) --- | Simplify a single body.  The @[Diet]@ only covers the value--- elements, because the context cannot be consumed.+-- | Simplify a single body. simplifyBody ::   SimplifiableRep rep =>-  [Diet] ->-  Body rep ->-  SimpleM rep (SimplifiedBody rep Result)-simplifyBody ds (Body _ stms res) =-  simplifyStms stms $ do-    res' <- simplifyResult ds res-    return (res', mempty)+  BlockPred (Wise rep) ->+  UT.UsageTable ->+  [UT.Usages] ->+  Body (Wise rep) ->+  SimpleM rep (Stms (Wise rep), Body (Wise rep))+simplifyBody blocker usage res_usages (Body _ stms res) = do+  (res', stms', hoisted) <-+    blockIf blocker stms $ do+      (res', res_usage) <- simplifyResult res_usages res+      pure (res', res_usage <> usage)+  body' <- constructBody stms' res'+  pure (hoisted, body') +-- | Simplify a single body.+simplifyBodyNoHoisting ::+  SimplifiableRep rep =>+  UT.UsageTable ->+  [UT.Usages] ->+  Body (Wise rep) ->+  SimpleM rep (Body (Wise rep))+simplifyBodyNoHoisting usage res_usages body =+  snd <$> simplifyBody (isFalse False) usage res_usages body++usageFromDiet :: Diet -> UT.Usages+usageFromDiet Consume = UT.consumedU+usageFromDiet _ = mempty+ -- | Simplify a single 'Result'.  The @[Diet]@ only covers the value -- elements, because the context cannot be consumed. simplifyResult ::-  SimplifiableRep rep => [Diet] -> Result -> SimpleM rep (Result, UT.UsageTable)-simplifyResult ds res = do+  SimplifiableRep rep => [UT.Usages] -> Result -> SimpleM rep (Result, UT.UsageTable)+simplifyResult usages res = do   res' <- mapM simplify res   vtable <- askVtable-  let consumption = consumeResult vtable $ zip ds res'-  return (res', UT.usages (freeIn res') <> consumption)+  let more_usages = mconcat $ do+        (u, Var v) <- zip usages $ map resSubExp res+        map (`UT.usage` u) $ v : namesToList (ST.lookupAliases v vtable)+  return (res', UT.usages (freeIn res') <> more_usages)  isDoLoopResult :: Result -> UT.UsageTable isDoLoopResult = mconcat . map checkForVar@@ -673,77 +736,49 @@  simplifyStms ::   SimplifiableRep rep =>-  Stms rep ->-  SimpleM rep (a, Stms (Wise rep)) ->-  SimpleM rep (a, Stms (Wise rep))-simplifyStms stms m =-  case stmsHead stms of-    Nothing -> inspectStms mempty m-    Just (Let pat (StmAux stm_cs attrs dec) e, stms') -> do-      stm_cs' <- simplify stm_cs-      ((e', e_stms), e_cs) <- collectCerts $ simplifyExp e-      (pat', pat_cs) <- collectCerts $ simplifyPat pat-      let cs = stm_cs' <> e_cs <> pat_cs-      inspectStms e_stms $-        inspectStm (mkWiseLetStm pat' (StmAux cs attrs dec) e') $-          simplifyStms stms' m--inspectStm ::-  SimplifiableRep rep =>-  Stm (Wise rep) ->-  SimpleM rep (a, Stms (Wise rep)) ->-  SimpleM rep (a, Stms (Wise rep))-inspectStm = inspectStms . oneStm+  Stms (Wise rep) ->+  SimpleM rep (Stms (Wise rep))+simplifyStms stms = do+  simplifyStmsWithUsage all_used stms+  where+    all_used =+      UT.usages (namesFromList (M.keys (scopeOf stms))) -inspectStms ::+simplifyStmsWithUsage ::   SimplifiableRep rep =>+  UT.UsageTable ->   Stms (Wise rep) ->-  SimpleM rep (a, Stms (Wise rep)) ->-  SimpleM rep (a, Stms (Wise rep))-inspectStms stms m =-  case stmsHead stms of-    Nothing -> m-    Just (stm, stms') -> do-      vtable <- askVtable-      rules <- asksEngineEnv envRules-      simplified <- topDownSimplifyStm rules vtable stm-      case simplified of-        Just newstms -> changed >> inspectStms (newstms <> stms') m-        Nothing -> do-          (x, stms'') <- localVtable (ST.insertStm stm) $ inspectStms stms' m-          return (x, oneStm stm <> stms'')+  SimpleM rep (Stms (Wise rep))+simplifyStmsWithUsage usage stms = do+  ((), stms', _) <- blockIf (isFalse False) stms $ pure ((), usage)+  pure stms' -simplifyOp :: Op rep -> SimpleM rep (Op (Wise rep), Stms (Wise rep))+simplifyOp :: Op (Wise rep) -> SimpleM rep (Op (Wise rep), Stms (Wise rep)) simplifyOp op = do   f <- asks $ simplifyOpS . fst   f op  simplifyExp ::   SimplifiableRep rep =>-  Exp rep ->+  UT.UsageTable ->+  Pat (Wise rep) ->+  Exp (Wise rep) ->   SimpleM rep (Exp (Wise rep), Stms (Wise rep))-simplifyExp (If cond tbranch fbranch (IfDec ts ifsort)) = do+simplifyExp usage (Pat pes) (If cond tbranch fbranch (IfDec ts ifsort)) = do   -- Here, we have to check whether 'cond' puts a bound on some free   -- variable, and if so, chomp it.  We should also try to do CSE   -- across branches.+  let pes_usages = map (fromMaybe mempty . (`UT.lookup` usage) . patElemName) pes   cond' <- simplify cond   ts' <- mapM simplify ts-  -- FIXME: we have to be conservative about the diet here, because we-  -- lack proper ifnormation.  Something is wrong with the order in-  -- which the simplifier does things - it should be purely bottom-up-  -- (or else, If expressions should indicate explicitly the diet of-  -- their return types).-  let ds = map (const Consume) ts-  tbranch' <- simplifyBody ds tbranch-  fbranch' <- simplifyBody ds fbranch-  (tbranch'', fbranch'', hoisted) <- hoistCommon cond' ifsort tbranch' fbranch'-  return (If cond' tbranch'' fbranch'' $ IfDec ts' ifsort, hoisted)-simplifyExp (DoLoop merge form loopbody) = do+  (tbranch', fbranch', hoisted) <-+    hoistCommon usage pes_usages cond' ifsort tbranch fbranch+  return (If cond' tbranch' fbranch' $ IfDec ts' ifsort, hoisted)+simplifyExp _ _ (DoLoop merge form loopbody) = do   let (params, args) = unzip merge   params' <- mapM (traverse simplify) params   args' <- mapM simplify args   let merge' = zip params' args'-      diets = map (diet . paramDeclType) params'   (form', boundnames, wrapbody) <- case form of     ForLoop loopvar it boundexp loopvars -> do       boundexp' <- simplify boundexp@@ -766,18 +801,22 @@           protectLoopHoisted merge' (WhileLoop cond')         )   seq_blocker <- asksEngineEnv $ blockHoistSeq . envHoistBlockers-  ((loopstms, loopres), hoisted) <-+  (loopres, loopstms, hoisted) <-     enterLoop . consumeMerge $-      bindMerge (zipWith withRes merge' (bodyResult loopbody)) $-        wrapbody $-          blockIf-            ( hasFree boundnames `orIf` isConsumed-                `orIf` seq_blocker-                `orIf` notWorthHoisting-            )-            $ do-              ((res, uses), stms) <- simplifyBody diets loopbody-              return ((res, uses <> isDoLoopResult res), stms)+      bindMerge (zipWith withRes merge' (bodyResult loopbody)) . wrapbody $+        blockIf+          ( hasFree boundnames `orIf` isConsumed+              `orIf` seq_blocker+              `orIf` notWorthHoisting+          )+          (bodyStms loopbody)+          $ do+            let params_usages =+                  map+                    (\p -> if unique (paramDeclType p) then UT.consumedU else mempty)+                    params'+            (res, uses) <- simplifyResult params_usages $ bodyResult loopbody+            pure (res, uses <> isDoLoopResult res)   loopbody' <- constructBody loopstms loopres   return (DoLoop merge' form' loopbody', hoisted)   where@@ -788,10 +827,10 @@     consumed_by_merge =       freeIn $ map snd $ filter (unique . paramDeclType . fst) merge     withRes (p, x) y = (p, x, y)-simplifyExp (Op op) = do+simplifyExp _ _ (Op op) = do   (op', stms) <- simplifyOp op   return (Op op', stms)-simplifyExp (WithAcc inputs lam) = do+simplifyExp usage _ (WithAcc inputs lam) = do   (inputs', inputs_stms) <- fmap unzip . forM inputs $ \(shape, arrs, op) -> do     (op', op_stms) <- case op of       Nothing ->@@ -801,49 +840,32 @@         nes' <- simplify nes         return (Just (op_lam', nes'), op_lam_stms)     (,op_stms) <$> ((,,op') <$> simplify shape <*> simplify arrs)-  (lam', lam_stms) <--    simplifyLambdaWithBody (isFalse True) lam $-      localVtable (ST.noteAccTokens (zip (map paramName (lambdaParams lam)) inputs')) $-        simplifyBody (map (const Observe) (lambdaReturnType lam)) $-          lambdaBody lam+  let noteAcc = ST.noteAccTokens (zip (map paramName (lambdaParams lam)) inputs')+  (lam', lam_stms) <- simplifyLambdaWith noteAcc (isFalse True) usage lam   pure (WithAcc inputs' lam', mconcat inputs_stms <> lam_stms)+simplifyExp _ _ e = do+  e' <- simplifyExpBase e+  pure (e', mempty) +-- The simple nonrecursive case that we can perform without bottom-up+-- information.+simplifyExpBase :: SimplifiableRep rep => Exp (Wise rep) -> SimpleM rep (Exp (Wise rep)) -- Special case for simplification of commutative BinOps where we -- arrange the operands in sorted order.  This can make expressions -- more identical, which helps CSE.-simplifyExp (BasicOp (BinOp op x y))+simplifyExpBase (BasicOp (BinOp op x y))   | commutativeBinOp op = do     x' <- simplify x     y' <- simplify y-    return (BasicOp $ BinOp op (min x' y') (max x' y'), mempty)-simplifyExp e = do-  e' <- simplifyExpBase e-  return (e', mempty)--simplifyExpBase ::-  SimplifiableRep rep =>-  Exp rep ->-  SimpleM rep (Exp (Wise rep))-simplifyExpBase = mapExpM hoist+    pure $ BasicOp $ BinOp op (min x' y') (max x' y')+simplifyExpBase e = mapExpM hoist e   where     hoist =-      Mapper-        { -- Bodies are handled explicitly because we need to-          -- provide their result diet.-          mapOnBody =-            error "Unhandled body in simplification engine.",-          mapOnSubExp = simplify,-          -- Lambdas are handled explicitly because we need to-          -- bind their parameters.+      identityMapper+        { mapOnSubExp = simplify,           mapOnVName = simplify,           mapOnRetType = simplify,-          mapOnBranchType = simplify,-          mapOnFParam =-            error "Unhandled FParam in simplification engine.",-          mapOnLParam =-            error "Unhandled LParam in simplification engine.",-          mapOnOp =-            error "Unhandled Op in simplification engine."+          mapOnBranchType = simplify         }  type SimplifiableRep rep =@@ -960,51 +982,47 @@  simplifyLambda ::   SimplifiableRep rep =>-  Lambda rep ->+  Lambda (Wise rep) ->   SimpleM rep (Lambda (Wise rep), Stms (Wise rep)) simplifyLambda lam = do   par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers-  simplifyLambdaMaybeHoist par_blocker lam+  simplifyLambdaMaybeHoist par_blocker mempty lam  simplifyLambdaNoHoisting ::   SimplifiableRep rep =>-  Lambda rep ->+  Lambda (Wise rep) ->   SimpleM rep (Lambda (Wise rep)) simplifyLambdaNoHoisting lam =-  fst <$> simplifyLambdaMaybeHoist (isFalse False) lam+  fst <$> simplifyLambdaMaybeHoist (isFalse False) mempty lam  simplifyLambdaMaybeHoist ::   SimplifiableRep rep =>   BlockPred (Wise rep) ->-  Lambda rep ->+  UT.UsageTable ->+  Lambda (Wise rep) ->   SimpleM rep (Lambda (Wise rep), Stms (Wise rep))-simplifyLambdaMaybeHoist blocked lam =-  simplifyLambdaWithBody blocked lam $-    simplifyBody (map (const Observe) (lambdaReturnType lam)) $ lambdaBody lam+simplifyLambdaMaybeHoist = simplifyLambdaWith id -simplifyLambdaWithBody ::+simplifyLambdaWith ::   SimplifiableRep rep =>+  (ST.SymbolTable (Wise rep) -> ST.SymbolTable (Wise rep)) ->   BlockPred (Wise rep) ->-  Lambda rep ->-  SimpleM rep (SimplifiedBody rep Result) ->+  UT.UsageTable ->+  Lambda (Wise rep) ->   SimpleM rep (Lambda (Wise rep), Stms (Wise rep))-simplifyLambdaWithBody blocked lam@(Lambda params _body rettype) m = do+simplifyLambdaWith f blocked usage lam@(Lambda params body rettype) = do   params' <- mapM (traverse simplify) params   let paramnames = namesFromList $ boundByLambda lam-  ((lamstms, lamres), hoisted) <--    enterLoop . bindLParams params' $-      blockIf (blocked `orIf` hasFree paramnames `orIf` isConsumed) m-  body' <- constructBody lamstms lamres+  (hoisted, body') <-+    bindLParams params' . localVtable f $+      simplifyBody+        (blocked `orIf` hasFree paramnames `orIf` isConsumed)+        usage+        (map (const mempty) rettype)+        body   rettype' <- simplify rettype   return (Lambda params' body' rettype', hoisted) -consumeResult :: ST.SymbolTable rep -> [(Diet, SubExpRes)] -> UT.UsageTable-consumeResult vtable = mconcat . map inspect-  where-    inspect (Consume, SubExpRes _ (Var v)) =-      mconcat $ map UT.consumedUsage $ v : namesToList (ST.lookupAliases v vtable)-    inspect _ = mempty- instance Simplifiable Certs where   simplify (Certs ocs) = Certs . nubOrd . concat <$> mapM check ocs     where@@ -1015,19 +1033,13 @@           Just (Var idd', _) -> return [idd']           _ -> return [idd] -insertAllStms ::-  SimplifiableRep rep =>-  SimpleM rep (SimplifiedBody rep Result) ->-  SimpleM rep (Body (Wise rep))-insertAllStms = uncurry constructBody . fst <=< blockIf (isFalse False)- simplifyFun ::   SimplifiableRep rep =>-  FunDef rep ->+  FunDef (Wise rep) ->   SimpleM rep (FunDef (Wise rep)) simplifyFun (FunDef entry attrs fname rettype params body) = do   rettype' <- simplify rettype   params' <- mapM (traverse simplify) params-  let ds = map (diet . declExtTypeOf) rettype'-  body' <- bindFParams params $ insertAllStms $ simplifyBody ds body-  return $ FunDef entry attrs fname rettype' params' body'+  let usages = map (usageFromDiet . diet . declExtTypeOf) rettype'+  body' <- bindFParams params $ simplifyBodyNoHoisting mempty usages body+  pure $ FunDef entry attrs fname rettype' params' body'

@@ -1,3 +1,4 @@+{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-}@@ -21,6 +22,12 @@     mkWiseLetStm,     mkWiseExpDec,     CanBeWise (..),++    -- * Constructing representation+    Informing,+    informLambda,+    informFunDef,+    informStms,   ) where @@ -258,6 +265,10 @@     let Body bodyrep _ _ = mkBody (fmap removeStmWisdom stms) res      in mkWiseBody bodyrep stms res +-- | Constraints that let us transform a representation into a 'Wise'+-- representation.+type Informing rep = (ASTRep rep, CanBeWise (Op rep))+ class   ( AliasedOp (OpWithWisdom op),     IsOp (OpWithWisdom op)@@ -266,7 +277,47 @@   where   type OpWithWisdom op :: Data.Kind.Type   removeOpWisdom :: OpWithWisdom op -> op+  addOpWisdom :: op -> OpWithWisdom op  instance CanBeWise () where   type OpWithWisdom () = ()   removeOpWisdom () = ()+  addOpWisdom () = ()++informStm :: Informing rep => Stm rep -> Stm (Wise rep)+informStm (Let pat aux e) = mkWiseLetStm pat aux $ informExp e++informStms :: Informing rep => Stms rep -> Stms (Wise rep)+informStms = fmap informStm++informBody :: Informing rep => Body rep -> Body (Wise rep)+informBody (Body dec stms res) = mkWiseBody dec (informStms stms) res++informLambda :: Informing rep => Lambda rep -> Lambda (Wise rep)+informLambda (Lambda ps body ret) = Lambda ps (informBody body) ret++informExp :: Informing rep => Exp rep -> Exp (Wise rep)+informExp (If cond tbranch fbranch (IfDec ts ifsort)) =+  If cond (informBody tbranch) (informBody fbranch) (IfDec ts ifsort)+informExp (DoLoop merge form loopbody) =+  let form' = case form of+        ForLoop i it bound params -> ForLoop i it bound params+        WhileLoop cond -> WhileLoop cond+   in DoLoop merge form' $ informBody loopbody+informExp e = runIdentity $ mapExpM mapper e+  where+    mapper =+      Mapper+        { mapOnBody = const $ pure . informBody,+          mapOnSubExp = pure,+          mapOnVName = pure,+          mapOnRetType = pure,+          mapOnBranchType = pure,+          mapOnFParam = pure,+          mapOnLParam = pure,+          mapOnOp = pure . addOpWisdom+        }++informFunDef :: Informing rep => FunDef rep -> FunDef (Wise rep)+informFunDef (FunDef entry attrs fname rettype params body) =+  FunDef entry attrs fname rettype params $ informBody body

@@ -279,6 +279,7 @@       foldMap consumedInStm $ fst $ Alias.analyseStms mempty prestms      mustBeInlinedExp (BasicOp (Index _ slice)) = not $ null $ sliceDims slice+    mustBeInlinedExp (BasicOp Iota {}) = True     mustBeInlinedExp (BasicOp Rotate {}) = True     mustBeInlinedExp (BasicOp Rearrange {}) = True     mustBeInlinedExp (BasicOp Reshape {}) = True@@ -437,11 +438,10 @@   tilingSegMap tiling "prelude" (scalarLevel tiling) ResultPrivate $     \in_bounds slice -> do       ts <- mapM lookupType prestms_live-      fmap varsRes $-        protectOutOfBounds "pre" in_bounds ts $ do-          addPrivStms slice privstms-          addStms prestms-          pure $ varsRes prestms_live+      fmap varsRes . protectOutOfBounds "pre" in_bounds ts $ do+        addPrivStms slice privstms+        addStms prestms+        pure $ varsRes prestms_live  liveSet :: FreeIn a => Stms GPU -> a -> Names liveSet stms after =@@ -795,7 +795,7 @@       fmap varsRes $         case kind of           TilePartial ->-            letTupExp "pre"+            letTupExp "pre1d"               =<< eIf                 (toExp $ pe64 j .<. pe64 w)                 (resultBody <$> mapM (fmap Var . readTileElem) arrs)@@ -1056,7 +1056,7 @@       fmap varsRes $         case kind of           TilePartial ->-            mapM (letExp "pre" <=< readTileElemIfInBounds) arrs_and_perms+            mapM (letExp "pre2d" <=< readTileElemIfInBounds) arrs_and_perms           TileFull ->             mapM readTileElem arrs_and_perms 

@@ -485,8 +485,7 @@     sliceKernelSizes num_threads variant_sizes kspace kstms   -- Note the recursive call to expand allocations inside the newly   -- produced kernels.-  (_, slice_stms_tmp) <--    simplifyStms =<< explicitAllocationsInStms slice_stms+  slice_stms_tmp <- simplifyStms =<< explicitAllocationsInStms slice_stms   slice_stms' <- transformStms slice_stms_tmp    let variant_allocs' :: [(VName, (SubExp, SubExp, Space))]

@@ -1095,7 +1095,7 @@     Mem rep inner   ) =>   (Engine.OpWithWisdom inner -> UT.UsageTable) ->-  (inner -> Engine.SimpleM rep (Engine.OpWithWisdom inner, Stms (Engine.Wise rep))) ->+  (Engine.OpWithWisdom inner -> Engine.SimpleM rep (Engine.OpWithWisdom inner, Stms (Engine.Wise rep))) ->   SimpleOps rep simplifiable innerUsage simplifyInnerOp =   SimpleOps mkExpDecS' mkBodyS' protectOp opUsage simplifyOp

@@ -393,7 +393,7 @@     let map_lam_sequential = soacsLambdaToGPU map_lam     lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt     addStms . fmap (certify cs)-      =<< segScan lvl res_pat w scan_ops map_lam_sequential arrs [] []+      =<< segScan lvl res_pat mempty w scan_ops map_lam_sequential arrs [] [] transformStm path (Let res_pat aux (Op (Screma w arrs form)))   | Just [Reduce comm red_fun nes] <- isReduceSOAC form,     let comm'@@ -401,7 +401,7 @@           | otherwise = comm,     Just do_irwim <- irwim res_pat w comm' red_fun $ zip nes arrs = do     types <- asksScope scopeForSOACs-    (_, stms) <- fst <$> runBuilderT (simplifyStms =<< collectStms_ (auxing aux do_irwim)) types+    stms <- fst <$> runBuilderT (simplifyStms =<< collectStms_ (auxing aux do_irwim)) types     transformStms path $ stmsToList stms transformStm path (Let pat aux@(StmAux cs _ _) (Op (Screma w arrs form)))   | Just (reds, map_lam) <- isRedomapSOAC form = do@@ -426,10 +426,8 @@           (mapstm, redstm) <-             redomapToMapAndReduce pat (w, reds, map_lam, arrs)           types <- asksScope scopeForSOACs-          transformStms path' . stmsToList <=< (`runBuilderT_` types) $ do-            (_, stms) <--              simplifyStms (stmsFromList [certify cs mapstm, certify cs redstm])-            addStms stms+          transformStms path' . stmsToList <=< (`runBuilderT_` types) $+            addStms =<< simplifyStms (stmsFromList [certify cs mapstm, certify cs redstm])          innerParallelBody path' =           renameBody

@@ -51,24 +51,25 @@  prepareRedOrScan ::   (MonadBuilder m, DistRep (Rep m)) =>+  Certs ->   SubExp ->   Lambda (Rep m) ->   [VName] ->   [(VName, SubExp)] ->   [KernelInput] ->   m (SegSpace, KernelBody (Rep m))-prepareRedOrScan w map_lam arrs ispace inps = do+prepareRedOrScan cs w map_lam arrs ispace inps = do   gtid <- newVName "gtid"   space <- mkSegSpace $ ispace ++ [(gtid, w)]   kbody <- fmap (uncurry (flip (KernelBody ()))) $     runBuilder $       localScope (scopeOfSegSpace space) $ do         mapM_ readKernelInput inps-        mapM_ readKernelInput $ do+        certifying cs . mapM_ readKernelInput $ do           (p, arr) <- zip (lambdaParams map_lam) arrs           pure $ KernelInput (paramName p) (paramType p) arr [Var gtid]         res <- bodyBind (lambdaBody map_lam)-        forM res $ \(SubExpRes cs se) -> pure $ Returns ResultMaySimplify cs se+        forM res $ \(SubExpRes res_cs se) -> pure $ Returns ResultMaySimplify res_cs se    return (space, kbody) @@ -76,6 +77,7 @@   (MonadFreshNames m, DistRep rep, HasScope rep m) =>   SegOpLevel rep ->   Pat rep ->+  Certs ->   SubExp -> -- segment size   [SegBinOp rep] ->   Lambda rep ->@@ -83,8 +85,8 @@   [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this reduction   [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace   m (Stms rep)-segRed lvl pat w ops map_lam arrs ispace inps = runBuilder_ $ do-  (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps+segRed lvl pat cs w ops map_lam arrs ispace inps = runBuilder_ $ do+  (kspace, kbody) <- prepareRedOrScan cs w map_lam arrs ispace inps   letBind pat $     Op $       segOp $@@ -94,6 +96,7 @@   (MonadFreshNames m, DistRep rep, HasScope rep m) =>   SegOpLevel rep ->   Pat rep ->+  Certs ->   SubExp -> -- segment size   [SegBinOp rep] ->   Lambda rep ->@@ -101,8 +104,8 @@   [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this scan   [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace   m (Stms rep)-segScan lvl pat w ops map_lam arrs ispace inps = runBuilder_ $ do-  (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps+segScan lvl pat cs w ops map_lam arrs ispace inps = runBuilder_ $ do+  (kspace, kbody) <- prepareRedOrScan cs w map_lam arrs ispace inps   letBind pat $     Op $       segOp $@@ -119,7 +122,7 @@   [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace   m (Stms rep) segMap lvl pat w map_lam arrs ispace inps = runBuilder_ $ do-  (kspace, kbody) <- prepareRedOrScan w map_lam arrs ispace inps+  (kspace, kbody) <- prepareRedOrScan mempty w map_lam arrs ispace inps   letBind pat $     Op $       segOp $@@ -163,7 +166,7 @@   m (Stms rep) nonSegRed lvl pat w ops map_lam arrs = runBuilder_ $ do   (pat', ispace, read_dummy) <- dummyDim pat-  addStms =<< segRed lvl pat' w ops map_lam arrs ispace []+  addStms =<< segRed lvl pat' mempty w ops map_lam arrs ispace []   read_dummy  segHist ::

@@ -341,7 +341,7 @@       types <- asksScope scopeForSOACs       stream_stms <-         snd <$> runBuilderT (sequentialStreamWholeArray pat w accs lam arrs) types-      (_, stream_stms') <-+      stream_stms' <-         runReaderT (copyPropagateInStms simpleSOACS types stream_stms) types       onStms acc $ stmsToList (fmap (certify cs) stream_stms') ++ stms     onStms acc (stm : stms) =@@ -406,8 +406,7 @@             -- (which are now innermost).             stms <-               (`runReaderT` types) $-                fmap snd . simplifyStms-                  =<< interchangeLoops nest' (SeqLoop perm pat merge form body)+                simplifyStms =<< interchangeLoops nest' (SeqLoop perm pat merge form body)             onTopLevelStms stms             return acc'       _ ->@@ -430,8 +429,7 @@             let branch = Branch perm pat cond tbranch fbranch ret             stms <-               (`runReaderT` types) $-                fmap snd . simplifyStms-                  =<< interchangeBranch nest' branch+                simplifyStms =<< interchangeBranch nest' branch             onTopLevelStms stms             return acc'       _ ->@@ -453,7 +451,7 @@             let withacc = WithAccStm perm pat inputs lam             stms <-               (`runReaderT` types) $-                fmap snd . simplifyStms =<< interchangeWithAcc nest' withacc+                simplifyStms =<< interchangeWithAcc nest' withacc             onTopLevelStms stms             return acc'       _ ->@@ -524,10 +522,18 @@             nest' <- expandKernelNest pat_unused nest             map_lam' <- soacsLambda map_lam             localScope (typeEnvFromDistAcc acc') $-              segmentedScanomapKernel nest' perm w lam map_lam' nes arrs-                >>= kernelOrNot cs stm acc kernels acc'+              segmentedScanomapKernel nest' perm cs w lam map_lam' nes arrs+                >>= kernelOrNot mempty stm acc kernels acc'       _ ->         addStmToAcc stm acc+-- If the map function of the reduction contains parallelism we split+-- it, so that the parallelism can be exploited.+maybeDistributeStm (Let pat aux (Op (Screma w arrs form))) acc+  | Just (reds, map_lam) <- isRedomapSOAC form,+    lambdaContainsParallelism map_lam = do+    (mapstm, redstm) <-+      redomapToMapAndReduce pat (w, reds, map_lam, arrs)+    distributeMapBodyStms acc $ oneStm mapstm {stmAux = aux} <> oneStm redstm -- if the reduction can be distributed by itself, we will turn it into a -- segmented reduce. --@@ -551,8 +557,8 @@                   | commutativeLambda lam = Commutative                   | otherwise = comm -            regularSegmentedRedomapKernel nest' perm w comm' lam' map_lam' nes arrs-              >>= kernelOrNot cs stm acc kernels acc'+            regularSegmentedRedomapKernel nest' perm cs w comm' lam' map_lam' nes arrs+              >>= kernelOrNot mempty stm acc kernels acc'       _ ->         addStmToAcc stm acc maybeDistributeStm (Let pat (StmAux cs _ _) (Op (Screma w arrs form))) acc = do@@ -1017,13 +1023,14 @@   (MonadFreshNames m, LocalScope rep m, DistRep rep) =>   KernelNest ->   [Int] ->+  Certs ->   SubExp ->   Lambda SOACS ->   Lambda rep ->   [SubExp] ->   [VName] ->   DistNestT rep m (Maybe (Stms rep))-segmentedScanomapKernel nest perm segment_size lam map_lam nes arrs = do+segmentedScanomapKernel nest perm cs segment_size lam map_lam nes arrs = do   mk_lvl <- asks distSegLevel   onLambda <- asks distOnSOACSLambda   let onLambda' = fmap fst . runBuilder . onLambda@@ -1034,12 +1041,13 @@       let scan_op = SegBinOp Noncommutative lam'' nes'' shape       lvl <- mk_lvl (segment_size : map snd ispace) "segscan" $ NoRecommendation SegNoVirt       addStms =<< traverse renameStm-        =<< segScan lvl pat segment_size [scan_op] map_lam arrs ispace inps+        =<< segScan lvl pat cs segment_size [scan_op] map_lam arrs ispace inps  regularSegmentedRedomapKernel ::   (MonadFreshNames m, LocalScope rep m, DistRep rep) =>   KernelNest ->   [Int] ->+  Certs ->   SubExp ->   Commutativity ->   Lambda rep ->@@ -1047,14 +1055,14 @@   [SubExp] ->   [VName] ->   DistNestT rep m (Maybe (Stms rep))-regularSegmentedRedomapKernel nest perm segment_size comm lam map_lam nes arrs = do+regularSegmentedRedomapKernel nest perm cs segment_size comm lam map_lam nes arrs = do   mk_lvl <- asks distSegLevel   isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $     \pat ispace inps nes' _ -> do       let red_op = SegBinOp comm lam nes' mempty       lvl <- mk_lvl (segment_size : map snd ispace) "segred" $ NoRecommendation SegNoVirt       addStms =<< traverse renameStm-        =<< segRed lvl pat segment_size [red_op] map_lam arrs ispace inps+        =<< segRed lvl pat cs segment_size [red_op] map_lam arrs ispace inps  isSegmentedOp ::   (MonadFreshNames m, LocalScope rep m, DistRep rep) =>

@@ -32,6 +32,7 @@     innermostKernelNesting,     pushKernelNesting,     pushInnerKernelNesting,+    scopeOfKernelNest,     kernelNestLoops,     kernelNestWidths,     boundInKernelNest,@@ -81,8 +82,7 @@ ppTargets (Targets target targets) =   unlines $ map ppTarget $ targets ++ [target]   where-    ppTarget (pat, res) =-      pretty pat ++ " <- " ++ pretty res+    ppTarget (pat, res) = pretty pat ++ " <- " ++ pretty res  singleTarget :: Target -> Targets singleTarget = flip Targets []@@ -127,7 +127,7 @@   }   deriving (Show) -scopeOfLoopNesting :: DistRep rep => LoopNesting -> Scope rep+scopeOfLoopNesting :: (LParamInfo rep ~ Type) => LoopNesting -> Scope rep scopeOfLoopNesting = scopeOfLParams . map fst . loopNestingParamsAndArrs  ppLoopNesting :: LoopNesting -> String@@ -141,10 +141,7 @@  instance FreeIn LoopNesting where   freeIn' (MapNesting pat aux w params_and_arrs) =-    freeIn' pat-      <> freeIn' aux-      <> freeIn' w-      <> freeIn' params_and_arrs+    freeIn' pat <> freeIn' aux <> freeIn' w <> freeIn' params_and_arrs  data Nesting = Nesting   { nestingLetBound :: Names,@@ -164,8 +161,7 @@ ppNestings (nesting, nestings) =   unlines $ map ppNesting $ nestings ++ [nesting]   where-    ppNesting (Nesting _ loop) =-      ppLoopNesting loop+    ppNesting (Nesting _ loop) = ppLoopNesting loop  singleNesting :: Nesting -> Nestings singleNesting = (,[])@@ -234,16 +230,15 @@ kernelNestLoops :: KernelNest -> [LoopNesting] kernelNestLoops (loop, loops) = loop : loops +scopeOfKernelNest :: LParamInfo rep ~ Type => KernelNest -> Scope rep+scopeOfKernelNest = foldMap scopeOfLoopNesting . kernelNestLoops+ boundInKernelNest :: KernelNest -> Names boundInKernelNest = mconcat . boundInKernelNests  boundInKernelNests :: KernelNest -> [Names] boundInKernelNests =-  map-    ( namesFromList-        . map (paramName . fst)-        . loopNestingParamsAndArrs-    )+  map (namesFromList . map (paramName . fst) . loopNestingParamsAndArrs)     . kernelNestLoops  kernelNestWidths :: KernelNest -> [SubExp]@@ -286,10 +281,7 @@ flatKernel ::   MonadFreshNames m =>   KernelNest ->-  m-    ( [(VName, SubExp)],-      [KernelInput]-    )+  m ([(VName, SubExp)], [KernelInput]) flatKernel (MapNesting _ _ nesting_w params_and_arrs, []) = do   i <- newVName "gtid"   let inps =@@ -376,9 +368,7 @@   runMaybeT $ fmap prepare $ recurse $ zip nests targets   where     prepare (x, _, z) = (z, x)-    bound_in_nest =-      mconcat $ map boundInNesting $ inner_nest : nests-+    bound_in_nest = mconcat $ map boundInNesting $ inner_nest : nests     distributableType =       (== mempty) . namesIntersection bound_in_nest . freeIn . arrayDims @@ -418,8 +408,7 @@                 case M.lookup pname identity_map of                   Nothing -> do                     arr <--                      newIdent (baseString pname ++ "_r") $-                        arrayOfRow ptype w+                      newIdent (baseString pname ++ "_r") $ arrayOfRow ptype w                     return                       ( Param mempty pname ptype,                         arr,@@ -462,9 +451,7 @@             addTarget (free_arrs_pat, varsRes $ map paramName free_params_pat)           ) -    recurse ::-      [(Nesting, Target)] ->-      MaybeT m (KernelNest, Names, Targets)+    recurse :: [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets)     recurse [] =       distributeAtNesting         inner_nest@@ -500,9 +487,7 @@   where     (params, arrs) = unzip params_and_arrs     (used_params, used_arrs) =-      unzip $-        filter ((`nameIn` used) . paramName . fst) $-          zip params arrs+      unzip $ filter ((`nameIn` used) . paramName . fst) $ zip params arrs  removeIdentityMappingGeneral ::   Names ->

@@ -27,14 +27,19 @@   ( KernelNest,     LoopNesting (..),     kernelNestLoops,+    scopeOfKernelNest,   ) import Futhark.Tools import Futhark.Transform.Rename+import Futhark.Util (splitFromEnd)  -- | An encoding of a sequential do-loop with no existential context, -- alongside its result pattern. data SeqLoop = SeqLoop [Int] Pat [(FParam, SubExp)] (LoopForm SOACS) Body +loopPerm :: SeqLoop -> [Int]+loopPerm (SeqLoop perm _ _ _ _) = perm+ seqLoopStm :: SeqLoop -> Stm seqLoopStm (SeqLoop _ pat merge form body) =   Let pat (defAux ()) $ DoLoop merge form body@@ -77,7 +82,7 @@         res = varsRes $ patNames loop_pat_expanded         pat' = Pat $ rearrangeShape perm $ patElems pat -    return $+    pure $       SeqLoop perm pat' merge_expanded form $         mkBody (pre_copy_stms <> oneStm map_stm) res     where@@ -91,7 +96,7 @@        expandedInit _ (Var v)         | Just arr <- isMapParameter v =-          return $ Var arr+          pure $ Var arr       expandedInit param_name se =         letSubExp (param_name <> "_expanded_init") $           BasicOp $ Replicate (Shape [w]) se@@ -99,8 +104,10 @@       expand (merge_param, merge_init) = do         expanded_param <-           newParam (param_name <> "_expanded") $-            arrayOf (paramDeclType merge_param) (Shape [w]) $-              uniqueness $ declTypeOf merge_param+            -- FIXME: Unique here is a hack to make sure the copy from+            -- makeCopyInitial is not prematurely simplified away.+            -- It'd be better to fix this somewhere else...+            arrayOf (paramDeclType merge_param) (Shape [w]) Unique         expanded_init <- expandedInit param_name merge_init         return (expanded_param, expanded_init)         where@@ -109,6 +116,39 @@       expandPatElem (PatElem name t) =         PatElem name $ arrayOfRow t w +-- We need to copy some initial arguments because otherwise the result+-- of the loop might alias the input (if the number of iterations is+-- 0), which is a problem if the result is consumed.+maybeCopyInitial ::+  (MonadBuilder m) =>+  (VName -> Bool) ->+  SeqLoop ->+  m SeqLoop+maybeCopyInitial isMapInput (SeqLoop perm loop_pat merge form body) =+  SeqLoop perm loop_pat <$> mapM f merge <*> pure form <*> pure body+  where+    f (p, Var arg)+      | isMapInput arg =+        (p,) <$> letSubExp (baseString (paramName p) <> "_inter_copy") (BasicOp $ Copy arg)+    f (p, arg) =+      pure (p, arg)++manifestMaps :: [LoopNesting] -> [VName] -> Stms SOACS -> ([VName], Stms SOACS)+manifestMaps [] res stms = (res, stms)+manifestMaps (n : ns) res stms =+  let (res', stms') = manifestMaps ns res stms+      (params, arrs) = unzip $ loopNestingParamsAndArrs n+      lam =+        Lambda+          params+          (mkBody stms' $ varsRes res')+          (map rowType $ patTypes (loopNestingPat n))+   in ( patNames $ loopNestingPat n,+        oneStm $+          Let (loopNestingPat n) (loopNestingAux n) $+            Op $ Screma (loopNestingWidth n) arrs (mapSOAC lam)+      )+ -- | Given a (parallel) map nesting and an inner sequential loop, move -- the maps inside the sequential loop.  The result is several -- statements - one of these will be the loop, which will then contain@@ -118,17 +158,29 @@   KernelNest ->   SeqLoop ->   m (Stms SOACS)-interchangeLoops nest loop = do-  (loop', stms) <--    runBuilder $-      foldM (interchangeLoop isMapParameter) loop $-        reverse $ kernelNestLoops nest-  return $ stms <> oneStm (seqLoopStm loop')+interchangeLoops full_nest = recurse (kernelNestLoops full_nest)   where-    isMapParameter v =-      fmap snd $-        find ((== v) . paramName . fst) $-          concatMap loopNestingParamsAndArrs $ kernelNestLoops nest+    recurse nest loop+      | (ns, [n]) <- splitFromEnd 1 nest = do+        let isMapParameter v =+              snd <$> find ((== v) . paramName . fst) (loopNestingParamsAndArrs n)+            isMapInput v =+              v `elem` map snd (loopNestingParamsAndArrs n)+        (loop', stms) <-+          runBuilder . localScope (scopeOfKernelNest full_nest) $+            maybeCopyInitial isMapInput+              =<< interchangeLoop isMapParameter loop n++        -- Only safe to continue interchanging if we didn't need to add+        -- any new statements; otherwise we manifest the remaining nests+        -- as Maps and hand them back to the flattener.+        if null stms+          then recurse ns loop'+          else+            let loop_stm = seqLoopStm loop'+                names = rearrangeShape (loopPerm loop') (patNames (stmPat loop_stm))+             in pure $ snd $ manifestMaps ns names $ stms <> oneStm loop_stm+      | otherwise = pure $ oneStm $ seqLoopStm loop  data Branch = Branch [Int] Pat SubExp Body Body (IfDec (BranchType SOACS)) 

@@ -251,7 +251,7 @@         let scanfun' = soacsLambdaToGPU scanfun             mapfun' = soacsLambdaToGPU mapfun         certifying (stmAuxCerts aux) $-          addStms =<< segScan lvl' pat w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] []+          addStms =<< segScan lvl' pat mempty w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' arrs [] []         parallelMin [w]     Op (Screma w arrs form)       | Just (reds, map_lam) <- isRedomapSOAC form,@@ -259,7 +259,7 @@         let red_lam' = soacsLambdaToGPU red_lam             map_lam' = soacsLambdaToGPU map_lam         certifying (stmAuxCerts aux) $-          addStms =<< segRed lvl' pat w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] []+          addStms =<< segRed lvl' pat mempty w [SegBinOp comm red_lam' nes mempty] map_lam' arrs [] []         parallelMin [w]     Op (Hist w arrs ops bucket_fun) -> do       ops' <- forM ops $ \(HistOp num_bins rf dests nes op) -> do

@@ -32,7 +32,7 @@   SimpleOps rep ->   Scope rep ->   Stms rep ->-  m (ST.SymbolTable (Wise rep), Stms rep)+  m (Stms rep) copyPropagateInStms simpl = simplifyStms simpl mempty neverHoist  -- | Run copy propagation on a function.

@@ -173,14 +173,12 @@           Nothing             | paramName p `nameIn` lam_cons -> do               p' <--                letExp (baseString (paramName p)) $-                  BasicOp $-                    Index arr $ fullSlice arr_t [DimFix $ Var i]+                letExp (baseString (paramName p)) . BasicOp $+                  Index arr $ fullSlice arr_t [DimFix $ Var i]               letBindNames [paramName p] $ BasicOp $ Copy p'             | otherwise ->               letBindNames [paramName p] $-                BasicOp $-                  Index arr $ fullSlice arr_t [DimFix $ Var i]+                BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var i]        -- Insert the statements of the lambda.  We have taken care to       -- ensure that the parameters are bound at this point.

@@ -1012,6 +1012,8 @@   let (mergepat, mergeexps) = unzip merge   mergeargs <- mapM checkArg mergeexps +  checkLoopArgs+   binding (scopeOf form) $ do     form_consumable <- checkForm mergeargs form @@ -1035,15 +1037,14 @@           checkBodyDec $ snd $ bodyDec loopbody            checkStms (bodyStms loopbody) $ do-            checkResult $ bodyResult loopbody+            context "In loop body result" $+              checkResult $ bodyResult loopbody              context "When matching result of body with loop parameters" $               matchLoopResult (map fst merge) $ bodyResult loopbody              let bound_here =-                  namesFromList $-                    M.keys $-                      scopeOf $ bodyStms loopbody+                  namesFromList $ M.keys $ scopeOf $ bodyStms loopbody             map (`namesSubtract` bound_here)               <$> mapM (subExpAliasesM . resSubExp) (bodyResult loopbody)   where@@ -1097,6 +1098,20 @@           paramts = map paramDeclType funparams       checkFuncall Nothing paramts mergeargs       pure mempty++    checkLoopArgs = do+      let (params, args) = unzip merge++      argtypes <- mapM subExpType args++      let expected = expectedTypes (map paramName params) params args+      unless (expected == argtypes) . bad . TypeError . pretty $+        "Loop parameters"+          </> indent 2 (ppTuple' params)+          </> "cannot accept initial values"+          </> indent 2 (ppTuple' args)+          </> "of types"+          </> indent 2 (ppTuple' argtypes) checkExp (WithAcc inputs lam) = do   unless (length (lambdaParams lam) == 2 * num_accs) $     bad . TypeError $

@@ -10,6 +10,7 @@ -- compatible). module Futhark.Util   ( nubOrd,+    nubByOrd,     mapAccumLM,     maxinum,     chunk,@@ -65,7 +66,7 @@ import Data.Char import Data.Either import Data.Function ((&))-import Data.List (foldl', genericDrop, genericSplitAt, sort)+import Data.List (foldl', genericDrop, genericSplitAt, sortBy) import qualified Data.List.NonEmpty as NE import qualified Data.Map as M import Data.Maybe@@ -89,7 +90,13 @@  -- | Like 'nub', but without the quadratic runtime. nubOrd :: Ord a => [a] -> [a]-nubOrd = map NE.head . NE.group . sort+nubOrd = nubByOrd compare++-- | Like 'nubBy', but without the quadratic runtime.+nubByOrd :: (a -> a -> Ordering) -> [a] -> [a]+nubByOrd cmp = map NE.head . NE.groupBy eq . sortBy cmp+  where+    eq x y = cmp x y == EQ  -- | Like 'Data.Traversable.mapAccumL', but monadic. mapAccumLM ::

@@ -751,20 +751,22 @@   Env ->   [TypeParam] ->   [Pat] ->-  StructType ->-  [VName] ->+  StructRetType ->   Exp ->   EvalM TermBinding-evalFunctionBinding env tparams ps ret retext fbody = do-  let ret' = evalType env ret+evalFunctionBinding env tparams ps ret fbody = do+  let ret' = evalType env $ retType ret       arrow (xp, xt) yt = Scalar $ Arrow () xp xt $ RetType [] yt       ftype = foldr (arrow . patternParam) ret' ps+      retext = case ps of+        [] -> retDims ret+        _ -> []    -- Distinguish polymorphic and non-polymorphic bindings here.   if null tparams     then       TermValue (Just $ T.BoundV [] ftype)-        <$> (returned env ret retext =<< evalFunction env [] ps fbody ret')+        <$> (returned env (retType ret) retext =<< evalFunction env [] ps fbody ret')     else return $       TermPoly (Just $ T.BoundV [] ftype) $ \ftype' -> do         let tparam_names = map typeParamName tparams@@ -777,7 +779,7 @@             missing_sizes =               filter (`M.notMember` envTerm env') $                 map typeParamName (filter isSizeParam tparams)-        returned env ret retext =<< evalFunction env' missing_sizes ps fbody ret'+        returned env (retType ret) retext =<< evalFunction env' missing_sizes ps fbody ret'  evalArg :: Env -> Exp -> Maybe VName -> EvalM Value evalArg env e ext = do@@ -863,8 +865,8 @@       v_s = valueShape v       env'' = env' <> i64Env (resolveExistentials (map sizeName sizes) p_t v_s)   eval env'' body-evalAppExp env (LetFun f (tparams, ps, _, Info (RetType _ ret), fbody) body _) = do-  binding <- evalFunctionBinding env tparams ps ret [] fbody+evalAppExp env (LetFun f (tparams, ps, _, Info ret, fbody) body _) = do+  binding <- evalFunctionBinding env tparams ps ret fbody   eval (env {envTerm = M.insert f binding $ envTerm env}) body evalAppExp   env@@ -1187,8 +1189,8 @@     _ -> error "Expected ModuleFun."  evalDec :: Env -> Dec -> EvalM Env-evalDec env (ValDec (ValBind _ v _ (Info (RetType _ ret, retext)) tparams ps fbody _ _ _)) = do-  binding <- evalFunctionBinding env tparams ps ret retext fbody+evalDec env (ValDec (ValBind _ v _ (Info ret) tparams ps fbody _ _ _)) = do+  binding <- evalFunctionBinding env tparams ps ret fbody   return $ env {envTerm = M.insert v binding $ envTerm env} evalDec env (OpenDec me _) = do   me' <- evalModExp env me

@@ -646,8 +646,6 @@          QualParens (QualName qs name, loc) $2 (srcspan $1 $>) }       -- Operator sections.-     | '(' '!' ')'-        { Var (qualName "!") NoInfo (srcspan $2 $>) }      | '(' '-' ')'         { OpSection (qualName (nameFromString "-")) NoInfo (srcspan $1 $>) }      | '(' Exp2 '-' ')'

@@ -469,7 +469,7 @@       fun         | isJust entry = "entry"         | otherwise = "let"-      retdecl' = case (ppr . fst <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of+      retdecl' = case (ppr <$> unAnnot rettype) `mplus` (ppr <$> retdecl) of         Just rettype' -> colon <+> align rettype'         Nothing -> mempty 

@@ -28,6 +28,7 @@     -- * Queries on expressions     typeOf,     valBindTypeScheme,+    valBindBound,     funType,      -- * Queries on patterns and params@@ -619,8 +620,18 @@ valBindTypeScheme :: ValBindBase Info VName -> ([TypeParamBase VName], StructType) valBindTypeScheme vb =   ( valBindTypeParams vb,-    funType (valBindParams vb) (fst (unInfo (valBindRetType vb)))+    funType (valBindParams vb) (unInfo (valBindRetType vb))   )++-- | The names that are brought into scope by this value binding (not+-- including its own parameter names, but including any existential+-- sizes).+valBindBound :: ValBindBase Info VName -> [VName]+valBindBound vb =+  valBindName vb :+  case valBindParams vb of+    [] -> retDims (unInfo (valBindRetType vb))+    _ -> []  -- | The type of a function with the given parameters and return type. funType :: [PatBase Info VName] -> StructRetType -> StructType

@@ -114,7 +114,7 @@   where     vbind_t =       foldFunType (map patternStructType (valBindParams vbind)) $-        fst $ unInfo $ valBindRetType vbind+        unInfo $ valBindRetType vbind  typeBindDefs :: TypeBind -> Defs typeBindDefs tbind =

@@ -1041,7 +1041,9 @@     valBindEntryPoint :: Maybe (f EntryPoint),     valBindName :: vn,     valBindRetDecl :: Maybe (TypeExp vn),-    valBindRetType :: f (StructRetType, [VName]),+    -- | If 'valBindParams' is null, then the 'retDims' are brought+    -- into scope at this point.+    valBindRetType :: f StructRetType,     valBindTypeParams :: [TypeParamBase vn],     valBindParams :: [PatBase f vn],     valBindBody :: ExpBase f vn,

@@ -633,7 +633,7 @@     typeError loc mempty $       withIndexLink "nested-entry" "Entry points may not be declared inside modules." -  (fname', tparams', params', maybe_tdecl', rettype@(RetType _ rettype_t), retext, body') <-+  (fname', tparams', params', maybe_tdecl', rettype@(RetType _ rettype_t), body') <-     checkFunDef (fname, maybe_tdecl, tparams, params, body, loc)    let (rettype_params, rettype') = unfoldFunType rettype_t@@ -668,7 +668,7 @@     _ -> return ()    attrs' <- mapM checkAttr attrs-  let vb = ValBind entry' fname' maybe_tdecl' (Info (rettype, retext)) tparams' params' body' doc attrs' loc+  let vb = ValBind entry' fname' maybe_tdecl' (Info rettype) tparams' params' body' doc attrs' loc   return     ( mempty         { envVtable =

@@ -496,7 +496,7 @@         Ident (sizeName size) (Info (Scalar $ Prim $ Signed Int64)) (srclocOf size) checkExp (AppExp (LetFun name (tparams, params, maybe_retdecl, NoInfo, e) body loc) _) =   sequentially (checkBinding (name, maybe_retdecl, tparams, params, e, loc)) $-    \(tparams', params', maybe_retdecl', rettype, _, e') closure -> do+    \(tparams', params', maybe_retdecl', rettype, e') closure -> do       closure' <- lexicalClosure params' closure        bindSpaced [(Term, name)] $ do@@ -1067,7 +1067,7 @@ checkOneExp e = fmap fst . runTermTypeM $ do   e' <- checkExp e   let t = toStruct $ typeOf e'-  (tparams, _, _, _) <-+  (tparams, _, _) <-     letGeneralise (nameFromString "<exp>") (srclocOf e) [] [] t   fixOverloadedTypes $ typeVars t   e'' <- updateTypes e'@@ -1236,12 +1236,11 @@       [Pat],       Maybe (TypeExp VName),       StructRetType,-      [VName],       Exp     ) checkFunDef (fname, maybe_retdecl, tparams, params, body, loc) =   fmap fst . runTermTypeM $ do-    (tparams', params', maybe_retdecl', RetType dims rettype', retext, body') <-+    (tparams', params', maybe_retdecl', RetType dims rettype', body') <-       checkBinding (fname, maybe_retdecl, tparams, params, body, loc)      -- Since this is a top-level function, we also resolve overloaded@@ -1270,7 +1269,7 @@         typeError loc mempty . withIndexLink "may-not-be-redefined" $           "The" <+> pprName fname <+> "operator may not be redefined." -      pure (fname', tparams', params'', maybe_retdecl'', RetType dims rettype'', retext, body'')+      pure (fname', tparams', params'', maybe_retdecl'', RetType dims rettype'', body'')  -- | This is "fixing" as in "setting them", not "correcting them".  We -- only make very conservative fixing.@@ -1354,7 +1353,6 @@       [Pat],       Maybe (TypeExp VName),       StructRetType,-      [VName],       Exp     ) checkBinding (fname, maybe_retdecl, tparams, params, body, loc) =@@ -1391,12 +1389,12 @@      verifyFunctionParams (Just fname) params'' -    (tparams'', params''', rettype'', retext) <-+    (tparams'', params''', rettype'') <-       letGeneralise fname loc tparams' params'' rettype      checkGlobalAliases params'' body_t loc -    pure (tparams'', params''', maybe_retdecl'', rettype'', retext, body')+    pure (tparams'', params''', maybe_retdecl'', rettype'', body')   where     checkReturnAlias rettp params' =       foldM_ (checkReturnAlias' params') S.empty . returnAliasing rettp@@ -1584,7 +1582,7 @@   [StructType] ->   StructType ->   Constraints ->-  TermTypeM ([TypeParam], StructRetType, [VName])+  TermTypeM ([TypeParam], StructRetType) closeOverTypes defname defloc tparams paramts ret substs = do   (more_tparams, retext) <-     partitionEithers . catMaybes@@ -1597,8 +1595,7 @@       mkExt AnyDim {} = error "closeOverTypes: AnyDim"   return     ( tparams ++ more_tparams,-      injectExt (mapMaybe mkExt (nestedDims ret)) ret,-      retext+      injectExt (retext ++ mapMaybe mkExt (nestedDims ret)) ret     )   where     t = foldFunType paramts $ RetType [] ret@@ -1638,7 +1635,7 @@   [TypeParam] ->   [Pat] ->   StructType ->-  TermTypeM ([TypeParam], [Pat], StructRetType, [VName])+  TermTypeM ([TypeParam], [Pat], StructRetType) letGeneralise defname defloc tparams params rettype =   onFailure (CheckingLetGeneralise defname) $ do     now_substs <- getConstraints@@ -1663,7 +1660,7 @@     let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= cur_lvl         new_substs = M.filterWithKey candidate now_substs -    (tparams', RetType ret_dims rettype', retext) <-+    (tparams', RetType ret_dims rettype') <-       closeOverTypes         defname         defloc@@ -1685,7 +1682,7 @@     -- let-generalisation.     modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` map typeParamName tparams' -    pure (tparams', params, RetType ret_dims rettype'', retext)+    pure (tparams', params, RetType ret_dims rettype'')  checkFunBody ::   [Pat] ->

@@ -791,16 +791,6 @@   mustBeOneOf ts (mkUsage (srclocOf e) why) . toStruct =<< expType e   pure e -renameRetType :: StructRetType -> TermTypeM StructRetType-renameRetType (RetType dims st)-  | dims /= mempty = do-    dims' <- mapM newName dims-    let m = M.fromList $ zip dims $ map (SizeSubst . NamedDim . qualName) dims'-        st' = applySubst (`M.lookup` m) st-    pure $ RetType dims' st'-  | otherwise =-    pure $ RetType dims st- termCheckTypeExp :: TypeExp Name -> TermTypeM (TypeExp VName, [VName], StructRetType) termCheckTypeExp te = do   (te', svars, rettype, _l) <- checkTypeExp te

@@ -7,6 +7,7 @@ -- | Type checker building blocks that do not involve unification. module Language.Futhark.TypeChecker.Types   ( checkTypeExp,+    renameRetType,     unifyTypesU,     subtypeOf,     subuniqueOf,@@ -109,14 +110,27 @@ subuniqueOf Nonunique Unique = False subuniqueOf _ _ = True +-- | Ensure that the dimensions of the RetType are unique by+-- generating new names for them.  This is to avoid name capture.+renameRetType :: MonadTypeChecker m => StructRetType -> m StructRetType+renameRetType (RetType dims st)+  | dims /= mempty = do+    dims' <- mapM newName dims+    let m = M.fromList $ zip dims $ map (SizeSubst . NamedDim . qualName) dims'+        st' = applySubst (`M.lookup` m) st+    pure $ RetType dims' st'+  | otherwise =+    pure $ RetType dims st+ evalTypeExp ::   MonadTypeChecker m =>   TypeExp Name ->   m (TypeExp VName, [VName], StructRetType, Liftedness) evalTypeExp (TEVar name loc) = do   (name', ps, t, l) <- lookupType loc name+  t' <- renameRetType t   case ps of-    [] -> pure (TEVar name' loc, [], t, l)+    [] -> pure (TEVar name' loc, [], t', l)     _ ->       typeError loc mempty $         "Type constructor" <+> pquote (spread (ppr name : map ppr ps))@@ -252,7 +266,8 @@     ) evalTypeExp ote@TEApply {} = do   (tname, tname_loc, targs) <- rootAndArgs ote-  (tname', ps, RetType t_dims t, l) <- lookupType tloc tname+  (tname', ps, tname_t, l) <- lookupType tloc tname+  RetType t_dims t <- renameRetType tname_t   if length ps /= length targs     then       typeError tloc mempty $

@@ -46,13 +46,20 @@                       TypeParamType Lifted "b_1102" mempty                     ]                     "a_1101 -> b_1102"+                ),+                ( "pair_1200",+                  TypeAbbr+                    SizeLifted+                    []+                    "?[n_1201][m_1202].([n_1201]i64, [m_1202]i64)"                 )               ]               <> envTypeTable initialEnv,           envNameMap =             M.fromList               [ ((Type, "square"), "square_1000"),-                ((Type, "fun"), "fun_1100")+                ((Type, "fun"), "fun_1100"),+                ((Type, "pair"), "pair_1200")               ]               <> envNameMap initialEnv         }@@ -114,7 +121,13 @@         ([], "bool -> ?[d_0].[d_0]i32 -> bool"),       evalTest         "bool -> fun bool ([]i32)"-        ([], "bool -> ?[d_0].bool -> [d_0]i32")+        ([], "bool -> ?[d_0].bool -> [d_0]i32"),+      evalTest+        "pair"+        ([], "?[n_0][m_1].([n_0]i64, [m_1]i64)"),+      evalTest+        "(pair,pair)"+        ([], "?[n_0][m_1][n_2][m_3].(([n_0]i64, [m_1]i64), ([n_2]i64, [m_3]i64))")     ]  substTest :: M.Map VName (Subst StructRetType) -> StructRetType -> StructRetType -> TestTree