futhark 0.22.2 → 0.22.3

raw patch · 58 files changed

+937/−733 lines, 58 filesPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

API changes (from Hackage documentation)

Files

• docs/installation.rst +2/−2
• futhark.cabal +1/−1
• src/Futhark/AD/Rev/Scan.hs +3/−9
• src/Futhark/Actions.hs +18/−0
• src/Futhark/Analysis/Interference.hs +1/−7
• src/Futhark/Analysis/LastUse.hs +20/−5
• src/Futhark/Analysis/PrimExp.hs +26/−3
• src/Futhark/Analysis/UsageTable.hs +10/−12
• src/Futhark/CLI/Dev.hs +5/−0
• src/Futhark/CodeGen/Backends/GenericC/CLI.hs +7/−10
• src/Futhark/CodeGen/ImpGen.hs +30/−12
• src/Futhark/CodeGen/ImpGen/GPU.hs +14/−16
• src/Futhark/CodeGen/ImpGen/GPU/Base.hs +31/−2
• src/Futhark/CodeGen/ImpGen/GPU/Group.hs +8/−13
• src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs +5/−5
• src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs +5/−2
• src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +8/−6
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs +3/−2
• src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +6/−5
• src/Futhark/Construct.hs +18/−0
• src/Futhark/IR/GPU/Op.hs +121/−55
• src/Futhark/IR/MC/Op.hs +1/−2
• src/Futhark/IR/Mem.hs +5/−5
• src/Futhark/IR/Parse.hs +30/−17
• src/Futhark/IR/Pretty.hs +7/−11
• src/Futhark/IR/Prop.hs +0/−1
• src/Futhark/IR/Prop/Names.hs +4/−0
• src/Futhark/IR/SOACS/SOAC.hs +1/−1
• src/Futhark/IR/SOACS/Simplify.hs +9/−9
• src/Futhark/IR/SegOp.hs +8/−43
• src/Futhark/IR/Syntax.hs +10/−2
• src/Futhark/IR/TypeCheck.hs +14/−7
• src/Futhark/Internalise/Defunctionalise.hs +24/−11
• src/Futhark/Internalise/Defunctorise.hs +3/−6
• src/Futhark/Internalise/Exps.hs +92/−73
• src/Futhark/Optimise/BlkRegTiling.hs +9/−6
• src/Futhark/Optimise/Fusion/GraphRep.hs +8/−0
• src/Futhark/Optimise/GenRedOpt.hs +3/−11
• src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +1/−4
• src/Futhark/Optimise/Simplify/Engine.hs +7/−6
• src/Futhark/Optimise/Simplify/Rules/Loop.hs +1/−30
• src/Futhark/Optimise/TileLoops.hs +130/−164
• src/Futhark/Optimise/TileLoops/Shared.hs +7/−9
• src/Futhark/Pass/ExpandAllocations.hs +55/−25
• src/Futhark/Pass/ExplicitAllocations/GPU.hs +39/−21
• src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +1/−0
• src/Futhark/Pass/ExtractKernels/DistributeNests.hs +2/−1
• src/Futhark/Pass/ExtractKernels/Intragroup.hs +25/−31
• src/Futhark/Pass/ExtractKernels/StreamKernel.hs +4/−2
• src/Futhark/Pass/ExtractKernels/ToGPU.hs +6/−4
• src/Futhark/Util.hs +9/−1
• src/Futhark/Util/Pretty.hs +5/−3
• src/Language/Futhark/Interpreter.hs +1/−1
• src/Language/Futhark/Pretty.hs +6/−8
• src/Language/Futhark/TypeChecker/Match.hs +42/−31
• src/Language/Futhark/TypeChecker/Modules.hs +10/−3
• src/Language/Futhark/TypeChecker/Terms.hs +11/−12
• src/Language/Futhark/TypeChecker/Terms/Monad.hs +5/−5

@@ -197,7 +197,7 @@ have installed otherwise.  * On NixOS, for OpenCL, you should import ``opencl-headers`` and-  ``opencl-icd``.  You also need some form of OpenCL backend.  If you+  ``ocl-icd``.  You also need some form of OpenCL backend.  If you   have an AMD GPU and use ROCm, you may also need   ``rocm-opencl-runtime``. @@ -212,4 +212,4 @@  These can be easily made available with e.g:: -  nix-shell -p opencl-headers -p opencl-icd+  nix-shell -p opencl-headers -p ocl-icd

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

@@ -79,12 +79,7 @@ -- but insert explicit indexing to reverse inside the map. mkScan2ndMaps :: SubExp -> (Type, VName, (VName, VName)) -> ADM VName mkScan2ndMaps w (arr_tp, y_adj, (ds, cs)) = do-  nm1 <- letSubExp "nm1" =<< toExp (pe64 w - 1)-  y_adj_last <--    letExp (baseString y_adj ++ "_last") $-      BasicOp $-        Index y_adj $-          fullSlice arr_tp [DimFix nm1]+  y_adj_last <- letExp (baseString y_adj <> "_last") =<< eLast y_adj    par_i <- newParam "i" $ Prim int64   lam <- mkLambda [par_i] $ do@@ -121,9 +116,8 @@           (resultBodyM $ map (Var . paramName) par_r)           ( buildBody_ $ do               im1 <- letSubExp "im1" =<< toExp (le64 i - 1)-              ys_im1 <- forM ys $ \y -> do-                y_t <- lookupType y-                letSubExp (baseString y ++ "_last") $ BasicOp $ Index y $ fullSlice y_t [DimFix im1]+              ys_im1 <- forM ys $ \y ->+                letSubExp (baseString y <> "_im1") =<< eIndex y (eSubExp im1)                lam_res <-                 mapM (letExp "const" . BasicOp . SubExp . resSubExp)

@@ -4,6 +4,7 @@   ( printAction,     printAliasesAction,     printLastUseGPU,+    printFusionGraph,     printInterferenceGPU,     printMemAliasGPU,     callGraphAction,@@ -54,6 +55,7 @@ import Futhark.IR.Prop.Aliases import Futhark.IR.SOACS (SOACS) import Futhark.IR.SeqMem (SeqMem)+import Futhark.Optimise.Fusion.GraphRep qualified import Futhark.Util (runProgramWithExitCode, unixEnvironment) import Futhark.Version (versionString) import System.Directory@@ -86,6 +88,22 @@     { actionName = "print last use gpu",       actionDescription = "Print last use information on gpu.",       actionProcedure = liftIO . print . LastUse.analyseGPUMem+    }++-- | Print fusion graph to stdout.+printFusionGraph :: Action SOACS+printFusionGraph =+  Action+    { actionName = "print fusion graph",+      actionDescription = "Print fusion graph in Graphviz format.",+      actionProcedure =+        liftIO+          . mapM_+            ( putStrLn+                . Futhark.Optimise.Fusion.GraphRep.pprg+                . Futhark.Optimise.Fusion.GraphRep.mkDepGraphForFun+            )+          . progFuns     }  -- | Print interference information to stdout.

@@ -16,7 +16,7 @@ import Futhark.Analysis.LastUse qualified as LastUse import Futhark.Analysis.MemAlias qualified as MemAlias import Futhark.IR.GPUMem-import Futhark.Util (invertMap)+import Futhark.Util (cartesian, invertMap)  -- | The set of 'VName' currently in use. type InUse = Names@@ -36,12 +36,6 @@ makeEdge v1 v2   | v1 == v2 = mempty   | otherwise = S.singleton (min v1 v2, max v1 v2)---- | Compute the cartesian product of two foldable collections, using the given--- combinator function.-cartesian :: (Monoid m, Foldable t) => (a -> a -> m) -> t a -> t a -> m-cartesian f xs ys =-  foldMap (uncurry f) [(x, y) | x <- toList xs, y <- toList ys]  analyseStm ::   LocalScope GPUMem m =>

@@ -136,6 +136,16 @@         used_acc <> unAliases aliases       ) +    addAliasesFromBodyRes (lumap_acc, used_acc) (PatElem {}, Constant _) = (lumap_acc, used_acc)+    addAliasesFromBodyRes (lumap_acc, used_acc) (PatElem name _, Var body_res) =+      -- Any aliases of `name` should have the same last-use as `name`+      ( case M.lookup name lumap_acc of+          Just name' ->+            insertNames name' (oneName body_res) lumap_acc+          Nothing -> lumap_acc,+        used_acc <> oneName body_res+      )+     pat_name = patElemName $ head $ patElems pat      analyseExp (lumap, used) (BasicOp _) = do@@ -149,16 +159,21 @@         bimap mconcat mconcat . unzip           <$> mapM (analyseBody lumap used . caseBody) cases       (lumap_defbody, used_defbody) <- analyseBody lumap used defbody-      let used' = used_cases <> used_defbody+      let (lumap', used') =+            (lumap_defbody <> lumap_cases, used_cases <> used_defbody)+              & flip (foldl addAliasesFromBodyRes) (zip (patElems pat) (map resSubExp $ bodyResult defbody))           nms = (freeIn ses <> freeIn dec) `namesSubtract` used'       pure-        ( insertNames pat_name nms (lumap_cases <> lumap_defbody),+        ( insertNames pat_name nms lumap',           used' <> nms         )     analyseExp (lumap, used) (DoLoop merge form body) = do-      (lumap', used') <- analyseBody lumap used body-      let nms = (freeIn merge <> freeIn form) `namesSubtract` used'-      pure (insertNames pat_name nms lumap', used' <> nms)+      let (lumap', used') =+            zip (patElems pat) (map resSubExp $ bodyResult body)+              & foldl addAliasesFromBodyRes (lumap, used)+      (lumap'', used'') <- analyseBody lumap' used' body+      let nms = (freeIn merge <> freeIn form) `namesSubtract` used''+      pure (insertNames pat_name nms lumap'', used'' <> nms)     analyseExp (lumap, used) (Op op) = do       onOp <- asks envLastUseOp       onOp pat_name (lumap, used) op

@@ -228,6 +228,20 @@   | oneIshExp y = y   | zeroIshExp x = y   | zeroIshExp y = x+constFoldPrimExp (UnOpExp Abs {} x)+  | not $ negativeIshExp x = x+constFoldPrimExp (BinOpExp UMod {} x y)+  | sameIshExp x y,+    IntType it <- primExpType x =+      ValueExp $ IntValue $ intValue it (0 :: Integer)+constFoldPrimExp (BinOpExp SMod {} x y)+  | sameIshExp x y,+    IntType it <- primExpType x =+      ValueExp $ IntValue $ intValue it (0 :: Integer)+constFoldPrimExp (BinOpExp SRem {} x y)+  | sameIshExp x y,+    IntType it <- primExpType x =+      ValueExp $ IntValue $ intValue it (0 :: Integer) constFoldPrimExp e = e  -- | The class of numeric types that can be used for constructing@@ -336,8 +350,8 @@     | otherwise = numBad "*" (x, y)    abs (TPrimExp x)-    | IntType t <- primExpType x = TPrimExp $ UnOpExp (Abs t) x-    | FloatType t <- primExpType x = TPrimExp $ UnOpExp (FAbs t) x+    | IntType t <- primExpType x = TPrimExp $ constFoldPrimExp $ UnOpExp (Abs t) x+    | FloatType t <- primExpType x = TPrimExp $ constFoldPrimExp $ UnOpExp (FAbs t) x     | otherwise = numBad "abs" x    signum (TPrimExp x)@@ -419,7 +433,7 @@    TPrimExp x `mod` TPrimExp y     | Just z <- msum [asIntOp (`SMod` Unsafe) x y] =-        TPrimExp z+        TPrimExp $ constFoldPrimExp z     | otherwise = numBad "mod" (x, y)    TPrimExp x `quot` TPrimExp y@@ -588,6 +602,15 @@ oneIshExp :: PrimExp v -> Bool oneIshExp (ValueExp v) = oneIsh v oneIshExp _ = False++-- | Is the expression a constant negative of some sort?+negativeIshExp :: PrimExp v -> Bool+negativeIshExp (ValueExp v) = negativeIsh v+negativeIshExp _ = False++sameIshExp :: PrimExp v -> PrimExp v -> Bool+sameIshExp (ValueExp v1) (ValueExp v2) = v1 == v2+sameIshExp _ _ = False  -- | If the given 'PrimExp' is a constant of the wrong integer type, -- coerce it to the given integer type.  This is a workaround for an

@@ -23,6 +23,7 @@     consumedU,     presentU,     usageInStm,+    usageInPat,   ) where @@ -119,7 +120,7 @@ -- | Construct a usage table where the given names have been used as -- an array or memory size. sizeUsages :: Names -> UsageTable-sizeUsages = UsageTable . IM.map (const sizeU) . namesIntMap+sizeUsages = UsageTable . IM.map (const (sizeU <> presentU)) . namesIntMap  -- | A description of how a single variable has been used. newtype Usages = Usages Int -- Bitmap representation for speed.@@ -155,20 +156,17 @@ usageInStm :: (ASTRep rep, Aliased rep) => Stm rep -> UsageTable usageInStm (Let pat rep e) =   mconcat-    [ usageInPat,-      usageInExpDec,+    [ usageInPat pat `without` patNames pat,+      usages $ freeIn rep,       usageInExp e,       usages (freeIn e)     ]-  where-    usageInPat =-      usages-        ( mconcat (map freeIn $ patElems pat)-            `namesSubtract` namesFromList (patNames pat)-        )-        <> sizeUsages (foldMap (freeIn . patElemType) (patElems pat))-    usageInExpDec =-      usages $ freeIn rep++-- | Usage table reflecting use in pattern.  In particular, free+-- variables in the decorations are considered used as sizes, even if+-- they are also bound in this pattern.+usageInPat :: FreeIn t => Pat t -> UsageTable+usageInPat = sizeUsages . foldMap freeIn . patElems  usageInExp :: Aliased rep => Exp rep -> UsageTable usageInExp (Apply _ args _ _) =

@@ -462,6 +462,11 @@       "Print the resulting IR with aliases.",     Option       []+      ["fusion-graph"]+      (NoArg $ Right $ \opts -> opts {futharkAction = SOACSAction printFusionGraph})+      "Print fusion graph.",+    Option+      []       ["print-last-use-gpu"]       ( NoArg $           Right $ \opts ->

@@ -282,6 +282,7 @@                  $ty:et' *arr = calloc($exp:num_elems, $id:info.size);                  assert(arr != NULL);                  assert($id:values_array(ctx, $exp:e, arr) == 0);+                 assert(futhark_context_sync(ctx) == 0);                  write_array(stdout, binary_output, &$id:info, arr,                              $id:shape_array(ctx, $exp:e), $int:rank);                  free(arr);@@ -304,10 +305,6 @@       printstms =         printResult manifest $ zip (map outputType outputs) output_vals -      ctx_ty = [C.cty|struct futhark_context|]-      sync_ctx = "futhark_context_sync" :: T.Text-      error_ctx = "futhark_context_get_error" :: T.Text-       cli_entry_point_function_name = "futrts_cli_entry_" ++ T.unpack entry_point_name        pause_profiling = "futhark_context_pause_profiling" :: T.Text@@ -320,8 +317,8 @@                 int r;                 // Run the program once.                 $stms:pack_input-                if ($id:sync_ctx(ctx) != 0) {-                  futhark_panic(1, "%s", $id:error_ctx(ctx));+                if (futhark_context_sync(ctx) != 0) {+                  futhark_panic(1, "%s", futhark_context_get_error(ctx));                 };                 // Only profile last run.                 if (profile_run) {@@ -332,10 +329,10 @@                              $args:(map addrOf output_vals),                              $args:input_args);                 if (r != 0) {-                  futhark_panic(1, "%s", $id:error_ctx(ctx));+                  futhark_panic(1, "%s", futhark_context_get_error(ctx));                 }-                if ($id:sync_ctx(ctx) != 0) {-                  futhark_panic(1, "%s", $id:error_ctx(ctx));+                if (futhark_context_sync(ctx) != 0) {+                  futhark_panic(1, "%s", futhark_context_get_error(ctx));                 };                 if (profile_run) {                   $id:pause_profiling(ctx);@@ -349,7 +346,7 @@                 $stms:free_input               |]    in ( [C.cedecl|-   static int $id:cli_entry_point_function_name($ty:ctx_ty *ctx) {+   static int $id:cli_entry_point_function_name(struct futhark_context *ctx) {      typename int64_t t_start, t_end;      int time_runs = 0, profile_run = 0;      int retval = 0;

@@ -109,6 +109,7 @@     sWrite,     sUpdate,     sLoopNest,+    sCopy,     sLoopSpace,     (<--),     (<~~),@@ -1497,7 +1498,7 @@    pure fname --- | Use an 'Imp.Copy' if possible, otherwise 'copyElementWise'.+-- | Use 'sCopy' if possible, otherwise 'copyElementWise'. defaultCopy :: CopyCompiler rep r op defaultCopy pt dest src   | Just (destoffset, srcoffset, num_arrays, size_x, size_y) <-@@ -1524,17 +1525,7 @@       destspace <- entryMemSpace <$> lookupMemory destmem       if isScalarSpace srcspace || isScalarSpace destspace         then copyElementWise pt dest src-        else-          emit-            $ Imp.Copy-              pt-              destmem-              (bytes destoffset)-              destspace-              srcmem-              (bytes srcoffset)-              srcspace-            $ num_elems `withElemType` pt+        else sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt   | otherwise =       copyElementWise pt dest src   where@@ -1926,6 +1917,33 @@   ([Imp.TExp Int64] -> ImpM rep r op ()) ->   ImpM rep r op () sLoopNest = sLoopSpace . map pe64 . shapeDims++sCopy ::+  VName ->+  Imp.TExp Int64 ->+  Space ->+  VName ->+  Imp.TExp Int64 ->+  Space ->+  Count Elements (Imp.TExp Int64) ->+  PrimType ->+  ImpM rep r op ()+sCopy destmem destoffset destspace srcmem srcoffset srcspace num_elems pt =+  if destmem == srcmem+    then sUnless (destoffset .==. srcoffset) the_copy+    else the_copy+  where+    the_copy =+      emit+        $ Imp.Copy+          pt+          destmem+          (bytes destoffset)+          destspace+          srcmem+          (bytes srcoffset)+          srcspace+        $ num_elems `withElemType` pt  -- | Untyped assignment. (<~~) :: VName -> Imp.Exp -> ImpM rep r op ()

@@ -16,7 +16,6 @@ import Data.List (foldl') import Data.Map qualified as M import Data.Maybe-import Futhark.CodeGen.ImpCode.GPU (bytes) import Futhark.CodeGen.ImpCode.GPU qualified as Imp import Futhark.CodeGen.ImpGen hiding (compileProg) import Futhark.CodeGen.ImpGen qualified@@ -166,12 +165,12 @@   CallKernelGen () segOpCompiler pat (SegMap lvl space _ kbody) =   compileSegMap pat lvl space kbody-segOpCompiler pat (SegRed lvl@SegThread {} space reds _ kbody) =+segOpCompiler pat (SegRed lvl@(SegThread _ _) space reds _ kbody) =   compileSegRed pat lvl space reds kbody-segOpCompiler pat (SegScan lvl@SegThread {} space scans _ kbody) =+segOpCompiler pat (SegScan lvl@(SegThread _ _) space scans _ kbody) =   compileSegScan pat lvl space scans kbody-segOpCompiler pat (SegHist (SegThread num_groups group_size _) space ops _ kbody) =-  compileSegHist pat num_groups group_size space ops kbody+segOpCompiler pat (SegHist lvl@(SegThread _ _) space ops _ kbody) =+  compileSegHist pat lvl space ops kbody segOpCompiler pat segop =   compilerBugS $ "segOpCompiler: unexpected " ++ prettyString (segLevel segop) ++ " for rhs of pattern " ++ prettyString pat @@ -298,17 +297,16 @@       let num_elems = Imp.elements $ product $ map pe64 srcshape       srcspace <- entryMemSpace <$> lookupMemory srcmem       destspace <- entryMemSpace <$> lookupMemory destmem-      emit-        $ Imp.Copy-          bt-          destmem-          (bytes $ sExt64 destoffset)-          destspace-          srcmem-          (bytes $ sExt64 srcoffset)-          srcspace-        $ num_elems `Imp.withElemType` bt-  | otherwise = sCopy bt destloc srcloc+      sCopy+        destmem+        (sExt64 destoffset)+        destspace+        srcmem+        (sExt64 srcoffset)+        srcspace+        num_elems+        bt+  | otherwise = sCopyKernel bt destloc srcloc  mapTransposeForType :: PrimType -> CallKernelGen Name mapTransposeForType bt = do

@@ -19,6 +19,8 @@     sKernelThread,     KernelAttrs (..),     defKernelAttrs,+    sCopyKernel,+    lvlKernelAttrs,     allocLocal,     kernelAlloc,     compileThreadResult,@@ -1076,6 +1078,33 @@       kAttrGroupSize = group_size     } +getSize :: String -> SizeClass -> CallKernelGen (TV Int64)+getSize desc size_class = do+  v <- dPrim desc int64+  fname <- askFunction+  let v_key = keyWithEntryPoint fname $ nameFromString $ prettyString $ tvVar v+  sOp $ Imp.GetSize (tvVar v) v_key size_class+  pure v++-- | Compute kernel attributes from 'SegLevel'; including synthesising+-- group-size and thread count if no grid is provided.+lvlKernelAttrs :: SegLevel -> CallKernelGen KernelAttrs+lvlKernelAttrs lvl =+  case lvl of+    SegThread _ Nothing -> mkGrid+    SegThread _ (Just (KernelGrid num_groups group_size)) ->+      pure $ defKernelAttrs num_groups group_size+    SegGroup _ Nothing -> mkGrid+    SegGroup _ (Just (KernelGrid num_groups group_size)) ->+      pure $ defKernelAttrs num_groups group_size+    SegThreadInGroup {} ->+      error "lvlKernelAttrs: SegThreadInGroup"+  where+    mkGrid = do+      group_size <- getSize "group_size" Imp.SizeGroup+      num_groups <- getSize "num_groups" Imp.SizeNumGroups+      pure $ defKernelAttrs (Count $ tvSize num_groups) (Count $ tvSize group_size)+ sKernel ::   Operations GPUMem KernelEnv Imp.KernelOp ->   (KernelConstants -> Imp.TExp Int32) ->@@ -1318,8 +1347,8 @@           [Imp.MemArg arr_mem, Imp.ExpArg $ untyped n, Imp.ExpArg x, Imp.ExpArg s]     else sIotaKernel arr n x s et -sCopy :: CopyCompiler GPUMem HostEnv Imp.HostOp-sCopy pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do+sCopyKernel :: CopyCompiler GPUMem HostEnv Imp.HostOp+sCopyKernel 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 pe64 srcdims

@@ -195,14 +195,8 @@   bimap (map fst) (map fst) $     partition ((`elem` seq_is) . snd) (zip (unSegSpace space) [0 ..]) -sanityCheckLevel :: SegLevel -> InKernelGen ()-sanityCheckLevel SegThread {} = pure ()-sanityCheckLevel SegGroup {} =-  error "compileGroupOp: unexpected group-level SegOp."--compileFlatId :: SegLevel -> SegSpace -> InKernelGen ()-compileFlatId lvl space = do-  sanityCheckLevel lvl+compileFlatId :: SegSpace -> InKernelGen ()+compileFlatId space = do   ltid <- kernelLocalThreadId . kernelConstants <$> askEnv   dPrimV_ (segFlat space) ltid @@ -342,7 +336,7 @@ compileGroupOp pat (Alloc size space) =   kernelAlloc pat size space compileGroupOp pat (Inner (SegOp (SegMap lvl space _ body))) = do-  compileFlatId lvl space+  compileFlatId space    groupCoverSegSpace (segVirt lvl) space $     compileStms mempty (kernelBodyStms body) $@@ -350,7 +344,7 @@         kernelBodyResult body   sOp $ Imp.ErrorSync Imp.FenceLocal compileGroupOp pat (Inner (SegOp (SegScan lvl space scans _ body))) = do-  compileFlatId lvl space+  compileFlatId space    let (ltids, dims) = unzip $ unSegSpace space       dims' = map pe64 dims@@ -397,7 +391,7 @@         (segBinOpLambda scan)         arrs_flat compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do-  compileFlatId lvl space+  compileFlatId space    let dims' = map pe64 dims       mkTempArr t =@@ -512,7 +506,7 @@        sOp $ Imp.Barrier Imp.FenceLocal compileGroupOp pat (Inner (SegOp (SegHist lvl space ops _ kbody))) = do-  compileFlatId lvl space+  compileFlatId space   let (ltids, _dims) = unzip $ unSegSpace space    -- We don't need the red_pes, because it is guaranteed by our type@@ -522,7 +516,8 @@       (_red_pes, map_pes) =         splitAt num_red_res $ patElems pat -  ops' <- prepareIntraGroupSegHist (segGroupSize lvl) ops+  group_size <- kernelGroupSizeCount . kernelConstants <$> askEnv+  ops' <- prepareIntraGroupSegHist group_size ops    -- Ensure that all locks have been initialised.   sOp $ Imp.Barrier Imp.FenceLocal

@@ -1046,13 +1046,13 @@ -- | Generate code for a segmented histogram called from the host. compileSegHist ::   Pat LetDecMem ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  SegLevel ->   SegSpace ->   [HistOp GPUMem] ->   KernelBody GPUMem ->   CallKernelGen ()-compileSegHist (Pat pes) num_groups group_size space ops kbody = do+compileSegHist (Pat pes) lvl space ops kbody = do+  KernelAttrs _ _ num_groups group_size <- lvlKernelAttrs lvl   -- Most of this function is not the histogram part itself, but   -- rather figuring out whether to use a local or global memory   -- strategy, as well as collapsing the subhistograms produced (which@@ -1152,7 +1152,7 @@          flat_gtid <- newVName "flat_gtid" -        let lvl = SegThread num_groups group_size SegVirt+        let grid = KernelGrid num_groups group_size             segred_space =               SegSpace flat_gtid $                 segment_dims@@ -1161,7 +1161,7 @@                   ++ [(subhistogram_id, Var $ tvVar num_histos)]          let segred_op = SegBinOp Commutative (histOp op) (histNeutral op) mempty-        compileSegRed' (Pat red_pes) lvl segred_space [segred_op] $ \red_cont ->+        compileSegRed' (Pat red_pes) grid segred_space [segred_op] $ \red_cont ->           red_cont . flip map subhistos $ \subhisto ->             ( Var subhisto,               map Imp.le64 $

@@ -23,10 +23,11 @@   KernelBody GPUMem ->   CallKernelGen () compileSegMap pat lvl space kbody = do+  attrs <- lvlKernelAttrs lvl+   let (is, dims) = unzip $ unSegSpace space       dims' = map pe64 dims-      group_size' = pe64 <$> segGroupSize lvl-      attrs = defKernelAttrs (segNumGroups lvl) (segGroupSize lvl)+      group_size' = pe64 <$> kAttrGroupSize attrs    emit $ Imp.DebugPrint "\n# SegMap" Nothing   case lvl of@@ -58,4 +59,6 @@             compileStms mempty (kernelBodyStms kbody) $               zipWithM_ (compileGroupResult space) (patElems pat) $                 kernelBodyResult kbody+    SegThreadInGroup {} ->+      error "compileSegMap: SegThreadInGroup"   emit $ Imp.DebugPrint "" Nothing

@@ -82,8 +82,10 @@   [SegBinOp GPUMem] ->   KernelBody GPUMem ->   CallKernelGen ()-compileSegRed pat lvl space reds body =-  compileSegRed' pat lvl space reds $ \red_cont ->+compileSegRed pat lvl space reds body = do+  KernelAttrs _ _ num_groups group_size <- lvlKernelAttrs lvl+  let grid = KernelGrid num_groups group_size+  compileSegRed' pat grid space reds $ \red_cont ->     compileStms mempty (kernelBodyStms body) $ do       let (red_res, map_res) = splitAt (segBinOpResults reds) $ kernelBodyResult body @@ -96,12 +98,12 @@ -- | Like 'compileSegRed', but where the body is a monadic action. compileSegRed' ::   Pat LetDecMem ->-  SegLevel ->+  KernelGrid ->   SegSpace ->   [SegBinOp GPUMem] ->   DoSegBody ->   CallKernelGen ()-compileSegRed' pat lvl space reds body+compileSegRed' pat grid space reds body   | genericLength reds > maxNumOps =       compilerLimitationS $         "compileSegRed': at most " ++ show maxNumOps ++ " reduction operators are supported."@@ -116,8 +118,8 @@         (smallSegmentsReduction pat num_groups group_size space reds body)         (largeSegmentsReduction pat num_groups group_size space reds body)   where-    num_groups = segNumGroups lvl-    group_size = segGroupSize lvl+    num_groups = gridNumGroups grid+    group_size = gridGroupSize grid  -- | Prepare intermediate arrays for the reduction.  Prim-typed -- arguments go in local memory (so we need to do the allocation of

@@ -215,6 +215,7 @@   KernelBody GPUMem ->   CallKernelGen () compileSegScan pat lvl space scanOp kbody = do+  attrs <- lvlKernelAttrs lvl   let Pat all_pes = pat       scanOpNe = segBinOpNeutral scanOp       tys = map (\(Prim pt) -> pt) $ lambdaReturnType $ segBinOpLambda scanOp@@ -233,7 +234,7 @@       mem_constraint = max k_mem sumT `div` maxT       reg_constraint = (k_reg - 1 - sumT') `div` (2 * sumT') -      group_size = segGroupSize lvl+      group_size = kAttrGroupSize attrs       group_size' = pe64 $ unCount group_size    num_groups <-@@ -274,7 +275,7 @@     constants <- kernelConstants <$> askEnv      (sharedId, transposedArrays, prefixArrays, warpscan, exchanges) <--      createLocalArrays (segGroupSize lvl) (intConst Int64 m) tys+      createLocalArrays (kAttrGroupSize attrs) (intConst Int64 m) tys      dynamicId <- dPrim "dynamic_id" int32     sWhen (kernelLocalThreadId constants .==. 0) $ do

@@ -491,6 +491,8 @@   KernelBody GPUMem ->   CallKernelGen () compileSegScan pat lvl space scans kbody = do+  attrs <- lvlKernelAttrs lvl+   -- Since stage 2 involves a group size equal to the number of groups   -- used for stage 1, we have to cap this number to the maximum group   -- size.@@ -501,14 +503,13 @@     fmap (Imp.Count . tvSize) $       dPrimV "stage1_num_groups" $         sMin64 (tvExp stage1_max_num_groups) $-          pe64 $-            Imp.unCount $-              segNumGroups lvl+          pe64 . Imp.unCount . kAttrNumGroups $+            attrs    (stage1_num_threads, elems_per_group, crossesSegment) <--    scanStage1 pat stage1_num_groups (segGroupSize lvl) space scans kbody+    scanStage1 pat stage1_num_groups (kAttrGroupSize attrs) space scans kbody    emit $ Imp.DebugPrint "elems_per_group" $ Just $ untyped elems_per_group    scanStage2 pat stage1_num_threads elems_per_group stage1_num_groups crossesSegment space scans-  scanStage3 pat (segNumGroups lvl) (segGroupSize lvl) elems_per_group crossesSegment space scans+  scanStage3 pat (kAttrNumGroups attrs) (kAttrGroupSize attrs) elems_per_group crossesSegment space scans

@@ -86,6 +86,8 @@     eAny,     eDimInBounds,     eOutOfBounds,+    eIndex,+    eLast,      -- * Other building blocks     asIntZ,@@ -413,6 +415,22 @@           BasicOp $             BinOp LogOr less_than_zero greater_than_size   foldBinOp LogOr (constant False) =<< zipWithM checkDim ws is'++-- | The array element at this index.+eIndex :: MonadBuilder m => VName -> m (Exp (Rep m)) -> m (Exp (Rep m))+eIndex arr i = do+  i' <- letSubExp "i" =<< i+  arr_t <- lookupType arr+  pure $ BasicOp $ Index arr $ fullSlice arr_t [DimFix i']++-- | The last element of the given array.+eLast :: MonadBuilder m => VName -> m (Exp (Rep m))+eLast arr = do+  n <- arraySize 0 <$> lookupType arr+  nm1 <-+    letSubExp "nm1" . BasicOp $+      BinOp (Sub Int64 OverflowUndef) n (intConst Int64 1)+  eIndex arr (eSubExp nm1)  -- | Construct an unspecified value of the given type. eBlank :: MonadBuilder m => Type -> m (Exp (Rep m))

@@ -12,6 +12,10 @@      -- * SegOp refinements     SegLevel (..),+    segVirt,+    SegVirt (..),+    SegSeqDims (..),+    KernelGrid (..),      -- * Reexports     module Futhark.IR.GPU.Sizes,@@ -44,68 +48,118 @@   ) import Futhark.Util.Pretty qualified as PP +-- | These dimensions (indexed from 0, outermost) of the corresponding+-- 'SegSpace' should not be parallelised, but instead iterated+-- sequentially.  For example, with a 'SegSeqDims' of @[0]@ and a+-- 'SegSpace' with dimensions @[n][m]@, there will be an outer loop+-- with @n@ iterations, while the @m@ dimension will be parallelised.+--+-- Semantically, this has no effect, but it may allow reductions in+-- memory usage or other low-level optimisations.  Operationally, the+-- guarantee is that for a SegSeqDims of e.g. @[i,j,k]@, threads+-- running at any given moment will always have the same indexes along+-- the dimensions specified by @[i,j,k]@.+--+-- At the moment, this is only supported for 'SegNoVirtFull'+-- intra-group parallelism in GPU code, as we have not yet found it+-- useful anywhere else.+newtype SegSeqDims = SegSeqDims {segSeqDims :: [Int]}+  deriving (Eq, Ord, Show)++-- | Do we need group-virtualisation when generating code for the+-- segmented operation?  In most cases, we do, but for some simple+-- kernels, we compute the full number of groups in advance, and then+-- virtualisation is an unnecessary (but generally very small)+-- overhead.  This only really matters for fairly trivial but very+-- wide @map@ kernels where each thread performs constant-time work on+-- scalars.+data SegVirt+  = SegVirt+  | SegNoVirt+  | -- | Not only do we not need virtualisation, but we _guarantee_+    -- that all physical threads participate in the work.  This can+    -- save some checks in code generation.+    SegNoVirtFull SegSeqDims+  deriving (Eq, Ord, Show)++-- | The actual, physical grid dimensions used for the GPU kernel+-- running this 'SegOp'.+data KernelGrid = KernelGrid+  { gridNumGroups :: Count NumGroups SubExp,+    gridGroupSize :: Count GroupSize SubExp+  }+  deriving (Eq, Ord, Show)+ -- | At which level the *body* of a t'SegOp' executes. data SegLevel-  = SegThread-      { segNumGroups :: Count NumGroups SubExp,-        segGroupSize :: Count GroupSize SubExp,-        segVirt :: SegVirt-      }-  | SegGroup-      { segNumGroups :: Count NumGroups SubExp,-        segGroupSize :: Count GroupSize SubExp,-        segVirt :: SegVirt-      }+  = SegThread SegVirt (Maybe KernelGrid)+  | SegGroup SegVirt (Maybe KernelGrid)+  | SegThreadInGroup SegVirt   deriving (Eq, Ord, Show) +-- | The 'SegVirt' of the 'SegLevel'.+segVirt :: SegLevel -> SegVirt+segVirt (SegThread v _) = v+segVirt (SegGroup v _) = v+segVirt (SegThreadInGroup v) = v++instance PP.Pretty SegVirt where+  pretty SegNoVirt = mempty+  pretty (SegNoVirtFull dims) = "full" <+> pretty (segSeqDims dims)+  pretty SegVirt = "virtualise"++instance PP.Pretty KernelGrid where+  pretty (KernelGrid num_groups group_size) =+    "groups=" <> pretty num_groups <> PP.semi+      <+> "groupsize=" <> pretty group_size+ instance PP.Pretty SegLevel where-  pretty lvl =-    PP.parens-      ( lvl' <> PP.semi-          <+> "#groups=" <> pretty (segNumGroups lvl) <> PP.semi-          <+> "groupsize=" <> pretty (segGroupSize lvl) <> virt-      )-    where-      lvl' = case lvl of-        SegThread {} -> "thread"-        SegGroup {} -> "group"-      virt = case segVirt lvl of-        SegNoVirt -> mempty-        SegNoVirtFull dims -> PP.semi <+> "full" <+> pretty (segSeqDims dims)-        SegVirt -> PP.semi <+> "virtualise"+  pretty (SegThread virt grid) =+    PP.parens ("thread" <> PP.semi <+> pretty virt <> PP.semi <+> pretty grid)+  pretty (SegGroup virt grid) =+    PP.parens ("group" <> PP.semi <+> pretty virt <> PP.semi <+> pretty grid)+  pretty (SegThreadInGroup virt) =+    PP.parens ("ingroup" <> PP.semi <+> pretty virt) -instance Engine.Simplifiable SegLevel where-  simplify (SegThread num_groups group_size virt) =-    SegThread-      <$> traverse Engine.simplify num_groups-      <*> traverse Engine.simplify group_size-      <*> pure virt-  simplify (SegGroup num_groups group_size virt) =-    SegGroup+instance Engine.Simplifiable KernelGrid where+  simplify (KernelGrid num_groups group_size) =+    KernelGrid       <$> traverse Engine.simplify num_groups       <*> traverse Engine.simplify group_size-      <*> pure virt -instance Substitute SegLevel where-  substituteNames substs (SegThread num_groups group_size virt) =-    SegThread-      (substituteNames substs num_groups)-      (substituteNames substs group_size)-      virt-  substituteNames substs (SegGroup num_groups group_size virt) =-    SegGroup+instance Engine.Simplifiable SegLevel where+  simplify (SegThread virt grid) =+    SegThread virt <$> Engine.simplify grid+  simplify (SegGroup virt grid) =+    SegGroup virt <$> Engine.simplify grid+  simplify (SegThreadInGroup virt) =+    pure $ SegThreadInGroup virt++instance Substitute KernelGrid where+  substituteNames substs (KernelGrid num_groups group_size) =+    KernelGrid       (substituteNames substs num_groups)       (substituteNames substs group_size)-      virt +instance Substitute SegLevel where+  substituteNames substs (SegThread virt grid) =+    SegThread virt (substituteNames substs grid)+  substituteNames substs (SegGroup virt grid) =+    SegGroup virt (substituteNames substs grid)+  substituteNames _ (SegThreadInGroup virt) =+    SegThreadInGroup virt+ instance Rename SegLevel where   rename = substituteRename +instance FreeIn KernelGrid where+  freeIn' (KernelGrid num_groups group_size) =+    freeIn' (num_groups, group_size)+ instance FreeIn SegLevel where-  freeIn' (SegThread num_groups group_size _) =-    freeIn' num_groups <> freeIn' group_size-  freeIn' (SegGroup num_groups group_size _) =-    freeIn' num_groups <> freeIn' group_size+  freeIn' (SegThread _virt grid) = freeIn' grid+  freeIn' (SegGroup _virt grid) = freeIn' grid+  freeIn' (SegThreadInGroup _virt) = mempty  -- | A simple size-level query or computation. data SizeOp@@ -305,22 +359,34 @@   opMetrics (SizeOp op) = opMetrics op   opMetrics (GPUBody _ body) = inside "GPUBody" $ bodyMetrics body +checkGrid :: TC.Checkable rep => KernelGrid -> TC.TypeM rep ()+checkGrid grid = do+  TC.require [Prim int64] $ unCount $ gridNumGroups grid+  TC.require [Prim int64] $ unCount $ gridGroupSize grid+ checkSegLevel ::   TC.Checkable rep =>   Maybe SegLevel ->   SegLevel ->   TC.TypeM rep ()-checkSegLevel Nothing lvl = do-  TC.require [Prim int64] $ unCount $ segNumGroups lvl-  TC.require [Prim int64] $ unCount $ segGroupSize lvl+checkSegLevel (Just SegGroup {}) (SegThreadInGroup _virt) =+  pure ()+checkSegLevel _ (SegThreadInGroup _virt) =+  TC.bad $ TC.TypeError "ingroup SegOp not in group SegOp." checkSegLevel (Just SegThread {}) _ =   TC.bad $ TC.TypeError "SegOps cannot occur when already at thread level."-checkSegLevel (Just x) y-  | x == y = TC.bad $ TC.TypeError $ "Already at at level " <> prettyText x-  | segNumGroups x /= segNumGroups y || segGroupSize x /= segGroupSize y =-      TC.bad $ TC.TypeError "Physical layout for SegLevel does not match parent SegLevel."-  | otherwise =-      pure ()+checkSegLevel (Just SegThreadInGroup {}) _ =+  TC.bad $ TC.TypeError "SegOps cannot occur when already at ingroup level."+checkSegLevel _ (SegThread _virt Nothing) =+  pure ()+checkSegLevel (Just _) SegThread {} =+  TC.bad $ TC.TypeError "thread-level SegOp cannot be nested"+checkSegLevel Nothing (SegThread _virt grid) =+  mapM_ checkGrid grid+checkSegLevel (Just _) SegGroup {} =+  TC.bad $ TC.TypeError "group-level SegOp cannot be nested"+checkSegLevel Nothing (SegGroup _virt grid) =+  mapM_ checkGrid grid  typeCheckHostOp ::   TC.Checkable rep =>

@@ -28,8 +28,7 @@ import Futhark.Transform.Rename import Futhark.Transform.Substitute import Futhark.Util.Pretty-  ( Pretty,-    nestedBlock,+  ( nestedBlock,     pretty,     (<+>),     (</>),

@@ -131,7 +131,7 @@ import Futhark.Transform.Rename import Futhark.Transform.Substitute import Futhark.Util-import Futhark.Util.Pretty (docText, indent, ppTuple', pretty, (<+>), (</>))+import Futhark.Util.Pretty (docText, indent, ppTupleLines', pretty, (<+>), (</>)) import Futhark.Util.Pretty qualified as PP import Prelude hiding (id, (.)) @@ -748,9 +748,9 @@       bad s =         TC.bad . TC.TypeError . docText $           "Return type"-            </> indent 2 (ppTuple' $ map pretty rettype)+            </> indent 2 (ppTupleLines' $ map pretty rettype)             </> "cannot match returns of results"-            </> indent 2 (ppTuple' $ map pretty ts)+            </> indent 2 (ppTupleLines' $ map pretty ts)             </> pretty s    either bad pure =<< runExceptT (zipWithM_ checkReturn rettype ts)@@ -772,9 +772,9 @@    unless ok . TC.bad . TC.TypeError . docText $     "Expression type:"-      </> indent 2 (ppTuple' $ map pretty rt)+      </> indent 2 (ppTupleLines' $ map pretty rt)       </> "cannot match pattern type:"-      </> indent 2 (ppTuple' $ map pretty val_ts)+      </> indent 2 (ppTupleLines' $ map pretty val_ts)   where     matches _ _ (MemPrim x) (MemPrim y) = x == y     matches _ _ (MemMem x_space) (MemMem y_space) =

@@ -49,7 +49,7 @@   lexeme . fmap nameFromString $     (:) <$> satisfy leading <*> many (satisfy constituent)   where-    leading c = isAlpha c || c `elem` ("_+-*/%=!<>|&^." :: String)+    leading c = isAlpha c || c `elem` ("_+-*/%=!<>|&^.#" :: String)  pVName :: Parser VName pVName = lexeme $ do@@ -910,23 +910,36 @@  pSegLevel :: Parser GPU.SegLevel pSegLevel =-  parens $-    choice-      [ keyword "thread" $> GPU.SegThread,-        keyword "group" $> GPU.SegGroup-      ]-      <*> (pSemi *> lexeme "#groups=" $> GPU.Count <*> pSubExp)-      <*> (pSemi *> lexeme "groupsize=" $> GPU.Count <*> pSubExp)-      <*> choice-        [ pSemi-            *> choice-              [ keyword "full"-                  $> SegOp.SegNoVirtFull-                  <*> (SegOp.SegSeqDims <$> brackets (pInt `sepBy` pComma)),-                keyword "virtualise" $> SegOp.SegVirt-              ],-          pure SegOp.SegNoVirt+  parens . choice $+    [ "thread"+        $> GPU.SegThread+        <* pSemi+        <*> pSegVirt+        <* pSemi+        <*> optional pKernelGrid,+      "group"+        $> GPU.SegGroup+        <* pSemi+        <*> pSegVirt+        <* pSemi+        <*> optional pKernelGrid,+      "ingroup" $> GPU.SegThreadInGroup <* pSemi <*> pSegVirt+    ]+  where+    pSegVirt =+      choice+        [ choice+            [ keyword "full"+                $> GPU.SegNoVirtFull+                <*> (GPU.SegSeqDims <$> brackets (pInt `sepBy` pComma)),+              keyword "virtualise" $> GPU.SegVirt+            ],+          pure GPU.SegNoVirt         ]+    pKernelGrid =+      GPU.KernelGrid+        <$> (lexeme "groups=" $> GPU.Count <*> pSubExp <* pSemi)+        <*> (lexeme "groupsize=" $> GPU.Count <*> pSubExp)  pHostOp :: PR rep -> Parser op -> Parser (GPU.HostOp rep op) pHostOp pr pOther =

@@ -278,7 +278,7 @@       <+> "else"       <+> maybeNest f       </> colon-      <+> ppTuple' (map pretty ret)+      <+> ppTupleLines' (map pretty ret)     where       info' = case ifsort of         MatchNormal -> mempty@@ -291,7 +291,7 @@       <+> "->"       <+> maybeNest defb       </> colon-      <+> ppTuple' (map pretty ret)+      <+> ppTupleLines' (map pretty ret)     where       info' = case ifsort of         MatchNormal -> mempty@@ -359,7 +359,7 @@   pretty (Lambda params body rettype) =     "\\"       <+> ppTuple' (map pretty params)-      </> indent 2 (colon <+> ppTupleLines' rettype <+> "->")+      </> indent 2 (colon <+> ppTupleLines' (map pretty rettype) <+> "->")       </> indent 2 (pretty body)  instance Pretty Signedness where@@ -385,8 +385,8 @@     annot (attrAnnots attrs) $       fun         </> indent 2 (pretty (nameToString name))-        <+> apply (map pretty fparams)-        </> indent 2 (colon <+> align (ppTuple' $ map pretty rettype))+        <+> parens (commastack $ map pretty fparams)+        </> indent 2 (colon <+> align (ppTupleLines' $ map pretty rettype))         <+> equals         <+> nestedBlock "{" "}" (pretty body)     where@@ -396,8 +396,8 @@           "entry"             <> (parens . align)               ( "\"" <> pretty p_name <> "\"" <> comma-                  </> ppTupleLines' p_entry <> comma-                  </> ppTupleLines' ret_entry+                  </> ppTupleLines' (map pretty p_entry) <> comma+                  </> ppTupleLines' (map pretty ret_entry)               )  instance Pretty OpaqueType where@@ -420,7 +420,3 @@ instance Pretty d => Pretty (DimIndex d) where   pretty (DimFix i) = pretty i   pretty (DimSlice i n s) = pretty i <+> ":+" <+> pretty n <+> "*" <+> pretty s---- | Like 'prettyTupleLines', but produces a 'Doc'.-ppTupleLines' :: Pretty a => [a] -> Doc b-ppTupleLines' = braces . stack . punctuate comma . map (align . pretty)

@@ -54,7 +54,6 @@ import Futhark.Transform.Rename (Rename, Renameable) import Futhark.Transform.Substitute (Substitutable, Substitute) import Futhark.Util (maybeNth)-import Futhark.Util.Pretty  -- | @isBuiltInFunction k@ is 'True' if @k@ is an element of 'builtInFunctions'. isBuiltInFunction :: Name -> Bool

@@ -47,6 +47,7 @@ import Data.IntMap.Strict qualified as IM import Data.IntSet qualified as IS import Data.Map.Strict qualified as M+import Data.Set qualified as S import Futhark.IR.Prop.Patterns import Futhark.IR.Prop.Scope import Futhark.IR.Syntax@@ -213,6 +214,9 @@   freeIn' (a, b, c, d) = freeIn' a <> freeIn' b <> freeIn' c <> freeIn' d  instance FreeIn a => FreeIn [a] where+  freeIn' = foldMap freeIn'++instance FreeIn a => FreeIn (S.Set a) where   freeIn' = foldMap freeIn'  instance

@@ -69,7 +69,7 @@ import Futhark.Transform.Rename import Futhark.Transform.Substitute import Futhark.Util (chunks, maybeNth)-import Futhark.Util.Pretty (Doc, Pretty, align, comma, commasep, docText, parens, ppTuple', pretty, (<+>), (</>))+import Futhark.Util.Pretty (Doc, align, comma, commasep, docText, parens, ppTuple', pretty, (<+>), (</>)) import Futhark.Util.Pretty qualified as PP import Prelude hiding (id, (.)) 

@@ -83,12 +83,12 @@   Simplify.SimplifiableRep rep =>   Simplify.SimplifyOp rep (SOAC (Wise rep)) simplifySOAC (VJP lam arr vec) = do-  (lam', hoisted) <- Engine.simplifyLambda lam+  (lam', hoisted) <- Engine.simplifyLambda mempty lam   arr' <- mapM Engine.simplify arr   vec' <- mapM Engine.simplify vec   pure (VJP lam' arr' vec', hoisted) simplifySOAC (JVP lam arr vec) = do-  (lam', hoisted) <- Engine.simplifyLambda lam+  (lam', hoisted) <- Engine.simplifyLambda mempty lam   arr' <- mapM Engine.simplify arr   vec' <- mapM Engine.simplify vec   pure (JVP lam' arr' vec', hoisted)@@ -96,11 +96,11 @@   outerdim' <- Engine.simplify outerdim   nes' <- mapM Engine.simplify nes   arr' <- mapM Engine.simplify arr-  (lam', lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam+  (lam', lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam   pure (Stream outerdim' arr' nes' lam', lam_hoisted) simplifySOAC (Scatter w ivs lam as) = do   w' <- Engine.simplify w-  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda lam+  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam   ivs' <- mapM Engine.simplify ivs   as' <- mapM Engine.simplify as   pure (Scatter w' ivs' lam' as', hoisted)@@ -112,25 +112,25 @@       rf' <- Engine.simplify rf       dests' <- Engine.simplify dests       nes' <- mapM Engine.simplify nes-      (op', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda op+      (op', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty op       pure (HistOp dests_w' rf' dests' nes' op', hoisted)   imgs' <- mapM Engine.simplify imgs-  (bfun', bfun_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda bfun+  (bfun', bfun_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty bfun   pure (Hist w' imgs' ops' bfun', mconcat hoisted <> bfun_hoisted) simplifySOAC (Screma w arrs (ScremaForm scans reds map_lam)) = do   (scans', scans_hoisted) <- fmap unzip $     forM scans $ \(Scan lam nes) -> do-      (lam', hoisted) <- Engine.simplifyLambda lam+      (lam', hoisted) <- Engine.simplifyLambda mempty lam       nes' <- Engine.simplify nes       pure (Scan lam' nes', hoisted)    (reds', reds_hoisted) <- fmap unzip $     forM reds $ \(Reduce comm lam nes) -> do-      (lam', hoisted) <- Engine.simplifyLambda lam+      (lam', hoisted) <- Engine.simplifyLambda mempty lam       nes' <- Engine.simplify nes       pure (Reduce comm lam' nes', hoisted) -  (map_lam', map_lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda map_lam+  (map_lam', map_lam_hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty map_lam    (,)     <$> ( Screma

@@ -6,8 +6,6 @@ -- over @iota@s (so there will be explicit indexing inside them). module Futhark.IR.SegOp   ( SegOp (..),-    SegVirt (..),-    SegSeqDims (..),     segLevel,     segBody,     segSpace,@@ -88,7 +86,6 @@ import Futhark.Util (chunks, maybeNth) import Futhark.Util.Pretty   ( Doc,-    Pretty,     apply,     hsep,     parens,@@ -436,40 +433,6 @@       onDim (dim, blk_tile, reg_tile) =         pretty dim <+> "/" <+> parens (pretty blk_tile <+> "*" <+> pretty reg_tile) --- | These dimensions (indexed from 0, outermost) of the corresponding--- 'SegSpace' should not be parallelised, but instead iterated--- sequentially.  For example, with a 'SegSeqDims' of @[0]@ and a--- 'SegSpace' with dimensions @[n][m]@, there will be an outer loop--- with @n@ iterations, while the @m@ dimension will be parallelised.------ Semantically, this has no effect, but it may allow reductions in--- memory usage or other low-level optimisations.  Operationally, the--- guarantee is that for a SegSeqDims of e.g. @[i,j,k]@, threads--- running at any given moment will always have the same indexes along--- the dimensions specified by @[i,j,k]@.------ At the moment, this is only supported for 'SegNoVirtFull'--- intra-group parallelism in GPU code, as we have not yet found it--- useful anywhere else.-newtype SegSeqDims = SegSeqDims {segSeqDims :: [Int]}-  deriving (Eq, Ord, Show)---- | Do we need group-virtualisation when generating code for the--- segmented operation?  In most cases, we do, but for some simple--- kernels, we compute the full number of groups in advance, and then--- virtualisation is an unnecessary (but generally very small)--- overhead.  This only really matters for fairly trivial but very--- wide @map@ kernels where each thread performs constant-time work on--- scalars.-data SegVirt-  = SegVirt-  | SegNoVirt-  | -- | Not only do we not need virtualisation, but we _guarantee_-    -- that all physical threads participate in the work.  This can-    -- save some checks in code generation.-    SegNoVirtFull SegSeqDims-  deriving (Eq, Ord, Show)- -- | Index space of a 'SegOp'. data SegSpace = SegSpace   { -- | Flat physical index corresponding to the@@ -1134,9 +1097,10 @@  simplifyLambda ::   Engine.SimplifiableRep rep =>+  Names ->   Lambda (Wise rep) ->   Engine.SimpleM rep (Lambda (Wise rep), Stms (Wise rep))-simplifyLambda = Engine.blockMigrated . Engine.simplifyLambda+simplifyLambda bound = Engine.blockMigrated . Engine.simplifyLambda bound  segSpaceSymbolTable :: ASTRep rep => SegSpace -> ST.SymbolTable rep segSpaceSymbolTable (SegSpace flat gtids_and_dims) =@@ -1146,12 +1110,13 @@  simplifySegBinOp ::   Engine.SimplifiableRep rep =>+  VName ->   SegBinOp (Wise rep) ->   Engine.SimpleM rep (SegBinOp (Wise rep), Stms (Wise rep))-simplifySegBinOp (SegBinOp comm lam nes shape) = do+simplifySegBinOp phys_id (SegBinOp comm lam nes shape) = do   (lam', hoisted) <-     Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $-      simplifyLambda lam+      simplifyLambda (oneName phys_id) lam   shape' <- Engine.simplify shape   nes' <- mapM Engine.simplify nes   pure (SegBinOp comm lam' nes' shape', hoisted)@@ -1175,7 +1140,7 @@   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)   (reds', reds_hoisted) <-     Engine.localVtable (<> scope_vtable) $-      unzip <$> mapM simplifySegBinOp reds+      unzip <$> mapM (simplifySegBinOp (segFlat space)) reds   (kbody', body_hoisted) <- simplifyKernelBody space kbody    pure@@ -1189,7 +1154,7 @@   (lvl', space', ts') <- Engine.simplify (lvl, space, ts)   (scans', scans_hoisted) <-     Engine.localVtable (<> scope_vtable) $-      unzip <$> mapM simplifySegBinOp scans+      unzip <$> mapM (simplifySegBinOp (segFlat space)) scans   (kbody', body_hoisted) <- simplifyKernelBody space kbody    pure@@ -1213,7 +1178,7 @@         (lam', op_hoisted) <-           Engine.localVtable (<> scope_vtable) $             Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $-              simplifyLambda lam+              simplifyLambda (oneName (segFlat space)) lam         pure           ( HistOp w' rf' arrs' nes' dims' lam',             op_hoisted

@@ -100,6 +100,7 @@   ( module Language.Futhark.Core,     prettyString,     prettyText,+    Pretty,     module Futhark.IR.Rep,     module Futhark.IR.Syntax.Core, @@ -171,7 +172,7 @@ import Data.Traversable (fmapDefault, foldMapDefault) import Futhark.IR.Rep import Futhark.IR.Syntax.Core-import Futhark.Util.Pretty (prettyString, prettyText)+import Futhark.Util.Pretty (Pretty, prettyString, prettyText) import Language.Futhark.Core import Prelude hiding (id, (.)) @@ -440,7 +441,14 @@  -- | For-loop or while-loop? data LoopForm rep-  = ForLoop VName IntType SubExp [(LParam rep, VName)]+  = ForLoop+      VName+      -- ^ The loop iterator var+      IntType+      -- ^ The type of the loop iterator var+      SubExp+      -- ^ The number of iterations.+      [(LParam rep, VName)]   | WhileLoop VName  deriving instance RepTypes rep => Eq (LoopForm rep)

@@ -57,7 +57,7 @@ import Futhark.Construct (instantiateShapes) import Futhark.IR.Aliases hiding (lookupAliases) import Futhark.Util-import Futhark.Util.Pretty (Pretty, align, docText, indent, ppTuple', pretty, (<+>), (</>))+import Futhark.Util.Pretty (align, docText, indent, ppTuple', pretty, (<+>), (</>))  -- | Information about an error during type checking.  The 'Show' -- instance for this type produces a human-readable description.@@ -384,9 +384,9 @@ checkConsumption (ConsumptionError e) = bad $ TypeError e checkConsumption (Consumption os) = pure os --- | Type check two mutually control flow branches.  Think @if@.  This--- interacts with consumption checking, as it is OK for an array to be--- consumed in both branches.+-- | Type check two mutually exclusive control flow branches.  Think+-- @if@.  This interacts with consumption checking, as it is OK for an+-- array to be consumed in both branches. alternative :: TypeM rep a -> TypeM rep b -> TypeM rep (a, b) alternative m1 m2 = do   (x, os1) <- collectOccurences m1@@ -394,6 +394,10 @@   tell $ Consumption $ os1 `altOccurences` os2   pure (x, y) +alternatives :: [TypeM rep ()] -> TypeM rep ()+alternatives [] = pure ()+alternatives (x : xs) = void $ x `alternative` alternatives xs+ -- | Permit consumption of only the specified names.  If one of these -- names is consumed, the consumption will be rewritten to be a -- consumption of the corresponding alias set.  Consumption of@@ -989,8 +993,9 @@ checkExp (BasicOp op) = checkBasicOp op checkExp (Match ses cases def_case info) = do   ses_ts <- mapM checkSubExp ses-  mapM_ (checkCase ses_ts) cases-  checkCaseBody def_case+  alternatives $+    context "in body of last case" (checkCaseBody def_case)+      : map (checkCase ses_ts) cases   where     checkVal t (Just v) = Prim (primValueType v) == t     checkVal _ Nothing = True@@ -1002,7 +1007,9 @@           </> "cannot match pattern"           </> indent 2 (ppTuple' $ map pretty vs)       context ("in body of case " <> prettyTuple vs) $ checkCaseBody body-    checkCaseBody = matchBranchType (matchReturns info)+    checkCaseBody body = do+      void $ checkBody body+      matchBranchType (matchReturns info) body checkExp (Apply fname args rettype_annot _) = do   (rettype_derived, paramtypes) <- lookupFun fname $ map fst args   argflows <- mapM (checkArg . fst) args

@@ -40,6 +40,7 @@     -- holes.)     DynamicFun (Exp, StaticVal) StaticVal   | IntrinsicSV+  | HoleSV SrcLoc   deriving (Show)  -- | The type is Just if this is a polymorphic binding that must be@@ -111,6 +112,8 @@         replaceStaticValSizes globals orig_substs sv2     IntrinsicSV ->       IntrinsicSV+    HoleSV loc ->+      HoleSV loc   where     tv substs =       identityMapper@@ -219,6 +222,7 @@     restrict' u (DynamicFun (e, sv1) sv2) =       DynamicFun (e, restrict' u sv1) $ restrict' u sv2     restrict' _ IntrinsicSV = IntrinsicSV+    restrict' _ (HoleSV loc) = HoleSV loc     restrict'' u (Binding t sv) = Binding t $ restrict' u sv  -- | Defunctionalization monad.  The Reader environment tracks both@@ -326,6 +330,8 @@   sizesToRename sv1 <> sizesToRename sv2 sizesToRename IntrinsicSV =   mempty+sizesToRename HoleSV {} =+  mempty sizesToRename Dynamic {} =   mempty sizesToRename (RecordSV fs) =@@ -498,11 +504,13 @@     IntrinsicSV -> do       (pats, body, tp) <- etaExpand (typeOf e) e       defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty+    HoleSV hole_loc ->+      pure (Hole (Info t) hole_loc, HoleSV hole_loc)     _ ->       let tp = typeFromSV sv        in pure (Var qn (Info tp) loc, sv) defuncExp (Hole (Info t) loc) =-  pure (Hole (Info t) loc, IntrinsicSV)+  pure (Hole (Info t) loc, HoleSV loc) defuncExp (Ascript e0 tydecl loc)   | orderZero (typeOf e0) = do       (e0', sv) <- defuncExp e0@@ -933,8 +941,18 @@       pure (apply_e, sv)     -- Propagate the 'IntrinsicsSV' until we reach the outermost application,     -- where we construct a dynamic static value with the appropriate type.-    IntrinsicSV-      | depth == 0 ->+    IntrinsicSV -> intrinsicOrHole argtypes e' sv1+    HoleSV _ -> intrinsicOrHole argtypes e' sv1+    _ ->+      error $+        "Application of an expression\n"+          ++ prettyString e1+          ++ "\nthat is neither a static lambda "+          ++ "nor a dynamic function, but has static value:\n"+          ++ show sv1+  where+    intrinsicOrHole argtypes e' sv+      | depth == 0 =           -- If the intrinsic is fully applied, then we are done.           -- Otherwise we need to eta-expand it and recursively           -- defunctionalise. XXX: might it be better to simply@@ -945,14 +963,7 @@             else do               (pats, body, tp) <- etaExpand (typeOf e') e'               defuncExp $ Lambda pats body Nothing (Info (mempty, tp)) mempty-      | otherwise -> pure (e', IntrinsicSV)-    _ ->-      error $-        "Application of an expression\n"-          ++ prettyString e1-          ++ "\nthat is neither a static lambda "-          ++ "nor a dynamic function, but has static value:\n"-          ++ show sv1+      | otherwise = pure (e', sv) defuncApply depth e@(Var qn (Info t) loc) = do   let (argtypes, _) = unfoldFunType t   sv <- lookupVar (toStruct t) (qualLeaf qn)@@ -1129,6 +1140,8 @@    in Scalar $ Sum $ M.insert name svs' $ M.fromList fields typeFromSV IntrinsicSV =   error "Tried to get the type from the static value of an intrinsic."+typeFromSV HoleSV {} =+  error "Tried to get the type from the static value of a hole."  -- | Construct the type for a fully-applied dynamic function from its -- static value and the original types of its arguments.

@@ -198,21 +198,18 @@         . 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+          let 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'+                  M.map (`lookupSubst` scopeSubsts (modScope arg_mod)) p_substs                     <> scopeSubsts f_closure                     <> scopeSubsts (modScope arg_mod)           extendScope             ( Scope                 abs_substs                 ( M.singleton (modParamName f_p) $-                    substituteInMod p_substs' arg_mod+                    substituteInMod p_substs arg_mod                 )             )             $ do

@@ -349,11 +349,31 @@        -- Some functions are magical (overloaded) and we handle that here.       case () of-        -- Overloaded functions never take array arguments (except-        -- equality, but those cannot be existential), so we can safely-        -- ignore the existential dimensions.         ()-          | Just internalise <- isOverloadedFunction qfname (map fst args) loc ->+          -- Short-circuiting operators are magical.+          | baseTag (qualLeaf qfname) <= maxIntrinsicTag,+            baseString (qualLeaf qfname) == "&&",+            [(x, _), (y, _)] <- args ->+              internaliseExp desc $+                E.AppExp+                  (E.If x y (E.Literal (E.BoolValue False) mempty) mempty)+                  (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])+          | baseTag (qualLeaf qfname) <= maxIntrinsicTag,+            baseString (qualLeaf qfname) == "||",+            [(x, _), (y, _)] <- args ->+              internaliseExp desc $+                E.AppExp+                  (E.If x (E.Literal (E.BoolValue True) mempty) y mempty)+                  (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])+          -- Overloaded and intrinsic functions never take array+          -- arguments (except equality, but those cannot be+          -- existential), so we can safely ignore the existential+          -- dimensions.+          | Just internalise <- isOverloadedFunction qfname desc loc -> do+              let prepareArg (arg, _) =+                    (E.toStruct (E.typeOf arg),) <$> internaliseExp "arg" arg+              internalise =<< mapM prepareArg args+          | Just internalise <- isIntrinsicFunction qfname (map fst args) loc ->               internalise desc           | baseTag (qualLeaf qfname) <= maxIntrinsicTag,             Just (rettype, _) <- M.lookup fname I.builtInFunctions -> do@@ -1274,6 +1294,10 @@   E.PrimType ->   E.PrimType ->   InternaliseM [I.SubExp]+internaliseBinOp _ desc E.LogAnd x y E.Bool _ =+  simpleBinOp desc I.LogAnd x y+internaliseBinOp _ desc E.LogOr x y E.Bool _ =+  simpleBinOp desc I.LogOr x y internaliseBinOp _ desc E.Plus x y (E.Signed t) _ =   simpleBinOp desc (I.Add t I.OverflowWrap) x y internaliseBinOp _ desc E.Plus x y (E.Unsigned t) _ =@@ -1459,82 +1483,30 @@       rettype       =<< bodyBind body --- | Some operators and functions are overloaded or otherwise special--- - we detect and treat them here.+-- | Overloaded operators are treated here. isOverloadedFunction ::   E.QualName VName ->-  [E.Exp] ->+  String ->   SrcLoc ->-  Maybe (String -> InternaliseM [SubExp])-isOverloadedFunction qname args loc = do+  Maybe ([(E.StructType, [SubExp])] -> InternaliseM [SubExp])+isOverloadedFunction qname desc loc = do   guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag-  let handlers =-        [ handleSign,-          handleIntrinsicOps,-          handleOps,-          handleSOACs,-          handleAccs,-          handleAD,-          handleRest-        ]-  msum [h args $ baseString $ qualLeaf qname | h <- handlers]+  handle $ baseString $ qualLeaf qname   where-    handleSign [x] "sign_i8" = Just $ toSigned I.Int8 x-    handleSign [x] "sign_i16" = Just $ toSigned I.Int16 x-    handleSign [x] "sign_i32" = Just $ toSigned I.Int32 x-    handleSign [x] "sign_i64" = Just $ toSigned I.Int64 x-    handleSign [x] "unsign_i8" = Just $ toUnsigned I.Int8 x-    handleSign [x] "unsign_i16" = Just $ toUnsigned I.Int16 x-    handleSign [x] "unsign_i32" = Just $ toUnsigned I.Int32 x-    handleSign [x] "unsign_i64" = Just $ toUnsigned I.Int64 x-    handleSign _ _ = Nothing--    handleIntrinsicOps [x] s-      | Just unop <- find ((== s) . prettyString) allUnOps = Just $ \desc -> do-          x' <- internaliseExp1 "x" x-          fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x'-    handleIntrinsicOps [TupLit [x, y] _] s-      | Just bop <- find ((== s) . prettyString) allBinOps = Just $ \desc -> do-          x' <- internaliseExp1 "x" x-          y' <- internaliseExp1 "y" y-          fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y'-      | Just cmp <- find ((== s) . prettyString) allCmpOps = Just $ \desc -> do-          x' <- internaliseExp1 "x" x-          y' <- internaliseExp1 "y" y-          fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y'-    handleIntrinsicOps [x] s-      | Just conv <- find ((== s) . prettyString) allConvOps = Just $ \desc -> do-          x' <- internaliseExp1 "x" x-          fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x'-    handleIntrinsicOps _ _ = Nothing--    -- Short-circuiting operators are magical.-    handleOps [x, y] "&&" = Just $ \desc ->-      internaliseExp desc $-        E.AppExp-          (E.If x y (E.Literal (E.BoolValue False) mempty) mempty)-          (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])-    handleOps [x, y] "||" = Just $ \desc ->-      internaliseExp desc $-        E.AppExp-          (E.If x (E.Literal (E.BoolValue True) mempty) y mempty)-          (Info $ AppRes (E.Scalar $ E.Prim E.Bool) [])     -- Handle equality and inequality specially, to treat the case of     -- arrays.-    handleOps [xe, ye] op-      | Just cmp_f <- isEqlOp op = Just $ \desc -> do-          xe' <- internaliseExp "x" xe-          ye' <- internaliseExp "y" ye-          rs <- zipWithM (doComparison desc) xe' ye'-          cmp_f desc =<< letSubExp "eq" =<< eAll rs+    handle op+      | Just cmp_f <- isEqlOp op = Just $ \[(_, xe'), (_, ye')] -> do+          rs <- zipWithM doComparison xe' ye'+          cmp_f =<< letSubExp "eq" =<< eAll rs       where-        isEqlOp "!=" = Just $ \desc eq ->+        isEqlOp "!=" = Just $ \eq ->           letTupExp' desc $ I.BasicOp $ I.UnOp I.Not eq-        isEqlOp "==" = Just $ \_ eq ->+        isEqlOp "==" = Just $ \eq ->           pure [eq]         isEqlOp _ = Nothing -        doComparison desc x y = do+        doComparison x y = do           x_t <- I.subExpType x           y_t <- I.subExpType y           case x_t of@@ -1572,15 +1544,62 @@                letSubExp "arrays_equal"                 =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False])-    handleOps [x, y] name+    handle name       | Just bop <- find ((name ==) . prettyString) [minBound .. maxBound :: E.BinOp] =-          Just $ \desc -> do-            x' <- internaliseExp1 "x" x-            y' <- internaliseExp1 "y" y-            case (E.typeOf x, E.typeOf y) of+          Just $ \[(x_t, [x']), (y_t, [y'])] ->+            case (x_t, y_t) of               (E.Scalar (E.Prim t1), E.Scalar (E.Prim t2)) ->                 internaliseBinOp loc desc bop x' y' t1 t2               _ -> error "Futhark.Internalise.internaliseExp: non-primitive type in BinOp."+    handle _ = Nothing++-- | Handle intrinsic functions.  These are only allowed to be called+-- in the prelude, and their internalisation may involve inspecting+-- the AST.+isIntrinsicFunction ::+  E.QualName VName ->+  [E.Exp] ->+  SrcLoc ->+  Maybe (String -> InternaliseM [SubExp])+isIntrinsicFunction qname args loc = do+  guard $ baseTag (qualLeaf qname) <= maxIntrinsicTag+  let handlers =+        [ handleSign,+          handleOps,+          handleSOACs,+          handleAccs,+          handleAD,+          handleRest+        ]+  msum [h args $ baseString $ qualLeaf qname | h <- handlers]+  where+    handleSign [x] "sign_i8" = Just $ toSigned I.Int8 x+    handleSign [x] "sign_i16" = Just $ toSigned I.Int16 x+    handleSign [x] "sign_i32" = Just $ toSigned I.Int32 x+    handleSign [x] "sign_i64" = Just $ toSigned I.Int64 x+    handleSign [x] "unsign_i8" = Just $ toUnsigned I.Int8 x+    handleSign [x] "unsign_i16" = Just $ toUnsigned I.Int16 x+    handleSign [x] "unsign_i32" = Just $ toUnsigned I.Int32 x+    handleSign [x] "unsign_i64" = Just $ toUnsigned I.Int64 x+    handleSign _ _ = Nothing++    handleOps [x] s+      | Just unop <- find ((== s) . prettyString) allUnOps = Just $ \desc -> do+          x' <- internaliseExp1 "x" x+          fmap pure $ letSubExp desc $ I.BasicOp $ I.UnOp unop x'+    handleOps [TupLit [x, y] _] s+      | Just bop <- find ((== s) . prettyString) allBinOps = Just $ \desc -> do+          x' <- internaliseExp1 "x" x+          y' <- internaliseExp1 "y" y+          fmap pure $ letSubExp desc $ I.BasicOp $ I.BinOp bop x' y'+      | Just cmp <- find ((== s) . prettyString) allCmpOps = Just $ \desc -> do+          x' <- internaliseExp1 "x" x+          y' <- internaliseExp1 "y" y+          fmap pure $ letSubExp desc $ I.BasicOp $ I.CmpOp cmp x' y'+    handleOps [x] s+      | Just conv <- find ((== s) . prettyString) allConvOps = Just $ \desc -> do+          x' <- internaliseExp1 "x" x+          fmap pure $ letSubExp desc $ I.BasicOp $ I.ConvOp conv x'     handleOps _ _ = Nothing      handleSOACs [TupLit [lam, arr] _] "map" = Just $ \desc -> do

@@ -211,7 +211,7 @@             is_inner_coal = isInnerCoal env inp_X load_X             str_A = baseString inp_X         x_loc <--          segScatter2D (str_A ++ "_glb2loc") loc_sz_X x_loc_init' segthd_lvl [r_par, tseq_div_tpar] (t_par, t_par) $+          segScatter2D (str_A ++ "_glb2loc") loc_sz_X x_loc_init' [r_par, tseq_div_tpar] (t_par, t_par) $             scatterFun is_inner_coal          pure (x_loc, copyLoc2Reg is_inner_coal str_A x_loc)@@ -406,7 +406,8 @@             (height_A, width_B, rem_outer_dims)             code2' -        let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt+        let grid = KernelGrid (Count grid_size) (Count group_size)+            level' = SegGroup SegNoVirt (Just grid)             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_t, gridDim_t), (gid_y, gridDim_y), (gid_x, gridDim_x)])             kbody' = KernelBody () stms_seggroup ret_seggroup         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'@@ -576,7 +577,8 @@             (height_A, width_B, rem_outer_dims)             code2' -        let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt+        let grid = KernelGrid (Count grid_size) (Count group_size)+            level' = SegGroup SegNoVirt (Just grid)             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_y, gridDim_y), (gid_x, gridDim_x)])             kbody' = KernelBody () stms_seggroup ret_seggroup         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'@@ -770,7 +772,7 @@ mkNewSegthdLvl tx ty grid_pexp = do   grid_size <- letSubExp "grid_size" =<< toExp grid_pexp   group_size <- letSubExp "group_size" =<< toExp (pe64 ty * pe64 tx)-  let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims []))+  let segthd_lvl = SegThreadInGroup (SegNoVirtFull (SegSeqDims []))   pure (grid_size, group_size, segthd_lvl)  mkGidsXYF :: Builder GPU (VName, VName, VName)@@ -1084,7 +1086,7 @@         let grid_pexp = product $ gridxyz_pexp : map (pe64 . snd) rem_outer_dims_rev         grid_size <- letSubExp "grid_size_tile3d" =<< toExp grid_pexp         group_size <- letSubExp "group_size_tile3d" =<< toExp (pe64 ty * pe64 tx)-        let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims []))+        let segthd_lvl = SegThreadInGroup (SegNoVirtFull (SegSeqDims []))          count_shmem <- letSubExp "count_shmem" =<< ceilDiv rz group_size @@ -1285,7 +1287,8 @@            pure $ map (RegTileReturns mempty regtile_ret_dims) epilogue_res'         -- END (ret_seggroup, stms_seggroup) <- runBuilder $ do-        let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt+        let grid = KernelGrid (Count grid_size) (Count group_size)+            level' = SegGroup SegNoVirt (Just grid)             space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_z, gridDim_z), (gid_y, gridDim_y), (gid_x, gridDim_x)])             kbody' = KernelBody () stms_seggroup ret_seggroup 

@@ -39,10 +39,12 @@      -- * Construction     mkDepGraph,+    mkDepGraphForFun,     pprg,   ) where +import Control.Monad.Reader import Data.Bifunctor (bimap) import Data.Foldable (foldlM) import Data.Graph.Inductive.Dot qualified as G@@ -330,6 +332,12 @@         makeAliasTable (bodyStms body),         initialGraphConstruction       ]++-- | Make a dependency graph corresponding to a function.+mkDepGraphForFun :: FunDef SOACS -> DepGraph+mkDepGraphForFun f = runReader (mkDepGraph (funDefBody f)) scope+  where+    scope = scopeOfFParams (funDefParams f) <> scopeOf (bodyStms (funDefBody f))  -- | Merges two contexts. mergedContext :: Ord b => a -> G.Context a b -> G.Context a b -> G.Context a b

@@ -89,12 +89,9 @@         Just stms_after -> pure $ Just $ stms_before <> stms_after         Nothing -> pure $ Just $ stms_before <> oneStm ker_snd -se1 :: SubExp-se1 = intConst Int64 1- genRed2Tile2d :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU)) genRed2Tile2d env kerstm@(Let pat_ker aux (Op (SegOp (SegMap seg_thd seg_space kres_tps old_kbody))))-  | (SegThread _ seg_group_size _novirt) <- seg_thd,+  | SegThread _novirt _ <- seg_thd,     -- novirt == SegNoVirtFull || novirt == SegNoVirt,     KernelBody () kstms kres <- old_kbody,     Just (css, r_ses) <- allGoodReturns kres,@@ -161,11 +158,7 @@       gid_flat_1 <- newVName "gid_flat"       let space1 = SegSpace gid_flat_1 gid_dims_new -      (grid_size, host_stms1) <- runBuilder $ do-        let grid_pexp = foldl (\x d -> x * pe64 d) (pe64 se1) $ map snd gid_dims_new-        dim_prod <- letSubExp "dim_prod" =<< toExp grid_pexp-        letSubExp "grid_size" =<< ceilDiv dim_prod (unCount seg_group_size)-      let level1 = SegThread (Count grid_size) seg_group_size (SegNoVirtFull (SegSeqDims [])) -- novirt ?+      let level1 = SegThread (SegNoVirtFull (SegSeqDims [])) Nothing -- novirt ?           kbody1 = KernelBody () ker1_stms [Returns ResultMaySimplify (Certs []) k1_res]        -- is it OK here to use the "aux" from the parrent kernel?@@ -177,7 +170,7 @@       ker2_exp <- renameExp $ Op (SegOp (SegMap seg_thd seg_space kres_tps ker2_body))       let ker2 = Let pat_ker aux ker2_exp       pure $-        Just (code1_tr_host <> host_stms1 <> oneStm ker1, ker2)+        Just (code1_tr_host <> oneStm ker1, ker2)   where     isIndVarToParDim _ (Constant _) _ = False     isIndVarToParDim variance (Var acc_ind) par_dim =@@ -198,7 +191,6 @@               (reverse acc_inds)        in invar_dims ++ inner_dims     ---    ceilDiv x y = pure $ BasicOp $ BinOp (SDivUp Int64 Unsafe) x y     getAccLambda acc_tp =       case acc_tp of         (Acc tp_id _shp el_tps _) ->

@@ -1108,10 +1108,7 @@     collectHostOp (OtherOp op) = collectFree op     collectHostOp GPUBody {} = pure () -    collectSegLevel (SegThread (Count num) (Count size) _) =-      collectSubExp num >> collectSubExp size-    collectSegLevel (SegGroup (Count num) (Count size) _) =-      collectSubExp num >> collectSubExp size+    collectSegLevel = mapM_ captureName . namesToList . freeIn      collectSegSpace space =       mapM_ collectSubExp (segSpaceDims space)

@@ -452,7 +452,7 @@   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+  ((e', e_hoisted), e_cs) <- collectCerts $ simplifyExp (usage <> UT.usageInPat pat) pat e   let aux' = StmAux (cs <> e_cs) attrs dec   pure (e_hoisted, mkWiseStm (removePatWisdom pat) aux' e') @@ -666,8 +666,8 @@   all (`nameIn` ST.availableAtClosestLoop vtable) . namesToList . freeIn  matchBlocker ::-  (ASTRep rep, CanBeWise (Op rep), FreeIn a) =>-  a ->+  (ASTRep rep, CanBeWise (Op rep)) =>+  [SubExp] ->   MatchDec rt ->   SimpleM rep (BlockPred (Wise rep)) matchBlocker cond (MatchDec _ ifsort) = do@@ -896,7 +896,7 @@       Nothing ->         pure (Nothing, mempty)       Just (op_lam, nes) -> do-        (op_lam', op_lam_stms) <- blockMigrated (simplifyLambda op_lam)+        (op_lam', op_lam_stms) <- blockMigrated (simplifyLambda mempty op_lam)         nes' <- simplify nes         pure (Just (op_lam', nes'), op_lam_stms)     (,op_stms) <$> ((,,op') <$> simplify shape <*> simplify arrs)@@ -1063,11 +1063,12 @@  simplifyLambda ::   SimplifiableRep rep =>+  Names ->   Lambda (Wise rep) ->   SimpleM rep (Lambda (Wise rep), Stms (Wise rep))-simplifyLambda lam = do+simplifyLambda extra_bound lam = do   par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers-  simplifyLambdaMaybeHoist par_blocker mempty lam+  simplifyLambdaMaybeHoist (par_blocker `orIf` hasFree extra_bound) mempty lam  simplifyLambdaNoHoisting ::   SimplifiableRep rep =>

@@ -238,34 +238,6 @@     notIndex _ = True simplifyLoopVariables _ _ _ _ = Skip --- If a for-loop with no loop variables has a counter of type Int64,--- and the bound is just a constant or sign-extended integer of--- smaller type, then change the loop to iterate over the smaller type--- instead.  We then move the sign extension inside the loop instead.--- This addresses loops of the form @for i in x..<y@ in the source--- language.-narrowLoopType :: (BuilderOps rep) => TopDownRuleDoLoop rep-narrowLoopType vtable pat aux (merge, ForLoop i Int64 n [], body)-  | Just (n', it', cs) <- smallerType =-      Simplify $ do-        i' <- newVName $ baseString i-        let form' = ForLoop i' it' n' []-        body' <- insertStmsM . inScopeOf form' $ do-          letBindNames [i] $ BasicOp $ ConvOp (SExt it' Int64) (Var i')-          pure body-        auxing aux $ certifying cs $ letBind pat $ DoLoop merge form' body'-  where-    smallerType-      | Var n' <- n,-        Just (ConvOp (SExt it' _) n'', cs) <- ST.lookupBasicOp n' vtable =-          Just (n'', it', cs)-      | Constant (IntValue (Int64Value n')) <- n,-        toInteger n' <= toInteger (maxBound :: Int32) =-          Just (intConst Int32 (toInteger n'), Int32, mempty)-      | otherwise =-          Nothing-narrowLoopType _ _ _ _ = Skip- unroll ::   BuilderOps rep =>   Integer ->@@ -313,8 +285,7 @@   [ RuleDoLoop hoistLoopInvariantMergeVariables,     RuleDoLoop simplifyClosedFormLoop,     RuleDoLoop simplifyKnownIterationLoop,-    RuleDoLoop simplifyLoopVariables,-    RuleDoLoop narrowLoopType+    RuleDoLoop simplifyLoopVariables   ]  bottomUpRules :: BuilderOps rep => [BottomUpRule rep]

@@ -126,13 +126,12 @@         not $ null $ tiledInputs inputs,         gtid_y : gtid_x : top_gtids_rev <- reverse gtids,         kdim_y : kdim_x : top_kdims_rev <- reverse kdims,-        (prestms', poststms') <-+        Just (prestms', poststms') <-           preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms),         used <- freeIn stm_to_tile <> freeIn poststms' <> freeIn stms_res =           Just . injectPrelude initial_space variance prestms' used             <$> tileGeneric               (tiling2d $ reverse $ zip top_gtids_rev top_kdims_rev)-              initial_lvl               res_ts               (stmPat stm_to_tile)               (gtid_x, gtid_y)@@ -148,13 +147,12 @@         inputs <- map (is1DTileable gtid variance) arrs,         not $ null $ tiledInputs inputs,         gtid `notNameIn` branch_variant,-        (prestms', poststms') <-+        Just (prestms', poststms') <-           preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms),         used <- freeIn stm_to_tile <> freeIn poststms' <> freeIn stms_res =           Just . injectPrelude initial_space variance prestms' used             <$> tileGeneric               (tiling1d $ reverse top_space_rev)-              initial_lvl               res_ts               (stmPat stm_to_tile)               gtid@@ -167,7 +165,7 @@       -- Tiling inside for-loop.       | DoLoop merge (ForLoop i it bound []) loopbody <- stmExp stm_to_tile,         not $ any ((`nameIn` freeIn merge) . paramName . fst) merge,-        (prestms', poststms') <-+        Just (prestms', poststms') <-           preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms) = do           let branch_variant' =                 branch_variant@@ -214,15 +212,22 @@             descend (prestms <> oneStm stm_to_tile) poststms  -- | Move statements from prelude to postlude if they are not used in--- the tiled statement anyway.+-- the tiled statement anyway.  Also, fail if the provided Stm uses+-- anything from the resulting prelude whose size is not free in the+-- prelude. preludeToPostlude ::   VarianceTable ->   Stms GPU ->   Stm GPU ->   Stms GPU ->-  (Stms GPU, Stms GPU)-preludeToPostlude variance prelude stm_to_tile postlude =-  (prelude_used, prelude_not_used <> postlude)+  Maybe (Stms GPU, Stms GPU)+preludeToPostlude variance prelude stm_to_tile postlude = do+  let prelude_sizes =+        freeIn $ foldMap (patTypes . stmPat) prelude_used+      prelude_bound =+        namesFromList $ foldMap (patNames . stmPat) prelude_used+  guard $ not $ prelude_sizes `namesIntersect` prelude_bound+  Just (prelude_used, prelude_not_used <> postlude)   where     used_in_tiled = freeIn stm_to_tile @@ -404,7 +409,7 @@         mergeinit' <-           fmap (map Var) $             certifying (stmAuxCerts aux) $-              tilingSegMap tiling "tiled_loopinit" (scalarLevel tiling) ResultPrivate $+              tilingSegMap tiling "tiled_loopinit" ResultPrivate $                 \in_bounds slice ->                   fmap varsRes $                     protectOutOfBounds "loopinit" in_bounds merge_ts $ do@@ -428,8 +433,7 @@          loopbody' <-           localScope (scopeOfFParams mergeparams') . runBodyBuilder $-            resultBody . map Var-              <$> tiledBody private' privstms'+            resultBody . map Var <$> tiledBody private' privstms'         accs' <-           letTupExp "tiled_inside_loop" $             DoLoop merge' (ForLoop i it bound []) loopbody'@@ -447,13 +451,12 @@ doPrelude :: Tiling -> PrivStms -> Stms GPU -> [VName] -> Builder GPU [VName] doPrelude tiling privstms prestms prestms_live =   -- Create a SegMap that takes care of the prelude for every thread.-  tilingSegMap tiling "prelude" (scalarLevel tiling) ResultPrivate $-    \in_bounds slice -> do-      ts <- mapM lookupType prestms_live-      fmap varsRes . protectOutOfBounds "pre" in_bounds ts $ do-        addPrivStms slice privstms-        addStms prestms-        pure $ varsRes prestms_live+  tilingSegMap tiling "prelude" 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  liveSet :: FreeIn a => Stms GPU -> a -> Names liveSet stms after =@@ -578,7 +581,6 @@ data Tiling = Tiling   { tilingSegMap ::       String ->-      SegLevel ->       ResultManifest ->       (PrimExp VName -> [DimIndex SubExp] -> Builder GPU Result) ->       Builder GPU [VName],@@ -604,13 +606,7 @@   }  type DoTiling gtids kdims =-  SegLevel -> gtids -> kdims -> SubExp -> Builder GPU Tiling--scalarLevel :: Tiling -> SegLevel-scalarLevel tiling =-  SegThread (segNumGroups lvl) (segGroupSize lvl) SegNoVirt-  where-    lvl = tilingLevel tiling+  gtids -> kdims -> SubExp -> Builder GPU Tiling  protectOutOfBounds ::   String ->@@ -643,7 +639,7 @@   [Type] ->   Builder GPU [VName] postludeGeneric tiling privstms pat accs' poststms poststms_res res_ts =-  tilingSegMap tiling "thread_res" (scalarLevel tiling) ResultPrivate $ \in_bounds slice -> do+  tilingSegMap tiling "thread_res" ResultPrivate $ \in_bounds slice -> do     -- Read our per-thread result from the tiled loop.     forM_ (zip (patNames pat) accs') $ \(us, everyone) -> do       everyone_t <- lookupType everyone@@ -664,7 +660,6 @@  tileGeneric ::   DoTiling gtids kdims ->-  SegLevel ->   [Type] ->   Pat Type ->   gtids ->@@ -675,8 +670,8 @@   Stms GPU ->   Result ->   TileM (Stms GPU, Tiling, TiledBody)-tileGeneric doTiling initial_lvl res_ts pat gtids kdims w form inputs poststms poststms_res = do-  (tiling, tiling_stms) <- runBuilder $ doTiling initial_lvl gtids kdims w+tileGeneric doTiling res_ts pat gtids kdims w form inputs poststms poststms_res = do+  (tiling, tiling_stms) <- runBuilder $ doTiling gtids kdims w    pure (tiling_stms, tiling, tiledBody tiling)   where@@ -690,7 +685,7 @@        -- We don't use a Replicate here, because we want to enforce a       -- scalar memory space.-      mergeinits <- tilingSegMap tiling "mergeinit" (scalarLevel tiling) ResultPrivate $ \in_bounds slice ->+      mergeinits <- tilingSegMap tiling "mergeinit" ResultPrivate $ \in_bounds slice ->         -- Constant neutral elements (a common case) do not need protection from OOB.         if freeIn red_nes == mempty           then pure $ subExpsRes red_nes@@ -736,10 +731,9 @@  mkReadPreludeValues :: [VName] -> [VName] -> ReadPrelude mkReadPreludeValues prestms_live_arrs prestms_live slice =-  fmap mconcat $-    forM (zip prestms_live_arrs prestms_live) $ \(arr, v) -> do-      arr_t <- lookupType arr-      letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice+  fmap mconcat . forM (zip prestms_live_arrs prestms_live) $ \(arr, v) -> do+    arr_t <- lookupType arr+    letBindNames [v] $ BasicOp $ Index arr $ fullSlice arr_t slice  tileReturns :: [(VName, SubExp)] -> [(SubExp, SubExp)] -> VName -> Builder GPU KernelResult tileReturns dims_on_top dims arr = do@@ -776,16 +770,15 @@   SubExp ->   VName ->   VName ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  KernelGrid ->   TileKind ->   PrivStms ->   SubExp ->   [InputArray] ->   Builder GPU [InputTile]-readTile1D tile_size gid gtid num_groups group_size kind privstms tile_id inputs =+readTile1D tile_size gid gtid (KernelGrid _num_groups group_size) kind privstms tile_id inputs =   fmap (inputsToTiles inputs)-    . segMap1D "full_tile" lvl ResultNoSimplify+    . segMap1D "full_tile" lvl ResultNoSimplify tile_size     $ \ltid -> do       j <-         letSubExp "j"@@ -813,18 +806,17 @@           TileFull ->             mapM readTileElem arrs   where-    lvl = SegThread num_groups group_size SegNoVirt+    lvl = SegThreadInGroup SegNoVirt  processTile1D ::   VName ->   VName ->   SubExp ->   SubExp ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  KernelGrid ->   ProcessTileArgs ->   Builder GPU [VName]-processTile1D gid gtid kdim tile_size num_groups group_size tile_args = do+processTile1D gid gtid kdim tile_size (KernelGrid _num_groups group_size) tile_args = do   let red_comm = processComm tile_args       privstms = processPrivStms tile_args       map_lam = processMapLam tile_args@@ -833,7 +825,7 @@       tile_id = processTileId tile_args       accs = processAcc tile_args -  segMap1D "acc" lvl ResultPrivate $ \ltid -> do+  segMap1D "acc" lvl ResultPrivate (unCount group_size) $ \ltid -> do     reconstructGtids1D group_size gtid gid ltid     addPrivStms [DimFix $ Var ltid] privstms @@ -857,18 +849,17 @@           (eBody [pure $ Op $ OtherOp $ Screma tile_size tiles' form'])           (resultBodyM thread_accs)   where-    lvl = SegThread num_groups group_size SegNoVirt+    lvl = SegThreadInGroup SegNoVirt  processResidualTile1D ::   VName ->   VName ->   SubExp ->   SubExp ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->+  KernelGrid ->   ResidualTileArgs ->   Builder GPU [VName]-processResidualTile1D gid gtid kdim tile_size num_groups group_size args = do+processResidualTile1D gid gtid kdim tile_size grid args = do   -- The number of residual elements that are not covered by   -- the whole tiles.   residual_input <-@@ -899,8 +890,7 @@           tile_size           gid           gtid-          num_groups-          group_size+          grid           TilePartial           privstms           num_whole_tiles@@ -921,61 +911,52 @@       let tile_args =             ProcessTileArgs privstms red_comm red_lam map_lam tiles accs num_whole_tiles       resultBody . map Var-        <$> processTile1D gid gtid kdim residual_input num_groups group_size tile_args+        <$> processTile1D gid gtid kdim residual_input grid tile_args  tiling1d :: [(VName, SubExp)] -> DoTiling VName SubExp-tiling1d dims_on_top initial_lvl gtid kdim w = do+tiling1d dims_on_top gtid kdim w = do   gid <- newVName "gid"   gid_flat <- newVName "gid_flat" -  (lvl, space) <--    if null dims_on_top-      then-        pure-          ( SegGroup (segNumGroups initial_lvl) (segGroupSize initial_lvl) $ segVirt initial_lvl,-            SegSpace gid_flat [(gid, unCount $ segNumGroups initial_lvl)]-          )-      else do-        group_size <--          letSubExp "computed_group_size" $-            BasicOp $-              BinOp (SMin Int64) (unCount (segGroupSize initial_lvl)) kdim+  tile_size_key <- nameFromString . prettyString <$> newVName "tile_size"+  tile_size <- letSubExp "tile_size" $ Op $ SizeOp $ GetSize tile_size_key SizeGroup+  let group_size = tile_size -        -- How many groups we need to exhaust the innermost dimension.-        ldim <--          letSubExp "ldim" $-            BasicOp $-              BinOp (SDivUp Int64 Unsafe) kdim group_size+  (grid, space) <- do+    -- How many groups we need to exhaust the innermost dimension.+    ldim <-+      letSubExp "ldim" . BasicOp $+        BinOp (SDivUp Int64 Unsafe) kdim group_size -        num_groups <--          letSubExp "computed_num_groups"-            =<< foldBinOp (Mul Int64 OverflowUndef) ldim (map snd dims_on_top)+    num_groups <-+      letSubExp "computed_num_groups"+        =<< foldBinOp (Mul Int64 OverflowUndef) ldim (map snd dims_on_top) -        pure-          ( SegGroup (Count num_groups) (Count group_size) SegNoVirt,-            SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)]-          )-  let tile_size = unCount $ segGroupSize lvl+    pure+      ( KernelGrid (Count num_groups) (Count group_size),+        SegSpace gid_flat $ dims_on_top ++ [(gid, ldim)]+      )+  let tiling_lvl = SegThreadInGroup SegNoVirt    pure     Tiling-      { tilingSegMap = \desc lvl' manifest f -> segMap1D desc lvl' manifest $ \ltid -> do+      { tilingSegMap = \desc manifest f -> segMap1D desc tiling_lvl manifest tile_size $ \ltid -> do           letBindNames [gtid]             =<< toExp (le64 gid * pe64 tile_size + le64 ltid)           f (untyped $ le64 gtid .<. pe64 kdim) [DimFix $ Var ltid],         tilingReadTile =-          readTile1D tile_size gid gtid (segNumGroups lvl) (segGroupSize lvl),+          readTile1D tile_size gid gtid grid,         tilingProcessTile =-          processTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl),+          processTile1D gid gtid kdim tile_size grid,         tilingProcessResidualTile =-          processResidualTile1D gid gtid kdim tile_size (segNumGroups lvl) (segGroupSize lvl),+          processResidualTile1D gid gtid kdim tile_size grid,         tilingTileReturns = tileReturns dims_on_top [(kdim, tile_size)],         tilingTileShape = Shape [tile_size],         tilingNumWholeTiles =           letSubExp "num_whole_tiles" $             BasicOp $               BinOp (SQuot Int64 Unsafe) w tile_size,-        tilingLevel = lvl,+        tilingLevel = SegGroup SegNoVirt (Just grid),         tilingSpace = space       } @@ -1015,18 +996,16 @@   (VName, VName) ->   (VName, VName) ->   SubExp ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->   TileKind ->   PrivStms ->   SubExp ->   [InputArray] ->   Builder GPU [InputTile]-readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size num_groups group_size kind privstms tile_id inputs =+readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size kind privstms tile_id inputs =   fmap (inputsToTiles inputs)     . segMap2D       "full_tile"-      (SegThread num_groups group_size (SegNoVirtFull (SegSeqDims [])))+      (SegThread (SegNoVirtFull (SegSeqDims [])) Nothing)       ResultNoSimplify       (tile_size, tile_size)     $ \(ltid_x, ltid_y) -> do@@ -1088,11 +1067,9 @@   (VName, VName) ->   (SubExp, SubExp) ->   SubExp ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->   ProcessTileArgs ->   Builder GPU [VName]-processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size num_groups group_size tile_args = do+processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size tile_args = do   let privstms = processPrivStms tile_args       red_comm = processComm tile_args       red_lam = processRedLam tile_args@@ -1106,7 +1083,7 @@    segMap2D     "acc"-    (SegThread num_groups group_size (SegNoVirtFull (SegSeqDims [])))+    (SegThreadInGroup (SegNoVirtFull (SegSeqDims [])))     ResultPrivate     (tile_size, tile_size)     $ \(ltid_x, ltid_y) -> do@@ -1146,82 +1123,69 @@   (VName, VName) ->   (SubExp, SubExp) ->   SubExp ->-  Count NumGroups SubExp ->-  Count GroupSize SubExp ->   ResidualTileArgs ->   Builder GPU [VName]-processResidualTile2D-  gids-  gtids-  kdims-  tile_size-  num_groups-  group_size-  args = do-    -- The number of residual elements that are not covered by-    -- the whole tiles.-    residual_input <--      letSubExp "residual_input" $-        BasicOp $-          BinOp (SRem Int64 Unsafe) w tile_size+processResidualTile2D gids gtids kdims tile_size args = do+  -- The number of residual elements that are not covered by+  -- the whole tiles.+  residual_input <-+    letSubExp "residual_input" $+      BasicOp $+        BinOp (SRem Int64 Unsafe) w tile_size -    letTupExp "acc_after_residual"-      =<< eIf-        (toExp $ pe64 residual_input .==. 0)-        (resultBodyM $ map Var accs)-        (nonemptyTile residual_input)-    where-      privstms = residualPrivStms args-      red_comm = residualComm args-      red_lam = residualRedLam args-      map_lam = residualMapLam args-      accs = residualAcc args-      inputs = residualInput args-      num_whole_tiles = residualNumWholeTiles args-      w = residualInputSize args+  letTupExp "acc_after_residual"+    =<< eIf+      (toExp $ pe64 residual_input .==. 0)+      (resultBodyM $ map Var accs)+      (nonemptyTile residual_input)+  where+    privstms = residualPrivStms args+    red_comm = residualComm args+    red_lam = residualRedLam args+    map_lam = residualMapLam args+    accs = residualAcc args+    inputs = residualInput args+    num_whole_tiles = residualNumWholeTiles args+    w = residualInputSize args -      nonemptyTile residual_input = renameBody <=< runBodyBuilder $ do-        -- Collectively construct a tile.  Threads that are out-of-bounds-        -- provide a blank dummy value.-        full_tile <--          readTile2D-            kdims-            gtids-            gids-            tile_size-            num_groups-            group_size-            TilePartial-            privstms-            num_whole_tiles-            inputs+    nonemptyTile residual_input = renameBody <=< runBodyBuilder $ do+      -- Collectively construct a tile.  Threads that are out-of-bounds+      -- provide a blank dummy value.+      full_tile <-+        readTile2D+          kdims+          gtids+          gids+          tile_size+          TilePartial+          privstms+          num_whole_tiles+          inputs -        let slice =-              DimSlice (intConst Int64 0) residual_input (intConst Int64 1)-        tiles <- forM full_tile $ \case-          InputTiled perm tile' ->-            InputTiled perm-              <$> letExp "partial_tile" (BasicOp $ Index tile' (Slice [slice, slice]))-          InputUntiled arr ->-            pure $ InputUntiled arr+      let slice =+            DimSlice (intConst Int64 0) residual_input (intConst Int64 1)+      tiles <- forM full_tile $ \case+        InputTiled perm tile' ->+          InputTiled perm+            <$> letExp "partial_tile" (BasicOp $ Index tile' (Slice [slice, slice]))+        InputUntiled arr ->+          pure $ InputUntiled arr -        let tile_args =-              ProcessTileArgs privstms red_comm red_lam map_lam tiles accs num_whole_tiles+      let tile_args =+            ProcessTileArgs privstms red_comm red_lam map_lam tiles accs num_whole_tiles -        -- Now each thread performs a traversal of the tile and-        -- updates its accumulator.-        resultBody . map Var-          <$> processTile2D-            gids-            gtids-            kdims-            tile_size-            num_groups-            group_size-            tile_args+      -- Now each thread performs a traversal of the tile and+      -- updates its accumulator.+      resultBody . map Var+        <$> processTile2D+          gids+          gtids+          kdims+          tile_size+          tile_args  tiling2d :: [(VName, SubExp)] -> DoTiling (VName, VName) (SubExp, SubExp)-tiling2d dims_on_top _initial_lvl (gtid_x, gtid_y) (kdim_x, kdim_y) w = do+tiling2d dims_on_top (gtid_x, gtid_y) (kdim_x, kdim_y) w = do   gid_x <- newVName "gid_x"   gid_y <- newVName "gid_y" @@ -1246,24 +1210,26 @@         (num_groups_y : map snd dims_on_top)    gid_flat <- newVName "gid_flat"-  let lvl = SegGroup (Count num_groups) (Count group_size) (SegNoVirtFull (SegSeqDims []))+  let grid = KernelGrid (Count num_groups) (Count group_size)+      lvl = SegGroup (SegNoVirtFull (SegSeqDims [])) (Just grid)       space =         SegSpace gid_flat $           dims_on_top ++ [(gid_x, num_groups_x), (gid_y, num_groups_y)]+      tiling_lvl = SegThreadInGroup SegNoVirt    pure     Tiling-      { tilingSegMap = \desc lvl' manifest f ->-          segMap2D desc lvl' manifest (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do+      { tilingSegMap = \desc manifest f ->+          segMap2D desc tiling_lvl manifest (tile_size, tile_size) $ \(ltid_x, ltid_y) -> do             reconstructGtids2D tile_size (gtid_x, gtid_y) (gid_x, gid_y) (ltid_x, ltid_y)             f               ( untyped $                   le64 gtid_x .<. pe64 kdim_x .&&. le64 gtid_y .<. pe64 kdim_y               )               [DimFix $ Var ltid_x, DimFix $ Var ltid_y],-        tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size (segNumGroups lvl) (segGroupSize lvl),-        tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl),-        tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size (segNumGroups lvl) (segGroupSize lvl),+        tilingReadTile = readTile2D (kdim_x, kdim_y) (gtid_x, gtid_y) (gid_x, gid_y) tile_size,+        tilingProcessTile = processTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size,+        tilingProcessResidualTile = processResidualTile2D (gid_x, gid_y) (gtid_x, gtid_y) (kdim_x, kdim_y) tile_size,         tilingTileReturns = tileReturns dims_on_top [(kdim_x, tile_size), (kdim_y, tile_size)],         tilingTileShape = Shape [tile_size, tile_size],         tilingNumWholeTiles =

@@ -85,12 +85,13 @@   String ->   SegLevel ->   ResultManifest ->+  SubExp -> -- dim_x   (VName -> Builder GPU Result) ->   Builder GPU [VName]-segMap1D desc lvl manifest f = do+segMap1D desc lvl manifest w f = do   ltid <- newVName "ltid"   ltid_flat <- newVName "ltid_flat"-  let space = SegSpace ltid_flat [(ltid, unCount $ segGroupSize lvl)]+  let space = SegSpace ltid_flat [(ltid, w)]    ((ts, res), stms) <- localScope (scopeOfSegSpace space) . runBuilder $ do     res <- f ltid@@ -164,12 +165,11 @@   String ->   SubExp ->   VName ->-  SegLevel -> -- lvl   [SubExp] -> -- dims of sequential loop on top   (SubExp, SubExp) -> -- (dim_y, dim_x)   ([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp)) -> -- f   Builder GPU VName-segScatter2D desc arr_size updt_arr lvl seq_dims (dim_x, dim_y) f = do+segScatter2D desc arr_size updt_arr seq_dims (dim_x, dim_y) f = do   ltid_flat <- newVName "ltid_flat"   ltid_y <- newVName "ltid_y"   ltid_x <- newVName "ltid_x"@@ -178,10 +178,8 @@   let seq_space = zip seq_is seq_dims    let segspace = SegSpace ltid_flat $ seq_space ++ [(ltid_y, dim_y), (ltid_x, dim_x)]-      lvl' =-        SegThread-          (segNumGroups lvl)-          (segGroupSize lvl)+      lvl =+        SegThreadInGroup           (SegNoVirtFull (SegSeqDims [0 .. length seq_dims - 1]))    ((t_v, res_v, res_i), stms) <- runBuilder $ do@@ -194,7 +192,7 @@   let ret = WriteReturns mempty (Shape [arr_size]) updt_arr [(Slice [DimFix res_i], res_v)]   let body = KernelBody () stms [ret] -  letExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl' segspace [t_v] body+  letExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl segspace [t_v] body  -- | The variance table keeps a mapping from a variable name -- (something produced by a 'Stm') to the kernel thread indices

@@ -110,33 +110,33 @@  transformExp :: Exp GPUMem -> ExpandM (Stms GPUMem, Exp GPUMem) transformExp (Op (Inner (SegOp (SegMap lvl space ts kbody)))) = do-  (alloc_stms, (_, kbody')) <- transformScanRed lvl space [] kbody+  (alloc_stms, (lvl', _, kbody')) <- transformScanRed lvl space [] kbody   pure     ( alloc_stms,-      Op $ Inner $ SegOp $ SegMap lvl space ts kbody'+      Op $ Inner $ SegOp $ SegMap lvl' space ts kbody'     ) transformExp (Op (Inner (SegOp (SegRed lvl space reds ts kbody)))) = do-  (alloc_stms, (lams, kbody')) <-+  (alloc_stms, (lvl', lams, kbody')) <-     transformScanRed lvl space (map segBinOpLambda reds) kbody   let reds' = zipWith (\red lam -> red {segBinOpLambda = lam}) reds lams   pure     ( alloc_stms,-      Op $ Inner $ SegOp $ SegRed lvl space reds' ts kbody'+      Op $ Inner $ SegOp $ SegRed lvl' space reds' ts kbody'     ) transformExp (Op (Inner (SegOp (SegScan lvl space scans ts kbody)))) = do-  (alloc_stms, (lams, kbody')) <-+  (alloc_stms, (lvl', lams, kbody')) <-     transformScanRed lvl space (map segBinOpLambda scans) kbody   let scans' = zipWith (\red lam -> red {segBinOpLambda = lam}) scans lams   pure     ( alloc_stms,-      Op $ Inner $ SegOp $ SegScan lvl space scans' ts kbody'+      Op $ Inner $ SegOp $ SegScan lvl' space scans' ts kbody'     ) transformExp (Op (Inner (SegOp (SegHist lvl space ops ts kbody)))) = do-  (alloc_stms, (lams', kbody')) <- transformScanRed lvl space lams kbody+  (alloc_stms, (lvl', lams', kbody')) <- transformScanRed lvl space lams kbody   let ops' = zipWith onOp ops lams'   pure     ( alloc_stms,-      Op $ Inner $ SegOp $ SegHist lvl space ops' ts kbody'+      Op $ Inner $ SegOp $ SegHist lvl' space ops' ts kbody'     )   where     lams = map histOp ops@@ -156,7 +156,7 @@       let -- XXX: fake a SegLevel, which we don't have here.  We will not           -- use it for anything, as we will not allow irregular           -- allocations inside the update function.-          lvl = SegThread (Count $ intConst Int64 0) (Count $ intConst Int64 0) SegNoVirt+          lvl = SegThread SegNoVirt Nothing           (op_lam', lam_allocs) =             extractLambdaAllocations (lvl, [0]) bound_outside mempty op_lam           variantAlloc (_, Var v, _) = v `notNameIn` bound_outside@@ -186,12 +186,33 @@ transformExp e =   pure (mempty, e) +ensureGridKnown :: SegLevel -> ExpandM (Stms GPUMem, SegLevel, KernelGrid)+ensureGridKnown lvl =+  case lvl of+    SegThread _ (Just grid) -> pure (mempty, lvl, grid)+    SegGroup _ (Just grid) -> pure (mempty, lvl, grid)+    SegThread virt Nothing -> mkGrid (SegThread virt)+    SegGroup virt Nothing -> mkGrid (SegGroup virt)+    SegThreadInGroup {} -> error "ensureGridKnown: SegThreadInGroup"+  where+    mkGrid f = do+      (grid, stms) <-+        runBuilder $+          KernelGrid+            <$> (Count <$> getSize "num_groups" SizeNumGroups)+            <*> (Count <$> getSize "group_size" SizeGroup)+      pure (stms, f $ Just grid, grid)++    getSize desc size_class = do+      size_key <- nameFromString . prettyString <$> newVName desc+      letSubExp desc $ Op $ Inner $ SizeOp $ GetSize size_key size_class+ transformScanRed ::   SegLevel ->   SegSpace ->   [Lambda GPUMem] ->   KernelBody GPUMem ->-  ExpandM (Stms GPUMem, ([Lambda GPUMem], KernelBody GPUMem))+  ExpandM (Stms GPUMem, (SegLevel, [Lambda GPUMem], KernelBody GPUMem)) transformScanRed lvl space ops kbody = do   bound_outside <- asks $ namesFromList . M.keys   let user = (lvl, [le64 $ segFlat space])@@ -221,10 +242,14 @@     _ ->       pure () -  allocsForBody variant_allocs invariant_allocs lvl space kbody' $ \alloc_stms kbody'' -> do-    ops'' <- forM ops' $ \op' ->-      localScope (scopeOf op') $ offsetMemoryInLambda op'-    pure (alloc_stms, (ops'', kbody''))+  if null variant_allocs && null invariant_allocs+    then pure (mempty, (lvl, ops, kbody))+    else do+      (lvl_stms, lvl', grid) <- ensureGridKnown lvl+      allocsForBody variant_allocs invariant_allocs grid space kbody' $ \alloc_stms kbody'' -> do+        ops'' <- forM ops' $ \op' ->+          localScope (scopeOf op') $ offsetMemoryInLambda op'+        pure (lvl_stms <> alloc_stms, (lvl', ops'', kbody''))   where     bound_in_kernel =       namesFromList (M.keys $ scopeOfSegSpace space)@@ -236,15 +261,15 @@ allocsForBody ::   Extraction ->   Extraction ->-  SegLevel ->+  KernelGrid ->   SegSpace ->   KernelBody GPUMem ->   (Stms GPUMem -> KernelBody GPUMem -> OffsetM b) ->   ExpandM b-allocsForBody variant_allocs invariant_allocs lvl space kbody' m = do+allocsForBody variant_allocs invariant_allocs grid space kbody' m = do   (alloc_offsets, alloc_stms) <-     memoryRequirements-      lvl+      grid       space       (kernelBodyStms kbody')       variant_allocs@@ -258,26 +283,26 @@       m alloc_stms kbody''  memoryRequirements ::-  SegLevel ->+  KernelGrid ->   SegSpace ->   Stms GPUMem ->   Extraction ->   Extraction ->   ExpandM (RebaseMap, Stms GPUMem)-memoryRequirements lvl space kstms variant_allocs invariant_allocs = do+memoryRequirements grid space kstms variant_allocs invariant_allocs = do   (num_threads, num_threads_stms) <-     runBuilder . letSubExp "num_threads" . BasicOp $       BinOp         (Mul Int64 OverflowUndef)-        (unCount $ segNumGroups lvl)-        (unCount $ segGroupSize lvl)+        (unCount $ gridNumGroups grid)+        (unCount $ gridGroupSize grid)    (invariant_alloc_stms, invariant_alloc_offsets) <-     inScopeOf num_threads_stms $       expandedInvariantAllocations         num_threads-        (segNumGroups lvl)-        (segGroupSize lvl)+        (gridNumGroups grid)+        (gridGroupSize grid)         invariant_allocs    (variant_alloc_stms, variant_alloc_offsets) <-@@ -437,7 +462,7 @@      untouched d = DimSlice 0 d 1 -    newBase user@(SegThread {}, _) (old_shape, _) =+    newBaseThread user (old_shape, _) =       let (users_shape, user_ids) = getNumUsers user           num_dims = length old_shape           perm = [num_dims .. num_dims + shapeRank users_shape - 1] ++ [0 .. num_dims - 1]@@ -448,7 +473,10 @@               Slice $                 map DimFix user_ids ++ map untouched old_shape        in offset_ixfun-    newBase user@(SegGroup {}, _) (old_shape, _) =++    newBase user@(SegThreadInGroup {}, _) = newBaseThread user+    newBase user@(SegThread {}, _) = newBaseThread user+    newBase user@(SegGroup {}, _) = \(old_shape, _) ->       let (users_shape, user_ids) = getNumUsers user           root_ixfun = IxFun.iota $ map pe64 (shapeDims users_shape) ++ old_shape           offset_ixfun =@@ -467,6 +495,8 @@   where     getNumUsers (SegThread {}, [gtid]) = (Shape [num_threads], [gtid])     getNumUsers (SegThread {}, [gid, ltid]) = (Shape [num_groups, group_size], [gid, ltid])+    getNumUsers (SegThreadInGroup {}, [gtid]) = (Shape [num_threads], [gtid])+    getNumUsers (SegThreadInGroup {}, [gid, ltid]) = (Shape [num_groups, group_size], [gid, ltid])     getNumUsers (SegGroup {}, [gid]) = (Shape [num_groups], [gid])     getNumUsers user = error $ "getNumUsers: unhandled " ++ show user 

@@ -20,32 +20,48 @@   opIsConst (SizeOp GetSizeMax {}) = True   opIsConst _ = False -allocAtLevel :: SegLevel -> AllocM fromrep trep a -> AllocM fromrep trep a+allocAtLevel :: SegLevel -> AllocM GPU GPUMem a -> AllocM GPU GPUMem a allocAtLevel lvl = local $ \env ->   env     { allocSpace = space,-      aggressiveReuse = True+      aggressiveReuse = True,+      allocInOp = handleHostOp (Just lvl)     }   where     space = case lvl of-      SegThread {} -> DefaultSpace       SegGroup {} -> Space "local"+      SegThread {} -> DefaultSpace+      SegThreadInGroup {} -> DefaultSpace  handleSegOp ::+  Maybe SegLevel ->   SegOp SegLevel GPU ->   AllocM GPU GPUMem (SegOp SegLevel GPUMem)-handleSegOp op = do+handleSegOp outer_lvl op = do   num_threads <--    letSubExp "num_threads" $-      BasicOp $-        BinOp-          (Mul Int64 OverflowUndef)-          (unCount (segNumGroups lvl))-          (unCount (segGroupSize lvl))-  allocAtLevel lvl $ mapSegOpM (mapper num_threads) op+    letSubExp "num_threads"+      =<< case maybe_grid of+        Just grid ->+          pure . BasicOp $+            BinOp+              (Mul Int64 OverflowUndef)+              (unCount (gridNumGroups grid))+              (unCount (gridGroupSize grid))+        Nothing ->+          foldBinOp+            (Mul Int64 OverflowUndef)+            (intConst Int64 1)+            (segSpaceDims $ segSpace op)+  allocAtLevel (segLevel op) $ mapSegOpM (mapper num_threads) op   where+    maybe_grid =+      case (outer_lvl, segLevel op) of+        (Just (SegThread _ (Just grid)), _) -> Just grid+        (Just (SegGroup _ (Just grid)), _) -> Just grid+        (_, SegThread _ (Just grid)) -> Just grid+        (_, SegGroup _ (Just grid)) -> Just grid+        _ -> Nothing     scope = scopeOfSegSpace $ segSpace op-    lvl = segLevel op     mapper num_threads =       identitySegOpMapper         { mapOnSegOpBody =@@ -56,20 +72,22 @@         }     f = case segLevel op of       SegThread {} -> inThread+      SegThreadInGroup {} -> inThread       SegGroup {} -> inGroup     inThread env = env {envExpHints = inThreadExpHints}     inGroup env = env {envExpHints = inGroupExpHints}  handleHostOp ::+  Maybe SegLevel ->   HostOp GPU (SOAC GPU) ->   AllocM GPU GPUMem (MemOp (HostOp GPUMem ()))-handleHostOp (SizeOp op) =+handleHostOp _ (SizeOp op) =   pure $ Inner $ SizeOp op-handleHostOp (OtherOp op) =+handleHostOp _ (OtherOp op) =   error $ "Cannot allocate memory in SOAC: " ++ prettyString op-handleHostOp (SegOp op) =-  Inner . SegOp <$> handleSegOp op-handleHostOp (GPUBody ts (Body _ stms res)) =+handleHostOp outer_lvl (SegOp op) =+  Inner . SegOp <$> handleSegOp outer_lvl op+handleHostOp _ (GPUBody ts (Body _ stms res)) =   fmap (Inner . GPUBody ts) . buildBody_ . allocInStms stms $ pure res  kernelExpHints :: Exp GPUMem -> AllocM GPU GPUMem [ExpHint]@@ -79,9 +97,9 @@       dims' = rearrangeShape perm dims       ixfun = IxFun.permute (IxFun.iota $ map pe64 dims') perm_inv   pure [Hint ixfun DefaultSpace]-kernelExpHints (Op (Inner (SegOp (SegMap lvl@SegThread {} space ts body)))) =+kernelExpHints (Op (Inner (SegOp (SegMap lvl@(SegThread _ _) space ts body)))) =   zipWithM (mapResultHint lvl space) ts $ kernelBodyResult body-kernelExpHints (Op (Inner (SegOp (SegRed lvl@SegThread {} space reds ts body)))) =+kernelExpHints (Op (Inner (SegOp (SegRed lvl@(SegThread _ _) space reds ts body)))) =   (map (const NoHint) red_res <>) <$> zipWithM (mapResultHint lvl space) (drop num_reds ts) map_res   where     num_reds = segBinOpResults reds@@ -166,11 +184,11 @@  -- | The pass from 'GPU' to 'GPUMem'. explicitAllocations :: Pass GPU GPUMem-explicitAllocations = explicitAllocationsGeneric handleHostOp kernelExpHints+explicitAllocations = explicitAllocationsGeneric (handleHostOp Nothing) kernelExpHints  -- | Convert some 'GPU' stms to 'GPUMem'. explicitAllocationsInStms ::   (MonadFreshNames m, HasScope GPUMem m) =>   Stms GPU ->   m (Stms GPUMem)-explicitAllocationsInStms = explicitAllocationsInStmsGeneric handleHostOp kernelExpHints+explicitAllocationsInStms = explicitAllocationsInStmsGeneric (handleHostOp Nothing) kernelExpHints

@@ -21,6 +21,7 @@ import Control.Monad.Writer import Futhark.Analysis.PrimExp import Futhark.IR+import Futhark.IR.GPU.Op (SegVirt (..)) import Futhark.IR.Prop.Aliases import Futhark.IR.SegOp import Futhark.MonadFreshNames

@@ -39,6 +39,7 @@ import Data.Map qualified as M import Data.Maybe import Futhark.IR+import Futhark.IR.GPU.Op (SegVirt (..)) import Futhark.IR.SOACS (SOACS) import Futhark.IR.SOACS qualified as SOACS import Futhark.IR.SOACS.SOAC hiding (HistOp, histDest)@@ -375,7 +376,7 @@         Nothing -> addStmToAcc stm acc         Just acc' -> distribute =<< onInnerMap (MapLoop pat (stmAux stm) w lam arrs) acc' maybeDistributeStm stm@(Let pat aux (DoLoop merge form@ForLoop {} body)) acc-  | all (`notNameIn` freeIn pat) (patNames pat),+  | all (`notNameIn` freeIn (patTypes pat)) (patNames pat),     bodyContainsParallelism body =       distributeSingleStm acc stm >>= \case         Just (kernels, res, nest, acc')

@@ -60,12 +60,9 @@   let body = lambdaBody lam    group_size <- newVName "computed_group_size"-  let intra_lvl = SegThread (Count num_groups) (Count $ Var group_size) SegNoVirt-   (wss_min, wss_avail, log, kbody) <--    lift $-      localScope (scopeOfLParams $ lambdaParams lam) $-        intraGroupParalleliseBody intra_lvl body+    lift . localScope (scopeOfLParams $ lambdaParams lam) $+      intraGroupParalleliseBody body    outside_scope <- lift askScope   -- outside_scope may also contain the inputs, even though those are@@ -116,12 +113,10 @@    let nested_pat = loopNestingPat first_nest       rts = map (length ispace `stripArray`) $ patTypes nested_pat-      lvl = SegGroup (Count num_groups) (Count $ Var group_size) SegNoVirt+      grid = KernelGrid (Count num_groups) (Count $ Var group_size)+      lvl = SegGroup SegNoVirt (Just grid)       kstm =-        Let nested_pat aux $-          Op $-            SegOp $-              SegMap lvl kspace rts kbody'+        Let nested_pat aux $ Op $ SegOp $ SegMap lvl kspace rts kbody'    let intra_min_par = intra_avail_par   pure@@ -184,21 +179,21 @@         accAvailPar = S.singleton ws       } -intraGroupBody :: SegLevel -> Body SOACS -> IntraGroupM (Body GPU)-intraGroupBody lvl body = do-  stms <- collectStms_ $ intraGroupStms lvl $ bodyStms body+intraGroupBody :: Body SOACS -> IntraGroupM (Body GPU)+intraGroupBody body = do+  stms <- collectStms_ $ intraGroupStms $ bodyStms body   pure $ mkBody stms $ bodyResult body -intraGroupStm :: SegLevel -> Stm SOACS -> IntraGroupM ()-intraGroupStm lvl stm@(Let pat aux e) = do+intraGroupStm :: Stm SOACS -> IntraGroupM ()+intraGroupStm stm@(Let pat aux e) = do   scope <- askScope-  let lvl' = SegThread (segNumGroups lvl) (segGroupSize lvl) SegNoVirt+  let lvl = SegThread SegNoVirt Nothing    case e of     DoLoop merge form loopbody ->       localScope (scopeOf form') $         localScope (scopeOfFParams $ map fst merge) $ do-          loopbody' <- intraGroupBody lvl loopbody+          loopbody' <- intraGroupBody loopbody           certifying (stmAuxCerts aux) $             letBind pat $               DoLoop merge form' loopbody'@@ -207,13 +202,13 @@           ForLoop i it bound inps -> ForLoop i it bound inps           WhileLoop cond -> WhileLoop cond     Match cond cases defbody ifdec -> do-      cases' <- mapM (traverse $ intraGroupBody lvl) cases-      defbody' <- intraGroupBody lvl defbody+      cases' <- mapM (traverse intraGroupBody) cases+      defbody' <- intraGroupBody defbody       certifying (stmAuxCerts aux) . letBind pat $         Match cond cases' defbody' ifdec     Op soac       | "sequential_outer" `inAttrs` stmAuxAttrs aux ->-          intraGroupStms lvl . fmap (certify (stmAuxCerts aux))+          intraGroupStms . fmap (certify (stmAuxCerts aux))             =<< runBuilder_ (FOT.transformSOAC pat soac)     Op (Screma w arrs form)       | Just lam <- isMapSOAC form -> do@@ -229,7 +224,7 @@                     distOnInnerMap =                       distributeMap,                     distOnTopLevelStms =-                      liftInner . collectStms_ . intraGroupStms lvl,+                      liftInner . collectStms_ . intraGroupStms,                     distSegLevel = \minw _ _ -> do                       lift $ parallelMin minw                       pure lvl,@@ -252,7 +247,7 @@           let scanfun' = soacsLambdaToGPU scanfun               mapfun' = soacsLambdaToGPU mapfun           certifying (stmAuxCerts aux) $-            addStms =<< segScan lvl' pat mempty 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,@@ -260,7 +255,7 @@           let red_lam' = soacsLambdaToGPU red_lam               map_lam' = soacsLambdaToGPU map_lam           certifying (stmAuxCerts aux) $-            addStms =<< segRed lvl' pat mempty 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@@ -270,7 +265,7 @@        let bucket_fun' = soacsLambdaToGPU bucket_fun       certifying (stmAuxCerts aux) $-        addStms =<< segHist lvl' pat w [] [] ops' bucket_fun' arrs+        addStms =<< segHist lvl pat w [] [] ops' bucket_fun' arrs       parallelMin [w]     Op (Stream w arrs accs lam)       | chunk_size_param : _ <- lambdaParams lam -> do@@ -281,7 +276,7 @@               replace se = se               replaceSets (IntraAcc x y log) =                 IntraAcc (S.map (map replace) x) (S.map (map replace) y) log-          censor replaceSets $ intraGroupStms lvl stream_stms+          censor replaceSets $ intraGroupStms stream_stms     Op (Scatter w ivs lam dests) -> do       write_i <- newVName "write_i"       space <- mkSegSpace [(write_i, w)]@@ -308,23 +303,22 @@       certifying (stmAuxCerts aux) $ do         let ts = zipWith (stripArray . length) dests_ws $ patTypes pat             body = KernelBody () kstms krets-        letBind pat $ Op $ SegOp $ SegMap lvl' space ts body+        letBind pat $ Op $ SegOp $ SegMap lvl space ts body        parallelMin [w]     _ ->       addStm $ soacsStmToGPU stm -intraGroupStms :: SegLevel -> Stms SOACS -> IntraGroupM ()-intraGroupStms lvl = mapM_ (intraGroupStm lvl)+intraGroupStms :: Stms SOACS -> IntraGroupM ()+intraGroupStms = mapM_ intraGroupStm  intraGroupParalleliseBody ::   (MonadFreshNames m, HasScope GPU m) =>-  SegLevel ->   Body SOACS ->   m ([[SubExp]], [[SubExp]], Log, KernelBody GPU)-intraGroupParalleliseBody lvl body = do+intraGroupParalleliseBody body = do   (IntraAcc min_ws avail_ws log, kstms) <--    runIntraGroupM $ intraGroupStms lvl $ bodyStms body+    runIntraGroupM $ intraGroupStms $ bodyStms body   pure     ( S.toList min_ws,       S.toList avail_ws,

@@ -74,7 +74,9 @@             (SDivUp Int64 Unsafe)             (eSubExp w)             (eSubExp =<< asIntS Int64 group_size)-      pure $ SegThread (Count usable_groups) (Count group_size) SegNoVirt+      let grid = KernelGrid (Count usable_groups) (Count group_size)+      pure $ SegThread SegNoVirt (Just grid)     NoRecommendation v -> do       (num_groups, _) <- numberOfGroups desc w group_size-      pure $ SegThread (Count num_groups) (Count group_size) v+      let grid = KernelGrid (Count num_groups) (Count group_size)+      pure $ SegThread v (Just grid)

@@ -34,10 +34,12 @@   String ->   m SegLevel segThread desc =-  SegThread-    <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups)-    <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup)-    <*> pure SegVirt+  SegThread SegVirt <$> (Just <$> kernelGrid)+  where+    kernelGrid =+      KernelGrid+        <$> (Count <$> getSize (desc ++ "_num_groups") SizeNumGroups)+        <*> (Count <$> getSize (desc ++ "_group_size") SizeGroup)  injectSOACS ::   ( Monad m,

@@ -43,6 +43,7 @@     zEncodeString,     atMostChars,     invertMap,+    cartesian,     traverseFold,     fixPoint,   )@@ -57,7 +58,7 @@ import Data.ByteString.Base16 qualified as Base16 import Data.Char import Data.Either-import Data.Foldable (fold)+import Data.Foldable (fold, toList) import Data.Function ((&)) import Data.List (foldl', genericDrop, genericSplitAt, sortBy) import Data.List.NonEmpty qualified as NE@@ -411,6 +412,13 @@   M.toList m     & fmap (swap . first S.singleton)     & foldr (uncurry $ M.insertWith (<>)) mempty++-- | Compute the cartesian product of two foldable collections, using the given+-- combinator function.+cartesian :: (Monoid m, Foldable t) => (a -> a -> m) -> t a -> t a -> m+cartesian f xs ys =+  [(x, y) | x <- toList xs, y <- toList ys]+    & foldMap (uncurry f)  -- | Applicatively fold a traversable. traverseFold :: (Monoid m, Traversable t, Applicative f) => (a -> f m) -> t a -> f m

@@ -35,6 +35,7 @@     stack,     parensIf,     ppTuple',+    ppTupleLines',      -- * Operators     (</>),@@ -113,15 +114,16 @@ ppTuple' :: [Doc a] -> Doc a ppTuple' ets = braces $ commasep $ map align ets +ppTupleLines' :: [Doc a] -> Doc a+ppTupleLines' ets = braces $ commastack $ map align ets+ -- | Prettyprint a list enclosed in curly braces. prettyTuple :: Pretty a => [a] -> Text prettyTuple = docText . ppTuple' . map pretty  -- | Like 'prettyTuple', but put a linebreak after every element. prettyTupleLines :: Pretty a => [a] -> Text-prettyTupleLines = docText . ppTupleLines'-  where-    ppTupleLines' = braces . hsep . punctuate comma . map (align . pretty)+prettyTupleLines = docText . ppTupleLines' . map pretty  -- | The document @'apply' ds@ separates @ds@ with commas and encloses them with -- parentheses.

@@ -1033,7 +1033,7 @@     onModule (Module (Env terms types _)) =       Module $ Env (replaceM onTerm terms) (replaceM onType types) mempty     onModule (ModuleFun f) =-      ModuleFun $ \m -> onModule <$> f (substituteInModule (M.mapMaybe maybeHead rev_substs) m)+      ModuleFun $ \m -> onModule <$> f (substituteInModule substs m)     onTerm (TermValue t v) = TermValue t v     onTerm (TermPoly t v) = TermPoly t v     onTerm (TermModule m) = TermModule $ onModule m

@@ -253,8 +253,7 @@       _ -> hasArrayLit e prettyAppExp _ (LetFun fname (tparams, params, retdecl, rettype, e) body _) =   "let"-    <+> prettyName fname-    <+> hsep (map pretty tparams ++ map pretty params)+    <+> hsep (prettyName fname : map pretty tparams ++ map pretty params)       <> retdecl'     <+> equals     </> indent 2 (pretty e)@@ -468,8 +467,7 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (TypeBindBase f vn) where   pretty (TypeBind name l params te rt _ _) =     "type" <> pretty l-      <+> prettyName name-      <+> hsep (map pretty params)+      <+> hsep (prettyName name : map pretty params)       <+> equals       <+> maybe (pretty te) pretty (unAnnot rt) @@ -502,9 +500,9 @@ instance (Eq vn, IsName vn, Annot f) => Pretty (SpecBase f vn) where   pretty (TypeAbbrSpec tpsig) = pretty tpsig   pretty (TypeSpec l name ps _ _) =-    "type" <> pretty l <+> prettyName name <+> hsep (map pretty ps)+    "type" <> pretty l <+> hsep (prettyName name : map pretty ps)   pretty (ValSpec name tparams vtype _ _ _) =-    "val" <+> prettyName name <+> hsep (map pretty tparams) <> colon <+> pretty vtype+    "val" <+> hsep (prettyName name : map pretty tparams) <> colon <+> pretty vtype   pretty (ModSpec name sig _ _) =     "module" <+> prettyName name <> colon <+> pretty sig   pretty (IncludeSpec e _) =@@ -531,11 +529,11 @@  instance (Eq vn, IsName vn, Annot f) => Pretty (ModBindBase f vn) where   pretty (ModBind name ps sig e _ _) =-    "module" <+> prettyName name <+> hsep (map pretty ps) <+> sig' <> " =" <+> pretty e+    "module" <+> hsep (prettyName name : map pretty ps) <> sig' <> " =" <+> pretty e     where       sig' = case sig of         Nothing -> mempty-        Just (s, _) -> colon <+> pretty s <> " "+        Just (s, _) -> " " <> colon <+> pretty s <> " "  ppBinOp :: IsName v => QualName v -> Doc a ppBinOp bop =

@@ -8,6 +8,7 @@   ) where +import Data.List qualified as L import Data.Map.Strict qualified as M import Data.Maybe import Futhark.Util (maybeHead, nubOrd)@@ -23,16 +24,16 @@   deriving (Eq, Ord, Show)  -- | A representation of the essentials of a pattern.-data Match-  = MatchWild StructType-  | MatchConstr Constr [Match] StructType+data Match t+  = MatchWild t+  | MatchConstr Constr [Match t] t   deriving (Eq, Ord, Show) -matchType :: Match -> StructType+matchType :: Match StructType -> StructType matchType (MatchWild t) = t matchType (MatchConstr _ _ t) = t -pprMatch :: Int -> Match -> Doc a+pprMatch :: Int -> Match t -> Doc a pprMatch _ MatchWild {} = "_" pprMatch _ (MatchConstr (ConstrLit l) _ _) = pretty l pprMatch p (MatchConstr (Constr c) ps _) =@@ -45,10 +46,10 @@   where     ppField name t = pretty (nameToString name) <> equals <> pprMatch (-1) t -instance Pretty Match where+instance Pretty (Match t) where   pretty = pprMatch (-1) -patternToMatch :: Pat -> Match+patternToMatch :: Pat -> Match StructType patternToMatch (Id _ (Info t) _) = MatchWild $ toStruct t patternToMatch (Wildcard (Info t) _) = MatchWild $ toStruct t patternToMatch (PatParens p _) = patternToMatch p@@ -67,29 +68,35 @@ patternToMatch (PatConstr c (Info t) args _) =   MatchConstr (Constr c) (map patternToMatch args) $ toStruct t -isConstr :: Match -> Maybe Name+isConstr :: Match t -> Maybe Name isConstr (MatchConstr (Constr c) _ _) = Just c isConstr _ = Nothing -complete :: [Match] -> Bool+isBool :: Match t -> Maybe Bool+isBool (MatchConstr (ConstrLit (PatLitPrim (BoolValue b))) _ _) = Just b+isBool _ = Nothing++complete :: [Match StructType] -> Bool complete xs   | Just x <- maybeHead xs,     Scalar (Sum all_cs) <- matchType x,     Just xs_cs <- mapM isConstr xs =       all (`elem` xs_cs) (M.keys all_cs)   | otherwise =-      (any (isBool True) xs && any (isBool False) xs)+      all (`elem` fromMaybe [] (mapM isBool xs)) [True, False]         || all isRecord xs         || all isTuple xs   where-    isBool b1 (MatchConstr (ConstrLit (PatLitPrim (BoolValue b2))) _ _) = b1 == b2-    isBool _ _ = False     isRecord (MatchConstr ConstrRecord {} _ _) = True     isRecord _ = False     isTuple (MatchConstr ConstrTuple _ _) = True     isTuple _ = False -specialise :: [StructType] -> Match -> [[Match]] -> [[Match]]+specialise ::+  [StructType] ->+  Match StructType ->+  [[Match StructType]] ->+  [[Match StructType]] specialise ats c1 = go   where     go ((c2 : row) : ps)@@ -109,14 +116,14 @@     match _ _ =       Nothing -defaultMat :: [[Match]] -> [[Match]]+defaultMat :: [[Match t]] -> [[Match t]] defaultMat = mapMaybe onRow   where     onRow (MatchConstr {} : _) = Nothing     onRow (MatchWild {} : ps) = Just ps     onRow [] = Nothing -- Should not happen. -findUnmatched :: [[Match]] -> Int -> [[Match]]+findUnmatched :: [[Match StructType]] -> Int -> [[Match ()]] findUnmatched pmat n   | ((p : _) : _) <- pmat,     Just heads <- mapM maybeHead pmat =@@ -133,40 +140,44 @@           pmat' = specialise ats c pmat       u <- findUnmatched pmat' (a_k + n - 1)       pure $ case c of-        MatchConstr c' _ t ->+        MatchConstr c' _ _ ->           let (r, p) = splitAt a_k u-           in MatchConstr c' r t : p-        MatchWild t ->-          MatchWild t : u+           in MatchConstr c' r () : p+        MatchWild _ ->+          MatchWild () : u      incompleteCase pt cs = do       u <- findUnmatched (defaultMat pmat) (n - 1)       if null cs-        then pure $ MatchWild pt : u+        then pure $ MatchWild () : u         else case pt of           Scalar (Sum all_cs) -> do             -- Figure out which constructors are missing.             let sigma = mapMaybe isConstr cs                 notCovered (k, _) = k `notElem` sigma             (cname, ts) <- filter notCovered $ M.toList all_cs-            pure $ MatchConstr (Constr cname) (map MatchWild ts) pt : u-          _ ->-            -- This is where we could have enumerated missing match-            -- values (e.g. for booleans), rather than just emitting a-            -- wildcard.-            pure $ MatchWild pt : u---- If we get here, then the number of columns must be zero.-findUnmatched [] _ = [[]]+            pure $ MatchConstr (Constr cname) (map (const (MatchWild ())) ts) () : u+          Scalar (Prim Bool) -> do+            -- Figure out which constants are missing.+            let sigma = mapMaybe isBool cs+            b <- filter (`notElem` sigma) [True, False]+            pure $ MatchConstr (ConstrLit (PatLitPrim (BoolValue b))) [] () : u+          _ -> do+            -- FIXME: this is wrong in the unlikely case where someone+            -- is pattern-matching every single possible number for+            -- some numeric type.  It should be handled more like Bool+            -- above.+            pure $ MatchWild () : u+findUnmatched [] n = [replicate n $ MatchWild ()] findUnmatched _ _ = []  {-# NOINLINE unmatched #-}  -- | Find the unmatched cases.-unmatched :: [Pat] -> [Match]+unmatched :: [Pat] -> [Match ()] unmatched orig_ps =   -- The algorithm may find duplicate example, which we filter away   -- here.   nubOrd $     mapMaybe maybeHead $-      findUnmatched (map ((: []) . patternToMatch) orig_ps) 1+      findUnmatched (map (L.singleton . patternToMatch) orig_ps) 1

@@ -375,8 +375,7 @@         <+> hsep (map pretty ps) ppTypeAbbr _ name (l, ps, t) =   "type" <> pretty l-    <+> pretty name-    <+> hsep (map pretty ps)+    <+> hsep (pretty name : map pretty ps)     <+> equals     <+> nest 2 (align (pretty t)) @@ -455,11 +454,19 @@       (ModFun (FunSig mod_abs mod_pmod mod_mod))       (ModFun (FunSig sig_abs sig_pmod sig_mod))       loc = do+        -- We need to use different substitutions when matching+        -- parameter and body signatures - this is because the+        -- concrete parameter must be *at least as* general as the+        -- ascripted parameter, while the concrete body must be *at+        -- most as* general as the ascripted body.         abs_substs <- resolveAbsTypes mod_abs mod_pmod sig_abs loc+        p_abs_substs <- resolveAbsTypes sig_abs sig_pmod mod_abs loc         let abs_subst_to_type =               old_abs_subst_to_type <> M.map (substFromAbbr . snd) abs_substs+            p_abs_subst_to_type =+              old_abs_subst_to_type <> M.map (substFromAbbr . snd) p_abs_substs             abs_name_substs = M.map (qualLeaf . fst) abs_substs-        pmod_substs <- matchMods abs_subst_to_type quals mod_pmod sig_pmod loc+        pmod_substs <- matchMods p_abs_subst_to_type quals sig_pmod mod_pmod loc         mod_substs <- matchMTys' abs_subst_to_type quals mod_mod sig_mod loc         pure (pmod_substs <> mod_substs <> abs_name_substs) 

@@ -895,10 +895,9 @@            in zeroOrderType (mkUsage argloc "potential consumption in expression") msg tp1         _ -> pure () -      occurs <- (dflow `seqOccurrences`) <$> consumeArg argloc argtype' (diet tp1')--      checkIfConsumable loc $ S.map AliasBound $ allConsumed occurs-      occur occurs+      arg_consumed <- consumedByArg argloc argtype' (diet tp1')+      checkIfConsumable loc $ mconcat arg_consumed+      occur $ dflow `seqOccurrences` map (`consumption` argloc) arg_consumed        -- Unification ignores uniqueness in higher-order arguments, so       -- we check for that here.@@ -1011,12 +1010,12 @@ maskAliases t FuncDiet {} = t maskAliases _ _ = error "Invalid arguments passed to maskAliases." -consumeArg :: SrcLoc -> PatType -> Diet -> TermTypeM [Occurrence]-consumeArg loc (Scalar (Record ets)) (RecordDiet ds) =-  concat . M.elems <$> traverse (uncurry $ consumeArg loc) (M.intersectionWith (,) ets ds)-consumeArg loc (Scalar (Sum ets)) (SumDiet ds) =-  concat <$> traverse (uncurry $ consumeArg loc) (concat $ M.elems $ M.intersectionWith zip ets ds)-consumeArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _)+consumedByArg :: SrcLoc -> PatType -> Diet -> TermTypeM [Aliasing]+consumedByArg loc (Scalar (Record ets)) (RecordDiet ds) =+  mconcat . M.elems <$> traverse (uncurry $ consumedByArg loc) (M.intersectionWith (,) ets ds)+consumedByArg loc (Scalar (Sum ets)) (SumDiet ds) =+  mconcat <$> traverse (uncurry $ consumedByArg loc) (concat $ M.elems $ M.intersectionWith zip ets ds)+consumedByArg loc (Scalar (Arrow _ _ t1 _)) (FuncDiet d _)   | not $ contravariantArg t1 d =       typeError loc mempty . withIndexLink "consuming-argument" $         "Non-consuming higher-order parameter passed consuming argument."@@ -1031,8 +1030,8 @@       contravariantArg tp dp && contravariantArg tr dr     contravariantArg _ _ =       True-consumeArg loc at Consume = pure [consumption (aliases at) loc]-consumeArg loc at _ = pure [observation (aliases at) loc]+consumedByArg _ at Consume = pure [aliases at]+consumedByArg _ _ _ = pure []  -- | Type-check a single expression in isolation.  This expression may -- turn out to be polymorphic, in which case the list of type

@@ -69,7 +69,6 @@     -- * Errors     useAfterConsume,     unusedSize,-    notConsumable,     uniqueReturnAliased,     returnAliased,     badLetWithValue,@@ -970,14 +969,15 @@ checkIfConsumable :: SrcLoc -> Aliasing -> TermTypeM () checkIfConsumable loc als = do   vtable <- asks $ scopeVtable . termScope-  let consumable v = case M.lookup v vtable of+  let boundAlias (AliasBound v) = Just v+      boundAlias (AliasFree _) = Nothing+      consumable v = case M.lookup v vtable of         Just (BoundV Local _ t)           | Scalar Arrow {} <- t -> False           | otherwise -> True         Just (BoundV l _ _) -> l == Local-        _ -> True-  -- The sort ensures that AliasBound vars are shown before AliasFree.-  case map aliasVar $ sort $ filter (not . consumable . aliasVar) $ S.toList als of+        _ -> False -- Implies name from module.+  case sort $ filter (not . consumable) $ mapMaybe boundAlias $ S.toList als of     v : _ -> notConsumable loc =<< describeVar loc v     [] -> pure ()