diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,6 +1,6 @@ cabal-version: 2.4 name: futhark -version: 0.21.13 +version: 0.21.14 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to @@ -37,6 +37,8 @@ -- Cabal's recompilation tracking doesn't work when we use wildcards -- here, so for now we spell out every single file. rts/c/atomics.h + rts/c/context.h + rts/c/context_prototypes.h rts/c/lock.h rts/c/timing.h rts/c/errors.h @@ -116,7 +118,6 @@ Futhark.Analysis.Metrics.Type Futhark.Analysis.PrimExp Futhark.Analysis.PrimExp.Convert - Futhark.Analysis.PrimExp.Generalize Futhark.Analysis.PrimExp.Parse Futhark.Analysis.PrimExp.Simplify Futhark.Analysis.Rephrase @@ -157,8 +158,12 @@ Futhark.CodeGen.Backends.COpenCL.Boilerplate Futhark.CodeGen.Backends.GenericC Futhark.CodeGen.Backends.GenericC.CLI + Futhark.CodeGen.Backends.GenericC.Code + Futhark.CodeGen.Backends.GenericC.EntryPoints + Futhark.CodeGen.Backends.GenericC.Monad Futhark.CodeGen.Backends.GenericC.Options Futhark.CodeGen.Backends.GenericC.Server + Futhark.CodeGen.Backends.GenericC.Types Futhark.CodeGen.Backends.GenericPython Futhark.CodeGen.Backends.GenericPython.AST Futhark.CodeGen.Backends.GenericPython.Options @@ -271,6 +276,7 @@ Futhark.Optimise.BlkRegTiling Futhark.Optimise.CSE Futhark.Optimise.DoubleBuffer + Futhark.Optimise.EntryPointMem Futhark.Optimise.Fusion Futhark.Optimise.Fusion.Composing Futhark.Optimise.Fusion.GraphRep @@ -296,6 +302,7 @@ Futhark.Optimise.Simplify.Rules.ClosedForm Futhark.Optimise.Simplify.Rules.Index Futhark.Optimise.Simplify.Rules.Loop + Futhark.Optimise.Simplify.Rules.Match Futhark.Optimise.Simplify.Rules.Simple Futhark.Optimise.Sink Futhark.Optimise.TileLoops @@ -388,7 +395,7 @@ aeson >=2.0.0.0 , ansi-terminal >=0.6.3.1 , array >=0.4 - , base >=4.13 && <5 + , base >=4.15 && <5 , base16-bytestring , binary >=0.8.3 , blaze-html >=0.9.0.1 diff --git a/prelude/array.fut b/prelude/array.fut --- a/prelude/array.fut +++ b/prelude/array.fut @@ -77,7 +77,12 @@ -- -- For example, if `b==rotate r a`, then `b[x] = a[x+r]`. -- --- **Complexity:** O(1). +-- **Work:** O(n). +-- +-- **Span:** O(1). +-- +-- Note: In most cases, `rotate` will be fused with subsequent +-- operations such as `map`, in which case it is free. def rotate [n] 't (r: i64) (xs: [n]t): [n]t = intrinsics.rotate (r, xs) -- | Construct an array of consecutive integers of the given length, diff --git a/rts/c/context.h b/rts/c/context.h new file mode 100644 --- /dev/null +++ b/rts/c/context.h @@ -0,0 +1,41 @@ +// Start of context.h + +// Eventually it would be nice to move the context definition in here +// instead of generating it in the compiler. For now it defines +// various helper functions that must be available. + +// Internal functions. + +static void set_error(struct futhark_context* ctx, char *error) { + lock_lock(&ctx->error_lock); + if (ctx->error == NULL) { + ctx->error = error; + } else { + free(error); + } + lock_unlock(&ctx->error_lock); +} + +// XXX: should be static, but used in ispc_util.h +void lexical_realloc_error(struct futhark_context* ctx, size_t new_size) { + set_error(ctx, + msgprintf("Failed to allocate memory.\nAttempted allocation: %12lld bytes\n", + (long long) new_size)); +} + +static int lexical_realloc(struct futhark_context *ctx, + unsigned char **ptr, + int64_t *old_size, + int64_t new_size) { + unsigned char *new = realloc(*ptr, (size_t)new_size); + if (new == NULL) { + lexical_realloc_error(ctx, new_size); + return FUTHARK_OUT_OF_MEMORY; + } else { + *ptr = new; + *old_size = new_size; + return FUTHARK_SUCCESS; + } +} + +// End of context.h diff --git a/rts/c/context_prototypes.h b/rts/c/context_prototypes.h new file mode 100644 --- /dev/null +++ b/rts/c/context_prototypes.h @@ -0,0 +1,11 @@ +// Start of context_prototypes.h +// +// Prototypes for the functions in prototypes.h that need to be +// available very early. + +struct futhark_context_config; +struct futhark_context; + +static void set_error(struct futhark_context* ctx, char *error); + +// End of of context_prototypes.h diff --git a/rts/c/ispc_util.h b/rts/c/ispc_util.h --- a/rts/c/ispc_util.h +++ b/rts/c/ispc_util.h @@ -164,16 +164,15 @@ } extern "C" unmasked uniform unsigned char * uniform realloc(uniform unsigned char * uniform ptr, uniform int64_t new_size); -extern "C" unmasked uniform char * uniform lexical_realloc_error(uniform int64_t new_size); -extern "C" unmasked uniform char * uniform * uniform futhark_get_error_ref(uniform struct futhark_context * uniform ctx); +extern "C" unmasked uniform char * uniform lexical_realloc_error(uniform struct futhark_context * uniform ctx, uniform int64_t new_size); -static inline uniform int lexical_realloc(uniform char * uniform * uniform error, +static inline uniform int lexical_realloc(uniform struct futhark_context * uniform ctx, unsigned char uniform * uniform * uniform ptr, int64_t uniform * uniform old_size, uniform int64_t new_size) { uniform unsigned char * uniform memptr = realloc(*ptr, new_size); if (memptr == NULL) { - *error = lexical_realloc_error(new_size); + lexical_realloc_error(ctx, new_size); return FUTHARK_OUT_OF_MEMORY; } else { *ptr = memptr; @@ -183,14 +182,14 @@ } -static inline uniform int lexical_realloc(uniform char * uniform * uniform error, +static inline uniform int lexical_realloc(uniform struct futhark_context *ctx, unsigned char uniform * uniform * uniform ptr, int64_t uniform * uniform old_size, varying int64_t new_size) { - return lexical_realloc(error, ptr, old_size, reduce_max(new_size)); + return lexical_realloc(ctx, ptr, old_size, reduce_max(new_size)); } -static inline uniform int lexical_realloc(uniform char * uniform * uniform error, +static inline uniform int lexical_realloc(uniform struct futhark_context * uniform ctx, unsigned char uniform * varying * uniform ptr, int64_t uniform * varying old_size, varying int64_t new_size) { @@ -198,7 +197,7 @@ foreach_active(i){ uniform unsigned char * uniform memptr = realloc(extract(*ptr,i), extract(new_size,i)); if (memptr == NULL) { - *error = lexical_realloc_error(extract(new_size,i)); + lexical_realloc_error(ctx, extract(new_size,i)); err = FUTHARK_OUT_OF_MEMORY; } else { *ptr = (uniform unsigned char * varying)insert((int64_t)*ptr, i, (uniform int64_t) memptr); @@ -208,7 +207,7 @@ return err; } -static inline uniform int lexical_realloc(uniform char * uniform * uniform error, +static inline uniform int lexical_realloc(uniform struct futhark_context * uniform ctx, unsigned char uniform * varying * uniform ptr, int64_t varying * uniform old_size, varying int64_t new_size) { @@ -216,7 +215,7 @@ foreach_active(i){ uniform unsigned char * uniform memptr = realloc(extract(*ptr,i), extract(new_size,i)); if (memptr == NULL) { - *error = lexical_realloc_error(extract(new_size,i)); + lexical_realloc_error(ctx, extract(new_size,i)); err = FUTHARK_OUT_OF_MEMORY; } else { *ptr = (uniform unsigned char * varying)insert((int64_t)*ptr, i, (uniform int64_t) memptr); @@ -226,14 +225,14 @@ return err; } -static inline uniform int lexical_realloc(uniform char * uniform * uniform error, +static inline uniform int lexical_realloc(uniform struct futhark_context * uniform ctx, unsigned char uniform * varying * uniform ptr, size_t varying * uniform old_size, varying int64_t new_size) { - return lexical_realloc(error, ptr, (varying int64_t * uniform)old_size, new_size); + return lexical_realloc(ctx, ptr, (varying int64_t * uniform)old_size, new_size); } -static inline uniform int lexical_realloc(uniform char * uniform * uniform error, +static inline uniform int lexical_realloc(uniform struct futhark_context * uniform ctx, unsigned char varying * uniform * uniform ptr, size_t varying * uniform old_size, uniform int64_t new_size) { @@ -241,7 +240,7 @@ uniform unsigned char * uniform memptr = realloc((uniform unsigned char * uniform )*ptr, new_size*programCount); if (memptr == NULL) { - *error = lexical_realloc_error(new_size); + lexical_realloc_error(ctx, new_size); err = FUTHARK_OUT_OF_MEMORY; } else { *ptr = (varying unsigned char * uniform)memptr; @@ -251,11 +250,11 @@ return err; } -static inline uniform int lexical_realloc(uniform char * uniform * uniform error, +static inline uniform int lexical_realloc(uniform struct futhark_context * uniform ctx, unsigned char varying * uniform * uniform ptr, size_t varying * uniform old_size, varying int64_t new_size) { - return lexical_realloc(error, ptr, old_size, reduce_max(new_size)); + return lexical_realloc(ctx, ptr, old_size, reduce_max(new_size)); } extern "C" unmasked uniform int memblock_unref(uniform struct futhark_context * uniform ctx, diff --git a/rts/c/util.h b/rts/c/util.h --- a/rts/c/util.h +++ b/rts/c/util.h @@ -135,22 +135,4 @@ b->used += needed; } -char * lexical_realloc_error(size_t new_size) { - return msgprintf("Failed to allocate memory.\nAttempted allocation: %12lld bytes\n", - (long long) new_size); -} - -static int lexical_realloc(char **error, unsigned char **ptr, int64_t *old_size, int64_t new_size) { - unsigned char *new = realloc(*ptr, (size_t)new_size); - if (new == NULL) { - *error = msgprintf("Failed to allocate memory.\nAttempted allocation: %12lld bytes\n", - (long long) new_size); - return FUTHARK_OUT_OF_MEMORY; - } else { - *ptr = new; - *old_size = new_size; - return FUTHARK_SUCCESS; - } -} - // End of util.h. diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs --- a/src/Futhark/AD/Fwd.hs +++ b/src/Futhark/AD/Fwd.hs @@ -253,9 +253,9 @@ addStm $ Let pat_tan aux $ BasicOp $ Replicate n x_tan Scratch t shape -> addStm $ Let pat_tan aux $ BasicOp $ Scratch t shape - Reshape reshape arr -> do + Reshape k reshape arr -> do arr_tan <- tangent arr - addStm $ Let pat_tan aux $ BasicOp $ Reshape reshape arr_tan + addStm $ Let pat_tan aux $ BasicOp $ Reshape k reshape arr_tan Rearrange perm arr -> do arr_tan <- tangent arr addStm $ Let pat_tan aux $ BasicOp $ Rearrange perm arr_tan @@ -422,12 +422,12 @@ e_t = primExpType e zipWithM_ (letBindNames . pure) (patNames pat_tan) =<< mapM toExp (zipWith (~*~) (map (convertTo ret) arg_tans) derivs) -fwdStm (Let pat aux (If cond t f (IfDec ret ifsort))) = do - t' <- slocal' $ fwdBody t - f' <- slocal' $ fwdBody f +fwdStm (Let pat aux (Match ses cases defbody (MatchDec ret ifsort))) = do + cases' <- slocal' $ mapM (traverse fwdBody) cases + defbody' <- slocal' $ fwdBody defbody pat' <- bundleNew pat ret' <- bundleTan ret - addStm $ Let pat' aux $ If cond t' f' $ IfDec ret' ifsort + addStm $ Let pat' aux $ Match ses cases' defbody' $ MatchDec ret' ifsort fwdStm (Let pat aux (DoLoop val_pats loop@(WhileLoop v) body)) = do val_pats' <- bundleNew val_pats pat' <- bundleNew pat diff --git a/src/Futhark/AD/Rev.hs b/src/Futhark/AD/Rev.hs --- a/src/Futhark/AD/Rev.hs +++ b/src/Futhark/AD/Rev.hs @@ -57,11 +57,11 @@ case t of FloatType ft -> update <=< letExp "contrib" $ - If - (Var pat_adj) - (resultBody [constant (floatValue ft (1 :: Int))]) + Match + [Var pat_adj] + [Case [Just $ BoolValue True] $ resultBody [constant (floatValue ft (1 :: Int))]] (resultBody [constant (floatValue ft (0 :: Int))]) - (IfDec [Prim (FloatType ft)] IfNormal) + (MatchDec [Prim (FloatType ft)] MatchNormal) IntType it -> update <=< letExp "contrib" $ BasicOp $ ConvOp (BToI it) (Var pat_adj) Bool -> @@ -130,13 +130,13 @@ (_pat_v, pat_adj) <- commonBasicOp pat aux e m returnSweepCode $ updateSubExpAdj se pat_adj -- - Reshape _ arr -> do + Reshape k _ arr -> do (_pat_v, pat_adj) <- commonBasicOp pat aux e m returnSweepCode $ do - arr_dims <- arrayDims <$> lookupType arr + arr_shape <- arrayShape <$> lookupType arr void $ updateAdj arr <=< letExp "adj_reshape" . BasicOp $ - Reshape (map DimNew arr_dims) pat_adj + Reshape k arr_shape pat_adj -- Rearrange perm arr -> do (_pat_v, pat_adj) <- commonBasicOp pat aux e m @@ -163,7 +163,7 @@ n <- letSubExp "rep_size" =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ns pat_adj_flat <- letExp (baseString pat_adj <> "_flat") . BasicOp $ - Reshape (map DimNew $ n : arrayDims x_t) pat_adj + Reshape ReshapeArbitrary (Shape $ n : arrayDims x_t) pat_adj reduce <- reduceSOAC [Reduce Commutative lam [ne]] updateSubExpAdj x =<< letExp "rep_contrib" (Op $ Screma n [pat_adj_flat] reduce) @@ -263,13 +263,13 @@ letExp "contrib" <=< toExp . convert ret argt $ pat_adj' ~*~ deriv zipWithM_ updateSubExpAdj (map fst args) contribs -diffStm stm@(Let pat _ (If cond tbody fbody _)) m = do +diffStm stm@(Let pat _ (Match ses cases defbody _)) m = do addStm stm m returnSweepCode $ do - let tbody_free = freeIn tbody - fbody_free = freeIn fbody - branches_free = namesToList $ tbody_free <> fbody_free + let cases_free = map freeIn cases + defbody_free = freeIn defbody + branches_free = namesToList $ mconcat $ defbody_free : cases_free adjs <- mapM lookupAdj $ patNames pat @@ -278,10 +278,10 @@ <=< letTupExp "branch_adj" <=< renameExp ) - =<< eIf - (eSubExp cond) - (diffBody adjs branches_free tbody) - (diffBody adjs branches_free fbody) + =<< eMatch + ses + (map (fmap $ diffBody adjs branches_free) cases) + (diffBody adjs branches_free defbody) zipWithM_ insAdj branches_free branches_free_adj diffStm (Let pat aux (Op soac)) m = vjpSOAC vjpOps pat aux soac m diff --git a/src/Futhark/AD/Rev/Monad.hs b/src/Futhark/AD/Rev/Monad.hs --- a/src/Futhark/AD/Rev/Monad.hs +++ b/src/Futhark/AD/Rev/Monad.hs @@ -105,8 +105,8 @@ OutOfBounds deriving (Eq, Ord, Show) --- | A symbolic representation of an array that is all zeroes, except at one --- index. +-- | A symbolic representation of an array that is all zeroes, except +-- at certain indexes. data Sparse = Sparse { -- | The shape of the array. sparseShape :: Shape, diff --git a/src/Futhark/Analysis/Alias.hs b/src/Futhark/Analysis/Alias.hs --- a/src/Futhark/Analysis/Alias.hs +++ b/src/Futhark/Analysis/Alias.hs @@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} -- | Alias analysis of a full Futhark program. Takes as input a -- program with an arbitrary rep and produces one with aliases. This @@ -90,12 +91,12 @@ Exp (Aliases rep) -- Would be better to put this in a BranchType annotation, but that -- requires a lot of other work. -analyseExp aliases (If cond tb fb dec) = - let Body ((tb_als, tb_cons), tb_dec) tb_stms tb_res = analyseBody aliases tb - Body ((fb_als, fb_cons), fb_dec) fb_stms fb_res = analyseBody aliases fb - cons = tb_cons <> fb_cons +analyseExp aliases (Match cond cases defbody matchdec) = + let cases' = map (fmap $ analyseBody aliases) cases + defbody' = analyseBody aliases defbody + all_cons = foldMap (snd . fst . bodyDec) $ defbody' : map caseBody cases' isConsumed v = - any (`nameIn` unAliases cons) $ + any (`nameIn` unAliases all_cons) $ v : namesToList (M.findWithDefault mempty v aliases) notConsumed = AliasDec @@ -103,11 +104,11 @@ . filter (not . isConsumed) . namesToList . unAliases - tb_als' = map notConsumed tb_als - fb_als' = map notConsumed fb_als - tb' = Body ((tb_als', tb_cons), tb_dec) tb_stms tb_res - fb' = Body ((fb_als', fb_cons), fb_dec) fb_stms fb_res - in If cond tb' fb' dec + onBody (Body ((als, cons), dec) stms res) = + Body ((map notConsumed als, cons), dec) stms res + cases'' = map (fmap onBody) cases' + defbody'' = onBody defbody' + in Match cond cases'' defbody'' matchdec analyseExp aliases e = mapExp analyse e where analyse = diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs --- a/src/Futhark/Analysis/DataDependencies.hs +++ b/src/Futhark/Analysis/DataDependencies.hs @@ -8,6 +8,7 @@ ) where +import qualified Data.List as L import qualified Data.Map.Strict as M import Futhark.IR @@ -28,14 +29,15 @@ Dependencies dataDependencies' startdeps = foldl grow startdeps . bodyStms where - grow deps (Let pat _ (If c tb fb _)) = - let tdeps = dataDependencies' deps tb - fdeps = dataDependencies' deps fb - cdeps = depsOf deps c - comb (pe, SubExpRes _ tres, SubExpRes _ fres) = + grow deps (Let pat _ (Match c cases defbody _)) = + let cases_deps = map (dataDependencies' deps . caseBody) cases + defbody_deps = dataDependencies' deps defbody + cdeps = foldMap (depsOf deps) c + comb (pe, se_cases_deps, se_defbody_deps) = ( patElemName pe, mconcat $ - [freeIn pe, cdeps, depsOf tdeps tres, depsOf fdeps fres] + se_cases_deps + ++ [freeIn pe, cdeps, se_defbody_deps] ++ map (depsOfVar deps) (namesToList $ freeIn pe) ) branchdeps = @@ -43,9 +45,11 @@ map comb $ zip3 (patElems pat) - (bodyResult tb) - (bodyResult fb) - in M.unions [branchdeps, deps, tdeps, fdeps] + ( L.transpose . zipWith (map . depsOf) cases_deps $ + map (map resSubExp . bodyResult . caseBody) cases + ) + (map (depsOf defbody_deps . resSubExp) (bodyResult defbody)) + in M.unions $ [branchdeps, deps, defbody_deps] ++ cases_deps grow deps (Let pat _ e) = let free = freeIn pat <> freeIn e freeDeps = mconcat $ map (depsOfVar deps) $ namesToList free diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs --- a/src/Futhark/Analysis/HORep/SOAC.hs +++ b/src/Futhark/Analysis/HORep/SOAC.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} -- | High-level representation of SOACs. When performing @@ -80,7 +79,6 @@ Rearrange, Replicate, Reshape, - Var, typeOf, ) import qualified Futhark.IR as Futhark @@ -103,11 +101,11 @@ = -- | A permutation of an otherwise valid input. Rearrange Certs [Int] | -- | A reshaping of an otherwise valid input. - Reshape Certs (ShapeChange SubExp) + Reshape Certs ReshapeKind Shape | -- | A reshaping of the outer dimension. - ReshapeOuter Certs (ShapeChange SubExp) + ReshapeOuter Certs ReshapeKind Shape | -- | A reshaping of everything but the outer dimension. - ReshapeInner Certs (ShapeChange SubExp) + ReshapeInner Certs ReshapeKind Shape | -- | Replicate the rows of the array a number of times. Replicate Certs Shape deriving (Show, Eq, Ord) @@ -115,12 +113,12 @@ instance Substitute ArrayTransform where substituteNames substs (Rearrange cs xs) = Rearrange (substituteNames substs cs) xs - substituteNames substs (Reshape cs ses) = - Reshape (substituteNames substs cs) (substituteNames substs ses) - substituteNames substs (ReshapeOuter cs ses) = - ReshapeOuter (substituteNames substs cs) (substituteNames substs ses) - substituteNames substs (ReshapeInner cs ses) = - ReshapeInner (substituteNames substs cs) (substituteNames substs ses) + substituteNames substs (Reshape cs k ses) = + Reshape (substituteNames substs cs) k (substituteNames substs ses) + substituteNames substs (ReshapeOuter cs k ses) = + ReshapeOuter (substituteNames substs cs) k (substituteNames substs ses) + substituteNames substs (ReshapeInner cs k ses) = + ReshapeInner (substituteNames substs cs) k (substituteNames substs ses) substituteNames substs (Replicate cs se) = Replicate (substituteNames substs cs) (substituteNames substs se) @@ -231,9 +229,9 @@ transformFromExp :: Certs -> Exp rep -> Maybe (VName, ArrayTransform) transformFromExp cs (BasicOp (Futhark.Rearrange perm v)) = Just (v, Rearrange cs perm) -transformFromExp cs (BasicOp (Futhark.Reshape shape v)) = - Just (v, Reshape cs shape) -transformFromExp cs (BasicOp (Futhark.Replicate shape (Futhark.Var v))) = +transformFromExp cs (BasicOp (Futhark.Reshape k shape v)) = + Just (v, Reshape cs k shape) +transformFromExp cs (BasicOp (Futhark.Replicate shape (Var v))) = Just (v, Replicate cs shape) transformFromExp _ _ = Nothing @@ -268,12 +266,12 @@ isVarInput (Input ts v _) | nullTransforms ts = Just v isVarInput _ = Nothing --- | If the given input is a plain variable input, with no non-vacuous transforms, --- return the variable. +-- | If the given input is a plain variable input, with no non-vacuous +-- transforms, return the variable. isVarishInput :: Input -> Maybe VName isVarishInput (Input ts v t) | nullTransforms ts = Just v - | Reshape cs [DimCoercion _] :< ts' <- viewf ts, + | Reshape cs ReshapeCoerce (Shape [_]) :< ts' <- viewf ts, cs == mempty = isVarishInput $ Input ts' v t isVarishInput _ = Nothing @@ -291,22 +289,22 @@ applyTransform :: MonadBuilder m => ArrayTransform -> VName -> m VName applyTransform (Replicate cs n) ia = certifying cs . letExp "repeat" . BasicOp $ - Futhark.Replicate n (Futhark.Var ia) + Futhark.Replicate n (Var ia) applyTransform (Rearrange cs perm) ia = do r <- arrayRank <$> lookupType ia certifying cs . letExp "rearrange" . BasicOp $ Futhark.Rearrange (perm ++ [length perm .. r - 1]) ia -applyTransform (Reshape cs shape) ia = +applyTransform (Reshape cs k shape) ia = certifying cs . letExp "reshape" . BasicOp $ - Futhark.Reshape shape ia -applyTransform (ReshapeOuter cs shape) ia = do + Futhark.Reshape k shape ia +applyTransform (ReshapeOuter cs k shape) ia = do shape' <- reshapeOuter shape 1 . arrayShape <$> lookupType ia certifying cs . letExp "reshape_outer" . BasicOp $ - Futhark.Reshape shape' ia -applyTransform (ReshapeInner cs shape) ia = do + Futhark.Reshape k shape' ia +applyTransform (ReshapeInner cs k shape) ia = do shape' <- reshapeInner shape 1 . arrayShape <$> lookupType ia certifying cs . letExp "reshape_inner" . BasicOp $ - Futhark.Reshape shape' ia + Futhark.Reshape k shape' ia applyTransforms :: MonadBuilder m => ArrayTransforms -> VName -> m VName applyTransforms (ArrayTransforms ts) a = foldlM (flip applyTransform) a ts @@ -337,14 +335,14 @@ arrayOfShape t shape transformType t (Rearrange _ perm) = rearrangeType perm t - transformType t (Reshape _ shape) = - t `setArrayShape` newShape shape - transformType t (ReshapeOuter _ shape) = + transformType t (Reshape _ _ shape) = + t `setArrayShape` shape + transformType t (ReshapeOuter _ _ shape) = let Shape oldshape = arrayShape t - in t `setArrayShape` Shape (newDims shape ++ drop 1 oldshape) - transformType t (ReshapeInner _ shape) = + in t `setArrayShape` Shape (shapeDims shape ++ drop 1 oldshape) + transformType t (ReshapeInner _ _ shape) = let Shape oldshape = arrayShape t - in t `setArrayShape` Shape (take 1 oldshape ++ newDims shape) + in t `setArrayShape` Shape (take 1 oldshape ++ shapeDims shape) -- | Return the row type of an input. Just a convenient alias. inputRowType :: Input -> Type @@ -362,8 +360,8 @@ where transformRows' inp (Rearrange cs perm) = addTransform (Rearrange cs (0 : map (+ 1) perm)) inp - transformRows' inp (Reshape cs shape) = - addTransform (ReshapeInner cs shape) inp + transformRows' inp (Reshape cs k shape) = + addTransform (ReshapeInner cs k shape) inp transformRows' inp (Replicate cs n) | inputRank inp == 1 = Rearrange mempty [1, 0] @@ -396,12 +394,18 @@ where f e (Rearrange cs perm) = text "rearrange" <> ppr cs <> PP.apply [PP.apply (map ppr perm), e] - f e (Reshape cs shape) = - text "reshape" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] - f e (ReshapeOuter cs shape) = - text "reshape_outer" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] - f e (ReshapeInner cs shape) = - text "reshape_inner" <> ppr cs <> PP.apply [PP.apply (map ppr shape), e] + f e (Reshape cs ReshapeArbitrary shape) = + text "reshape" <> ppr cs <> PP.apply [ppr shape, e] + f e (ReshapeOuter cs ReshapeArbitrary shape) = + text "reshape_outer" <> ppr cs <> PP.apply [ppr shape, e] + f e (ReshapeInner cs ReshapeArbitrary shape) = + text "reshape_inner" <> ppr cs <> PP.apply [ppr shape, e] + f e (Reshape cs ReshapeCoerce shape) = + text "coerce" <> ppr cs <> PP.apply [ppr shape, e] + f e (ReshapeOuter cs ReshapeCoerce shape) = + text "coerce_outer" <> ppr cs <> PP.apply [ppr shape, e] + f e (ReshapeInner cs ReshapeCoerce shape) = + text "coerce_inner" <> ppr cs <> PP.apply [ppr shape, e] f e (Replicate cs ne) = text "replicate" <> ppr cs <> PP.apply [ppr ne, e] @@ -526,12 +530,17 @@ -- Returns the Stream SOAC and the -- extra-accumulator body-result ident if any. soacToStream :: - (MonadFreshNames m, Buildable rep, Op rep ~ Futhark.SOAC rep) => + ( HasScope rep m, + MonadFreshNames m, + Buildable rep, + BuilderOps rep, + Op rep ~ Futhark.SOAC rep + ) => SOAC rep -> m (SOAC rep, [Ident]) soacToStream soac = do chunk_param <- newParam "chunk" $ Prim int64 - let chvar = Futhark.Var $ paramName chunk_param + let chvar = Var $ paramName chunk_param (lam, inps) = (lambda soac, inputs soac) w = width soac lam' <- renameLambda lam @@ -554,7 +563,7 @@ Futhark.Screma chvar (map paramName strm_inpids) $ Futhark.mapSOAC lam' insstm = mkLet strm_resids $ Op insoac - strmbdy = mkBody (oneStm insstm) $ map (subExpRes . Futhark.Var . identName) strm_resids + strmbdy = mkBody (oneStm insstm) $ map (subExpRes . Var . identName) strm_resids strmpar = chunk_param : strm_inpids strmlam = Lambda strmpar strmbdy loutps empty_lam = Lambda [] (mkBody mempty []) [] @@ -573,77 +582,51 @@ -- {acc', strm_resids, map_resids} -- the array and accumulator result types let scan_arr_ts = map (`arrayOfRow` chvar) $ lambdaReturnType scan_lam - map_arr_ts = drop (length nes) loutps accrtps = lambdaReturnType scan_lam - -- array result and input IDs of the stream's lambda - strm_resids <- mapM (newIdent "res") scan_arr_ts - scan0_ids <- mapM (newIdent "resarr0") scan_arr_ts - map_resids <- mapM (newIdent "map_res") map_arr_ts - - lastel_ids <- mapM (newIdent "lstel") accrtps - lastel_tmp_ids <- mapM (newIdent "lstel_tmp") accrtps - empty_arr <- newIdent "empty_arr" $ Prim Bool inpacc_ids <- mapM (newParam "inpacc") accrtps - outszm1id <- newIdent "szm1" $ Prim int64 - -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam,nes,map_lam,a_ch) - let insstm = - mkLet (scan0_ids ++ map_resids) . Op $ - Futhark.Screma chvar (map paramName strm_inpids) $ - Futhark.scanomapSOAC [Futhark.Scan scan_lam nes] lam' - -- 2. let outerszm1id = chunksize - 1 - outszm1stm = - mkLet [outszm1id] . BasicOp $ - BinOp - (Sub Int64 OverflowUndef) - (Futhark.Var $ paramName chunk_param) - (constant (1 :: Int64)) - -- 3. let lasteel_ids = ... - empty_arr_stm = - mkLet [empty_arr] . BasicOp $ - CmpOp - (CmpSlt Int64) - (Futhark.Var $ identName outszm1id) - (constant (0 :: Int64)) - leltmpstms = - zipWith - ( \lid arrid -> - mkLet [lid] . BasicOp $ - Index (identName arrid) $ - fullSlice - (identType arrid) - [DimFix $ Futhark.Var $ identName outszm1id] - ) - lastel_tmp_ids - scan0_ids - lelstm = - mkLet lastel_ids - $ If - (Futhark.Var $ identName empty_arr) - (mkBody mempty $ subExpsRes nes) - ( mkBody (stmsFromList leltmpstms) $ - varsRes $ - map identName lastel_tmp_ids - ) - $ ifCommon - $ map identType lastel_tmp_ids - -- 4. let strm_resids = map (acc `+`,nes, scan0_ids) - maplam <- mkMapPlusAccLam (map (Futhark.Var . paramName) inpacc_ids) scan_lam - let mapstm = - mkLet strm_resids . Op $ - Futhark.Screma chvar (map identName scan0_ids) (Futhark.mapSOAC maplam) - -- 5. let acc' = acc + lasteel_ids - addlelbdy <- - mkPlusBnds scan_lam $ - map Futhark.Var $ - map paramName inpacc_ids ++ map identName lastel_ids + maplam <- mkMapPlusAccLam (map (Var . paramName) inpacc_ids) scan_lam -- Finally, construct the stream - let (addlelstm, addlelres) = (bodyStms addlelbdy, bodyResult addlelbdy) - strmbdy = - mkBody (stmsFromList [insstm, outszm1stm, empty_arr_stm, lelstm, mapstm] <> addlelstm) $ - addlelres ++ map (subExpRes . Futhark.Var . identName) (strm_resids ++ map_resids) - strmpar = chunk_param : inpacc_ids ++ strm_inpids - strmlam = Lambda strmpar strmbdy (accrtps ++ loutps) + let strmpar = chunk_param : inpacc_ids ++ strm_inpids + strmlam <- fmap fst . runBuilder . mkLambda strmpar $ do + -- 1. let (scan0_ids,map_resids) = scanomap(scan_lam,nes,map_lam,a_ch) + (scan0_ids, map_resids) <- + fmap (splitAt (length scan_arr_ts)) . letTupExp "scan" . Op $ + Futhark.Screma chvar (map paramName strm_inpids) $ + Futhark.scanomapSOAC [Futhark.Scan scan_lam nes] lam' + -- 2. let outerszm1id = chunksize - 1 + outszm1id <- + letSubExp "outszm1" . BasicOp $ + BinOp + (Sub Int64 OverflowUndef) + (Var $ paramName chunk_param) + (constant (1 :: Int64)) + empty_arr <- + letExp "empty_arr" . BasicOp $ + CmpOp + (CmpSlt Int64) + outszm1id + (constant (0 :: Int64)) + -- 3. let lasteel_ids = ... + let indexLast arr = do + arr_t <- lookupType arr + pure $ BasicOp . Index arr $ fullSlice arr_t [DimFix outszm1id] + lastel_ids <- + letTupExp "lastel" + =<< eIf + (eSubExp $ Var empty_arr) + (resultBodyM nes) + (eBody $ map indexLast scan0_ids) + addlelbdy <- + mkPlusBnds scan_lam $ map Var $ map paramName inpacc_ids ++ lastel_ids + let (addlelstm, addlelres) = (bodyStms addlelbdy, bodyResult addlelbdy) + -- 4. let strm_resids = map (acc `+`,nes, scan0_ids) + strm_resids <- + letTupExp "strm_res" . Op $ + Futhark.Screma chvar scan0_ids (Futhark.mapSOAC maplam) + -- 5. let acc' = acc + lasteel_ids + addStms addlelstm + pure $ addlelres ++ map (subExpRes . Var) (strm_resids ++ map_resids) pure ( Stream w Sequential strmlam nes inps, map paramIdent inpacc_ids @@ -674,13 +657,13 @@ -- 2. let acc' = acc + acc0_ids in addaccbdy <- mkPlusBnds lamin $ - map Futhark.Var $ + map Var $ map paramName inpacc_ids ++ map identName acc0_ids -- Construct the stream let (addaccstm, addaccres) = (bodyStms addaccbdy, bodyResult addaccbdy) strmbdy = mkBody (oneStm insstm <> addaccstm) $ - addaccres ++ map (subExpRes . Futhark.Var . identName) strm_resids + addaccres ++ map (subExpRes . Var . identName) strm_resids strmpar = chunk_param : inpacc_ids ++ strm_inpids strmlam = Lambda strmpar strmbdy (accrtps ++ loutps') lam0 <- renameLambda lamin diff --git a/src/Futhark/Analysis/Interference.hs b/src/Futhark/Analysis/Interference.hs --- a/src/Futhark/Analysis/Interference.hs +++ b/src/Futhark/Analysis/Interference.hs @@ -117,10 +117,10 @@ m (InUse, LastUsed, Graph VName) analyseExp lumap inuse_outside expr = case expr of - If _ then_body else_body _ -> do - res1 <- analyseBody lumap inuse_outside then_body - res2 <- analyseBody lumap inuse_outside else_body - pure $ res1 <> res2 + Match _ cases defbody _ -> + fmap mconcat $ + mapM (analyseBody lumap inuse_outside) $ + defbody : map caseBody cases DoLoop merge _ body -> analyseLoopParams merge <$> analyseBody lumap inuse_outside body Op (Inner (SegOp segop)) -> do @@ -287,10 +287,8 @@ <$> mapM memSizesStm $ stmsToList $ kernelBodyStms body - memSizesExp (If _ then_body else_body _) = do - then_res <- memSizes $ bodyStms then_body - else_res <- memSizes $ bodyStms else_body - pure $ then_res <> else_res + memSizesExp (Match _ cases defbody _) = do + mconcat <$> mapM (memSizes . bodyStms) (defbody : map caseBody cases) memSizesExp (DoLoop _ _ body) = memSizes $ bodyStms body memSizesExp _ = pure mempty @@ -306,9 +304,8 @@ getSpacesStm (Let _ _ (Op (Alloc _ _))) = error "impossible" getSpacesStm (Let _ _ (Op (Inner (SegOp segop)))) = foldMap getSpacesStm $ kernelBodyStms $ segBody segop - getSpacesStm (Let _ _ (If _ then_body else_body _)) = - foldMap getSpacesStm (bodyStms then_body) - <> foldMap getSpacesStm (bodyStms else_body) + getSpacesStm (Let _ _ (Match _ cases defbody _)) = + foldMap (foldMap getSpacesStm . bodyStms) $ defbody : map caseBody cases getSpacesStm (Let _ _ (DoLoop _ _ body)) = foldMap getSpacesStm (bodyStms body) getSpacesStm _ = mempty @@ -327,11 +324,10 @@ m (InUse, LastUsed, Graph VName) helper stm@Let {stmExp = Op (Inner (SegOp segop))} = inScopeOf stm $ analyseSegOp lumap mempty segop - helper stm@Let {stmExp = If _ then_body else_body _} = - inScopeOf stm $ do - res1 <- analyseGPU' lumap (bodyStms then_body) - res2 <- analyseGPU' lumap (bodyStms else_body) - pure (res1 <> res2) + helper stm@Let {stmExp = Match _ cases defbody _} = + inScopeOf stm $ + mconcat + <$> mapM (analyseGPU' lumap . bodyStms) (defbody : map caseBody cases) helper stm@Let {stmExp = DoLoop merge _ body} = fmap (analyseLoopParams merge) . inScopeOf stm $ analyseGPU' lumap $ diff --git a/src/Futhark/Analysis/LastUse.hs b/src/Futhark/Analysis/LastUse.hs --- a/src/Futhark/Analysis/LastUse.hs +++ b/src/Futhark/Analysis/LastUse.hs @@ -15,7 +15,7 @@ where import Control.Monad.Reader -import Data.Bifunctor (first) +import Data.Bifunctor (bimap, first) import Data.Foldable import Data.Function ((&)) import Data.Map (Map) @@ -148,12 +148,17 @@ analyseExp (lumap, used) (Apply _ args _ _) = do let nms = freeIn $ map fst args pure (insertNames pat_name nms lumap, used <> nms) - analyseExp (lumap, used) (If cse then_body else_body dec) = do - (lumap_then, used_then) <- analyseBody lumap used then_body - (lumap_else, used_else) <- analyseBody lumap used else_body - let used' = used_then <> used_else - nms = (freeIn cse <> freeIn dec) `namesSubtract` used' - pure (insertNames pat_name nms (lumap_then <> lumap_else), used' <> nms) + analyseExp (lumap, used) (Match ses cases defbody dec) = do + (lumap_cases, used_cases) <- + bimap mconcat mconcat . unzip + <$> mapM (analyseBody lumap used . caseBody) cases + (lumap_defbody, used_defbody) <- analyseBody lumap used defbody + let used' = used_cases <> used_defbody + nms = (freeIn ses <> freeIn dec) `namesSubtract` used' + pure + ( insertNames pat_name nms (lumap_cases <> lumap_defbody), + used' <> nms + ) analyseExp (lumap, used) (DoLoop merge form body) = do (lumap', used') <- analyseBody lumap used body let nms = (freeIn merge <> freeIn form) `namesSubtract` used' diff --git a/src/Futhark/Analysis/MemAlias.hs b/src/Futhark/Analysis/MemAlias.hs --- a/src/Futhark/Analysis/MemAlias.hs +++ b/src/Futhark/Analysis/MemAlias.hs @@ -85,12 +85,10 @@ analyzeStm m (Let _ _ (Op (Inner inner))) = do on_inner <- asks onInner on_inner m inner -analyzeStm m (Let pat _ (If _ then_body else_body _)) = do - m' <- - analyzeStms (bodyStms then_body) m - >>= analyzeStms (bodyStms else_body) - zip (patNames pat) (map resSubExp $ bodyResult then_body) - <> zip (patNames pat) (map resSubExp $ bodyResult else_body) +analyzeStm m (Let pat _ (Match _ cases defbody _)) = do + let bodies = defbody : map caseBody cases + m' <- foldM (flip analyzeStms) m $ map bodyStms bodies + foldMap (zip (patNames pat) . map resSubExp . bodyResult) bodies & mapMaybe (filterFun m') & foldr (uncurry addAlias) m' & pure diff --git a/src/Futhark/Analysis/Metrics.hs b/src/Futhark/Analysis/Metrics.hs --- a/src/Futhark/Analysis/Metrics.hs +++ b/src/Futhark/Analysis/Metrics.hs @@ -107,10 +107,15 @@ inside "DoLoop" $ seen "ForLoop" >> bodyMetrics body expMetrics (DoLoop _ WhileLoop {} body) = inside "DoLoop" $ seen "WhileLoop" >> bodyMetrics body -expMetrics (If _ tb fb _) = +expMetrics (Match _ [Case [Just (BoolValue True)] tb] fb _) = inside "If" $ do inside "True" $ bodyMetrics tb inside "False" $ bodyMetrics fb +expMetrics (Match _ cases defbody _) = + inside "Match" $ do + forM_ (zip [0 ..] cases) $ \(i, c) -> + inside (T.pack (show (i :: Int))) $ bodyMetrics $ caseBody c + inside "default" $ bodyMetrics defbody expMetrics Apply {} = seen "Apply" expMetrics (WithAcc _ lam) = diff --git a/src/Futhark/Analysis/PrimExp.hs b/src/Futhark/Analysis/PrimExp.hs --- a/src/Futhark/Analysis/PrimExp.hs +++ b/src/Futhark/Analysis/PrimExp.hs @@ -67,6 +67,7 @@ (~/~), (~+~), (~-~), + (~==~), ) where @@ -754,6 +755,14 @@ Bool -> LogOr Unit -> LogOr +-- | Equality of untyped 'PrimExp's, which must have the same type. +(~==~) :: PrimExp v -> PrimExp v -> PrimExp v +x ~==~ y = CmpOpExp (CmpEq t) x y + where + t = primExpType x + infix 7 ~*~, ~/~ infix 6 ~+~, ~-~ + +infix 4 ~==~ diff --git a/src/Futhark/Analysis/PrimExp/Generalize.hs b/src/Futhark/Analysis/PrimExp/Generalize.hs deleted file mode 100644 --- a/src/Futhark/Analysis/PrimExp/Generalize.hs +++ /dev/null @@ -1,73 +0,0 @@ --- | Generalization (anti-unification) of 'PrimExp's. -module Futhark.Analysis.PrimExp.Generalize - ( leastGeneralGeneralization, - ) -where - -import Data.List (elemIndex) -import Futhark.Analysis.PrimExp -import Futhark.IR.Syntax.Core (Ext (..)) - --- | Generalize two 'PrimExp's of the the same type. -leastGeneralGeneralization :: - (Eq v) => - [(PrimExp v, PrimExp v)] -> - PrimExp v -> - PrimExp v -> - (PrimExp (Ext v), [(PrimExp v, PrimExp v)]) -leastGeneralGeneralization m exp1@(LeafExp v1 t1) exp2@(LeafExp v2 _) = - if v1 == v2 - then (LeafExp (Free v1) t1, m) - else generalize m exp1 exp2 -leastGeneralGeneralization m exp1@(ValueExp v1) exp2@(ValueExp v2) = - if v1 == v2 - then (ValueExp v1, m) - else generalize m exp1 exp2 -leastGeneralGeneralization m exp1@(BinOpExp op1 e11 e12) exp2@(BinOpExp op2 e21 e22) = - if op1 == op2 - then - let (e1, m1) = leastGeneralGeneralization m e11 e21 - (e2, m2) = leastGeneralGeneralization m1 e12 e22 - in (BinOpExp op1 e1 e2, m2) - else generalize m exp1 exp2 -leastGeneralGeneralization m exp1@(CmpOpExp op1 e11 e12) exp2@(CmpOpExp op2 e21 e22) = - if op1 == op2 - then - let (e1, m1) = leastGeneralGeneralization m e11 e21 - (e2, m2) = leastGeneralGeneralization m1 e12 e22 - in (CmpOpExp op1 e1 e2, m2) - else generalize m exp1 exp2 -leastGeneralGeneralization m exp1@(UnOpExp op1 e1) exp2@(UnOpExp op2 e2) = - if op1 == op2 - then - let (e, m1) = leastGeneralGeneralization m e1 e2 - in (UnOpExp op1 e, m1) - else generalize m exp1 exp2 -leastGeneralGeneralization m exp1@(ConvOpExp op1 e1) exp2@(ConvOpExp op2 e2) = - if op1 == op2 - then - let (e, m1) = leastGeneralGeneralization m e1 e2 - in (ConvOpExp op1 e, m1) - else generalize m exp1 exp2 -leastGeneralGeneralization m exp1@(FunExp s1 args1 t1) exp2@(FunExp s2 args2 _) = - if s1 == s2 && length args1 == length args2 - then - let (args, m') = - foldl - ( \(arg_acc, m_acc) (a1, a2) -> - let (a, m'') = leastGeneralGeneralization m_acc a1 a2 - in (a : arg_acc, m'') - ) - ([], m) - (zip args1 args2) - in (FunExp s1 (reverse args) t1, m') - else generalize m exp1 exp2 -leastGeneralGeneralization m exp1 exp2 = - generalize m exp1 exp2 - -generalize :: Eq v => [(PrimExp v, PrimExp v)] -> PrimExp v -> PrimExp v -> (PrimExp (Ext v), [(PrimExp v, PrimExp v)]) -generalize m exp1 exp2 = - let t = primExpType exp1 - in case elemIndex (exp1, exp2) m of - Just i -> (LeafExp (Ext i) t, m) - Nothing -> (LeafExp (Ext $ length m) t, m ++ [(exp1, exp2)]) diff --git a/src/Futhark/Analysis/SymbolTable.hs b/src/Futhark/Analysis/SymbolTable.hs --- a/src/Futhark/Analysis/SymbolTable.hs +++ b/src/Futhark/Analysis/SymbolTable.hs @@ -352,12 +352,13 @@ indexExp table (BasicOp (Replicate (Shape [_]) (Var v))) _ (_ : is) = do guard $ v `available` table index' v is table -indexExp table (BasicOp (Reshape newshape v)) _ is +indexExp table (BasicOp (Reshape _ newshape v)) _ is | Just oldshape <- arrayDims <$> lookupType v table = + -- TODO: handle coercions more efficiently. let is' = reshapeIndex (map pe64 oldshape) - (map pe64 $ newDims newshape) + (map pe64 $ shapeDims newshape) is in index' v is' table indexExp table (BasicOp (Index v slice)) _ is = do diff --git a/src/Futhark/Analysis/UsageTable.hs b/src/Futhark/Analysis/UsageTable.hs --- a/src/Futhark/Analysis/UsageTable.hs +++ b/src/Futhark/Analysis/UsageTable.hs @@ -179,8 +179,8 @@ ] usageInExp e@DoLoop {} = foldMap consumedUsage $ namesToList $ consumedInExp e -usageInExp (If _ tbranch fbranch _) = - usageInBody tbranch <> usageInBody fbranch +usageInExp (Match _ cases defbody _) = + foldMap (usageInBody . caseBody) cases <> usageInBody defbody usageInExp (WithAcc inputs lam) = foldMap inputUsage inputs <> usageInBody (lambdaBody lam) where diff --git a/src/Futhark/Builder.hs b/src/Futhark/Builder.hs --- a/src/Futhark/Builder.hs +++ b/src/Futhark/Builder.hs @@ -123,7 +123,7 @@ pure x instance - (ASTRep rep, MonadFreshNames m, BuilderOps rep) => + (MonadFreshNames m, BuilderOps rep) => MonadBuilder (BuilderT rep m) where type Rep (BuilderT rep m) = rep diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs --- a/src/Futhark/CLI/Dev.hs +++ b/src/Futhark/CLI/Dev.hs @@ -430,28 +430,28 @@ "Run this compiler backend on pipeline result.", Option [] - ["compile-imperative"] + ["compile-imp-seq"] ( NoArg $ Right $ \opts -> opts {futharkAction = SeqMemAction $ \_ _ _ -> impCodeGenAction} ) - "Translate program into the imperative IL and write it on standard output.", + "Translate pipeline result to ImpSequential and write it on stdout.", Option [] - ["compile-imperative-kernels"] + ["compile-imp-gpu"] ( NoArg $ Right $ \opts -> opts {futharkAction = GPUMemAction $ \_ _ _ -> kernelImpCodeGenAction} ) - "Translate program into the imperative IL with kernels and write it on standard output.", + "Translate pipeline result to ImpGPU and write it on stdout.", Option [] - ["compile-imperative-multicore"] + ["compile-imp-multicore"] ( NoArg $ Right $ \opts -> opts {futharkAction = MCMemAction $ \_ _ _ -> multicoreImpCodeGenAction} ) - "Translate program into the imperative IL with kernels and write it on standard output.", + "Translate pipeline result to ImpMC write it on stdout.", Option "p" ["print"] diff --git a/src/Futhark/CLI/Main.hs b/src/Futhark/CLI/Main.hs --- a/src/Futhark/CLI/Main.hs +++ b/src/Futhark/CLI/Main.hs @@ -78,7 +78,7 @@ ("literate", (Literate.main, "Process a literate Futhark program.")), ("lsp", (LSP.main, "Run LSP server.")), ("thanks", (Misc.mainThanks, "Express gratitude.")), - ("tokens", (Misc.mainTokens, "Express gratitude.")) + ("tokens", (Misc.mainTokens, "Print tokens from Futhark file.")) ] msg :: String diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs --- a/src/Futhark/CLI/Test.hs +++ b/src/Futhark/CLI/Test.hs @@ -497,19 +497,6 @@ moveCursorToTableTop :: IO () moveCursorToTableTop = cursorUpLine tableLines -reportText :: TestStatus -> IO () -reportText ts = - putStr $ - "(" - ++ show (testStatusFail ts) - ++ " failed, " - ++ show (testStatusPass ts) - ++ " passed, " - ++ show num_remain - ++ " to go).\n" - where - num_remain = length $ testStatusRemain ts - runTests :: TestConfig -> [FilePath] -> IO () runTests config paths = do -- We force line buffering to ensure that we produce running output. @@ -533,10 +520,10 @@ report | fancy = reportTable - | otherwise = reportText + | otherwise = const (pure ()) clear | fancy = clearFromCursorToScreenEnd - | otherwise = putStr "\n" + | otherwise = pure () numTestCases tc = case testAction $ testCaseTest tc of @@ -552,10 +539,7 @@ report ts msg <- takeMVar reportmvar case msg of - TestStarted test -> do - unless fancy $ - putStr $ - "Started testing " <> testCaseProgram test <> " " + TestStarted test -> getResults $ ts {testStatusRun = test : testStatusRun ts} TestDone test res -> do let ts' = @@ -573,9 +557,6 @@ { testStatusRunPass = testStatusRunPass ts' + numTestCases test } - unless fancy $ - putStr $ - "Finished testing " <> testCaseProgram test <> " " getResults $ ts'' {testStatusPass = testStatusPass ts + 1} Failure s -> do when fancy moveCursorToTableTop diff --git a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs @@ -336,6 +336,7 @@ int logging; typename lock_t lock; char *error; + typename lock_t error_lock; typename FILE *log; $sdecls:fields $sdecls:kernel_fields @@ -374,6 +375,7 @@ ctx->profiling_paused = 0; ctx->logging = cfg->cu_cfg.logging; ctx->error = NULL; + create_lock(&ctx->error_lock); ctx->log = stderr; ctx->cuda.profiling_records_capacity = 200; ctx->cuda.profiling_records_used = 0; diff --git a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs @@ -319,6 +319,7 @@ int logging; typename lock_t lock; char *error; + typename lock_t error_lock; typename FILE *log; $sdecls:fields $sdecls:ctx_opencl_fields @@ -342,6 +343,7 @@ ctx->profiling_paused = 0; ctx->logging = cfg->opencl.logging; ctx->error = NULL; + create_lock(&ctx->error_lock); ctx->log = stderr; ctx->opencl.profiling_records_capacity = 200; ctx->opencl.profiling_records_used = 0; diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -5,2539 +5,746 @@ {-# LANGUAGE Trustworthy #-} {-# LANGUAGE TupleSections #-} --- | C code generator framework. -module Futhark.CodeGen.Backends.GenericC - ( compileProg, - compileProg', - CParts (..), - asLibrary, - asExecutable, - asServer, - - -- * Pluggable compiler - Operations (..), - defaultOperations, - OpCompiler, - ErrorCompiler, - CallCompiler, - PointerQuals, - MemoryType, - WriteScalar, - writeScalarPointerWithQuals, - ReadScalar, - readScalarPointerWithQuals, - Allocate, - Deallocate, - CopyBarrier (..), - Copy, - StaticArray, - - -- * Monadic compiler interface - CompilerM, - CompilerState (compUserState, compNameSrc, compDeclaredMem), - CompilerEnv (envCachedMem), - getUserState, - modifyUserState, - contextContents, - contextFinalInits, - runCompilerM, - inNewFunction, - cachingMemory, - compileFun, - compileCode, - compileExp, - compilePrimExp, - compileExpToName, - rawMem, - item, - items, - stm, - stms, - decl, - atInit, - headerDecl, - publicDef, - publicDef_, - profileReport, - onClear, - HeaderSection (..), - libDecl, - earlyDecl, - publicName, - contextField, - contextFieldDyn, - memToCType, - cacheMem, - fatMemory, - rawMemCType, - cproduct, - fatMemType, - declAllocatedMem, - freeAllocatedMem, - collect, - - -- * Building Blocks - primTypeToCType, - intTypeToCType, - copyMemoryDefaultSpace, - linearCode, - derefPointer, - fatMemAlloc, - fatMemSet, - fatMemUnRef, - errorMsgString, - ) -where - -import Control.Monad.Identity -import Control.Monad.Reader -import Control.Monad.State -import Data.Bifunctor (first) -import Data.Char (isAlpha, isAlphaNum, isDigit) -import qualified Data.DList as DL -import Data.List (unzip4) -import Data.Loc -import qualified Data.Map.Strict as M -import Data.Maybe -import qualified Data.Text as T -import Futhark.CodeGen.Backends.GenericC.CLI (cliDefs) -import Futhark.CodeGen.Backends.GenericC.Options -import Futhark.CodeGen.Backends.GenericC.Server (serverDefs) -import Futhark.CodeGen.Backends.SimpleRep -import Futhark.CodeGen.ImpCode -import Futhark.CodeGen.RTS.C (cacheH, errorsH, halfH, lockH, timingH, utilH) -import Futhark.IR.Prop (isBuiltInFunction) -import qualified Futhark.Manifest as Manifest -import Futhark.MonadFreshNames -import Futhark.Util (chunks, mapAccumLM, zEncodeString) -import Futhark.Util.Pretty (prettyText) -import qualified Language.C.Quote.OpenCL as C -import qualified Language.C.Syntax as C -import NeatInterpolation (untrimming) - --- How public an array type definition sould be. Public types show up --- in the generated API, while private types are used only to --- implement the members of opaques. -data Publicness = Private | Public - deriving (Eq, Ord, Show) - -type ArrayType = (Signedness, PrimType, Int) - -data CompilerState s = CompilerState - { compArrayTypes :: M.Map ArrayType Publicness, - compEarlyDecls :: DL.DList C.Definition, - compInit :: [C.Stm], - compNameSrc :: VNameSource, - compUserState :: s, - compHeaderDecls :: M.Map HeaderSection (DL.DList C.Definition), - compLibDecls :: DL.DList C.Definition, - compCtxFields :: DL.DList (C.Id, C.Type, Maybe C.Exp, Maybe C.Stm), - compProfileItems :: DL.DList C.BlockItem, - compClearItems :: DL.DList C.BlockItem, - compDeclaredMem :: [(VName, Space)], - compItems :: DL.DList C.BlockItem - } - -newCompilerState :: VNameSource -> s -> CompilerState s -newCompilerState src s = - CompilerState - { compArrayTypes = mempty, - compEarlyDecls = mempty, - compInit = [], - compNameSrc = src, - compUserState = s, - compHeaderDecls = mempty, - compLibDecls = mempty, - compCtxFields = mempty, - compProfileItems = mempty, - compClearItems = mempty, - compDeclaredMem = mempty, - compItems = mempty - } - --- | In which part of the header file we put the declaration. This is --- to ensure that the header file remains structured and readable. -data HeaderSection - = ArrayDecl String - | OpaqueTypeDecl String - | OpaqueDecl String - | EntryDecl - | MiscDecl - | InitDecl - deriving (Eq, Ord) - --- | A substitute expression compiler, tried before the main --- compilation function. -type OpCompiler op s = op -> CompilerM op s () - -type ErrorCompiler op s = ErrorMsg Exp -> String -> CompilerM op s () - --- | The address space qualifiers for a pointer of the given type with --- the given annotation. -type PointerQuals op s = String -> CompilerM op s [C.TypeQual] - --- | The type of a memory block in the given memory space. -type MemoryType op s = SpaceId -> CompilerM op s C.Type - --- | Write a scalar to the given memory block with the given element --- index and in the given memory space. -type WriteScalar op s = - C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s () - --- | Read a scalar from the given memory block with the given element --- index and in the given memory space. -type ReadScalar op s = - C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp - --- | Allocate a memory block of the given size and with the given tag --- in the given memory space, saving a reference in the given variable --- name. -type Allocate op s = - C.Exp -> - C.Exp -> - C.Exp -> - SpaceId -> - CompilerM op s () - --- | De-allocate the given memory block with the given tag, which is --- in the given memory space. -type Deallocate op s = C.Exp -> C.Exp -> SpaceId -> CompilerM op s () - --- | Create a static array of values - initialised at load time. -type StaticArray op s = VName -> SpaceId -> PrimType -> ArrayContents -> CompilerM op s () - --- | Whether a copying operation should implicitly function as a --- barrier regarding further operations on the source. This is a --- rather subtle detail and is mostly useful for letting some --- device/GPU copies be asynchronous (#1664). -data CopyBarrier - = CopyBarrier - | -- | Explicit context synchronisation should be done - -- before the source or target is used. - CopyNoBarrier - deriving (Eq, Show) - --- | Copy from one memory block to another. -type Copy op s = - CopyBarrier -> - C.Exp -> - C.Exp -> - Space -> - C.Exp -> - C.Exp -> - Space -> - C.Exp -> - CompilerM op s () - --- | Call a function. -type CallCompiler op s = [VName] -> Name -> [C.Exp] -> CompilerM op s () - -data Operations op s = Operations - { opsWriteScalar :: WriteScalar op s, - opsReadScalar :: ReadScalar op s, - opsAllocate :: Allocate op s, - opsDeallocate :: Deallocate op s, - opsCopy :: Copy op s, - opsStaticArray :: StaticArray op s, - opsMemoryType :: MemoryType op s, - opsCompiler :: OpCompiler op s, - opsError :: ErrorCompiler op s, - opsCall :: CallCompiler op s, - -- | If true, use reference counting. Otherwise, bare - -- pointers. - opsFatMemory :: Bool, - -- | Code to bracket critical sections. - opsCritical :: ([C.BlockItem], [C.BlockItem]) - } - -errorMsgString :: ErrorMsg Exp -> CompilerM op s (String, [C.Exp]) -errorMsgString (ErrorMsg parts) = do - let boolStr e = [C.cexp|($exp:e) ? "true" : "false"|] - asLongLong e = [C.cexp|(long long int)$exp:e|] - asDouble e = [C.cexp|(double)$exp:e|] - onPart (ErrorString s) = pure ("%s", [C.cexp|$string:s|]) - onPart (ErrorVal Bool x) = ("%s",) . boolStr <$> compileExp x - onPart (ErrorVal Unit _) = pure ("%s", [C.cexp|"()"|]) - onPart (ErrorVal (IntType Int8) x) = ("%hhd",) <$> compileExp x - onPart (ErrorVal (IntType Int16) x) = ("%hd",) <$> compileExp x - onPart (ErrorVal (IntType Int32) x) = ("%d",) <$> compileExp x - onPart (ErrorVal (IntType Int64) x) = ("%lld",) . asLongLong <$> compileExp x - onPart (ErrorVal (FloatType Float16) x) = ("%f",) . asDouble <$> compileExp x - onPart (ErrorVal (FloatType Float32) x) = ("%f",) . asDouble <$> compileExp x - onPart (ErrorVal (FloatType Float64) x) = ("%f",) <$> compileExp x - (formatstrs, formatargs) <- unzip <$> mapM onPart parts - pure (mconcat formatstrs, formatargs) - -freeAllocatedMem :: CompilerM op s [C.BlockItem] -freeAllocatedMem = collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem - -declAllocatedMem :: CompilerM op s [C.BlockItem] -declAllocatedMem = collect $ mapM_ f =<< gets compDeclaredMem - where - f (name, space) = do - ty <- memToCType name space - decl [C.cdecl|$ty:ty $id:name;|] - resetMem name space - -defError :: ErrorCompiler op s -defError msg stacktrace = do - (formatstr, formatargs) <- errorMsgString msg - let formatstr' = "Error: " <> formatstr <> "\n\nBacktrace:\n%s" - items - [C.citems|if (ctx->error == NULL) - ctx->error = msgprintf($string:formatstr', $args:formatargs, $string:stacktrace); - err = FUTHARK_PROGRAM_ERROR; - goto cleanup;|] - -defCall :: CallCompiler op s -defCall dests fname args = do - let out_args = [[C.cexp|&$id:d|] | d <- dests] - args' - | isBuiltInFunction fname = args - | otherwise = [C.cexp|ctx|] : out_args ++ args - case dests of - [dest] - | isBuiltInFunction fname -> - stm [C.cstm|$id:dest = $id:(funName fname)($args:args');|] - _ -> - item [C.citem|if ($id:(funName fname)($args:args') != 0) { err = 1; goto cleanup; }|] - --- | A set of operations that fail for every operation involving --- non-default memory spaces. Uses plain pointers and @malloc@ for --- memory management. -defaultOperations :: Operations op s -defaultOperations = - Operations - { opsWriteScalar = defWriteScalar, - opsReadScalar = defReadScalar, - opsAllocate = defAllocate, - opsDeallocate = defDeallocate, - opsCopy = defCopy, - opsStaticArray = defStaticArray, - opsMemoryType = defMemoryType, - opsCompiler = defCompiler, - opsFatMemory = True, - opsError = defError, - opsCall = defCall, - opsCritical = mempty - } - where - defWriteScalar _ _ _ _ _ = - error "Cannot write to non-default memory space because I am dumb" - defReadScalar _ _ _ _ = - error "Cannot read from non-default memory space" - defAllocate _ _ _ = - error "Cannot allocate in non-default memory space" - defDeallocate _ _ = - error "Cannot deallocate in non-default memory space" - defCopy _ destmem destoffset DefaultSpace srcmem srcoffset DefaultSpace size = - copyMemoryDefaultSpace destmem destoffset srcmem srcoffset size - defCopy _ _ _ _ _ _ _ _ = - error "Cannot copy to or from non-default memory space" - defStaticArray _ _ _ _ = - error "Cannot create static array in non-default memory space" - defMemoryType _ = - error "Has no type for non-default memory space" - defCompiler _ = - error "The default compiler cannot compile extended operations" - -data CompilerEnv op s = CompilerEnv - { envOperations :: Operations op s, - -- | Mapping memory blocks to sizes. These memory blocks are CPU - -- memory that we know are used in particularly simple ways (no - -- reference counting necessary). To cut down on allocator - -- pressure, we keep these allocations around for a long time, and - -- record their sizes so we can reuse them if possible (and - -- realloc() when needed). - envCachedMem :: M.Map C.Exp VName - } - -envOpCompiler :: CompilerEnv op s -> OpCompiler op s -envOpCompiler = opsCompiler . envOperations - -envMemoryType :: CompilerEnv op s -> MemoryType op s -envMemoryType = opsMemoryType . envOperations - -envReadScalar :: CompilerEnv op s -> ReadScalar op s -envReadScalar = opsReadScalar . envOperations - -envWriteScalar :: CompilerEnv op s -> WriteScalar op s -envWriteScalar = opsWriteScalar . envOperations - -envAllocate :: CompilerEnv op s -> Allocate op s -envAllocate = opsAllocate . envOperations - -envDeallocate :: CompilerEnv op s -> Deallocate op s -envDeallocate = opsDeallocate . envOperations - -envCopy :: CompilerEnv op s -> Copy op s -envCopy = opsCopy . envOperations - -envStaticArray :: CompilerEnv op s -> StaticArray op s -envStaticArray = opsStaticArray . envOperations - -envFatMemory :: CompilerEnv op s -> Bool -envFatMemory = opsFatMemory . envOperations - -declsCode :: (HeaderSection -> Bool) -> CompilerState s -> T.Text -declsCode p = - T.unlines - . map prettyText - . concatMap (DL.toList . snd) - . filter (p . fst) - . M.toList - . compHeaderDecls - -initDecls, arrayDecls, opaqueDecls, opaqueTypeDecls, entryDecls, miscDecls :: CompilerState s -> T.Text -initDecls = declsCode (== InitDecl) -arrayDecls = declsCode isArrayDecl - where - isArrayDecl ArrayDecl {} = True - isArrayDecl _ = False -opaqueTypeDecls = declsCode isOpaqueTypeDecl - where - isOpaqueTypeDecl OpaqueTypeDecl {} = True - isOpaqueTypeDecl _ = False -opaqueDecls = declsCode isOpaqueDecl - where - isOpaqueDecl OpaqueDecl {} = True - isOpaqueDecl _ = False -entryDecls = declsCode (== EntryDecl) -miscDecls = declsCode (== MiscDecl) - -contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm], [C.Stm]) -contextContents = do - (field_names, field_types, field_values, field_frees) <- - gets $ unzip4 . DL.toList . compCtxFields - let fields = - [ [C.csdecl|$ty:ty $id:name;|] - | (name, ty) <- zip field_names field_types - ] - init_fields = - [ [C.cstm|ctx->$id:name = $exp:e;|] - | (name, Just e) <- zip field_names field_values - ] - pure (fields, init_fields, catMaybes field_frees) - -contextFinalInits :: CompilerM op s [C.Stm] -contextFinalInits = gets compInit - -newtype CompilerM op s a - = CompilerM (ReaderT (CompilerEnv op s) (State (CompilerState s)) a) - deriving - ( Functor, - Applicative, - Monad, - MonadState (CompilerState s), - MonadReader (CompilerEnv op s) - ) - -instance MonadFreshNames (CompilerM op s) where - getNameSource = gets compNameSrc - putNameSource src = modify $ \s -> s {compNameSrc = src} - -runCompilerM :: - Operations op s -> - VNameSource -> - s -> - CompilerM op s a -> - (a, CompilerState s) -runCompilerM ops src userstate (CompilerM m) = - runState - (runReaderT m (CompilerEnv ops mempty)) - (newCompilerState src userstate) - -getUserState :: CompilerM op s s -getUserState = gets compUserState - -modifyUserState :: (s -> s) -> CompilerM op s () -modifyUserState f = modify $ \compstate -> - compstate {compUserState = f $ compUserState compstate} - -atInit :: C.Stm -> CompilerM op s () -atInit x = modify $ \s -> - s {compInit = compInit s ++ [x]} - -collect :: CompilerM op s () -> CompilerM op s [C.BlockItem] -collect m = snd <$> collect' m - -collect' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) -collect' m = do - old <- gets compItems - modify $ \s -> s {compItems = mempty} - x <- m - new <- gets compItems - modify $ \s -> s {compItems = old} - pure (x, DL.toList new) - --- | Used when we, inside an existing 'CompilerM' action, want to --- generate code for a new function. Use this so that the compiler --- understands that previously declared memory doesn't need to be --- freed inside this action. -inNewFunction :: CompilerM op s a -> CompilerM op s a -inNewFunction m = do - old_mem <- gets compDeclaredMem - modify $ \s -> s {compDeclaredMem = mempty} - x <- local noCached m - modify $ \s -> s {compDeclaredMem = old_mem} - pure x - where - noCached env = env {envCachedMem = mempty} - -item :: C.BlockItem -> CompilerM op s () -item x = modify $ \s -> s {compItems = DL.snoc (compItems s) x} - -items :: [C.BlockItem] -> CompilerM op s () -items xs = modify $ \s -> s {compItems = DL.append (compItems s) (DL.fromList xs)} - -fatMemory :: Space -> CompilerM op s Bool -fatMemory ScalarSpace {} = pure False -fatMemory _ = asks envFatMemory - -cacheMem :: C.ToExp a => a -> CompilerM op s (Maybe VName) -cacheMem a = asks $ M.lookup (C.toExp a noLoc) . envCachedMem - --- | Construct a publicly visible definition using the specified name --- as the template. The first returned definition is put in the --- header file, and the second is the implementation. Returns the public --- name. -publicDef :: - String -> - HeaderSection -> - (String -> (C.Definition, C.Definition)) -> - CompilerM op s String -publicDef s h f = do - s' <- publicName s - let (pub, priv) = f s' - headerDecl h pub - earlyDecl priv - pure s' - --- | As 'publicDef', but ignores the public name. -publicDef_ :: - String -> - HeaderSection -> - (String -> (C.Definition, C.Definition)) -> - CompilerM op s () -publicDef_ s h f = void $ publicDef s h f - -headerDecl :: HeaderSection -> C.Definition -> CompilerM op s () -headerDecl sec def = modify $ \s -> - s - { compHeaderDecls = - M.unionWith - (<>) - (compHeaderDecls s) - (M.singleton sec (DL.singleton def)) - } - -libDecl :: C.Definition -> CompilerM op s () -libDecl def = modify $ \s -> - s {compLibDecls = compLibDecls s <> DL.singleton def} - -earlyDecl :: C.Definition -> CompilerM op s () -earlyDecl def = modify $ \s -> - s {compEarlyDecls = compEarlyDecls s <> DL.singleton def} - -contextField :: C.Id -> C.Type -> Maybe C.Exp -> CompilerM op s () -contextField name ty initial = modify $ \s -> - s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial, Nothing)} - -contextFieldDyn :: C.Id -> C.Type -> Maybe C.Exp -> C.Stm -> CompilerM op s () -contextFieldDyn name ty initial free = modify $ \s -> - s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial, Just free)} - -profileReport :: C.BlockItem -> CompilerM op s () -profileReport x = modify $ \s -> - s {compProfileItems = compProfileItems s <> DL.singleton x} - -onClear :: C.BlockItem -> CompilerM op s () -onClear x = modify $ \s -> - s {compClearItems = compClearItems s <> DL.singleton x} - -stm :: C.Stm -> CompilerM op s () -stm s = item [C.citem|$stm:s|] - -stms :: [C.Stm] -> CompilerM op s () -stms = mapM_ stm - -decl :: C.InitGroup -> CompilerM op s () -decl x = item [C.citem|$decl:x;|] - --- | Public names must have a consitent prefix. -publicName :: String -> CompilerM op s String -publicName s = pure $ "futhark_" ++ s - --- | The generated code must define a context struct with this name. -contextType :: CompilerM op s C.Type -contextType = do - name <- publicName "context" - pure [C.cty|struct $id:name|] - --- | The generated code must define a configuration struct with this --- name. -configType :: CompilerM op s C.Type -configType = do - name <- publicName "context_config" - pure [C.cty|struct $id:name|] - -memToCType :: VName -> Space -> CompilerM op s C.Type -memToCType v space = do - refcount <- fatMemory space - cached <- isJust <$> cacheMem v - if refcount && not cached - then pure $ fatMemType space - else rawMemCType space - -rawMemCType :: Space -> CompilerM op s C.Type -rawMemCType DefaultSpace = pure defaultMemBlockType -rawMemCType (Space sid) = join $ asks envMemoryType <*> pure sid -rawMemCType (ScalarSpace [] t) = - pure [C.cty|$ty:(primTypeToCType t)[1]|] -rawMemCType (ScalarSpace ds t) = - pure [C.cty|$ty:(primTypeToCType t)[$exp:(cproduct ds')]|] - where - ds' = map (`C.toExp` noLoc) ds - -fatMemType :: Space -> C.Type -fatMemType space = - [C.cty|struct $id:name|] - where - name = case space of - Space sid -> "memblock_" ++ sid - _ -> "memblock" - -fatMemSet :: Space -> String -fatMemSet (Space sid) = "memblock_set_" ++ sid -fatMemSet _ = "memblock_set" - -fatMemAlloc :: Space -> String -fatMemAlloc (Space sid) = "memblock_alloc_" ++ sid -fatMemAlloc _ = "memblock_alloc" - -fatMemUnRef :: Space -> String -fatMemUnRef (Space sid) = "memblock_unref_" ++ sid -fatMemUnRef _ = "memblock_unref" - -rawMem :: VName -> CompilerM op s C.Exp -rawMem v = rawMem' <$> fat <*> pure v - where - fat = asks ((&&) . envFatMemory) <*> (isNothing <$> cacheMem v) - -rawMem' :: C.ToExp a => Bool -> a -> C.Exp -rawMem' True e = [C.cexp|$exp:e.mem|] -rawMem' False e = [C.cexp|$exp:e|] - -allocRawMem :: - (C.ToExp a, C.ToExp b, C.ToExp c) => - a -> - b -> - Space -> - c -> - CompilerM op s () -allocRawMem dest size space desc = case space of - Space sid -> - join $ - asks envAllocate - <*> pure [C.cexp|$exp:dest|] - <*> pure [C.cexp|$exp:size|] - <*> pure [C.cexp|$exp:desc|] - <*> pure sid - _ -> - stm [C.cstm|$exp:dest = (unsigned char*) malloc((size_t)$exp:size);|] - -freeRawMem :: - (C.ToExp a, C.ToExp b) => - a -> - Space -> - b -> - CompilerM op s () -freeRawMem mem space desc = - case space of - Space sid -> do - free_mem <- asks envDeallocate - free_mem [C.cexp|$exp:mem|] [C.cexp|$exp:desc|] sid - _ -> item [C.citem|free($exp:mem);|] - -defineMemorySpace :: Space -> CompilerM op s (C.Definition, [C.Definition], C.BlockItem) -defineMemorySpace space = do - rm <- rawMemCType space - let structdef = - [C.cedecl|struct $id:sname { int *references; - $ty:rm mem; - typename int64_t size; - const char *desc; };|] - - contextField peakname [C.cty|typename int64_t|] $ Just [C.cexp|0|] - contextField usagename [C.cty|typename int64_t|] $ Just [C.cexp|0|] - - -- Unreferencing a memory block consists of decreasing its reference - -- count and freeing the corresponding memory if the count reaches - -- zero. - free <- collect $ freeRawMem [C.cexp|block->mem|] space [C.cexp|desc|] - ctx_ty <- contextType - let unrefdef = - [C.cedecl|int $id:(fatMemUnRef space) ($ty:ctx_ty *ctx, $ty:mty *block, const char *desc) { - if (block->references != NULL) { - *(block->references) -= 1; - if (ctx->detail_memory) { - fprintf(ctx->log, "Unreferencing block %s (allocated as %s) in %s: %d references remaining.\n", - desc, block->desc, $string:spacedesc, *(block->references)); - } - if (*(block->references) == 0) { - ctx->$id:usagename -= block->size; - $items:free - free(block->references); - if (ctx->detail_memory) { - fprintf(ctx->log, "%lld bytes freed (now allocated: %lld bytes)\n", - (long long) block->size, (long long) ctx->$id:usagename); - } - } - block->references = NULL; - } - return 0; -}|] - - -- When allocating a memory block we initialise the reference count to 1. - alloc <- - collect $ - allocRawMem [C.cexp|block->mem|] [C.cexp|size|] space [C.cexp|desc|] - let allocdef = - [C.cedecl|int $id:(fatMemAlloc space) ($ty:ctx_ty *ctx, $ty:mty *block, typename int64_t size, const char *desc) { - if (size < 0) { - futhark_panic(1, "Negative allocation of %lld bytes attempted for %s in %s.\n", - (long long)size, desc, $string:spacedesc, ctx->$id:usagename); - } - int ret = $id:(fatMemUnRef space)(ctx, block, desc); - - if (ret != FUTHARK_SUCCESS) { - return ret; - } - - if (ctx->detail_memory) { - fprintf(ctx->log, "Allocating %lld bytes for %s in %s (then allocated: %lld bytes)", - (long long) size, - desc, $string:spacedesc, - (long long) ctx->$id:usagename + size); - } - if (ctx->$id:usagename > ctx->$id:peakname) { - ctx->$id:peakname = ctx->$id:usagename; - if (ctx->detail_memory) { - fprintf(ctx->log, " (new peak).\n"); - } - } else if (ctx->detail_memory) { - fprintf(ctx->log, ".\n"); - } - - $items:alloc - - if (ctx->error == NULL) { - block->references = (int*) malloc(sizeof(int)); - *(block->references) = 1; - block->size = size; - block->desc = desc; - ctx->$id:usagename += size; - return FUTHARK_SUCCESS; - } else { - // We are naively assuming that any memory allocation error is due to OOM. - // We preserve the original error so that a savvy user can perhaps find - // glory despite our naiveté. - - char *old_error = ctx->error; - ctx->error = msgprintf("Failed to allocate memory in %s.\nAttempted allocation: %12lld bytes\nCurrently allocated: %12lld bytes\n%s", - $string:spacedesc, - (long long) size, - (long long) ctx->$id:usagename, - old_error); - free(old_error); - return FUTHARK_OUT_OF_MEMORY; - } - }|] - - -- Memory setting - unreference the destination and increase the - -- count of the source by one. - let setdef = - [C.cedecl|int $id:(fatMemSet space) ($ty:ctx_ty *ctx, $ty:mty *lhs, $ty:mty *rhs, const char *lhs_desc) { - int ret = $id:(fatMemUnRef space)(ctx, lhs, lhs_desc); - if (rhs->references != NULL) { - (*(rhs->references))++; - } - *lhs = *rhs; - return ret; -} -|] - - onClear [C.citem|ctx->$id:peakname = 0;|] - - let peakmsg = "Peak memory usage for " ++ spacedesc ++ ": %lld bytes.\n" - pure - ( structdef, - [unrefdef, allocdef, setdef], - -- Do not report memory usage for DefaultSpace (CPU memory), - -- because it would not be accurate anyway. This whole - -- tracking probably needs to be rethought. - if space == DefaultSpace - then [C.citem|{}|] - else [C.citem|str_builder(&builder, $string:peakmsg, (long long) ctx->$id:peakname);|] - ) - where - mty = fatMemType space - (peakname, usagename, sname, spacedesc) = case space of - Space sid -> - ( C.toIdent ("peak_mem_usage_" ++ sid) noLoc, - C.toIdent ("cur_mem_usage_" ++ sid) noLoc, - C.toIdent ("memblock_" ++ sid) noLoc, - "space '" ++ sid ++ "'" - ) - _ -> - ( "peak_mem_usage_default", - "cur_mem_usage_default", - "memblock", - "default space" - ) - -declMem :: VName -> Space -> CompilerM op s () -declMem name space = do - cached <- isJust <$> cacheMem name - fat <- fatMemory space - unless cached $ - if fat - then modify $ \s -> s {compDeclaredMem = (name, space) : compDeclaredMem s} - else do - ty <- memToCType name space - decl [C.cdecl|$ty:ty $id:name;|] - -resetMem :: C.ToExp a => a -> Space -> CompilerM op s () -resetMem mem space = do - refcount <- fatMemory space - cached <- isJust <$> cacheMem mem - if cached - then stm [C.cstm|$exp:mem = NULL;|] - else - when refcount $ - stm [C.cstm|$exp:mem.references = NULL;|] - -setMem :: (C.ToExp a, C.ToExp b) => a -> b -> Space -> CompilerM op s () -setMem dest src space = do - refcount <- fatMemory space - let src_s = pretty $ C.toExp src noLoc - if refcount - then - stm - [C.cstm|if ($id:(fatMemSet space)(ctx, &$exp:dest, &$exp:src, - $string:src_s) != 0) { - return 1; - }|] - else case space of - ScalarSpace ds _ -> do - i' <- newVName "i" - let i = C.toIdent i' - it = primTypeToCType $ IntType Int32 - ds' = map (`C.toExp` noLoc) ds - bound = cproduct ds' - stm - [C.cstm|for ($ty:it $id:i = 0; $id:i < $exp:bound; $id:i++) { - $exp:dest[$id:i] = $exp:src[$id:i]; - }|] - _ -> stm [C.cstm|$exp:dest = $exp:src;|] - -unRefMem :: C.ToExp a => a -> Space -> CompilerM op s () -unRefMem mem space = do - refcount <- fatMemory space - cached <- isJust <$> cacheMem mem - let mem_s = pretty $ C.toExp mem noLoc - when (refcount && not cached) $ - stm - [C.cstm|if ($id:(fatMemUnRef space)(ctx, &$exp:mem, $string:mem_s) != 0) { - return 1; - }|] - -allocMem :: - (C.ToExp a, C.ToExp b) => - a -> - b -> - Space -> - C.Stm -> - CompilerM op s () -allocMem mem size space on_failure = do - refcount <- fatMemory space - let mem_s = pretty $ C.toExp mem noLoc - if refcount - then - stm - [C.cstm|if ($id:(fatMemAlloc space)(ctx, &$exp:mem, $exp:size, - $string:mem_s)) { - $stm:on_failure - }|] - else do - freeRawMem mem space mem_s - allocRawMem mem size space [C.cexp|desc|] - -copyMemoryDefaultSpace :: - C.Exp -> - C.Exp -> - C.Exp -> - C.Exp -> - C.Exp -> - CompilerM op s () -copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = - stm - [C.cstm|if ($exp:nbytes > 0) { - memmove($exp:destmem + $exp:destidx, - $exp:srcmem + $exp:srcidx, - $exp:nbytes); - }|] - ---- Entry points. - -criticalSection :: Operations op s -> [C.BlockItem] -> [C.BlockItem] -criticalSection ops x = - [C.citems|lock_lock(&ctx->lock); - $items:(fst (opsCritical ops)) - $items:x - $items:(snd (opsCritical ops)) - lock_unlock(&ctx->lock); - |] - -arrayLibraryFunctions :: - Publicness -> - Space -> - PrimType -> - Signedness -> - Int -> - CompilerM op s Manifest.ArrayOps -arrayLibraryFunctions pub space pt signed rank = do - let pt' = primAPIType signed pt - name = arrayName pt signed rank - arr_name = "futhark_" ++ name - array_type = [C.cty|struct $id:arr_name|] - - new_array <- publicName $ "new_" ++ name - new_raw_array <- publicName $ "new_raw_" ++ name - free_array <- publicName $ "free_" ++ name - values_array <- publicName $ "values_" ++ name - values_raw_array <- publicName $ "values_raw_" ++ name - shape_array <- publicName $ "shape_" ++ name - - let shape_names = ["dim" ++ show i | i <- [0 .. rank - 1]] - shape_params = [[C.cparam|typename int64_t $id:k|] | k <- shape_names] - arr_size = cproduct [[C.cexp|$id:k|] | k <- shape_names] - arr_size_array = cproduct [[C.cexp|arr->shape[$int:i]|] | i <- [0 .. rank - 1]] - copy <- asks envCopy - - memty <- rawMemCType space - - let prepare_new = do - resetMem [C.cexp|arr->mem|] space - allocMem - [C.cexp|arr->mem|] - [C.cexp|$exp:arr_size * $int:(primByteSize pt::Int)|] - space - [C.cstm|return NULL;|] - forM_ [0 .. rank - 1] $ \i -> - let dim_s = "dim" ++ show i - in stm [C.cstm|arr->shape[$int:i] = $id:dim_s;|] - - new_body <- collect $ do - prepare_new - copy - CopyNoBarrier - [C.cexp|arr->mem.mem|] - [C.cexp|0|] - space - [C.cexp|data|] - [C.cexp|0|] - DefaultSpace - [C.cexp|((size_t)$exp:arr_size) * $int:(primByteSize pt::Int)|] - - new_raw_body <- collect $ do - prepare_new - copy - CopyNoBarrier - [C.cexp|arr->mem.mem|] - [C.cexp|0|] - space - [C.cexp|data|] - [C.cexp|offset|] - space - [C.cexp|((size_t)$exp:arr_size) * $int:(primByteSize pt::Int)|] - - free_body <- collect $ unRefMem [C.cexp|arr->mem|] space - - values_body <- - collect $ - copy - CopyNoBarrier - [C.cexp|data|] - [C.cexp|0|] - DefaultSpace - [C.cexp|arr->mem.mem|] - [C.cexp|0|] - space - [C.cexp|((size_t)$exp:arr_size_array) * $int:(primByteSize pt::Int)|] - - ctx_ty <- contextType - ops <- asks envOperations - - let proto = case pub of - Public -> headerDecl (ArrayDecl name) - Private -> libDecl - - proto - [C.cedecl|struct $id:arr_name;|] - proto - [C.cedecl|$ty:array_type* $id:new_array($ty:ctx_ty *ctx, const $ty:pt' *data, $params:shape_params);|] - proto - [C.cedecl|$ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, const $ty:memty data, typename int64_t offset, $params:shape_params);|] - proto - [C.cedecl|int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - proto - [C.cedecl|int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data);|] - proto - [C.cedecl|$ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - proto - [C.cedecl|const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] - - mapM_ - libDecl - [C.cunit| - $ty:array_type* $id:new_array($ty:ctx_ty *ctx, const $ty:pt' *data, $params:shape_params) { - $ty:array_type* bad = NULL; - $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); - if (arr == NULL) { - return bad; - } - $items:(criticalSection ops new_body) - return arr; - } - - $ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, const $ty:memty data, typename int64_t offset, - $params:shape_params) { - $ty:array_type* bad = NULL; - $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); - if (arr == NULL) { - return bad; - } - $items:(criticalSection ops new_raw_body) - return arr; - } - - int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr) { - $items:(criticalSection ops free_body) - free(arr); - return 0; - } - - int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data) { - $items:(criticalSection ops values_body) - return 0; - } - - $ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr) { - (void)ctx; - return arr->mem.mem; - } - - const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr) { - (void)ctx; - return arr->shape; - } - |] - - pure $ - Manifest.ArrayOps - { Manifest.arrayFree = T.pack free_array, - Manifest.arrayShape = T.pack shape_array, - Manifest.arrayValues = T.pack values_array, - Manifest.arrayNew = T.pack new_array - } - -lookupOpaqueType :: String -> OpaqueTypes -> OpaqueType -lookupOpaqueType v (OpaqueTypes types) = - case lookup v types of - Just t -> t - Nothing -> error $ "Unknown opaque type: " ++ show v - -opaquePayload :: OpaqueTypes -> OpaqueType -> [ValueType] -opaquePayload _ (OpaqueType ts) = ts -opaquePayload types (OpaqueRecord fs) = concatMap f fs - where - f (_, TypeOpaque s) = opaquePayload types $ lookupOpaqueType s types - f (_, TypeTransparent v) = [v] - -opaqueToCType :: String -> CompilerM op s C.Type -opaqueToCType desc = do - name <- publicName $ opaqueName desc - pure [C.cty|struct $id:name|] - -valueTypeToCType :: Publicness -> ValueType -> CompilerM op s C.Type -valueTypeToCType _ (ValueType signed (Rank 0) pt) = - pure $ primAPIType signed pt -valueTypeToCType pub (ValueType signed (Rank rank) pt) = do - name <- publicName $ arrayName pt signed rank - let add = M.insertWith max (signed, pt, rank) pub - modify $ \s -> s {compArrayTypes = add $ compArrayTypes s} - pure [C.cty|struct $id:name|] - -entryPointTypeToCType :: Publicness -> EntryPointType -> CompilerM op s C.Type -entryPointTypeToCType _ (TypeOpaque desc) = opaqueToCType desc -entryPointTypeToCType pub (TypeTransparent vt) = valueTypeToCType pub vt - -entryTypeName :: EntryPointType -> Manifest.TypeName -entryTypeName (TypeOpaque desc) = T.pack desc -entryTypeName (TypeTransparent vt) = prettyText vt - --- | Figure out which of the members of an opaque type corresponds to --- which fields. -recordFieldPayloads :: OpaqueTypes -> [EntryPointType] -> [a] -> [[a]] -recordFieldPayloads types = chunks . map typeLength - where - typeLength (TypeTransparent _) = 1 - typeLength (TypeOpaque desc) = - length $ opaquePayload types $ lookupOpaqueType desc types - -opaqueProjectFunctions :: - OpaqueTypes -> - String -> - [(Name, EntryPointType)] -> - [ValueType] -> - CompilerM op s [Manifest.RecordField] -opaqueProjectFunctions types desc fs vds = do - opaque_type <- opaqueToCType desc - ctx_ty <- contextType - ops <- asks envOperations - let mkProject (TypeTransparent (ValueType sign (Rank 0) pt)) [(i, _)] = do - pure - ( primAPIType sign pt, - [C.citems|v = obj->$id:(tupleField i);|] - ) - mkProject (TypeTransparent vt) [(i, _)] = do - ct <- valueTypeToCType Public vt - pure - ( [C.cty|$ty:ct *|], - criticalSection - ops - [C.citems|v = malloc(sizeof($ty:ct)); - memcpy(v, obj->$id:(tupleField i), sizeof($ty:ct)); - (void)(*(v->mem.references))++;|] - ) - mkProject (TypeTransparent _) rep = - error $ "mkProject: invalid representation of transparent type: " ++ show rep - mkProject (TypeOpaque f_desc) components = do - ct <- opaqueToCType f_desc - let setField j (i, ValueType _ (Rank r) _) = - if r == 0 - then [C.citems|v->$id:(tupleField j) = obj->$id:(tupleField i);|] - else - [C.citems|v->$id:(tupleField j) = malloc(sizeof(*v->$id:(tupleField j))); - *v->$id:(tupleField j) = *obj->$id:(tupleField i); - (void)(*(v->$id:(tupleField j)->mem.references))++;|] - pure - ( [C.cty|$ty:ct *|], - criticalSection - ops - [C.citems|v = malloc(sizeof($ty:ct)); - $items:(concat (zipWith setField [0..] components))|] - ) - let onField ((f, et), elems) = do - project <- publicName $ "project_" ++ opaqueName desc ++ "_" ++ nameToString f - (et_ty, project_items) <- mkProject et elems - headerDecl - (OpaqueDecl desc) - [C.cedecl|int $id:project($ty:ctx_ty *ctx, $ty:et_ty *out, const $ty:opaque_type *obj);|] - libDecl - [C.cedecl|int $id:project($ty:ctx_ty *ctx, $ty:et_ty *out, const $ty:opaque_type *obj) { - (void)ctx; - $ty:et_ty v; - $items:project_items - *out = v; - return 0; - }|] - pure $ Manifest.RecordField (nameToText f) (entryTypeName et) (T.pack project) - - mapM onField . zip fs . recordFieldPayloads types (map snd fs) $ - zip [0 ..] vds - -opaqueNewFunctions :: - OpaqueTypes -> - String -> - [(Name, EntryPointType)] -> - [ValueType] -> - CompilerM op s Manifest.CFuncName -opaqueNewFunctions types desc fs vds = do - opaque_type <- opaqueToCType desc - ctx_ty <- contextType - ops <- asks envOperations - new <- publicName $ "new_" ++ opaqueName desc - - (params, new_stms) <- - fmap (unzip . snd) - . mapAccumLM onField 0 - . zip fs - . recordFieldPayloads types (map snd fs) - $ vds - - headerDecl - (OpaqueDecl desc) - [C.cedecl|int $id:new($ty:ctx_ty *ctx, $ty:opaque_type** out, $params:params);|] - libDecl - [C.cedecl|int $id:new($ty:ctx_ty *ctx, $ty:opaque_type** out, $params:params) { - $ty:opaque_type* v = malloc(sizeof($ty:opaque_type)); - $items:(criticalSection ops new_stms) - *out = v; - return 0; - }|] - pure $ T.pack new - where - onField offset ((f, et), f_vts) = do - let param_name = - if all isDigit (nameToString f) - then C.toIdent ("v" <> f) mempty - else C.toIdent f mempty - case et of - TypeTransparent (ValueType sign (Rank 0) pt) -> do - let ct = primAPIType sign pt - pure - ( offset + 1, - ( [C.cparam|const $ty:ct $id:param_name|], - [C.citem|v->$id:(tupleField offset) = $id:param_name;|] - ) - ) - TypeTransparent vt -> do - ct <- valueTypeToCType Public vt - pure - ( offset + 1, - ( [C.cparam|const $ty:ct* $id:param_name|], - [C.citem|{v->$id:(tupleField offset) = malloc(sizeof($ty:ct)); - *v->$id:(tupleField offset) = *$id:param_name; - (void)(*(v->$id:(tupleField offset)->mem.references))++;}|] - ) - ) - TypeOpaque f_desc -> do - ct <- opaqueToCType f_desc - let param_fields = do - i <- [0 ..] - pure [C.cexp|$id:param_name->$id:(tupleField i)|] - pure - ( offset + length f_vts, - ( [C.cparam|const $ty:ct* $id:param_name|], - [C.citem|{$stms:(zipWith3 setFieldField [offset ..] param_fields f_vts)}|] - ) - ) - - setFieldField i e (ValueType _ (Rank r) _) - | r == 0 = - [C.cstm|v->$id:(tupleField i) = $exp:e;|] - | otherwise = - [C.cstm|{v->$id:(tupleField i) = malloc(sizeof(*$exp:e)); - *v->$id:(tupleField i) = *$exp:e; - (void)(*(v->$id:(tupleField i)->mem.references))++;}|] - -processOpaqueRecord :: - OpaqueTypes -> - String -> - OpaqueType -> - [ValueType] -> - CompilerM op s (Maybe Manifest.RecordOps) -processOpaqueRecord _ _ (OpaqueType _) _ = pure Nothing -processOpaqueRecord types desc (OpaqueRecord fs) vds = - Just - <$> ( Manifest.RecordOps - <$> opaqueProjectFunctions types desc fs vds - <*> opaqueNewFunctions types desc fs vds - ) - -opaqueLibraryFunctions :: - OpaqueTypes -> - String -> - OpaqueType -> - CompilerM op s (Manifest.OpaqueOps, Maybe Manifest.RecordOps) -opaqueLibraryFunctions types desc ot = do - name <- publicName $ opaqueName desc - free_opaque <- publicName $ "free_" ++ opaqueName desc - store_opaque <- publicName $ "store_" ++ opaqueName desc - restore_opaque <- publicName $ "restore_" ++ opaqueName desc - - let opaque_type = [C.cty|struct $id:name|] - - freeComponent i (ValueType signed (Rank rank) pt) = unless (rank == 0) $ do - let field = tupleField i - free_array <- publicName $ "free_" ++ arrayName pt signed rank - -- Protect against NULL here, because we also want to use this - -- to free partially loaded opaques. - stm - [C.cstm|if (obj->$id:field != NULL && (tmp = $id:free_array(ctx, obj->$id:field)) != 0) { - ret = tmp; - }|] - - storeComponent i (ValueType sign (Rank 0) pt) = - let field = tupleField i - in ( storageSize pt 0 [C.cexp|NULL|], - storeValueHeader sign pt 0 [C.cexp|NULL|] [C.cexp|out|] - ++ [C.cstms|memcpy(out, &obj->$id:field, sizeof(obj->$id:field)); - out += sizeof(obj->$id:field);|] - ) - storeComponent i (ValueType sign (Rank rank) pt) = - let arr_name = arrayName pt sign rank - field = tupleField i - shape_array = "futhark_shape_" ++ arr_name - values_array = "futhark_values_" ++ arr_name - shape' = [C.cexp|$id:shape_array(ctx, obj->$id:field)|] - num_elems = cproduct [[C.cexp|$exp:shape'[$int:j]|] | j <- [0 .. rank - 1]] - in ( storageSize pt rank shape', - storeValueHeader sign pt rank shape' [C.cexp|out|] - ++ [C.cstms|ret |= $id:values_array(ctx, obj->$id:field, (void*)out); - out += $exp:num_elems * $int:(primByteSize pt::Int);|] - ) - - ctx_ty <- contextType - - let vds = opaquePayload types ot - free_body <- collect $ zipWithM_ freeComponent [0 ..] vds - - store_body <- collect $ do - let (sizes, stores) = unzip $ zipWith storeComponent [0 ..] vds - size_vars = map (("size_" ++) . show) [0 .. length sizes - 1] - size_sum = csum [[C.cexp|$id:size|] | size <- size_vars] - forM_ (zip size_vars sizes) $ \(v, e) -> - item [C.citem|typename int64_t $id:v = $exp:e;|] - stm [C.cstm|*n = $exp:size_sum;|] - stm [C.cstm|if (p != NULL && *p == NULL) { *p = malloc(*n); }|] - stm [C.cstm|if (p != NULL) { unsigned char *out = *p; $stms:(concat stores) }|] - - let restoreComponent i (ValueType sign (Rank 0) pt) = do - let field = tupleField i - dataptr = "data_" ++ show i - stms $ loadValueHeader sign pt 0 [C.cexp|NULL|] [C.cexp|src|] - item [C.citem|const void* $id:dataptr = src;|] - stm [C.cstm|src += sizeof(obj->$id:field);|] - pure [C.cstms|memcpy(&obj->$id:field, $id:dataptr, sizeof(obj->$id:field));|] - restoreComponent i (ValueType sign (Rank rank) pt) = do - let field = tupleField i - arr_name = arrayName pt sign rank - new_array = "futhark_new_" ++ arr_name - dataptr = "data_" ++ show i - shapearr = "shape_" ++ show i - dims = [[C.cexp|$id:shapearr[$int:j]|] | j <- [0 .. rank - 1]] - num_elems = cproduct dims - item [C.citem|typename int64_t $id:shapearr[$int:rank] = {0};|] - stms $ loadValueHeader sign pt rank [C.cexp|$id:shapearr|] [C.cexp|src|] - item [C.citem|const void* $id:dataptr = src;|] - stm [C.cstm|obj->$id:field = NULL;|] - stm [C.cstm|src += $exp:num_elems * $int:(primByteSize pt::Int);|] - pure - [C.cstms| - obj->$id:field = $id:new_array(ctx, $id:dataptr, $args:dims); - if (obj->$id:field == NULL) { err = 1; }|] - - load_body <- collect $ do - loads <- concat <$> zipWithM restoreComponent [0 ..] (opaquePayload types ot) - stm - [C.cstm|if (err == 0) { - $stms:loads - }|] - - headerDecl - (OpaqueTypeDecl desc) - [C.cedecl|struct $id:name;|] - headerDecl - (OpaqueDecl desc) - [C.cedecl|int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj);|] - headerDecl - (OpaqueDecl desc) - [C.cedecl|int $id:store_opaque($ty:ctx_ty *ctx, const $ty:opaque_type *obj, void **p, size_t *n);|] - headerDecl - (OpaqueDecl desc) - [C.cedecl|$ty:opaque_type* $id:restore_opaque($ty:ctx_ty *ctx, const void *p);|] - - record <- processOpaqueRecord types desc ot vds - - -- We do not need to enclose most bodies in a critical section, - -- because when we operate on the components of the opaque, we are - -- calling public API functions that do their own locking. The - -- exception is projection, where we fiddle with reference counts. - mapM_ - libDecl - [C.cunit| - int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj) { - (void)ctx; - int ret = 0, tmp; - $items:free_body - free(obj); - return ret; - } - - int $id:store_opaque($ty:ctx_ty *ctx, - const $ty:opaque_type *obj, void **p, size_t *n) { - (void)ctx; - int ret = 0; - $items:store_body - return ret; - } - - $ty:opaque_type* $id:restore_opaque($ty:ctx_ty *ctx, - const void *p) { - int err = 0; - const unsigned char *src = p; - $ty:opaque_type* obj = malloc(sizeof($ty:opaque_type)); - $items:load_body - if (err != 0) { - int ret = 0, tmp; - $items:free_body - free(obj); - obj = NULL; - } - return obj; - } - |] - - pure - ( Manifest.OpaqueOps - { Manifest.opaqueFree = T.pack free_opaque, - Manifest.opaqueStore = T.pack store_opaque, - Manifest.opaqueRestore = T.pack restore_opaque - }, - record - ) - -valueDescToType :: ValueDesc -> ValueType -valueDescToType (ScalarValue pt signed _) = - ValueType signed (Rank 0) pt -valueDescToType (ArrayValue _ _ pt signed shape) = - ValueType signed (Rank (length shape)) pt - -generateArray :: - Space -> - ((Signedness, PrimType, Int), Publicness) -> - CompilerM op s (Maybe (T.Text, Manifest.Type)) -generateArray space ((signed, pt, rank), pub) = do - name <- publicName $ arrayName pt signed rank - let memty = fatMemType space - libDecl [C.cedecl|struct $id:name { $ty:memty mem; typename int64_t shape[$int:rank]; };|] - ops <- arrayLibraryFunctions pub space pt signed rank - let pt_name = T.pack $ prettySigned (signed == Unsigned) pt - pretty_name = mconcat (replicate rank "[]") <> pt_name - arr_type = [C.cty|struct $id:name*|] - case pub of - Public -> - pure $ - Just - ( pretty_name, - Manifest.TypeArray (prettyText arr_type) pt_name rank ops - ) - Private -> - pure Nothing - -generateOpaque :: - OpaqueTypes -> - (String, OpaqueType) -> - CompilerM op s (T.Text, Manifest.Type) -generateOpaque types (desc, ot) = do - name <- publicName $ opaqueName desc - members <- zipWithM field (opaquePayload types ot) [(0 :: Int) ..] - libDecl [C.cedecl|struct $id:name { $sdecls:members };|] - (ops, record) <- opaqueLibraryFunctions types desc ot - let opaque_type = [C.cty|struct $id:name*|] - pure (T.pack desc, Manifest.TypeOpaque (prettyText opaque_type) ops record) - where - field vt@(ValueType _ (Rank r) _) i = do - ct <- valueTypeToCType Private vt - pure $ - if r == 0 - then [C.csdecl|$ty:ct $id:(tupleField i);|] - else [C.csdecl|$ty:ct *$id:(tupleField i);|] - -generateAPITypes :: Space -> OpaqueTypes -> CompilerM op s (M.Map T.Text Manifest.Type) -generateAPITypes arr_space types@(OpaqueTypes opaques) = do - mapM_ (findNecessaryArrays . snd) opaques - array_ts <- mapM (generateArray arr_space) . M.toList =<< gets compArrayTypes - opaque_ts <- mapM (generateOpaque types) opaques - pure $ M.fromList $ catMaybes array_ts <> opaque_ts - where - -- Ensure that array types will be generated before the opaque - -- records that allow projection of them. This is because the - -- projection functions somewhat uglily directly poke around in - -- the innards to increment reference counts. - findNecessaryArrays (OpaqueType _) = - pure () - findNecessaryArrays (OpaqueRecord fs) = - mapM_ (entryPointTypeToCType Public . snd) fs - -allTrue :: [C.Exp] -> C.Exp -allTrue [] = [C.cexp|true|] -allTrue [x] = x -allTrue (x : xs) = [C.cexp|$exp:x && $exp:(allTrue xs)|] - -prepareEntryInputs :: - [ExternalValue] -> - CompilerM op s ([(C.Param, Maybe C.Exp)], [C.BlockItem]) -prepareEntryInputs args = collect' $ zipWithM prepare [(0 :: Int) ..] args - where - arg_names = namesFromList $ concatMap evNames args - evNames (OpaqueValue _ vds) = map vdName vds - evNames (TransparentValue vd) = [vdName vd] - vdName (ArrayValue v _ _ _ _) = v - vdName (ScalarValue _ _ v) = v - - prepare pno (TransparentValue vd) = do - let pname = "in" ++ show pno - (ty, check) <- prepareValue Public [C.cexp|$id:pname|] vd - pure - ( [C.cparam|const $ty:ty $id:pname|], - if null check then Nothing else Just $ allTrue check - ) - prepare pno (OpaqueValue desc vds) = do - ty <- opaqueToCType desc - let pname = "in" ++ show pno - field i ScalarValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] - field i ArrayValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] - checks <- map snd <$> zipWithM (prepareValue Private) (zipWith field [0 ..] vds) vds - pure - ( [C.cparam|const $ty:ty *$id:pname|], - if all null checks - then Nothing - else Just $ allTrue $ concat checks - ) - - prepareValue _ src (ScalarValue pt signed name) = do - let pt' = primAPIType signed pt - src' = fromStorage pt $ C.toExp src mempty - stm [C.cstm|$id:name = $exp:src';|] - pure (pt', []) - prepareValue pub src vd@(ArrayValue mem _ _ _ shape) = do - ty <- valueTypeToCType pub $ valueDescToType vd - - stm [C.cstm|$exp:mem = $exp:src->mem;|] - - let rank = length shape - maybeCopyDim (Var d) i - | d `notNameIn` arg_names = - ( Just [C.cstm|$id:d = $exp:src->shape[$int:i];|], - [C.cexp|$id:d == $exp:src->shape[$int:i]|] - ) - maybeCopyDim x i = - ( Nothing, - [C.cexp|$exp:x == $exp:src->shape[$int:i]|] - ) - - let (sets, checks) = - unzip $ zipWith maybeCopyDim shape [0 .. rank - 1] - stms $ catMaybes sets - - pure ([C.cty|$ty:ty*|], checks) - -prepareEntryOutputs :: [ExternalValue] -> CompilerM op s ([C.Param], [C.BlockItem]) -prepareEntryOutputs = collect' . zipWithM prepare [(0 :: Int) ..] - where - prepare pno (TransparentValue vd) = do - let pname = "out" ++ show pno - ty <- valueTypeToCType Public $ valueDescToType vd - - case vd of - ArrayValue {} -> do - stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] - prepareValue [C.cexp|*$id:pname|] vd - pure [C.cparam|$ty:ty **$id:pname|] - ScalarValue {} -> do - prepareValue [C.cexp|*$id:pname|] vd - pure [C.cparam|$ty:ty *$id:pname|] - prepare pno (OpaqueValue desc vds) = do - let pname = "out" ++ show pno - ty <- opaqueToCType desc - vd_ts <- mapM (valueTypeToCType Private . valueDescToType) vds - - stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] - - forM_ (zip3 [0 ..] vd_ts vds) $ \(i, ct, vd) -> do - let field = [C.cexp|((*$id:pname)->$id:(tupleField i))|] - case vd of - ScalarValue {} -> pure () - ArrayValue {} -> do - stm [C.cstm|assert(($exp:field = ($ty:ct*) malloc(sizeof($ty:ct))) != NULL);|] - prepareValue field vd - - pure [C.cparam|$ty:ty **$id:pname|] - - prepareValue dest (ScalarValue t _ name) = - let name' = toStorage t $ C.toExp name mempty - in stm [C.cstm|$exp:dest = $exp:name';|] - prepareValue dest (ArrayValue mem _ _ _ shape) = do - stm [C.cstm|$exp:dest->mem = $id:mem;|] - - let rank = length shape - maybeCopyDim (Constant x) i = - [C.cstm|$exp:dest->shape[$int:i] = $exp:x;|] - maybeCopyDim (Var d) i = - [C.cstm|$exp:dest->shape[$int:i] = $id:d;|] - stms $ zipWith maybeCopyDim shape [0 .. rank - 1] - -isValidCName :: Name -> Bool -isValidCName = check . nameToString - where - check [] = True -- academic - check (c : cs) = isAlpha c && all constituent cs - constituent c = isAlphaNum c || c == '_' - -entryName :: Name -> String -entryName v - | isValidCName v = "entry_" <> nameToString v - | otherwise = "entry_" <> zEncodeString (nameToString v) - -onEntryPoint :: - [C.BlockItem] -> - Name -> - Function op -> - CompilerM op s (Maybe (C.Definition, (T.Text, Manifest.EntryPoint))) -onEntryPoint _ _ (Function Nothing _ _ _) = pure Nothing -onEntryPoint get_consts fname (Function (Just (EntryPoint ename results args)) outputs inputs _) = inNewFunction $ do - let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs - in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs - - inputdecls <- collect $ mapM_ stubParam inputs - outputdecls <- collect $ mapM_ stubParam outputs - decl_mem <- declAllocatedMem - - entry_point_function_name <- publicName $ entryName ename - - (inputs', unpack_entry_inputs) <- prepareEntryInputs $ map snd args - let (entry_point_input_params, entry_point_input_checks) = unzip inputs' - - (entry_point_output_params, pack_entry_outputs) <- - prepareEntryOutputs $ map snd results - - ctx_ty <- contextType - - headerDecl - EntryDecl - [C.cedecl|int $id:entry_point_function_name - ($ty:ctx_ty *ctx, - $params:entry_point_output_params, - $params:entry_point_input_params);|] - - let checks = catMaybes entry_point_input_checks - check_input = - if null checks - then [] - else - [C.citems| - if (!($exp:(allTrue (catMaybes entry_point_input_checks)))) { - ret = 1; - if (!ctx->error) { - ctx->error = msgprintf("Error: entry point arguments have invalid sizes.\n"); - } - }|] - - critical = - [C.citems| - $items:decl_mem - $items:unpack_entry_inputs - $items:check_input - if (ret == 0) { - ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args); - if (ret == 0) { - $items:get_consts - - $items:pack_entry_outputs - } - } - |] - - ops <- asks envOperations - - let cdef = - [C.cedecl| - int $id:entry_point_function_name - ($ty:ctx_ty *ctx, - $params:entry_point_output_params, - $params:entry_point_input_params) { - $items:inputdecls - $items:outputdecls - - int ret = 0; - - $items:(criticalSection ops critical) - - return ret; - }|] - - manifest = - Manifest.EntryPoint - { Manifest.entryPointCFun = T.pack entry_point_function_name, - -- Note that our convention about what is "input/output" - -- and what is "results/args" is different between the - -- manifest and ImpCode. - Manifest.entryPointOutputs = map outputManifest results, - Manifest.entryPointInputs = map inputManifest args - } - - pure $ Just (cdef, (nameToText ename, manifest)) - where - stubParam (MemParam name space) = - declMem name space - stubParam (ScalarParam name ty) = do - let ty' = primTypeToCType ty - decl [C.cdecl|$ty:ty' $id:name;|] - - vdType (TransparentValue (ScalarValue pt signed _)) = - T.pack $ prettySigned (signed == Unsigned) pt - vdType (TransparentValue (ArrayValue _ _ pt signed shape)) = - T.pack $ - mconcat (replicate (length shape) "[]") - <> prettySigned (signed == Unsigned) pt - vdType (OpaqueValue name _) = - T.pack name - - outputManifest (u, vd) = - Manifest.Output - { Manifest.outputType = vdType vd, - Manifest.outputUnique = u == Unique - } - inputManifest ((v, u), vd) = - Manifest.Input - { Manifest.inputName = nameToText v, - Manifest.inputType = vdType vd, - Manifest.inputUnique = u == Unique - } - --- | The result of compilation to C is multiple parts, which can be --- put together in various ways. The obvious way is to concatenate --- all of them, which yields a CLI program. Another is to compile the --- library part by itself, and use the header file to call into it. -data CParts = CParts - { cHeader :: T.Text, - -- | Utility definitions that must be visible - -- to both CLI and library parts. - cUtils :: T.Text, - cCLI :: T.Text, - cServer :: T.Text, - cLib :: T.Text, - -- | The manifest, in JSON format. - cJsonManifest :: T.Text - } - -gnuSource :: T.Text -gnuSource = - [untrimming| -// We need to define _GNU_SOURCE before -// _any_ headers files are imported to get -// the usage statistics of a thread (i.e. have RUSAGE_THREAD) on GNU/Linux -// https://manpages.courier-mta.org/htmlman2/getrusage.2.html -#ifndef _GNU_SOURCE // Avoid possible double-definition warning. -#define _GNU_SOURCE -#endif -|] - --- We may generate variables that are never used (e.g. for --- certificates) or functions that are never called (e.g. unused --- intrinsics), and generated code may have other cosmetic issues that --- compilers warn about. We disable these warnings to not clutter the --- compilation logs. -disableWarnings :: T.Text -disableWarnings = - [untrimming| -#ifdef __clang__ -#pragma clang diagnostic ignored "-Wunused-function" -#pragma clang diagnostic ignored "-Wunused-variable" -#pragma clang diagnostic ignored "-Wparentheses" -#pragma clang diagnostic ignored "-Wunused-label" -#elif __GNUC__ -#pragma GCC diagnostic ignored "-Wunused-function" -#pragma GCC diagnostic ignored "-Wunused-variable" -#pragma GCC diagnostic ignored "-Wparentheses" -#pragma GCC diagnostic ignored "-Wunused-label" -#pragma GCC diagnostic ignored "-Wunused-but-set-variable" -#endif -|] - --- | Produce header, implementation, and manifest files. -asLibrary :: CParts -> (T.Text, T.Text, T.Text) -asLibrary parts = - ( "#pragma once\n\n" <> cHeader parts, - gnuSource <> disableWarnings <> cHeader parts <> cUtils parts <> cLib parts, - cJsonManifest parts - ) - --- | As executable with command-line interface. -asExecutable :: CParts -> T.Text -asExecutable parts = - gnuSource <> disableWarnings <> cHeader parts <> cUtils parts <> cCLI parts <> cLib parts - --- | As server executable. -asServer :: CParts -> T.Text -asServer parts = - gnuSource <> disableWarnings <> cHeader parts <> cUtils parts <> cServer parts <> cLib parts - -compileProg' :: - MonadFreshNames m => - T.Text -> - T.Text -> - Operations op s -> - s -> - CompilerM op s () -> - T.Text -> - (Space, [Space]) -> - [Option] -> - Definitions op -> - m (CParts, CompilerState s) -compileProg' backend version ops def extra header_extra (arr_space, spaces) options prog = do - src <- getNameSource - let ((prototypes, definitions, entry_point_decls, manifest), endstate) = - runCompilerM ops src def compileProgAction - initdecls = initDecls endstate - entrydecls = entryDecls endstate - arraydecls = arrayDecls endstate - opaquetypedecls = opaqueTypeDecls endstate - opaquedecls = opaqueDecls endstate - miscdecls = miscDecls endstate - - let headerdefs = - [untrimming| -// Headers -#include -#include -#include -#include -#include -$header_extra -#ifdef __cplusplus -extern "C" { -#endif - -// Initialisation -$initdecls - -// Arrays -$arraydecls - -// Opaque values -$opaquetypedecls -$opaquedecls - -// Entry points -$entrydecls - -// Miscellaneous -$miscdecls -#define FUTHARK_BACKEND_$backend -$errorsH - -#ifdef __cplusplus -} -#endif -|] - - let utildefs = - [untrimming| -#include -#include -#include -#include -#include -// If NDEBUG is set, the assert() macro will do nothing. Since Futhark -// (unfortunately) makes use of assert() for error detection (and even some -// side effects), we want to avoid that. -#undef NDEBUG -#include -#include -$utilH -$cacheH -$halfH -$timingH -|] - - let early_decls = T.unlines $ map prettyText $ DL.toList $ compEarlyDecls endstate - lib_decls = T.unlines $ map prettyText $ DL.toList $ compLibDecls endstate - clidefs = cliDefs options manifest - serverdefs = serverDefs options manifest - libdefs = - [untrimming| -#ifdef _MSC_VER -#define inline __inline -#endif -#include -#include -#include -#include -#include - -$header_extra - -$lockH - -#define FUTHARK_F64_ENABLED - -$cScalarDefs - -$early_decls - -$prototypes - -$lib_decls - -$definitions - -$entry_point_decls - |] - - pure - ( CParts - { cHeader = headerdefs, - cUtils = utildefs, - cCLI = clidefs, - cServer = serverdefs, - cLib = libdefs, - cJsonManifest = Manifest.manifestToJSON manifest - }, - endstate - ) - where - Definitions types consts (Functions funs) = prog - - compileProgAction = do - (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces - - get_consts <- compileConstants consts - - ctx_ty <- contextType - - (prototypes, functions) <- - unzip <$> mapM (compileFun get_consts [[C.cparam|$ty:ctx_ty *ctx|]]) funs - - mapM_ earlyDecl memstructs - (entry_points, entry_points_manifest) <- - unzip . catMaybes <$> mapM (uncurry (onEntryPoint get_consts)) funs - - extra - - mapM_ earlyDecl $ concat memfuns - - type_funs <- generateAPITypes arr_space types - generateCommonLibFuns memreport - - pure - ( T.unlines $ map prettyText prototypes, - T.unlines $ map (prettyText . funcToDef) functions, - T.unlines $ map prettyText entry_points, - Manifest.Manifest (M.fromList entry_points_manifest) type_funs backend version - ) - - funcToDef func = C.FuncDef func loc - where - loc = case func of - C.OldFunc _ _ _ _ _ _ l -> l - C.Func _ _ _ _ _ l -> l - --- | Compile imperative program to a C program. Always uses the --- function named "main" as entry point, so make sure it is defined. -compileProg :: - MonadFreshNames m => - T.Text -> - T.Text -> - Operations op () -> - CompilerM op () () -> - T.Text -> - (Space, [Space]) -> - [Option] -> - Definitions op -> - m CParts -compileProg backend version ops extra header_extra (arr_space, spaces) options prog = - fst <$> compileProg' backend version ops () extra header_extra (arr_space, spaces) options prog - -generateCommonLibFuns :: [C.BlockItem] -> CompilerM op s () -generateCommonLibFuns memreport = do - ctx <- contextType - cfg <- configType - ops <- asks envOperations - profilereport <- gets $ DL.toList . compProfileItems - - publicDef_ "context_config_set_cache_file" MiscDecl $ \s -> - ( [C.cedecl|void $id:s($ty:cfg* cfg, const char *f);|], - [C.cedecl|void $id:s($ty:cfg* cfg, const char *f) { - cfg->cache_fname = f; - }|] - ) - - publicDef_ "get_tuning_param_count" InitDecl $ \s -> - ( [C.cedecl|int $id:s(void);|], - [C.cedecl|int $id:s(void) { - return sizeof(tuning_param_names)/sizeof(tuning_param_names[0]); - }|] - ) - - publicDef_ "get_tuning_param_name" InitDecl $ \s -> - ( [C.cedecl|const char* $id:s(int);|], - [C.cedecl|const char* $id:s(int i) { - return tuning_param_names[i]; - }|] - ) - - publicDef_ "get_tuning_param_class" InitDecl $ \s -> - ( [C.cedecl|const char* $id:s(int);|], - [C.cedecl|const char* $id:s(int i) { - return tuning_param_classes[i]; - }|] - ) - - sync <- publicName "context_sync" - publicDef_ "context_report" MiscDecl $ \s -> - ( [C.cedecl|char* $id:s($ty:ctx *ctx);|], - [C.cedecl|char* $id:s($ty:ctx *ctx) { - if ($id:sync(ctx) != 0) { - return NULL; - } - - struct str_builder builder; - str_builder_init(&builder); - $items:memreport - if (ctx->profiling) { - $items:profilereport - } - return builder.str; - }|] - ) - - publicDef_ "context_get_error" MiscDecl $ \s -> - ( [C.cedecl|char* $id:s($ty:ctx* ctx);|], - [C.cedecl|char* $id:s($ty:ctx* ctx) { - char* error = ctx->error; - ctx->error = NULL; - return error; - }|] - ) - - publicDef_ "context_set_logging_file" MiscDecl $ \s -> - ( [C.cedecl|void $id:s($ty:ctx* ctx, typename FILE* f);|], - [C.cedecl|void $id:s($ty:ctx* ctx, typename FILE* f) { - ctx->log = f; - }|] - ) - - publicDef_ "context_pause_profiling" MiscDecl $ \s -> - ( [C.cedecl|void $id:s($ty:ctx* ctx);|], - [C.cedecl|void $id:s($ty:ctx* ctx) { - ctx->profiling_paused = 1; - }|] - ) - - publicDef_ "context_unpause_profiling" MiscDecl $ \s -> - ( [C.cedecl|void $id:s($ty:ctx* ctx);|], - [C.cedecl|void $id:s($ty:ctx* ctx) { - ctx->profiling_paused = 0; - }|] - ) - - clears <- gets $ DL.toList . compClearItems - publicDef_ "context_clear_caches" MiscDecl $ \s -> - ( [C.cedecl|int $id:s($ty:ctx* ctx);|], - [C.cedecl|int $id:s($ty:ctx* ctx) { - $items:(criticalSection ops clears) - return ctx->error != NULL; - }|] - ) - -compileConstants :: Constants op -> CompilerM op s [C.BlockItem] -compileConstants (Constants ps init_consts) = do - ctx_ty <- contextType - const_fields <- mapM constParamField ps - -- Avoid an empty struct, as that is apparently undefined behaviour. - let const_fields' - | null const_fields = [[C.csdecl|int dummy;|]] - | otherwise = const_fields - contextField "constants" [C.cty|struct { $sdecls:const_fields' }|] Nothing - earlyDecl [C.cedecl|static int init_constants($ty:ctx_ty*);|] - earlyDecl [C.cedecl|static int free_constants($ty:ctx_ty*);|] - - inNewFunction $ do - -- We locally define macros for the constants, so that when we - -- generate assignments to local variables, we actually assign into - -- the constants struct. This is not needed for functions, because - -- they can only read constants, not write them. - let (defs, undefs) = unzip $ map constMacro ps - init_consts' <- collect $ do - mapM_ resetMemConst ps - compileCode init_consts - decl_mem <- declAllocatedMem - free_mem <- freeAllocatedMem - libDecl - [C.cedecl|static int init_constants($ty:ctx_ty *ctx) { - (void)ctx; - int err = 0; - $items:defs - $items:decl_mem - $items:init_consts' - $items:free_mem - $items:undefs - cleanup: - return err; - }|] - - inNewFunction $ do - free_consts <- collect $ mapM_ freeConst ps - libDecl - [C.cedecl|static int free_constants($ty:ctx_ty *ctx) { - (void)ctx; - $items:free_consts - return 0; - }|] - - mapM getConst ps - where - constParamField (ScalarParam name bt) = do - let ctp = primTypeToCType bt - pure [C.csdecl|$ty:ctp $id:name;|] - constParamField (MemParam name space) = do - ty <- memToCType name space - pure [C.csdecl|$ty:ty $id:name;|] - - constMacro p = ([C.citem|$escstm:def|], [C.citem|$escstm:undef|]) - where - p' = pretty (C.toIdent (paramName p) mempty) - def = "#define " ++ p' ++ " (" ++ "ctx->constants." ++ p' ++ ")" - undef = "#undef " ++ p' - - resetMemConst ScalarParam {} = pure () - resetMemConst (MemParam name space) = resetMem name space - - freeConst ScalarParam {} = pure () - freeConst (MemParam name space) = unRefMem [C.cexp|ctx->constants.$id:name|] space - - getConst (ScalarParam name bt) = do - let ctp = primTypeToCType bt - pure [C.citem|$ty:ctp $id:name = ctx->constants.$id:name;|] - getConst (MemParam name space) = do - ty <- memToCType name space - pure [C.citem|$ty:ty $id:name = ctx->constants.$id:name;|] - -cachingMemory :: - M.Map VName Space -> - ([C.BlockItem] -> [C.Stm] -> CompilerM op s a) -> - CompilerM op s a -cachingMemory lexical f = do - -- We only consider lexical 'DefaultSpace' memory blocks to be - -- cached. This is not a deep technical restriction, but merely a - -- heuristic based on GPU memory usually involving larger - -- allocations, that do not suffer from the overhead of reference - -- counting. - let cached = M.keys $ M.filter (== DefaultSpace) lexical - - cached' <- forM cached $ \mem -> do - size <- newVName $ pretty mem <> "_cached_size" - pure (mem, size) - - let lexMem env = - env - { envCachedMem = - M.fromList (map (first (`C.toExp` noLoc)) cached') - <> envCachedMem env - } - - declCached (mem, size) = - [ [C.citem|typename int64_t $id:size = 0;|], - [C.citem|$ty:defaultMemBlockType $id:mem = NULL;|] - ] - - freeCached (mem, _) = - [C.cstm|free($id:mem);|] - - local lexMem $ f (concatMap declCached cached') (map freeCached cached') - -compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func) -compileFun get_constants extra (fname, func@(Function _ outputs inputs body)) = inNewFunction $ do - (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs - inparams <- mapM compileInput inputs - - cachingMemory (lexicalMemoryUsage func) $ \decl_cached free_cached -> do - body' <- collect $ compileFunBody out_ptrs outputs body - decl_mem <- declAllocatedMem - free_mem <- freeAllocatedMem - - pure - ( [C.cedecl|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams);|], - [C.cfun|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams) { - $stms:ignores - int err = 0; - $items:decl_cached - $items:decl_mem - $items:get_constants - $items:body' - cleanup: - { - $stms:free_cached - $items:free_mem - } - return err; - }|] - ) - where - -- Ignore all the boilerplate parameters, just in case we don't - -- actually need to use them. - ignores = [[C.cstm|(void)$id:p;|] | C.Param (Just p) _ _ _ <- extra] - - compileInput (ScalarParam name bt) = do - let ctp = primTypeToCType bt - pure [C.cparam|$ty:ctp $id:name|] - compileInput (MemParam name space) = do - ty <- memToCType name space - pure [C.cparam|$ty:ty $id:name|] - - compileOutput (ScalarParam name bt) = do - let ctp = primTypeToCType bt - p_name <- newVName $ "out_" ++ baseString name - pure ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|]) - compileOutput (MemParam name space) = do - ty <- memToCType name space - p_name <- newVName $ baseString name ++ "_p" - pure ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|]) - -derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp -derefPointer ptr i res_t = - [C.cexp|(($ty:res_t)$exp:ptr)[$exp:i]|] - -volQuals :: Volatility -> [C.TypeQual] -volQuals Volatile = [C.ctyquals|volatile|] -volQuals Nonvolatile = [] - -writeScalarPointerWithQuals :: PointerQuals op s -> WriteScalar op s -writeScalarPointerWithQuals quals_f dest i elemtype space vol v = do - quals <- quals_f space - let quals' = volQuals vol ++ quals - deref = - derefPointer - dest - i - [C.cty|$tyquals:quals' $ty:elemtype*|] - stm [C.cstm|$exp:deref = $exp:v;|] - -readScalarPointerWithQuals :: PointerQuals op s -> ReadScalar op s -readScalarPointerWithQuals quals_f dest i elemtype space vol = do - quals <- quals_f space - let quals' = volQuals vol ++ quals - pure $ derefPointer dest i [C.cty|$tyquals:quals' $ty:elemtype*|] - -compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName -compileExpToName _ _ (LeafExp v _) = - pure v -compileExpToName desc t e = do - desc' <- newVName desc - e' <- compileExp e - decl [C.cdecl|$ty:(primTypeToCType t) $id:desc' = $e';|] - pure desc' - -compileExp :: Exp -> CompilerM op s C.Exp -compileExp = compilePrimExp $ \v -> pure [C.cexp|$id:v|] - --- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. -compilePrimExp :: Monad m => (v -> m C.Exp) -> PrimExp v -> m C.Exp -compilePrimExp _ (ValueExp val) = - pure $ C.toExp val mempty -compilePrimExp f (LeafExp v _) = - f v -compilePrimExp f (UnOpExp Complement {} x) = do - x' <- compilePrimExp f x - pure [C.cexp|~$exp:x'|] -compilePrimExp f (UnOpExp Not {} x) = do - x' <- compilePrimExp f x - pure [C.cexp|!$exp:x'|] -compilePrimExp f (UnOpExp (FAbs Float32) x) = do - x' <- compilePrimExp f x - pure [C.cexp|(float)fabs($exp:x')|] -compilePrimExp f (UnOpExp (FAbs Float64) x) = do - x' <- compilePrimExp f x - pure [C.cexp|fabs($exp:x')|] -compilePrimExp f (UnOpExp SSignum {} x) = do - x' <- compilePrimExp f x - pure [C.cexp|($exp:x' > 0 ? 1 : 0) - ($exp:x' < 0 ? 1 : 0)|] -compilePrimExp f (UnOpExp USignum {} x) = do - x' <- compilePrimExp f x - pure [C.cexp|($exp:x' > 0 ? 1 : 0) - ($exp:x' < 0 ? 1 : 0) != 0|] -compilePrimExp f (UnOpExp op x) = do - x' <- compilePrimExp f x - pure [C.cexp|$id:(pretty op)($exp:x')|] -compilePrimExp f (CmpOpExp cmp x y) = do - x' <- compilePrimExp f x - y' <- compilePrimExp f y - pure $ case cmp of - CmpEq {} -> [C.cexp|$exp:x' == $exp:y'|] - FCmpLt {} -> [C.cexp|$exp:x' < $exp:y'|] - FCmpLe {} -> [C.cexp|$exp:x' <= $exp:y'|] - CmpLlt {} -> [C.cexp|$exp:x' < $exp:y'|] - CmpLle {} -> [C.cexp|$exp:x' <= $exp:y'|] - _ -> [C.cexp|$id:(pretty cmp)($exp:x', $exp:y')|] -compilePrimExp f (ConvOpExp conv x) = do - x' <- compilePrimExp f x - pure [C.cexp|$id:(pretty conv)($exp:x')|] -compilePrimExp f (BinOpExp bop x y) = do - x' <- compilePrimExp f x - y' <- compilePrimExp f y - -- Note that integer addition, subtraction, and multiplication with - -- OverflowWrap are not handled by explicit operators, but rather by - -- functions. This is because we want to implicitly convert them to - -- unsigned numbers, so we can do overflow without invoking - -- undefined behaviour. - pure $ case bop of - Add _ OverflowUndef -> [C.cexp|$exp:x' + $exp:y'|] - Sub _ OverflowUndef -> [C.cexp|$exp:x' - $exp:y'|] - Mul _ OverflowUndef -> [C.cexp|$exp:x' * $exp:y'|] - FAdd {} -> [C.cexp|$exp:x' + $exp:y'|] - FSub {} -> [C.cexp|$exp:x' - $exp:y'|] - FMul {} -> [C.cexp|$exp:x' * $exp:y'|] - FDiv {} -> [C.cexp|$exp:x' / $exp:y'|] - Xor {} -> [C.cexp|$exp:x' ^ $exp:y'|] - And {} -> [C.cexp|$exp:x' & $exp:y'|] - Or {} -> [C.cexp|$exp:x' | $exp:y'|] - LogAnd {} -> [C.cexp|$exp:x' && $exp:y'|] - LogOr {} -> [C.cexp|$exp:x' || $exp:y'|] - _ -> [C.cexp|$id:(pretty bop)($exp:x', $exp:y')|] -compilePrimExp f (FunExp h args _) = do - args' <- mapM (compilePrimExp f) args - pure [C.cexp|$id:(funName (nameFromString h))($args:args')|] - -linearCode :: Code op -> [Code op] -linearCode = reverse . go [] - where - go acc (x :>>: y) = - go (go acc x) y - go acc x = x : acc - -compileCode :: Code op -> CompilerM op s () -compileCode (Op op) = - join $ asks envOpCompiler <*> pure op -compileCode Skip = pure () -compileCode (Comment s code) = do - xs <- collect $ compileCode code - let comment = "// " ++ s - stm - [C.cstm|$comment:comment - { $items:xs } - |] -compileCode (TracePrint msg) = do - (formatstr, formatargs) <- errorMsgString msg - stm [C.cstm|fprintf(ctx->log, $string:formatstr, $args:formatargs);|] -compileCode (DebugPrint s (Just e)) = do - e' <- compileExp e - stm - [C.cstm|if (ctx->debugging) { - fprintf(ctx->log, $string:fmtstr, $exp:s, ($ty:ety)$exp:e', '\n'); - }|] - where - (fmt, ety) = case primExpType e of - IntType _ -> ("llu", [C.cty|long long int|]) - FloatType _ -> ("f", [C.cty|double|]) - _ -> ("d", [C.cty|int|]) - fmtstr = "%s: %" ++ fmt ++ "%c" -compileCode (DebugPrint s Nothing) = - stm - [C.cstm|if (ctx->debugging) { - fprintf(ctx->log, "%s\n", $exp:s); - }|] --- :>>: is treated in a special way to detect declare-set pairs in --- order to generate prettier code. -compileCode (c1 :>>: c2) = go (linearCode (c1 :>>: c2)) - where - go (DeclareScalar name vol t : SetScalar dest e : code) - | name == dest = do - let ct = primTypeToCType t - e' <- compileExp e - item [C.citem|$tyquals:(volQuals vol) $ty:ct $id:name = $exp:e';|] - go code - go (x : xs) = compileCode x >> go xs - go [] = pure () -compileCode (Assert e msg (loc, locs)) = do - e' <- compileExp e - err <- - collect . join $ - asks (opsError . envOperations) <*> pure msg <*> pure stacktrace - stm [C.cstm|if (!$exp:e') { $items:err }|] - where - stacktrace = prettyStacktrace 0 $ map locStr $ loc : locs -compileCode (Allocate _ _ ScalarSpace {}) = - -- Handled by the declaration of the memory block, which is - -- translated to an actual array. - pure () -compileCode (Allocate name (Count (TPrimExp e)) space) = do - size <- compileExp e - cached <- cacheMem name - case cached of - Just cur_size -> - stm - [C.cstm|if ($exp:cur_size < $exp:size) { - err = lexical_realloc(&ctx->error, &$exp:name, &$exp:cur_size, $exp:size); - if (err != FUTHARK_SUCCESS) { - goto cleanup; - } - }|] - _ -> - allocMem name size space [C.cstm|{err = 1; goto cleanup;}|] -compileCode (Free name space) = do - cached <- isJust <$> cacheMem name - unless cached $ unRefMem name space -compileCode (For i bound body) = do - let i' = C.toIdent i - t = primTypeToCType $ primExpType bound - bound' <- compileExp bound - body' <- collect $ compileCode body - stm - [C.cstm|for ($ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) { - $items:body' - }|] -compileCode (While cond body) = do - cond' <- compileExp $ untyped cond - body' <- collect $ compileCode body - stm - [C.cstm|while ($exp:cond') { - $items:body' - }|] -compileCode (If cond tbranch fbranch) = do - cond' <- compileExp $ untyped cond - tbranch' <- collect $ compileCode tbranch - fbranch' <- collect $ compileCode fbranch - stm $ case (tbranch', fbranch') of - (_, []) -> - [C.cstm|if ($exp:cond') { $items:tbranch' }|] - ([], _) -> - [C.cstm|if (!($exp:cond')) { $items:fbranch' }|] - _ -> - [C.cstm|if ($exp:cond') { $items:tbranch' } else { $items:fbranch' }|] -compileCode (Copy _ dest (Count destoffset) DefaultSpace src (Count srcoffset) DefaultSpace (Count size)) = - join $ - copyMemoryDefaultSpace - <$> rawMem dest - <*> compileExp (untyped destoffset) - <*> rawMem src - <*> compileExp (untyped srcoffset) - <*> compileExp (untyped size) -compileCode (Copy _ dest (Count destoffset) destspace src (Count srcoffset) srcspace (Count size)) = do - copy <- asks envCopy - join $ - copy CopyBarrier - <$> rawMem dest - <*> compileExp (untyped destoffset) - <*> pure destspace - <*> rawMem src - <*> compileExp (untyped srcoffset) - <*> pure srcspace - <*> compileExp (untyped size) -compileCode (Write _ _ Unit _ _ _) = pure () -compileCode (Write dest (Count idx) elemtype DefaultSpace vol elemexp) = do - dest' <- rawMem dest - deref <- - derefPointer dest' - <$> compileExp (untyped idx) - <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType elemtype)*|] - elemexp' <- toStorage elemtype <$> compileExp elemexp - stm [C.cstm|$exp:deref = $exp:elemexp';|] -compileCode (Write dest (Count idx) _ ScalarSpace {} _ elemexp) = do - idx' <- compileExp (untyped idx) - elemexp' <- compileExp elemexp - stm [C.cstm|$id:dest[$exp:idx'] = $exp:elemexp';|] -compileCode (Write dest (Count idx) elemtype (Space space) vol elemexp) = - join $ - asks envWriteScalar - <*> rawMem dest - <*> compileExp (untyped idx) - <*> pure (primStorageType elemtype) - <*> pure space - <*> pure vol - <*> (toStorage elemtype <$> compileExp elemexp) -compileCode (Read x _ _ Unit __ _) = - stm [C.cstm|$id:x = $exp:(UnitValue);|] -compileCode (Read x src (Count iexp) restype DefaultSpace vol) = do - src' <- rawMem src - e <- - fmap (fromStorage restype) $ - derefPointer src' - <$> compileExp (untyped iexp) - <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType restype)*|] - stm [C.cstm|$id:x = $exp:e;|] -compileCode (Read x src (Count iexp) restype (Space space) vol) = do - e <- - fmap (fromStorage restype) . join $ - asks envReadScalar - <*> rawMem src - <*> compileExp (untyped iexp) - <*> pure (primStorageType restype) - <*> pure space - <*> pure vol - stm [C.cstm|$id:x = $exp:e;|] -compileCode (Read x src (Count iexp) _ ScalarSpace {} _) = do - iexp' <- compileExp $ untyped iexp - stm [C.cstm|$id:x = $id:src[$exp:iexp'];|] -compileCode (DeclareMem name space) = - declMem name space -compileCode (DeclareScalar name vol t) = do - let ct = primTypeToCType t - decl [C.cdecl|$tyquals:(volQuals vol) $ty:ct $id:name;|] -compileCode (DeclareArray name ScalarSpace {} _ _) = - error $ "Cannot declare array " ++ pretty name ++ " in scalar space." -compileCode (DeclareArray name DefaultSpace t vs) = do - name_realtype <- newVName $ baseString name ++ "_realtype" - let ct = primTypeToCType t - case vs of - ArrayValues vs' -> do - let vs'' = [[C.cinit|$exp:v|] | v <- vs'] - earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs')] = {$inits:vs''};|] - ArrayZeros n -> - earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|] - -- Fake a memory block. - contextField - (C.toIdent name noLoc) - [C.cty|struct memblock|] - $ Just - [C.cexp|(struct memblock){NULL, - (unsigned char*)$id:name_realtype, - 0, - $string:(pretty name)}|] - item [C.citem|struct memblock $id:name = ctx->$id:name;|] -compileCode (DeclareArray name (Space space) t vs) = - join $ - asks envStaticArray - <*> pure name - <*> pure space - <*> pure t - <*> pure vs --- For assignments of the form 'x = x OP e', we generate C assignment --- operators to make the resulting code slightly nicer. This has no --- effect on performance. -compileCode (SetScalar dest (BinOpExp op (LeafExp x _) y)) - | dest == x, - Just f <- assignmentOperator op = do - y' <- compileExp y - stm [C.cstm|$exp:(f dest y');|] -compileCode (SetScalar dest src) = do - src' <- compileExp src - stm [C.cstm|$id:dest = $exp:src';|] -compileCode (SetMem dest src space) = - setMem dest src space -compileCode (Call dests fname args) = - join $ - asks (opsCall . envOperations) - <*> pure dests - <*> pure fname - <*> mapM compileArg args - where - compileArg (MemArg m) = pure [C.cexp|$exp:m|] - compileArg (ExpArg e) = compileExp e - -compileFunBody :: [C.Exp] -> [Param] -> Code op -> CompilerM op s () -compileFunBody output_ptrs outputs code = do - mapM_ declareOutput outputs - compileCode code - zipWithM_ setRetVal' output_ptrs outputs - where - declareOutput (MemParam name space) = - declMem name space - declareOutput (ScalarParam name pt) = do - let ctp = primTypeToCType pt - decl [C.cdecl|$ty:ctp $id:name;|] - - setRetVal' p (MemParam name space) = do - resetMem [C.cexp|*$exp:p|] space - setMem [C.cexp|*$exp:p|] name space - setRetVal' p (ScalarParam name _) = - stm [C.cstm|*$exp:p = $id:name;|] - -assignmentOperator :: BinOp -> Maybe (VName -> C.Exp -> C.Exp) -assignmentOperator Add {} = Just $ \d e -> [C.cexp|$id:d += $exp:e|] -assignmentOperator Sub {} = Just $ \d e -> [C.cexp|$id:d -= $exp:e|] -assignmentOperator Mul {} = Just $ \d e -> [C.cexp|$id:d *= $exp:e|] -assignmentOperator _ = Nothing +-- | C code generation for whole programs, built on +-- "Futhark.CodeGen.Backends.GenericC.Monad". Most of this module is +-- concerned with constructing the C API. +module Futhark.CodeGen.Backends.GenericC + ( compileProg, + compileProg', + compileFun, + defaultOperations, + CParts (..), + asLibrary, + asExecutable, + asServer, + module Futhark.CodeGen.Backends.GenericC.Monad, + module Futhark.CodeGen.Backends.GenericC.Code, + ) +where + +import Control.Monad.Reader +import Control.Monad.State +import qualified Data.DList as DL +import Data.Loc +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Text as T +import Futhark.CodeGen.Backends.GenericC.CLI (cliDefs) +import Futhark.CodeGen.Backends.GenericC.Code +import Futhark.CodeGen.Backends.GenericC.EntryPoints +import Futhark.CodeGen.Backends.GenericC.Monad +import Futhark.CodeGen.Backends.GenericC.Options +import Futhark.CodeGen.Backends.GenericC.Server (serverDefs) +import Futhark.CodeGen.Backends.GenericC.Types +import Futhark.CodeGen.ImpCode +import Futhark.CodeGen.RTS.C (cacheH, contextH, contextPrototypesH, errorsH, halfH, lockH, timingH, utilH) +import Futhark.IR.Prop (isBuiltInFunction) +import qualified Futhark.Manifest as Manifest +import Futhark.MonadFreshNames +import Futhark.Util.Pretty (prettyText) +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C +import NeatInterpolation (untrimming) + +defCall :: CallCompiler op s +defCall dests fname args = do + let out_args = [[C.cexp|&$id:d|] | d <- dests] + args' + | isBuiltInFunction fname = args + | otherwise = [C.cexp|ctx|] : out_args ++ args + case dests of + [dest] + | isBuiltInFunction fname -> + stm [C.cstm|$id:dest = $id:(funName fname)($args:args');|] + _ -> + item [C.citem|if ($id:(funName fname)($args:args') != 0) { err = 1; goto cleanup; }|] + +defError :: ErrorCompiler op s +defError msg stacktrace = do + (formatstr, formatargs) <- errorMsgString msg + let formatstr' = "Error: " <> formatstr <> "\n\nBacktrace:\n%s" + items + [C.citems|set_error(ctx, msgprintf($string:formatstr', $args:formatargs, $string:stacktrace)); + err = FUTHARK_PROGRAM_ERROR; + goto cleanup;|] + +-- | A set of operations that fail for every operation involving +-- non-default memory spaces. Uses plain pointers and @malloc@ for +-- memory management. +defaultOperations :: Operations op s +defaultOperations = + Operations + { opsWriteScalar = defWriteScalar, + opsReadScalar = defReadScalar, + opsAllocate = defAllocate, + opsDeallocate = defDeallocate, + opsCopy = defCopy, + opsStaticArray = defStaticArray, + opsMemoryType = defMemoryType, + opsCompiler = defCompiler, + opsFatMemory = True, + opsError = defError, + opsCall = defCall, + opsCritical = mempty + } + where + defWriteScalar _ _ _ _ _ = + error "Cannot write to non-default memory space because I am dumb" + defReadScalar _ _ _ _ = + error "Cannot read from non-default memory space" + defAllocate _ _ _ = + error "Cannot allocate in non-default memory space" + defDeallocate _ _ = + error "Cannot deallocate in non-default memory space" + defCopy _ destmem destoffset DefaultSpace srcmem srcoffset DefaultSpace size = + copyMemoryDefaultSpace destmem destoffset srcmem srcoffset size + defCopy _ _ _ _ _ _ _ _ = + error "Cannot copy to or from non-default memory space" + defStaticArray _ _ _ _ = + error "Cannot create static array in non-default memory space" + defMemoryType _ = + error "Has no type for non-default memory space" + defCompiler _ = + error "The default compiler cannot compile extended operations" + +compileFunBody :: [C.Exp] -> [Param] -> Code op -> CompilerM op s () +compileFunBody output_ptrs outputs code = do + mapM_ declareOutput outputs + compileCode code + zipWithM_ setRetVal' output_ptrs outputs + where + declareOutput (MemParam name space) = + declMem name space + declareOutput (ScalarParam name pt) = do + let ctp = primTypeToCType pt + decl [C.cdecl|$ty:ctp $id:name;|] + + setRetVal' p (MemParam name space) = do + resetMem [C.cexp|*$exp:p|] space + setMem [C.cexp|*$exp:p|] name space + setRetVal' p (ScalarParam name _) = + stm [C.cstm|*$exp:p = $id:name;|] + +compileFun :: [C.BlockItem] -> [C.Param] -> (Name, Function op) -> CompilerM op s (C.Definition, C.Func) +compileFun get_constants extra (fname, func@(Function _ outputs inputs body)) = inNewFunction $ do + (outparams, out_ptrs) <- unzip <$> mapM compileOutput outputs + inparams <- mapM compileInput inputs + + cachingMemory (lexicalMemoryUsage func) $ \decl_cached free_cached -> do + body' <- collect $ compileFunBody out_ptrs outputs body + decl_mem <- declAllocatedMem + free_mem <- freeAllocatedMem + + pure + ( [C.cedecl|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams);|], + [C.cfun|static int $id:(funName fname)($params:extra, $params:outparams, $params:inparams) { + $stms:ignores + int err = 0; + $items:decl_cached + $items:decl_mem + $items:get_constants + $items:body' + cleanup: + { + $stms:free_cached + $items:free_mem + } + return err; + }|] + ) + where + -- Ignore all the boilerplate parameters, just in case we don't + -- actually need to use them. + ignores = [[C.cstm|(void)$id:p;|] | C.Param (Just p) _ _ _ <- extra] + + compileInput (ScalarParam name bt) = do + let ctp = primTypeToCType bt + pure [C.cparam|$ty:ctp $id:name|] + compileInput (MemParam name space) = do + ty <- memToCType name space + pure [C.cparam|$ty:ty $id:name|] + + compileOutput (ScalarParam name bt) = do + let ctp = primTypeToCType bt + p_name <- newVName $ "out_" ++ baseString name + pure ([C.cparam|$ty:ctp *$id:p_name|], [C.cexp|$id:p_name|]) + compileOutput (MemParam name space) = do + ty <- memToCType name space + p_name <- newVName $ baseString name ++ "_p" + pure ([C.cparam|$ty:ty *$id:p_name|], [C.cexp|$id:p_name|]) + +declsCode :: (HeaderSection -> Bool) -> CompilerState s -> T.Text +declsCode p = + T.unlines + . map prettyText + . concatMap (DL.toList . snd) + . filter (p . fst) + . M.toList + . compHeaderDecls + +initDecls, arrayDecls, opaqueDecls, opaqueTypeDecls, entryDecls, miscDecls :: CompilerState s -> T.Text +initDecls = declsCode (== InitDecl) +arrayDecls = declsCode isArrayDecl + where + isArrayDecl ArrayDecl {} = True + isArrayDecl _ = False +opaqueTypeDecls = declsCode isOpaqueTypeDecl + where + isOpaqueTypeDecl OpaqueTypeDecl {} = True + isOpaqueTypeDecl _ = False +opaqueDecls = declsCode isOpaqueDecl + where + isOpaqueDecl OpaqueDecl {} = True + isOpaqueDecl _ = False +entryDecls = declsCode (== EntryDecl) +miscDecls = declsCode (== MiscDecl) + +defineMemorySpace :: Space -> CompilerM op s (C.Definition, [C.Definition], C.BlockItem) +defineMemorySpace space = do + rm <- rawMemCType space + let structdef = + [C.cedecl|struct $id:sname { int *references; + $ty:rm mem; + typename int64_t size; + const char *desc; };|] + + contextField peakname [C.cty|typename int64_t|] $ Just [C.cexp|0|] + contextField usagename [C.cty|typename int64_t|] $ Just [C.cexp|0|] + + -- Unreferencing a memory block consists of decreasing its reference + -- count and freeing the corresponding memory if the count reaches + -- zero. + free <- collect $ freeRawMem [C.cexp|block->mem|] space [C.cexp|desc|] + ctx_ty <- contextType + let unrefdef = + [C.cedecl|int $id:(fatMemUnRef space) ($ty:ctx_ty *ctx, $ty:mty *block, const char *desc) { + if (block->references != NULL) { + *(block->references) -= 1; + if (ctx->detail_memory) { + fprintf(ctx->log, "Unreferencing block %s (allocated as %s) in %s: %d references remaining.\n", + desc, block->desc, $string:spacedesc, *(block->references)); + } + if (*(block->references) == 0) { + ctx->$id:usagename -= block->size; + $items:free + free(block->references); + if (ctx->detail_memory) { + fprintf(ctx->log, "%lld bytes freed (now allocated: %lld bytes)\n", + (long long) block->size, (long long) ctx->$id:usagename); + } + } + block->references = NULL; + } + return 0; +}|] + + -- When allocating a memory block we initialise the reference count to 1. + alloc <- + collect $ + allocRawMem [C.cexp|block->mem|] [C.cexp|size|] space [C.cexp|desc|] + let allocdef = + [C.cedecl|int $id:(fatMemAlloc space) ($ty:ctx_ty *ctx, $ty:mty *block, typename int64_t size, const char *desc) { + if (size < 0) { + futhark_panic(1, "Negative allocation of %lld bytes attempted for %s in %s.\n", + (long long)size, desc, $string:spacedesc, ctx->$id:usagename); + } + int ret = $id:(fatMemUnRef space)(ctx, block, desc); + + if (ret != FUTHARK_SUCCESS) { + return ret; + } + + if (ctx->detail_memory) { + fprintf(ctx->log, "Allocating %lld bytes for %s in %s (then allocated: %lld bytes)", + (long long) size, + desc, $string:spacedesc, + (long long) ctx->$id:usagename + size); + } + if (ctx->$id:usagename > ctx->$id:peakname) { + ctx->$id:peakname = ctx->$id:usagename; + if (ctx->detail_memory) { + fprintf(ctx->log, " (new peak).\n"); + } + } else if (ctx->detail_memory) { + fprintf(ctx->log, ".\n"); + } + + $items:alloc + + if (ctx->error == NULL) { + block->references = (int*) malloc(sizeof(int)); + *(block->references) = 1; + block->size = size; + block->desc = desc; + ctx->$id:usagename += size; + return FUTHARK_SUCCESS; + } else { + // We are naively assuming that any memory allocation error is due to OOM. + // We preserve the original error so that a savvy user can perhaps find + // glory despite our naiveté. + + char *old_error = ctx->error; + set_error(ctx, msgprintf("Failed to allocate memory in %s.\nAttempted allocation: %12lld bytes\nCurrently allocated: %12lld bytes\n%s", + $string:spacedesc, + (long long) size, + (long long) ctx->$id:usagename, + old_error)); + free(old_error); + return FUTHARK_OUT_OF_MEMORY; + } + }|] + + -- Memory setting - unreference the destination and increase the + -- count of the source by one. + let setdef = + [C.cedecl|int $id:(fatMemSet space) ($ty:ctx_ty *ctx, $ty:mty *lhs, $ty:mty *rhs, const char *lhs_desc) { + int ret = $id:(fatMemUnRef space)(ctx, lhs, lhs_desc); + if (rhs->references != NULL) { + (*(rhs->references))++; + } + *lhs = *rhs; + return ret; +} +|] + + onClear [C.citem|ctx->$id:peakname = 0;|] + + let peakmsg = "Peak memory usage for " ++ spacedesc ++ ": %lld bytes.\n" + pure + ( structdef, + [unrefdef, allocdef, setdef], + -- Do not report memory usage for DefaultSpace (CPU memory), + -- because it would not be accurate anyway. This whole + -- tracking probably needs to be rethought. + if space == DefaultSpace + then [C.citem|{}|] + else [C.citem|str_builder(&builder, $string:peakmsg, (long long) ctx->$id:peakname);|] + ) + where + mty = fatMemType space + (peakname, usagename, sname, spacedesc) = case space of + Space sid -> + ( C.toIdent ("peak_mem_usage_" ++ sid) noLoc, + C.toIdent ("cur_mem_usage_" ++ sid) noLoc, + C.toIdent ("memblock_" ++ sid) noLoc, + "space '" ++ sid ++ "'" + ) + _ -> + ( "peak_mem_usage_default", + "cur_mem_usage_default", + "memblock", + "default space" + ) + +-- | The result of compilation to C is multiple parts, which can be +-- put together in various ways. The obvious way is to concatenate +-- all of them, which yields a CLI program. Another is to compile the +-- library part by itself, and use the header file to call into it. +data CParts = CParts + { cHeader :: T.Text, + -- | Utility definitions that must be visible + -- to both CLI and library parts. + cUtils :: T.Text, + cCLI :: T.Text, + cServer :: T.Text, + cLib :: T.Text, + -- | The manifest, in JSON format. + cJsonManifest :: T.Text + } + +gnuSource :: T.Text +gnuSource = + [untrimming| +// We need to define _GNU_SOURCE before +// _any_ headers files are imported to get +// the usage statistics of a thread (i.e. have RUSAGE_THREAD) on GNU/Linux +// https://manpages.courier-mta.org/htmlman2/getrusage.2.html +#ifndef _GNU_SOURCE // Avoid possible double-definition warning. +#define _GNU_SOURCE +#endif +|] + +-- We may generate variables that are never used (e.g. for +-- certificates) or functions that are never called (e.g. unused +-- intrinsics), and generated code may have other cosmetic issues that +-- compilers warn about. We disable these warnings to not clutter the +-- compilation logs. +disableWarnings :: T.Text +disableWarnings = + [untrimming| +#ifdef __clang__ +#pragma clang diagnostic ignored "-Wunused-function" +#pragma clang diagnostic ignored "-Wunused-variable" +#pragma clang diagnostic ignored "-Wparentheses" +#pragma clang diagnostic ignored "-Wunused-label" +#elif __GNUC__ +#pragma GCC diagnostic ignored "-Wunused-function" +#pragma GCC diagnostic ignored "-Wunused-variable" +#pragma GCC diagnostic ignored "-Wparentheses" +#pragma GCC diagnostic ignored "-Wunused-label" +#pragma GCC diagnostic ignored "-Wunused-but-set-variable" +#endif +|] + +-- | Produce header, implementation, and manifest files. +asLibrary :: CParts -> (T.Text, T.Text, T.Text) +asLibrary parts = + ( "#pragma once\n\n" <> cHeader parts, + gnuSource <> disableWarnings <> cHeader parts <> cUtils parts <> cLib parts, + cJsonManifest parts + ) + +-- | As executable with command-line interface. +asExecutable :: CParts -> T.Text +asExecutable parts = + gnuSource <> disableWarnings <> cHeader parts <> cUtils parts <> cCLI parts <> cLib parts + +-- | As server executable. +asServer :: CParts -> T.Text +asServer parts = + gnuSource <> disableWarnings <> cHeader parts <> cUtils parts <> cServer parts <> cLib parts + +compileProg' :: + MonadFreshNames m => + T.Text -> + T.Text -> + Operations op s -> + s -> + CompilerM op s () -> + T.Text -> + (Space, [Space]) -> + [Option] -> + Definitions op -> + m (CParts, CompilerState s) +compileProg' backend version ops def extra header_extra (arr_space, spaces) options prog = do + src <- getNameSource + let ((prototypes, definitions, entry_point_decls, manifest), endstate) = + runCompilerM ops src def compileProgAction + initdecls = initDecls endstate + entrydecls = entryDecls endstate + arraydecls = arrayDecls endstate + opaquetypedecls = opaqueTypeDecls endstate + opaquedecls = opaqueDecls endstate + miscdecls = miscDecls endstate + + let headerdefs = + [untrimming| +// Headers +#include +#include +#include +#include +#include +$header_extra +#ifdef __cplusplus +extern "C" { +#endif + +// Initialisation +$initdecls + +// Arrays +$arraydecls + +// Opaque values +$opaquetypedecls +$opaquedecls + +// Entry points +$entrydecls + +// Miscellaneous +$miscdecls +#define FUTHARK_BACKEND_$backend +$errorsH + +#ifdef __cplusplus +} +#endif +|] + + let utildefs = + [untrimming| +#include +#include +#include +#include +#include +// If NDEBUG is set, the assert() macro will do nothing. Since Futhark +// (unfortunately) makes use of assert() for error detection (and even some +// side effects), we want to avoid that. +#undef NDEBUG +#include +#include +$utilH +$cacheH +$halfH +$timingH +|] + + let early_decls = T.unlines $ map prettyText $ DL.toList $ compEarlyDecls endstate + lib_decls = T.unlines $ map prettyText $ DL.toList $ compLibDecls endstate + clidefs = cliDefs options manifest + serverdefs = serverDefs options manifest + libdefs = + [untrimming| +#ifdef _MSC_VER +#define inline __inline +#endif +#include +#include +#include +#include +#include + +$header_extra + +$lockH + +#define FUTHARK_F64_ENABLED + +$cScalarDefs + +$contextPrototypesH + +$early_decls + +$contextH + +$prototypes + +$lib_decls + +$definitions + +$entry_point_decls + |] + + pure + ( CParts + { cHeader = headerdefs, + cUtils = utildefs, + cCLI = clidefs, + cServer = serverdefs, + cLib = libdefs, + cJsonManifest = Manifest.manifestToJSON manifest + }, + endstate + ) + where + Definitions types consts (Functions funs) = prog + + compileProgAction = do + (memstructs, memfuns, memreport) <- unzip3 <$> mapM defineMemorySpace spaces + + get_consts <- compileConstants consts + + ctx_ty <- contextType + + (prototypes, functions) <- + unzip <$> mapM (compileFun get_consts [[C.cparam|$ty:ctx_ty *ctx|]]) funs + + mapM_ earlyDecl memstructs + (entry_points, entry_points_manifest) <- + unzip . catMaybes <$> mapM (uncurry (onEntryPoint get_consts)) funs + + extra + + mapM_ earlyDecl $ concat memfuns + + type_funs <- generateAPITypes arr_space types + generateCommonLibFuns memreport + + pure + ( T.unlines $ map prettyText prototypes, + T.unlines $ map (prettyText . funcToDef) functions, + T.unlines $ map prettyText entry_points, + Manifest.Manifest (M.fromList entry_points_manifest) type_funs backend version + ) + + funcToDef func = C.FuncDef func loc + where + loc = case func of + C.OldFunc _ _ _ _ _ _ l -> l + C.Func _ _ _ _ _ l -> l + +-- | Compile imperative program to a C program. Always uses the +-- function named "main" as entry point, so make sure it is defined. +compileProg :: + MonadFreshNames m => + T.Text -> + T.Text -> + Operations op () -> + CompilerM op () () -> + T.Text -> + (Space, [Space]) -> + [Option] -> + Definitions op -> + m CParts +compileProg backend version ops extra header_extra (arr_space, spaces) options prog = + fst <$> compileProg' backend version ops () extra header_extra (arr_space, spaces) options prog + +generateCommonLibFuns :: [C.BlockItem] -> CompilerM op s () +generateCommonLibFuns memreport = do + ctx <- contextType + cfg <- configType + ops <- asks envOperations + profilereport <- gets $ DL.toList . compProfileItems + + publicDef_ "context_config_set_cache_file" MiscDecl $ \s -> + ( [C.cedecl|void $id:s($ty:cfg* cfg, const char *f);|], + [C.cedecl|void $id:s($ty:cfg* cfg, const char *f) { + cfg->cache_fname = f; + }|] + ) + + publicDef_ "get_tuning_param_count" InitDecl $ \s -> + ( [C.cedecl|int $id:s(void);|], + [C.cedecl|int $id:s(void) { + return sizeof(tuning_param_names)/sizeof(tuning_param_names[0]); + }|] + ) + + publicDef_ "get_tuning_param_name" InitDecl $ \s -> + ( [C.cedecl|const char* $id:s(int);|], + [C.cedecl|const char* $id:s(int i) { + return tuning_param_names[i]; + }|] + ) + + publicDef_ "get_tuning_param_class" InitDecl $ \s -> + ( [C.cedecl|const char* $id:s(int);|], + [C.cedecl|const char* $id:s(int i) { + return tuning_param_classes[i]; + }|] + ) + + sync <- publicName "context_sync" + publicDef_ "context_report" MiscDecl $ \s -> + ( [C.cedecl|char* $id:s($ty:ctx *ctx);|], + [C.cedecl|char* $id:s($ty:ctx *ctx) { + if ($id:sync(ctx) != 0) { + return NULL; + } + + struct str_builder builder; + str_builder_init(&builder); + $items:memreport + if (ctx->profiling) { + $items:profilereport + } + return builder.str; + }|] + ) + + publicDef_ "context_get_error" MiscDecl $ \s -> + ( [C.cedecl|char* $id:s($ty:ctx* ctx);|], + [C.cedecl|char* $id:s($ty:ctx* ctx) { + char* error = ctx->error; + ctx->error = NULL; + return error; + }|] + ) + + publicDef_ "context_set_logging_file" MiscDecl $ \s -> + ( [C.cedecl|void $id:s($ty:ctx* ctx, typename FILE* f);|], + [C.cedecl|void $id:s($ty:ctx* ctx, typename FILE* f) { + ctx->log = f; + }|] + ) + + publicDef_ "context_pause_profiling" MiscDecl $ \s -> + ( [C.cedecl|void $id:s($ty:ctx* ctx);|], + [C.cedecl|void $id:s($ty:ctx* ctx) { + ctx->profiling_paused = 1; + }|] + ) + + publicDef_ "context_unpause_profiling" MiscDecl $ \s -> + ( [C.cedecl|void $id:s($ty:ctx* ctx);|], + [C.cedecl|void $id:s($ty:ctx* ctx) { + ctx->profiling_paused = 0; + }|] + ) + + clears <- gets $ DL.toList . compClearItems + publicDef_ "context_clear_caches" MiscDecl $ \s -> + ( [C.cedecl|int $id:s($ty:ctx* ctx);|], + [C.cedecl|int $id:s($ty:ctx* ctx) { + $items:(criticalSection ops clears) + return ctx->error != NULL; + }|] + ) + +compileConstants :: Constants op -> CompilerM op s [C.BlockItem] +compileConstants (Constants ps init_consts) = do + ctx_ty <- contextType + const_fields <- mapM constParamField ps + -- Avoid an empty struct, as that is apparently undefined behaviour. + let const_fields' + | null const_fields = [[C.csdecl|int dummy;|]] + | otherwise = const_fields + contextField "constants" [C.cty|struct { $sdecls:const_fields' }|] Nothing + earlyDecl [C.cedecl|static int init_constants($ty:ctx_ty*);|] + earlyDecl [C.cedecl|static int free_constants($ty:ctx_ty*);|] + + inNewFunction $ do + -- We locally define macros for the constants, so that when we + -- generate assignments to local variables, we actually assign into + -- the constants struct. This is not needed for functions, because + -- they can only read constants, not write them. + let (defs, undefs) = unzip $ map constMacro ps + init_consts' <- collect $ do + mapM_ resetMemConst ps + compileCode init_consts + decl_mem <- declAllocatedMem + free_mem <- freeAllocatedMem + libDecl + [C.cedecl|static int init_constants($ty:ctx_ty *ctx) { + (void)ctx; + int err = 0; + $items:defs + $items:decl_mem + $items:init_consts' + $items:free_mem + $items:undefs + cleanup: + return err; + }|] + + inNewFunction $ do + free_consts <- collect $ mapM_ freeConst ps + libDecl + [C.cedecl|static int free_constants($ty:ctx_ty *ctx) { + (void)ctx; + $items:free_consts + return 0; + }|] + + mapM getConst ps + where + constParamField (ScalarParam name bt) = do + let ctp = primTypeToCType bt + pure [C.csdecl|$ty:ctp $id:name;|] + constParamField (MemParam name space) = do + ty <- memToCType name space + pure [C.csdecl|$ty:ty $id:name;|] + + constMacro p = ([C.citem|$escstm:def|], [C.citem|$escstm:undef|]) + where + p' = pretty (C.toIdent (paramName p) mempty) + def = "#define " ++ p' ++ " (" ++ "ctx->constants." ++ p' ++ ")" + undef = "#undef " ++ p' + + resetMemConst ScalarParam {} = pure () + resetMemConst (MemParam name space) = resetMem name space + + freeConst ScalarParam {} = pure () + freeConst (MemParam name space) = unRefMem [C.cexp|ctx->constants.$id:name|] space + + getConst (ScalarParam name bt) = do + let ctp = primTypeToCType bt + pure [C.citem|$ty:ctp $id:name = ctx->constants.$id:name;|] + getConst (MemParam name space) = do + ty <- memToCType name space + pure [C.citem|$ty:ty $id:name = ctx->constants.$id:name;|] diff --git a/src/Futhark/CodeGen/Backends/GenericC/Code.hs b/src/Futhark/CodeGen/Backends/GenericC/Code.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericC/Code.hs @@ -0,0 +1,356 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} + +-- | Translation of ImpCode Exp and Code to C. +module Futhark.CodeGen.Backends.GenericC.Code + ( compilePrimExp, + compileExp, + compileExpToName, + compileCode, + errorMsgString, + linearCode, + ) +where + +import Control.Monad.Reader +import Data.Loc +import Data.Maybe +import Futhark.CodeGen.Backends.GenericC.Monad +import Futhark.CodeGen.ImpCode +import Futhark.MonadFreshNames +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C + +errorMsgString :: ErrorMsg Exp -> CompilerM op s (String, [C.Exp]) +errorMsgString (ErrorMsg parts) = do + let boolStr e = [C.cexp|($exp:e) ? "true" : "false"|] + asLongLong e = [C.cexp|(long long int)$exp:e|] + asDouble e = [C.cexp|(double)$exp:e|] + onPart (ErrorString s) = pure ("%s", [C.cexp|$string:s|]) + onPart (ErrorVal Bool x) = ("%s",) . boolStr <$> compileExp x + onPart (ErrorVal Unit _) = pure ("%s", [C.cexp|"()"|]) + onPart (ErrorVal (IntType Int8) x) = ("%hhd",) <$> compileExp x + onPart (ErrorVal (IntType Int16) x) = ("%hd",) <$> compileExp x + onPart (ErrorVal (IntType Int32) x) = ("%d",) <$> compileExp x + onPart (ErrorVal (IntType Int64) x) = ("%lld",) . asLongLong <$> compileExp x + onPart (ErrorVal (FloatType Float16) x) = ("%f",) . asDouble <$> compileExp x + onPart (ErrorVal (FloatType Float32) x) = ("%f",) . asDouble <$> compileExp x + onPart (ErrorVal (FloatType Float64) x) = ("%f",) <$> compileExp x + (formatstrs, formatargs) <- unzip <$> mapM onPart parts + pure (mconcat formatstrs, formatargs) + +compileExpToName :: String -> PrimType -> Exp -> CompilerM op s VName +compileExpToName _ _ (LeafExp v _) = + pure v +compileExpToName desc t e = do + desc' <- newVName desc + e' <- compileExp e + decl [C.cdecl|$ty:(primTypeToCType t) $id:desc' = $e';|] + pure desc' + +compileExp :: Exp -> CompilerM op s C.Exp +compileExp = compilePrimExp $ \v -> pure [C.cexp|$id:v|] + +-- | Tell me how to compile a @v@, and I'll Compile any @PrimExp v@ for you. +compilePrimExp :: Monad m => (v -> m C.Exp) -> PrimExp v -> m C.Exp +compilePrimExp _ (ValueExp val) = + pure $ C.toExp val mempty +compilePrimExp f (LeafExp v _) = + f v +compilePrimExp f (UnOpExp Complement {} x) = do + x' <- compilePrimExp f x + pure [C.cexp|~$exp:x'|] +compilePrimExp f (UnOpExp Not {} x) = do + x' <- compilePrimExp f x + pure [C.cexp|!$exp:x'|] +compilePrimExp f (UnOpExp (FAbs Float32) x) = do + x' <- compilePrimExp f x + pure [C.cexp|(float)fabs($exp:x')|] +compilePrimExp f (UnOpExp (FAbs Float64) x) = do + x' <- compilePrimExp f x + pure [C.cexp|fabs($exp:x')|] +compilePrimExp f (UnOpExp SSignum {} x) = do + x' <- compilePrimExp f x + pure [C.cexp|($exp:x' > 0 ? 1 : 0) - ($exp:x' < 0 ? 1 : 0)|] +compilePrimExp f (UnOpExp USignum {} x) = do + x' <- compilePrimExp f x + pure [C.cexp|($exp:x' > 0 ? 1 : 0) - ($exp:x' < 0 ? 1 : 0) != 0|] +compilePrimExp f (UnOpExp op x) = do + x' <- compilePrimExp f x + pure [C.cexp|$id:(pretty op)($exp:x')|] +compilePrimExp f (CmpOpExp cmp x y) = do + x' <- compilePrimExp f x + y' <- compilePrimExp f y + pure $ case cmp of + CmpEq {} -> [C.cexp|$exp:x' == $exp:y'|] + FCmpLt {} -> [C.cexp|$exp:x' < $exp:y'|] + FCmpLe {} -> [C.cexp|$exp:x' <= $exp:y'|] + CmpLlt {} -> [C.cexp|$exp:x' < $exp:y'|] + CmpLle {} -> [C.cexp|$exp:x' <= $exp:y'|] + _ -> [C.cexp|$id:(pretty cmp)($exp:x', $exp:y')|] +compilePrimExp f (ConvOpExp conv x) = do + x' <- compilePrimExp f x + pure [C.cexp|$id:(pretty conv)($exp:x')|] +compilePrimExp f (BinOpExp bop x y) = do + x' <- compilePrimExp f x + y' <- compilePrimExp f y + -- Note that integer addition, subtraction, and multiplication with + -- OverflowWrap are not handled by explicit operators, but rather by + -- functions. This is because we want to implicitly convert them to + -- unsigned numbers, so we can do overflow without invoking + -- undefined behaviour. + pure $ case bop of + Add _ OverflowUndef -> [C.cexp|$exp:x' + $exp:y'|] + Sub _ OverflowUndef -> [C.cexp|$exp:x' - $exp:y'|] + Mul _ OverflowUndef -> [C.cexp|$exp:x' * $exp:y'|] + FAdd {} -> [C.cexp|$exp:x' + $exp:y'|] + FSub {} -> [C.cexp|$exp:x' - $exp:y'|] + FMul {} -> [C.cexp|$exp:x' * $exp:y'|] + FDiv {} -> [C.cexp|$exp:x' / $exp:y'|] + Xor {} -> [C.cexp|$exp:x' ^ $exp:y'|] + And {} -> [C.cexp|$exp:x' & $exp:y'|] + Or {} -> [C.cexp|$exp:x' | $exp:y'|] + LogAnd {} -> [C.cexp|$exp:x' && $exp:y'|] + LogOr {} -> [C.cexp|$exp:x' || $exp:y'|] + _ -> [C.cexp|$id:(pretty bop)($exp:x', $exp:y')|] +compilePrimExp f (FunExp h args _) = do + args' <- mapM (compilePrimExp f) args + pure [C.cexp|$id:(funName (nameFromString h))($args:args')|] + +linearCode :: Code op -> [Code op] +linearCode = reverse . go [] + where + go acc (x :>>: y) = + go (go acc x) y + go acc x = x : acc + +assignmentOperator :: BinOp -> Maybe (VName -> C.Exp -> C.Exp) +assignmentOperator Add {} = Just $ \d e -> [C.cexp|$id:d += $exp:e|] +assignmentOperator Sub {} = Just $ \d e -> [C.cexp|$id:d -= $exp:e|] +assignmentOperator Mul {} = Just $ \d e -> [C.cexp|$id:d *= $exp:e|] +assignmentOperator _ = Nothing + +compileCode :: Code op -> CompilerM op s () +compileCode (Op op) = + join $ asks (opsCompiler . envOperations) <*> pure op +compileCode Skip = pure () +compileCode (Comment s code) = do + xs <- collect $ compileCode code + let comment = "// " ++ s + stm + [C.cstm|$comment:comment + { $items:xs } + |] +compileCode (TracePrint msg) = do + (formatstr, formatargs) <- errorMsgString msg + stm [C.cstm|fprintf(ctx->log, $string:formatstr, $args:formatargs);|] +compileCode (DebugPrint s (Just e)) = do + e' <- compileExp e + stm + [C.cstm|if (ctx->debugging) { + fprintf(ctx->log, $string:fmtstr, $exp:s, ($ty:ety)$exp:e', '\n'); + }|] + where + (fmt, ety) = case primExpType e of + IntType _ -> ("llu", [C.cty|long long int|]) + FloatType _ -> ("f", [C.cty|double|]) + _ -> ("d", [C.cty|int|]) + fmtstr = "%s: %" ++ fmt ++ "%c" +compileCode (DebugPrint s Nothing) = + stm + [C.cstm|if (ctx->debugging) { + fprintf(ctx->log, "%s\n", $exp:s); + }|] +-- :>>: is treated in a special way to detect declare-set pairs in +-- order to generate prettier code. +compileCode (c1 :>>: c2) = go (linearCode (c1 :>>: c2)) + where + go (DeclareScalar name vol t : SetScalar dest e : code) + | name == dest = do + let ct = primTypeToCType t + e' <- compileExp e + item [C.citem|$tyquals:(volQuals vol) $ty:ct $id:name = $exp:e';|] + go code + go (x : xs) = compileCode x >> go xs + go [] = pure () +compileCode (Assert e msg (loc, locs)) = do + e' <- compileExp e + err <- + collect . join $ + asks (opsError . envOperations) <*> pure msg <*> pure stacktrace + stm [C.cstm|if (!$exp:e') { $items:err }|] + where + stacktrace = prettyStacktrace 0 $ map locStr $ loc : locs +compileCode (Allocate _ _ ScalarSpace {}) = + -- Handled by the declaration of the memory block, which is + -- translated to an actual array. + pure () +compileCode (Allocate name (Count (TPrimExp e)) space) = do + size <- compileExp e + cached <- cacheMem name + case cached of + Just cur_size -> + stm + [C.cstm|if ($exp:cur_size < $exp:size) { + err = lexical_realloc(ctx, &$exp:name, &$exp:cur_size, $exp:size); + if (err != FUTHARK_SUCCESS) { + goto cleanup; + } + }|] + _ -> + allocMem name size space [C.cstm|{err = 1; goto cleanup;}|] +compileCode (Free name space) = do + cached <- isJust <$> cacheMem name + unless cached $ unRefMem name space +compileCode (For i bound body) = do + let i' = C.toIdent i + t = primTypeToCType $ primExpType bound + bound' <- compileExp bound + body' <- collect $ compileCode body + stm + [C.cstm|for ($ty:t $id:i' = 0; $id:i' < $exp:bound'; $id:i'++) { + $items:body' + }|] +compileCode (While cond body) = do + cond' <- compileExp $ untyped cond + body' <- collect $ compileCode body + stm + [C.cstm|while ($exp:cond') { + $items:body' + }|] +compileCode (If cond tbranch fbranch) = do + cond' <- compileExp $ untyped cond + tbranch' <- collect $ compileCode tbranch + fbranch' <- collect $ compileCode fbranch + stm $ case (tbranch', fbranch') of + (_, []) -> + [C.cstm|if ($exp:cond') { $items:tbranch' }|] + ([], _) -> + [C.cstm|if (!($exp:cond')) { $items:fbranch' }|] + (_, [C.BlockStm x@C.If {}]) -> + [C.cstm|if ($exp:cond') { $items:tbranch' } else $stm:x|] + _ -> + [C.cstm|if ($exp:cond') { $items:tbranch' } else { $items:fbranch' }|] +compileCode (Copy _ dest (Count destoffset) DefaultSpace src (Count srcoffset) DefaultSpace (Count size)) = + join $ + copyMemoryDefaultSpace + <$> rawMem dest + <*> compileExp (untyped destoffset) + <*> rawMem src + <*> compileExp (untyped srcoffset) + <*> compileExp (untyped size) +compileCode (Copy _ dest (Count destoffset) destspace src (Count srcoffset) srcspace (Count size)) = do + copy <- asks $ opsCopy . envOperations + join $ + copy CopyBarrier + <$> rawMem dest + <*> compileExp (untyped destoffset) + <*> pure destspace + <*> rawMem src + <*> compileExp (untyped srcoffset) + <*> pure srcspace + <*> compileExp (untyped size) +compileCode (Write _ _ Unit _ _ _) = pure () +compileCode (Write dest (Count idx) elemtype DefaultSpace vol elemexp) = do + dest' <- rawMem dest + deref <- + derefPointer dest' + <$> compileExp (untyped idx) + <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType elemtype)*|] + elemexp' <- toStorage elemtype <$> compileExp elemexp + stm [C.cstm|$exp:deref = $exp:elemexp';|] +compileCode (Write dest (Count idx) _ ScalarSpace {} _ elemexp) = do + idx' <- compileExp (untyped idx) + elemexp' <- compileExp elemexp + stm [C.cstm|$id:dest[$exp:idx'] = $exp:elemexp';|] +compileCode (Write dest (Count idx) elemtype (Space space) vol elemexp) = + join $ + asks (opsWriteScalar . envOperations) + <*> rawMem dest + <*> compileExp (untyped idx) + <*> pure (primStorageType elemtype) + <*> pure space + <*> pure vol + <*> (toStorage elemtype <$> compileExp elemexp) +compileCode (Read x _ _ Unit __ _) = + stm [C.cstm|$id:x = $exp:(UnitValue);|] +compileCode (Read x src (Count iexp) restype DefaultSpace vol) = do + src' <- rawMem src + e <- + fmap (fromStorage restype) $ + derefPointer src' + <$> compileExp (untyped iexp) + <*> pure [C.cty|$tyquals:(volQuals vol) $ty:(primStorageType restype)*|] + stm [C.cstm|$id:x = $exp:e;|] +compileCode (Read x src (Count iexp) restype (Space space) vol) = do + e <- + fmap (fromStorage restype) . join $ + asks (opsReadScalar . envOperations) + <*> rawMem src + <*> compileExp (untyped iexp) + <*> pure (primStorageType restype) + <*> pure space + <*> pure vol + stm [C.cstm|$id:x = $exp:e;|] +compileCode (Read x src (Count iexp) _ ScalarSpace {} _) = do + iexp' <- compileExp $ untyped iexp + stm [C.cstm|$id:x = $id:src[$exp:iexp'];|] +compileCode (DeclareMem name space) = + declMem name space +compileCode (DeclareScalar name vol t) = do + let ct = primTypeToCType t + decl [C.cdecl|$tyquals:(volQuals vol) $ty:ct $id:name;|] +compileCode (DeclareArray name ScalarSpace {} _ _) = + error $ "Cannot declare array " ++ pretty name ++ " in scalar space." +compileCode (DeclareArray name DefaultSpace t vs) = do + name_realtype <- newVName $ baseString name ++ "_realtype" + let ct = primTypeToCType t + case vs of + ArrayValues vs' -> do + let vs'' = [[C.cinit|$exp:v|] | v <- vs'] + earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:(length vs')] = {$inits:vs''};|] + ArrayZeros n -> + earlyDecl [C.cedecl|static $ty:ct $id:name_realtype[$int:n];|] + -- Fake a memory block. + contextField + (C.toIdent name noLoc) + [C.cty|struct memblock|] + $ Just + [C.cexp|(struct memblock){NULL, + (unsigned char*)$id:name_realtype, + 0, + $string:(pretty name)}|] + item [C.citem|struct memblock $id:name = ctx->$id:name;|] +compileCode (DeclareArray name (Space space) t vs) = + join $ + asks (opsStaticArray . envOperations) + <*> pure name + <*> pure space + <*> pure t + <*> pure vs +-- For assignments of the form 'x = x OP e', we generate C assignment +-- operators to make the resulting code slightly nicer. This has no +-- effect on performance. +compileCode (SetScalar dest (BinOpExp op (LeafExp x _) y)) + | dest == x, + Just f <- assignmentOperator op = do + y' <- compileExp y + stm [C.cstm|$exp:(f dest y');|] +compileCode (SetScalar dest src) = do + src' <- compileExp src + stm [C.cstm|$id:dest = $exp:src';|] +compileCode (SetMem dest src space) = + setMem dest src space +compileCode (Call dests fname args) = + join $ + asks (opsCall . envOperations) + <*> pure dests + <*> pure fname + <*> mapM compileArg args + where + compileArg (MemArg m) = pure [C.cexp|$exp:m|] + compileArg (ExpArg e) = compileExp e diff --git a/src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs b/src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericC/EntryPoints.hs @@ -0,0 +1,264 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} + +-- | Generate the entry point packing/unpacking code. +module Futhark.CodeGen.Backends.GenericC.EntryPoints + ( onEntryPoint, + ) +where + +import Control.Monad.Reader +import Data.Char (isAlpha, isAlphaNum) +import Data.Maybe +import qualified Data.Text as T +import Futhark.CodeGen.Backends.GenericC.Monad +import Futhark.CodeGen.Backends.GenericC.Types (opaqueToCType, valueTypeToCType) +import Futhark.CodeGen.ImpCode +import qualified Futhark.Manifest as Manifest +import Futhark.Util (zEncodeString) +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C + +valueDescToType :: ValueDesc -> ValueType +valueDescToType (ScalarValue pt signed _) = + ValueType signed (Rank 0) pt +valueDescToType (ArrayValue _ _ pt signed shape) = + ValueType signed (Rank (length shape)) pt + +allTrue :: [C.Exp] -> C.Exp +allTrue [] = [C.cexp|true|] +allTrue [x] = x +allTrue (x : xs) = [C.cexp|$exp:x && $exp:(allTrue xs)|] + +prepareEntryInputs :: + [ExternalValue] -> + CompilerM op s ([(C.Param, Maybe C.Exp)], [C.BlockItem]) +prepareEntryInputs args = collect' $ zipWithM prepare [(0 :: Int) ..] args + where + arg_names = namesFromList $ concatMap evNames args + evNames (OpaqueValue _ vds) = map vdName vds + evNames (TransparentValue vd) = [vdName vd] + vdName (ArrayValue v _ _ _ _) = v + vdName (ScalarValue _ _ v) = v + + prepare pno (TransparentValue vd) = do + let pname = "in" ++ show pno + (ty, check) <- prepareValue Public [C.cexp|$id:pname|] vd + pure + ( [C.cparam|const $ty:ty $id:pname|], + if null check then Nothing else Just $ allTrue check + ) + prepare pno (OpaqueValue desc vds) = do + ty <- opaqueToCType desc + let pname = "in" ++ show pno + field i ScalarValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] + field i ArrayValue {} = [C.cexp|$id:pname->$id:(tupleField i)|] + checks <- map snd <$> zipWithM (prepareValue Private) (zipWith field [0 ..] vds) vds + pure + ( [C.cparam|const $ty:ty *$id:pname|], + if all null checks + then Nothing + else Just $ allTrue $ concat checks + ) + + prepareValue _ src (ScalarValue pt signed name) = do + let pt' = primAPIType signed pt + src' = fromStorage pt $ C.toExp src mempty + stm [C.cstm|$id:name = $exp:src';|] + pure (pt', []) + prepareValue pub src vd@(ArrayValue mem _ _ _ shape) = do + ty <- valueTypeToCType pub $ valueDescToType vd + + stm [C.cstm|$exp:mem = $exp:src->mem;|] + + let rank = length shape + maybeCopyDim (Var d) i + | d `notNameIn` arg_names = + ( Just [C.cstm|$id:d = $exp:src->shape[$int:i];|], + [C.cexp|$id:d == $exp:src->shape[$int:i]|] + ) + maybeCopyDim x i = + ( Nothing, + [C.cexp|$exp:x == $exp:src->shape[$int:i]|] + ) + + let (sets, checks) = + unzip $ zipWith maybeCopyDim shape [0 .. rank - 1] + stms $ catMaybes sets + + pure ([C.cty|$ty:ty*|], checks) + +prepareEntryOutputs :: [ExternalValue] -> CompilerM op s ([C.Param], [C.BlockItem]) +prepareEntryOutputs = collect' . zipWithM prepare [(0 :: Int) ..] + where + prepare pno (TransparentValue vd) = do + let pname = "out" ++ show pno + ty <- valueTypeToCType Public $ valueDescToType vd + + case vd of + ArrayValue {} -> do + stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] + prepareValue [C.cexp|*$id:pname|] vd + pure [C.cparam|$ty:ty **$id:pname|] + ScalarValue {} -> do + prepareValue [C.cexp|*$id:pname|] vd + pure [C.cparam|$ty:ty *$id:pname|] + prepare pno (OpaqueValue desc vds) = do + let pname = "out" ++ show pno + ty <- opaqueToCType desc + vd_ts <- mapM (valueTypeToCType Private . valueDescToType) vds + + stm [C.cstm|assert((*$id:pname = ($ty:ty*) malloc(sizeof($ty:ty))) != NULL);|] + + forM_ (zip3 [0 ..] vd_ts vds) $ \(i, ct, vd) -> do + let field = [C.cexp|((*$id:pname)->$id:(tupleField i))|] + case vd of + ScalarValue {} -> pure () + ArrayValue {} -> do + stm [C.cstm|assert(($exp:field = ($ty:ct*) malloc(sizeof($ty:ct))) != NULL);|] + prepareValue field vd + + pure [C.cparam|$ty:ty **$id:pname|] + + prepareValue dest (ScalarValue t _ name) = + let name' = toStorage t $ C.toExp name mempty + in stm [C.cstm|$exp:dest = $exp:name';|] + prepareValue dest (ArrayValue mem _ _ _ shape) = do + stm [C.cstm|$exp:dest->mem = $id:mem;|] + + let rank = length shape + maybeCopyDim (Constant x) i = + [C.cstm|$exp:dest->shape[$int:i] = $exp:x;|] + maybeCopyDim (Var d) i = + [C.cstm|$exp:dest->shape[$int:i] = $id:d;|] + stms $ zipWith maybeCopyDim shape [0 .. rank - 1] + +isValidCName :: Name -> Bool +isValidCName = check . nameToString + where + check [] = True -- academic + check (c : cs) = isAlpha c && all constituent cs + constituent c = isAlphaNum c || c == '_' + +entryName :: Name -> String +entryName v + | isValidCName v = "entry_" <> nameToString v + | otherwise = "entry_" <> zEncodeString (nameToString v) + +onEntryPoint :: + [C.BlockItem] -> + Name -> + Function op -> + CompilerM op s (Maybe (C.Definition, (T.Text, Manifest.EntryPoint))) +onEntryPoint _ _ (Function Nothing _ _ _) = pure Nothing +onEntryPoint get_consts fname (Function (Just (EntryPoint ename results args)) outputs inputs _) = inNewFunction $ do + let out_args = map (\p -> [C.cexp|&$id:(paramName p)|]) outputs + in_args = map (\p -> [C.cexp|$id:(paramName p)|]) inputs + + inputdecls <- collect $ mapM_ stubParam inputs + outputdecls <- collect $ mapM_ stubParam outputs + decl_mem <- declAllocatedMem + + entry_point_function_name <- publicName $ entryName ename + + (inputs', unpack_entry_inputs) <- prepareEntryInputs $ map snd args + let (entry_point_input_params, entry_point_input_checks) = unzip inputs' + + (entry_point_output_params, pack_entry_outputs) <- + prepareEntryOutputs $ map snd results + + ctx_ty <- contextType + + headerDecl + EntryDecl + [C.cedecl|int $id:entry_point_function_name + ($ty:ctx_ty *ctx, + $params:entry_point_output_params, + $params:entry_point_input_params);|] + + let checks = catMaybes entry_point_input_checks + check_input = + if null checks + then [] + else + [C.citems| + if (!($exp:(allTrue (catMaybes entry_point_input_checks)))) { + ret = 1; + set_error(ctx, msgprintf("Error: entry point arguments have invalid sizes.\n")); + }|] + + critical = + [C.citems| + $items:decl_mem + $items:unpack_entry_inputs + $items:check_input + if (ret == 0) { + ret = $id:(funName fname)(ctx, $args:out_args, $args:in_args); + if (ret == 0) { + $items:get_consts + + $items:pack_entry_outputs + } + } + |] + + ops <- asks envOperations + + let cdef = + [C.cedecl| + int $id:entry_point_function_name + ($ty:ctx_ty *ctx, + $params:entry_point_output_params, + $params:entry_point_input_params) { + $items:inputdecls + $items:outputdecls + + int ret = 0; + + $items:(criticalSection ops critical) + + return ret; + }|] + + manifest = + Manifest.EntryPoint + { Manifest.entryPointCFun = T.pack entry_point_function_name, + -- Note that our convention about what is "input/output" + -- and what is "results/args" is different between the + -- manifest and ImpCode. + Manifest.entryPointOutputs = map outputManifest results, + Manifest.entryPointInputs = map inputManifest args + } + + pure $ Just (cdef, (nameToText ename, manifest)) + where + stubParam (MemParam name space) = + declMem name space + stubParam (ScalarParam name ty) = do + let ty' = primTypeToCType ty + decl [C.cdecl|$ty:ty' $id:name;|] + + vdType (TransparentValue (ScalarValue pt signed _)) = + T.pack $ prettySigned (signed == Unsigned) pt + vdType (TransparentValue (ArrayValue _ _ pt signed shape)) = + T.pack $ + mconcat (replicate (length shape) "[]") + <> prettySigned (signed == Unsigned) pt + vdType (OpaqueValue name _) = + T.pack name + + outputManifest (u, vd) = + Manifest.Output + { Manifest.outputType = vdType vd, + Manifest.outputUnique = u == Unique + } + inputManifest ((v, u), vd) = + Manifest.Input + { Manifest.inputName = nameToText v, + Manifest.inputType = vdType vd, + Manifest.inputUnique = u == Unique + } diff --git a/src/Futhark/CodeGen/Backends/GenericC/Monad.hs b/src/Futhark/CodeGen/Backends/GenericC/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericC/Monad.hs @@ -0,0 +1,671 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} + +-- | C code generator framework. +module Futhark.CodeGen.Backends.GenericC.Monad + ( -- * Pluggable compiler + Operations (..), + Publicness (..), + OpCompiler, + ErrorCompiler, + CallCompiler, + PointerQuals, + MemoryType, + WriteScalar, + writeScalarPointerWithQuals, + ReadScalar, + readScalarPointerWithQuals, + Allocate, + Deallocate, + CopyBarrier (..), + Copy, + StaticArray, + + -- * Monadic compiler interface + CompilerM, + CompilerState (..), + CompilerEnv (..), + getUserState, + modifyUserState, + contextContents, + contextFinalInits, + runCompilerM, + inNewFunction, + cachingMemory, + volQuals, + rawMem, + item, + items, + stm, + stms, + decl, + atInit, + headerDecl, + publicDef, + publicDef_, + profileReport, + onClear, + HeaderSection (..), + libDecl, + earlyDecl, + publicName, + contextField, + contextFieldDyn, + memToCType, + cacheMem, + fatMemory, + rawMemCType, + freeRawMem, + allocRawMem, + fatMemType, + declAllocatedMem, + freeAllocatedMem, + collect, + collect', + contextType, + configType, + + -- * Building Blocks + copyMemoryDefaultSpace, + derefPointer, + setMem, + allocMem, + unRefMem, + declMem, + resetMem, + fatMemAlloc, + fatMemSet, + fatMemUnRef, + criticalSection, + module Futhark.CodeGen.Backends.SimpleRep, + ) +where + +import Control.Monad.Identity +import Control.Monad.Reader +import Control.Monad.State +import Data.Bifunctor (first) +import qualified Data.DList as DL +import Data.List (unzip4) +import Data.Loc +import qualified Data.Map.Strict as M +import Data.Maybe +import Futhark.CodeGen.Backends.SimpleRep +import Futhark.CodeGen.ImpCode +import Futhark.MonadFreshNames +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C + +-- How public an array type definition sould be. Public types show up +-- in the generated API, while private types are used only to +-- implement the members of opaques. +data Publicness = Private | Public + deriving (Eq, Ord, Show) + +type ArrayType = (Signedness, PrimType, Int) + +data CompilerState s = CompilerState + { compArrayTypes :: M.Map ArrayType Publicness, + compEarlyDecls :: DL.DList C.Definition, + compInit :: [C.Stm], + compNameSrc :: VNameSource, + compUserState :: s, + compHeaderDecls :: M.Map HeaderSection (DL.DList C.Definition), + compLibDecls :: DL.DList C.Definition, + compCtxFields :: DL.DList (C.Id, C.Type, Maybe C.Exp, Maybe C.Stm), + compProfileItems :: DL.DList C.BlockItem, + compClearItems :: DL.DList C.BlockItem, + compDeclaredMem :: [(VName, Space)], + compItems :: DL.DList C.BlockItem + } + +newCompilerState :: VNameSource -> s -> CompilerState s +newCompilerState src s = + CompilerState + { compArrayTypes = mempty, + compEarlyDecls = mempty, + compInit = [], + compNameSrc = src, + compUserState = s, + compHeaderDecls = mempty, + compLibDecls = mempty, + compCtxFields = mempty, + compProfileItems = mempty, + compClearItems = mempty, + compDeclaredMem = mempty, + compItems = mempty + } + +-- | In which part of the header file we put the declaration. This is +-- to ensure that the header file remains structured and readable. +data HeaderSection + = ArrayDecl String + | OpaqueTypeDecl String + | OpaqueDecl String + | EntryDecl + | MiscDecl + | InitDecl + deriving (Eq, Ord) + +-- | A substitute expression compiler, tried before the main +-- compilation function. +type OpCompiler op s = op -> CompilerM op s () + +type ErrorCompiler op s = ErrorMsg Exp -> String -> CompilerM op s () + +-- | The address space qualifiers for a pointer of the given type with +-- the given annotation. +type PointerQuals op s = String -> CompilerM op s [C.TypeQual] + +-- | The type of a memory block in the given memory space. +type MemoryType op s = SpaceId -> CompilerM op s C.Type + +-- | Write a scalar to the given memory block with the given element +-- index and in the given memory space. +type WriteScalar op s = + C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> C.Exp -> CompilerM op s () + +-- | Read a scalar from the given memory block with the given element +-- index and in the given memory space. +type ReadScalar op s = + C.Exp -> C.Exp -> C.Type -> SpaceId -> Volatility -> CompilerM op s C.Exp + +-- | Allocate a memory block of the given size and with the given tag +-- in the given memory space, saving a reference in the given variable +-- name. +type Allocate op s = + C.Exp -> + C.Exp -> + C.Exp -> + SpaceId -> + CompilerM op s () + +-- | De-allocate the given memory block with the given tag, which is +-- in the given memory space. +type Deallocate op s = C.Exp -> C.Exp -> SpaceId -> CompilerM op s () + +-- | Create a static array of values - initialised at load time. +type StaticArray op s = VName -> SpaceId -> PrimType -> ArrayContents -> CompilerM op s () + +-- | Whether a copying operation should implicitly function as a +-- barrier regarding further operations on the source. This is a +-- rather subtle detail and is mostly useful for letting some +-- device/GPU copies be asynchronous (#1664). +data CopyBarrier + = CopyBarrier + | -- | Explicit context synchronisation should be done + -- before the source or target is used. + CopyNoBarrier + deriving (Eq, Show) + +-- | Copy from one memory block to another. +type Copy op s = + CopyBarrier -> + C.Exp -> + C.Exp -> + Space -> + C.Exp -> + C.Exp -> + Space -> + C.Exp -> + CompilerM op s () + +-- | Call a function. +type CallCompiler op s = [VName] -> Name -> [C.Exp] -> CompilerM op s () + +data Operations op s = Operations + { opsWriteScalar :: WriteScalar op s, + opsReadScalar :: ReadScalar op s, + opsAllocate :: Allocate op s, + opsDeallocate :: Deallocate op s, + opsCopy :: Copy op s, + opsStaticArray :: StaticArray op s, + opsMemoryType :: MemoryType op s, + opsCompiler :: OpCompiler op s, + opsError :: ErrorCompiler op s, + opsCall :: CallCompiler op s, + -- | If true, use reference counting. Otherwise, bare + -- pointers. + opsFatMemory :: Bool, + -- | Code to bracket critical sections. + opsCritical :: ([C.BlockItem], [C.BlockItem]) + } + +freeAllocatedMem :: CompilerM op s [C.BlockItem] +freeAllocatedMem = collect $ mapM_ (uncurry unRefMem) =<< gets compDeclaredMem + +declAllocatedMem :: CompilerM op s [C.BlockItem] +declAllocatedMem = collect $ mapM_ f =<< gets compDeclaredMem + where + f (name, space) = do + ty <- memToCType name space + decl [C.cdecl|$ty:ty $id:name;|] + resetMem name space + +data CompilerEnv op s = CompilerEnv + { envOperations :: Operations op s, + -- | Mapping memory blocks to sizes. These memory blocks are CPU + -- memory that we know are used in particularly simple ways (no + -- reference counting necessary). To cut down on allocator + -- pressure, we keep these allocations around for a long time, and + -- record their sizes so we can reuse them if possible (and + -- realloc() when needed). + envCachedMem :: M.Map C.Exp VName + } + +contextContents :: CompilerM op s ([C.FieldGroup], [C.Stm], [C.Stm]) +contextContents = do + (field_names, field_types, field_values, field_frees) <- + gets $ unzip4 . DL.toList . compCtxFields + let fields = + [ [C.csdecl|$ty:ty $id:name;|] + | (name, ty) <- zip field_names field_types + ] + init_fields = + [ [C.cstm|ctx->$id:name = $exp:e;|] + | (name, Just e) <- zip field_names field_values + ] + pure (fields, init_fields, catMaybes field_frees) + +contextFinalInits :: CompilerM op s [C.Stm] +contextFinalInits = gets compInit + +newtype CompilerM op s a + = CompilerM (ReaderT (CompilerEnv op s) (State (CompilerState s)) a) + deriving + ( Functor, + Applicative, + Monad, + MonadState (CompilerState s), + MonadReader (CompilerEnv op s) + ) + +instance MonadFreshNames (CompilerM op s) where + getNameSource = gets compNameSrc + putNameSource src = modify $ \s -> s {compNameSrc = src} + +runCompilerM :: + Operations op s -> + VNameSource -> + s -> + CompilerM op s a -> + (a, CompilerState s) +runCompilerM ops src userstate (CompilerM m) = + runState + (runReaderT m (CompilerEnv ops mempty)) + (newCompilerState src userstate) + +getUserState :: CompilerM op s s +getUserState = gets compUserState + +modifyUserState :: (s -> s) -> CompilerM op s () +modifyUserState f = modify $ \compstate -> + compstate {compUserState = f $ compUserState compstate} + +atInit :: C.Stm -> CompilerM op s () +atInit x = modify $ \s -> + s {compInit = compInit s ++ [x]} + +collect :: CompilerM op s () -> CompilerM op s [C.BlockItem] +collect m = snd <$> collect' m + +collect' :: CompilerM op s a -> CompilerM op s (a, [C.BlockItem]) +collect' m = do + old <- gets compItems + modify $ \s -> s {compItems = mempty} + x <- m + new <- gets compItems + modify $ \s -> s {compItems = old} + pure (x, DL.toList new) + +-- | Used when we, inside an existing 'CompilerM' action, want to +-- generate code for a new function. Use this so that the compiler +-- understands that previously declared memory doesn't need to be +-- freed inside this action. +inNewFunction :: CompilerM op s a -> CompilerM op s a +inNewFunction m = do + old_mem <- gets compDeclaredMem + modify $ \s -> s {compDeclaredMem = mempty} + x <- local noCached m + modify $ \s -> s {compDeclaredMem = old_mem} + pure x + where + noCached env = env {envCachedMem = mempty} + +item :: C.BlockItem -> CompilerM op s () +item x = modify $ \s -> s {compItems = DL.snoc (compItems s) x} + +items :: [C.BlockItem] -> CompilerM op s () +items xs = modify $ \s -> s {compItems = DL.append (compItems s) (DL.fromList xs)} + +fatMemory :: Space -> CompilerM op s Bool +fatMemory ScalarSpace {} = pure False +fatMemory _ = asks $ opsFatMemory . envOperations + +cacheMem :: C.ToExp a => a -> CompilerM op s (Maybe VName) +cacheMem a = asks $ M.lookup (C.toExp a noLoc) . envCachedMem + +-- | Construct a publicly visible definition using the specified name +-- as the template. The first returned definition is put in the +-- header file, and the second is the implementation. Returns the public +-- name. +publicDef :: + String -> + HeaderSection -> + (String -> (C.Definition, C.Definition)) -> + CompilerM op s String +publicDef s h f = do + s' <- publicName s + let (pub, priv) = f s' + headerDecl h pub + earlyDecl priv + pure s' + +-- | As 'publicDef', but ignores the public name. +publicDef_ :: + String -> + HeaderSection -> + (String -> (C.Definition, C.Definition)) -> + CompilerM op s () +publicDef_ s h f = void $ publicDef s h f + +headerDecl :: HeaderSection -> C.Definition -> CompilerM op s () +headerDecl sec def = modify $ \s -> + s + { compHeaderDecls = + M.unionWith + (<>) + (compHeaderDecls s) + (M.singleton sec (DL.singleton def)) + } + +libDecl :: C.Definition -> CompilerM op s () +libDecl def = modify $ \s -> + s {compLibDecls = compLibDecls s <> DL.singleton def} + +earlyDecl :: C.Definition -> CompilerM op s () +earlyDecl def = modify $ \s -> + s {compEarlyDecls = compEarlyDecls s <> DL.singleton def} + +contextField :: C.Id -> C.Type -> Maybe C.Exp -> CompilerM op s () +contextField name ty initial = modify $ \s -> + s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial, Nothing)} + +contextFieldDyn :: C.Id -> C.Type -> Maybe C.Exp -> C.Stm -> CompilerM op s () +contextFieldDyn name ty initial free = modify $ \s -> + s {compCtxFields = compCtxFields s <> DL.singleton (name, ty, initial, Just free)} + +profileReport :: C.BlockItem -> CompilerM op s () +profileReport x = modify $ \s -> + s {compProfileItems = compProfileItems s <> DL.singleton x} + +onClear :: C.BlockItem -> CompilerM op s () +onClear x = modify $ \s -> + s {compClearItems = compClearItems s <> DL.singleton x} + +stm :: C.Stm -> CompilerM op s () +stm s = item [C.citem|$stm:s|] + +stms :: [C.Stm] -> CompilerM op s () +stms = mapM_ stm + +decl :: C.InitGroup -> CompilerM op s () +decl x = item [C.citem|$decl:x;|] + +-- | Public names must have a consitent prefix. +publicName :: String -> CompilerM op s String +publicName s = pure $ "futhark_" ++ s + +memToCType :: VName -> Space -> CompilerM op s C.Type +memToCType v space = do + refcount <- fatMemory space + cached <- isJust <$> cacheMem v + if refcount && not cached + then pure $ fatMemType space + else rawMemCType space + +rawMemCType :: Space -> CompilerM op s C.Type +rawMemCType DefaultSpace = pure defaultMemBlockType +rawMemCType (Space sid) = join $ asks (opsMemoryType . envOperations) <*> pure sid +rawMemCType (ScalarSpace [] t) = + pure [C.cty|$ty:(primTypeToCType t)[1]|] +rawMemCType (ScalarSpace ds t) = + pure [C.cty|$ty:(primTypeToCType t)[$exp:(cproduct ds')]|] + where + ds' = map (`C.toExp` noLoc) ds + +fatMemType :: Space -> C.Type +fatMemType space = + [C.cty|struct $id:name|] + where + name = case space of + Space sid -> "memblock_" ++ sid + _ -> "memblock" + +fatMemSet :: Space -> String +fatMemSet (Space sid) = "memblock_set_" ++ sid +fatMemSet _ = "memblock_set" + +fatMemAlloc :: Space -> String +fatMemAlloc (Space sid) = "memblock_alloc_" ++ sid +fatMemAlloc _ = "memblock_alloc" + +fatMemUnRef :: Space -> String +fatMemUnRef (Space sid) = "memblock_unref_" ++ sid +fatMemUnRef _ = "memblock_unref" + +rawMem :: VName -> CompilerM op s C.Exp +rawMem v = rawMem' <$> fat <*> pure v + where + fat = asks ((&&) . opsFatMemory . envOperations) <*> (isNothing <$> cacheMem v) + +rawMem' :: C.ToExp a => Bool -> a -> C.Exp +rawMem' True e = [C.cexp|$exp:e.mem|] +rawMem' False e = [C.cexp|$exp:e|] + +allocRawMem :: + (C.ToExp a, C.ToExp b, C.ToExp c) => + a -> + b -> + Space -> + c -> + CompilerM op s () +allocRawMem dest size space desc = case space of + Space sid -> + join $ + asks (opsAllocate . envOperations) + <*> pure [C.cexp|$exp:dest|] + <*> pure [C.cexp|$exp:size|] + <*> pure [C.cexp|$exp:desc|] + <*> pure sid + _ -> + stm [C.cstm|$exp:dest = (unsigned char*) malloc((size_t)$exp:size);|] + +freeRawMem :: + (C.ToExp a, C.ToExp b) => + a -> + Space -> + b -> + CompilerM op s () +freeRawMem mem space desc = + case space of + Space sid -> do + free_mem <- asks (opsDeallocate . envOperations) + free_mem [C.cexp|$exp:mem|] [C.cexp|$exp:desc|] sid + _ -> item [C.citem|free($exp:mem);|] + +declMem :: VName -> Space -> CompilerM op s () +declMem name space = do + cached <- isJust <$> cacheMem name + fat <- fatMemory space + unless cached $ + if fat + then modify $ \s -> s {compDeclaredMem = (name, space) : compDeclaredMem s} + else do + ty <- memToCType name space + decl [C.cdecl|$ty:ty $id:name;|] + +resetMem :: C.ToExp a => a -> Space -> CompilerM op s () +resetMem mem space = do + refcount <- fatMemory space + cached <- isJust <$> cacheMem mem + if cached + then stm [C.cstm|$exp:mem = NULL;|] + else + when refcount $ + stm [C.cstm|$exp:mem.references = NULL;|] + +setMem :: (C.ToExp a, C.ToExp b) => a -> b -> Space -> CompilerM op s () +setMem dest src space = do + refcount <- fatMemory space + let src_s = pretty $ C.toExp src noLoc + if refcount + then + stm + [C.cstm|if ($id:(fatMemSet space)(ctx, &$exp:dest, &$exp:src, + $string:src_s) != 0) { + return 1; + }|] + else case space of + ScalarSpace ds _ -> do + i' <- newVName "i" + let i = C.toIdent i' + it = primTypeToCType $ IntType Int32 + ds' = map (`C.toExp` noLoc) ds + bound = cproduct ds' + stm + [C.cstm|for ($ty:it $id:i = 0; $id:i < $exp:bound; $id:i++) { + $exp:dest[$id:i] = $exp:src[$id:i]; + }|] + _ -> stm [C.cstm|$exp:dest = $exp:src;|] + +unRefMem :: C.ToExp a => a -> Space -> CompilerM op s () +unRefMem mem space = do + refcount <- fatMemory space + cached <- isJust <$> cacheMem mem + let mem_s = pretty $ C.toExp mem noLoc + when (refcount && not cached) $ + stm + [C.cstm|if ($id:(fatMemUnRef space)(ctx, &$exp:mem, $string:mem_s) != 0) { + return 1; + }|] + +allocMem :: + (C.ToExp a, C.ToExp b) => + a -> + b -> + Space -> + C.Stm -> + CompilerM op s () +allocMem mem size space on_failure = do + refcount <- fatMemory space + let mem_s = pretty $ C.toExp mem noLoc + if refcount + then + stm + [C.cstm|if ($id:(fatMemAlloc space)(ctx, &$exp:mem, $exp:size, + $string:mem_s)) { + $stm:on_failure + }|] + else do + freeRawMem mem space mem_s + allocRawMem mem size space [C.cexp|desc|] + +copyMemoryDefaultSpace :: + C.Exp -> + C.Exp -> + C.Exp -> + C.Exp -> + C.Exp -> + CompilerM op s () +copyMemoryDefaultSpace destmem destidx srcmem srcidx nbytes = + stm + [C.cstm|if ($exp:nbytes > 0) { + memmove($exp:destmem + $exp:destidx, + $exp:srcmem + $exp:srcidx, + $exp:nbytes); + }|] + +cachingMemory :: + M.Map VName Space -> + ([C.BlockItem] -> [C.Stm] -> CompilerM op s a) -> + CompilerM op s a +cachingMemory lexical f = do + -- We only consider lexical 'DefaultSpace' memory blocks to be + -- cached. This is not a deep technical restriction, but merely a + -- heuristic based on GPU memory usually involving larger + -- allocations, that do not suffer from the overhead of reference + -- counting. + let cached = M.keys $ M.filter (== DefaultSpace) lexical + + cached' <- forM cached $ \mem -> do + size <- newVName $ pretty mem <> "_cached_size" + pure (mem, size) + + let lexMem env = + env + { envCachedMem = + M.fromList (map (first (`C.toExp` noLoc)) cached') + <> envCachedMem env + } + + declCached (mem, size) = + [ [C.citem|typename int64_t $id:size = 0;|], + [C.citem|$ty:defaultMemBlockType $id:mem = NULL;|] + ] + + freeCached (mem, _) = + [C.cstm|free($id:mem);|] + + local lexMem $ f (concatMap declCached cached') (map freeCached cached') + +derefPointer :: C.Exp -> C.Exp -> C.Type -> C.Exp +derefPointer ptr i res_t = + [C.cexp|(($ty:res_t)$exp:ptr)[$exp:i]|] + +volQuals :: Volatility -> [C.TypeQual] +volQuals Volatile = [C.ctyquals|volatile|] +volQuals Nonvolatile = [] + +writeScalarPointerWithQuals :: PointerQuals op s -> WriteScalar op s +writeScalarPointerWithQuals quals_f dest i elemtype space vol v = do + quals <- quals_f space + let quals' = volQuals vol ++ quals + deref = + derefPointer + dest + i + [C.cty|$tyquals:quals' $ty:elemtype*|] + stm [C.cstm|$exp:deref = $exp:v;|] + +readScalarPointerWithQuals :: PointerQuals op s -> ReadScalar op s +readScalarPointerWithQuals quals_f dest i elemtype space vol = do + quals <- quals_f space + let quals' = volQuals vol ++ quals + pure $ derefPointer dest i [C.cty|$tyquals:quals' $ty:elemtype*|] + +criticalSection :: Operations op s -> [C.BlockItem] -> [C.BlockItem] +criticalSection ops x = + [C.citems|lock_lock(&ctx->lock); + $items:(fst (opsCritical ops)) + $items:x + $items:(snd (opsCritical ops)) + lock_unlock(&ctx->lock); + |] + +-- | The generated code must define a context struct with this name. +contextType :: CompilerM op s C.Type +contextType = do + name <- publicName "context" + pure [C.cty|struct $id:name|] + +-- | The generated code must define a configuration struct with this +-- name. +configType :: CompilerM op s C.Type +configType = do + name <- publicName "context_config" + pure [C.cty|struct $id:name|] diff --git a/src/Futhark/CodeGen/Backends/GenericC/Types.hs b/src/Futhark/CodeGen/Backends/GenericC/Types.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/GenericC/Types.hs @@ -0,0 +1,586 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} + +-- | Code generation for public API types. +module Futhark.CodeGen.Backends.GenericC.Types + ( generateAPITypes, + valueTypeToCType, + opaqueToCType, + ) +where + +import Control.Monad.Reader +import Control.Monad.State +import Data.Char (isDigit) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Text as T +import Futhark.CodeGen.Backends.GenericC.Monad +import Futhark.CodeGen.ImpCode +import qualified Futhark.Manifest as Manifest +import Futhark.Util (chunks, mapAccumLM) +import Futhark.Util.Pretty (prettyText) +import qualified Language.C.Quote.OpenCL as C +import qualified Language.C.Syntax as C + +opaqueToCType :: String -> CompilerM op s C.Type +opaqueToCType desc = do + name <- publicName $ opaqueName desc + pure [C.cty|struct $id:name|] + +valueTypeToCType :: Publicness -> ValueType -> CompilerM op s C.Type +valueTypeToCType _ (ValueType signed (Rank 0) pt) = + pure $ primAPIType signed pt +valueTypeToCType pub (ValueType signed (Rank rank) pt) = do + name <- publicName $ arrayName pt signed rank + let add = M.insertWith max (signed, pt, rank) pub + modify $ \s -> s {compArrayTypes = add $ compArrayTypes s} + pure [C.cty|struct $id:name|] + +arrayLibraryFunctions :: + Publicness -> + Space -> + PrimType -> + Signedness -> + Int -> + CompilerM op s Manifest.ArrayOps +arrayLibraryFunctions pub space pt signed rank = do + let pt' = primAPIType signed pt + name = arrayName pt signed rank + arr_name = "futhark_" ++ name + array_type = [C.cty|struct $id:arr_name|] + + new_array <- publicName $ "new_" ++ name + new_raw_array <- publicName $ "new_raw_" ++ name + free_array <- publicName $ "free_" ++ name + values_array <- publicName $ "values_" ++ name + values_raw_array <- publicName $ "values_raw_" ++ name + shape_array <- publicName $ "shape_" ++ name + + let shape_names = ["dim" ++ show i | i <- [0 .. rank - 1]] + shape_params = [[C.cparam|typename int64_t $id:k|] | k <- shape_names] + arr_size = cproduct [[C.cexp|$id:k|] | k <- shape_names] + arr_size_array = cproduct [[C.cexp|arr->shape[$int:i]|] | i <- [0 .. rank - 1]] + copy <- asks $ opsCopy . envOperations + + memty <- rawMemCType space + + let prepare_new = do + resetMem [C.cexp|arr->mem|] space + allocMem + [C.cexp|arr->mem|] + [C.cexp|$exp:arr_size * $int:(primByteSize pt::Int)|] + space + [C.cstm|return NULL;|] + forM_ [0 .. rank - 1] $ \i -> + let dim_s = "dim" ++ show i + in stm [C.cstm|arr->shape[$int:i] = $id:dim_s;|] + + new_body <- collect $ do + prepare_new + copy + CopyNoBarrier + [C.cexp|arr->mem.mem|] + [C.cexp|0|] + space + [C.cexp|data|] + [C.cexp|0|] + DefaultSpace + [C.cexp|((size_t)$exp:arr_size) * $int:(primByteSize pt::Int)|] + + new_raw_body <- collect $ do + prepare_new + copy + CopyNoBarrier + [C.cexp|arr->mem.mem|] + [C.cexp|0|] + space + [C.cexp|data|] + [C.cexp|offset|] + space + [C.cexp|((size_t)$exp:arr_size) * $int:(primByteSize pt::Int)|] + + free_body <- collect $ unRefMem [C.cexp|arr->mem|] space + + values_body <- + collect $ + copy + CopyNoBarrier + [C.cexp|data|] + [C.cexp|0|] + DefaultSpace + [C.cexp|arr->mem.mem|] + [C.cexp|0|] + space + [C.cexp|((size_t)$exp:arr_size_array) * $int:(primByteSize pt::Int)|] + + ctx_ty <- contextType + ops <- asks envOperations + + let proto = case pub of + Public -> headerDecl (ArrayDecl name) + Private -> libDecl + + proto + [C.cedecl|struct $id:arr_name;|] + proto + [C.cedecl|$ty:array_type* $id:new_array($ty:ctx_ty *ctx, const $ty:pt' *data, $params:shape_params);|] + proto + [C.cedecl|$ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, const $ty:memty data, typename int64_t offset, $params:shape_params);|] + proto + [C.cedecl|int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + proto + [C.cedecl|int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data);|] + proto + [C.cedecl|$ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + proto + [C.cedecl|const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr);|] + + mapM_ + libDecl + [C.cunit| + $ty:array_type* $id:new_array($ty:ctx_ty *ctx, const $ty:pt' *data, $params:shape_params) { + $ty:array_type* bad = NULL; + $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); + if (arr == NULL) { + return bad; + } + $items:(criticalSection ops new_body) + return arr; + } + + $ty:array_type* $id:new_raw_array($ty:ctx_ty *ctx, const $ty:memty data, typename int64_t offset, + $params:shape_params) { + $ty:array_type* bad = NULL; + $ty:array_type *arr = ($ty:array_type*) malloc(sizeof($ty:array_type)); + if (arr == NULL) { + return bad; + } + $items:(criticalSection ops new_raw_body) + return arr; + } + + int $id:free_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + $items:(criticalSection ops free_body) + free(arr); + return 0; + } + + int $id:values_array($ty:ctx_ty *ctx, $ty:array_type *arr, $ty:pt' *data) { + $items:(criticalSection ops values_body) + return 0; + } + + $ty:memty $id:values_raw_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + (void)ctx; + return arr->mem.mem; + } + + const typename int64_t* $id:shape_array($ty:ctx_ty *ctx, $ty:array_type *arr) { + (void)ctx; + return arr->shape; + } + |] + + pure $ + Manifest.ArrayOps + { Manifest.arrayFree = T.pack free_array, + Manifest.arrayShape = T.pack shape_array, + Manifest.arrayValues = T.pack values_array, + Manifest.arrayNew = T.pack new_array + } + +lookupOpaqueType :: String -> OpaqueTypes -> OpaqueType +lookupOpaqueType v (OpaqueTypes types) = + case lookup v types of + Just t -> t + Nothing -> error $ "Unknown opaque type: " ++ show v + +opaquePayload :: OpaqueTypes -> OpaqueType -> [ValueType] +opaquePayload _ (OpaqueType ts) = ts +opaquePayload types (OpaqueRecord fs) = concatMap f fs + where + f (_, TypeOpaque s) = opaquePayload types $ lookupOpaqueType s types + f (_, TypeTransparent v) = [v] + +entryPointTypeToCType :: Publicness -> EntryPointType -> CompilerM op s C.Type +entryPointTypeToCType _ (TypeOpaque desc) = opaqueToCType desc +entryPointTypeToCType pub (TypeTransparent vt) = valueTypeToCType pub vt + +entryTypeName :: EntryPointType -> Manifest.TypeName +entryTypeName (TypeOpaque desc) = T.pack desc +entryTypeName (TypeTransparent vt) = prettyText vt + +-- | Figure out which of the members of an opaque type corresponds to +-- which fields. +recordFieldPayloads :: OpaqueTypes -> [EntryPointType] -> [a] -> [[a]] +recordFieldPayloads types = chunks . map typeLength + where + typeLength (TypeTransparent _) = 1 + typeLength (TypeOpaque desc) = + length $ opaquePayload types $ lookupOpaqueType desc types + +opaqueProjectFunctions :: + OpaqueTypes -> + String -> + [(Name, EntryPointType)] -> + [ValueType] -> + CompilerM op s [Manifest.RecordField] +opaqueProjectFunctions types desc fs vds = do + opaque_type <- opaqueToCType desc + ctx_ty <- contextType + ops <- asks envOperations + let mkProject (TypeTransparent (ValueType sign (Rank 0) pt)) [(i, _)] = do + pure + ( primAPIType sign pt, + [C.citems|v = obj->$id:(tupleField i);|] + ) + mkProject (TypeTransparent vt) [(i, _)] = do + ct <- valueTypeToCType Public vt + pure + ( [C.cty|$ty:ct *|], + criticalSection + ops + [C.citems|v = malloc(sizeof($ty:ct)); + memcpy(v, obj->$id:(tupleField i), sizeof($ty:ct)); + (void)(*(v->mem.references))++;|] + ) + mkProject (TypeTransparent _) rep = + error $ "mkProject: invalid representation of transparent type: " ++ show rep + mkProject (TypeOpaque f_desc) components = do + ct <- opaqueToCType f_desc + let setField j (i, ValueType _ (Rank r) _) = + if r == 0 + then [C.citems|v->$id:(tupleField j) = obj->$id:(tupleField i);|] + else + [C.citems|v->$id:(tupleField j) = malloc(sizeof(*v->$id:(tupleField j))); + *v->$id:(tupleField j) = *obj->$id:(tupleField i); + (void)(*(v->$id:(tupleField j)->mem.references))++;|] + pure + ( [C.cty|$ty:ct *|], + criticalSection + ops + [C.citems|v = malloc(sizeof($ty:ct)); + $items:(concat (zipWith setField [0..] components))|] + ) + let onField ((f, et), elems) = do + project <- publicName $ "project_" ++ opaqueName desc ++ "_" ++ nameToString f + (et_ty, project_items) <- mkProject et elems + headerDecl + (OpaqueDecl desc) + [C.cedecl|int $id:project($ty:ctx_ty *ctx, $ty:et_ty *out, const $ty:opaque_type *obj);|] + libDecl + [C.cedecl|int $id:project($ty:ctx_ty *ctx, $ty:et_ty *out, const $ty:opaque_type *obj) { + (void)ctx; + $ty:et_ty v; + $items:project_items + *out = v; + return 0; + }|] + pure $ Manifest.RecordField (nameToText f) (entryTypeName et) (T.pack project) + + mapM onField . zip fs . recordFieldPayloads types (map snd fs) $ + zip [0 ..] vds + +opaqueNewFunctions :: + OpaqueTypes -> + String -> + [(Name, EntryPointType)] -> + [ValueType] -> + CompilerM op s Manifest.CFuncName +opaqueNewFunctions types desc fs vds = do + opaque_type <- opaqueToCType desc + ctx_ty <- contextType + ops <- asks envOperations + new <- publicName $ "new_" ++ opaqueName desc + + (params, new_stms) <- + fmap (unzip . snd) + . mapAccumLM onField 0 + . zip fs + . recordFieldPayloads types (map snd fs) + $ vds + + headerDecl + (OpaqueDecl desc) + [C.cedecl|int $id:new($ty:ctx_ty *ctx, $ty:opaque_type** out, $params:params);|] + libDecl + [C.cedecl|int $id:new($ty:ctx_ty *ctx, $ty:opaque_type** out, $params:params) { + $ty:opaque_type* v = malloc(sizeof($ty:opaque_type)); + $items:(criticalSection ops new_stms) + *out = v; + return 0; + }|] + pure $ T.pack new + where + onField offset ((f, et), f_vts) = do + let param_name = + if all isDigit (nameToString f) + then C.toIdent ("v" <> f) mempty + else C.toIdent f mempty + case et of + TypeTransparent (ValueType sign (Rank 0) pt) -> do + let ct = primAPIType sign pt + pure + ( offset + 1, + ( [C.cparam|const $ty:ct $id:param_name|], + [C.citem|v->$id:(tupleField offset) = $id:param_name;|] + ) + ) + TypeTransparent vt -> do + ct <- valueTypeToCType Public vt + pure + ( offset + 1, + ( [C.cparam|const $ty:ct* $id:param_name|], + [C.citem|{v->$id:(tupleField offset) = malloc(sizeof($ty:ct)); + *v->$id:(tupleField offset) = *$id:param_name; + (void)(*(v->$id:(tupleField offset)->mem.references))++;}|] + ) + ) + TypeOpaque f_desc -> do + ct <- opaqueToCType f_desc + let param_fields = do + i <- [0 ..] + pure [C.cexp|$id:param_name->$id:(tupleField i)|] + pure + ( offset + length f_vts, + ( [C.cparam|const $ty:ct* $id:param_name|], + [C.citem|{$stms:(zipWith3 setFieldField [offset ..] param_fields f_vts)}|] + ) + ) + + setFieldField i e (ValueType _ (Rank r) _) + | r == 0 = + [C.cstm|v->$id:(tupleField i) = $exp:e;|] + | otherwise = + [C.cstm|{v->$id:(tupleField i) = malloc(sizeof(*$exp:e)); + *v->$id:(tupleField i) = *$exp:e; + (void)(*(v->$id:(tupleField i)->mem.references))++;}|] + +processOpaqueRecord :: + OpaqueTypes -> + String -> + OpaqueType -> + [ValueType] -> + CompilerM op s (Maybe Manifest.RecordOps) +processOpaqueRecord _ _ (OpaqueType _) _ = pure Nothing +processOpaqueRecord types desc (OpaqueRecord fs) vds = + Just + <$> ( Manifest.RecordOps + <$> opaqueProjectFunctions types desc fs vds + <*> opaqueNewFunctions types desc fs vds + ) + +opaqueLibraryFunctions :: + OpaqueTypes -> + String -> + OpaqueType -> + CompilerM op s (Manifest.OpaqueOps, Maybe Manifest.RecordOps) +opaqueLibraryFunctions types desc ot = do + name <- publicName $ opaqueName desc + free_opaque <- publicName $ "free_" ++ opaqueName desc + store_opaque <- publicName $ "store_" ++ opaqueName desc + restore_opaque <- publicName $ "restore_" ++ opaqueName desc + + let opaque_type = [C.cty|struct $id:name|] + + freeComponent i (ValueType signed (Rank rank) pt) = unless (rank == 0) $ do + let field = tupleField i + free_array <- publicName $ "free_" ++ arrayName pt signed rank + -- Protect against NULL here, because we also want to use this + -- to free partially loaded opaques. + stm + [C.cstm|if (obj->$id:field != NULL && (tmp = $id:free_array(ctx, obj->$id:field)) != 0) { + ret = tmp; + }|] + + storeComponent i (ValueType sign (Rank 0) pt) = + let field = tupleField i + in ( storageSize pt 0 [C.cexp|NULL|], + storeValueHeader sign pt 0 [C.cexp|NULL|] [C.cexp|out|] + ++ [C.cstms|memcpy(out, &obj->$id:field, sizeof(obj->$id:field)); + out += sizeof(obj->$id:field);|] + ) + storeComponent i (ValueType sign (Rank rank) pt) = + let arr_name = arrayName pt sign rank + field = tupleField i + shape_array = "futhark_shape_" ++ arr_name + values_array = "futhark_values_" ++ arr_name + shape' = [C.cexp|$id:shape_array(ctx, obj->$id:field)|] + num_elems = cproduct [[C.cexp|$exp:shape'[$int:j]|] | j <- [0 .. rank - 1]] + in ( storageSize pt rank shape', + storeValueHeader sign pt rank shape' [C.cexp|out|] + ++ [C.cstms|ret |= $id:values_array(ctx, obj->$id:field, (void*)out); + out += $exp:num_elems * $int:(primByteSize pt::Int);|] + ) + + ctx_ty <- contextType + + let vds = opaquePayload types ot + free_body <- collect $ zipWithM_ freeComponent [0 ..] vds + + store_body <- collect $ do + let (sizes, stores) = unzip $ zipWith storeComponent [0 ..] vds + size_vars = map (("size_" ++) . show) [0 .. length sizes - 1] + size_sum = csum [[C.cexp|$id:size|] | size <- size_vars] + forM_ (zip size_vars sizes) $ \(v, e) -> + item [C.citem|typename int64_t $id:v = $exp:e;|] + stm [C.cstm|*n = $exp:size_sum;|] + stm [C.cstm|if (p != NULL && *p == NULL) { *p = malloc(*n); }|] + stm [C.cstm|if (p != NULL) { unsigned char *out = *p; $stms:(concat stores) }|] + + let restoreComponent i (ValueType sign (Rank 0) pt) = do + let field = tupleField i + dataptr = "data_" ++ show i + stms $ loadValueHeader sign pt 0 [C.cexp|NULL|] [C.cexp|src|] + item [C.citem|const void* $id:dataptr = src;|] + stm [C.cstm|src += sizeof(obj->$id:field);|] + pure [C.cstms|memcpy(&obj->$id:field, $id:dataptr, sizeof(obj->$id:field));|] + restoreComponent i (ValueType sign (Rank rank) pt) = do + let field = tupleField i + arr_name = arrayName pt sign rank + new_array = "futhark_new_" ++ arr_name + dataptr = "data_" ++ show i + shapearr = "shape_" ++ show i + dims = [[C.cexp|$id:shapearr[$int:j]|] | j <- [0 .. rank - 1]] + num_elems = cproduct dims + item [C.citem|typename int64_t $id:shapearr[$int:rank] = {0};|] + stms $ loadValueHeader sign pt rank [C.cexp|$id:shapearr|] [C.cexp|src|] + item [C.citem|const void* $id:dataptr = src;|] + stm [C.cstm|obj->$id:field = NULL;|] + stm [C.cstm|src += $exp:num_elems * $int:(primByteSize pt::Int);|] + pure + [C.cstms| + obj->$id:field = $id:new_array(ctx, $id:dataptr, $args:dims); + if (obj->$id:field == NULL) { err = 1; }|] + + load_body <- collect $ do + loads <- concat <$> zipWithM restoreComponent [0 ..] (opaquePayload types ot) + stm + [C.cstm|if (err == 0) { + $stms:loads + }|] + + headerDecl + (OpaqueTypeDecl desc) + [C.cedecl|struct $id:name;|] + headerDecl + (OpaqueDecl desc) + [C.cedecl|int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj);|] + headerDecl + (OpaqueDecl desc) + [C.cedecl|int $id:store_opaque($ty:ctx_ty *ctx, const $ty:opaque_type *obj, void **p, size_t *n);|] + headerDecl + (OpaqueDecl desc) + [C.cedecl|$ty:opaque_type* $id:restore_opaque($ty:ctx_ty *ctx, const void *p);|] + + record <- processOpaqueRecord types desc ot vds + + -- We do not need to enclose most bodies in a critical section, + -- because when we operate on the components of the opaque, we are + -- calling public API functions that do their own locking. The + -- exception is projection, where we fiddle with reference counts. + mapM_ + libDecl + [C.cunit| + int $id:free_opaque($ty:ctx_ty *ctx, $ty:opaque_type *obj) { + (void)ctx; + int ret = 0, tmp; + $items:free_body + free(obj); + return ret; + } + + int $id:store_opaque($ty:ctx_ty *ctx, + const $ty:opaque_type *obj, void **p, size_t *n) { + (void)ctx; + int ret = 0; + $items:store_body + return ret; + } + + $ty:opaque_type* $id:restore_opaque($ty:ctx_ty *ctx, + const void *p) { + int err = 0; + const unsigned char *src = p; + $ty:opaque_type* obj = malloc(sizeof($ty:opaque_type)); + $items:load_body + if (err != 0) { + int ret = 0, tmp; + $items:free_body + free(obj); + obj = NULL; + } + return obj; + } + |] + + pure + ( Manifest.OpaqueOps + { Manifest.opaqueFree = T.pack free_opaque, + Manifest.opaqueStore = T.pack store_opaque, + Manifest.opaqueRestore = T.pack restore_opaque + }, + record + ) + +generateArray :: + Space -> + ((Signedness, PrimType, Int), Publicness) -> + CompilerM op s (Maybe (T.Text, Manifest.Type)) +generateArray space ((signed, pt, rank), pub) = do + name <- publicName $ arrayName pt signed rank + let memty = fatMemType space + libDecl [C.cedecl|struct $id:name { $ty:memty mem; typename int64_t shape[$int:rank]; };|] + ops <- arrayLibraryFunctions pub space pt signed rank + let pt_name = T.pack $ prettySigned (signed == Unsigned) pt + pretty_name = mconcat (replicate rank "[]") <> pt_name + arr_type = [C.cty|struct $id:name*|] + case pub of + Public -> + pure $ + Just + ( pretty_name, + Manifest.TypeArray (prettyText arr_type) pt_name rank ops + ) + Private -> + pure Nothing + +generateOpaque :: + OpaqueTypes -> + (String, OpaqueType) -> + CompilerM op s (T.Text, Manifest.Type) +generateOpaque types (desc, ot) = do + name <- publicName $ opaqueName desc + members <- zipWithM field (opaquePayload types ot) [(0 :: Int) ..] + libDecl [C.cedecl|struct $id:name { $sdecls:members };|] + (ops, record) <- opaqueLibraryFunctions types desc ot + let opaque_type = [C.cty|struct $id:name*|] + pure (T.pack desc, Manifest.TypeOpaque (prettyText opaque_type) ops record) + where + field vt@(ValueType _ (Rank r) _) i = do + ct <- valueTypeToCType Private vt + pure $ + if r == 0 + then [C.csdecl|$ty:ct $id:(tupleField i);|] + else [C.csdecl|$ty:ct *$id:(tupleField i);|] + +generateAPITypes :: Space -> OpaqueTypes -> CompilerM op s (M.Map T.Text Manifest.Type) +generateAPITypes arr_space types@(OpaqueTypes opaques) = do + mapM_ (findNecessaryArrays . snd) opaques + array_ts <- mapM (generateArray arr_space) . M.toList =<< gets compArrayTypes + opaque_ts <- mapM (generateOpaque types) opaques + pure $ M.fromList $ catMaybes array_ts <> opaque_ts + where + -- Ensure that array types will be generated before the opaque + -- records that allow projection of them. This is because the + -- projection functions somewhat uglily directly poke around in + -- the innards to increment reference counts. + findNecessaryArrays (OpaqueType _) = + pure () + findNecessaryArrays (OpaqueRecord fs) = + mapM_ (entryPointTypeToCType Public . snd) fs diff --git a/src/Futhark/CodeGen/Backends/MulticoreC.hs b/src/Futhark/CodeGen/Backends/MulticoreC.hs --- a/src/Futhark/CodeGen/Backends/MulticoreC.hs +++ b/src/Futhark/CodeGen/Backends/MulticoreC.hs @@ -144,6 +144,7 @@ int logging; typename lock_t lock; char *error; + typename lock_t error_lock; typename FILE *log; int total_runs; long int total_runtime; @@ -174,6 +175,7 @@ ctx->profiling_paused = 0; ctx->logging = 0; ctx->error = NULL; + create_lock(&ctx->error_lock); ctx->log = stderr; create_lock(&ctx->lock); diff --git a/src/Futhark/CodeGen/Backends/MulticoreISPC.hs b/src/Futhark/CodeGen/Backends/MulticoreISPC.hs --- a/src/Futhark/CodeGen/Backends/MulticoreISPC.hs +++ b/src/Futhark/CodeGen/Backends/MulticoreISPC.hs @@ -76,7 +76,6 @@ operations (ISPCState mempty mempty) ( do - GC.libDecl [C.cedecl|char** futhark_get_error_ref(struct futhark_context* ctx) { return &ctx->error; }|] MC.generateContext mapM_ compileBuiltinFun funs ) @@ -379,8 +378,7 @@ shim <- MC.multicoreDef "assert_shim" $ \s -> do pure [C.cedecl|void $id:s(struct futhark_context* ctx, $params:params) { - if (ctx->error == NULL) - ctx->error = msgprintf($string:formatstr', $args:formatargs', $string:stacktrace); + set_error(ctx, msgprintf($string:formatstr', $args:formatargs', $string:stacktrace)); }|] ispcDecl [C.cedecl|extern "C" $tyqual:unmasked void $id:shim($tyqual:uniform struct futhark_context* $tyqual:uniform, $params:params_uni);|] @@ -544,7 +542,7 @@ Just cur_size -> GC.stm [C.cstm|if ($exp:cur_size < $exp:size) { - err = lexical_realloc(futhark_get_error_ref(ctx), &$exp:name, &$exp:cur_size, $exp:size); + err = lexical_realloc(ctx, &$exp:name, &$exp:cur_size, $exp:size); if (err != FUTHARK_SUCCESS) { $escstm:("unmasked { return err; }") } diff --git a/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs b/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs @@ -73,6 +73,7 @@ int logging; typename lock_t lock; char *error; + typename lock_t error_lock; typename FILE *log; int profiling_paused; $sdecls:fields @@ -95,6 +96,7 @@ ctx->profiling = cfg->debugging; ctx->logging = cfg->debugging; ctx->error = NULL; + create_lock(&ctx->error_lock); ctx->log = stderr; create_lock(&ctx->lock); $stms:init_fields diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs --- a/src/Futhark/CodeGen/ImpCode.hs +++ b/src/Futhark/CodeGen/ImpCode.hs @@ -572,9 +572,11 @@ <+> ppr cond <+> "then {" indent 2 (ppr tbranch) - "} else {" - indent 2 (ppr fbranch) - "}" + "} else" + <+> case fbranch of + If {} -> ppr fbranch + _ -> + "{" indent 2 (ppr fbranch) "}" ppr (Call dests fname args) = commasep (map ppr dests) <+> "<-" diff --git a/src/Futhark/CodeGen/ImpGen.hs b/src/Futhark/CodeGen/ImpGen.hs --- a/src/Futhark/CodeGen/ImpGen.hs +++ b/src/Futhark/CodeGen/ImpGen.hs @@ -82,6 +82,7 @@ typeSize, inBounds, isMapTransposeCopy, + caseMatch, -- * Constructing code. dLParams, @@ -97,6 +98,7 @@ dPrimVE, dIndexSpace, dIndexSpace', + rotateIndex, sFor, sWhile, sComment, @@ -155,7 +157,7 @@ import Futhark.Util.IntegralExp import Futhark.Util.Loc (noLoc) import Language.Futhark.Warnings -import Prelude hiding (quot) +import Prelude hiding (mod, quot) -- | How to compile an t'Op'. type OpCompiler rep r op = Pat (LetDec rep) -> Op rep -> ImpM rep r op () @@ -795,13 +797,23 @@ ec <- asks envExpCompiler ec pat e +-- | Generate an expression that is true if the subexpressions match +-- the case pasttern. +caseMatch :: [SubExp] -> [Maybe PrimValue] -> Imp.TExp Bool +caseMatch ses vs = foldl (.&&.) true (zipWith cmp ses vs) + where + cmp se (Just v) = isBool $ toExp' (primValueType v) se ~==~ ValueExp v + cmp _ Nothing = true + defCompileExp :: (Mem rep inner) => Pat (LetDec rep) -> Exp rep -> ImpM rep r op () -defCompileExp pat (If cond tbranch fbranch _) = - sIf (toBoolExp cond) (compileBody pat tbranch) (compileBody pat fbranch) +defCompileExp pat (Match ses cases defbody _) = + foldl f (compileBody pat defbody) cases + where + f rest (Case vs body) = sIf (caseMatch ses vs) (compileBody pat body) rest defCompileExp pat (Apply fname args _ _) = do dest <- destinationFromPat pat targets <- funcallTargets dest @@ -920,7 +932,7 @@ uncurry warn loc "Safety check required at run-time." defCompileBasicOp (Pat [pe]) (Index src slice) | Just idxs <- sliceIndices slice = - copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . toInt64Exp) idxs + copyDWIM (patElemName pe) [] (Var src) $ map (DimFix . pe64) idxs defCompileBasicOp _ Index {} = pure () defCompileBasicOp (Pat [pe]) (Update safety _ slice se) = @@ -928,8 +940,8 @@ Unsafe -> write Safe -> sWhen (inBounds slice' dims) write where - slice' = fmap toInt64Exp slice - dims = map toInt64Exp $ arrayDims $ patElemType pe + slice' = fmap pe64 slice + dims = map pe64 $ arrayDims $ patElemType pe write = sUpdate (patElemName pe) slice' se defCompileBasicOp _ FlatIndex {} = pure () @@ -938,7 +950,7 @@ v_loc <- entryArrayLoc <$> lookupArray v copy (elemType (patElemType pe)) (flatSliceMemLoc pe_loc slice') v_loc where - slice' = fmap toInt64Exp slice + slice' = fmap pe64 slice defCompileBasicOp (Pat [pe]) (Replicate (Shape ds) se) | Acc {} <- patElemType pe = pure () | otherwise = do @@ -951,7 +963,7 @@ defCompileBasicOp (Pat [pe]) (Iota n e s it) = do e' <- toExp e s' <- toExp s - sFor "i" (toInt64Exp n) $ \i -> do + sFor "i" (pe64 n) $ \i -> do let i' = sExt it $ untyped i x <- dPrimV "x" . TPrimExp $ @@ -969,10 +981,10 @@ y_dims <- arrayDims <$> lookupType y let rows = case drop i y_dims of [] -> error $ "defCompileBasicOp Concat: empty array shape for " ++ pretty y - r : _ -> toInt64Exp r + r : _ -> pe64 r skip_dims = take i y_dims sliceAllDim d = DimSlice 0 d 1 - skip_slices = map (sliceAllDim . toInt64Exp) skip_dims + skip_slices = map (sliceAllDim . pe64) skip_dims destslice = skip_slices ++ [DimSlice (tvExp offs_glb) rows 1] copyDWIM (patElemName pe) destslice (Var y) [] offs_glb <-- tvExp offs_glb + rows @@ -997,8 +1009,13 @@ isLiteral _ = Nothing defCompileBasicOp _ Rearrange {} = pure () -defCompileBasicOp _ Rotate {} = - pure () +defCompileBasicOp (Pat [pe]) (Rotate rs arr) = do + shape <- arrayShape <$> lookupType arr + sLoopNest shape $ \is -> do + is' <- sequence $ zipWith3 rotate (shapeDims shape) rs is + copyDWIMFix (patElemName pe) is (Var arr) is' + where + rotate d r i = dPrimVE "rot_i" $ rotateIndex (pe64 d) (pe64 r) i defCompileBasicOp _ Reshape {} = pure () defCompileBasicOp _ (UpdateAcc acc is vs) = sComment "UpdateAcc" $ do @@ -1007,7 +1024,7 @@ -- we might otherwise end up declaring lambda parameters (if any) -- multiple times, as they are duplicated every time we do an -- UpdateAcc for the same accumulator. - let is' = map toInt64Exp is + let is' = map pe64 is -- We need to figure out whether we are updating a scatter-like -- accumulator or a generalised reduction. This also binds the @@ -1227,19 +1244,13 @@ -- | Compile things to 'Imp.Exp'. class ToExp a where - -- | Compile to an 'Imp.Exp', where the type (must must still be a + -- | Compile to an 'Imp.Exp', where the type (which must still be a -- primitive) is deduced monadically. toExp :: a -> ImpM rep r op Imp.Exp -- | Compile where we know the type in advance. toExp' :: PrimType -> a -> Imp.Exp - toInt64Exp :: a -> Imp.TExp Int64 - toInt64Exp = TPrimExp . toExp' int64 - - toBoolExp :: a -> Imp.TExp Bool - toBoolExp = TPrimExp . toExp' Bool - instance ToExp SubExp where toExp (Constant v) = pure $ Imp.ValueExp v @@ -1368,12 +1379,12 @@ ( acc, space, arrs, - map toInt64Exp (shapeDims ispace), + map pe64 (shapeDims ispace), Just op {lambdaParams = ps} ) Just (arrs@(arr : _), Nothing) -> do space <- lookupArraySpace arr - pure (acc, space, arrs, map toInt64Exp (shapeDims ispace), Nothing) + pure (acc, space, arrs, map pe64 (shapeDims ispace), Nothing) Nothing -> error $ "ImpGen.lookupAcc: unlisted accumulator: " ++ pretty name _ -> error $ "ImpGen.lookupAcc: not an accumulator: " ++ pretty name @@ -1588,8 +1599,8 @@ emit $ Imp.Read tmp srcmem srcoffset bt srcspace vol emit $ Imp.Write targetmem targetoffset bt destspace vol $ Imp.var tmp bt | otherwise = do - let destslice' = fullSliceNum (map toInt64Exp destshape) destslice - srcslice' = fullSliceNum (map toInt64Exp srcshape) srcslice + let destslice' = fullSliceNum (map pe64 destshape) destslice + srcslice' = fullSliceNum (map pe64 srcshape) srcslice destrank = length $ sliceDims destslice' srcrank = length $ sliceDims srcslice' destlocation' = sliceMemLoc destlocation destslice' @@ -1749,7 +1760,7 @@ compileAlloc :: Mem rep inner => Pat (LetDec rep) -> SubExp -> Space -> ImpM rep r op () compileAlloc (Pat [mem]) e space = do - let e' = Imp.bytes $ toInt64Exp e + let e' = Imp.bytes $ pe64 e allocator <- asks $ M.lookup space . envAllocCompilers case allocator of Nothing -> emit $ Imp.Allocate (patElemName mem) e' space @@ -1761,7 +1772,7 @@ -- straightforward contiguous format, as an t'Int64' expression. typeSize :: Type -> Count Bytes (Imp.TExp Int64) typeSize t = - Imp.bytes $ primByteSize (elemType t) * product (map toInt64Exp (arrayDims t)) + Imp.bytes $ primByteSize (elemType t) * product (map pe64 (arrayDims t)) -- | Is this indexing in-bounds for an array of the given shape? This -- is useful for things like scatter, which ignores out-of-bounds @@ -1776,6 +1787,13 @@ --- Building blocks for constructing code. +rotateIndex :: + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.TExp Int64 -> + Imp.TExp Int64 +rotateIndex d r i = (i + r) `mod` d + sFor' :: VName -> Imp.Exp -> ImpM rep r op () -> ImpM rep r op () sFor' i bound body = do let it = case primExpType bound of @@ -1913,7 +1931,7 @@ Shape -> ([Imp.TExp Int64] -> ImpM rep r op ()) -> ImpM rep r op () -sLoopNest = sLoopSpace . map toInt64Exp . shapeDims +sLoopNest = sLoopSpace . map pe64 . shapeDims -- | Untyped assignment. (<~~) :: VName -> Imp.Exp -> ImpM rep r op () diff --git a/src/Futhark/CodeGen/ImpGen/GPU.hs b/src/Futhark/CodeGen/ImpGen/GPU.hs --- a/src/Futhark/CodeGen/ImpGen/GPU.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU.hs @@ -135,7 +135,7 @@ -- The calculations are done with 64-bit integers to avoid overflow -- issues. let num_groups_maybe_zero = - sMin64 (toInt64Exp w64 `divUp` toInt64Exp group_size) $ + sMin64 (pe64 w64 `divUp` pe64 group_size) $ sExt64 (tvExp max_num_groups) -- We also don't want zero groups. let num_groups = sMax64 1 num_groups_maybe_zero @@ -247,29 +247,34 @@ x' <- toExp x s' <- toExp s - sIota (patElemName pe) (toInt64Exp n) x' s' et + sIota (patElemName pe) (pe64 n) x' s' et expCompiler (Pat [pe]) (BasicOp (Replicate _ se)) | Acc {} <- patElemType pe = pure () | otherwise = sReplicate (patElemName pe) se +expCompiler (Pat [pe]) (BasicOp (Rotate rs arr)) + | Acc {} <- patElemType pe = pure () + | otherwise = + sRotateKernel (patElemName pe) (map pe64 rs) arr -- Allocation in the "local" space is just a placeholder. expCompiler _ (Op (Alloc _ (Space "local"))) = pure () expCompiler pat (WithAcc inputs lam) = withAcc pat inputs lam --- This is a multi-versioning If created by incremental flattening. +-- This is a multi-versioning Match created by incremental flattening. -- We need to augment the conditional with a check that any local -- memory requirements in tbranch are compatible with the hardware. --- We do not check anything for fbranch, as we assume that it will +-- We do not check anything for defbody, as we assume that it will -- always be safe (and what would we do if none of the branches would -- work?). -expCompiler dest (If cond tbranch fbranch (IfDec _ IfEquiv)) = do - tcode <- collect $ compileBody dest tbranch - fcode <- collect $ compileBody dest fbranch +expCompiler dest (Match cond (first_case : cases) defbranch sort@(MatchDec _ MatchEquiv)) = do + tcode <- collect $ compileBody dest $ caseBody first_case + fcode <- collect $ expCompiler dest $ Match cond cases defbranch sort check <- checkLocalMemoryReqs tcode + let matches = caseMatch cond (casePat first_case) emit $ case check of Nothing -> fcode - Just ok -> Imp.If (ok .&&. toBoolExp cond) tcode fcode + Just ok -> Imp.If (matches .&&. ok) tcode fcode expCompiler dest e = defCompileExp dest e @@ -293,7 +298,7 @@ | bt_size <- primByteSize bt, Just destoffset <- IxFun.linearWithOffset destIxFun bt_size, Just srcoffset <- IxFun.linearWithOffset srcIxFun bt_size = do - let num_elems = Imp.elements $ product $ map toInt64Exp srcshape + let num_elems = Imp.elements $ product $ map pe64 srcshape srcspace <- entryMemSpace <$> lookupMemory srcmem destspace <- entryMemSpace <$> lookupMemory destmem emit diff --git a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/Base.hs @@ -21,6 +21,7 @@ defKernelAttrs, sReplicate, sIota, + sRotateKernel, sCopy, compileThreadResult, compileGroupResult, @@ -97,7 +98,6 @@ kernelGroupSize :: Imp.TExp Int64, kernelNumThreads :: Imp.TExp Int32, kernelWaveSize :: Imp.TExp Int32, - kernelThreadActive :: Imp.TExp Bool, -- | A mapping from dimensions of nested SegOps to already -- computed local thread IDs. Only valid in non-virtualised case. kernelLocalIdMap :: M.Map [SubExp] [Imp.TExp Int32], @@ -121,8 +121,8 @@ _ -> S.singleton $ map snd $ unSegSpace $ segSpace op onExp (BasicOp (Replicate shape _)) = S.singleton $ shapeDims shape - onExp (If _ tbranch fbranch _) = - onStms (bodyStms tbranch) <> onStms (bodyStms fbranch) + onExp (Match _ cases defbody _) = + foldMap (onStms . bodyStms . caseBody) cases <> onStms (bodyStms defbody) onExp (DoLoop _ _ body) = onStms (bodyStms body) onExp _ = mempty @@ -161,7 +161,7 @@ localEnv f m where mkMap ltid dims = do - let dims' = map toInt64Exp dims + let dims' = map pe64 dims ids' <- dIndexSpace' "ltid_pre" dims' (sExt64 ltid) pure (dims, map sExt32 ids') @@ -183,23 +183,22 @@ -- translated to an actual scalar variable during C code generation. pure () kernelAlloc (Pat [mem]) size (Space "local") = - allocLocal (patElemName mem) $ Imp.bytes $ toInt64Exp size + allocLocal (patElemName mem) $ Imp.bytes $ pe64 size kernelAlloc (Pat [mem]) _ _ = compilerLimitationS $ "Cannot allocate memory block " ++ pretty mem ++ " in kernel." kernelAlloc dest _ _ = error $ "Invalid target for in-kernel allocation: " ++ show dest splitSpace :: - (ToExp w, ToExp i, ToExp elems_per_thread) => Pat LetDecMem -> SplitOrdering -> - w -> - i -> - elems_per_thread -> + SubExp -> + SubExp -> + SubExp -> ImpM rep r op () splitSpace (Pat [size]) o w i elems_per_thread = do num_elements <- Imp.elements . TPrimExp <$> toExp w - let i' = toInt64Exp i + let i' = pe64 i elems_per_thread' <- Imp.elements . TPrimExp <$> toExp elems_per_thread computeThreadChunkSize o i' elems_per_thread' num_elements (mkTV (patElemName size) int64) splitSpace pat _ _ _ _ = @@ -208,7 +207,7 @@ updateAcc :: VName -> [SubExp] -> [SubExp] -> InKernelGen () updateAcc acc is vs = sComment "UpdateAcc" $ do -- See the ImpGen implementation of UpdateAcc for general notes. - let is' = map toInt64Exp is + let is' = map pe64 is (c, space, arrs, dims, op) <- lookupAcc acc is' sWhen (inBounds (Slice (map DimFix is')) dims) $ case op of @@ -308,7 +307,7 @@ localThreadIDs :: [SubExp] -> InKernelGen [Imp.TExp Int64] localThreadIDs dims = do ltid <- sExt64 . kernelLocalThreadId . kernelConstants <$> askEnv - let dims' = map toInt64Exp dims + let dims' = map pe64 dims maybe (dIndexSpace' "ltid" dims' ltid) (pure . map sExt64) . M.lookup dims . kernelLocalIdMap @@ -374,6 +373,14 @@ groupCoverSegSpace SegVirt (SegSpace flat $ zip is $ shapeDims ds) $ copyDWIMFix (patElemName dest) is' se [] sOp $ Imp.Barrier Imp.FenceLocal +compileGroupExp (Pat [dest]) (BasicOp (Rotate rs arr)) = do + ds <- map pe64 . arrayDims <$> lookupType arr + groupCoverSpace ds $ \is -> do + is' <- sequence $ zipWith3 rotate ds rs is + copyDWIMFix (patElemName dest) is (Var arr) is' + sOp $ Imp.Barrier Imp.FenceLocal + where + rotate d r i = dPrimVE "rot_i" $ rotateIndex d (pe64 r) i compileGroupExp (Pat [dest]) (BasicOp (Iota n e s it)) = do n' <- toExp n e' <- toExp e @@ -401,8 +408,8 @@ Safe -> sWhen (inBounds slice' dims) write sOp $ Imp.Barrier Imp.FenceLocal where - slice' = fmap toInt64Exp slice - dims = map toInt64Exp $ arrayDims $ patElemType pe + slice' = fmap pe64 slice + dims = map pe64 $ arrayDims $ patElemType pe write = copyDWIM (patElemName pe) (unSlice slice') se [] compileGroupExp dest e = defCompileExp dest e @@ -439,8 +446,8 @@ (Nothing, AtomicLocking f) -> do locks <- newVName "locks" - let num_locks = toInt64Exp $ unCount group_size - dims = map toInt64Exp $ shapeDims (histOpShape op <> histShape op) + let num_locks = pe64 $ unCount group_size + dims = map pe64 $ shapeDims (histOpShape op <> histShape op) l' = Locking locks 0 1 0 (pure . (`rem` num_locks) . flattenIndex dims) locks_t = Array int32 (Shape [unCount group_size]) NoUniqueness @@ -512,7 +519,7 @@ let flat_shape = Shape $ Var (tvVar flat) : drop k (memLocShape arr_loc) sArray (baseString arr ++ "_flat") pt flat_shape (memLocName arr_loc) $ IxFun.reshape (memLocIxFun arr_loc) $ - map (DimNew . pe64) $ + map pe64 $ shapeDims flat_shape -- | @applyLambda lam dests args@ emits code that: @@ -608,7 +615,7 @@ compileFlatId lvl space let (ltids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims + dims' = map pe64 dims groupCoverSegSpace (segVirt lvl) space $ compileStms mempty (kernelBodyStms body) $ @@ -654,7 +661,7 @@ compileGroupOp pat (Inner (SegOp (SegRed lvl space ops _ body))) = do compileFlatId lvl space - let dims' = map toInt64Exp dims + let dims' = map pe64 dims mkTempArr t = sAllocArray "red_arr" (elemType t) (Shape dims <> arrayShape t) $ Space "local" @@ -792,8 +799,8 @@ forM_ (zip4 red_is vs_per_op ops' ops) $ \(bin, op_vs, do_op, HistOp dest_shape _ _ _ shape lam) -> do - let bin' = toInt64Exp bin - dest_shape' = map toInt64Exp $ shapeDims dest_shape + let bin' = pe64 bin + dest_shape' = map pe64 $ shapeDims dest_shape bin_in_bounds = inBounds (Slice (map DimFix [bin'])) dest_shape' bin_is = map Imp.le64 (init ltids) ++ [bin'] vs_params = takeLast (length op_vs) $ lambdaParams lam @@ -1106,7 +1113,7 @@ chunk_var <-- sMin64 (Imp.unCount elements_per_thread) - ((Imp.unCount num_elements - thread_index) `divUp` toInt64Exp stride) + ((Imp.unCount num_elements - thread_index) `divUp` pe64 stride) computeThreadChunkSize SplitContiguous thread_index elements_per_thread num_elements chunk_var = do starting_point <- dPrimV "starting_point" $ @@ -1160,7 +1167,6 @@ kernelGroupSize = group_size', kernelNumThreads = sExt32 (group_size' * num_groups'), kernelWaveSize = Imp.le32 wave_size, - kernelThreadActive = true, kernelLocalIdMap = mempty, kernelChunkItersMap = mempty } @@ -1185,7 +1191,7 @@ x : xs -> foldl (.&&.) x xs where (is, ws) = unzip limit - actives = zipWith active is $ map toInt64Exp ws + actives = zipWith active is $ map pe64 ws active i = (Imp.le64 i .<.) -- | Change every memory block to be in the global address space, @@ -1553,28 +1559,61 @@ | otherwise = copyDWIM (paramName y) [] (Var $ paramName x) [] -computeMapKernelGroups :: +simpleKernelGroups :: Imp.TExp Int64 -> - CallKernelGen (Count NumGroups SubExp, Count GroupSize SubExp) -computeMapKernelGroups kernel_size = do + Imp.TExp Int64 -> + CallKernelGen (Imp.TExp Int32, Count NumGroups SubExp, Count GroupSize SubExp) +simpleKernelGroups max_num_groups kernel_size = do group_size <- dPrim "group_size" int64 fname <- askFunction let group_size_key = keyWithEntryPoint fname $ nameFromString $ pretty $ tvVar group_size sOp $ Imp.GetSize (tvVar group_size) group_size_key Imp.SizeGroup - num_groups <- dPrimV "num_groups" $ kernel_size `divUp` tvExp group_size - pure (Count $ tvSize num_groups, Count $ tvSize group_size) + virt_num_groups <- dPrimVE "virt_num_groups" $ kernel_size `divUp` tvExp group_size + num_groups <- dPrimV "num_groups" $ virt_num_groups `sMin64` max_num_groups + pure (sExt32 virt_num_groups, Count $ tvSize num_groups, Count $ tvSize group_size) simpleKernelConstants :: Imp.TExp Int64 -> String -> - CallKernelGen (KernelConstants, InKernelGen ()) + CallKernelGen + ( (Imp.TExp Int64 -> InKernelGen ()) -> InKernelGen (), + KernelConstants + ) simpleKernelConstants kernel_size desc = do + -- For performance reasons, codegen assumes that the thread count is + -- never more than will fit in an i32. This means we need to cap + -- the number of groups here. The cap is set much higher than any + -- GPU will possibly need. Feel free to come back and laugh at me + -- in the future. + let max_num_groups = 1024 * 1024 thread_gtid <- newVName $ desc ++ "_gtid" thread_ltid <- newVName $ desc ++ "_ltid" group_id <- newVName $ desc ++ "_gid" inner_group_size <- newVName "group_size" - (num_groups, group_size) <- computeMapKernelGroups kernel_size - let set_constants = do + (virt_num_groups, num_groups, group_size) <- + simpleKernelGroups max_num_groups kernel_size + let group_size' = Imp.pe64 $ unCount group_size + num_groups' = Imp.pe64 $ unCount num_groups + + constants = + KernelConstants + { kernelGlobalThreadId = Imp.le32 thread_gtid, + kernelLocalThreadId = Imp.le32 thread_ltid, + kernelGroupId = Imp.le32 group_id, + kernelGlobalThreadIdVar = thread_gtid, + kernelLocalThreadIdVar = thread_ltid, + kernelGroupIdVar = group_id, + kernelNumGroupsCount = num_groups, + kernelGroupSizeCount = group_size, + kernelNumGroups = num_groups', + kernelGroupSize = group_size', + kernelNumThreads = sExt32 (group_size' * num_groups'), + kernelWaveSize = 0, + kernelLocalIdMap = mempty, + kernelChunkItersMap = mempty + } + + wrapKernel m = do dPrim_ thread_ltid int32 dPrim_ inner_group_size int64 dPrim_ group_id int32 @@ -1582,29 +1621,14 @@ sOp (Imp.GetLocalSize inner_group_size 0) sOp (Imp.GetGroupId group_id 0) dPrimV_ thread_gtid $ le32 group_id * le32 inner_group_size + le32 thread_ltid - group_size' = Imp.pe64 $ unCount group_size - num_groups' = Imp.pe64 $ unCount num_groups + virtualiseGroups SegVirt virt_num_groups $ \virt_group_id -> do + global_tid <- + dPrimVE "global_tid" $ + sExt64 virt_group_id * sExt64 (le32 inner_group_size) + + sExt64 (kernelLocalThreadId constants) + m global_tid - pure - ( KernelConstants - { kernelGlobalThreadId = Imp.le32 thread_gtid, - kernelLocalThreadId = Imp.le32 thread_ltid, - kernelGroupId = Imp.le32 group_id, - kernelGlobalThreadIdVar = thread_gtid, - kernelLocalThreadIdVar = thread_ltid, - kernelGroupIdVar = group_id, - kernelNumGroupsCount = num_groups, - kernelGroupSizeCount = group_size, - kernelNumGroups = num_groups', - kernelGroupSize = group_size', - kernelNumThreads = sExt32 (group_size' * num_groups'), - kernelWaveSize = 0, - kernelThreadActive = Imp.le64 thread_gtid .<. kernel_size, - kernelLocalIdMap = mempty, - kernelChunkItersMap = mempty - }, - set_constants - ) + pure (wrapKernel, constants) -- | For many kernels, we may not have enough physical groups to cover -- the logical iteration space. Some groups thus have to perform @@ -1792,9 +1816,9 @@ t <- subExpType se ds <- dropLast (arrayRank t) . arrayDims <$> lookupType arr - let dims = map toInt64Exp $ ds ++ arrayDims t - (constants, set_constants) <- - simpleKernelConstants (product $ map sExt64 dims) "replicate" + let dims = map pe64 $ ds ++ arrayDims t + n <- dPrimVE "replicate_n" $ product $ map sExt64 dims + (virtualise, constants) <- simpleKernelConstants n "replicate" fname <- askFunction let name = @@ -1802,12 +1826,12 @@ nameFromString $ "replicate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - is' <- dIndexSpace' "rep_i" dims $ sExt64 $ kernelGlobalThreadId constants - sWhen (kernelThreadActive constants) $ - copyDWIMFix arr is' se $ - drop (length ds) is' + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> do + is' <- dIndexSpace' "rep_i" dims gtid + sWhen (gtid .<. n) $ + copyDWIMFix arr is' se $ + drop (length ds) is' replicateName :: PrimType -> String replicateName bt = "replicate_" ++ pretty bt @@ -1824,16 +1848,13 @@ let params = [ Imp.MemParam mem (Space "device"), - Imp.ScalarParam num_elems int32, + Imp.ScalarParam num_elems int64, Imp.ScalarParam val bt ] shape = Shape [Var num_elems] function fname [] params $ do arr <- - sArray "arr" bt shape mem $ - IxFun.iota $ - map pe64 $ - shapeDims shape + sArray "arr" bt shape mem $ IxFun.iota $ map pe64 $ shapeDims shape sReplicateKernel arr $ Var val pure fname @@ -1852,7 +1873,7 @@ [] fname [ Imp.MemArg arr_mem, - Imp.ExpArg $ untyped $ product $ map toInt64Exp arr_shape, + Imp.ExpArg $ untyped $ product $ map pe64 arr_shape, Imp.ExpArg $ toExp' v_t' v ] _ -> pure Nothing @@ -1878,7 +1899,7 @@ CallKernelGen () sIotaKernel arr n x s et = do destloc <- entryArrayLoc <$> lookupArray arr - (constants, set_constants) <- simpleKernelConstants n "iota" + (virtualise, constants) <- simpleKernelConstants n "iota" fname <- askFunction let name = @@ -1889,18 +1910,17 @@ ++ "_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - sKernelFailureTolerant True threadOperations constants name $ do - set_constants - let gtid = sExt64 $ kernelGlobalThreadId constants - sWhen (kernelThreadActive constants) $ do - (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> + sWhen (gtid .<. n) $ do + (destmem, destspace, destidx) <- fullyIndexArray' destloc [gtid] - emit $ - Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ - BinOpExp - (Add et OverflowWrap) - (BinOpExp (Mul et OverflowWrap) (Imp.sExt et $ untyped gtid) s) - x + emit $ + Imp.Write destmem destidx (IntType et) destspace Imp.Nonvolatile $ + BinOpExp + (Add et OverflowWrap) + (BinOpExp (Mul et OverflowWrap) (Imp.sExt et $ untyped gtid) s) + x iotaName :: IntType -> String iotaName bt = "iota_" ++ pretty bt @@ -1961,10 +1981,10 @@ sCopy pt destloc@(MemLoc destmem _ _) srcloc@(MemLoc srcmem srcdims _) = do -- Note that the shape of the destination and the source are -- necessarily the same. - let shape = map toInt64Exp srcdims + let shape = map pe64 srcdims kernel_size = product shape - (constants, set_constants) <- simpleKernelConstants kernel_size "copy" + (virtualise, constants) <- simpleKernelConstants kernel_size "copy" fname <- askFunction let name = @@ -1972,48 +1992,68 @@ nameFromString $ "copy_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) - sKernelFailureTolerant True threadOperations constants name $ do - set_constants + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> do + is <- dIndexSpace' "copy_i" shape gtid - let gtid = sExt64 $ kernelGlobalThreadId constants - is <- dIndexSpace' "copy_i" shape gtid + (_, destspace, destidx) <- fullyIndexArray' destloc is + (_, srcspace, srcidx) <- fullyIndexArray' srcloc is - (_, destspace, destidx) <- fullyIndexArray' destloc is - (_, srcspace, srcidx) <- fullyIndexArray' srcloc is + sWhen (gtid .<. kernel_size) $ do + tmp <- tvVar <$> dPrim "tmp" pt + emit $ Imp.Read tmp srcmem srcidx pt srcspace Imp.Nonvolatile + emit $ Imp.Write destmem destidx pt destspace Imp.Nonvolatile $ Imp.var tmp pt - sWhen (gtid .<. kernel_size) $ do - tmp <- tvVar <$> dPrim "tmp" pt - emit $ Imp.Read tmp srcmem srcidx pt srcspace Imp.Nonvolatile - emit $ Imp.Write destmem destidx pt destspace Imp.Nonvolatile $ Imp.var tmp pt +-- | Perform a Rotate with a kernel. +sRotateKernel :: VName -> [Imp.TExp Int64] -> VName -> CallKernelGen () +sRotateKernel dest rs src = do + t <- lookupType src + let ds = map pe64 $ arrayDims t + n <- dPrimVE "rotate_n" $ product ds + (virtualise, constants) <- simpleKernelConstants n "rotate" + fname <- askFunction + let name = + keyWithEntryPoint fname $ + nameFromString $ + "rotate_" ++ show (baseTag $ kernelGlobalThreadIdVar constants) + + sKernelFailureTolerant True threadOperations constants name $ + virtualise $ \gtid -> sWhen (gtid .<. n) $ do + is' <- dIndexSpace' "rep_i" ds gtid + is'' <- sequence $ zipWith3 rotate ds rs is' + copyDWIMFix dest is' (Var src) is'' + where + rotate d r i = dPrimVE "rot_i" $ rotateIndex d r i + compileGroupResult :: SegSpace -> PatElem LetDecMem -> KernelResult -> InKernelGen () compileGroupResult _ pe (TileReturns _ [(w, per_group_elems)] what) = do - n <- toInt64Exp . arraySize 0 <$> lookupType what + n <- pe64 . arraySize 0 <$> lookupType what constants <- kernelConstants <$> askEnv let ltid = sExt64 $ kernelLocalThreadId constants offset = - toInt64Exp per_group_elems + pe64 per_group_elems * sExt64 (kernelGroupId constants) -- Avoid loop for the common case where each thread is statically -- known to write at most one element. localOps threadOperations $ - if toInt64Exp per_group_elems == kernelGroupSize constants + if pe64 per_group_elems == kernelGroupSize constants then - sWhen (ltid + offset .<. toInt64Exp w) $ + sWhen (ltid + offset .<. pe64 w) $ copyDWIMFix (patElemName pe) [ltid + offset] (Var what) [ltid] else sFor "i" (n `divUp` kernelGroupSize constants) $ \i -> do j <- dPrimVE "j" $ kernelGroupSize constants * i + ltid - sWhen (j + offset .<. toInt64Exp w) $ + sWhen (j + offset .<. pe64 w) $ copyDWIMFix (patElemName pe) [j + offset] (Var what) [j] compileGroupResult space pe (TileReturns _ dims what) = do let gids = map fst $ unSegSpace space - out_tile_sizes = map (toInt64Exp . snd) dims + out_tile_sizes = map (pe64 . snd) dims group_is = zipWith (*) (map Imp.le64 gids) out_tile_sizes local_is <- localThreadIDs $ map snd dims is_for_thread <- @@ -2028,8 +2068,8 @@ let gids = map fst $ unSegSpace space (dims, group_tiles, reg_tiles) = unzip3 dims_n_tiles - group_tiles' = map toInt64Exp group_tiles - reg_tiles' = map toInt64Exp reg_tiles + group_tiles' = map pe64 group_tiles + reg_tiles' = map pe64 reg_tiles -- Which group tile is this group responsible for? let group_tile_is = map Imp.le64 gids @@ -2057,7 +2097,7 @@ sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do let dest_is = fixSlice reg_tile_slices is_in_reg_tile src_is = reg_tile_is ++ is_in_reg_tile - sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map toInt64Exp dims) $ + sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map pe64 dims) $ copyDWIMFix (patElemName pe) dest_is (Var what) src_is compileGroupResult space pe (Returns _ _ what) = do constants <- kernelConstants <$> askEnv @@ -2092,20 +2132,19 @@ compileThreadResult _ pe (ConcatReturns _ SplitContiguous _ per_thread_elems what) = do constants <- kernelConstants <$> askEnv let offset = - toInt64Exp per_thread_elems + pe64 per_thread_elems * sExt64 (kernelGlobalThreadId constants) - n <- toInt64Exp . arraySize 0 <$> lookupType what + n <- pe64 . arraySize 0 <$> lookupType what copyDWIM (patElemName pe) [DimSlice offset n 1] (Var what) [] compileThreadResult _ pe (ConcatReturns _ (SplitStrided stride) _ _ what) = do offset <- sExt64 . kernelGlobalThreadId . kernelConstants <$> askEnv - n <- toInt64Exp . arraySize 0 <$> lookupType what - copyDWIM (patElemName pe) [DimSlice offset n $ toInt64Exp stride] (Var what) [] + n <- pe64 . arraySize 0 <$> lookupType what + copyDWIM (patElemName pe) [DimSlice offset n $ pe64 stride] (Var what) [] compileThreadResult _ pe (WriteReturns _ (Shape rws) _arr dests) = do - constants <- kernelConstants <$> askEnv - let rws' = map toInt64Exp rws + let rws' = map pe64 rws forM_ dests $ \(slice, e) -> do - let slice' = fmap toInt64Exp slice - write = kernelThreadActive constants .&&. inBounds slice' rws' + let slice' = fmap pe64 slice + write = inBounds slice' rws' sWhen write $ copyDWIM (patElemName pe) (unSlice slice') e [] compileThreadResult _ _ TileReturns {} = compilerBugS "compileThreadResult: TileReturns unhandled." diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegHist.hs @@ -76,7 +76,7 @@ histOp op histSize :: HistOp GPUMem -> Imp.TExp Int64 -histSize = product . map toInt64Exp . shapeDims . histShape +histSize = product . map pe64 . shapeDims . histShape histRank :: HistOp GPUMem -> Int histRank = shapeRank . histShape @@ -127,7 +127,7 @@ Space "device" multiHistoCase = do - let num_elems = product $ map toInt64Exp $ shapeDims subhistos_shape + let num_elems = product $ map pe64 $ shapeDims subhistos_shape subhistos_mem_size = Imp.bytes $ Imp.unCount (Imp.elements num_elems `Imp.withElemType` elemType dest_t) @@ -136,15 +136,15 @@ sReplicate subhistos ne subhistos_t <- lookupType subhistos let slice = - fullSliceNum (map toInt64Exp $ arrayDims subhistos_t) $ - map (unitSlice 0 . toInt64Exp . snd) segment_dims + fullSliceNum (map pe64 $ arrayDims subhistos_t) $ + map (unitSlice 0 . pe64 . snd) segment_dims ++ [DimFix 0] sUpdate subhistos slice $ Var dest sIf (tvExp num_subhistos .==. 1) unitHistoCase multiHistoCase let h = histSpaceUsage op - segmented_h = h * product (map (Imp.bytes . toInt64Exp) $ init $ segSpaceDims space) + segmented_h = h * product (map (Imp.bytes . pe64) $ init $ segSpaceDims space) atomics <- hostAtomics <$> askEnv @@ -181,7 +181,7 @@ -- algorithm to ensure good distribution of locks. let num_locks = 100151 dims = - map toInt64Exp $ + map pe64 $ shapeDims (histOpShape (slugOp slug)) ++ [tvSize (slugNumSubhistos slug)] ++ shapeDims (histShape (slugOp slug)) @@ -225,7 +225,7 @@ hist_RF <- dPrimVE "hist_RF" $ - sum (map (r64 . toInt64Exp . histRaceFactor . slugOp) slugs) + sum (map (r64 . pe64 . histRaceFactor . slugOp) slugs) / genericLength slugs hist_el_size <- dPrimVE "hist_el_size" $ sum $ map slugElAvgSize slugs @@ -399,7 +399,7 @@ CallKernelGen () histKernelGlobalPass map_pes num_groups group_size space slugs kbody histograms hist_S chk_i = do let (space_is, space_sizes) = unzip $ unSegSpace space - space_sizes_64 = map (sExt64 . toInt64Exp) space_sizes + space_sizes_64 = map (sExt64 . pe64) space_sizes total_w_64 = product space_sizes_64 hist_H_chks <- forM (map (histSize . slugOp) slugs) $ \w -> @@ -456,8 +456,8 @@ hist_H_chk ) -> do let chk_beg = sExt64 chk_i * hist_H_chk - bucket' = map toInt64Exp bucket - dest_shape' = map toInt64Exp $ shapeDims dest_shape + bucket' = map pe64 bucket + dest_shape' = map pe64 $ shapeDims dest_shape flat_bucket = flattenIndex dest_shape' bucket' bucket_in_bounds = chk_beg .<=. flat_bucket @@ -485,8 +485,8 @@ KernelBody GPUMem -> CallKernelGen () histKernelGlobal map_pes num_groups group_size space slugs kbody = do - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size + let num_groups' = fmap pe64 num_groups + group_size' = fmap pe64 group_size let (_space_is, space_sizes) = unzip $ unSegSpace space num_threads = sExt32 $ unCount num_groups' * unCount group_size' @@ -496,7 +496,7 @@ prepareIntermediateArraysGlobal (bodyPassage kbody) num_threads - (toInt64Exp $ last space_sizes) + (pe64 $ last space_sizes) slugs sFor "chk_i" hist_S $ \chk_i -> @@ -549,7 +549,7 @@ : shapeDims (histOpShape op) ++ [hist_H_chk] - let dims = map toInt64Exp $ shapeDims lock_shape + let dims = map pe64 $ shapeDims lock_shape locks <- sAllocArray "locks" int32 lock_shape $ Space "local" @@ -614,12 +614,12 @@ segment_dims = init space_sizes (i_in_segment, segment_size) = last $ unSegSpace space num_subhistos_per_group = tvExp num_subhistos_per_group_var - segment_size' = toInt64Exp segment_size + segment_size' = pe64 segment_size num_segments <- dPrimVE "num_segments" $ product $ - map toInt64Exp segment_dims + map pe64 segment_dims hist_H_chks <- forM (map slugOp slugs) $ \op -> dPrimV "hist_H_chk" $ histSize op `divUp` sExt64 hist_S @@ -627,7 +627,7 @@ histo_sizes <- forM (zip slugs hist_H_chks) $ \(slug, hist_H_chk) -> do let histo_dims = tvExp hist_H_chk - : map toInt64Exp (shapeDims (histOpShape (slugOp slug))) + : map pe64 (shapeDims (histOpShape (slugOp slug))) histo_size <- dPrimVE "histo_size" $ product histo_dims let group_hists_size = @@ -656,7 +656,7 @@ -- Set segment indices. zipWithM_ dPrimV_ segment_is $ - unflattenIndex (map toInt64Exp segment_dims) $ + unflattenIndex (map pe64 segment_dims) $ sExt64 flat_segment_id histograms <- forM (zip init_histograms hist_H_chks) $ @@ -694,7 +694,7 @@ local_subhisto_i <- dPrimVE "local_subhisto_i" $ j `quot` sExt32 histo_size let local_bucket_is = unflattenIndex histo_dims $ sExt64 $ j `rem` sExt32 histo_size nested_hist_size = - map toInt64Exp $ shapeDims $ histShape $ slugOp slug + map pe64 $ shapeDims $ histShape $ slugOp slug global_bucket_is = unflattenIndex @@ -757,8 +757,8 @@ (bucket, vs') ) -> do let chk_beg = sExt64 chk_i * tvExp hist_H_chk - bucket' = map toInt64Exp bucket - dest_shape' = map toInt64Exp $ shapeDims dest_shape + bucket' = map pe64 bucket + dest_shape' = map pe64 $ shapeDims dest_shape flat_bucket = flattenIndex dest_shape' bucket' bucket_in_bounds = inBounds (Slice (map DimFix bucket')) dest_shape' @@ -785,7 +785,7 @@ histSize (slugOp slug) - sExt64 chk_i * head histo_dims let trunc_histo_dims = tvExp trunc_H - : map toInt64Exp (shapeDims (histOpShape (slugOp slug))) + : map pe64 (shapeDims (histOpShape (slugOp slug))) trunc_histo_size <- dPrimVE "histo_size" $ sExt32 $ product trunc_histo_dims sFor "local_i" bins_per_thread $ \i -> do @@ -798,7 +798,7 @@ -- we immediately unflatten. let local_bucket_is = unflattenIndex histo_dims $ sExt64 j nested_hist_size = - map toInt64Exp $ shapeDims $ histShape $ slugOp slug + map pe64 $ shapeDims $ histShape $ slugOp slug global_bucket_is = unflattenIndex nested_hist_size @@ -908,9 +908,9 @@ num_groups <- fmap (Imp.Count . tvSize) $ dPrimV "num_groups" $ - sExt64 hist_T `divUp` toInt64Exp (unCount group_size) - let num_groups' = toInt64Exp <$> num_groups - group_size' = toInt64Exp <$> group_size + sExt64 hist_T `divUp` pe64 (unCount group_size) + let num_groups' = pe64 <$> num_groups + group_size' = pe64 <$> group_size let r64 = isF64 . ConvOpExp (SIToFP Int64 Float64) . untyped t64 = isInt64 . ConvOpExp (FPToSI Float64 Int64) . untyped @@ -937,9 +937,9 @@ let q_small = 2 -- The number of segments/histograms produced.. - hist_Nout <- dPrimVE "hist_Nout" $ product $ map toInt64Exp segment_dims + hist_Nout <- dPrimVE "hist_Nout" $ product $ map pe64 segment_dims - hist_Nin <- dPrimVE "hist_Nin" $ toInt64Exp $ last space_sizes + hist_Nin <- dPrimVE "hist_Nin" $ pe64 $ last space_sizes -- Maximum M for work efficiency. work_asymp_M_max <- @@ -1059,9 +1059,9 @@ -- rather figuring out whether to use a local or global memory -- strategy, as well as collapsing the subhistograms produced (which -- are always in global memory, but their number may vary). - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size - dims = map toInt64Exp $ segSpaceDims space + let num_groups' = fmap pe64 num_groups + group_size' = fmap pe64 group_size + dims = map pe64 $ segSpaceDims space num_red_res = length ops + sum (map (length . histNeutral) ops) (all_red_pes, map_pes) = splitAt num_red_res pes @@ -1097,7 +1097,7 @@ hist_RF <- dPrimVE "hist_RF" $ sExt32 $ - sum (map (toInt64Exp . histRaceFactor . slugOp) slugs) + sum (map (pe64 . histRaceFactor . slugOp) slugs) `quot` genericLength slugs let hist_T = sExt32 $ unCount num_groups' * unCount group_size' @@ -1111,7 +1111,7 @@ Just $ untyped $ product $ - map (toInt64Exp . snd) segment_dims + map (pe64 . snd) segment_dims emit $ Imp.DebugPrint "Histogram element size (el_size)" $ Just $ untyped hist_el_size emit $ Imp.DebugPrint "Race factor (RF)" $ Just $ untyped hist_RF emit $ Imp.DebugPrint "Memory per set of subhistograms per segment" $ Just $ untyped h diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegMap.hs @@ -24,8 +24,8 @@ CallKernelGen () compileSegMap pat lvl space kbody = do let (is, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - group_size' = toInt64Exp <$> segGroupSize lvl + dims' = map pe64 dims + group_size' = pe64 <$> segGroupSize lvl attrs = defKernelAttrs (segNumGroups lvl) (segGroupSize lvl) emit $ Imp.DebugPrint "\n# SegMap" Nothing diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs @@ -109,8 +109,8 @@ | [(_, Constant (IntValue (Int64Value 1))), _] <- unSegSpace space = nonsegmentedReduction pat num_groups group_size space reds body | otherwise = do - let group_size' = toInt64Exp $ unCount group_size - segment_size = toInt64Exp $ last $ segSpaceDims space + let group_size' = pe64 $ unCount group_size + segment_size = pe64 $ last $ segSpaceDims space use_small_segments = segment_size * 2 .<. group_size' sIf use_small_segments @@ -182,9 +182,9 @@ CallKernelGen () nonsegmentedReduction segred_pat num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims - num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size + dims' = map pe64 dims + num_groups' = fmap pe64 num_groups + group_size' = fmap pe64 group_size global_tid = Imp.le64 $ segFlat space w = last dims' @@ -266,15 +266,15 @@ CallKernelGen () smallSegmentsReduction (Pat segred_pes) num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims + dims' = map pe64 dims segment_size = last dims' -- Careful to avoid division by zero now. segment_size_nonzero <- dPrimVE "segment_size_nonzero" $ sMax64 1 segment_size - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size + let num_groups' = fmap pe64 num_groups + group_size' = fmap pe64 group_size num_threads <- dPrimV "num_threads" $ unCount num_groups' * unCount group_size' let num_segments = product $ init dims' segments_per_group = unCount group_size' `quot` segment_size_nonzero @@ -383,11 +383,11 @@ CallKernelGen () largeSegmentsReduction segred_pat num_groups group_size space reds body = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims + dims' = map pe64 dims num_segments = product $ init dims' segment_size = last dims' - num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size + num_groups' = fmap pe64 num_groups + group_size' = fmap pe64 group_size (groups_per_segment, elems_per_thread) <- groupsPerSegmentAndElementsPerThread @@ -458,7 +458,7 @@ let first_group_for_segment = sExt64 flat_segment_id * groups_per_segment dIndexSpace (zip segment_gtids (init dims')) $ sExt64 flat_segment_id dPrim_ (last gtids) int64 - let num_elements = Imp.elements $ toInt64Exp w + let num_elements = Imp.elements $ pe64 w slugs <- mapM (segBinOpSlug local_tid group_id) $ diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan.hs @@ -64,4 +64,4 @@ SinglePass.compileSegScan pat lvl space scan' kbody _ -> TwoPass.compileSegScan pat lvl space scans kbody where - n = product $ map toInt64Exp $ segSpaceDims space + n = product $ map pe64 $ segSpaceDims space diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/SinglePass.hs @@ -34,8 +34,8 @@ [PrimType] -> InKernelGen (VName, [VName], [VName], VName, [VName]) createLocalArrays (Count groupSize) m types = do - let groupSizeE = toInt64Exp groupSize - workSize = toInt64Exp m * groupSizeE + let groupSizeE = pe64 groupSize + workSize = pe64 m * groupSizeE prefixArraysSize = foldl (\acc tySize -> alignTo acc tySize + tySize * groupSizeE) 0 $ map primByteSize types @@ -55,7 +55,7 @@ byteOffsets <- mapM (fmap varTE . dPrimV "byte_offsets") $ - scanl (\off tySize -> alignTo off tySize + toInt64Exp groupSize * tySize) 0 $ + scanl (\off tySize -> alignTo off tySize + pe64 groupSize * tySize) 0 $ map primByteSize types warpByteOffsets <- @@ -245,7 +245,7 @@ dPrimVE "num_threads" $ num_groups' * group_size' let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims + dims' = map pe64 dims segmented = length dims' > 1 not_segmented_e = if segmented then false else true segment_size = last dims' diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegScan/TwoPass.hs @@ -156,12 +156,12 @@ KernelBody GPUMem -> CallKernelGen (TV Int32, Imp.TExp Int64, CrossesSegment) scanStage1 (Pat all_pes) num_groups group_size space scans kbody = do - let num_groups' = fmap toInt64Exp num_groups - group_size' = fmap toInt64Exp group_size + let num_groups' = fmap pe64 num_groups + group_size' = fmap pe64 group_size num_threads <- dPrimV "num_threads" $ sExt32 $ unCount num_groups' * unCount group_size' let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims + dims' = map pe64 dims let num_elements = product dims' elems_per_thread = num_elements `divUp` sExt64 (tvExp num_threads) elems_per_group = unCount group_size' * elems_per_thread @@ -332,7 +332,7 @@ CallKernelGen () scanStage2 (Pat all_pes) stage1_num_threads elems_per_group num_groups crossesSegment space scans = do let (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims + dims' = map pe64 dims -- Our group size is the number of groups for the stage 1 kernel. let group_size = Count $ unCount num_groups @@ -407,9 +407,9 @@ [SegBinOp GPUMem] -> CallKernelGen () scanStage3 (Pat all_pes) num_groups group_size elems_per_group crossesSegment space scans = do - let group_size' = fmap toInt64Exp group_size + let group_size' = fmap pe64 group_size (gtids, dims) = unzip $ unSegSpace space - dims' = map toInt64Exp dims + dims' = map pe64 dims required_groups <- dPrimVE "required_groups" $ sExt32 $ @@ -502,7 +502,7 @@ fmap (Imp.Count . tvSize) $ dPrimV "stage1_num_groups" $ sMin64 (tvExp stage1_max_num_groups) $ - toInt64Exp $ + pe64 $ Imp.unCount $ segNumGroups lvl diff --git a/src/Futhark/CodeGen/ImpGen/Multicore.hs b/src/Futhark/CodeGen/ImpGen/Multicore.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore.hs @@ -57,7 +57,7 @@ updateAcc :: VName -> [SubExp] -> [SubExp] -> MulticoreGen () updateAcc acc is vs = sComment "UpdateAcc" $ do -- See the ImpGen implementation of UpdateAcc for general notes. - let is' = map toInt64Exp is + let is' = map pe64 is (c, _space, arrs, dims, op) <- lookupAcc acc is' sWhen (inBounds (Slice (map DimFix is')) dims) $ case op of diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs @@ -96,11 +96,11 @@ getIterationDomain :: SegOp () MCMem -> SegSpace -> MulticoreGen (Imp.TExp Int64) getIterationDomain SegMap {} space = do let ns = map snd $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns pure $ product ns_64 getIterationDomain _ space = do let ns = map snd $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns case unSegSpace space of [_] -> pure $ product ns_64 -- A segmented SegOp is over the segments @@ -346,9 +346,9 @@ where sLoopNest' is [] f = f $ reverse is sLoopNest' is [d] f = - sForVectorized "nest_i" (toInt64Exp d) $ \i -> sLoopNest' (i : is) [] f + sForVectorized "nest_i" (pe64 d) $ \i -> sLoopNest' (i : is) [] f sLoopNest' is (d : ds) f = - sFor "nest_i" (toInt64Exp d) $ \i -> sLoopNest' (i : is) ds f + sFor "nest_i" (pe64 d) $ \i -> sLoopNest' (i : is) ds f ------------------------------- ------- SegHist helpers ------- diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegHist.hs @@ -35,7 +35,7 @@ segHistOpChunks = chunks . map (length . histNeutral) histSize :: HistOp MCMem -> Imp.TExp Int64 -histSize = product . map toInt64Exp . shapeDims . histShape +histSize = product . map pe64 . shapeDims . histShape genHistOpParams :: HistOp MCMem -> MulticoreGen () genHistOpParams histops = @@ -56,7 +56,7 @@ MulticoreGen Imp.MCCode nonsegmentedHist pat space histops kbody num_histos = do let ns = map snd $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns num_histos' = tvExp num_histos hist_width = histSize $ head histops use_subhistogram = sExt64 num_histos' * hist_width .<=. product ns_64 @@ -94,7 +94,7 @@ -- Allocate a static array of locks -- as in the GPU backend let num_locks = 100151 -- This number is taken from the GPU backend - dims = map toInt64Exp $ shapeDims (histOpShape op <> histShape op) + dims = map pe64 $ shapeDims (histOpShape op <> histShape op) locks <- sStaticArray "hist_locks" DefaultSpace int32 $ Imp.ArrayZeros num_locks @@ -109,7 +109,7 @@ MulticoreGen () atomicHistogram pat space histops kbody = do let (is, ns) = unzip $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns let num_red_res = length histops + sum (map (length . histNeutral) histops) (all_red_pes, map_pes) = splitAt num_red_res $ patElems pat @@ -129,8 +129,8 @@ forM_ (zip4 histops red_res_split atomicOps pes_per_op) $ \(HistOp dest_shape _ _ _ shape lam, (bucket, vs'), do_op, dest_res) -> do let (_is_params, vs_params) = splitAt (length vs') $ lambdaParams lam - dest_shape' = map toInt64Exp $ shapeDims dest_shape - bucket' = map toInt64Exp bucket + dest_shape' = map pe64 $ shapeDims dest_shape + bucket' = map pe64 bucket bucket_in_bounds = inBounds (Slice (map DimFix bucket')) dest_shape' sComment "save map-out results" $ @@ -187,7 +187,7 @@ emit $ Imp.DebugPrint "subHistogram segHist" Nothing let (is, ns) = unzip $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns let pes = patElems pat num_red_res = length histops + sum (map (length . histNeutral) histops) @@ -247,8 +247,8 @@ ) -> do histop' <- renameHistop histop - let bucket' = map toInt64Exp bucket - dest_shape' = map toInt64Exp $ shapeDims dest_shape + let bucket' = map pe64 bucket + dest_shape' = map pe64 $ shapeDims dest_shape acc_params' = (lambdaParams . histOp) histop' vs_params' = takeLast (length vs') $ lambdaParams $ histOp histop' @@ -310,7 +310,7 @@ map fst segment_dims ++ [subhistogram_id] ++ bucket_ids ) - let ns_red = map (toInt64Exp . snd) $ unSegSpace segred_space + let ns_red = map (pe64 . snd) $ unSegSpace segred_space iterations = product $ init ns_red -- The segmented reduction is sequential over the inner most dimension scheduler_info = Imp.SchedulerInfo (untyped iterations) Imp.Static red_task = Imp.ParallelTask red_code @@ -346,7 +346,7 @@ MulticoreGen Imp.MCCode compileSegHistBody pat space histops kbody = collect $ do let (is, ns) = unzip $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns let num_red_res = length histops + sum (map (length . histNeutral) histops) map_pes = drop num_red_res $ patElems pat @@ -369,8 +369,8 @@ forM_ (zip3 per_red_pes histops (splitHistResults histops red_res)) $ \(red_pes, HistOp dest_shape _ _ _ shape lam, (bucket, vs')) -> do let (is_params, vs_params) = splitAt (length vs') $ lambdaParams lam - bucket' = map toInt64Exp bucket - dest_shape' = map toInt64Exp $ shapeDims dest_shape + bucket' = map pe64 bucket + dest_shape' = map pe64 $ shapeDims dest_shape bucket_in_bounds = inBounds (Slice (map DimFix bucket')) dest_shape' sComment "save map-out results" $ diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegMap.hs @@ -20,9 +20,9 @@ copyDWIMFix (patElemName pe) (map Imp.le64 is) se [] writeResult _ pe (WriteReturns _ (Shape rws) _ idx_vals) = do let (iss, vs) = unzip idx_vals - rws' = map toInt64Exp rws + rws' = map pe64 rws forM_ (zip iss vs) $ \(slice, v) -> do - let slice' = fmap toInt64Exp slice + let slice' = fmap pe64 slice sWhen (inBounds slice' rws') $ copyDWIM (patElemName pe) (unSlice slice') v [] writeResult _ _ res = @@ -35,7 +35,7 @@ MulticoreGen Imp.MCCode compileSegMapBody pat space (KernelBody _ kstms kres) = collect $ do let (is, ns) = unzip $ unSegSpace space - ns' = map toInt64Exp ns + ns' = map pe64 ns dPrim_ (segFlat space) int64 sOp $ Imp.GetTaskId (segFlat space) kstms' <- mapM renameStm kstms diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegRed.hs @@ -210,7 +210,7 @@ MulticoreGen () genReductionLoop typ kbodymap slugs slug_local_accs space i = do let (is, ns) = unzip $ unSegSpace space - ns' = map toInt64Exp ns + ns' = map pe64 ns zipWithM_ dPrimV_ is $ unflattenIndex ns' i kbodymap $ \all_red_res' -> do forM_ (zip3 all_red_res' slugs slug_local_accs) $ \(red_res, slug, local_accs) -> @@ -389,7 +389,7 @@ MulticoreGen Imp.MCCode compileSegRedBody pat space reds kbody = do let (is, ns) = unzip $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns inner_bound = last ns_64 dPrim_ (segFlat space) int64 sOp $ Imp.GetTaskId (segFlat space) diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/SegScan.hs @@ -142,7 +142,7 @@ per_scan_res = segBinOpChunks scan_ops all_scan_res per_scan_pes = segBinOpChunks scan_ops $ patElems pat let (is, ns) = unzip $ unSegSpace space - ns' = map toInt64Exp ns + ns' = map pe64 ns zipWithM_ dPrimV_ is $ unflattenIndex ns' i compileStms mempty (kernelBodyStms kbody) $ do @@ -242,7 +242,7 @@ scanStage2 pat nsubtasks space scan_ops kbody = do emit $ Imp.DebugPrint "nonsegmentedScan stage 2" Nothing let (is, ns) = unzip $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns per_scan_pes = segBinOpChunks scan_ops $ patElems pat nsubtasks' = tvExp nsubtasks @@ -401,7 +401,7 @@ MulticoreGen Imp.MCCode compileSegScanBody pat space scan_ops kbody = collect $ do let (is, ns) = unzip $ unSegSpace space - ns_64 = map toInt64Exp ns + ns_64 = map pe64 ns dPrim_ (segFlat space) int64 sOp $ Imp.GetTaskId (segFlat space) diff --git a/src/Futhark/CodeGen/RTS/C.hs b/src/Futhark/CodeGen/RTS/C.hs --- a/src/Futhark/CodeGen/RTS/C.hs +++ b/src/Futhark/CodeGen/RTS/C.hs @@ -3,6 +3,8 @@ -- | Code snippets used by the C backends. module Futhark.CodeGen.RTS.C ( atomicsH, + contextH, + contextPrototypesH, cudaH, freeListH, halfH, @@ -118,3 +120,13 @@ cacheH :: T.Text cacheH = $(embedStringFile "rts/c/cache.h") {-# NOINLINE cacheH #-} + +-- | @rts/c/context.h@ +contextH :: T.Text +contextH = $(embedStringFile "rts/c/context.h") +{-# NOINLINE contextH #-} + +-- | @rts/c/context_prototypes.h@ +contextPrototypesH :: T.Text +contextPrototypesH = $(embedStringFile "rts/c/context_prototypes.h") +{-# NOINLINE contextPrototypesH #-} diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -72,9 +72,12 @@ -- * Monadic expression builders eSubExp, eParam, + eMatch', + eMatch, eIf, eIf', eBinOp, + eUnOp, eCmpOp, eConvOp, eSignum, @@ -85,6 +88,7 @@ eSliceArray, eBlank, eAll, + eAny, eDimInBounds, eOutOfBounds, @@ -106,7 +110,6 @@ fullSliceNum, isFullSlice, sliceAt, - ifCommon, -- * Result types instantiateShapes, @@ -122,7 +125,7 @@ import Control.Monad.Identity import Control.Monad.State -import Data.List (sortOn) +import Data.List (foldl', sortOn, transpose) import qualified Data.Map.Strict as M import Futhark.Builder import Futhark.IR @@ -209,43 +212,76 @@ m (Exp (Rep m)) eParam = eSubExp . Var . paramName --- | Construct an 'If' expression from a monadic condition and monadic --- branches. 'eBody' might be convenient for constructing the --- branches. +removeRedundantScrutinees :: [SubExp] -> [Case b] -> ([SubExp], [Case b]) +removeRedundantScrutinees ses cases = + let (ses', vs) = + unzip $ filter interesting $ zip ses $ transpose (map casePat cases) + in (ses', zipWith Case (transpose vs) $ map caseBody cases) + where + interesting = any (/= Nothing) . snd + +-- | As 'eMatch', but an 'MatchSort' can be given. +eMatch' :: + (MonadBuilder m, BranchType (Rep m) ~ ExtType) => + [SubExp] -> + [Case (m (Body (Rep m)))] -> + m (Body (Rep m)) -> + MatchSort -> + m (Exp (Rep m)) +eMatch' ses cases_m defbody_m sort = do + cases <- mapM (traverse insertStmsM) cases_m + defbody <- insertStmsM defbody_m + ts <- + foldl' generaliseExtTypes + <$> bodyExtType defbody + <*> mapM (bodyExtType . caseBody) cases + cases' <- mapM (traverse $ addContextForBranch ts) cases + defbody' <- addContextForBranch ts defbody + let ts' = replicate (length (shapeContext ts)) (Prim int64) ++ ts + (ses', cases'') = removeRedundantScrutinees ses cases' + pure $ Match ses' cases'' defbody' $ MatchDec ts' sort + where + addContextForBranch ts (Body _ stms val_res) = do + body_ts <- extendedScope (traverse subExpResType val_res) stmsscope + let ctx_res = + map snd $ sortOn fst $ M.toList $ shapeExtMapping ts body_ts + mkBodyM stms $ subExpsRes ctx_res ++ val_res + where + stmsscope = scopeOf stms + +-- | Construct a 'Match' expression. The main convenience here is +-- that the existential context of the return type is automatically +-- deduced, and the necessary elements added to the branches. +eMatch :: + (MonadBuilder m, BranchType (Rep m) ~ ExtType) => + [SubExp] -> + [Case (m (Body (Rep m)))] -> + m (Body (Rep m)) -> + m (Exp (Rep m)) +eMatch ses cases_m defbody_m = eMatch' ses cases_m defbody_m MatchNormal + +-- | Construct a 'Match' modelling an if-expression from a monadic +-- condition and monadic branches. 'eBody' might be convenient for +-- constructing the branches. eIf :: (MonadBuilder m, BranchType (Rep m) ~ ExtType) => m (Exp (Rep m)) -> m (Body (Rep m)) -> m (Body (Rep m)) -> m (Exp (Rep m)) -eIf ce te fe = eIf' ce te fe IfNormal +eIf ce te fe = eIf' ce te fe MatchNormal --- | As 'eIf', but an 'IfSort' can be given. +-- | As 'eIf', but an 'MatchSort' can be given. eIf' :: (MonadBuilder m, BranchType (Rep m) ~ ExtType) => m (Exp (Rep m)) -> m (Body (Rep m)) -> m (Body (Rep m)) -> - IfSort -> + MatchSort -> m (Exp (Rep m)) eIf' ce te fe if_sort = do ce' <- letSubExp "cond" =<< ce - te' <- insertStmsM te - fe' <- insertStmsM fe - -- We need to construct the context. - ts <- generaliseExtTypes <$> bodyExtType te' <*> bodyExtType fe' - te'' <- addContextForBranch ts te' - fe'' <- addContextForBranch ts fe' - let ts' = replicate (length (shapeContext ts)) (Prim int64) ++ ts - pure $ If ce' te'' fe'' $ IfDec ts' if_sort - where - addContextForBranch ts (Body _ stms val_res) = do - body_ts <- extendedScope (traverse subExpResType val_res) stmsscope - let ctx_res = - map snd $ sortOn fst $ M.toList $ shapeExtMapping ts body_ts - mkBodyM stms $ subExpsRes ctx_res ++ val_res - where - stmsscope = scopeOf stms + eMatch' [ce'] [Case [Just $ BoolValue True] te] fe if_sort -- The type of a body. Watch out: this only works for the degenerate -- case where the body does not already return its context. @@ -268,6 +304,14 @@ y' <- letSubExp "y" =<< y pure $ BasicOp $ BinOp op x' y' +-- | Construct a v'UnOp' expression with the given operator. +eUnOp :: + MonadBuilder m => + UnOp -> + m (Exp (Rep m)) -> + m (Exp (Rep m)) +eUnOp op x = BasicOp . UnOp op <$> (letSubExp "x" =<< x) + -- | Construct a v'CmpOp' expression with the given comparison. eCmpOp :: MonadBuilder m => @@ -456,8 +500,15 @@ -- | True if all operands are true. eAll :: MonadBuilder m => [SubExp] -> m (Exp (Rep m)) eAll [] = pure $ BasicOp $ SubExp $ constant True +eAll [x] = eSubExp x eAll (x : xs) = foldBinOp LogAnd x xs +-- | True if any operand is true. +eAny :: MonadBuilder m => [SubExp] -> m (Exp (Rep m)) +eAny [] = pure $ BasicOp $ SubExp $ constant False +eAny [x] = eSubExp x +eAny (x : xs) = foldBinOp LogOr x xs + -- | Create a two-parameter lambda whose body applies the given binary -- operation to its arguments. It is assumed that both argument and -- result types are the same. (This assumption should be fixed at @@ -547,10 +598,6 @@ allOfIt (Constant v) DimFix {} = oneIsh v allOfIt d (DimSlice _ n _) = d == n allOfIt _ _ = False - --- | Produce the common case of an 'IfDec'. -ifCommon :: [Type] -> IfDec ExtType -ifCommon ts = IfDec (staticShapes ts) IfNormal -- | Conveniently construct a body that contains no bindings. resultBody :: Buildable rep => [SubExp] -> Body rep diff --git a/src/Futhark/IR/Aliases.hs b/src/Futhark/IR/Aliases.hs --- a/src/Futhark/IR/Aliases.hs +++ b/src/Futhark/IR/Aliases.hs @@ -393,4 +393,9 @@ let Body bodyrep _ _ = mkBody (fmap removeStmAliases stms) res in mkAliasedBody bodyrep stms res -instance (ASTRep (Aliases rep), Buildable (Aliases rep)) => BuilderOps (Aliases rep) +instance + ( ASTRep rep, + CanBeAliased (Op rep), + Buildable (Aliases rep) + ) => + BuilderOps (Aliases rep) diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs --- a/src/Futhark/IR/Mem.hs +++ b/src/Futhark/IR/Mem.hs @@ -84,6 +84,7 @@ lookupMemInfo, subExpMemInfo, lookupArraySummary, + lookupMemSpace, existentialiseIxFun, -- * Type checking parts @@ -778,12 +779,13 @@ ( length val_ts == length rt && and (zipWith (matches ctx_map_ids ctx_map_exts) val_ts rt) ) - $ TC.bad - $ TC.TypeError - $ "Expression type:\n " - ++ prettyTuple rt - ++ "\ncannot match pattern type:\n " - ++ prettyTuple val_ts + . TC.bad + . TC.TypeError + . pretty + $ "Expression type:" + indent 2 (ppTuple' rt) + "cannot match pattern type:" + indent 2 (ppTuple' val_ts) where matches _ _ (MemPrim x) (MemPrim y) = x == y matches _ _ (MemMem x_space) (MemMem y_space) = @@ -864,6 +866,22 @@ ++ " to be array but bound to:\n" ++ pretty summary +lookupMemSpace :: + (Mem rep inner, HasScope rep m, Monad m) => + VName -> + m Space +lookupMemSpace name = do + summary <- lookupMemInfo name + case summary of + MemMem space -> + pure space + _ -> + error $ + "Expected " + ++ pretty name + ++ " to be memory but bound to:\n" + ++ pretty summary + checkMemInfo :: TC.Checkable rep => VName -> @@ -996,26 +1014,25 @@ -- | The return information of an expression. This can be seen as the -- "return type with memory annotations" of the expression. expReturns :: - ( Monad m, - LocalScope rep m, - Mem rep inner - ) => + (LocalScope rep m, Mem rep inner) => Exp rep -> m [ExpReturns] expReturns (BasicOp (SubExp se)) = pure <$> subExpReturns se expReturns (BasicOp (Opaque _ (Var v))) = pure <$> varReturns v -expReturns (BasicOp (Reshape newshape v)) = do +expReturns (BasicOp (Reshape k newshape v)) = do (et, _, mem, ixfun) <- arrayVarReturns v pure - [ MemArray et (Shape $ map (Free . newDim) newshape) NoUniqueness $ - Just $ - ReturnsInBlock mem $ - existentialiseIxFun [] $ - IxFun.reshape ixfun $ - map (fmap pe64) newshape + [ MemArray et (fmap Free newshape) NoUniqueness $ + Just . ReturnsInBlock mem . existentialiseIxFun [] $ + reshaper ixfun $ + map pe64 (shapeDims newshape) ] + where + reshaper = case k of + ReshapeArbitrary -> IxFun.reshape + ReshapeCoerce -> IxFun.coerce expReturns (BasicOp (Rearrange perm v)) = do (et, Shape dims, mem, ixfun) <- arrayVarReturns v let ixfun' = IxFun.permute ixfun perm @@ -1026,16 +1043,6 @@ ReturnsInBlock mem $ existentialiseIxFun [] ixfun' ] -expReturns (BasicOp (Rotate offsets v)) = do - (et, Shape dims, mem, ixfun) <- arrayVarReturns v - let offsets' = map pe64 offsets - ixfun' = IxFun.rotate ixfun offsets' - pure - [ MemArray et (Shape $ map Free dims) NoUniqueness $ - Just $ - ReturnsInBlock mem $ - existentialiseIxFun [] ixfun' - ] expReturns (BasicOp (Index v slice)) = do pure . varInfoToExpReturns <$> sliceInfo v slice expReturns (BasicOp (Update _ v _ _)) = @@ -1074,7 +1081,7 @@ mergevars = map fst merge expReturns (Apply _ _ ret _) = pure $ map funReturnsToExpReturns ret -expReturns (If _ _ _ (IfDec ret _)) = +expReturns (Match _ _ _ (MatchDec ret _)) = pure $ map bodyReturnsToExpReturns ret expReturns (Op op) = opReturns op diff --git a/src/Futhark/IR/Mem/IxFun.hs b/src/Futhark/IR/Mem/IxFun.hs --- a/src/Futhark/IR/Mem/IxFun.hs +++ b/src/Futhark/IR/Mem/IxFun.hs @@ -6,15 +6,17 @@ -- linear-memory accessor descriptors; see Zhu, Hoeflinger and David work. module Futhark.IR.Mem.IxFun ( IxFun (..), + Shape, LMAD (..), LMADDim (..), Monotonicity (..), index, + mkExistential, iota, iotaOffset, permute, - rotate, reshape, + coerce, slice, flatSlice, rebase, @@ -25,7 +27,6 @@ isDirect, isLinear, substituteInIxFun, - leastGeneralGeneralization, existentialize, closeEnough, equivalent, @@ -38,21 +39,18 @@ import Control.Monad.State import Control.Monad.Writer import Data.Function (on, (&)) -import Data.List (sort, sortBy, zip4, zip5, zipWith5) +import Data.List (sort, sortBy, zip4, zipWith4) import Data.List.NonEmpty (NonEmpty (..)) import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as M import Data.Maybe (isJust) import Futhark.Analysis.PrimExp import Futhark.Analysis.PrimExp.Convert (substituteInPrimExp) -import qualified Futhark.Analysis.PrimExp.Generalize as PEG import Futhark.IR.Prop import Futhark.IR.Syntax - ( DimChange (..), - DimIndex (..), + ( DimIndex (..), FlatDimIndex (..), FlatSlice (..), - ShapeChange, Slice (..), dimFix, flatSliceDims, @@ -66,6 +64,7 @@ import Futhark.Util.Pretty import Prelude hiding (id, mod, (.)) +-- | The shape of an index function. type Shape num = [num] type Indices num = [num] @@ -86,7 +85,6 @@ -- | A single dimension in an 'LMAD'. data LMADDim num = LMADDim { ldStride :: num, - ldRotate :: num, ldShape :: num, ldPerm :: Int, ldMon :: Monotonicity @@ -94,7 +92,7 @@ deriving (Show, Eq) -- | LMAD's representation consists of a general offset and for each dimension a --- stride, rotate factor, number of elements (or shape), permutation, and +-- stride, number of elements (or shape), permutation, and -- monotonicity. Note that the permutation is not strictly necessary in that the -- permutation can be performed directly on LMAD dimensions, but then it is -- difficult to extract the permutation back from an LMAD. @@ -137,7 +135,7 @@ { ixfunLMADs :: NonEmpty (LMAD num), base :: Shape num, -- | ignoring permutations, is the index function contiguous? - ixfunContig :: Bool + contiguous :: Bool } deriving (Show, Eq) @@ -150,7 +148,6 @@ semisep [ "offset: " <> oneLine (ppr offset), "strides: " <> p ldStride, - "rotates: " <> p ldRotate, "shape: " <> p ldShape, "permutation: " <> p ldPerm, "monotonicity: " <> p ldMon @@ -201,9 +198,10 @@ traverse f (LMAD offset dims) = LMAD <$> f offset <*> traverse f' dims where - f' (LMADDim s r n p m) = - LMADDim <$> f s <*> f r <*> f n <*> pure p <*> pure m + f' (LMADDim s n p m) = LMADDim <$> f s <*> f n <*> pure p <*> pure m +-- It is important that the traversal order here is the same as in +-- mkExistential. instance Traversable IxFun where traverse f (IxFun lmads oshp cg) = IxFun <$> traverse (traverse f) lmads <*> traverse f oshp <*> pure cg @@ -241,10 +239,9 @@ let offset' = substituteInPrimExp tab offset dims' = map - ( \(LMADDim s r n p m) -> + ( \(LMADDim s n p m) -> LMADDim (substituteInPrimExp tab s) - (substituteInPrimExp tab r) (substituteInPrimExp tab n) p m @@ -274,9 +271,7 @@ && length oshp == length dims && offset == 0 && all - ( \(LMADDim s r n p _, m, d, se) -> - s == se && r == 0 && n == d && p == m - ) + (\(LMADDim s n p _, m, d, se) -> s == se && n == d && p == m) (zip4 dims [0 .. length dims - 1] oshp strides_expected) isDirect _ = False @@ -330,20 +325,31 @@ sum $ zipWith flatOneDim - (map (\(LMADDim s r n _ _) -> (s, r, n)) dims) + (map ldStride dims) (permuteInv (lmadPermutation lmad) inds) in off + prod -- | iota with offset. iotaOffset :: IntegralExp num => num -> Shape num -> IxFun num -iotaOffset o ns = - let rs = replicate (length ns) 0 - in IxFun (makeRotIota Inc o (zip rs ns) :| []) ns True +iotaOffset o ns = IxFun (makeRotIota Inc o ns :| []) ns True -- | iota. iota :: IntegralExp num => Shape num -> IxFun num iota = iotaOffset 0 +-- | Create a contiguous single-LMAD index function that is +-- existential in everything, with the provided permutation, +-- monotonicity, and contiguousness. +mkExistential :: Int -> [(Int, Monotonicity)] -> Bool -> Int -> IxFun (Ext a) +mkExistential basis_rank perm contig start = + IxFun (NE.singleton lmad) basis contig + where + basis = take basis_rank $ map Ext [start + 1 + dims_rank * 2 ..] + dims_rank = length perm + lmad = LMAD (Ext start) $ zipWith onDim perm [0 ..] + onDim (p, mon) i = + LMADDim (Ext (start + 1 + i * 2)) (Ext (start + 2 + i * 2)) p mon + -- | Permute dimensions. permute :: IntegralExp num => @@ -355,24 +361,6 @@ perm = map (perm_cur !!) perm_new in IxFun (setLMADPermutation perm lmad :| lmads) oshp cg --- | Rotate an index function. -rotate :: - (Eq num, IntegralExp num) => - IxFun num -> - Indices num -> - IxFun num -rotate (IxFun (lmad@(LMAD off dims) :| lmads) oshp cg) offs = - let dims' = - zipWith - ( \(LMADDim s r n p f) o -> - if s == 0 - then LMADDim 0 0 n p Unknown - else LMADDim s (r + o) n p f - ) - dims - (permuteInv (lmadPermutation lmad) offs) - in IxFun (LMAD off dims' :| lmads) oshp cg - -- | Handle the case where a slice can stay within a single LMAD. sliceOneLMAD :: (Eq num, IntegralExp num) => @@ -383,7 +371,6 @@ let perm = lmadPermutation lmad is' = permuteInv perm is cg' = cg && slicePreservesContiguous lmad (Slice is') - guard $ harmlessRotation lmad (Slice is') let lmad' = foldl sliceOne (LMAD (lmadOffset lmad) []) $ zip is' ldims -- need to remove the fixed dims from the permutation perm' = @@ -405,28 +392,6 @@ d = foldl f 0 inds in [p - d | d /= -1] - harmlessRotation' :: - (Eq num, IntegralExp num) => - LMADDim num -> - DimIndex num -> - Bool - harmlessRotation' _ (DimFix _) = True - harmlessRotation' (LMADDim 0 _ _ _ _) _ = True - harmlessRotation' (LMADDim _ 0 _ _ _) _ = True - harmlessRotation' (LMADDim _ _ n _ _) dslc - | dslc == DimSlice (n - 1) n (-1) - || dslc == unitSlice 0 n = - True - harmlessRotation' _ _ = False - - harmlessRotation :: - (Eq num, IntegralExp num) => - LMAD num -> - Slice num -> - Bool - harmlessRotation (LMAD _ dims) (Slice iss) = - and $ zipWith harmlessRotation' dims iss - -- XXX: TODO: what happens to r on a negative-stride slice; is there -- such a case? sliceOne :: @@ -434,26 +399,24 @@ LMAD num -> (DimIndex num, LMADDim num) -> LMAD num - sliceOne (LMAD off dims) (DimFix i, LMADDim s r n _ _) = - LMAD (off + flatOneDim (s, r, n) i) dims - sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ _ p _) = - LMAD off (dims ++ [LMADDim 0 0 ne p Unknown]) - sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ _ n _ _)) + sliceOne (LMAD off dims) (DimFix i, LMADDim s _x _ _) = + LMAD (off + flatOneDim s i) dims + sliceOne (LMAD off dims) (DimSlice _ ne _, LMADDim 0 _ p _) = + LMAD off (dims ++ [LMADDim 0 ne p Unknown]) + sliceOne (LMAD off dims) (dmind, dim@(LMADDim _ n _ _)) | dmind == unitSlice 0 n = LMAD off (dims ++ [dim]) - sliceOne (LMAD off dims) (dmind, LMADDim s r n p m) + sliceOne (LMAD off dims) (dmind, LMADDim s n p m) | dmind == DimSlice (n - 1) n (-1) = - let r' = if r == 0 then 0 else n - r - off' = off + flatOneDim (s, 0, n) (n - 1) - in LMAD off' (dims ++ [LMADDim (s * (-1)) r' n p (invertMonotonicity m)]) - sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s r n p _) = - LMAD (off + flatOneDim (s, r, n) b) (dims ++ [LMADDim 0 0 ne p Unknown]) - sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s 0 _ p m) = + let off' = off + flatOneDim s (n - 1) + in LMAD off' (dims ++ [LMADDim (s * (-1)) n p (invertMonotonicity m)]) + sliceOne (LMAD off dims) (DimSlice b ne 0, LMADDim s _ p _) = + LMAD (off + flatOneDim s b) (dims ++ [LMADDim 0 ne p Unknown]) + sliceOne (LMAD off dims) (DimSlice bs ns ss, LMADDim s _ p m) = let m' = case sgn ss of Just 1 -> m Just (-1) -> invertMonotonicity m _ -> Unknown - in LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) 0 ns p m']) - sliceOne _ _ = error "slice: reached impossible case" + in LMAD (off + s * bs) (dims ++ [LMADDim (ss * s) ns p m']) slicePreservesContiguous :: (Eq num, IntegralExp num) => @@ -473,17 +436,17 @@ map normIndex slc -- Check that: -- 1. a clean split point exists between Fixed and Sliced dims - -- 2. the outermost sliced dim has +/- 1 stride AND is unrotated or full. + -- 2. the outermost sliced dim has +/- 1 stride. -- 3. the rest of inner sliced dims are full. (_, success) = foldl - ( \(found, res) (slcdim, LMADDim _ r n _ _) -> + ( \(found, res) (slcdim, LMADDim _ n _ _) -> case (slcdim, found) of (DimFix {}, True) -> (found, False) (DimFix {}, False) -> (found, res) - (DimSlice _ ne ds, False) -> + (DimSlice _ _ ds, False) -> -- outermost sliced dim: +/-1 stride - let res' = (r == 0 || n == ne) && (ds == 1 || ds == -1) + let res' = (ds == 1 || ds == -1) in (True, res && res') (DimSlice _ ne ds, True) -> -- inner sliced dim: needs to be full @@ -525,54 +488,27 @@ FlatSlice num -> IxFun num flatSlice ixfun@(IxFun (LMAD offset (dim : dims) :| lmads) oshp cg) (FlatSlice new_offset is) - | hasContiguousPerm ixfun, - ldRotate dim == 0 = + | hasContiguousPerm ixfun = let lmad = LMAD (offset + new_offset * ldStride dim) - ( map (helper $ ldStride dim) is - <> dims - ) + (map (helper $ ldStride dim) is <> dims) & setLMADPermutation [0 ..] in IxFun (lmad :| lmads) oshp cg where helper s0 (FlatDimIndex n s) = let new_mon = if s0 * s == 1 then Inc else Unknown - in LMADDim (s0 * s) 0 n 0 new_mon + in LMADDim (s0 * s) n 0 new_mon flatSlice (IxFun (lmad :| lmads) oshp cg) s@(FlatSlice new_offset _) = IxFun (LMAD (new_offset * base_stride) (new_dims <> tail_dims) :| lmad : lmads) oshp cg where tail_shapes = tail $ lmadShape lmad base_stride = product tail_shapes tail_strides = tail $ scanr (*) 1 tail_shapes - tail_dims = zipWith5 LMADDim tail_strides (repeat 0) tail_shapes [length new_shapes ..] (repeat Inc) + tail_dims = zipWith4 LMADDim tail_strides tail_shapes [length new_shapes ..] (repeat Inc) new_shapes = flatSliceDims s new_strides = map (* base_stride) $ flatSliceStrides s - new_dims = zipWith5 LMADDim new_strides (repeat 0) new_shapes [0 ..] (repeat Inc) - --- | Handle the simple case where all reshape dimensions are coercions. -reshapeCoercion :: - (Eq num, IntegralExp num) => - IxFun num -> - ShapeChange num -> - Maybe (IxFun num) -reshapeCoercion (IxFun (lmad@(LMAD off dims) :| lmads) oldbase cg) newshape = do - let perm = lmadPermutation lmad - (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape - let hd_len = length head_coercions - num_coercions = hd_len + length tail_coercions - dims' = permuteFwd perm dims - mid_dims = take (length dims - num_coercions) $ drop hd_len dims' - num_rshps = length reshapes - guard (num_rshps == 0 || (num_rshps == 1 && length mid_dims == 1)) - let dims'' = - permuteInv perm $ - zipWith - (\ld n -> ld {ldShape = n}) - dims' - (newDims newshape) - lmad' = LMAD off dims'' - pure $ IxFun (lmad' :| lmads) oldbase cg + new_dims = zipWith4 LMADDim new_strides new_shapes [0 ..] (repeat Inc) -- | Handle the case where a reshape operation can stay inside a single LMAD. -- @@ -583,9 +519,7 @@ -- the LMAD dimensions that were *not* reshape coercions. -- (2) the repetition of dimensions of the underlying LMAD must -- refer only to the coerced-dimensions of the reshape operation. --- (3) similarly, the rotated dimensions must refer only to --- dimensions that are coerced by the reshape operation. --- (4) finally, the underlying memory is contiguous (and monotonous). +-- (3) finally, the underlying memory is contiguous (and monotonous). -- -- If any of these conditions do not hold, then the reshape operation will -- conservatively add a new LMAD to the list, leading to a representation that @@ -593,79 +527,51 @@ reshapeOneLMAD :: (Eq num, IntegralExp num) => IxFun num -> - ShapeChange num -> + Shape num -> Maybe (IxFun num) reshapeOneLMAD ixfun@(IxFun (lmad@(LMAD off dims) :| lmads) oldbase cg) newshape = do let perm = lmadPermutation lmad - (head_coercions, reshapes, tail_coercions) <- splitCoercions newshape - let hd_len = length head_coercions - num_coercions = hd_len + length tail_coercions dims_perm = permuteFwd perm dims - mid_dims = take (length dims - num_coercions) $ drop hd_len dims_perm - -- Ignore rotates, as we only care about not having rotates in the - -- dimensions that aren't coercions (@mid_dims@), which we check - -- separately. - mon = ixfunMonotonicityRots True ixfun + mid_dims = take (length dims) dims_perm + mon = ixfunMonotonicity ixfun guard $ - -- checking conditions (2) and (3) - all (\(LMADDim s r _ _ _) -> s /= 0 && r == 0) mid_dims + -- checking conditions (2) + all (\(LMADDim s _ _ _) -> s /= 0) mid_dims && -- checking condition (1) - consecutive hd_len (map ldPerm mid_dims) + consecutive 0 (map ldPerm mid_dims) && - -- checking condition (4) + -- checking condition (3) hasContiguousPerm ixfun && cg && (mon == Inc || mon == Dec) -- make new permutation - let rsh_len = length reshapes + let rsh_len = length newshape diff = length newshape - length dims iota_shape = [0 .. length newshape - 1] perm' = map ( \i -> - let ind = - if i < hd_len - then i - else i - diff - in if (i >= hd_len) && (i < hd_len + rsh_len) + let ind = i - diff + in if (i >= 0) && (i < rsh_len) then i -- already checked mid_dims not affected - else - let p = ldPerm (dims !! ind) - in if p < hd_len - then p - else p + diff + else ldPerm (dims !! ind) + diff ) iota_shape -- split the dimensions (support_inds, repeat_inds) = foldl - ( \(sup, rpt) (i, shpdim, ip) -> - case (i < hd_len, i >= hd_len + rsh_len, shpdim) of - (True, _, DimCoercion n) -> - case dims_perm !! i of - (LMADDim 0 _ _ _ _) -> (sup, (ip, n) : rpt) - (LMADDim _ r _ _ _) -> ((ip, (r, n)) : sup, rpt) - (_, True, DimCoercion n) -> - case dims_perm !! (i - diff) of - (LMADDim 0 _ _ _ _) -> (sup, (ip, n) : rpt) - (LMADDim _ r _ _ _) -> ((ip, (r, n)) : sup, rpt) - (False, False, _) -> - ((ip, (0, newDim shpdim)) : sup, rpt) - -- already checked that the reshaped - -- dims cannot be rotates - _ -> error "reshape: reached impossible case" - ) + (\(sup, rpt) (shpdim, ip) -> ((ip, shpdim) : sup, rpt)) ([], []) $ reverse - $ zip3 iota_shape newshape perm' + $ zip newshape perm' (sup_inds, support) = unzip $ sortBy (compare `on` fst) support_inds (rpt_inds, repeats) = unzip repeat_inds LMAD off' dims_sup = makeRotIota mon off support - repeats' = map (\n -> LMADDim 0 0 n 0 Unknown) repeats + repeats' = map (\n -> LMADDim 0 n 0 Unknown) repeats dims' = map snd $ sortBy (compare `on` fst) $ @@ -677,33 +583,32 @@ consecutive i [p] = i == p consecutive i ps = and $ zipWith (==) ps [i, i + 1 ..] -splitCoercions :: - (Eq num, IntegralExp num) => - ShapeChange num -> - Maybe (ShapeChange num, ShapeChange num, ShapeChange num) -splitCoercions newshape' = do - let (head_coercions, newshape'') = span isCoercion newshape' - (reshapes, tail_coercions) = break isCoercion newshape'' - guard (all isCoercion tail_coercions) - pure (head_coercions, reshapes, tail_coercions) - where - isCoercion DimCoercion {} = True - isCoercion _ = False - -- | Reshape an index function. reshape :: (Eq num, IntegralExp num) => IxFun num -> - ShapeChange num -> + Shape num -> IxFun num reshape ixfun new_shape - | Just ixfun' <- reshapeCoercion ixfun new_shape = ixfun' | Just ixfun' <- reshapeOneLMAD ixfun new_shape = ixfun' reshape (IxFun (lmad0 :| lmad0s) oshp cg) new_shape = - case iota (newDims new_shape) of + case iota new_shape of IxFun (lmad :| []) _ _ -> IxFun (lmad :| lmad0 : lmad0s) oshp cg _ -> error "reshape: reached impossible case" +-- | Coerce an index function to look like it has a new shape. +-- Dynamically the shape must be the same. +coerce :: + (Eq num, IntegralExp num) => + IxFun num -> + Shape num -> + IxFun num +coerce (IxFun (lmad :| lmads) oshp cg) new_shape = + IxFun (onLMAD lmad :| lmads) oshp cg + where + onLMAD (LMAD offset dims) = LMAD offset $ zipWith onDim dims new_shape + onDim ld d = ld {ldShape = d} + -- | The number of dimensions in the domain of the input function. rank :: IntegralExp num => @@ -788,16 +693,11 @@ (dims_base', offs_contrib) = unzip $ zipWith - ( \(LMADDim s1 r1 n1 p1 _) (LMADDim _ r2 _ _ m2) -> + ( \(LMADDim s1 n1 p1 _) (LMADDim _ _ _ m2) -> let (s', off') | m2 == Inc = (s1, 0) | otherwise = (s1 * (-1), s1 * (n1 - 1)) - r' - | m2 == Inc = if r2 == 0 then r1 else r1 + r2 - | r1 == 0 = r2 - | r2 == 0 = n1 - r1 - | otherwise = n1 - r1 + r2 - in (LMADDim s' r' n1 (p1 - n_fewer_dims) Inc, off') + in (LMADDim s' n1 (p1 - n_fewer_dims) Inc, off') ) -- If @dims@ is morally a slice, it might have fewer dimensions than -- @dims_base@. Drop extraneous outer dimensions. @@ -835,13 +735,10 @@ if base ixfun == shape new_base then (lmads_base, shp_base) else - let IxFun lmads' shp_base'' _ = reshape new_base $ map DimCoercion shp + let IxFun lmads' shp_base'' _ = reshape new_base shp in (lmads', shp_base'') in IxFun (lmads @++@ lmads_base') shp_base' (cg && cg_base) -ixfunMonotonicity :: (Eq num, IntegralExp num) => IxFun num -> Monotonicity -ixfunMonotonicity = ixfunMonotonicityRots False - -- | If the memory support of the index function is contiguous and row-major -- (i.e., no transpositions, repetitions, rotates, etc.), then this should -- return the offset from which the memory-support of this index function @@ -888,13 +785,12 @@ flatOneDim :: (Eq num, IntegralExp num) => - (num, num, num) -> num -> + num -> num -flatOneDim (s, r, n) i +flatOneDim s i | s == 0 = 0 - | r == 0 = i * s - | otherwise = ((i + r) `mod` n) * s + | otherwise = i * s -- | Generalised iota with user-specified offset and rotates. makeRotIota :: @@ -902,13 +798,12 @@ Monotonicity -> -- | Offset num -> - -- | Pairs of shape and rotation - [(num, num)] -> + -- | Shape + [num] -> LMAD num -makeRotIota mon off support +makeRotIota mon off ns | mon == Inc || mon == Dec = - let rk = length support - (rs, ns) = unzip support + let rk = length ns ss0 = reverse $ take rk $ scanl (*) 1 $ reverse ns ss = if mon == Inc @@ -916,16 +811,15 @@ else map (* (-1)) ss0 ps = map fromIntegral [0 .. rk - 1] fi = replicate rk mon - in LMAD off $ zipWith5 LMADDim ss rs ns ps fi + in LMAD off $ zipWith4 LMADDim ss ns ps fi | otherwise = error "makeRotIota: requires Inc or Dec" -- | Check monotonicity of an index function. -ixfunMonotonicityRots :: +ixfunMonotonicity :: (Eq num, IntegralExp num) => - Bool -> IxFun num -> Monotonicity -ixfunMonotonicityRots ignore_rots (IxFun (lmad :| lmads) _ _) = +ixfunMonotonicity (IxFun (lmad :| lmads) _ _) = let mon0 = lmadMonotonicityRots lmad in if all ((== mon0) . lmadMonotonicityRots) lmads then mon0 @@ -945,93 +839,19 @@ Monotonicity -> LMADDim num -> Bool - isMonDim mon (LMADDim s r _ _ ldmon) = - s == 0 || ((ignore_rots || r == 0) && mon == ldmon) - --- | Generalization (anti-unification) --- --- Anti-unification of two index functions is supported under the following conditions: --- 0. Both index functions are represented by ONE lmad (assumed common case!) --- 1. The support array of the two indexfuns have the same dimensionality --- (we can relax this condition if we use a 1D support, as we probably should!) --- 2. The contiguous property and the per-dimension monotonicity are the same --- (otherwise we might loose important information; this can be relaxed!) --- 3. Most importantly, both index functions correspond to the same permutation --- (since the permutation is represented by INTs, this restriction cannot --- be relaxed, unless we move to a gated-LMAD representation!) -leastGeneralGeneralization :: - Eq v => - IxFun (PrimExp v) -> - IxFun (PrimExp v) -> - Maybe (IxFun (PrimExp (Ext v)), [(PrimExp v, PrimExp v)]) -leastGeneralGeneralization (IxFun (lmad1 :| []) oshp1 ctg1) (IxFun (lmad2 :| []) oshp2 ctg2) = do - guard $ - length oshp1 == length oshp2 - && ctg1 == ctg2 - && map ldPerm (lmadDims lmad1) == map ldPerm (lmadDims lmad2) - && lmadDMon lmad1 == lmadDMon lmad2 - let (ctg, dperm, dmon) = (ctg1, lmadPermutation lmad1, lmadDMon lmad1) - (dshp, m1) <- generalize [] (lmadDShp lmad1) (lmadDShp lmad2) - (oshp, m2) <- generalize m1 oshp1 oshp2 - (dstd, m3) <- generalize m2 (lmadDSrd lmad1) (lmadDSrd lmad2) - (drot, m4) <- generalize m3 (lmadDRot lmad1) (lmadDRot lmad2) - let (offt, m5) = PEG.leastGeneralGeneralization m4 (lmadOffset lmad1) (lmadOffset lmad2) - let lmad_dims = - map (\(a, b, c, d, e) -> LMADDim a b c d e) $ - zip5 dstd drot dshp dperm dmon - lmad = LMAD offt lmad_dims - pure (IxFun (lmad :| []) oshp ctg, m5) - where - lmadDMon = map ldMon . lmadDims - lmadDSrd = map ldStride . lmadDims - lmadDShp = map ldShape . lmadDims - lmadDRot = map ldRotate . lmadDims - generalize m l1 l2 = - foldM - ( \(l_acc, m') (pe1, pe2) -> do - let (e, m'') = PEG.leastGeneralGeneralization m' pe1 pe2 - pure (l_acc ++ [e], m'') - ) - ([], m) - (zip l1 l2) -leastGeneralGeneralization _ _ = Nothing - -isSequential :: [Int] -> Bool -isSequential xs = - all (uncurry (==)) $ zip xs [0 ..] - -existentializeExp :: TPrimExp t v -> State [TPrimExp t v] (TPrimExp t (Ext v)) -existentializeExp e = do - i <- gets length - modify (++ [e]) - let t = primExpType $ untyped e - pure $ TPrimExp $ LeafExp (Ext i) t + isMonDim mon (LMADDim s _ _ ldmon) = + s == 0 || mon == ldmon --- | Try to turn all the leaves of the index function into 'Ext's. We --- require that there's only one LMAD, that the index function is --- contiguous, and the base shape has only one dimension. +-- | Turn all the leaves of the index function into 'Ext's. existentialize :: - (IntExp t, Eq v, Pretty v) => - IxFun (TPrimExp t v) -> - State [TPrimExp t v] (Maybe (IxFun (TPrimExp t (Ext v)))) -existentialize (IxFun (lmad :| []) oshp True) - | all ((== 0) . ldRotate) (lmadDims lmad), - length (lmadShape lmad) == length oshp, - isSequential (map ldPerm $ lmadDims lmad) = do - oshp' <- mapM existentializeExp oshp - lmadOffset' <- existentializeExp $ lmadOffset lmad - lmadDims' <- mapM existentializeLMADDim $ lmadDims lmad - let lmad' = LMAD lmadOffset' lmadDims' - pure $ Just $ IxFun (lmad' :| []) oshp' True + IxFun (TPrimExp Int64 a) -> + IxFun (TPrimExp Int64 (Ext b)) +existentialize ixfun = evalState (traverse (const mkExt) ixfun) 0 where - existentializeLMADDim :: - LMADDim (TPrimExp t v) -> - State [TPrimExp t v] (LMADDim (TPrimExp t (Ext v))) - existentializeLMADDim (LMADDim str rot shp perm mon) = do - stride' <- existentializeExp str - shape' <- existentializeExp shp - pure $ LMADDim stride' (fmap Free rot) shape' perm mon -existentialize _ = pure Nothing + mkExt = do + i <- get + put $ i + 1 + pure $ TPrimExp $ LeafExp (Ext i) int64 -- | When comparing index functions as part of the type check in KernelsMem, -- we may run into problems caused by the simplifier. As index functions can be @@ -1048,6 +868,8 @@ (length (base ixf1) == length (base ixf2)) && (NE.length (ixfunLMADs ixf1) == NE.length (ixfunLMADs ixf2)) && all closeEnoughLMADs (NE.zip (ixfunLMADs ixf1) (ixfunLMADs ixf2)) + -- This treats ixf1 as the "declared type" that we are matching against. + && (contiguous ixf1 <= contiguous ixf2) where closeEnoughLMADs :: (LMAD num, LMAD num) -> Bool closeEnoughLMADs (lmad1, lmad2) = @@ -1072,8 +894,6 @@ == lmadOffset lmad2 && map ldStride (lmadDims lmad1) == map ldStride (lmadDims lmad2) - && map ldRotate (lmadDims lmad1) - == map ldRotate (lmadDims lmad2) -- | Dynamically determine if two 'LMADDim' are equal. -- @@ -1081,7 +901,6 @@ dynamicEqualsLMADDim :: Eq num => LMADDim (TPrimExp t num) -> LMADDim (TPrimExp t num) -> TPrimExp Bool num dynamicEqualsLMADDim dim1 dim2 = ldStride dim1 .==. ldStride dim2 - .&&. ldRotate dim1 .==. ldRotate dim2 .&&. ldShape dim1 .==. ldShape dim2 .&&. fromBool (ldPerm dim1 == ldPerm dim2) .&&. fromBool (ldMon dim1 == ldMon dim2) diff --git a/src/Futhark/IR/Mem/Simplify.hs b/src/Futhark/IR/Mem/Simplify.hs --- a/src/Futhark/IR/Mem/Simplify.hs +++ b/src/Futhark/IR/Mem/Simplify.hs @@ -101,6 +101,7 @@ BodyDec rep ~ (), CanBeWise (Op rep), BuilderOps (Wise rep), + OpReturns (OpWithWisdom inner), Mem rep inner ) @@ -109,7 +110,7 @@ standardRules <> ruleBook [ RuleBasicOp copyCopyToCopy, - RuleIf unExistentialiseMemory, + RuleMatch unExistentialiseMemory, RuleOp decertifySafeAlloc ] [] @@ -118,8 +119,8 @@ -- the array is not existential, and the index function of the array -- does not refer to any names in the pattern, then we can create a -- block of the proper size and always return there. -unExistentialiseMemory :: SimplifyMemory rep inner => TopDownRuleIf (Wise rep) -unExistentialiseMemory vtable pat _ (cond, tbranch, fbranch, ifdec) +unExistentialiseMemory :: SimplifyMemory rep inner => TopDownRuleMatch (Wise rep) +unExistentialiseMemory vtable pat _ (cond, cases, defbody, ifdec) | ST.simplifyMemory vtable, fixable <- foldl hasConcretisableMemory mempty $ patElems pat, not $ null fixable = Simplify $ do @@ -149,9 +150,9 @@ pure $ SubExpRes cs $ Var mem updateResult _ se = pure se - tbranch' <- updateBody tbranch - fbranch' <- updateBody fbranch - letBind pat $ If cond tbranch' fbranch' ifdec + cases' <- mapM (traverse updateBody) cases + defbody' <- updateBody defbody + letBind pat $ Match cond cases' defbody' ifdec where onlyUsedIn name here = not . any ((name `nameIn`) . freeIn) . filter ((/= here) . patElemName) $ @@ -167,13 +168,13 @@ <$> find ((mem ==) . patElemName . snd) (zip [(0 :: Int) ..] $ patElems pat), - Just tse <- maybeNth j $ bodyResult tbranch, - Just fse <- maybeNth j $ bodyResult fbranch, + Just cases_ses <- mapM (maybeNth j . bodyResult . caseBody) cases, + Just defbody_se <- maybeNth j $ bodyResult defbody, mem `onlyUsedIn` patElemName pat_elem, length (IxFun.base ixfun) == shapeRank shape, -- See #1325 all knownSize (shapeDims shape), not $ freeIn ixfun `namesIntersect` namesFromList (patNames pat), - fse /= tse = + any (defbody_se /=) cases_ses = let mem_size = untyped $ product $ primByteSize pt : map sExt64 (IxFun.base ixfun) in (pat_elem, mem_size, mem, space) : fixable diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -16,7 +16,7 @@ import Data.Char (isAlpha) import Data.Functor -import Data.List (zipWith5) +import Data.List (zipWith4) import Data.List.NonEmpty (NonEmpty (..)) import qualified Data.List.NonEmpty as NE import qualified Data.Set as S @@ -238,15 +238,6 @@ pErrorLoc :: Parser (SrcLoc, [SrcLoc]) pErrorLoc = (,mempty) <$> pSrcLoc -pShapeChange :: Parser (ShapeChange SubExp) -pShapeChange = parens $ pDimChange `sepBy` pComma - where - pDimChange = - choice - [ "~" $> DimCoercion <*> pSubExp, - DimNew <$> pSubExp - ] - pIota :: Parser BasicOp pIota = choice $ map p allIntTypes @@ -286,7 +277,9 @@ keyword "replicate" *> parens (Replicate <$> pShape <* pComma <*> pSubExp), keyword "reshape" - *> parens (Reshape <$> pShapeChange <* pComma <*> pVName), + *> parens (Reshape ReshapeArbitrary <$> pShape <* pComma <*> pVName), + keyword "coerce" + *> parens (Reshape ReshapeCoerce <$> pShape <* pComma <*> pVName), keyword "scratch" *> parens (Scratch <$> pPrimType <*> many (pComma *> pSubExp)), keyword "rearrange" @@ -415,30 +408,55 @@ pResult :: Parser Result pResult = braces $ pSubExpRes `sepBy` pComma +pMatchSort :: Parser MatchSort +pMatchSort = + choice + [ lexeme "" $> MatchFallback, + lexeme "" $> MatchEquiv, + pure MatchNormal + ] + +pBranchBody :: PR rep -> Parser (Body rep) +pBranchBody pr = + choice + [ try $ Body (pBodyDec pr) mempty <$> pResult, + braces (pBody pr) + ] + pIf :: PR rep -> Parser (Exp rep) pIf pr = keyword "if" $> f - <*> pSort + <*> pMatchSort <*> pSubExp - <*> (keyword "then" *> pBranchBody) - <*> (keyword "else" *> pBranchBody) + <*> (keyword "then" *> pBranchBody pr) + <*> (keyword "else" *> pBranchBody pr) <*> (lexeme ":" *> pBranchTypes pr) where - pSort = - choice - [ lexeme "" $> IfFallback, - lexeme "" $> IfEquiv, - pure IfNormal - ] f sort cond tbranch fbranch t = - If cond tbranch fbranch $ IfDec t sort - pBranchBody = - choice - [ try $ Body (pBodyDec pr) mempty <$> pResult, - braces (pBody pr) - ] + Match [cond] [Case [Just $ BoolValue True] tbranch] fbranch $ MatchDec t sort +pMatch :: PR rep -> Parser (Exp rep) +pMatch pr = + keyword "match" + $> f + <*> pMatchSort + <*> braces (pSubExp `sepBy` pComma) + <*> many pCase + <*> (keyword "default" *> lexeme "->" *> pBranchBody pr) + <*> (lexeme ":" *> pBranchTypes pr) + where + f sort cond cases defbody t = + Match cond cases defbody $ MatchDec t sort + pCase = + keyword "case" + $> Case + <*> braces (pMaybeValue `sepBy` pComma) + <* lexeme "->" + <*> pBranchBody pr + pMaybeValue = + choice [lexeme "_" $> Nothing, Just <$> pPrimValue] + pApply :: PR rep -> Parser (Exp rep) pApply pr = keyword "apply" *> (p =<< choice [lexeme "" $> Unsafe, pure Safe]) @@ -534,6 +552,7 @@ pExp pr = choice [ pIf pr, + pMatch pr, pApply pr, pLoop pr, pWithAcc pr, @@ -1012,11 +1031,10 @@ pLMAD = braces $ do offset <- pLab "offset" pNum <* pSemi strides <- pLab "strides" $ brackets (pNum `sepBy` pComma) <* pSemi - rotates <- pLab "rotates" $ brackets (pNum `sepBy` pComma) <* pSemi shape <- pLab "shape" $ brackets (pNum `sepBy` pComma) <* pSemi perm <- pLab "permutation" $ brackets (pInt `sepBy` pComma) <* pSemi mon <- pLab "monotonicity" $ brackets (pMon `sepBy` pComma) - pure $ IxFun.LMAD offset $ zipWith5 IxFun.LMADDim strides rotates shape perm mon + pure $ IxFun.LMAD offset $ zipWith4 IxFun.LMADDim strides shape perm mon pPrimExpLeaf :: Parser VName pPrimExpLeaf = pVName diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs --- a/src/Futhark/IR/Pretty.hs +++ b/src/Futhark/IR/Pretty.hs @@ -217,8 +217,10 @@ text "replicate" <> apply [ppr ne, align (ppr ve)] ppr (Scratch t shape) = text "scratch" <> apply (ppr t : map ppr shape) - ppr (Reshape shape e) = - text "reshape" <> apply [apply (map ppr shape), ppr e] + ppr (Reshape ReshapeArbitrary shape e) = + text "reshape" <> apply [ppr shape, ppr e] + ppr (Reshape ReshapeCoerce shape e) = + text "coerce" <> apply [ppr shape, ppr e] ppr (Rearrange perm e) = text "rearrange" <> apply [apply (map ppr perm), ppr e] ppr (Rotate es e) = @@ -238,8 +240,17 @@ p (ErrorString s) = text $ show s p (ErrorVal t x) = ppr x <+> colon <+> ppr t +maybeNest :: PrettyRep rep => Body rep -> Doc +maybeNest b + | null $ bodyStms b = ppr b + | otherwise = nestedBlock "{" "}" $ ppr b + +instance PrettyRep rep => Pretty (Case (Body rep)) where + ppr (Case vs b) = + "case" <+> ppTuple' (map (maybe "_" ppr) vs) <+> "->" <+> maybeNest b + instance PrettyRep rep => Pretty (Exp rep) where - ppr (If c t f (IfDec ret ifsort)) = + ppr (Match [c] [Case [Just (BoolValue True)] t] f (MatchDec ret ifsort)) = text "if" <+> info' <+> ppr c @@ -251,12 +262,22 @@ <+> ppTuple' ret where info' = case ifsort of - IfNormal -> mempty - IfFallback -> text "" - IfEquiv -> text "" - maybeNest b - | null $ bodyStms b = ppr b - | otherwise = nestedBlock "{" "}" $ ppr b + MatchNormal -> mempty + MatchFallback -> text "" + MatchEquiv -> text "" + ppr (Match ses cs defb (MatchDec ret ifsort)) = + ("match" <+> info' <+> ppTuple' ses) + stack (map ppr cs) + "default" + <+> "->" + <+> maybeNest defb + colon + <+> ppTuple' ret + where + info' = case ifsort of + MatchNormal -> mempty + MatchFallback -> text "" + MatchEquiv -> text "" ppr (BasicOp op) = ppr op ppr (Apply fname args ret (safety, _, _)) = applykw @@ -376,10 +397,6 @@ instance PrettyRep rep => Pretty (Prog rep) where ppr (Prog types consts funs) = stack $ punctuate line $ ppr types : ppr consts : map ppr funs - -instance Pretty d => Pretty (DimChange d) where - ppr (DimCoercion se) = text "~" <> ppr se - ppr (DimNew se) = ppr se instance Pretty d => Pretty (DimIndex d) where ppr (DimFix i) = ppr i diff --git a/src/Futhark/IR/Prop.hs b/src/Futhark/IR/Prop.hs --- a/src/Futhark/IR/Prop.hs +++ b/src/Futhark/IR/Prop.hs @@ -126,9 +126,9 @@ safeExp (DoLoop _ _ body) = safeBody body safeExp (Apply fname _ _ _) = isBuiltInFunction fname -safeExp (If _ tbranch fbranch _) = - all (safeExp . stmExp) (bodyStms tbranch) - && all (safeExp . stmExp) (bodyStms fbranch) +safeExp (Match _ cases def_case _) = + all (all (safeExp . stmExp) . bodyStms . caseBody) cases + && all (safeExp . stmExp) (bodyStms def_case) safeExp WithAcc {} = True -- Although unlikely to matter. safeExp (Op op) = safeOp op diff --git a/src/Futhark/IR/Prop/Aliases.hs b/src/Futhark/IR/Prop/Aliases.hs --- a/src/Futhark/IR/Prop/Aliases.hs +++ b/src/Futhark/IR/Prop/Aliases.hs @@ -33,7 +33,7 @@ import Data.Bifunctor (first, second) import qualified Data.Kind -import Data.List (find) +import Data.List (find, transpose) import qualified Data.Map as M import Futhark.IR.Prop (IsOp, NameInfo (..), Scope) import Futhark.IR.Prop.Names @@ -72,7 +72,7 @@ basicOpAliases Iota {} = [mempty] basicOpAliases Replicate {} = [mempty] basicOpAliases Scratch {} = [mempty] -basicOpAliases (Reshape _ e) = [vnameAliases e] +basicOpAliases (Reshape _ _ e) = [vnameAliases e] basicOpAliases (Rearrange _ e) = [vnameAliases e] basicOpAliases (Rotate _ e) = [vnameAliases e] basicOpAliases Concat {} = [mempty] @@ -81,11 +81,11 @@ basicOpAliases Assert {} = [mempty] basicOpAliases UpdateAcc {} = [mempty] -ifAliases :: ([Names], Names) -> ([Names], Names) -> [Names] -ifAliases (als1, cons1) (als2, cons2) = - map (`namesSubtract` cons) $ zipWith mappend als1 als2 +matchAliases :: [([Names], Names)] -> [Names] +matchAliases l = + map ((`namesSubtract` mconcat conses) . mconcat) $ transpose alses where - cons = cons1 <> cons2 + (alses, conses) = unzip l funcallAliases :: [(SubExp, Diet)] -> [TypeBase shape Uniqueness] -> [Names] funcallAliases args t = @@ -93,14 +93,10 @@ -- | The aliases of an expression, one per non-context value returned. expAliases :: (Aliased rep) => Exp rep -> [Names] -expAliases (If _ tb fb dec) = - drop (length all_aliases - length ts) all_aliases +expAliases (Match _ cases defbody _) = + matchAliases $ onBody defbody : map (onBody . caseBody) cases where - ts = ifReturns dec - all_aliases = - ifAliases - (bodyAliases tb, consumedInBody tb) - (bodyAliases fb, consumedInBody fb) + onBody body = (bodyAliases body, consumedInBody body) expAliases (BasicOp op) = basicOpAliases op expAliases (DoLoop merge _ loopbody) = do (p, als) <- @@ -160,8 +156,8 @@ where consumeArg (als, Consume) = als consumeArg _ = mempty -consumedInExp (If _ tb fb _) = - consumedInBody tb <> consumedInBody fb +consumedInExp (Match _ cases defbody _) = + foldMap (consumedInBody . caseBody) cases <> consumedInBody defbody consumedInExp (DoLoop merge form body) = mconcat ( map (subExpAliases . snd) $ diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs --- a/src/Futhark/IR/Prop/Names.hs +++ b/src/Futhark/IR/Prop/Names.hs @@ -306,6 +306,9 @@ instance FreeIn (Stm rep) => FreeIn (Stms rep) where freeIn' = foldMap freeIn' +instance FreeIn body => FreeIn (Case body) where + freeIn' = freeIn' . caseBody + instance FreeIn Names where freeIn' = fvNames @@ -356,9 +359,6 @@ freeIn' (ForLoop _ _ bound loop_vars) = freeIn' bound <> freeIn' loop_vars freeIn' (WhileLoop cond) = freeIn' cond -instance FreeIn d => FreeIn (DimChange d) where - freeIn' = Data.Foldable.foldMap freeIn' - instance FreeIn d => FreeIn (DimIndex d) where freeIn' = Data.Foldable.foldMap freeIn' @@ -389,8 +389,8 @@ instance FreeIn dec => FreeIn (StmAux dec) where freeIn' (StmAux cs attrs dec) = freeIn' cs <> freeIn' attrs <> freeIn' dec -instance FreeIn a => FreeIn (IfDec a) where - freeIn' (IfDec r _) = freeIn' r +instance FreeIn a => FreeIn (MatchDec a) where + freeIn' (MatchDec r _) = freeIn' r -- | Either return precomputed free names stored in the attribute, or -- the freshly computed names. Relies on lazy evaluation to avoid the diff --git a/src/Futhark/IR/Prop/Reshape.hs b/src/Futhark/IR/Prop/Reshape.hs --- a/src/Futhark/IR/Prop/Reshape.hs +++ b/src/Futhark/IR/Prop/Reshape.hs @@ -1,24 +1,14 @@ -- | Facilities for creating, inspecting, and simplifying reshape and -- coercion operations. module Futhark.IR.Prop.Reshape - ( -- * Basic tools - newDim, - newDims, - newShape, - - -- * Construction + ( -- * Construction shapeCoerce, -- * Execution reshapeOuter, reshapeInner, - -- * Inspection - shapeCoercion, - -- * Simplification - fuseReshape, - informReshape, -- * Shape calculations reshapeIndex, @@ -33,86 +23,23 @@ import Futhark.Util.IntegralExp import Prelude hiding (product, quot, sum) --- | The new dimension. -newDim :: DimChange d -> d -newDim (DimCoercion se) = se -newDim (DimNew se) = se - --- | The new dimensions resulting from a reshape operation. -newDims :: ShapeChange d -> [d] -newDims = map newDim - --- | The new shape resulting from a reshape operation. -newShape :: ShapeChange SubExp -> Shape -newShape = Shape . newDims - --- | Construct a 'Reshape' where all dimension changes are --- 'DimCoercion's. +-- | Construct a 'Reshape' that is a 'ReshapeCoerce'. shapeCoerce :: [SubExp] -> VName -> Exp rep shapeCoerce newdims arr = - BasicOp $ Reshape (map DimCoercion newdims) arr + BasicOp $ Reshape ReshapeCoerce (Shape newdims) arr -- | @reshapeOuter newshape n oldshape@ returns a 'Reshape' expression -- that replaces the outer @n@ dimensions of @oldshape@ with @newshape@. -reshapeOuter :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp +reshapeOuter :: Shape -> Int -> Shape -> Shape reshapeOuter newshape n oldshape = - newshape ++ map coercion_or_new (drop n (shapeDims oldshape)) - where - coercion_or_new - | length newshape == n = DimCoercion - | otherwise = DimNew + newshape <> Shape (drop n (shapeDims oldshape)) -- | @reshapeInner newshape n oldshape@ returns a 'Reshape' expression -- that replaces the inner @m-n@ dimensions (where @m@ is the rank of -- @oldshape@) of @src@ with @newshape@. -reshapeInner :: ShapeChange SubExp -> Int -> Shape -> ShapeChange SubExp +reshapeInner :: Shape -> Int -> Shape -> Shape reshapeInner newshape n oldshape = - map coercion_or_new (take n (shapeDims oldshape)) ++ newshape - where - coercion_or_new - | length newshape == m - n = DimCoercion - | otherwise = DimNew - m = shapeRank oldshape - --- | If the shape change is nothing but shape coercions, return the new dimensions. Otherwise, return --- 'Nothing'. -shapeCoercion :: ShapeChange d -> Maybe [d] -shapeCoercion = mapM dimCoercion - where - dimCoercion (DimCoercion d) = Just d - dimCoercion (DimNew _) = Nothing - --- | @fuseReshape s1 s2@ creates a new 'ShapeChange' that is --- semantically the same as first applying @s1@ and then @s2@. This --- may take advantage of properties of 'DimCoercion' versus 'DimNew' --- to preserve information. -fuseReshape :: Eq d => ShapeChange d -> ShapeChange d -> ShapeChange d -fuseReshape s1 s2 - | length s1 == length s2 = - zipWith comb s1 s2 - where - comb (DimNew _) (DimCoercion d2) = - DimNew d2 - comb (DimCoercion d1) (DimNew d2) - | d1 == d2 = DimCoercion d2 - | otherwise = DimNew d2 - comb _ d2 = - d2 --- TODO: intelligently handle case where s1 is a prefix of s2. -fuseReshape _ s2 = s2 - --- | Given concrete information about the shape of the source array, --- convert some 'DimNew's into 'DimCoercion's. -informReshape :: Eq d => [d] -> ShapeChange d -> ShapeChange d -informReshape shape sc - | length shape == length sc = - zipWith inform shape sc - where - inform d1 (DimNew d2) - | d1 == d2 = DimCoercion d2 - inform _ dc = - dc -informReshape _ sc = sc + Shape (take n (shapeDims oldshape)) <> newshape -- | @reshapeIndex to_dims from_dims is@ transforms the index list -- @is@ (which is into an array of shape @from_dims@) into an index diff --git a/src/Futhark/IR/Prop/Scope.hs b/src/Futhark/IR/Prop/Scope.hs --- a/src/Futhark/IR/Prop/Scope.hs +++ b/src/Futhark/IR/Prop/Scope.hs @@ -102,7 +102,7 @@ asksScope f = f <$> askScope instance - (Applicative m, Monad m, RepTypes rep) => + (Monad m, RepTypes rep) => HasScope rep (ReaderT (Scope rep) m) where askScope = ask @@ -111,13 +111,13 @@ askScope = lift askScope instance - (Applicative m, Monad m, Monoid w, RepTypes rep) => + (Monad m, Monoid w, RepTypes rep) => HasScope rep (Control.Monad.RWS.Strict.RWST (Scope rep) w s m) where askScope = ask instance - (Applicative m, Monad m, Monoid w, RepTypes rep) => + (Monad m, Monoid w, RepTypes rep) => HasScope rep (Control.Monad.RWS.Lazy.RWST (Scope rep) w s m) where askScope = ask @@ -131,23 +131,23 @@ -- does not replace it. localScope :: Scope rep -> m a -> m a -instance (Monad m, LocalScope rep m) => LocalScope rep (ExceptT e m) where +instance (LocalScope rep m) => LocalScope rep (ExceptT e m) where localScope = mapExceptT . localScope instance - (Applicative m, Monad m, RepTypes rep) => + (Monad m, RepTypes rep) => LocalScope rep (ReaderT (Scope rep) m) where localScope = local . M.union instance - (Applicative m, Monad m, Monoid w, RepTypes rep) => + (Monad m, Monoid w, RepTypes rep) => LocalScope rep (Control.Monad.RWS.Strict.RWST (Scope rep) w s m) where localScope = local . M.union instance - (Applicative m, Monad m, Monoid w, RepTypes rep) => + (Monad m, Monoid w, RepTypes rep) => LocalScope rep (Control.Monad.RWS.Lazy.RWST (Scope rep) w s m) where localScope = local . M.union diff --git a/src/Futhark/IR/Prop/TypeOf.hs b/src/Futhark/IR/Prop/TypeOf.hs --- a/src/Futhark/IR/Prop/TypeOf.hs +++ b/src/Futhark/IR/Prop/TypeOf.hs @@ -37,7 +37,6 @@ import Data.List.NonEmpty (NonEmpty (..)) import Futhark.IR.Prop.Constants -import Futhark.IR.Prop.Reshape import Futhark.IR.Prop.Scope import Futhark.IR.Prop.Types import Futhark.IR.RetType @@ -102,14 +101,14 @@ pure . flip arrayOfShape shape <$> subExpType e basicOpType (Scratch t shape) = pure [arrayOf (Prim t) (Shape shape) NoUniqueness] -basicOpType (Reshape [] e) = +basicOpType (Reshape _ (Shape []) e) = result <$> lookupType e where result t = [Prim $ elemType t] -basicOpType (Reshape shape e) = +basicOpType (Reshape _ shape e) = result <$> lookupType e where - result t = [t `setArrayShape` newShape shape] + result t = [t `setArrayShape` shape] basicOpType (Rearrange perm e) = result <$> lookupType e where @@ -135,7 +134,7 @@ Exp rep -> m [ExtType] expExtType (Apply _ _ rt _) = pure $ map (fromDecl . declExtTypeOf) rt -expExtType (If _ _ _ rt) = pure $ map extTypeOf $ ifReturns rt +expExtType (Match _ _ _ rt) = pure $ map extTypeOf $ matchReturns rt expExtType (DoLoop merge _ _) = pure $ loopExtType $ map fst merge expExtType (BasicOp op) = staticShapes <$> basicOpType op diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -66,7 +66,8 @@ lamUsesAD lam || any (lamUsesAD . histOp) ops expUsesAD (Op (Scatter _ _ lam _)) = lamUsesAD lam - expUsesAD (If _ tbody fbody _) = bodyUsesAD tbody || bodyUsesAD fbody + expUsesAD (Match _ cases def_case _) = + any (bodyUsesAD . caseBody) cases || bodyUsesAD def_case expUsesAD (DoLoop _ _ body) = bodyUsesAD body expUsesAD (WithAcc _ lam) = lamUsesAD lam expUsesAD BasicOp {} = False diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs --- a/src/Futhark/IR/SOACS/SOAC.hs +++ b/src/Futhark/IR/SOACS/SOAC.hs @@ -402,7 +402,7 @@ -- SOAC. The mapping does not descend recursively into subexpressions -- and is done left-to-right. mapSOACM :: - (Applicative m, Monad m) => + Monad m => SOACMapper frep trep m -> SOAC frep -> m (SOAC trep) diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -353,7 +353,7 @@ -- | Remove all arguments to the map that are simply replicates. -- These can be turned into free variables instead. removeReplicateMapping :: - (Aliased rep, Buildable rep, BuilderOps rep, HasSOAC rep) => + (Aliased rep, BuilderOps rep, HasSOAC rep) => TopDownRuleOp rep removeReplicateMapping vtable pat aux op | Just (Screma w arrs form) <- asSOAC op, @@ -484,15 +484,12 @@ -- Mapping some operations becomes an extension of that operation. mapOpToOp :: BottomUpRuleOp (Wise SOACS) mapOpToOp (_, used) pat aux1 e - | Just (map_pe, cs, w, BasicOp (Reshape newshape reshape_arr), [p], [arr]) <- + | Just (map_pe, cs, w, BasicOp (Reshape k newshape reshape_arr), [p], [arr]) <- isMapWithOp pat e, paramName p == reshape_arr, not $ UT.isConsumed (patElemName map_pe) used = Simplify $ do - let redim - | isJust $ shapeCoercion newshape = DimCoercion w - | otherwise = DimNew w certifying (stmAuxCerts aux1 <> cs) . letBind pat . BasicOp $ - Reshape (redim : newshape) arr + Reshape k (Shape [w] <> newshape) arr | Just (_, cs, _, BasicOp (Concat d (arr :| arrs) dw), ps, outer_arr : outer_arrs) <- isMapWithOp pat e, (arr : arrs) == map paramName ps = @@ -628,8 +625,7 @@ y_ws <- mapM sizeOf ys guard $ all (x_w ==) y_ws pure (x_w, x : ys, cs) - Just (BasicOp (Reshape reshape arr), cs) -> do - guard $ isJust $ shapeCoercion reshape + Just (BasicOp (Reshape ReshapeCoerce _ arr), cs) -> do (a, b, cs') <- isConcat arr pure (a, b, cs <> cs') _ -> Nothing @@ -710,7 +706,7 @@ = ArrayIndexing Certs VName (Slice SubExp) | ArrayRearrange Certs VName [Int] | ArrayRotate Certs VName [SubExp] - | ArrayReshape Certs VName (ShapeChange SubExp) + | ArrayReshape Certs VName ReshapeKind Shape | ArrayCopy Certs VName | -- | Never constructed. ArrayVar Certs VName @@ -720,7 +716,7 @@ arrayOpArr (ArrayIndexing _ arr _) = arr arrayOpArr (ArrayRearrange _ arr _) = arr arrayOpArr (ArrayRotate _ arr _) = arr -arrayOpArr (ArrayReshape _ arr _) = arr +arrayOpArr (ArrayReshape _ arr _ _) = arr arrayOpArr (ArrayCopy _ arr) = arr arrayOpArr (ArrayVar _ arr) = arr @@ -728,7 +724,7 @@ arrayOpCerts (ArrayIndexing cs _ _) = cs arrayOpCerts (ArrayRearrange cs _ _) = cs arrayOpCerts (ArrayRotate cs _ _) = cs -arrayOpCerts (ArrayReshape cs _ _) = cs +arrayOpCerts (ArrayReshape cs _ _ _) = cs arrayOpCerts (ArrayCopy cs _) = cs arrayOpCerts (ArrayVar cs _) = cs @@ -739,8 +735,8 @@ Just $ ArrayRearrange cs arr perm isArrayOp cs (BasicOp (Rotate rots arr)) = Just $ ArrayRotate cs arr rots -isArrayOp cs (BasicOp (Reshape new_shape arr)) = - Just $ ArrayReshape cs arr new_shape +isArrayOp cs (BasicOp (Reshape k new_shape arr)) = + Just $ ArrayReshape cs arr k new_shape isArrayOp cs (BasicOp (Copy arr)) = Just $ ArrayCopy cs arr isArrayOp _ _ = @@ -750,7 +746,7 @@ fromArrayOp (ArrayIndexing cs arr slice) = (cs, BasicOp $ Index arr slice) fromArrayOp (ArrayRearrange cs arr perm) = (cs, BasicOp $ Rearrange perm arr) fromArrayOp (ArrayRotate cs arr rots) = (cs, BasicOp $ Rotate rots arr) -fromArrayOp (ArrayReshape cs arr new_shape) = (cs, BasicOp $ Reshape new_shape arr) +fromArrayOp (ArrayReshape cs arr k new_shape) = (cs, BasicOp $ Reshape k new_shape arr) fromArrayOp (ArrayCopy cs arr) = (cs, BasicOp $ Copy arr) fromArrayOp (ArrayVar cs arr) = (cs, BasicOp $ SubExp $ Var arr) @@ -949,7 +945,7 @@ arrayIsMapParam (_, ArrayRotate cs arr rots) = arr `elem` map_param_names && all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn rots) - arrayIsMapParam (_, ArrayReshape cs arr new_shape) = + arrayIsMapParam (_, ArrayReshape cs arr _ new_shape) = arr `elem` map_param_names && all (`ST.elem` vtable) (namesToList $ freeIn cs <> freeIn new_shape) arrayIsMapParam (_, ArrayCopy cs arr) = @@ -972,8 +968,8 @@ BasicOp $ Rearrange (0 : map (+ 1) perm) arr ArrayRotate _ _ rots -> BasicOp $ Rotate (intConst Int64 0 : rots) arr - ArrayReshape _ _ new_shape -> - BasicOp $ Reshape (DimCoercion w : new_shape) arr + ArrayReshape _ _ k new_shape -> + BasicOp $ Reshape k (Shape [w] <> new_shape) arr ArrayCopy {} -> BasicOp $ Copy arr ArrayVar {} -> diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs --- a/src/Futhark/IR/SegOp.hs +++ b/src/Futhark/IR/SegOp.hs @@ -830,7 +830,7 @@ -- | Apply a 'SegOpMapper' to the given 'SegOp'. mapSegOpM :: - (Applicative m, Monad m) => + Monad m => SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep) @@ -922,10 +922,7 @@ where renamer = SegOpMapper rename rename rename rename rename -instance - (ASTRep rep, FreeIn (LParamInfo rep), FreeIn lvl) => - FreeIn (SegOp lvl rep) - where +instance (ASTRep rep, FreeIn lvl) => FreeIn (SegOp lvl rep) where freeIn' e = fvBind (namesFromList $ M.keys $ scopeOfSegSpace (segSpace e)) $ flip execState mempty $ diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs --- a/src/Futhark/IR/Syntax.hs +++ b/src/Futhark/IR/Syntax.hs @@ -132,13 +132,13 @@ CmpOp (..), ConvOp (..), OpaqueOp (..), - DimChange (..), - ShapeChange, + ReshapeKind (..), WithAccInput, Exp (..), + Case (..), LoopForm (..), - IfDec (..), - IfSort (..), + MatchDec (..), + MatchSort (..), Safety (..), Lambda (..), @@ -298,40 +298,6 @@ deriving instance RepTypes rep => Eq (Body rep) --- | The new dimension in a 'Reshape'-like operation. This allows us to --- disambiguate "real" reshapes, that change the actual shape of the --- array, from type coercions that are just present to make the types --- work out. The two constructors are considered equal for purposes of 'Eq'. -data DimChange d - = -- | The new dimension is guaranteed to be numerically - -- equal to the old one. - DimCoercion d - | -- | The new dimension is not necessarily numerically - -- equal to the old one. - DimNew d - deriving (Ord, Show) - -instance Eq d => Eq (DimChange d) where - DimCoercion x == DimNew y = x == y - DimCoercion x == DimCoercion y = x == y - DimNew x == DimCoercion y = x == y - DimNew x == DimNew y = x == y - -instance Functor DimChange where - fmap f (DimCoercion d) = DimCoercion $ f d - fmap f (DimNew d) = DimNew $ f d - -instance Foldable DimChange where - foldMap f (DimCoercion d) = f d - foldMap f (DimNew d) = f d - -instance Traversable DimChange where - traverse f (DimCoercion d) = DimCoercion <$> f d - traverse f (DimNew d) = DimNew <$> f d - --- | A list of 'DimChange's, indicating the new dimensions of an array. -type ShapeChange d = [DimChange d] - -- | Apart from being Opaque, what else is going on here? data OpaqueOp = -- | No special operation. @@ -340,6 +306,14 @@ OpaqueTrace String deriving (Eq, Ord, Show) +-- | Which kind of reshape is this? +data ReshapeKind + = -- | Any kind of reshaping. + ReshapeCoerce + | -- | New shape is dynamically same as original. + ReshapeArbitrary + deriving (Eq, Ord, Show) + -- | A primitive operation that returns something of known size and -- does not itself contain any bindings. data BasicOp @@ -366,9 +340,7 @@ | -- | Turn a boolean into a certificate, halting the program with the -- given error message if the boolean is false. Assert SubExp (ErrorMsg SubExp) (SrcLoc, [SrcLoc]) - | -- Primitive array operations - - -- | The certificates for bounds-checking are part of the 'Stm'. + | -- | The certificates for bounds-checking are part of the 'Stm'. Index VName (Slice SubExp) | -- | An in-place update of the given array at the given position. -- Consumes the array. If 'Safe', perform a run-time bounds check @@ -401,10 +373,8 @@ Replicate Shape SubExp | -- | Create array of given type and shape, with undefined elements. Scratch PrimType [SubExp] - | -- Array index space transformation. - - -- | 1st arg is the new shape, 2nd arg is the input array. - Reshape (ShapeChange SubExp) VName + | -- | 1st arg is the new shape, 2nd arg is the input array. + Reshape ReshapeKind Shape VName | -- | Permute the dimensions of the input array. The list -- of integers is a list of dimensions (0-indexed), which -- must be a permutation of @[0,n-1]@, where @n@ is the @@ -425,6 +395,22 @@ type WithAccInput rep = (Shape, [VName], Maybe (Lambda rep, [SubExp])) +-- | A non-default case in a 'Match' statement. The number of +-- elements in the pattern must match the number of scrutinees. A +-- 'Nothing' value indicates that we don't care about it (i.e. a +-- wildcard). +data Case body = Case {casePat :: [Maybe PrimValue], caseBody :: body} + deriving (Eq, Ord, Show) + +instance Functor Case where + fmap = fmapDefault + +instance Foldable Case where + foldMap = foldMapDefault + +instance Traversable Case where + traverse f (Case vs b) = Case vs <$> f b + -- | The root Futhark expression type. The v'Op' constructor contains -- a rep-specific operation. Do-loops, branches and function calls -- are special. Everything else is a simple t'BasicOp'. @@ -432,7 +418,11 @@ = -- | A simple (non-recursive) operation. BasicOp BasicOp | Apply Name [(SubExp, Diet)] [RetType rep] (Safety, SrcLoc, [SrcLoc]) - | If SubExp (Body rep) (Body rep) (IfDec (BranchType rep)) + | -- | A match statement picks a branch by comparing the given + -- subexpressions (called the /scrutinee/) with the pattern in + -- each of the cases. If none of the cases match, the /default + -- body/ is picked. + Match [SubExp] [Case (Body rep)] (Body rep) (MatchDec (BranchType rep)) | -- | @loop {a} = {v} (for i < n|while b) do b@. DoLoop [(FParam rep, SubExp)] (LoopForm rep) (Body rep) | -- | Create accumulators backed by the given arrays (which are @@ -463,29 +453,29 @@ deriving instance RepTypes rep => Ord (LoopForm rep) -- | Data associated with a branch. -data IfDec rt = IfDec - { ifReturns :: [rt], - ifSort :: IfSort +data MatchDec rt = MatchDec + { matchReturns :: [rt], + matchSort :: MatchSort } deriving (Eq, Show, Ord) -- | What kind of branch is this? This has no semantic meaning, but -- provides hints to simplifications. -data IfSort +data MatchSort = -- | An ordinary branch. - IfNormal + MatchNormal | -- | A branch where the "true" case is what we are -- actually interested in, and the "false" case is only -- present as a fallback for when the true case cannot -- be safely evaluated. The compiler is permitted to -- optimise away the branch if the true case contains -- only safe statements. - IfFallback + MatchFallback | -- | Both of these branches are semantically equivalent, -- and it is fine to eliminate one if it turns out to -- have problems (e.g. contain things we cannot generate -- code for). - IfEquiv + MatchEquiv deriving (Eq, Show, Ord) -- | Anonymous function for use in a SOAC. diff --git a/src/Futhark/IR/Traversals.hs b/src/Futhark/IR/Traversals.hs --- a/src/Futhark/IR/Traversals.hs +++ b/src/Futhark/IR/Traversals.hs @@ -83,7 +83,7 @@ -- expression. Importantly, the mapping does not descend recursively -- into subexpressions. The mapping is done left-to-right. mapExpM :: - (Applicative m, Monad m) => + Monad m => Mapper frep trep m -> Exp frep -> m (Exp trep) @@ -103,12 +103,14 @@ BasicOp <$> (ConvOp conv <$> mapOnSubExp tv x) mapExpM tv (BasicOp (UnOp op x)) = BasicOp <$> (UnOp op <$> mapOnSubExp tv x) -mapExpM tv (If c texp fexp (IfDec ts s)) = - If - <$> mapOnSubExp tv c - <*> mapOnBody tv mempty texp - <*> mapOnBody tv mempty fexp - <*> (IfDec <$> mapM (mapOnBranchType tv) ts <*> pure s) +mapExpM tv (Match ses cases defbody (MatchDec ts s)) = + Match + <$> mapM (mapOnSubExp tv) ses + <*> mapM mapOnCase cases + <*> mapOnBody tv mempty defbody + <*> (MatchDec <$> mapM (mapOnBranchType tv) ts <*> pure s) + where + mapOnCase (Case vs body) = Case vs <$> mapOnBody tv mempty body mapExpM tv (Apply fname args ret loc) = do args' <- forM args $ \(arg, d) -> (,) <$> mapOnSubExp tv arg <*> pure d @@ -145,10 +147,10 @@ BasicOp <$> (Replicate <$> mapOnShape tv shape <*> mapOnSubExp tv vexp) mapExpM tv (BasicOp (Scratch t shape)) = BasicOp <$> (Scratch t <$> mapM (mapOnSubExp tv) shape) -mapExpM tv (BasicOp (Reshape shape arrexp)) = +mapExpM tv (BasicOp (Reshape kind shape arrexp)) = BasicOp - <$> ( Reshape - <$> mapM (traverse (mapOnSubExp tv)) shape + <$> ( Reshape kind + <$> mapM (mapOnSubExp tv) shape <*> mapOnVName tv arrexp ) mapExpM tv (BasicOp (Rearrange perm e)) = @@ -301,10 +303,10 @@ walkOnSubExp tv x walkExpM tv (BasicOp (UnOp _ x)) = walkOnSubExp tv x -walkExpM tv (If c texp fexp (IfDec ts _)) = do - walkOnSubExp tv c - walkOnBody tv mempty texp - walkOnBody tv mempty fexp +walkExpM tv (Match ses cases defbody (MatchDec ts _)) = do + mapM_ (walkOnSubExp tv) ses + mapM_ (walkOnBody tv mempty . caseBody) cases + walkOnBody tv mempty defbody mapM_ (walkOnBranchType tv) ts walkExpM tv (Apply _ args ret _) = mapM_ (walkOnSubExp tv . fst) args >> mapM_ (walkOnRetType tv) ret @@ -326,8 +328,8 @@ walkOnShape tv shape >> walkOnSubExp tv vexp walkExpM tv (BasicOp (Scratch _ shape)) = mapM_ (walkOnSubExp tv) shape -walkExpM tv (BasicOp (Reshape shape arrexp)) = - mapM_ (traverse_ (walkOnSubExp tv)) shape >> walkOnVName tv arrexp +walkExpM tv (BasicOp (Reshape _ shape arrexp)) = + mapM_ (walkOnSubExp tv) shape >> walkOnVName tv arrexp walkExpM tv (BasicOp (Rearrange _ e)) = walkOnVName tv e walkExpM tv (BasicOp (Rotate es e)) = diff --git a/src/Futhark/IR/TypeCheck.hs b/src/Futhark/IR/TypeCheck.hs --- a/src/Futhark/IR/TypeCheck.hs +++ b/src/Futhark/IR/TypeCheck.hs @@ -904,27 +904,15 @@ void $ checkSubExp valexp checkBasicOp (Scratch _ shape) = mapM_ checkSubExp shape -checkBasicOp (Reshape newshape arrexp) = do +checkBasicOp (Reshape k newshape arrexp) = do rank <- shapeRank . fst <$> checkArrIdent arrexp - mapM_ (require [Prim int64] . newDim) newshape - zipWithM_ (checkDimChange rank) newshape [0 ..] - where - checkDimChange _ (DimNew _) _ = + mapM_ (require [Prim int64]) $ shapeDims newshape + case k of + ReshapeCoerce -> + when (shapeRank newshape /= rank) . bad $ + TypeError "Coercion changes rank of array." + ReshapeArbitrary -> pure () - checkDimChange rank (DimCoercion se) i - | i >= rank = - bad . TypeError $ - "Asked to coerce dimension " - ++ show i - ++ " to " - ++ pretty se - ++ ", but array " - ++ pretty arrexp - ++ " has only " - ++ pretty rank - ++ " dimensions" - | otherwise = - pure () checkBasicOp (Rearrange perm arr) = do arrt <- lookupType arr let rank = arrayRank arrt @@ -1014,13 +1002,22 @@ Exp (Aliases rep) -> TypeM rep () checkExp (BasicOp op) = checkBasicOp op -checkExp (If e1 e2 e3 info) = do - require [Prim Bool] e1 - _ <- - context "in true branch" (checkBody e2) - `alternative` context "in false branch" (checkBody e3) - context "in true branch" $ matchBranchType (ifReturns info) e2 - context "in false branch" $ matchBranchType (ifReturns info) e3 +checkExp (Match ses cases def_case info) = do + ses_ts <- mapM checkSubExp ses + mapM_ (checkCase ses_ts) cases + checkCaseBody def_case + where + checkVal t (Just v) = Prim (primValueType v) == t + checkVal _ Nothing = True + checkCase ses_ts (Case vs body) = do + let ok = length vs == length ses_ts && and (zipWith checkVal ses_ts vs) + unless ok . bad . TypeError . pretty $ + "Scrutinee" + indent 2 (ppTuple' ses) + "cannot match pattern" + indent 2 (ppTuple' vs) + context ("in body of case " <> prettyTuple vs) $ checkCaseBody body + checkCaseBody = matchBranchType (matchReturns info) checkExp (Apply fname args rettype_annot _) = do (rettype_derived, paramtypes) <- lookupFun fname $ map fst args argflows <- mapM (checkArg . fst) args diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -9,7 +9,7 @@ module Futhark.Internalise.Exps (transformProg) where import Control.Monad.Reader -import Data.List (find, intercalate, intersperse, transpose) +import Data.List (elemIndex, find, intercalate, intersperse, transpose) import Data.List.NonEmpty (NonEmpty (..)) import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as M @@ -160,11 +160,6 @@ letValExp' _ (BasicOp (SubExp se)) = pure [se] letValExp' name ses = map I.Var <$> letValExp name ses -eValBody :: [InternaliseM (I.Exp SOACS)] -> InternaliseM (I.Body SOACS) -eValBody es = buildBody_ $ do - es' <- sequence es - varsRes . concat <$> mapM (letValExp "x") es' - internaliseAppExp :: String -> E.AppRes -> E.AppExp -> InternaliseM [I.SubExp] internaliseAppExp desc _ (E.Index e idxs loc) = do vs <- internaliseExpToVars "indexed" e @@ -279,18 +274,16 @@ bounds_invalid <- letSubExp "bounds_invalid" - $ I.If - downwards - (resultBody [bounds_invalid_downwards]) - (resultBody [bounds_invalid_upwards]) - $ ifCommon [I.Prim I.Bool] + =<< eIf + (eSubExp downwards) + (resultBodyM [bounds_invalid_downwards]) + (resultBodyM [bounds_invalid_upwards]) distance_exclusive <- letSubExp "distance_exclusive" - $ I.If - downwards - (resultBody [distance_downwards_exclusive]) - (resultBody [distance_upwards_exclusive]) - $ ifCommon [I.Prim $ IntType it] + =<< eIf + (eSubExp downwards) + (resultBodyM [distance_downwards_exclusive]) + (resultBodyM [distance_upwards_exclusive]) distance_exclusive_i64 <- asIntS Int64 distance_exclusive distance <- letSubExp "distance" $ @@ -374,7 +367,7 @@ args' <- concat . reverse <$> mapM (internaliseArg arg_desc) (reverse args) fst <$> funcall desc qfname args' loc internaliseAppExp desc _ (E.LetPat sizes pat e body _) = - internalisePat desc sizes pat e body (internaliseExp desc) + internalisePat desc sizes pat e $ internaliseExp desc body internaliseAppExp _ _ (E.LetFun ofname _ _ _) = error $ "Unexpected LetFun " ++ pretty ofname internaliseAppExp desc _ (E.DoLoop sparams mergepat mergeexp form loopbody loc) = do @@ -489,7 +482,7 @@ case se of I.Var v | not $ primType $ paramType p -> - Reshape (map DimCoercion $ arrayDims $ paramType p) v + Reshape I.ReshapeCoerce (I.arrayShape $ paramType p) v _ -> SubExp se internaliseExp1 "loop_cond" cond @@ -515,7 +508,7 @@ case se of I.Var v | not $ primType $ paramType p -> - Reshape (map DimCoercion $ arrayDims $ paramType p) v + Reshape I.ReshapeCoerce (I.arrayShape $ paramType p) v _ -> SubExp se subExpsRes <$> internaliseExp "loop_cond" cond loop_end_cond <- bodyBind loop_end_cond_body @@ -536,21 +529,20 @@ E.AppExp (E.LetPat [] pat e body loc) (Info (AppRes (E.typeOf body) mempty)) -internaliseAppExp desc _ (E.Match e cs _) = do +internaliseAppExp desc _ (E.Match e orig_cs _) = do ses <- internaliseExp (desc ++ "_scrutinee") e + cs <- mapM (onCase ses) orig_cs case NE.uncons cs of - (CasePat pCase eCase _, Nothing) -> do - (_, pertinent) <- generateCond pCase ses - internalisePat' [] pCase pertinent eCase (internaliseExp desc) - (c, Just cs') -> do - let CasePat pLast eLast _ = NE.last cs' - bFalse <- do - (_, pertinent) <- generateCond pLast ses - eLast' <- internalisePat' [] pLast pertinent eLast (internaliseBody desc) - foldM (\bf c' -> eValBody $ pure $ generateCaseIf ses c' bf) eLast' $ - reverse $ - NE.init cs' - letValExp' desc =<< generateCaseIf ses c bFalse + (I.Case _ body, Nothing) -> + fmap (map resSubExp) $ bodyBind =<< body + _ -> do + letValExp' desc =<< eMatch ses (NE.init cs) (I.caseBody $ NE.last cs) + where + onCase ses (E.CasePat p case_e _) = do + (cmps, pertinent) <- generateCond p ses + pure . I.Case cmps $ + internalisePat' [] p pertinent $ + internaliseBody "case" case_e internaliseAppExp desc _ (E.If ce te fe _) = letValExp' desc =<< eIf @@ -625,10 +617,10 @@ flat_arr_t <- lookupType flat_arr let new_shape' = reshapeOuter - (map (DimNew . intConst Int64 . toInteger) new_shape) + (I.Shape $ map (intConst Int64 . toInteger) new_shape) 1 $ I.arrayShape flat_arr_t - letSubExp desc $ I.BasicOp $ I.Reshape new_shape' flat_arr + letSubExp desc $ I.BasicOp $ I.Reshape I.ReshapeArbitrary new_shape' flat_arr | otherwise = do es' <- mapM (internaliseExp "arr_elem") es let arr_t_ext = internaliseType $ E.toStruct arr_t @@ -836,33 +828,45 @@ _ -> pure () pure arg' -subExpPrimType :: I.SubExp -> InternaliseM I.PrimType -subExpPrimType = fmap I.elemType . subExpType +internalisePatLit :: E.PatLit -> E.PatType -> I.PrimValue +internalisePatLit (E.PatLitPrim v) _ = + internalisePrimValue v +internalisePatLit (E.PatLitInt x) (E.Scalar (E.Prim (E.Signed it))) = + I.IntValue $ intValue it x +internalisePatLit (E.PatLitInt x) (E.Scalar (E.Prim (E.Unsigned it))) = + I.IntValue $ intValue it x +internalisePatLit (E.PatLitFloat x) (E.Scalar (E.Prim (E.FloatType ft))) = + I.FloatValue $ floatValue ft x +internalisePatLit l t = + error $ "Nonsensical pattern and type: " ++ show (l, t) -generateCond :: E.Pat -> [I.SubExp] -> InternaliseM (I.SubExp, [I.SubExp]) +generateCond :: + E.Pat -> + [I.SubExp] -> + InternaliseM ([Maybe I.PrimValue], [I.SubExp]) generateCond orig_p orig_ses = do (cmps, pertinent, _) <- compares orig_p orig_ses - cmp <- letSubExp "matches" =<< eAll cmps - pure (cmp, pertinent) + pure (cmps, pertinent) where - -- Literals are always primitive values. - compares (E.PatLit l t _) (se : ses) = do - e' <- case l of - PatLitPrim v -> pure $ constant $ internalisePrimValue v - PatLitInt x -> internaliseExp1 "constant" $ E.IntLit x t mempty - PatLitFloat x -> internaliseExp1 "constant" $ E.FloatLit x t mempty - t' <- subExpPrimType se - cmp <- letSubExp "match_lit" $ I.BasicOp $ I.CmpOp (I.CmpEq t') e' se - pure ([cmp], [se], ses) - compares (E.PatConstr c (Info (E.Scalar (E.Sum fs))) pats _) (se : ses) = do + compares (E.PatLit l (Info t) _) (se : ses) = + pure ([Just $ internalisePatLit l t], [se], ses) + compares (E.PatConstr c (Info (E.Scalar (E.Sum fs))) pats _) (_ : ses) = do (payload_ts, m) <- internaliseSumType $ M.map (map toStruct) fs case M.lookup c m of - Just (i, payload_is) -> do - let i' = intConst Int8 $ toInteger i + Just (tag, payload_is) -> do let (payload_ses, ses') = splitAt (length payload_ts) ses - cmp <- letSubExp "match_constr" $ I.BasicOp $ I.CmpOp (I.CmpEq int8) i' se - (cmps, pertinent, _) <- comparesMany pats $ map (payload_ses !!) payload_is - pure (cmp : cmps, pertinent, ses') + (cmps, pertinent, _) <- + comparesMany pats $ map (payload_ses !!) payload_is + let missingCmps i _ = + case i `elemIndex` payload_is of + Just j -> cmps !! j + Nothing -> Nothing + pure + ( Just (I.IntValue $ intValue Int8 $ toInteger tag) + : zipWith missingCmps [0 ..] payload_ses, + pertinent, + ses' + ) Nothing -> error "generateCond: missing constructor" compares (E.PatConstr _ (Info t) _ _) _ = @@ -871,7 +875,7 @@ compares (E.Wildcard t loc) ses compares (E.Wildcard (Info t) _) ses = do let (id_ses, rest_ses) = splitAt (internalisedTypeSize $ E.toStruct t) ses - pure ([], id_ses, rest_ses) + pure (map (const Nothing) id_ses, id_ses, rest_ses) compares (E.PatParens pat _) ses = compares pat ses compares (E.PatAttr _ pat _) ses = @@ -900,23 +904,16 @@ ses'' ) -generateCaseIf :: [I.SubExp] -> Case -> I.Body SOACS -> InternaliseM (I.Exp SOACS) -generateCaseIf ses (CasePat p eCase _) bFail = do - (cond, pertinent) <- generateCond p ses - eCase' <- internalisePat' [] p pertinent eCase (internaliseBody "case") - eIf (eSubExp cond) (pure eCase') (pure bFail) - internalisePat :: String -> [E.SizeBinder VName] -> E.Pat -> E.Exp -> - E.Exp -> - (E.Exp -> InternaliseM a) -> + InternaliseM a -> InternaliseM a -internalisePat desc sizes p e body m = do +internalisePat desc sizes p e m = do ses <- internaliseExp desc' e - internalisePat' sizes p ses body m + internalisePat' sizes p ses m where desc' = case S.toList $ E.patIdents p of [v] -> baseString $ E.identName v @@ -926,16 +923,15 @@ [E.SizeBinder VName] -> E.Pat -> [I.SubExp] -> - E.Exp -> - (E.Exp -> InternaliseM a) -> + InternaliseM a -> InternaliseM a -internalisePat' sizes p ses body m = do +internalisePat' sizes p ses m = do ses_ts <- mapM subExpType ses stmPat p ses_ts $ \pat_names -> do bindExtSizes (AppRes (E.patternType p) (map E.sizeName sizes)) ses forM_ (zip pat_names ses) $ \(v, se) -> letBindNames [v] $ I.BasicOp $ I.SubExp se - m body + m internaliseSlice :: SrcLoc -> @@ -994,18 +990,16 @@ w_minus_1 <- letSubExp "w_minus_1" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) w one let i_def = letSubExp "i_def" - $ I.If - backwards - (resultBody [w_minus_1]) - (resultBody [zero]) - $ ifCommon [I.Prim int64] + =<< eIf + (eSubExp backwards) + (resultBodyM [w_minus_1]) + (resultBodyM [zero]) j_def = letSubExp "j_def" - $ I.If - backwards - (resultBody [negone]) - (resultBody [w]) - $ ifCommon [I.Prim int64] + =<< eIf + (eSubExp backwards) + (resultBodyM [negone]) + (resultBodyM [w]) i' <- maybe i_def (fmap fst . internaliseSizeExp "i") i j' <- maybe j_def (fmap fst . internaliseSizeExp "j") j j_m_i <- letSubExp "j_m_i" $ BasicOp $ I.BinOp (Sub Int64 I.OverflowWrap) j' i' @@ -1062,11 +1056,10 @@ slice_ok <- letSubExp "slice_ok" - $ I.If - backwards - (resultBody [backwards_ok]) - (resultBody [forwards_ok]) - $ ifCommon [I.Prim I.Bool] + =<< eIf + (eSubExp backwards) + (resultBodyM [backwards_ok]) + (resultBodyM [forwards_ok]) ok_or_empty <- letSubExp "ok_or_empty" $ @@ -1634,32 +1627,33 @@ dims_match <- forM (zip x_dims y_dims) $ \(x_dim, y_dim) -> letSubExp "dim_eq" $ I.BasicOp $ I.CmpOp (I.CmpEq int64) x_dim y_dim shapes_match <- letSubExp "shapes_match" =<< eAll dims_match - compare_elems_body <- runBodyBuilder $ do - -- Flatten both x and y. - x_num_elems <- - letSubExp "x_num_elems" - =<< foldBinOp (I.Mul Int64 I.OverflowUndef) (constant (1 :: Int64)) x_dims - x' <- letExp "x" $ I.BasicOp $ I.SubExp x - y' <- letExp "x" $ I.BasicOp $ I.SubExp y - x_flat <- letExp "x_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] x' - y_flat <- letExp "y_flat" $ I.BasicOp $ I.Reshape [I.DimNew x_num_elems] y' + let compare_elems_body = runBodyBuilder $ do + -- Flatten both x and y. + x_num_elems <- + letSubExp "x_num_elems" + =<< foldBinOp (I.Mul Int64 I.OverflowUndef) (constant (1 :: Int64)) x_dims + x' <- letExp "x" $ I.BasicOp $ I.SubExp x + y' <- letExp "x" $ I.BasicOp $ I.SubExp y + x_flat <- + letExp "x_flat" $ I.BasicOp $ I.Reshape I.ReshapeArbitrary (I.Shape [x_num_elems]) x' + y_flat <- + letExp "y_flat" $ I.BasicOp $ I.Reshape I.ReshapeArbitrary (I.Shape [x_num_elems]) y' - -- Compare the elements. - cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t) - cmps <- - letExp "cmps" $ - I.Op $ - I.Screma x_num_elems [x_flat, y_flat] (I.mapSOAC cmp_lam) + -- Compare the elements. + cmp_lam <- cmpOpLambda $ I.CmpEq (elemType x_t) + cmps <- + letExp "cmps" $ + I.Op $ + I.Screma x_num_elems [x_flat, y_flat] (I.mapSOAC cmp_lam) - -- Check that all were equal. - and_lam <- binOpLambda I.LogAnd I.Bool - reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]] - all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems [cmps] reduce - pure $ resultBody [all_equal] + -- Check that all were equal. + and_lam <- binOpLambda I.LogAnd I.Bool + reduce <- I.reduceSOAC [Reduce Commutative and_lam [constant True]] + all_equal <- letSubExp "all_equal" $ I.Op $ I.Screma x_num_elems [cmps] reduce + pure $ resultBody [all_equal] - letSubExp "arrays_equal" $ - I.If shapes_match compare_elems_body (resultBody [constant False]) $ - ifCommon [I.Prim I.Bool] + letSubExp "arrays_equal" + =<< eIf (eSubExp shapes_match) compare_elems_body (resultBodyM [constant False]) handleOps [x, y] name | Just bop <- find ((name ==) . pretty) [minBound .. maxBound :: E.BinOp] = Just $ \desc -> do @@ -1776,9 +1770,11 @@ certifying dim_ok_cert $ forM arrs $ \arr' -> do arr_t <- lookupType arr' - letSubExp desc $ - I.BasicOp $ - I.Reshape (reshapeOuter [DimNew n', DimNew m'] 1 $ I.arrayShape arr_t) arr' + letSubExp desc . I.BasicOp $ + I.Reshape + I.ReshapeArbitrary + (reshapeOuter (I.Shape [n', m']) 1 $ I.arrayShape arr_t) + arr' handleRest [arr] "flatten" = Just $ \desc -> do arrs <- internaliseExpToVars "flatten_arr" arr forM arrs $ \arr' -> do @@ -1786,9 +1782,11 @@ let n = arraySize 0 arr_t m = arraySize 1 arr_t k <- letSubExp "flat_dim" $ I.BasicOp $ I.BinOp (Mul Int64 I.OverflowUndef) n m - letSubExp desc $ - I.BasicOp $ - I.Reshape (reshapeOuter [DimNew k] 2 $ I.arrayShape arr_t) arr' + letSubExp desc . I.BasicOp $ + I.Reshape + I.ReshapeArbitrary + (reshapeOuter (I.Shape [k]) 2 $ I.arrayShape arr_t) + arr' handleRest [TupLit [x, y] _] "concat" = Just $ \desc -> do xs <- internaliseExpToVars "concat_x" x ys <- internaliseExpToVars "concat_y" y @@ -1841,11 +1839,10 @@ case E.typeOf e of E.Scalar (E.Prim E.Bool) -> letTupExp' desc - $ I.If - e' - (resultBody [intConst int_to 1]) - (resultBody [intConst int_to 0]) - $ ifCommon [I.Prim $ I.IntType int_to] + =<< eIf + (eSubExp e') + (resultBodyM [intConst int_to 1]) + (resultBodyM [intConst int_to 0]) E.Scalar (E.Prim (E.Signed int_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.SExt int_from int_to) e' E.Scalar (E.Prim (E.Unsigned int_from)) -> @@ -1859,11 +1856,10 @@ case E.typeOf e of E.Scalar (E.Prim E.Bool) -> letTupExp' desc - $ I.If - e' - (resultBody [intConst int_to 1]) - (resultBody [intConst int_to 0]) - $ ifCommon [I.Prim $ I.IntType int_to] + =<< eIf + (eSubExp e') + (resultBodyM [intConst int_to 1]) + (resultBodyM [intConst int_to 0]) E.Scalar (E.Prim (E.Signed int_from)) -> letTupExp' desc $ I.BasicOp $ I.ConvOp (I.ZExt int_from int_to) e' E.Scalar (E.Prim (E.Unsigned int_from)) -> @@ -1897,9 +1893,8 @@ "length of index and value array does not match" loc certifying c $ - letExp (baseString sv ++ "_write_sv") $ - I.BasicOp $ - I.Reshape (reshapeOuter [DimCoercion si_w] 1 sv_shape) sv + letExp (baseString sv ++ "_write_sv") . I.BasicOp $ + I.Reshape I.ReshapeCoerce (reshapeOuter (I.Shape [si_w]) 1 sv_shape) sv indexType <- fmap rowType <$> mapM lookupType si' indexName <- mapM (\_ -> newVName "write_index") indexType @@ -2139,18 +2134,14 @@ -- the total sizes, which are the last elements in the offests. We -- just have to be careful in case the array is empty. last_index <- letSubExp "last_index" $ I.BasicOp $ I.BinOp (I.Sub Int64 OverflowUndef) w $ constant (1 :: Int64) - nonempty_body <- runBodyBuilder $ - fmap resultBody $ - forM all_offsets $ \offset_array -> - letSubExp "last_offset" $ I.BasicOp $ I.Index offset_array $ Slice [I.DimFix last_index] - let empty_body = resultBody $ replicate k $ constant (0 :: Int64) + let nonempty_body = runBodyBuilder $ + fmap resultBody $ + forM all_offsets $ \offset_array -> + letSubExp "last_offset" $ I.BasicOp $ I.Index offset_array $ Slice [I.DimFix last_index] + empty_body = resultBodyM $ replicate k $ constant (0 :: Int64) is_empty <- letSubExp "is_empty" $ I.BasicOp $ I.CmpOp (CmpEq int64) w $ constant (0 :: Int64) sizes <- - letTupExp "partition_size" $ - I.If is_empty empty_body nonempty_body $ - ifCommon $ - replicate k $ - I.Prim int64 + letTupExp "partition_size" =<< eIf (eSubExp is_empty) empty_body nonempty_body -- The total size of all partitions must necessarily be equal to the -- size of the input array. @@ -2218,11 +2209,10 @@ (constant (-1 :: Int64)) (I.Var (I.paramName p) : take i sizes) letSubExp "total_res" - $ I.If - is_this_one - (resultBody [this_one]) - (resultBody [next_one]) - $ ifCommon [I.Prim int64] + =<< eIf + (eSubExp is_this_one) + (resultBodyM [this_one]) + (resultBodyM [next_one]) typeExpForError :: E.TypeExp VName -> InternaliseM [ErrorMsgPart SubExp] typeExpForError (E.TEVar qn _) = diff --git a/src/Futhark/LSP/Handlers.hs b/src/Futhark/LSP/Handlers.hs --- a/src/Futhark/LSP/Handlers.hs +++ b/src/Futhark/LSP/Handlers.hs @@ -49,7 +49,6 @@ logStringStderr <& "Got custom request: onFocusTextDocument" let NotificationMessage _ _ (Array vector_param) = msg String focused_uri = V.head vector_param -- only one parameter passed from the client - logStringStderr <& show focused_uri tryReCompile state_mvar (uriToFilePath (Uri focused_uri)) goToDefinitionHandler :: IORef State -> Handlers (LspM ()) diff --git a/src/Futhark/MonadFreshNames.hs b/src/Futhark/MonadFreshNames.hs --- a/src/Futhark/MonadFreshNames.hs +++ b/src/Futhark/MonadFreshNames.hs @@ -41,27 +41,27 @@ -- getNameSource = get -- putNameSource = put -- @ -class (Applicative m, Monad m) => MonadFreshNames m where +class Monad m => MonadFreshNames m where getNameSource :: m VNameSource putNameSource :: VNameSource -> m () -instance (Applicative im, Monad im) => MonadFreshNames (Control.Monad.State.Lazy.StateT VNameSource im) where +instance Monad im => MonadFreshNames (Control.Monad.State.Lazy.StateT VNameSource im) where getNameSource = Control.Monad.State.Lazy.get putNameSource = Control.Monad.State.Lazy.put -instance (Applicative im, Monad im) => MonadFreshNames (Control.Monad.State.Strict.StateT VNameSource im) where +instance Monad im => MonadFreshNames (Control.Monad.State.Strict.StateT VNameSource im) where getNameSource = Control.Monad.State.Strict.get putNameSource = Control.Monad.State.Strict.put instance - (Applicative im, Monad im, Monoid w) => + (Monad im, Monoid w) => MonadFreshNames (Control.Monad.RWS.Lazy.RWST r w VNameSource im) where getNameSource = Control.Monad.RWS.Lazy.get putNameSource = Control.Monad.RWS.Lazy.put instance - (Applicative im, Monad im, Monoid w) => + (Monad im, Monoid w) => MonadFreshNames (Control.Monad.RWS.Strict.RWST r w VNameSource im) where getNameSource = Control.Monad.RWS.Strict.get diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs --- a/src/Futhark/Optimise/BlkRegTiling.hs +++ b/src/Futhark/Optimise/BlkRegTiling.hs @@ -634,8 +634,9 @@ epilogue_t <- lookupType epilogue_res let (block_dims, rest_dims) = splitAt 2 $ arrayDims epilogue_t ones = map (const $ intConst Int64 1) rem_outer_dims - new_shape = concat [ones, block_dims, ones, rest_dims] - letExp "res_reshaped" $ BasicOp $ Reshape (map DimNew new_shape) epilogue_res + new_shape = Shape $ concat [ones, block_dims, ones, rest_dims] + letExp "res_reshaped" . BasicOp $ + Reshape ReshapeArbitrary new_shape epilogue_res pure [RegTileReturns mempty regtile_ret_dims epilogue_res'] mmBlkRegTilingNrm _ _ = pure Nothing @@ -1279,8 +1280,9 @@ res_tp' <- lookupType res let (block_dims, rest_dims) = splitAt 2 $ arrayDims res_tp' ones = map (const se1) rem_outer_dims - new_shape = concat [ones, block_dims, ones, rest_dims] - letExp "res_reshaped" $ BasicOp $ Reshape (map DimNew new_shape) res + new_shape = Shape $ concat [ones, block_dims, ones, rest_dims] + letExp "res_reshaped" . BasicOp $ + Reshape ReshapeArbitrary new_shape res pure $ map (RegTileReturns mempty regtile_ret_dims) epilogue_res' -- END (ret_seggroup, stms_seggroup) <- runBuilder $ do diff --git a/src/Futhark/Optimise/EntryPointMem.hs b/src/Futhark/Optimise/EntryPointMem.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/EntryPointMem.hs @@ -0,0 +1,81 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +-- | We require that entry points return arrays with zero offset in +-- row-major order. "Futhark.Pass.ExplicitAllocations" is +-- conservative and inserts copies to ensure this is the case. After +-- simplification, it may turn out that those copies are redundant. +-- This pass removes them. It's a pretty simple pass, as it only has +-- to look at the top level of entry points. +module Futhark.Optimise.EntryPointMem + ( entryPointMemGPU, + entryPointMemMC, + entryPointMemSeq, + ) +where + +import Data.List (find) +import qualified Data.Map.Strict as M +import Futhark.IR.GPUMem (GPUMem) +import Futhark.IR.MCMem (MCMem) +import Futhark.IR.Mem +import Futhark.IR.SeqMem (SeqMem) +import Futhark.Pass +import Futhark.Pass.ExplicitAllocations.GPU () +import Futhark.Transform.Substitute + +type Table rep = M.Map VName (Stm rep) + +mkTable :: Stms rep -> Table rep +mkTable = foldMap f + where + f stm = M.fromList $ zip (patNames (stmPat stm)) $ repeat stm + +varInfo :: Mem rep inner => VName -> Table rep -> Maybe (LetDecMem, Exp rep) +varInfo v table = do + Let pat _ e <- M.lookup v table + PatElem _ info <- find ((== v) . patElemName) (patElems pat) + Just (letDecMem info, e) + +optimiseFun :: Mem rep inner => Table rep -> FunDef rep -> FunDef rep +optimiseFun consts_table fd = + fd {funDefBody = onBody $ funDefBody fd} + where + table = consts_table <> mkTable (bodyStms (funDefBody fd)) + mkSubst (Var v0) + | Just (MemArray _ _ _ (ArrayIn mem0 ixfun0), BasicOp (Manifest _ v1)) <- + varInfo v0 table, + Just (MemArray _ _ _ (ArrayIn mem1 ixfun1), _) <- + varInfo v1 table, + ixfun0 == ixfun1 = + M.fromList [(mem0, mem1), (v0, v1)] + mkSubst _ = mempty + onBody (Body dec stms res) = + let substs = mconcat $ map (mkSubst . resSubExp) res + in Body dec stms $ substituteNames substs res + +entryPointMem :: Mem rep inner => Pass rep rep +entryPointMem = + Pass + { passName = "Entry point memory optimisation", + passDescription = "Remove redundant copies of entry point results.", + passFunction = intraproceduralTransformationWithConsts pure onFun + } + where + onFun consts fd = pure $ optimiseFun (mkTable consts) fd + +-- | The pass for GPU representation. +entryPointMemGPU :: Pass GPUMem GPUMem +entryPointMemGPU = entryPointMem + +-- | The pass for MC representation. +entryPointMemMC :: Pass MCMem MCMem +entryPointMemMC = entryPointMem + +-- | The pass for Seq representation. +entryPointMemSeq :: Pass SeqMem SeqMem +entryPointMemSeq = entryPointMem diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs --- a/src/Futhark/Optimise/Fusion.hs +++ b/src/Futhark/Optimise/Fusion.hs @@ -98,7 +98,7 @@ DoNode stm lst -> do lst' <- mapM (finalizeNode . fst) lst pure $ mconcat lst' <> oneStm stm - IfNode stm lst -> do + MatchNode stm lst -> do lst' <- mapM (finalizeNode . fst) lst pure $ mconcat lst' <> oneStm stm FinalNode stms1 nt' stms2 -> do @@ -253,10 +253,10 @@ -- First node is producer, second is consumer. vFuseNodeT :: [EdgeT] -> [VName] -> (NodeT, [EdgeT], [EdgeT]) -> (NodeT, [EdgeT]) -> FusionM (Maybe NodeT) -vFuseNodeT _ infusible (s1, _, e1s) (IfNode stm2 dfused, _) +vFuseNodeT _ infusible (s1, _, e1s) (MatchNode stm2 dfused, _) | isRealNode s1, null infusible = - pure $ Just $ IfNode stm2 $ (s1, e1s) : dfused + pure $ Just $ MatchNode stm2 $ (s1, e1s) : dfused vFuseNodeT _ infusible (StmNode stm1, _, _) (SoacNode ots2 pats2 soac2 aux2, _) | null infusible, [stm1_out] <- patNames $ stmPat stm1, @@ -398,11 +398,10 @@ doFuseScans . localScope (scopeOfFParams (map fst params) <> scopeOf form) $ do b <- doFusionWithDelayed body to_fuse pure (incoming, node, DoNode (Let pat aux (DoLoop params form b)) [], outgoing) - IfNode (Let pat aux (If sz b1 b2 dec)) to_fuse -> doFuseScans $ do - b1' <- doFusionWithDelayed b1 to_fuse - b2' <- doFusionWithDelayed b2 to_fuse - rb2' <- renameBody b2' - pure (incoming, node, IfNode (Let pat aux (If sz b1' rb2' dec)) [], outgoing) + MatchNode (Let pat aux (Match cond cases defbody dec)) to_fuse -> doFuseScans $ do + cases' <- mapM (traverse $ renameBody <=< (`doFusionWithDelayed` to_fuse)) cases + defbody' <- doFusionWithDelayed defbody to_fuse + pure (incoming, node, MatchNode (Let pat aux (Match cond cases' defbody' dec)) [], outgoing) StmNode (Let pat aux (Op (Futhark.VJP lam args vec))) -> doFuseScans $ do lam' <- doFusionLambda lam pure (incoming, node, StmNode (Let pat aux (Op (Futhark.VJP lam' args vec))), outgoing) diff --git a/src/Futhark/Optimise/Fusion/GraphRep.hs b/src/Futhark/Optimise/Fusion/GraphRep.hs --- a/src/Futhark/Optimise/Fusion/GraphRep.hs +++ b/src/Futhark/Optimise/Fusion/GraphRep.hs @@ -3,7 +3,8 @@ -- | A graph representation of a sequence of Futhark statements -- (i.e. a 'Body'), built to handle fusion. Could perhaps be made -- more general. An important property is that it does not handle --- "nested bodies" (e.g. 'If'); these are represented as single nodes. +-- "nested bodies" (e.g. 'Match'); these are represented as single +-- nodes. -- -- This is all implemented on top of the graph representation provided -- by the @fgl@ package ("Data.Graph.Inductive"). The graph provided @@ -83,7 +84,7 @@ -- Unclear whether we actually need these. FreeNode VName | FinalNode (Stms SOACS) NodeT (Stms SOACS) - | IfNode (Stm SOACS) [(NodeT, [EdgeT])] + | MatchNode (Stm SOACS) [(NodeT, [EdgeT])] | DoNode (Stm SOACS) [(NodeT, [EdgeT])] deriving (Eq) @@ -101,7 +102,7 @@ show (FinalNode _ nt _) = show nt show (ResNode name) = pretty $ "Res: " ++ pretty name show (FreeNode name) = pretty $ "Input: " ++ pretty name - show (IfNode stm _) = "If: " ++ L.intercalate ", " (map pretty $ stmNames stm) + show (MatchNode stm _) = "Match: " ++ L.intercalate ", " (map pretty $ stmNames stm) show (DoNode stm _) = "Do: " ++ L.intercalate ", " (map pretty $ stmNames stm) -- | The name that this edge depends on. @@ -290,8 +291,8 @@ Left H.NotSOAC -> pure n DoLoop {} -> pure $ DoNode s [] - If {} -> - pure $ IfNode s [] + Match {} -> + pure $ MatchNode s [] _ -> pure n nodeToSoacNode n = pure n @@ -375,8 +376,10 @@ bodyInputs (Body _ stms res) = foldMap stmInputs stms <> freeClassifications res expInputs :: Exp SOACS -> Classifications -expInputs (If cond b1 b2 attr) = - bodyInputs b1 <> bodyInputs b2 <> freeClassifications (cond, attr) +expInputs (Match cond cases defbody attr) = + foldMap (bodyInputs . caseBody) cases + <> bodyInputs defbody + <> freeClassifications (cond, attr) expInputs (DoLoop params form b1) = freeClassifications (params, form) <> bodyInputs b1 expInputs (Op soac) = case soac of @@ -407,7 +410,7 @@ (StmNode stm) -> stmNames stm (ResNode _) -> [] (FreeNode name) -> [name] - (IfNode stm _) -> stmNames stm + (MatchNode stm _) -> stmNames stm (DoNode stm _) -> stmNames stm FinalNode {} -> error "Final nodes cannot generate edges" (SoacNode _ pat _ _) -> patNames pat diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs --- a/src/Futhark/Optimise/Fusion/TryFusion.hs +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -746,16 +746,16 @@ pullReshape :: SOAC -> SOAC.ArrayTransforms -> TryFusion (SOAC, SOAC.ArrayTransforms) pullReshape (SOAC.Screma _ form inps) ots | Just maplam <- Futhark.isMapSOAC form, - SOAC.Reshape cs shape SOAC.:< ots' <- SOAC.viewf ots, + SOAC.Reshape cs k shape SOAC.:< ots' <- SOAC.viewf ots, all primType $ lambdaReturnType maplam = do - let mapw' = case reverse $ newDims shape of + let mapw' = case reverse $ shapeDims shape of [] -> intConst Int64 0 d : _ -> d trInput inp | arrayRank (SOAC.inputType inp) == 1 = - SOAC.addTransform (SOAC.Reshape cs shape) inp + SOAC.addTransform (SOAC.Reshape cs k shape) inp | otherwise = - SOAC.addTransform (SOAC.ReshapeOuter cs shape) inp + SOAC.addTransform (SOAC.ReshapeOuter cs k shape) inp inputs' = map trInput inps inputTypes = map SOAC.inputType inputs' @@ -784,12 +784,9 @@ op' <- foldM outersoac (SOAC.Screma mapw' $ Futhark.mapSOAC maplam) $ - zip (drop 1 $ reverse $ newDims shape) $ - drop 1 $ - reverse $ - drop 1 $ - tails $ - newDims shape + zip (drop 1 $ reverse $ shapeDims shape) $ + drop 1 . reverse . drop 1 . tails $ + shapeDims shape pure (op' inputs', ots') pullReshape _ _ = fail "Cannot pull reshape" diff --git a/src/Futhark/Optimise/GenRedOpt.hs b/src/Futhark/Optimise/GenRedOpt.hs --- a/src/Futhark/Optimise/GenRedOpt.hs +++ b/src/Futhark/Optimise/GenRedOpt.hs @@ -414,8 +414,8 @@ costRedundantStmt (Let _ _ DoLoop {}) = Big costRedundantStmt (Let _ _ Apply {}) = Big costRedundantStmt (Let _ _ WithAcc {}) = Big -costRedundantStmt (Let _ _ (If _cond b_then b_else _)) = - maxCost (costBody b_then) (costBody b_else) +costRedundantStmt (Let _ _ (Match _ cases defbody _)) = + L.foldl' maxCost (costBody defbody) $ map (costBody . caseBody) cases costRedundantStmt (Let _ _ (BasicOp (ArrayLit _ Array {}))) = Big costRedundantStmt (Let _ _ (BasicOp (ArrayLit _ _))) = Small 1 costRedundantStmt (Let _ _ (BasicOp (Index _ slc))) = diff --git a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs --- a/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs +++ b/src/Futhark/Optimise/InPlaceLowering/SubstituteIndices.hs @@ -37,8 +37,7 @@ ( MonadFreshNames m, BuilderOps rep, Buildable rep, - Aliased rep, - LParamInfo rep ~ Type + Aliased rep ) => IndexSubstitutions -> Stms rep -> diff --git a/src/Futhark/Optimise/MemoryBlockMerging.hs b/src/Futhark/Optimise/MemoryBlockMerging.hs --- a/src/Futhark/Optimise/MemoryBlockMerging.hs +++ b/src/Futhark/Optimise/MemoryBlockMerging.hs @@ -30,9 +30,8 @@ getAllocsStm (Let (Pat [PatElem name _]) _ (Op (Alloc se sp))) = M.singleton name (se, sp) getAllocsStm (Let _ _ (Op (Alloc _ _))) = error "impossible" -getAllocsStm (Let _ _ (If _ then_body else_body _)) = - foldMap getAllocsStm (bodyStms then_body) - <> foldMap getAllocsStm (bodyStms else_body) +getAllocsStm (Let _ _ (Match _ cases defbody _)) = + foldMap (foldMap getAllocsStm . bodyStms) $ defbody : map caseBody cases getAllocsStm (Let _ _ (DoLoop _ _ body)) = foldMap getAllocsStm (bodyStms body) getAllocsStm _ = mempty @@ -54,15 +53,10 @@ setAllocsStm _ stm@(Let _ _ (Op (Alloc _ _))) = stm setAllocsStm m stm@(Let _ _ (Op (Inner (SegOp segop)))) = stm {stmExp = Op $ Inner $ SegOp $ setAllocsSegOp m segop} -setAllocsStm m stm@(Let _ _ (If cse then_body else_body dec)) = - stm - { stmExp = - If - cse - (then_body {bodyStms = setAllocsStm m <$> bodyStms then_body}) - (else_body {bodyStms = setAllocsStm m <$> bodyStms else_body}) - dec - } +setAllocsStm m stm@(Let _ _ (Match cond cases defbody dec)) = + stm {stmExp = Match cond (map (fmap onBody) cases) (onBody defbody) dec} + where + onBody (Body () stms res) = Body () (setAllocsStm m <$> stms) res setAllocsStm m stm@(Let _ _ (DoLoop merge form body)) = stm { stmExp = @@ -172,23 +166,18 @@ (SegOp SegLevel GPUMem -> m (SegOp SegLevel GPUMem)) -> Stms GPUMem -> m (Stms GPUMem) -onKernels f stms = inScopeOf stms $ mapM helper stms +onKernels f orig_stms = inScopeOf orig_stms $ mapM helper orig_stms where helper stm@Let {stmExp = Op (Inner (SegOp segop))} = do exp' <- f segop pure $ stm {stmExp = Op $ Inner $ SegOp exp'} - helper stm@Let {stmExp = If c then_body else_body dec} = do - then_body_stms <- f `onKernels` bodyStms then_body - else_body_stms <- f `onKernels` bodyStms else_body - pure $ - stm - { stmExp = - If - c - (then_body {bodyStms = then_body_stms}) - (else_body {bodyStms = else_body_stms}) - dec - } + helper stm@Let {stmExp = Match c cases defbody dec} = do + cases' <- mapM (traverse onBody) cases + defbody' <- onBody defbody + pure $ stm {stmExp = Match c cases' defbody' dec} + where + onBody (Body () stms res) = + Body () <$> f `onKernels` stms <*> pure res helper stm@Let {stmExp = DoLoop merge form body} = do body_stms <- f `onKernels` bodyStms body pure $ stm {stmExp = DoLoop merge form (body {bodyStms = body_stms})} diff --git a/src/Futhark/Optimise/MergeGPUBodies.hs b/src/Futhark/Optimise/MergeGPUBodies.hs --- a/src/Futhark/Optimise/MergeGPUBodies.hs +++ b/src/Futhark/Optimise/MergeGPUBodies.hs @@ -12,6 +12,7 @@ import Control.Monad import Control.Monad.Trans.Class import Control.Monad.Trans.State.Strict hiding (State) +import Data.Bifunctor (first) import Data.Foldable import qualified Data.IntMap as IM import Data.IntSet ((\\)) @@ -173,11 +174,13 @@ case e of BasicOp {} -> pure (removeExpAliases e, depsOf e) Apply {} -> pure (removeExpAliases e, depsOf e) - If c tbody fbody dec -> do - (tbody', t_deps) <- transformBody aliases tbody - (fbody', f_deps) <- transformBody aliases fbody - let deps = depsOf c <> t_deps <> f_deps <> depsOf dec - pure (If c tbody' fbody' dec, deps) + Match ses cases defbody dec -> do + let transformCase (Case vs body) = + first (Case vs) <$> transformBody aliases body + (cases', cases_deps) <- unzip <$> mapM transformCase cases + (defbody', defbody_deps) <- transformBody aliases defbody + let deps = depsOf ses <> mconcat cases_deps <> defbody_deps <> depsOf dec + pure (Match ses cases' defbody' dec, deps) DoLoop merge lform body -> do -- What merge and lform aliases outside the loop is irrelevant as those -- cannot be consumed within the loop. diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs.hs b/src/Futhark/Optimise/ReduceDeviceSyncs.hs --- a/src/Futhark/Optimise/ReduceDeviceSyncs.hs +++ b/src/Futhark/Optimise/ReduceDeviceSyncs.hs @@ -15,11 +15,11 @@ import Data.Bifunctor (second) import Data.Foldable import qualified Data.IntMap.Strict as IM -import Data.List (unzip4, zip4) +import Data.List (transpose) import qualified Data.Map.Strict as M import Data.Sequence ((<|), (><), (|>)) import qualified Data.Text as T -import Futhark.Construct (fullSlice, sliceDim) +import Futhark.Construct (fullSlice, mkBody, sliceDim) import Futhark.Error import Futhark.IR.GPU import Futhark.MonadFreshNames @@ -149,50 +149,53 @@ pure (out |> stm) Apply {} -> pure (out |> stm) - If cond (Body _ tstms0 tres) (Body _ fstms0 fres) (IfDec btypes sort) -> - do - -- Rewrite branches. - tstms1 <- optimizeStms tstms0 - fstms1 <- optimizeStms fstms0 + Match ses cases defbody (MatchDec btypes sort) -> do + -- Rewrite branches. + cases_stms <- mapM (optimizeStms . bodyStms . caseBody) cases + let cases_res = map (bodyResult . caseBody) cases + defbody_stms <- optimizeStms $ bodyStms defbody + let defbody_res = bodyResult defbody - -- Ensure return values and types match if one or both branches - -- return a result that now reside on device. - let bmerge (res, tstms, fstms) (pe, tr, fr, bt) = - do - let onHost (Var v) = (v ==) <$> resolveName v - onHost _ = pure True + -- Ensure return values and types match if one or both branches + -- return a result that now reside on device. + let bmerge (acc, all_stms) (pe, reses, bt) = do + let onHost (Var v) = (v ==) <$> resolveName v + onHost _ = pure True - tr_on_host <- onHost (resSubExp tr) - fr_on_host <- onHost (resSubExp fr) + on_host <- and <$> mapM (onHost . resSubExp) reses - if tr_on_host && fr_on_host - then -- No result resides on device ==> nothing to do. - pure ((pe, tr, fr, bt) : res, tstms, fstms) - else -- Otherwise, ensure both results are migrated. - do - let t = patElemType pe - (tstms', tarr) <- storeScalar tstms (resSubExp tr) t - (fstms', farr) <- storeScalar fstms (resSubExp fr) t + if on_host + then -- No result resides on device ==> nothing to do. + pure ((pe, reses, bt) : acc, all_stms) + else do + -- Otherwise, ensure all results are migrated. + (all_stms', arrs) <- + fmap unzip $ forM (zip all_stms reses) $ \(stms, res) -> + storeScalar stms (resSubExp res) (patElemType pe) - pe' <- arrayizePatElem pe - let bt' = staticShapes1 (patElemType pe') - let tr' = tr {resSubExp = Var tarr} - let fr' = fr {resSubExp = Var farr} - pure ((pe', tr', fr', bt') : res, tstms', fstms') + pe' <- arrayizePatElem pe + let bt' = staticShapes1 (patElemType pe') + reses' = zipWith SubExpRes (map resCerts reses) (map Var arrs) + pure ((pe', reses', bt') : acc, all_stms') - let pes = patElems (stmPat stm) - let zipped = zip4 pes tres fres btypes - (zipped', tstms2, fstms2) <- foldM bmerge ([], tstms1, fstms1) zipped - let (pes', tres', fres', btypes') = unzip4 (reverse zipped') + pes = patElems (stmPat stm) + (acc, ~(defbody_stms' : cases_stms')) <- + foldM bmerge ([], defbody_stms : cases_stms) $ + zip3 pes (transpose $ defbody_res : cases_res) btypes + let (pes', reses, btypes') = unzip3 (reverse acc) - -- Rewrite statement. - let tbranch' = Body () tstms2 tres' - let fbranch' = Body () fstms2 fres' - let e' = If cond tbranch' fbranch' (IfDec btypes' sort) - let stm' = Let (Pat pes') (stmAux stm) e' + -- Rewrite statement. + let cases' = + zipWith Case (map casePat cases) $ + zipWith mkBody cases_stms' $ + drop 1 $ + transpose reses + defbody' = mkBody defbody_stms' $ map head reses + e' = Match ses cases' defbody' (MatchDec btypes' sort) + stm' = Let (Pat pes') (stmAux stm) e' - -- Read migrated scalars that are used on host. - foldM addRead (out |> stm') (zip pes pes') + -- Read migrated scalars that are used on host. + foldM addRead (out |> stm') (zip pes pes') DoLoop ps lf b -> do -- Enable the migration of for-in loop variables. (params, lform, body) <- rewriteForIn (ps, lf, b) diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs --- a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +++ b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs @@ -62,6 +62,7 @@ import Data.Foldable import qualified Data.IntMap.Strict as IM import qualified Data.IntSet as IS +import qualified Data.List as L import qualified Data.Map.Strict as M import Data.Maybe (fromJust, fromMaybe, isJust, isNothing) import qualified Data.Sequence as SQ @@ -124,8 +125,8 @@ statusOf n mt /= StayOnHost shouldMoveStm (Let (Pat ((PatElem n _) : _)) _ Apply {}) mt = statusOf n mt /= StayOnHost -shouldMoveStm (Let _ _ (If (Var n) _ _ _)) mt = - statusOf n mt == MoveToDevice +shouldMoveStm (Let _ _ (Match cond _ _ _)) mt = + all ((== MoveToDevice) . (`statusOf` mt)) $ subExpVars cond shouldMoveStm (Let _ _ (DoLoop _ (ForLoop _ _ (Var n) _) _)) mt = statusOf n mt == MoveToDevice shouldMoveStm (Let _ _ (DoLoop _ (WhileLoop n) _)) mt = @@ -213,15 +214,13 @@ checkExp (WithAcc _ _) = hostOnly checkExp (Op _) = hostOnly checkExp (Apply fn _ _ _) = Just (S.singleton fn) - checkExp (If _ tbranch fbranch _) = do - calls1 <- checkBody tbranch - calls2 <- checkBody fbranch - pure (calls1 <> calls2) + checkExp (Match _ cases defbody _) = + mconcat <$> mapM checkBody (defbody : map caseBody cases) checkExp (DoLoop params lform body) = do checkLParams params checkLoopForm lform checkBody body - checkExp _ = Just S.empty + checkExp BasicOp {} = Just S.empty -------------------------------------------------------------------------------- -- MIGRATION ANALYSIS -- @@ -467,8 +466,8 @@ -- Can be replaced with 'graphHostOnly e' to disable migration. -- A fix can be verified by enabling tests/migration/reuse4_scratch.fut graphInefficientReturn s e - BasicOp (Reshape s arr) -> do - graphInefficientReturn (newDims s) e + BasicOp (Reshape _ s arr) -> do + graphInefficientReturn (shapeDims s) e one bs `reuses` arr BasicOp (Rearrange _ arr) -> do graphInefficientReturn [] e @@ -515,8 +514,8 @@ graphUpdateAcc (one bs) e Apply fn _ _ _ -> graphApply fn bs e - If cond tbody fbody _ -> - graphIf bs cond tbody fbody + Match ses cases defbody _ -> + graphMatch bs ses cases defbody DoLoop params lform body -> graphLoop bs params lform body WithAcc inputs f -> @@ -609,28 +608,29 @@ then graphHostOnly e else graphSimple bs e --- | Graph an if statement. -graphIf :: [Binding] -> SubExp -> Body GPU -> Body GPU -> Grapher () -graphIf bs cond tbody fbody = do +-- | Graph a Match statement. +graphMatch :: [Binding] -> [SubExp] -> [Case (Body GPU)] -> Body GPU -> Grapher () +graphMatch bs ses cases defbody = do body_host_only <- - incForkDepthFor - ( do - tstats <- captureBodyStats (graphBody tbody) - fstats <- captureBodyStats (graphBody fbody) - pure $ bodyHostOnly tstats || bodyHostOnly fstats - ) + incForkDepthFor $ + any bodyHostOnly + <$> mapM (captureBodyStats . graphBody) (defbody : map caseBody cases) + let branch_results = results defbody : map (results . caseBody) cases + -- Record aliases for copyable memory backing returned arrays. - may_copy_results <- reusesBranches bs (results tbody) (results fbody) + may_copy_results <- reusesBranches bs branch_results let may_migrate = not body_host_only && may_copy_results - cond_id <- case (may_migrate, cond) of - (False, Var n) -> - -- The migration status of the condition is what determines whether the - -- statement may be migrated as a whole or not. See 'shouldMoveStm'. - connectToSink (nameToId n) >> pure IS.empty - (True, Var n) -> onlyGraphedScalar n - (_, _) -> pure IS.empty + cond_id <- + if may_migrate + then onlyGraphedScalars $ subExpVars ses + else do + -- The migration status of the condition is what determines + -- whether the statement may be migrated as a whole or + -- not. See 'shouldMoveStm'. + mapM_ (connectToSink . nameToId) (subExpVars ses) + pure IS.empty tellOperands cond_id @@ -659,15 +659,12 @@ -- of host-device reads it means that some reads may needlessly be delayed -- out of branches. The overhead as measured on futhark-benchmarks appears -- to be neglible though. - ret <- zipWithM (comb cond_id) (bodyResult tbody) (bodyResult fbody) + ret <- mapM (comb cond_id) $ L.transpose branch_results mapM_ (uncurry createNode) (zip bs ret) where results = map resSubExp . bodyResult - comb ci a b = (ci <>) <$> onlyGraphedScalars (toSet a <> toSet b) - - toSet (SubExpRes _ (Var n)) = S.singleton n - toSet _ = S.empty + comb ci a = (ci <>) <$> onlyGraphedScalars (S.fromList $ subExpVars a) ----------------------------------------------------- -- These type aliases are only used by 'graphLoop' -- @@ -702,7 +699,7 @@ -- Does the loop return any arrays which prevent it from being migrated? let args = map snd params let results = map resSubExp (bodyResult body) - may_copy_results <- reusesBranches (b : bs) args results + may_copy_results <- reusesBranches (b : bs) [args, results] -- Connect loop condition to a sink if the loop cannot be migrated. -- The migration status of the condition is what determines whether the @@ -1017,8 +1014,10 @@ collectBasic b collect (Apply _ params _ _) = mapM_ (collectSubExp . fst) params - collect (If cond tbranch fbranch _) = - collectSubExp cond >> collectBody tbranch >> collectBody fbranch + collect (Match ses cases defbody _) = do + mapM_ collectSubExp ses + mapM_ (collectBody . caseBody) cases + collectBody defbody collect (DoLoop params lform body) = do mapM_ (collectSubExp . snd) params collectLForm lform @@ -1257,19 +1256,21 @@ | otherwise = pure onlyCopyable --- @reusesBranches bs b1 b2@ records each array binding in @bs@ as reusing --- copyable memory if each corresponding return value in the lists @b1@ and @b2@ --- are backed by copyable memory. +-- @reusesBranches bs seses@ records each array binding in @bs@ as +-- reusing copyable memory if each corresponding return value in the +-- lists in @ses@ are backed by copyable memory. Each list is the +-- result of a branch body (i.e. for 'if' the list has two elements). -- --- If every array binding is registered as being backed by copyable memory then --- the function returns @True@, otherwise it returns @False@. -reusesBranches :: [Binding] -> [SubExp] -> [SubExp] -> Grapher Bool -reusesBranches bs b1 b2 = do +-- If every array binding is registered as being backed by copyable +-- memory then the function returns @True@, otherwise it returns +-- @False@. +reusesBranches :: [Binding] -> [[SubExp]] -> Grapher Bool +reusesBranches bs seses = do body_depth <- metaBodyDepth <$> getMeta - foldM (reuse body_depth) True $ zip3 bs b1 b2 + foldM (reuse body_depth) True $ zip bs $ L.transpose seses where - reuse :: Int -> Bool -> (Binding, SubExp, SubExp) -> Grapher Bool - reuse body_depth onlyCopyable (b, se1, se2) + reuse :: Int -> Bool -> (Binding, [SubExp]) -> Grapher Bool + reuse body_depth onlyCopyable (b, ses) | all (== intConst Int64 1) (arrayDims $ snd b) = -- Single element arrays are immediately recognizable as copyable so -- don't bother recording those. Note that this case also matches @@ -1277,19 +1278,16 @@ pure onlyCopyable | (i, t) <- b, isArray t, - Var n1 <- se1, - Var n2 <- se2 = - do - body_depth_1 <- outermostCopyableArray n1 - body_depth_2 <- outermostCopyableArray n2 - case (body_depth_1, body_depth_2) of - (Just bd1, Just bd2) -> do - let inner = min bd1 bd2 - recordCopyableMemory i (min body_depth inner) - let returns_free_var = inner <= body_depth - pure (onlyCopyable && not returns_free_var) - _ -> - pure False + Just ns <- mapM subExpVar ses = do + body_depths <- mapM outermostCopyableArray ns + case sequence body_depths of + Just bds -> do + let inner = minimum bds + recordCopyableMemory i (min body_depth inner) + let returns_free_var = inner <= body_depth + pure (onlyCopyable && not returns_free_var) + _ -> + pure False | otherwise = pure onlyCopyable diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -72,7 +72,7 @@ import Control.Monad.Reader import Control.Monad.State.Strict import Data.Either -import Data.List (find, foldl', mapAccumL) +import Data.List (find, foldl', inits, mapAccumL) import qualified Data.Map as M import Data.Maybe import qualified Futhark.Analysis.SymbolTable as ST @@ -117,7 +117,7 @@ } -- | A function that protects a hoisted operation (if possible). The --- first operand is the condition of the 'If' we have hoisted out of +-- first operand is the condition of the 'Case' we have hoisted out of -- (or equivalently, a boolean indicating whether a loop has nonzero -- trip count). type Protect m = SubExp -> Pat (LetDec (Rep m)) -> Op (Rep m) -> Maybe (m ()) @@ -140,6 +140,10 @@ -- actually be used. protectHoistedOpS :: Protect (Builder (Wise rep)), opUsageS :: Op (Wise rep) -> UT.UsageTable, + simplifyPatFromExpS :: + Pat (LetDec rep) -> + Exp (Wise rep) -> + SimpleM rep (Pat (LetDec rep)), simplifyOpS :: SimplifyOp rep (Op (Wise rep)) } @@ -148,11 +152,12 @@ SimplifyOp rep (Op (Wise rep)) -> SimpleOps rep bindableSimpleOps = - SimpleOps mkExpDecS' mkBodyS' protectHoistedOpS' (const mempty) + SimpleOps mkExpDecS' mkBodyS' protectHoistedOpS' (const mempty) simplifyPatFromExp where mkExpDecS' _ pat e = pure $ mkExpDec pat e mkBodyS' _ stms res = pure $ mkBody stms res protectHoistedOpS' _ _ _ = Nothing + simplifyPatFromExp pat _ = traverse simplify pat newtype SimpleM rep a = SimpleM @@ -262,37 +267,80 @@ bindLoopVar var it bound = localVtable $ ST.insertLoopVar var it bound --- | We are willing to hoist potentially unsafe statements out of --- branches, but they most be protected by adding a branch on top of --- them. (This means such hoisting is not worth it unless they are in --- turn hoisted out of a loop somewhere.) -protectIfHoisted :: - SimplifiableRep rep => - -- | Branch condition. - SubExp -> - -- | Which side of the branch are we - -- protecting here? - Bool -> - SimpleM rep (Stms (Wise rep), a) -> - SimpleM rep (Stms (Wise rep), a) -protectIfHoisted cond side m = do - (hoisted, x) <- m - ops <- asks $ protectHoistedOpS . fst - hoisted' <- runBuilder_ $ do - if not $ all (safeExp . stmExp) hoisted - then do - cond' <- - if side - then pure cond - else letSubExp "cond_neg" $ BasicOp $ UnOp Not cond - mapM_ (protectIf ops unsafeOrCostly cond') hoisted - else addStms hoisted - pure (hoisted', x) +makeSafe :: Exp rep -> Maybe (Exp rep) +makeSafe (BasicOp (BinOp (SDiv t _) x y)) = + Just $ BasicOp (BinOp (SDiv t Safe) x y) +makeSafe (BasicOp (BinOp (SDivUp t _) x y)) = + Just $ BasicOp (BinOp (SDivUp t Safe) x y) +makeSafe (BasicOp (BinOp (SQuot t _) x y)) = + Just $ BasicOp (BinOp (SQuot t Safe) x y) +makeSafe (BasicOp (BinOp (UDiv t _) x y)) = + Just $ BasicOp (BinOp (UDiv t Safe) x y) +makeSafe (BasicOp (BinOp (UDivUp t _) x y)) = + Just $ BasicOp (BinOp (UDivUp t Safe) x y) +makeSafe (BasicOp (BinOp (SMod t _) x y)) = + Just $ BasicOp (BinOp (SMod t Safe) x y) +makeSafe (BasicOp (BinOp (SRem t _) x y)) = + Just $ BasicOp (BinOp (SRem t Safe) x y) +makeSafe (BasicOp (BinOp (UMod t _) x y)) = + Just $ BasicOp (BinOp (UMod t Safe) x y) +makeSafe _ = + Nothing + +emptyOfType :: MonadBuilder m => [VName] -> Type -> m (Exp (Rep m)) +emptyOfType _ Mem {} = + error "emptyOfType: Cannot hoist non-existential memory." +emptyOfType _ Acc {} = + error "emptyOfType: Cannot hoist accumulator." +emptyOfType _ (Prim pt) = + pure $ BasicOp $ SubExp $ Constant $ blankPrimValue pt +emptyOfType ctx_names (Array et shape _) = do + let dims = map zeroIfContext $ shapeDims shape + pure $ BasicOp $ Scratch et dims where - unsafeOrCostly e = not (safeExp e) || not (cheapExp e) + zeroIfContext (Var v) | v `elem` ctx_names = intConst Int64 0 + zeroIfContext se = se +protectIf :: + MonadBuilder m => + Protect m -> + (Exp (Rep m) -> Bool) -> + SubExp -> + Stm (Rep m) -> + m () +protectIf _ _ taken (Let pat aux (Match [cond] [Case [Just (BoolValue True)] taken_body] untaken_body (MatchDec if_ts MatchFallback))) = do + cond' <- letSubExp "protect_cond_conj" $ BasicOp $ BinOp LogAnd taken cond + auxing aux . letBind pat $ + Match [cond'] [Case [Just (BoolValue True)] taken_body] untaken_body $ + MatchDec if_ts MatchFallback +protectIf _ _ taken (Let pat aux (BasicOp (Assert cond msg loc))) = do + not_taken <- letSubExp "loop_not_taken" $ BasicOp $ UnOp Not taken + cond' <- letSubExp "protect_assert_disj" $ BasicOp $ BinOp LogOr not_taken cond + auxing aux $ letBind pat $ BasicOp $ Assert cond' msg loc +protectIf protect _ taken (Let pat aux (Op op)) + | Just m <- protect taken pat op = + auxing aux m +protectIf _ f taken (Let pat aux e) + | f e = + case makeSafe e of + Just e' -> + auxing aux $ letBind pat e' + Nothing -> do + taken_body <- eBody [pure e] + untaken_body <- + eBody $ map (emptyOfType $ patNames pat) (patTypes pat) + if_ts <- expTypesFromPat pat + auxing aux . letBind pat + $ Match + [taken] + [Case [Just $ BoolValue True] taken_body] + untaken_body + $ MatchDec if_ts MatchFallback +protectIf _ _ _ stm = + addStm stm + -- | We are willing to hoist potentially unsafe statements out of --- loops, but they most be protected by adding a branch on top of +-- loops, but they must be protected by adding a branch on top of -- them. protectLoopHoisted :: SimplifiableRep rep => @@ -323,74 +371,61 @@ BasicOp $ CmpOp (CmpSlt it) (intConst it 0) bound -protectIf :: - MonadBuilder m => - Protect m -> - (Exp (Rep m) -> Bool) -> - SubExp -> - Stm (Rep m) -> - m () -protectIf _ _ taken (Let pat aux (If cond taken_body untaken_body (IfDec if_ts IfFallback))) = do - cond' <- letSubExp "protect_cond_conj" $ BasicOp $ BinOp LogAnd taken cond - auxing aux . letBind pat $ - If cond' taken_body untaken_body $ - IfDec if_ts IfFallback -protectIf _ _ taken (Let pat aux (BasicOp (Assert cond msg loc))) = do - not_taken <- letSubExp "loop_not_taken" $ BasicOp $ UnOp Not taken - cond' <- letSubExp "protect_assert_disj" $ BasicOp $ BinOp LogOr not_taken cond - auxing aux $ letBind pat $ BasicOp $ Assert cond' msg loc -protectIf protect _ taken (Let pat aux (Op op)) - | Just m <- protect taken pat op = - auxing aux m -protectIf _ f taken (Let pat aux e) - | f e = - case makeSafe e of - Just e' -> - auxing aux $ letBind pat e' - Nothing -> do - taken_body <- eBody [pure e] - untaken_body <- - eBody $ map (emptyOfType $ patNames pat) (patTypes pat) - if_ts <- expTypesFromPat pat - auxing aux . letBind pat $ - If taken taken_body untaken_body $ - IfDec if_ts IfFallback -protectIf _ _ _ stm = - addStm stm +-- Produces a true subexpression if the pattern (as in a 'Case') +-- matches the subexpression. +matching :: + BuilderOps rep => + [(SubExp, Maybe PrimValue)] -> + Builder rep SubExp +matching = letSubExp "match" <=< eAll <=< sequence . mapMaybe cmp + where + cmp (se, Just (BoolValue True)) = + Just $ pure se + cmp (se, Just v) = + Just . letSubExp "match_val" . BasicOp $ + CmpOp (CmpEq (primValueType v)) se (Constant v) + cmp (_, Nothing) = Nothing -makeSafe :: Exp rep -> Maybe (Exp rep) -makeSafe (BasicOp (BinOp (SDiv t _) x y)) = - Just $ BasicOp (BinOp (SDiv t Safe) x y) -makeSafe (BasicOp (BinOp (SDivUp t _) x y)) = - Just $ BasicOp (BinOp (SDivUp t Safe) x y) -makeSafe (BasicOp (BinOp (SQuot t _) x y)) = - Just $ BasicOp (BinOp (SQuot t Safe) x y) -makeSafe (BasicOp (BinOp (UDiv t _) x y)) = - Just $ BasicOp (BinOp (UDiv t Safe) x y) -makeSafe (BasicOp (BinOp (UDivUp t _) x y)) = - Just $ BasicOp (BinOp (UDivUp t Safe) x y) -makeSafe (BasicOp (BinOp (SMod t _) x y)) = - Just $ BasicOp (BinOp (SMod t Safe) x y) -makeSafe (BasicOp (BinOp (SRem t _) x y)) = - Just $ BasicOp (BinOp (SRem t Safe) x y) -makeSafe (BasicOp (BinOp (UMod t _) x y)) = - Just $ BasicOp (BinOp (UMod t Safe) x y) -makeSafe _ = - Nothing +matchingExactlyThis :: + BuilderOps rep => + [SubExp] -> + [[Maybe PrimValue]] -> + [Maybe PrimValue] -> + Builder rep SubExp +matchingExactlyThis ses prior this = do + prior_matches <- mapM (matching . zip ses) prior + letSubExp "matching_just_this" + =<< eBinOp + LogAnd + (eUnOp Not (eAny prior_matches)) + (eSubExp =<< matching (zip ses this)) -emptyOfType :: MonadBuilder m => [VName] -> Type -> m (Exp (Rep m)) -emptyOfType _ Mem {} = - error "emptyOfType: Cannot hoist non-existential memory." -emptyOfType _ Acc {} = - error "emptyOfType: Cannot hoist accumulator." -emptyOfType _ (Prim pt) = - pure $ BasicOp $ SubExp $ Constant $ blankPrimValue pt -emptyOfType ctx_names (Array et shape _) = do - let dims = map zeroIfContext $ shapeDims shape - pure $ BasicOp $ Scratch et dims +-- | We are willing to hoist potentially unsafe statements out of +-- matches, but they must be protected by adding a branch on top of +-- them. (This means such hoisting is not worth it unless they are in +-- turn hoisted out of a loop somewhere.) +protectCaseHoisted :: + SimplifiableRep rep => + -- | Scrutinee. + [SubExp] -> + -- | Pattern of previosu cases. + [[Maybe PrimValue]] -> + -- | Pattern of this case. + [Maybe PrimValue] -> + SimpleM rep (Stms (Wise rep), a) -> + SimpleM rep (Stms (Wise rep), a) +protectCaseHoisted ses prior vs m = do + (hoisted, x) <- m + ops <- asks $ protectHoistedOpS . fst + hoisted' <- runBuilder_ $ do + if not $ all (safeExp . stmExp) hoisted + then do + cond' <- matchingExactlyThis ses prior vs + mapM_ (protectIf ops unsafeOrCostly cond') hoisted + else addStms hoisted + pure (hoisted', x) where - zeroIfContext (Var v) | v `elem` ctx_names = intConst Int64 0 - zeroIfContext se = se + unsafeOrCostly e = not (safeExp e) || not (cheapExp e) -- | Statements that are not worth hoisting out of loops, because they -- are unsafe, and added safety (by 'protectLoopHoisted') may inhibit @@ -407,9 +442,11 @@ SimpleM rep (Stm (Wise rep)) nonrecSimplifyStm (Let pat (StmAux cs attrs (_, dec)) e) = do cs' <- simplify cs - (pat', pat_cs) <- collectCerts $ simplifyPat $ removePatWisdom pat + e' <- simplifyExpBase e + simplifyPat <- asks $ simplifyPatFromExpS . fst + (pat', pat_cs) <- collectCerts $ simplifyPat (removePatWisdom pat) e' let aux' = StmAux (cs' <> pat_cs) attrs dec - mkWiseStm pat' aux' <$> simplifyExpBase e + pure $ mkWiseStm pat' aux' e' -- Bottom-up simplify a statement. Recurses into sub-Bodies and Ops. -- Does not copy-propagate into the pattern and similar, as it is @@ -623,9 +660,9 @@ cheapExp (BasicOp Concat {}) = False cheapExp (BasicOp Manifest {}) = False cheapExp DoLoop {} = False -cheapExp (If _ tbranch fbranch _) = - all cheapStm (bodyStms tbranch) - && all cheapStm (bodyStms fbranch) +cheapExp (Match _ cases defbranch _) = + all (all cheapStm . bodyStms . caseBody) cases + && all cheapStm (bodyStms defbranch) cheapExp (Op op) = cheapOp op cheapExp _ = True -- Used to be False, but -- let's try it out. @@ -634,21 +671,12 @@ loopInvariantStm vtable = all (`nameIn` ST.availableAtClosestLoop vtable) . namesToList . freeIn -hoistCommon :: - SimplifiableRep rep => - UT.UsageTable -> - [UT.Usages] -> - SubExp -> - IfDec (BranchType rep) -> - Body (Wise rep) -> - Body (Wise rep) -> - SimpleM - rep - ( Body (Wise rep), - Body (Wise rep), - Stms (Wise rep) - ) -hoistCommon res_usage res_usages cond (IfDec _ ifsort) body1 body2 = do +matchBlocker :: + (ASTRep rep, CanBeWise (Op rep), FreeIn a) => + a -> + MatchDec rt -> + SimpleM rep (BlockPred (Wise rep)) +matchBlocker cond (MatchDec _ ifsort) = do is_alloc_fun <- asksEngineEnv $ isAllocation . envHoistBlockers branch_blocker <- asksEngineEnv $ blockHoistBranch . envHoistBlockers vtable <- askVtable @@ -667,13 +695,13 @@ is_alloc_fun stm || ( ST.loopDepth vtable > 0 && cond_loop_invariant - && ifsort /= IfFallback + && ifsort /= MatchFallback && loopInvariantStm vtable stm -- Avoid hoisting out something that might change the -- asymptotics of the program. && all primType (patTypes (stmPat stm)) ) - || ( ifsort /= IfFallback + || ( ifsort /= MatchFallback && any (`UT.isSize` usage) (patNames (stmPat stm)) && all primType (patTypes (stmPat stm)) ) @@ -694,7 +722,7 @@ | is_alloc_fun stm = False isNotHoistableBnd _ _ _ = -- Hoist aggressively out of versioning branches. - ifsort /= IfEquiv + ifsort /= MatchEquiv block = branch_blocker @@ -702,14 +730,7 @@ `andAlso` notDesirableToHoist ) `orIf` isConsuming - - (hoisted1, body1') <- - protectIfHoisted cond True $ - simplifyBody block res_usage res_usages body1 - (hoisted2, body2') <- - protectIfHoisted cond False $ - simplifyBody block res_usage res_usages body2 - pure (body1', body2', hoisted1 <> hoisted2) + pure block -- | Simplify a single body. simplifyBody :: @@ -797,16 +818,27 @@ Pat (LetDec (Wise rep)) -> Exp (Wise rep) -> SimpleM rep (Exp (Wise rep), Stms (Wise rep)) -simplifyExp usage (Pat pes) (If cond tbranch fbranch ifdec@(IfDec ts ifsort)) = do - -- Here, we have to check whether 'cond' puts a bound on some free - -- variable, and if so, chomp it. We should also try to do CSE - -- across branches. +simplifyExp usage (Pat pes) (Match ses cases defbody ifdec@(MatchDec ts ifsort)) = do let pes_usages = map (fromMaybe mempty . (`UT.lookup` usage) . patElemName) pes - cond' <- simplify cond + ses' <- mapM simplify ses ts' <- mapM simplify ts - (tbranch', fbranch', hoisted) <- - hoistCommon usage pes_usages cond' ifdec tbranch fbranch - pure (If cond' tbranch' fbranch' $ IfDec ts' ifsort, hoisted) + let pats = map casePat cases + block <- matchBlocker ses ifdec + (cases_hoisted, cases') <- + unzip <$> zipWithM (simplifyCase block ses' pes_usages) (inits pats) cases + (defbody_hoisted, defbody') <- + protectCaseHoisted ses' pats [] $ + simplifyBody block usage pes_usages defbody + pure + ( Match ses' cases' defbody' $ MatchDec ts' ifsort, + mconcat $ defbody_hoisted : cases_hoisted + ) + where + simplifyCase block ses' pes_usages prior (Case vs body) = do + (hoisted, body') <- + protectCaseHoisted ses' prior vs $ + simplifyBody block usage pes_usages body + pure (hoisted, Case vs body') simplifyExp _ _ (DoLoop merge form loopbody) = do let (params, args) = unzip merge params' <- mapM (traverse simplify) params @@ -990,15 +1022,6 @@ cs' <- simplify cs (se', se_cs) <- collectCerts $ simplify se pure $ SubExpRes (se_cs <> cs') se' - -simplifyPat :: - (SimplifiableRep rep, Simplifiable dec) => - Pat dec -> - SimpleM rep (Pat dec) -simplifyPat (Pat xs) = - Pat <$> mapM inspect xs - where - inspect (PatElem name rep) = PatElem name <$> simplify rep instance Simplifiable () where simplify = pure diff --git a/src/Futhark/Optimise/Simplify/Rep.hs b/src/Futhark/Optimise/Simplify/Rep.hs --- a/src/Futhark/Optimise/Simplify/Rep.hs +++ b/src/Futhark/Optimise/Simplify/Rep.hs @@ -315,8 +315,8 @@ -- | Construct a 'Wise' expression. informExp :: Informing rep => Exp rep -> Exp (Wise rep) -informExp (If cond tbranch fbranch (IfDec ts ifsort)) = - If cond (informBody tbranch) (informBody fbranch) (IfDec ts ifsort) +informExp (Match cond cases defbody (MatchDec ts ifsort)) = + Match cond (map (fmap informBody) cases) (informBody defbody) (MatchDec ts ifsort) informExp (DoLoop merge form loopbody) = let form' = case form of ForLoop i it bound params -> ForLoop i it bound params diff --git a/src/Futhark/Optimise/Simplify/Rule.hs b/src/Futhark/Optimise/Simplify/Rule.hs --- a/src/Futhark/Optimise/Simplify/Rule.hs +++ b/src/Futhark/Optimise/Simplify/Rule.hs @@ -24,7 +24,7 @@ SimplificationRule (..), RuleGeneric, RuleBasicOp, - RuleIf, + RuleMatch, RuleDoLoop, -- * Top-down rules @@ -32,7 +32,7 @@ TopDownRule, TopDownRuleGeneric, TopDownRuleBasicOp, - TopDownRuleIf, + TopDownRuleMatch, TopDownRuleDoLoop, TopDownRuleOp, @@ -41,7 +41,7 @@ BottomUpRule, BottomUpRuleGeneric, BottomUpRuleBasicOp, - BottomUpRuleIf, + BottomUpRuleMatch, BottomUpRuleDoLoop, BottomUpRuleOp, @@ -72,7 +72,7 @@ LocalScope rep ) -instance (ASTRep rep, BuilderOps rep) => MonadBuilder (RuleM rep) where +instance (BuilderOps rep) => MonadBuilder (RuleM rep) where type Rep (RuleM rep) = rep mkExpDecM pat e = RuleM $ mkExpDecM pat e mkBodyM stms res = RuleM $ mkBodyM stms res @@ -116,14 +116,14 @@ Rule rep ) -type RuleIf rep a = +type RuleMatch rep a = a -> Pat (LetDec rep) -> StmAux (ExpDec rep) -> - ( SubExp, - Body rep, + ( [SubExp], + [Case (Body rep)], Body rep, - IfDec (BranchType rep) + MatchDec (BranchType rep) ) -> Rule rep @@ -149,7 +149,7 @@ data SimplificationRule rep a = RuleGeneric (RuleGeneric rep a) | RuleBasicOp (RuleBasicOp rep a) - | RuleIf (RuleIf rep a) + | RuleMatch (RuleMatch rep a) | RuleDoLoop (RuleDoLoop rep a) | RuleOp (RuleOp rep a) @@ -158,7 +158,7 @@ data Rules rep a = Rules { rulesAny :: [SimplificationRule rep a], rulesBasicOp :: [SimplificationRule rep a], - rulesIf :: [SimplificationRule rep a], + rulesMatch :: [SimplificationRule rep a], rulesDoLoop :: [SimplificationRule rep a], rulesOp :: [SimplificationRule rep a] } @@ -178,7 +178,7 @@ type TopDownRuleBasicOp rep = RuleBasicOp rep (TopDown rep) -type TopDownRuleIf rep = RuleIf rep (TopDown rep) +type TopDownRuleMatch rep = RuleMatch rep (TopDown rep) type TopDownRuleDoLoop rep = RuleDoLoop rep (TopDown rep) @@ -194,7 +194,7 @@ type BottomUpRuleBasicOp rep = RuleBasicOp rep (BottomUp rep) -type BottomUpRuleIf rep = RuleIf rep (BottomUp rep) +type BottomUpRuleMatch rep = RuleMatch rep (BottomUp rep) type BottomUpRuleDoLoop rep = RuleDoLoop rep (BottomUp rep) @@ -231,19 +231,20 @@ groupRules :: [SimplificationRule m a] -> Rules m a groupRules rs = Rules - rs - (filter forBasicOp rs) - (filter forIf rs) - (filter forDoLoop rs) - (filter forOp rs) + { rulesAny = rs, + rulesBasicOp = filter forBasicOp rs, + rulesMatch = filter forMatch rs, + rulesDoLoop = filter forDoLoop rs, + rulesOp = filter forOp rs + } forBasicOp RuleBasicOp {} = True forBasicOp RuleGeneric {} = True forBasicOp _ = False - forIf RuleIf {} = True - forIf RuleGeneric {} = True - forIf _ = False + forMatch RuleMatch {} = True + forMatch RuleGeneric {} = True + forMatch _ = False forDoLoop RuleDoLoop {} = True forDoLoop RuleGeneric {} = True @@ -283,7 +284,7 @@ BasicOp {} -> rulesBasicOp DoLoop {} -> rulesDoLoop Op {} -> rulesOp - If {} -> rulesIf + Match {} -> rulesMatch _ -> rulesAny applyRule :: SimplificationRule rep a -> a -> Stm rep -> Rule rep @@ -291,8 +292,8 @@ applyRule (RuleBasicOp f) a (Let pat aux (BasicOp e)) = f a pat aux e applyRule (RuleDoLoop f) a (Let pat aux (DoLoop merge form body)) = f a pat aux (merge, form, body) -applyRule (RuleIf f) a (Let pat aux (If cond tbody fbody ifsort)) = - f a pat aux (cond, tbody, fbody, ifsort) +applyRule (RuleMatch f) a (Let pat aux (Match cond cases defbody ifsort)) = + f a pat aux (cond, cases, defbody, ifsort) applyRule (RuleOp f) a (Let pat aux (Op op)) = f a pat aux op applyRule _ _ _ = diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs --- a/src/Futhark/Optimise/Simplify/Rules.hs +++ b/src/Futhark/Optimise/Simplify/Rules.hs @@ -1,5 +1,4 @@ {-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE TypeFamilies #-} @@ -22,7 +21,6 @@ import Control.Monad import Control.Monad.State -import Data.Either import Data.List (insert, unzip4, zip4) import qualified Data.Map.Strict as M import Data.Maybe @@ -36,20 +34,18 @@ import Futhark.Optimise.Simplify.Rules.BasicOp import Futhark.Optimise.Simplify.Rules.Index import Futhark.Optimise.Simplify.Rules.Loop +import Futhark.Optimise.Simplify.Rules.Match import Futhark.Util topDownRules :: BuilderOps rep => [TopDownRule rep] topDownRules = [ RuleGeneric constantFoldPrimFun, - RuleIf ruleIf, - RuleIf hoistBranchInvariant, RuleGeneric withAccTopDown ] bottomUpRules :: (BuilderOps rep, TraverseOpStms rep) => [BottomUpRule rep] bottomUpRules = - [ RuleIf removeDeadBranchResult, - RuleGeneric withAccBottomUp, + [ RuleGeneric withAccBottomUp, RuleBasicOp simplifyIndex ] @@ -57,7 +53,11 @@ -- functional semantics, and so probably should not be applied after -- memory block merging. standardRules :: (BuilderOps rep, TraverseOpStms rep, Aliased rep) => RuleBook rep -standardRules = ruleBook topDownRules bottomUpRules <> loopRules <> basicOpRules +standardRules = + ruleBook topDownRules bottomUpRules + <> loopRules + <> basicOpRules + <> matchRules -- | Turn @copy(x)@ into @x@ iff @x@ is not used after this copy -- statement and it can be consumed. @@ -135,211 +135,6 @@ seType (Constant v) = Just $ Prim $ primValueType v simplifyIndex _ _ _ _ = Skip -ruleIf :: BuilderOps rep => TopDownRuleIf rep -ruleIf _ pat _ (e1, tb, fb, IfDec _ ifsort) - | Just branch <- checkBranch, - ifsort /= IfFallback || isCt1 e1 = Simplify $ do - let ses = bodyResult branch - addStms $ bodyStms branch - sequence_ - [ certifying cs $ letBindNames [patElemName p] $ BasicOp $ SubExp se - | (p, SubExpRes cs se) <- zip (patElems pat) ses - ] - where - checkBranch - | isCt1 e1 = Just tb - | isCt0 e1 = Just fb - | otherwise = Nothing - --- IMPROVE: the following two rules can be generalised to work in more --- cases, especially when the branches have bindings, or return more --- than one value. --- --- if c then True else v == c || v -ruleIf - _ - pat - _ - ( cond, - Body _ tstms [SubExpRes tcs (Constant (BoolValue True))], - Body _ fstms [SubExpRes fcs se], - IfDec ts _ - ) - | null tstms, - null fstms, - [Prim Bool] <- map extTypeOf ts = - Simplify $ certifying (tcs <> fcs) $ letBind pat $ BasicOp $ BinOp LogOr cond se --- When type(x)==bool, if c then x else y == (c && x) || (!c && y) -ruleIf _ pat _ (cond, tb, fb, IfDec ts _) - | Body _ tstms [SubExpRes tcs tres] <- tb, - Body _ fstms [SubExpRes fcs fres] <- fb, - all (safeExp . stmExp) $ tstms <> fstms, - all ((== Prim Bool) . extTypeOf) ts = Simplify $ do - addStms tstms - addStms fstms - e <- - eBinOp - LogOr - (pure $ BasicOp $ BinOp LogAnd cond tres) - ( eBinOp - LogAnd - (pure $ BasicOp $ UnOp Not cond) - (pure $ BasicOp $ SubExp fres) - ) - certifying (tcs <> fcs) $ letBind pat e -ruleIf _ pat _ (_, tbranch, _, IfDec _ IfFallback) - | all (safeExp . stmExp) $ bodyStms tbranch = Simplify $ do - let ses = bodyResult tbranch - addStms $ bodyStms tbranch - sequence_ - [ certifying cs $ letBindNames [patElemName p] $ BasicOp $ SubExp se - | (p, SubExpRes cs se) <- zip (patElems pat) ses - ] -ruleIf _ pat _ (cond, tb, fb, _) - | Body _ _ [SubExpRes tcs (Constant (IntValue t))] <- tb, - Body _ _ [SubExpRes fcs (Constant (IntValue f))] <- fb = - if oneIshInt t && zeroIshInt f && tcs == mempty && fcs == mempty - then - Simplify $ - letBind pat $ - BasicOp $ - ConvOp (BToI (intValueType t)) cond - else - if zeroIshInt t && oneIshInt f - then Simplify $ do - cond_neg <- letSubExp "cond_neg" $ BasicOp $ UnOp Not cond - letBind pat $ BasicOp $ ConvOp (BToI (intValueType t)) cond_neg - else Skip --- Simplify --- --- let z = if c then x else y --- --- to --- --- let z = y --- --- in the case where 'x' is a loop parameter with initial value 'y' --- and the new value of the loop parameter is 'z'. ('x' and 'y' can --- be flipped.) -ruleIf vtable (Pat [pe]) aux (_c, tb, fb, IfDec [_] _) - | Body _ tstms [SubExpRes xcs x] <- tb, - null tstms, - Body _ fstms [SubExpRes ycs y] <- fb, - null fstms, - matches x y || matches y x = - Simplify $ - certifying (stmAuxCerts aux <> xcs <> ycs) $ - letBind (Pat [pe]) $ - BasicOp $ - SubExp y - where - z = patElemName pe - matches (Var x) y - | Just (initial, res) <- ST.lookupLoopParam x vtable = - initial == y && res == Var z - matches _ _ = False -ruleIf _ _ _ _ = Skip - --- | Move out results of a conditional expression whose computation is --- either invariant to the branches (only done for results used for --- existentials), or the same in both branches. -hoistBranchInvariant :: BuilderOps rep => TopDownRuleIf rep -hoistBranchInvariant _ pat _ (cond, tb, fb, IfDec ret ifsort) = Simplify $ do - let tses = bodyResult tb - fses = bodyResult fb - (hoistings, (pes, ts, res)) <- - fmap (fmap unzip3 . partitionEithers) . mapM branchInvariant $ - zip4 [0 ..] (patElems pat) ret (zip tses fses) - let ctx_fixes = catMaybes hoistings - (tses', fses') = unzip res - tb' = tb {bodyResult = tses'} - fb' = fb {bodyResult = fses'} - ret' = foldr (uncurry fixExt) ts ctx_fixes - if not $ null hoistings -- Was something hoisted? - then do - -- We may have to add some reshapes if we made the type - -- less existential. - tb'' <- reshapeBodyResults tb' $ map extTypeOf ret' - fb'' <- reshapeBodyResults fb' $ map extTypeOf ret' - letBind (Pat pes) $ If cond tb'' fb'' (IfDec ret' ifsort) - else cannotSimplify - where - bound_in_branches = - namesFromList . concatMap (patNames . stmPat) $ - bodyStms tb <> bodyStms fb - invariant Constant {} = True - invariant (Var v) = v `notNameIn` bound_in_branches - - branchInvariant (i, pe, t, (tse, fse)) - -- Do both branches return the same value? - | tse == fse = do - certifying (resCerts tse <> resCerts fse) $ - letBindNames [patElemName pe] $ - BasicOp $ - SubExp $ - resSubExp tse - hoisted i pe - - -- Do both branches return values that are free in the - -- branch, and are we not the only pattern element? The - -- latter is to avoid infinite application of this rule. - | invariant $ resSubExp tse, - invariant $ resSubExp fse, - patSize pat > 1, - Prim _ <- patElemType pe = do - bt <- expTypesFromPat $ Pat [pe] - letBindNames [patElemName pe] - =<< ( If cond - <$> resultBodyM [resSubExp tse] - <*> resultBodyM [resSubExp fse] - <*> pure (IfDec bt ifsort) - ) - hoisted i pe - | otherwise = - pure $ Right (pe, t, (tse, fse)) - - hoisted i pe = pure $ Left $ Just (i, Var $ patElemName pe) - - reshapeBodyResults body rets = buildBody_ $ do - ses <- bodyBind body - let (ctx_ses, val_ses) = splitFromEnd (length rets) ses - (ctx_ses ++) <$> zipWithM reshapeResult val_ses rets - reshapeResult (SubExpRes cs (Var v)) t@Array {} = do - v_t <- lookupType v - let newshape = arrayDims $ removeExistentials t v_t - SubExpRes cs - <$> if newshape /= arrayDims v_t - then letSubExp "branch_ctx_reshaped" (shapeCoerce newshape v) - else pure $ Var v - reshapeResult se _ = - pure se - --- | Remove the return values of a branch, that are not actually used --- after a branch. Standard dead code removal can remove the branch --- if *none* of the return values are used, but this rule is more --- precise. -removeDeadBranchResult :: BuilderOps rep => BottomUpRuleIf rep -removeDeadBranchResult (_, used) pat _ (e1, tb, fb, IfDec rettype ifsort) - | -- Only if there is no existential binding... - all (`notNameIn` foldMap freeIn (patElems pat)) (patNames pat), - -- Figure out which of the names in 'pat' are used... - patused <- map (`UT.isUsedDirectly` used) $ patNames pat, - -- If they are not all used, then this rule applies. - not (and patused) = - -- Remove the parts of the branch-results that correspond to dead - -- return value bindings. Note that this leaves dead code in the - -- branch bodies, but that will be removed later. - let tses = bodyResult tb - fses = bodyResult fb - pick :: [a] -> [a] - pick = map snd . filter fst . zip patused - tb' = tb {bodyResult = pick tses} - fb' = fb {bodyResult = pick fses} - pat' = pick $ patElems pat - rettype' = pick rettype - in Simplify $ letBind (Pat pat') $ If e1 tb' fb' $ IfDec rettype' ifsort - | otherwise = Skip - withAccTopDown :: BuilderOps rep => TopDownRuleGeneric rep -- A WithAcc with no accumulators is sent to Valhalla. withAccTopDown _ (Let pat aux (WithAcc [] lam)) = Simplify . auxing aux $ do @@ -465,16 +260,6 @@ getRidOf (pes, _) = not $ any ((`UT.used` utable) . patElemName) pes keepNonAccRes (pe, _) = patElemName pe `UT.used` utable withAccBottomUp _ _ = Skip - --- Some helper functions - -isCt1 :: SubExp -> Bool -isCt1 (Constant v) = oneIsh v -isCt1 _ = False - -isCt0 :: SubExp -> Bool -isCt0 (Constant v) = zeroIsh v -isCt0 _ = False -- Note [Dead Code Elimination for WithAcc] -- diff --git a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs --- a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +++ b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs @@ -196,9 +196,8 @@ case se of Var v | not $ null $ sliceDims is -> do v_reshaped <- - letExp (baseString v ++ "_reshaped") $ - BasicOp $ - Reshape (map DimNew $ arrayDims dest_t) v + letExp (baseString v ++ "_reshaped") . BasicOp $ + Reshape ReshapeArbitrary (arrayShape dest_t) v letBind pat $ BasicOp $ Copy v_reshaped _ -> letBind pat $ BasicOp $ ArrayLit [se] $ rowType dest_t ruleBasicOp vtable pat (StmAux cs1 attrs _) (Update safety1 dest1 is1 (Var v1)) @@ -217,7 +216,7 @@ where simplifyWith (Var v) x | Just stm <- ST.lookupStm v vtable, - If p tbranch fbranch _ <- stmExp stm, + Match [p] [Case [Just (BoolValue True)] tbranch] fbranch _ <- stmExp stm, Just (y, z) <- returns v (stmPat stm) tbranch fbranch, not $ boundInBody tbranch `namesIntersect` freeIn y, @@ -264,23 +263,20 @@ in letBind pat $ BasicOp $ Replicate (Shape [n]) se ruleBasicOp vtable pat aux (Index idd slice) | Just inds <- sliceIndices slice, - Just (BasicOp (Reshape newshape idd2), idd_cs) <- ST.lookupExp idd vtable, + Just (BasicOp (Reshape k newshape idd2), idd_cs) <- ST.lookupExp idd vtable, length newshape == length inds = Simplify $ - case shapeCoercion newshape of - Just _ -> - certifying idd_cs $ - auxing aux $ - letBind pat $ - BasicOp $ - Index idd2 slice - Nothing -> do + case k of + ReshapeCoerce -> + certifying idd_cs . auxing aux . letBind pat . BasicOp $ + Index idd2 slice + ReshapeArbitrary -> do -- Linearise indices and map to old index space. oldshape <- arrayDims <$> lookupType idd2 let new_inds = reshapeIndex (map pe64 oldshape) - (map pe64 $ newDims newshape) + (map pe64 $ shapeDims newshape) (map pe64 inds) new_inds' <- mapM (toSubExp "new_index") new_inds diff --git a/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs b/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs --- a/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs +++ b/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs @@ -39,7 +39,7 @@ -- | @foldClosedForm look foldfun accargs arrargs@ determines whether -- each of the results of @foldfun@ can be expressed in a closed form. foldClosedForm :: - (ASTRep rep, BuilderOps rep) => + (BuilderOps rep) => VarLookup rep -> Pat (LetDec rep) -> Lambda rep -> @@ -68,10 +68,10 @@ BasicOp $ CmpOp (CmpEq int64) inputsize (intConst Int64 0) letBind pat - =<< ( If (Var isEmpty) - <$> resultBodyM accs + =<< ( Match [Var isEmpty] + <$> (pure . Case [Just $ BoolValue True] <$> resultBodyM accs) <*> renameBody closedBody - <*> pure (IfDec [primBodyType t] IfNormal) + <*> pure (MatchDec [primBodyType t] MatchNormal) ) where knownBnds = determineKnownBindings look lam accs arrs @@ -79,7 +79,7 @@ -- | @loopClosedForm pat respat merge bound bodys@ determines whether -- the do-loop can be expressed in a closed form. loopClosedForm :: - (ASTRep rep, BuilderOps rep) => + (BuilderOps rep) => Pat (LetDec rep) -> [(FParam rep, SubExp)] -> Names -> @@ -108,10 +108,10 @@ CmpOp (CmpSlt it) bound (intConst it 0) letBind pat - =<< ( If (Var isEmpty) - <$> resultBodyM mergeexp + =<< ( Match [Var isEmpty] + <$> (pure . Case [Just (BoolValue True)] <$> resultBodyM mergeexp) <*> renameBody closedBody - <*> pure (IfDec [primBodyType t] IfNormal) + <*> pure (MatchDec [primBodyType t] MatchNormal) ) where (mergepat, mergeexp) = unzip merge diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs --- a/src/Futhark/Optimise/Simplify/Rules/Index.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Index.hs @@ -154,18 +154,16 @@ not consuming, ST.available src vtable -> Just $ pure $ IndexResult cs src $ Slice inds - Just (Reshape newshape src, cs) - | Just newdims <- shapeCoercion newshape, - Just olddims <- arrayDims <$> seType (Var src), - changed_dims <- zipWith (/=) newdims olddims, + Just (Reshape ReshapeCoerce newshape src, cs) + | Just olddims <- arrayDims <$> seType (Var src), + changed_dims <- zipWith (/=) (shapeDims newshape) olddims, not $ or $ drop (length inds) changed_dims -> Just $ pure $ IndexResult cs src $ Slice inds - | Just newdims <- shapeCoercion newshape, - Just olddims <- arrayDims <$> seType (Var src), + | Just olddims <- arrayDims <$> seType (Var src), length newshape == length inds, - length olddims == length newdims -> + length olddims == length (shapeDims newshape) -> Just $ pure $ IndexResult cs src $ Slice inds - Just (Reshape [_] v2, cs) + Just (Reshape _ (Shape [_]) v2, cs) | Just [_] <- arrayDims <$> seType (Var v2) -> Just $ pure $ IndexResult cs v2 $ Slice inds Just (Concat d (x :| xs) _, cs) @@ -197,14 +195,13 @@ (thisres, thisstms) <- collectStms $ do i' <- letSubExp "index_concat_i" $ BasicOp $ BinOp (Sub Int64 OverflowWrap) i start letSubExp "index_concat" . BasicOp . Index x' $ - Slice $ - ibef ++ DimFix i' : iaft + Slice (ibef ++ DimFix i' : iaft) thisbody <- mkBodyM thisstms [subExpRes thisres] (altres, altstms) <- collectStms $ mkBranch xs_and_starts' altbody <- mkBodyM altstms [subExpRes altres] letSubExp "index_concat_branch" $ - If cmp thisbody altbody $ - IfDec [primBodyType res_t] IfNormal + Match [cmp] [Case [Just $ BoolValue True] thisbody] altbody $ + MatchDec [primBodyType res_t] MatchNormal SubExpResult cs <$> mkBranch xs_and_starts Just (ArrayLit ses _, cs) | DimFix (Constant (IntValue (Int64Value i))) : inds' <- inds, diff --git a/src/Futhark/Optimise/Simplify/Rules/Match.hs b/src/Futhark/Optimise/Simplify/Rules/Match.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rules/Match.hs @@ -0,0 +1,251 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Match simplification rules. +module Futhark.Optimise.Simplify.Rules.Match (matchRules) where + +import Control.Monad +import Data.Either +import Data.List (partition, transpose, unzip4, zip5) +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Construct +import Futhark.IR +import Futhark.Optimise.Simplify.Rule +import Futhark.Util + +-- Does this case always match the scrutinees? +caseAlwaysMatches :: [SubExp] -> Case a -> Bool +caseAlwaysMatches ses = and . zipWith match ses . casePat + where + match se (Just v) = se == Constant v + match _ Nothing = True + +-- Can this case never match the scrutinees? +caseNeverMatches :: [SubExp] -> Case a -> Bool +caseNeverMatches ses = or . zipWith impossible ses . casePat + where + impossible (Constant v1) (Just v2) = v1 /= v2 + impossible _ _ = False + +ruleMatch :: BuilderOps rep => TopDownRuleMatch rep +-- Remove impossible cases. +ruleMatch _ pat _ (cond, cases, defbody, ifdec) + | (impossible, cases') <- partition (caseNeverMatches cond) cases, + not $ null impossible = + Simplify $ letBind pat $ Match cond cases' defbody ifdec +-- Find new default case. +ruleMatch _ pat _ (cond, cases, _, ifdec) + | (always_matches, cases') <- partition (caseAlwaysMatches cond) cases, + new_default : _ <- reverse always_matches = + Simplify $ letBind pat $ Match cond cases' (caseBody new_default) ifdec +-- Remove caseless match. +ruleMatch _ pat (StmAux cs _ _) (_, [], defbody, _) = Simplify $ do + defbody_res <- bodyBind defbody + certifying cs $ forM_ (zip (patElems pat) defbody_res) $ \(pe, res) -> + certifying (resCerts res) . letBind (Pat [pe]) $ + BasicOp (SubExp $ resSubExp res) +-- IMPROVE: the following two rules can be generalised to work in more +-- cases, especially when the branches have bindings, or return more +-- than one value. +-- +-- if c then True else v == c || v +ruleMatch + _ + pat + _ + ( [cond], + [ Case + [Just (BoolValue True)] + (Body _ tstms [SubExpRes tcs (Constant (BoolValue True))]) + ], + Body _ fstms [SubExpRes fcs se], + MatchDec ts _ + ) + | null tstms, + null fstms, + [Prim Bool] <- map extTypeOf ts = + Simplify $ certifying (tcs <> fcs) $ letBind pat $ BasicOp $ BinOp LogOr cond se +-- When type(x)==bool, if c then x else y == (c && x) || (!c && y) +ruleMatch _ pat _ ([cond], [Case [Just (BoolValue True)] tb], fb, MatchDec ts _) + | Body _ tstms [SubExpRes tcs tres] <- tb, + Body _ fstms [SubExpRes fcs fres] <- fb, + all (safeExp . stmExp) $ tstms <> fstms, + all ((== Prim Bool) . extTypeOf) ts = Simplify $ do + addStms tstms + addStms fstms + e <- + eBinOp + LogOr + (pure $ BasicOp $ BinOp LogAnd cond tres) + ( eBinOp + LogAnd + (pure $ BasicOp $ UnOp Not cond) + (pure $ BasicOp $ SubExp fres) + ) + certifying (tcs <> fcs) $ letBind pat e +ruleMatch _ pat _ (_, [Case _ tbranch], _, MatchDec _ MatchFallback) + | all (safeExp . stmExp) $ bodyStms tbranch = Simplify $ do + let ses = bodyResult tbranch + addStms $ bodyStms tbranch + sequence_ + [ certifying cs $ letBindNames [patElemName p] $ BasicOp $ SubExp se + | (p, SubExpRes cs se) <- zip (patElems pat) ses + ] +ruleMatch _ pat _ ([cond], [Case [Just (BoolValue True)] tb], fb, _) + | Body _ _ [SubExpRes tcs (Constant (IntValue t))] <- tb, + Body _ _ [SubExpRes fcs (Constant (IntValue f))] <- fb = + if oneIshInt t && zeroIshInt f && tcs == mempty && fcs == mempty + then + Simplify . letBind pat . BasicOp $ + ConvOp (BToI (intValueType t)) cond + else + if zeroIshInt t && oneIshInt f + then Simplify $ do + cond_neg <- letSubExp "cond_neg" $ BasicOp $ UnOp Not cond + letBind pat $ BasicOp $ ConvOp (BToI (intValueType t)) cond_neg + else Skip +-- Simplify +-- +-- let z = if c then x else y +-- +-- to +-- +-- let z = y +-- +-- in the case where 'x' is a loop parameter with initial value 'y' +-- and the new value of the loop parameter is 'z'. ('x' and 'y' can +-- be flipped.) +ruleMatch vtable (Pat [pe]) aux (_c, [Case _ tb], fb, MatchDec [_] _) + | Body _ tstms [SubExpRes xcs x] <- tb, + null tstms, + Body _ fstms [SubExpRes ycs y] <- fb, + null fstms, + matches x y || matches y x = + Simplify . certifying (stmAuxCerts aux <> xcs <> ycs) $ + letBind (Pat [pe]) (BasicOp $ SubExp y) + where + z = patElemName pe + matches (Var x) y + | Just (initial, res) <- ST.lookupLoopParam x vtable = + initial == y && res == Var z + matches _ _ = False +ruleMatch _ _ _ _ = Skip + +-- | Move out results of a conditional expression whose computation is +-- either invariant to the branches (only done for results used for +-- existentials), or the same in both branches. +hoistBranchInvariant :: BuilderOps rep => TopDownRuleMatch rep +hoistBranchInvariant _ pat _ (cond, cases, defbody, MatchDec ret ifsort) = + let case_reses = map (bodyResult . caseBody) cases + defbody_res = bodyResult defbody + (hoistings, (pes, ts, case_reses_tr, defbody_res')) = + (fmap unzip4 . partitionEithers) . map branchInvariant $ + zip5 [0 ..] (patElems pat) ret (transpose case_reses) defbody_res + in if null hoistings + then Skip + else Simplify $ do + ctx_fixes <- sequence hoistings + let onCase (Case vs body) case_res = Case vs $ body {bodyResult = case_res} + cases' = zipWith onCase cases $ transpose case_reses_tr + defbody' = defbody {bodyResult = defbody_res'} + ret' = foldr (uncurry fixExt) ts ctx_fixes + -- We may have to add some reshapes if we made the type + -- less existential. + cases'' <- mapM (traverse $ reshapeBodyResults $ map extTypeOf ret') cases' + defbody'' <- reshapeBodyResults (map extTypeOf ret') defbody' + letBind (Pat pes) $ Match cond cases'' defbody'' (MatchDec ret' ifsort) + where + bound_in_branches = + namesFromList . concatMap (patNames . stmPat) $ + foldMap (bodyStms . caseBody) cases <> bodyStms defbody + + branchInvariant (i, pe, t, case_reses, defres) + -- If just one branch has a variant result, then we give up. + | namesIntersect bound_in_branches $ freeIn $ defres : case_reses = + noHoisting + -- Do all branches return the same value? + | all ((== resSubExp defres) . resSubExp) case_reses = Left $ do + certifying (foldMap resCerts case_reses <> resCerts defres) $ + letBindNames [patElemName pe] . BasicOp . SubExp $ + resSubExp defres + hoisted i pe + + -- Do all branches return values that are free in the + -- branch, and are we not the only pattern element? The + -- latter is to avoid infinite application of this rule. + | not $ namesIntersect bound_in_branches $ freeIn $ defres : case_reses, + patSize pat > 1, + Prim _ <- patElemType pe = Left $ do + bt <- expTypesFromPat $ Pat [pe] + letBindNames [patElemName pe] + =<< ( Match cond + <$> ( zipWith Case (map casePat cases) + <$> mapM (resultBodyM . pure . resSubExp) case_reses + ) + <*> resultBodyM [resSubExp defres] + <*> pure (MatchDec bt ifsort) + ) + hoisted i pe + | otherwise = noHoisting + where + noHoisting = Right (pe, t, case_reses, defres) + + hoisted i pe = pure (i, Var $ patElemName pe) + + reshapeBodyResults rets body = buildBody_ $ do + ses <- bodyBind body + let (ctx_ses, val_ses) = splitFromEnd (length rets) ses + (ctx_ses ++) <$> zipWithM reshapeResult val_ses rets + reshapeResult (SubExpRes cs (Var v)) t@Array {} = do + v_t <- lookupType v + let newshape = arrayDims $ removeExistentials t v_t + SubExpRes cs + <$> if newshape /= arrayDims v_t + then letSubExp "branch_ctx_reshaped" (shapeCoerce newshape v) + else pure $ Var v + reshapeResult se _ = + pure se + +-- | Remove the return values of a branch, that are not actually used +-- after a branch. Standard dead code removal can remove the branch +-- if *none* of the return values are used, but this rule is more +-- precise. +removeDeadBranchResult :: BuilderOps rep => BottomUpRuleMatch rep +removeDeadBranchResult (_, used) pat _ (cond, cases, defbody, MatchDec rettype ifsort) + | -- Only if there is no existential binding... + all (`notNameIn` foldMap freeIn (patElems pat)) (patNames pat), + -- Figure out which of the names in 'pat' are used... + patused <- map (`UT.isUsedDirectly` used) $ patNames pat, + -- If they are not all used, then this rule applies. + not (and patused) = do + -- Remove the parts of the branch-results that correspond to dead + -- return value bindings. Note that this leaves dead code in the + -- branch bodies, but that will be removed later. + let pick :: [a] -> [a] + pick = map snd . filter fst . zip patused + pat' = pick $ patElems pat + rettype' = pick rettype + Simplify $ do + cases' <- mapM (traverse $ onBody pick) cases + defbody' <- onBody pick defbody + letBind (Pat pat') $ Match cond cases' defbody' $ MatchDec rettype' ifsort + | otherwise = Skip + where + onBody pick (Body _ stms res) = mkBodyM stms $ pick res + +topDownRules :: BuilderOps rep => [TopDownRule rep] +topDownRules = + [ RuleMatch ruleMatch, + RuleMatch hoistBranchInvariant + ] + +bottomUpRules :: (BuilderOps rep) => [BottomUpRule rep] +bottomUpRules = + [ RuleMatch removeDeadBranchResult + ] + +matchRules :: (BuilderOps rep) => RuleBook rep +matchRules = ruleBook topDownRules bottomUpRules diff --git a/src/Futhark/Optimise/Simplify/Rules/Simple.hs b/src/Futhark/Optimise/Simplify/Rules/Simple.hs --- a/src/Futhark/Optimise/Simplify/Rules/Simple.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Simple.hs @@ -264,48 +264,48 @@ simplifyAssert _ _ _ = Nothing +-- No-op reshape. simplifyIdentityReshape :: SimpleRule rep -simplifyIdentityReshape _ seType (Reshape newshape v) +simplifyIdentityReshape _ seType (Reshape _ newshape v) | Just t <- seType $ Var v, - newDims newshape == arrayDims t -- No-op reshape. - = + newshape == arrayShape t = resIsSubExp $ Var v simplifyIdentityReshape _ _ _ = Nothing simplifyReshapeReshape :: SimpleRule rep -simplifyReshapeReshape defOf _ (Reshape newshape v) - | Just (BasicOp (Reshape oldshape v2), v_cs) <- defOf v = - Just (Reshape (fuseReshape oldshape newshape) v2, v_cs) +simplifyReshapeReshape defOf _ (Reshape k1 newshape v) + | Just (BasicOp (Reshape k2 _ v2), v_cs) <- defOf v = + Just (Reshape (max k1 k2) newshape v2, v_cs) simplifyReshapeReshape _ _ _ = Nothing simplifyReshapeScratch :: SimpleRule rep -simplifyReshapeScratch defOf _ (Reshape newshape v) +simplifyReshapeScratch defOf _ (Reshape _ newshape v) | Just (BasicOp (Scratch bt _), v_cs) <- defOf v = - Just (Scratch bt $ newDims newshape, v_cs) + Just (Scratch bt $ shapeDims newshape, v_cs) simplifyReshapeScratch _ _ _ = Nothing simplifyReshapeReplicate :: SimpleRule rep -simplifyReshapeReplicate defOf seType (Reshape newshape v) +simplifyReshapeReplicate defOf seType (Reshape _ newshape v) | Just (BasicOp (Replicate _ se), v_cs) <- defOf v, Just oldshape <- arrayShape <$> seType se, - shapeDims oldshape `isSuffixOf` newDims newshape = + shapeDims oldshape `isSuffixOf` shapeDims newshape = let new = take (length newshape - shapeRank oldshape) $ - newDims newshape + shapeDims newshape in Just (Replicate (Shape new) se, v_cs) simplifyReshapeReplicate _ _ _ = Nothing simplifyReshapeIota :: SimpleRule rep -simplifyReshapeIota defOf _ (Reshape newshape v) +simplifyReshapeIota defOf _ (Reshape _ newshape v) | Just (BasicOp (Iota _ offset stride it), v_cs) <- defOf v, - [n] <- newDims newshape = + [n] <- shapeDims newshape = Just (Iota n offset stride it, v_cs) simplifyReshapeIota _ _ _ = Nothing simplifyReshapeConcat :: SimpleRule rep -simplifyReshapeConcat defOf seType (Reshape newshape v) = do +simplifyReshapeConcat defOf seType (Reshape ReshapeCoerce newshape v) = do (BasicOp (Concat d arrs _), v_cs) <- defOf v - (bef, w', aft) <- focusNth d =<< shapeCoercion newshape + (bef, w', aft) <- focusNth d $ shapeDims newshape (arr_bef, _, arr_aft) <- focusNth d <=< fmap arrayDims $ seType $ Var $ NE.head arrs guard $ arr_bef == bef @@ -323,10 +323,9 @@ -- If we are size-coercing a slice, then we might as well just use a -- different slice instead. simplifyReshapeIndex :: SimpleRule rep -simplifyReshapeIndex defOf _ (Reshape newshape v) - | Just ds <- shapeCoercion newshape, - Just (BasicOp (Index v' slice), v_cs) <- defOf v, - slice' <- Slice $ reshapeSlice (unSlice slice) ds, +simplifyReshapeIndex defOf _ (Reshape ReshapeCoerce newshape v) + | Just (BasicOp (Index v' slice), v_cs) <- defOf v, + slice' <- Slice $ reshapeSlice (unSlice slice) $ shapeDims newshape, slice' /= slice = Just (Index v' slice', v_cs) simplifyReshapeIndex _ _ _ = Nothing @@ -335,22 +334,13 @@ -- instead use the original array and update the slice dimensions. simplifyUpdateReshape :: SimpleRule rep simplifyUpdateReshape defOf seType (Update safety dest slice (Var v)) - | Just (BasicOp (Reshape newshape v'), v_cs) <- defOf v, - Just _ <- shapeCoercion newshape, + | Just (BasicOp (Reshape ReshapeCoerce _ v'), v_cs) <- defOf v, Just ds <- arrayDims <$> seType (Var v'), slice' <- Slice $ reshapeSlice (unSlice slice) ds, slice' /= slice = Just (Update safety dest slice' $ Var v', v_cs) simplifyUpdateReshape _ _ _ = Nothing -improveReshape :: SimpleRule rep -improveReshape _ seType (Reshape newshape v) - | Just t <- seType $ Var v, - newshape' <- informReshape (arrayDims t) newshape, - newshape' /= newshape = - Just (Reshape newshape' v, mempty) -improveReshape _ _ _ = Nothing - -- | If we are copying a scratch array (possibly indirectly), just turn it into a scratch by -- itself. copyScratchToScratch :: SimpleRule rep @@ -364,7 +354,7 @@ case asBasicOp . fst =<< defOf v of Just Scratch {} -> True Just (Rearrange _ v') -> isActuallyScratch v' - Just (Reshape _ v') -> isActuallyScratch v' + Just (Reshape _ _ v') -> isActuallyScratch v' _ -> False copyScratchToScratch _ _ _ = Nothing @@ -384,8 +374,7 @@ simplifyReshapeIota, simplifyReshapeConcat, simplifyReshapeIndex, - simplifyUpdateReshape, - improveReshape + simplifyUpdateReshape ] -- | Try to simplify the given t'BasicOp', returning a new t'BasicOp' diff --git a/src/Futhark/Optimise/Sink.hs b/src/Futhark/Optimise/Sink.hs --- a/src/Futhark/Optimise/Sink.hs +++ b/src/Futhark/Optimise/Sink.hs @@ -81,13 +81,16 @@ multiplicity :: Constraints rep => Stm rep -> M.Map VName Int multiplicity stm = case stmExp stm of - If cond tbranch fbranch _ -> - free cond 1 `comb` free tbranch 1 `comb` free fbranch 1 - Op {} -> free stm 2 - DoLoop {} -> free stm 2 - _ -> free stm 1 + Match cond cases defbody _ -> + foldl comb mempty $ + free 1 cond + : free 1 defbody + : map (free 1 . caseBody) cases + Op {} -> free 2 stm + DoLoop {} -> free 2 stm + _ -> free 1 stm where - free x k = M.fromList $ zip (namesToList $ freeIn x) $ repeat k + free k x = M.fromList $ zip (namesToList $ freeIn x) $ repeat k comb = M.unionWith (+) optimiseBranch :: @@ -172,12 +175,15 @@ in if patElemName pe `nameIn` sunk then (stms', sunk) else (stm : stms', sunk) - | If cond tbranch fbranch ret <- stmExp stm = - let (tbranch', tsunk) = optimiseBranch onOp vtable sinking tbranch - (fbranch', fsunk) = optimiseBranch onOp vtable sinking fbranch + | Match cond cases defbody ret <- stmExp stm = + let onCase (Case vs body) = + let (body', body_sunk) = optimiseBranch onOp vtable sinking body + in (Case vs body', body_sunk) + (cases', cases_sunk) = unzip $ map onCase cases + (defbody', defbody_sunk) = optimiseBranch onOp vtable sinking defbody (stms', sunk) = optimiseStms' vtable' sinking stms - in ( stm {stmExp = If cond tbranch' fbranch' ret} : stms', - tsunk <> fsunk <> sunk + in ( stm {stmExp = Match cond cases' defbody' ret} : stms', + mconcat cases_sunk <> defbody_sunk <> sunk ) | DoLoop merge lform body <- stmExp stm = let comps = (merge, lform, body) diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs --- a/src/Futhark/Optimise/TileLoops.hs +++ b/src/Futhark/Optimise/TileLoops.hs @@ -750,8 +750,9 @@ if null dims_on_top || null (arrayDims arr_t) -- Second check is for accumulators. then pure arr else do - let new_shape = unit_dims ++ arrayDims arr_t - letExp (baseString arr) $ BasicOp $ Reshape (map DimNew new_shape) arr + let new_shape = Shape $ unit_dims ++ arrayDims arr_t + letExp (baseString arr) . BasicOp $ + Reshape ReshapeArbitrary new_shape arr let tile_dims = zip (map snd dims_on_top) unit_dims ++ dims pure $ TileReturns mempty tile_dims arr' diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs --- a/src/Futhark/Optimise/TileLoops/Shared.hs +++ b/src/Futhark/Optimise/TileLoops/Shared.hs @@ -311,8 +311,10 @@ _ -> pure env changeIxFnEnv :: IxFnEnv -> VName -> Exp GPU -> TileM IxFnEnv -changeIxFnEnv env y (BasicOp (Reshape shp_chg x)) = - composeIxfuns env y x (`IxFun.reshape` map (fmap ExpMem.pe64) shp_chg) +changeIxFnEnv env y (BasicOp (Reshape ReshapeArbitrary shp_chg x)) = + composeIxfuns env y x (`IxFun.reshape` fmap ExpMem.pe64 (shapeDims shp_chg)) +changeIxFnEnv env y (BasicOp (Reshape ReshapeCoerce shp_chg x)) = + composeIxfuns env y x (`IxFun.coerce` fmap ExpMem.pe64 (shapeDims shp_chg)) changeIxFnEnv env y (BasicOp (Manifest perm x)) = do tp <- lookupType x case tp of @@ -323,8 +325,6 @@ _ -> error "In TileLoops/Shared.hs, changeIxFnEnv: manifest applied to a non-array!" changeIxFnEnv env y (BasicOp (Rearrange perm x)) = composeIxfuns env y x (`IxFun.permute` perm) -changeIxFnEnv env y (BasicOp (Rotate rs x)) = - composeIxfuns env y x (`IxFun.rotate` fmap ExpMem.pe64 rs) changeIxFnEnv env y (BasicOp (Index x slc)) = composeIxfuns env y x (`IxFun.slice` (Slice $ map (fmap ExpMem.pe64) $ unSlice slc)) changeIxFnEnv env y (BasicOp (Opaque _ (Var x))) = diff --git a/src/Futhark/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs --- a/src/Futhark/Pass/ExpandAllocations.hs +++ b/src/Futhark/Pass/ExpandAllocations.hs @@ -9,6 +9,7 @@ import Control.Monad.Reader import Control.Monad.State import Control.Monad.Writer +import Data.Either (rights) import Data.List (find, foldl') import qualified Data.Map.Strict as M import Data.Maybe @@ -86,26 +87,20 @@ transformStm :: Stm GPUMem -> ExpandM (Stms GPUMem) -- It is possible that we are unable to expand allocations in some -- code versions. If so, we can remove the offending branch. Only if --- both versions fail do we propagate the error. -transformStm (Let pat aux (If cond tbranch fbranch (IfDec ts IfEquiv))) = do - tbranch' <- (Right <$> transformBody tbranch) `catchError` (pure . Left) - fbranch' <- (Right <$> transformBody fbranch) `catchError` (pure . Left) - case (tbranch', fbranch') of - (Left _, Right fbranch'') -> - pure $ useBranch fbranch'' - (Right tbranch'', Left _) -> - pure $ useBranch tbranch'' - (Right tbranch'', Right fbranch'') -> - pure $ oneStm $ Let pat aux $ If cond tbranch'' fbranch'' (IfDec ts IfEquiv) - (Left e, _) -> +-- all versions fail do we propagate the error. +-- FIXME: this can remove safety checks if the default branch fails! +transformStm (Let pat aux (Match cond cases defbody (MatchDec ts MatchEquiv))) = do + let onCase (Case vs body) = + (Right . Case vs <$> transformBody body) `catchError` (pure . Left) + cases' <- rights <$> mapM onCase cases + defbody' <- (Right <$> transformBody defbody) `catchError` (pure . Left) + case (cases', defbody') of + ([], Left e) -> throwError e - where - bindRes pe (SubExpRes cs se) = - certify cs $ Let (Pat [pe]) (defAux ()) $ BasicOp $ SubExp se - - useBranch b = - bodyStms b - <> stmsFromList (zipWith bindRes (patElems pat) (bodyResult b)) + (_ : _, Left _) -> + pure $ oneStm $ Let pat aux $ Match cond (init cases') (caseBody $ last cases') (MatchDec ts MatchEquiv) + (_, Right defbody'') -> + pure $ oneStm $ Let pat aux $ Match cond cases' defbody'' (MatchDec ts MatchEquiv) transformStm (Let pat aux e) = do (stms, e') <- transformExp =<< mapExpM transform e pure $ stms <> oneStm (Let pat aux e') @@ -531,11 +526,9 @@ offset_ixfun = IxFun.slice root_ixfun . Slice $ [DimSlice 0 num_threads' 1, DimFix gtid] - shapechange = - if length old_shape == 1 - then map DimCoercion old_shape - else map DimNew old_shape - in IxFun.reshape offset_ixfun shapechange + in if length old_shape == 1 + then IxFun.coerce offset_ixfun old_shape + else IxFun.reshape offset_ixfun old_shape -- | A map from memory block names to new index function bases. type RebaseMap = M.Map VName (([TPrimExp Int64 VName], PrimType) -> IxFun) diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -45,7 +45,11 @@ import Control.Monad.Reader import Control.Monad.State import Control.Monad.Writer -import Data.List (foldl', partition, zip5) +import Data.Bifunctor (first) +import Data.Either (partitionEithers) +import Data.Foldable (toList) +import Data.List (foldl', transpose, zip4) +import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as M import Data.Maybe import qualified Data.Set as S @@ -58,7 +62,7 @@ import Futhark.Optimise.Simplify.Rep (mkWiseBody) import Futhark.Pass import Futhark.Tools -import Futhark.Util (maybeNth, splitAt3, splitFromEnd, takeLast) +import Futhark.Util (maybeNth, splitAt3) -- | The subexpression giving the number of elements we should -- allocate space for. See 'ChunkMap' comment. @@ -310,7 +314,7 @@ pure $ MemAcc acc ispace ts u summaryForBindage def_space chunkmap t@(Array pt shape u) NoHint = do m <- allocForArray' chunkmap t def_space - pure $ directIxFun pt shape u m t + pure $ MemArray pt shape u $ ArrayIn m $ IxFun.iota $ map pe64 $ arrayDims t summaryForBindage _ _ t@(Array pt _ _) (Hint ixfun space) = do bytes <- letSubExp "bytes" <=< toExp . untyped $ @@ -321,27 +325,15 @@ m <- letExp "mem" $ Op $ Alloc bytes space pure $ MemArray pt (arrayShape t) NoUniqueness $ ArrayIn m ixfun -lookupMemSpace :: (HasScope rep m, Monad m) => VName -> m Space -lookupMemSpace v = do - t <- lookupType v - case t of - Mem space -> pure space - _ -> error $ "lookupMemSpace: " ++ pretty v ++ " is not a memory block." - -directIxFun :: PrimType -> Shape -> u -> VName -> Type -> MemBound u -directIxFun bt shape u mem t = - let ixf = IxFun.iota $ map pe64 $ arrayDims t - in MemArray bt shape u $ ArrayIn mem ixf - allocInFParams :: (Allocable fromrep torep inner) => [(FParam fromrep, Space)] -> ([FParam torep] -> AllocM fromrep torep a) -> AllocM fromrep torep a allocInFParams params m = do - (valparams, (ctxparams, memparams)) <- + (valparams, (memparams, ctxparams)) <- runWriterT $ mapM (uncurry allocInFParam) params - let params' = ctxparams <> memparams <> valparams + let params' = memparams <> ctxparams <> valparams summary = scopeOfFParams params' localScope summary $ m params' @@ -359,7 +351,7 @@ let memname = baseString (paramName param) <> "_mem" ixfun = IxFun.iota $ map pe64 $ shapeDims shape mem <- lift $ newVName memname - tell ([], [Param (paramAttrs param) mem $ MemMem pspace]) + tell ([Param (paramAttrs param) mem $ MemMem pspace], []) pure param {paramDec = MemArray pt shape u $ ArrayIn mem ixfun} Prim pt -> pure param {paramDec = MemPrim pt} @@ -368,6 +360,38 @@ Acc acc ispace ts u -> pure param {paramDec = MemAcc acc ispace ts u} +ensureRowMajorArray :: + (Allocable fromrep torep inner) => + Maybe Space -> + VName -> + AllocM fromrep torep (VName, VName) +ensureRowMajorArray space_ok v = do + (mem, ixfun) <- lookupArraySummary v + mem_space <- lookupMemSpace mem + default_space <- askDefaultSpace + let space = fromMaybe default_space space_ok + if numLMADs ixfun == 1 + && ixFunPerm ixfun == [0 .. IxFun.rank ixfun - 1] + && length (IxFun.base ixfun) == IxFun.rank ixfun + && maybe True (== mem_space) space_ok + && IxFun.contiguous ixfun + then pure (mem, v) + else allocLinearArray space (baseString v) v + +ensureArrayIn :: + (Allocable fromrep torep inner) => + Space -> + SubExp -> + WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp +ensureArrayIn _ (Constant v) = + error $ "ensureArrayIn: " ++ pretty v ++ " cannot be an array." +ensureArrayIn space (Var v) = do + (mem', v') <- lift $ ensureRowMajorArray (Just space) v + (_, ixfun) <- lift $ lookupArraySummary v' + ctx <- lift $ mapM (letSubExp "ixfun_arg" <=< toExp) (toList ixfun) + tell ([Var mem'], ctx) + pure $ Var v' + allocInMergeParams :: (Allocable fromrep torep inner) => [(FParam fromrep, SubExp)] -> @@ -377,21 +401,19 @@ ) -> AllocM fromrep torep a allocInMergeParams merge m = do - ((valparams, valargs, handle_loop_subexps), (ctx_params, mem_params)) <- + ((valparams, valargs, handle_loop_subexps), (mem_params, ctx_params)) <- runWriterT $ unzip3 <$> mapM allocInMergeParam merge - let mergeparams' = ctx_params <> mem_params <> valparams + let mergeparams' = mem_params <> ctx_params <> valparams summary = scopeOfFParams mergeparams' mk_loop_res ses = do - (ses', (ctxargs, memargs)) <- + (ses', (memargs, ctxargs)) <- runWriterT $ zipWithM ($) handle_loop_subexps ses - pure (ctxargs <> memargs, ses') + pure (memargs <> ctxargs, ses') (valctx_args, valargs') <- mk_loop_res valargs let merge' = - zip - (ctx_params <> mem_params <> valparams) - (valctx_args <> valargs') + zip (mem_params <> ctx_params <> valparams) (valctx_args <> valargs') localScope summary $ m merge' mk_loop_res where param_names = namesFromList $ map (paramName . fst) merge @@ -407,7 +429,7 @@ if (res_mem_space, res_ixfun) == (v_mem_space, v_ixfun) then pure (res_mem, res) else lift $ arrayWithIxFun chunkmap v_mem_space v_ixfun (fromDecl param_t) res - tell ([], [Var res_mem']) + tell ([Var res_mem'], []) pure $ Var res' scalarRes _ _ _ se = pure se @@ -439,7 +461,7 @@ allocInMergeParam (mergeparam, Var v') else do p <- newParam "mem_param" $ MemMem v_mem_space - tell ([], [p]) + tell ([p], []) pure ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName p) v_ixfun}, @@ -447,28 +469,27 @@ scalarRes param_t v_mem_space v_ixfun ) _ -> do - (v', ext_ixfun, substs, v_mem') <- - lift $ existentializeArray v_mem_space v + (v_mem', v') <- lift $ ensureRowMajorArray Nothing v + (_, v_ixfun') <- lift $ lookupArraySummary v' v_mem_space' <- lift $ lookupMemSpace v_mem' - (ctx_params, param_ixfun_substs) <- - fmap unzip . forM substs $ \e -> do - p <- newParam "ctx_param_ext" $ MemPrim $ primExpType $ untyped e - pure (p, fmap Free $ le64 $ paramName p) - - tell (ctx_params, []) + ctx_params <- + replicateM (length v_ixfun') $ + newParam "ctx_param_ext" (MemPrim int64) param_ixfun <- instantiateIxFun $ IxFun.substituteInIxFun - (M.fromList $ zip (fmap Ext [0 ..]) param_ixfun_substs) - ext_ixfun + ( M.fromList . zip (fmap Ext [0 ..]) $ + map (le64 . Free . paramName) ctx_params + ) + (IxFun.existentialize v_ixfun') mem_param <- newParam "mem_param" $ MemMem v_mem_space' - tell ([], [mem_param]) + tell ([mem_param], ctx_params) pure ( mergeparam {paramDec = MemArray pt shape u $ ArrayIn (paramName mem_param) param_ixfun}, - v', + Var v', ensureArrayIn v_mem_space' ) allocInMergeParam (mergeparam, se) = doDefault mergeparam se =<< lift askDefaultSpace @@ -477,26 +498,6 @@ mergeparam' <- allocInFParam mergeparam space pure (mergeparam', se, linearFuncallArg (paramType mergeparam) space) --- Returns the existentialized index function, the list of substituted values and the memory location. -existentializeArray :: - (Allocable fromrep torep inner) => - Space -> - VName -> - AllocM fromrep torep (SubExp, ExtIxFun, [TPrimExp Int64 VName], VName) -existentializeArray space v = do - (mem', ixfun) <- lookupArraySummary v - sp <- lookupMemSpace mem' - - let (ext_ixfun', substs') = runState (IxFun.existentialize ixfun) [] - - case (ext_ixfun', sp == space) of - (Just x, True) -> pure (Var v, x, substs', mem') - _ -> do - (mem, v') <- allocLinearArray space (baseString v) v - ixfun' <- fromJust <$> lookupIxFun v' - let (ext_ixfun, substs) = runState (IxFun.existentialize ixfun') [] - pure (Var v', fromJust ext_ixfun, substs, mem) - arrayWithIxFun :: (MonadBuilder m, Op (Rep m) ~ MemOp inner, LetDec (Rep m) ~ LetDecMem) => ChunkMap -> @@ -512,28 +513,6 @@ letBind (Pat [PatElem v_copy $ MemArray pt shape u $ ArrayIn mem ixfun]) $ BasicOp $ Copy v pure (mem, v_copy) -ensureArrayIn :: - (Allocable fromrep torep inner) => - Space -> - SubExp -> - WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp -ensureArrayIn _ (Constant v) = - error $ "ensureArrayIn: " ++ pretty v ++ " cannot be an array." -ensureArrayIn space (Var v) = do - (sub_exp, _, substs, mem) <- lift $ existentializeArray space v - (ctx_vals, _) <- - unzip - <$> mapM - ( \s -> do - vname <- lift $ letExp "ctx_val" =<< toExp s - pure (Var vname, fmap Free $ primExpFromSubExp int64 $ Var vname) - ) - substs - - tell (ctx_vals, [Var mem]) - - pure sub_exp - ensureDirectArray :: (Allocable fromrep torep inner) => Maybe Space -> @@ -552,25 +531,39 @@ -- binding for the size of the memory block. allocLinearArray space (baseString v) v -allocLinearArray :: +allocPermArray :: (Allocable fromrep torep inner) => Space -> + [Int] -> String -> VName -> AllocM fromrep torep (VName, VName) -allocLinearArray space s v = do +allocPermArray space perm s v = do t <- lookupType v case t of Array pt shape u -> do mem <- allocForArray t space - v' <- newVName $ s <> "_linear" - let ixfun = directIxFun pt shape u mem t - pat = Pat [PatElem v' ixfun] - addStm $ Let pat (defAux ()) $ BasicOp $ Copy v + v' <- newVName $ s <> "_desired_form" + let info = + MemArray pt shape u . ArrayIn mem $ + IxFun.permute (IxFun.iota $ map pe64 $ arrayDims t) perm + pat = Pat [PatElem v' info] + addStm $ Let pat (defAux ()) $ BasicOp $ Manifest perm v pure (mem, v') _ -> - error $ "allocLinearArray: " ++ pretty t + error $ "allocPermArray: " ++ pretty t +allocLinearArray :: + (Allocable fromrep torep inner) => + Space -> + String -> + VName -> + AllocM fromrep torep (VName, VName) +allocLinearArray space s v = do + t <- lookupType v + let perm = [0 .. arrayRank t - 1] + allocPermArray space perm s v + funcallArgs :: (Allocable fromrep torep inner) => [(SubExp, Diet)] -> @@ -592,7 +585,7 @@ WriterT ([SubExp], [SubExp]) (AllocM fromrep torep) SubExp linearFuncallArg Array {} space (Var v) = do (mem, arg') <- lift $ ensureDirectArray (Just space) v - tell ([], [Var mem]) + tell ([Var mem], []) pure $ Var arg' linearFuncallArg _ _ arg = pure arg @@ -736,6 +729,191 @@ allocInLambda params body = mkLambda params . allocInStms (bodyStms body) $ pure $ bodyResult body +numLMADs :: IxFun -> Int +numLMADs = length . IxFun.ixfunLMADs + +ixFunPerm :: IxFun -> [Int] +ixFunPerm = map IxFun.ldPerm . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs + +ixFunMon :: IxFun -> [IxFun.Monotonicity] +ixFunMon = map IxFun.ldMon . IxFun.lmadDims . NE.head . IxFun.ixfunLMADs + +data MemReq + = MemReq Space [Int] [IxFun.Monotonicity] Rank Bool + | NeedsLinearisation Space + deriving (Eq, Show) + +combMemReqs :: MemReq -> MemReq -> MemReq +combMemReqs x@NeedsLinearisation {} _ = x +combMemReqs _ y@NeedsLinearisation {} = y +combMemReqs x@(MemReq x_space _ _ _ _) y@MemReq {} = + if x == y then x else NeedsLinearisation x_space + +type MemReqType = MemInfo (Ext SubExp) NoUniqueness MemReq + +combMemReqTypes :: MemReqType -> MemReqType -> MemReqType +combMemReqTypes (MemArray pt shape u x) (MemArray _ _ _ y) = + MemArray pt shape u $ combMemReqs x y +combMemReqTypes x _ = x + +contextRets :: MemReqType -> [MemInfo d u r] +contextRets (MemArray _ shape _ (MemReq space _ _ (Rank base_rank) _)) = + -- Memory + offset + base_rank + (stride,size)*rank. + MemMem space + : MemPrim int64 + : replicate base_rank (MemPrim int64) + ++ replicate (2 * shapeRank shape) (MemPrim int64) +contextRets (MemArray _ shape _ (NeedsLinearisation space)) = + -- Memory + offset + (base,stride,size)*rank. + MemMem space + : MemPrim int64 + : replicate (3 * shapeRank shape) (MemPrim int64) +contextRets _ = [] + +-- Add memory information to the body, but do not return memory/ixfun +-- information. Instead, return restrictions on what the index +-- function should look like. We will then (crudely) unify these +-- restrictions across all bodies. +allocInMatchBody :: + (Allocable fromrep torep inner) => + [ExtType] -> + Body fromrep -> + AllocM fromrep torep (Body torep, [MemReqType]) +allocInMatchBody rets (Body _ stms res) = + buildBody . allocInStms stms $ do + restrictions <- zipWithM restriction rets (map resSubExp res) + pure (res, restrictions) + where + restriction t se = do + v_info <- subExpMemInfo se + case (t, v_info) of + (Array pt shape u, MemArray _ _ _ (ArrayIn mem ixfun)) -> do + space <- lookupMemSpace mem + pure . MemArray pt shape u $ + if numLMADs ixfun == 1 + then + MemReq + space + (ixFunPerm ixfun) + (ixFunMon ixfun) + (Rank $ length $ IxFun.base ixfun) + (IxFun.contiguous ixfun) + else NeedsLinearisation space + (_, MemMem space) -> pure $ MemMem space + (_, MemPrim pt) -> pure $ MemPrim pt + (_, MemAcc acc ispace ts u) -> pure $ MemAcc acc ispace ts u + _ -> error $ "allocInMatchBody: mismatch: " ++ show (t, v_info) + +mkBranchRet :: [MemReqType] -> [BranchTypeMem] +mkBranchRet reqs = + let (ctx_rets, res_rets) = foldl helper ([], []) $ zip reqs offsets + in ctx_rets ++ res_rets + where + numCtxNeeded = length . contextRets + + offsets = scanl (+) 0 $ map numCtxNeeded reqs + num_new_ctx = last offsets + + helper (ctx_rets_acc, res_rets_acc) (req, ctx_offset) = + ( ctx_rets_acc ++ contextRets req, + res_rets_acc ++ [inspect ctx_offset req] + ) + + arrayInfo rank (NeedsLinearisation space) = + (space, [0 .. rank - 1], repeat IxFun.Inc, rank, True) + arrayInfo _ (MemReq space perm mon (Rank base_rank) contig) = + (space, perm, mon, base_rank, contig) + + inspect ctx_offset (MemArray pt shape u req) = + let shape' = fmap (adjustExt num_new_ctx) shape + (space, perm, mon, base_rank, contig) = arrayInfo (shapeRank shape) req + in MemArray pt shape' u . ReturnsNewBlock space ctx_offset $ + convert + <$> IxFun.mkExistential base_rank (zip perm mon) contig (ctx_offset + 1) + inspect _ (MemAcc acc ispace ts u) = MemAcc acc ispace ts u + inspect _ (MemPrim pt) = MemPrim pt + inspect _ (MemMem space) = MemMem space + + convert (Ext i) = le64 (Ext i) + convert (Free v) = Free <$> pe64 v + + adjustExt :: Int -> Ext a -> Ext a + adjustExt _ (Free v) = Free v + adjustExt k (Ext i) = Ext (k + i) + +addCtxToMatchBody :: + (Allocable fromrep torep inner) => + [MemReqType] -> + Body torep -> + AllocM fromrep torep (Body torep) +addCtxToMatchBody reqs body = buildBody_ $ do + res <- zipWithM linearIfNeeded reqs =<< bodyBind body + ctx <- concat <$> mapM resCtx res + pure $ ctx ++ res + where + linearIfNeeded (MemArray _ _ _ (NeedsLinearisation space)) (SubExpRes cs (Var v)) = + SubExpRes cs . Var . snd <$> ensureRowMajorArray (Just space) v + linearIfNeeded _ res = + pure res + + resCtx (SubExpRes _ Constant {}) = + pure [] + resCtx (SubExpRes _ (Var v)) = do + info <- lookupMemInfo v + case info of + MemPrim {} -> pure [] + MemAcc {} -> pure [] + MemMem {} -> pure [] -- should not happen + MemArray _ _ _ (ArrayIn mem ixfun) -> do + ixfun_exts <- mapM (letSubExp "ixfun_ext" <=< toExp) $ toList ixfun + pure $ subExpRes (Var mem) : subExpsRes ixfun_exts + +-- Do a a simple form of invariance analysis to simplify a Match. It +-- is unfortunate that we have to do it here, but functions such as +-- scalarRes will look carefully at the index functions before the +-- simplifier has a chance to run. In a perfect world we would +-- simplify away those copies afterwards. XXX; this should be fixed by +-- a more general copy-removal pass. See +-- Futhark.Optimise.EntryPointMem for a very specialised version of +-- the idea, but which could perhaps be generalised. +simplifyMatch :: + Mem rep inner => + [Case (Body rep)] -> + Body rep -> + [BranchTypeMem] -> + ( [Case (Body rep)], + Body rep, + [BranchTypeMem] + ) +simplifyMatch cases defbody ts = + let case_reses = map (bodyResult . caseBody) cases + defbody_res = bodyResult defbody + (ctx_fixes, variant) = + partitionEithers . map branchInvariant $ + zip4 [0 ..] (transpose case_reses) defbody_res ts + (cases_reses, defbody_reses, ts') = unzip3 variant + in ( zipWith onCase cases (transpose cases_reses), + onBody defbody defbody_reses, + foldr (uncurry fixExt) ts' ctx_fixes + ) + where + bound_in_branches = + namesFromList . concatMap (patNames . stmPat) $ + foldMap (bodyStms . caseBody) cases <> bodyStms defbody + + onCase c res = fmap (`onBody` res) c + onBody body res = body {bodyResult = res} + + branchInvariant (i, case_reses, defres, t) + -- If even one branch has a variant result, then we give up. + | namesIntersect bound_in_branches $ freeIn $ defres : case_reses = + Right (case_reses, defres, t) + -- Do all branches return the same value? + | all ((== resSubExp defres) . resSubExp) case_reses = + Left (i, resSubExp defres) + | otherwise = + Right (case_reses, defres, t) + allocInExp :: (Allocable fromrep torep inner) => Exp fromrep -> @@ -746,8 +924,8 @@ localScope (scopeOf form') $ do body' <- buildBody_ . allocInStms bodystms $ do - (val_ses, valres') <- mk_loop_val $ map resSubExp bodyres - pure $ subExpsRes val_ses <> zipWith SubExpRes (map resCerts bodyres) valres' + (valctx, valres') <- mk_loop_val $ map resSubExp bodyres + pure $ subExpsRes valctx <> zipWith SubExpRes (map resCerts bodyres) valres' pure $ DoLoop merge' form' body' allocInExp (Apply fname args rettype loc) = do args' <- funcallArgs args @@ -756,76 +934,17 @@ where mems = replicate num_arrays (MemMem DefaultSpace) num_arrays = length $ filter ((> 0) . arrayRank . declExtTypeOf) rettype -allocInExp (If cond tbranch0 fbranch0 (IfDec rets ifsort)) = do - let num_rets = length rets - -- switch to the explicit-mem rep, but do nothing about results - (tbranch, tm_ixfs) <- allocInIfBody num_rets tbranch0 - (fbranch, fm_ixfs) <- allocInIfBody num_rets fbranch0 - tspaces <- mkSpaceOks num_rets tbranch - fspaces <- mkSpaceOks num_rets fbranch - -- try to generalize (antiunify) the index functions of the then and else bodies - let sp_substs = zipWith generalize (zip tspaces tm_ixfs) (zip fspaces fm_ixfs) - (spaces, subs) = unzip sp_substs - tsubs = map (selectSub fst) subs - fsubs = map (selectSub snd) subs - (tbranch', trets) <- addResCtxInIfBody rets tbranch spaces tsubs - (fbranch', frets) <- addResCtxInIfBody rets fbranch spaces fsubs - if frets /= trets - then error "In allocInExp, IF case: antiunification of then/else produce different ExtInFn!" - else do - -- above is a sanity check; implementation continues on else branch - let res_then = bodyResult tbranch' - res_else = bodyResult fbranch' - size_ext = length res_then - length trets - (ind_ses0, r_then_else) = - partition (\(r_then, r_else, _) -> r_then == r_else) $ - zip3 res_then res_else [0 .. size_ext - 1] - (r_then_ext, r_else_ext, _) = unzip3 r_then_else - ind_ses = - zipWith - (\(se, _, i) k -> (i - k, se)) - ind_ses0 - [0 .. length ind_ses0 - 1] - rets'' = foldl (\acc (i, SubExpRes _ se) -> fixExt i se acc) trets ind_ses - tbranch'' = tbranch' {bodyResult = r_then_ext ++ drop size_ext res_then} - fbranch'' = fbranch' {bodyResult = r_else_ext ++ drop size_ext res_else} - res_if_expr = If cond tbranch'' fbranch'' $ IfDec rets'' ifsort - pure res_if_expr +allocInExp (Match ses cases defbody (MatchDec rets ifsort)) = do + (defbody', def_reqs) <- allocInMatchBody rets defbody + (cases', cases_reqs) <- unzip <$> mapM onCase cases + let reqs = zipWith (foldl combMemReqTypes) def_reqs (transpose cases_reqs) + defbody'' <- addCtxToMatchBody reqs defbody' + cases'' <- mapM (traverse $ addCtxToMatchBody reqs) cases' + let (cases''', defbody''', rets') = + simplifyMatch cases'' defbody'' $ mkBranchRet reqs + pure $ Match ses cases''' defbody''' $ MatchDec rets' ifsort where - generalize :: - (Maybe Space, Maybe IxFun) -> - (Maybe Space, Maybe IxFun) -> - (Maybe Space, Maybe (ExtIxFun, [(TPrimExp Int64 VName, TPrimExp Int64 VName)])) - generalize (Just sp1, Just ixf1) (Just sp2, Just ixf2) = - if sp1 /= sp2 - then (Just sp1, Nothing) - else case IxFun.leastGeneralGeneralization (fmap untyped ixf1) (fmap untyped ixf2) of - Just (ixf, m) -> - ( Just sp1, - Just - ( fmap TPrimExp ixf, - zip (map (TPrimExp . fst) m) (map (TPrimExp . snd) m) - ) - ) - Nothing -> (Just sp1, Nothing) - generalize (mbsp1, _) _ = (mbsp1, Nothing) - - selectSub :: - ((a, a) -> a) -> - Maybe (ExtIxFun, [(a, a)]) -> - Maybe (ExtIxFun, [a]) - selectSub f (Just (ixfn, m)) = Just (ixfn, map f m) - selectSub _ Nothing = Nothing - allocInIfBody :: - (Allocable fromrep torep inner) => - Int -> - Body fromrep -> - AllocM fromrep torep (Body torep, [Maybe IxFun]) - allocInIfBody num_vals (Body _ stms res) = - buildBody . allocInStms stms $ do - let (_, val_res) = splitFromEnd num_vals res - mem_ixfs <- mapM (subExpIxFun . resSubExp) val_res - pure (res, mem_ixfs) + onCase (Case vs body) = first (Case vs) <$> allocInMatchBody rets body allocInExp (WithAcc inputs bodylam) = WithAcc <$> mapM onInput inputs <*> onLambda bodylam where @@ -893,120 +1012,6 @@ handle op } -lookupIxFun :: - (Allocable fromrep torep inner) => - VName -> - AllocM fromrep torep (Maybe IxFun) -lookupIxFun v = do - info <- lookupMemInfo v - case info of - MemArray _ptp _shp _u (ArrayIn _ ixf) -> pure $ Just ixf - _ -> pure Nothing - -subExpIxFun :: - (Allocable fromrep torep inner) => - SubExp -> - AllocM fromrep torep (Maybe IxFun) -subExpIxFun Constant {} = pure Nothing -subExpIxFun (Var v) = lookupIxFun v - -addResCtxInIfBody :: - (Allocable fromrep torep inner) => - [ExtType] -> - Body torep -> - [Maybe Space] -> - [Maybe (ExtIxFun, [TPrimExp Int64 VName])] -> - AllocM fromrep torep (Body torep, [BodyReturns]) -addResCtxInIfBody ifrets (Body _ stms res) spaces substs = buildBody $ do - mapM_ addStm stms - let offsets = scanl (+) 0 $ zipWith numCtxNeeded ifrets substs - num_new_ctx = last offsets - (ctx, ctx_rets, res', res_rets) <- - foldM (helper num_new_ctx) ([], [], [], []) $ - zip5 ifrets res substs spaces offsets - pure (ctx <> res', ctx_rets ++ res_rets) - where - numCtxNeeded Array {} Nothing = 1 - numCtxNeeded Array {} (Just (_, m)) = length m + 1 - numCtxNeeded _ _ = 0 - - helper - num_new_ctx - (ctx_acc, ctx_rets_acc, res_acc, res_rets_acc) - (ifr, r, mbixfsub, sp, ctx_offset) = - case mbixfsub of - Nothing -> do - -- does NOT generalize/antiunify; ensure direct - r' <- ensureDirect sp r - (mem_ctx_ses, mem_ctx_rets) <- unzip <$> bodyReturnMemCtx r' - let body_ret = inspect num_new_ctx ctx_offset ifr sp - pure - ( ctx_acc ++ mem_ctx_ses, - ctx_rets_acc ++ mem_ctx_rets, - res_acc ++ [r'], - res_rets_acc ++ [body_ret] - ) - Just (ixfn, m) -> do - -- generalizes - let i = length m - ext_ses <- mapM (toSubExp "ixfn_exist") m - (mem_ctx_ses, mem_ctx_rets) <- unzip <$> bodyReturnMemCtx r - let sp' = fromMaybe DefaultSpace sp - ixfn' = fmap (adjustExtPE ctx_offset) ixfn - exttp = case ifr of - Array pt shape u -> - MemArray pt (fmap (adjustExt num_new_ctx) shape) u $ - ReturnsNewBlock sp' (ctx_offset + i) ixfn' - _ -> error "Impossible case reached in addResCtxInIfBody" - pure - ( ctx_acc ++ subExpsRes ext_ses ++ mem_ctx_ses, - ctx_rets_acc ++ map (const (MemPrim int64)) ext_ses ++ mem_ctx_rets, - res_acc ++ [r], - res_rets_acc ++ [exttp] - ) - - inspect num_new_ctx k (Array pt shape u) space = - let space' = fromMaybe DefaultSpace space - shape' = fmap (adjustExt num_new_ctx) shape - bodyret = - MemArray pt shape' u . ReturnsNewBlock space' k $ - IxFun.iota $ - map convert $ - shapeDims shape' - in bodyret - inspect _ _ (Acc acc ispace ts u) _ = MemAcc acc ispace ts u - inspect _ _ (Prim pt) _ = MemPrim pt - inspect _ _ (Mem space) _ = MemMem space - - convert (Ext i) = le64 (Ext i) - convert (Free v) = Free <$> pe64 v - - adjustExt :: Int -> Ext a -> Ext a - adjustExt _ (Free v) = Free v - adjustExt k (Ext i) = Ext (k + i) - - adjustExtPE :: Int -> TPrimExp t (Ext VName) -> TPrimExp t (Ext VName) - adjustExtPE k = fmap (adjustExt k) - -mkSpaceOks :: - (Mem torep inner, LocalScope torep m) => - Int -> - Body torep -> - m [Maybe Space] -mkSpaceOks num_vals (Body _ stms res) = - inScopeOf stms $ mapM (mkSpaceOK . resSubExp) $ takeLast num_vals res - where - mkSpaceOK (Var v) = do - v_info <- lookupMemInfo v - case v_info of - MemArray _ _ _ (ArrayIn mem _) -> do - mem_info <- lookupMemInfo mem - case mem_info of - MemMem space -> pure $ Just space - _ -> pure Nothing - _ -> pure Nothing - mkSpaceOK _ = pure Nothing - allocInLoopForm :: (Allocable fromrep torep inner) => LoopForm fromrep -> @@ -1020,9 +1025,7 @@ case paramType p of Array pt shape u -> do dims <- map pe64 . arrayDims <$> lookupType a - let ixfun' = - IxFun.slice ixfun $ - fullSliceNum dims [DimFix $ le64 i] + let ixfun' = IxFun.slice ixfun $ fullSliceNum dims [DimFix $ le64 i] pure (p {paramDec = MemArray pt shape u $ ArrayIn mem ixfun'}, a) Prim bt -> pure (p {paramDec = MemPrim bt}, a) @@ -1072,8 +1075,7 @@ nohints = map (const NoHint) names mkLetNamesB'' :: - ( BuilderOps rep, - Mem rep inner, + ( Mem rep inner, LetDec rep ~ LetDecMem, OpReturns (Engine.OpWithWisdom inner), ExpDec rep ~ (), @@ -1097,13 +1099,15 @@ ( Engine.SimplifiableRep rep, ExpDec rep ~ (), BodyDec rep ~ (), + LetDec rep ~ LetDecMem, + OpReturns (Engine.OpWithWisdom inner), Mem rep inner ) => (Engine.OpWithWisdom inner -> UT.UsageTable) -> (Engine.OpWithWisdom inner -> Engine.SimpleM rep (Engine.OpWithWisdom inner, Stms (Engine.Wise rep))) -> SimpleOps rep simplifiable innerUsage simplifyInnerOp = - SimpleOps mkExpDecS' mkBodyS' protectOp opUsage simplifyOp + SimpleOps mkExpDecS' mkBodyS' protectOp opUsage simplifyPat simplifyOp where mkExpDecS' _ pat e = pure $ Engine.mkWiseExpDec pat () e @@ -1115,8 +1119,8 @@ fbody <- resultBodyM [intConst Int64 0] size' <- letSubExp "hoisted_alloc_size" $ - If taken tbody fbody $ - IfDec [MemPrim int64] IfFallback + Match [taken] [Case [Just $ BoolValue True] tbody] fbody $ + MatchDec [MemPrim int64] MatchFallback letBind pat $ Op $ Alloc size' space protectOp _ _ _ = Nothing @@ -1132,6 +1136,26 @@ simplifyOp (Inner k) = do (k', hoisted) <- simplifyInnerOp k pure (Inner k', hoisted) + + simplifyPat (Pat pes) e = do + rets <- expReturns e + Pat <$> zipWithM update pes rets + where + names = map patElemName pes + update + (PatElem pe_v (MemArray pt shape u (ArrayIn mem _))) + (MemArray _ _ _ (Just (ReturnsInBlock _ ixfun))) + | Just ixfun' <- traverse (traverse inst) ixfun = + PatElem pe_v + <$> ( MemArray pt + <$> Engine.simplify shape + <*> pure u + <*> (ArrayIn <$> Engine.simplify mem <*> pure ixfun') + ) + where + inst (Ext i) = maybeNth i names + inst (Free v) = Just v + update pe _ = traverse Engine.simplify pe data ExpHint = NoHint diff --git a/src/Futhark/Pass/ExplicitAllocations/GPU.hs b/src/Futhark/Pass/ExplicitAllocations/GPU.hs --- a/src/Futhark/Pass/ExplicitAllocations/GPU.hs +++ b/src/Futhark/Pass/ExplicitAllocations/GPU.hs @@ -126,7 +126,7 @@ hint Prim {} (ConcatReturns _ SplitContiguous w elems_per_thread _) = do let ixfun_base = IxFun.iota [sExt64 num_threads, pe64 elems_per_thread] ixfun_tr = IxFun.permute ixfun_base [1, 0] - ixfun = IxFun.reshape ixfun_tr $ map (DimNew . pe64) [w] + ixfun = IxFun.reshape ixfun_tr [pe64 w] pure $ Hint ixfun DefaultSpace hint _ _ = pure NoHint diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs --- a/src/Futhark/Pass/ExtractKernels.hs +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -289,10 +289,11 @@ unbalancedStm _ DoLoop {} = False unbalancedStm bound (WithAcc _ lam) = unbalancedBody bound (lambdaBody lam) - unbalancedStm bound (If cond tbranch fbranch _) = - cond - `subExpBound` bound - && (unbalancedBody bound tbranch || unbalancedBody bound fbranch) + unbalancedStm bound (Match ses cases defbody _) = + any (`subExpBound` bound) ses + && ( any (unbalancedBody bound . caseBody) cases + || unbalancedBody bound defbody + ) unbalancedStm _ (BasicOp _) = False unbalancedStm _ (Apply fname _ _ _) = @@ -342,9 +343,8 @@ alt_stms <- kernelAlternatives alts_pat default_body alts let alt_body = mkBody alt_stms $ varsRes $ patNames alts_pat - letBind pat $ - If cond alt alt_body $ - IfDec (staticShapes (patTypes pat)) IfEquiv + letBind pat . Match [cond] [Case [Just $ BoolValue True] alt] alt_body $ + MatchDec (staticShapes (patTypes pat)) MatchEquiv transformLambda :: KernelPath -> Lambda SOACS -> DistribM (Lambda GPU) transformLambda path (Lambda params body ret) = @@ -360,10 +360,10 @@ | "sequential_outer" `inAttrs` stmAuxAttrs aux = transformStms path . stmsToList . fmap (certify (stmAuxCerts aux)) =<< runBuilder_ (FOT.transformSOAC pat soac) -transformStm path (Let pat aux (If c tb fb rt)) = do - tb' <- transformBody path tb - fb' <- transformBody path fb - pure $ oneStm $ Let pat aux $ If c tb' fb' rt +transformStm path (Let pat aux (Match c cases defbody rt)) = do + cases' <- mapM (traverse $ transformBody path) cases + defbody' <- transformBody path defbody + pure $ oneStm $ Let pat aux $ Match c cases' defbody' rt transformStm path (Let pat aux (WithAcc inputs lam)) = oneStm . Let pat aux <$> (WithAcc (map transformInput inputs) <$> transformLambda path lam) @@ -611,8 +611,11 @@ mapLike w lam' | DoLoop _ _ body <- stmExp stm = bodyInterest body * 10 - | If _ tbody fbody _ <- stmExp stm = - max (bodyInterest tbody) (bodyInterest fbody) + | Match _ cases defbody _ <- stmExp stm = + foldl + max + (bodyInterest defbody) + (map (bodyInterest . caseBody) cases) | Op (Screma w _ (ScremaForm _ _ lam')) <- stmExp stm = zeroIfTooSmall w + bodyInterest (lambdaBody lam') | Op (Stream _ _ Sequential _ lam') <- stmExp stm = diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs --- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +++ b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs @@ -129,7 +129,7 @@ SegMap lvl kspace (lambdaReturnType map_lam) kbody dummyDim :: - (MonadFreshNames m, MonadBuilder m) => + MonadBuilder m => Pat Type -> m (Pat Type, [(VName, SubExp)], m ()) dummyDim pat = do diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs --- a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs +++ b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs @@ -322,7 +322,7 @@ && not ("sequential" `inAttrs` stmAuxAttrs (stmAux stm)) isMap BasicOp {} = False isMap Apply {} = False - isMap If {} = False + isMap Match {} = False isMap (DoLoop _ ForLoop {} body) = bodyContainsParallelism body isMap (DoLoop _ WhileLoop {} _) = False isMap (WithAcc _ lam) = bodyContainsParallelism $ lambdaBody lam @@ -413,11 +413,10 @@ pure acc' _ -> addStmToAcc stm acc -maybeDistributeStm stm@(Let pat _ (If cond tbranch fbranch ret)) acc +maybeDistributeStm stm@(Let pat _ (Match cond cases defbody ret)) acc | all (`notNameIn` freeIn pat) (patNames pat), - bodyContainsParallelism tbranch - || bodyContainsParallelism fbranch - || not (all primType (ifReturns ret)) = + any bodyContainsParallelism (defbody : map caseBody cases) + || not (all primType (matchReturns ret)) = distributeSingleStm acc stm >>= \case Just (kernels, res, nest, acc') | not $ @@ -429,7 +428,7 @@ nest' <- expandKernelNest pat_unused nest addPostStms kernels types <- asksScope scopeForSOACs - let branch = Branch perm pat cond tbranch fbranch ret + let branch = Branch perm pat cond cases defbody ret stms <- (`runReaderT` types) $ simplifyStms . oneStm =<< interchangeBranch nest' branch @@ -606,12 +605,10 @@ [ Let (Pat [PatElem arr' arr_t]) aux $ BasicOp $ Copy arr, Let outerpat aux $ BasicOp $ Rearrange perm' arr' ] -maybeDistributeStm stm@(Let _ aux (BasicOp (Reshape reshape stm_arr))) acc = +maybeDistributeStm stm@(Let _ aux (BasicOp (Reshape k reshape stm_arr))) acc = distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do - let reshape' = - map DimNew (kernelNestWidths nest) - ++ map DimNew (newDims reshape) - pure $ oneStm $ Let outerpat aux $ BasicOp $ Reshape reshape' arr + let reshape' = Shape (kernelNestWidths nest) <> reshape + pure $ oneStm $ Let outerpat aux $ BasicOp $ Reshape k reshape' arr maybeDistributeStm stm@(Let _ aux (BasicOp (Rotate rots stm_arr))) acc = distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do let rots' = map (const $ intConst Int64 0) (kernelNestWidths nest) ++ rots diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs --- a/src/Futhark/Pass/ExtractKernels/Interchange.hs +++ b/src/Futhark/Pass/ExtractKernels/Interchange.hs @@ -189,11 +189,11 @@ -- | An encoding of a branch with alongside its result pattern. data Branch - = Branch [Int] (Pat Type) SubExp (Body SOACS) (Body SOACS) (IfDec (BranchType SOACS)) + = Branch [Int] (Pat Type) [SubExp] [Case (Body SOACS)] (Body SOACS) (MatchDec (BranchType SOACS)) branchStm :: Branch -> Stm SOACS -branchStm (Branch _ pat cond tbranch fbranch ret) = - Let pat (defAux ()) $ If cond tbranch fbranch ret +branchStm (Branch _ pat cond cases defbody ret) = + Let pat (defAux ()) $ Match cond cases defbody ret interchangeBranch1 :: (MonadFreshNames m, HasScope SOACS m) => @@ -201,7 +201,7 @@ LoopNesting -> m Branch interchangeBranch1 - (Branch perm branch_pat cond tbranch fbranch (IfDec ret if_sort)) + (Branch perm branch_pat cond cases defbody (MatchDec ret if_sort)) (MapNesting pat aux w params_and_arrs) = do let ret' = map (`arrayOfRow` Free w) ret pat' = Pat $ rearrangeShape perm $ patElems pat @@ -218,11 +218,10 @@ map_stm = Let branch_pat' aux $ Op $ Screma w arrs $ mapSOAC lam pure $ mkBody (oneStm map_stm) res - tbranch' <- runBodyBuilder $ mkBranch tbranch - fbranch' <- runBodyBuilder $ mkBranch fbranch - pure $ - Branch [0 .. patSize pat - 1] pat' cond tbranch' fbranch' $ - IfDec ret' if_sort + cases' <- mapM (traverse $ runBodyBuilder . mkBranch) cases + defbody' <- runBodyBuilder $ mkBranch defbody + pure . Branch [0 .. patSize pat - 1] pat' cond cases' defbody' $ + MatchDec ret' if_sort -- | Given a (parallel) map nesting and an inner branch, move the maps -- inside the branch. The result is the resulting branch expression, diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs --- a/src/Futhark/Pass/ExtractKernels/Intragroup.hs +++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs @@ -209,12 +209,11 @@ form' = case form of ForLoop i it bound inps -> ForLoop i it bound inps WhileLoop cond -> WhileLoop cond - If cond tbody fbody ifdec -> do - tbody' <- intraGroupBody lvl tbody - fbody' <- intraGroupBody lvl fbody - certifying (stmAuxCerts aux) $ - letBind pat $ - If cond tbody' fbody' ifdec + Match cond cases defbody ifdec -> do + cases' <- mapM (traverse $ intraGroupBody lvl) cases + defbody' <- intraGroupBody lvl defbody + certifying (stmAuxCerts aux) . letBind pat $ + Match cond cases' defbody' ifdec Op soac | "sequential_outer" `inAttrs` stmAuxAttrs aux -> intraGroupStms lvl . fmap (certify (stmAuxCerts aux)) diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs --- a/src/Futhark/Pass/ExtractMulticore.hs +++ b/src/Futhark/Pass/ExtractMulticore.hs @@ -121,9 +121,11 @@ localScope (scopeOfFParams (map fst merge) <> scopeOf form') $ transformBody body pure $ oneStm $ Let pat aux $ DoLoop merge form' body' -transformStm (Let pat aux (If cond tbranch fbranch ret)) = +transformStm (Let pat aux (Match ses cases defbody ret)) = oneStm . Let pat aux - <$> (If cond <$> transformBody tbranch <*> transformBody fbranch <*> pure ret) + <$> (Match ses <$> mapM transformCase cases <*> transformBody defbody <*> pure ret) + where + transformCase (Case vs body) = Case vs <$> transformBody body transformStm (Let pat aux (WithAcc inputs lam)) = oneStm . Let pat aux <$> (WithAcc <$> mapM transformInput inputs <*> transformLambda lam) diff --git a/src/Futhark/Pass/KernelBabysitting.hs b/src/Futhark/Pass/KernelBabysitting.hs --- a/src/Futhark/Pass/KernelBabysitting.hs +++ b/src/Futhark/Pass/KernelBabysitting.hs @@ -65,7 +65,7 @@ case M.lookup name m of Just (Let _ _ (BasicOp (Opaque _ (Var arr)))) -> nonlinearInMemory arr m Just (Let _ _ (BasicOp (Rearrange perm _))) -> Just $ Just $ rearrangeInverse perm - Just (Let _ _ (BasicOp (Reshape _ arr))) -> nonlinearInMemory arr m + Just (Let _ _ (BasicOp (Reshape _ _ arr))) -> nonlinearInMemory arr m Just (Let _ _ (BasicOp (Manifest perm _))) -> Just $ Just perm Just (Let pat _ (Op (SegOp (SegMap _ _ ts _)))) -> nonlinear @@ -150,7 +150,7 @@ m (Maybe (VName, Slice SubExp)) traverseKernelBodyArrayIndexes :: - (Applicative f, Monad f) => + Monad f => Names -> Names -> Scope GPU -> @@ -474,12 +474,10 @@ let arr_shape = arrayShape arr_t padding_shape = setDim d arr_shape padding arr_padding <- - letExp (baseString arr <> "_padding") $ - BasicOp $ - Scratch (elemType arr_t) (shapeDims padding_shape) - letExp (baseString arr <> "_padded") $ - BasicOp $ - Concat d (arr :| [arr_padding]) w_padded + letExp (baseString arr <> "_padding") . BasicOp $ + Scratch (elemType arr_t) (shapeDims padding_shape) + letExp (baseString arr <> "_padded") . BasicOp $ + Concat d (arr :| [arr_padding]) w_padded rearrange num_chunks' w_padded per_chunk arr_name arr_padded arr_t = do let arr_dims = arrayDims arr_t @@ -488,19 +486,17 @@ extradim_shape = Shape $ pre_dims ++ [num_chunks', per_chunk] ++ post_dims tr_perm = [0 .. d - 1] ++ map (+ d) ([1] ++ [2 .. shapeRank extradim_shape - 1 - d] ++ [0]) arr_extradim <- - letExp (arr_name <> "_extradim") $ - BasicOp $ - Reshape (map DimNew $ shapeDims extradim_shape) arr_padded + letExp (arr_name <> "_extradim") . BasicOp $ + Reshape ReshapeArbitrary extradim_shape arr_padded arr_extradim_tr <- - letExp (arr_name <> "_extradim_tr") $ - BasicOp $ - Manifest tr_perm arr_extradim + letExp (arr_name <> "_extradim_tr") . BasicOp $ + Manifest tr_perm arr_extradim arr_inv_tr <- - letExp (arr_name <> "_inv_tr") $ - BasicOp $ - Reshape - (map DimCoercion pre_dims ++ map DimNew (w_padded : post_dims)) - arr_extradim_tr + letExp (arr_name <> "_inv_tr") . BasicOp $ + Reshape + ReshapeArbitrary + (Shape $ pre_dims ++ w_padded : post_dims) + arr_extradim_tr letExp (arr_name <> "_inv_tr_init") =<< eSliceArray d arr_inv_tr (eSubExp $ constant (0 :: Int64)) (eSubExp w) diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -22,6 +22,7 @@ import Futhark.IR.SeqMem (SeqMem) import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer +import Futhark.Optimise.EntryPointMem import Futhark.Optimise.Fusion import Futhark.Optimise.GenRedOpt import Futhark.Optimise.HistAccs @@ -125,6 +126,8 @@ >>> onePass Seq.explicitAllocations >>> passes [ performCSE False, + simplifySeqMem, + entryPointMemSeq, simplifySeqMem ] @@ -138,6 +141,7 @@ [ simplifyGPUMem, performCSE False, simplifyGPUMem, + entryPointMemGPU, doubleBufferGPU, simplifyGPUMem, MemoryBlockMerging.optimise, @@ -170,6 +174,7 @@ [ simplifyMCMem, performCSE False, simplifyMCMem, + entryPointMemMC, doubleBufferMC, simplifyMCMem ] diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -126,9 +126,8 @@ -- Finally, the array parameters are set to the arrays (but reshaped -- to make the types work out; this will be simplified rapidly). forM_ (zip arr_params arrs) $ \(p, arr) -> - letBindNames [paramName p] $ - BasicOp $ - Reshape (map DimCoercion $ arrayDims $ paramType p) arr + letBindNames [paramName p] . BasicOp $ + Reshape ReshapeCoerce (arrayShape $ paramType p) arr -- Then we just inline the lambda body. mapM_ addStm $ bodyStms $ lambdaBody lam @@ -140,7 +139,7 @@ certifying cs $ case (arrayDims $ patElemType pe, se) of (dims, Var v) | not $ null dims -> - letBindNames [patElemName pe] $ BasicOp $ Reshape (map DimCoercion dims) v + letBindNames [patElemName pe] $ BasicOp $ Reshape ReshapeCoerce (Shape dims) v _ -> letBindNames [patElemName pe] $ BasicOp $ SubExp se -- | Split the parameters of a stream reduction lambda into the chunk diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs --- a/src/Futhark/Transform/Substitute.hs +++ b/src/Futhark/Transform/Substitute.hs @@ -183,9 +183,6 @@ identType = substituteNames substs $ identType v } -instance Substitute d => Substitute (DimChange d) where - substituteNames substs = fmap $ substituteNames substs - instance Substitute d => Substitute (DimIndex d) where substituteNames substs = fmap $ substituteNames substs diff --git a/src/Futhark/Util/Log.hs b/src/Futhark/Util/Log.hs --- a/src/Futhark/Util/Log.hs +++ b/src/Futhark/Util/Log.hs @@ -51,11 +51,11 @@ -- | Append an entire log. addLog :: Log -> m () -instance (Applicative m, Monad m) => MonadLogger (WriterT Log m) where +instance Monad m => MonadLogger (WriterT Log m) where addLog = tell -instance (Applicative m, Monad m) => MonadLogger (Control.Monad.RWS.Lazy.RWST r Log s m) where +instance Monad m => MonadLogger (Control.Monad.RWS.Lazy.RWST r Log s m) where addLog = tell -instance (Applicative m, Monad m) => MonadLogger (Control.Monad.RWS.Strict.RWST r Log s m) where +instance Monad m => MonadLogger (Control.Monad.RWS.Strict.RWST r Log s m) where addLog = tell diff --git a/src/Language/Futhark/FreeVars.hs b/src/Language/Futhark/FreeVars.hs --- a/src/Language/Futhark/FreeVars.hs +++ b/src/Language/Futhark/FreeVars.hs @@ -11,7 +11,6 @@ import qualified Data.Map.Strict as M import qualified Data.Set as S -import Futhark.IR.Pretty () import Language.Futhark.Prop import Language.Futhark.Syntax diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -37,6 +37,7 @@ ( find, foldl', genericLength, + genericTake, intercalate, isPrefixOf, transpose, @@ -1917,8 +1918,9 @@ pure $ toArray (ShapeDim m (ShapeDim n shape)) $ map (toArray (ShapeDim n shape)) $ - transpose $ - map (snd . fromArray) xs' + -- Slight hack to work around empty dimensions. + genericTake m $ + transpose (map (snd . fromArray) xs') ++ repeat [] def "rotate" = Just $ fun2t $ \i xs -> do let (shape, xs') = fromArray xs diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs --- a/src/Language/Futhark/TypeChecker.hs +++ b/src/Language/Futhark/TypeChecker.hs @@ -158,7 +158,8 @@ typeError loc1 mempty $ "Duplicate definition of" <+> ppr space - <+> pprName name <> ". Previously defined at" + <+> pprName name <> "." + "Previously defined at" <+> text (locStr loc2) <> "." checkForDuplicateDecs :: [DecBase NoInfo Name] -> TypeM () diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -130,7 +130,12 @@ vtable <- asks $ scopeVtable . termScope let isGlobal v = case v `M.lookup` vtable of Just (BoundV Global _ _) -> True - _ -> False + Just EqualityF {} -> True + Just OverloadedF {} -> True + Just (BoundV Local _ _) -> False + Just (BoundV Nonlocal _ _) -> False + Just WasConsumed {} -> False + Nothing -> False pure . S.map AliasBound . S.filter (not . isGlobal) $ allOccurring closure S.\\ mconcat (map patNames params) diff --git a/unittests/Futhark/IR/Mem/IxFun/Alg.hs b/unittests/Futhark/IR/Mem/IxFun/Alg.hs --- a/unittests/Futhark/IR/Mem/IxFun/Alg.hs +++ b/unittests/Futhark/IR/Mem/IxFun/Alg.hs @@ -5,8 +5,8 @@ iota, offsetIndex, permute, - rotate, reshape, + coerce, slice, flatSlice, rebase, @@ -15,14 +15,11 @@ ) where -import Futhark.IR.Pretty () import Futhark.IR.Prop import Futhark.IR.Syntax - ( DimChange (..), - DimIndex (..), + ( DimIndex (..), FlatDimIndex (..), FlatSlice (..), - ShapeChange, Slice (..), flatSliceDims, sliceDims, @@ -41,10 +38,10 @@ data IxFun num = Direct (Shape num) | Permute (IxFun num) Permutation - | Rotate (IxFun num) (Indices num) | Index (IxFun num) (Slice num) | FlatIndex (IxFun num) (FlatSlice num) - | Reshape (IxFun num) (ShapeChange num) + | Reshape (IxFun num) (Shape num) + | Coerce (IxFun num) (Shape num) | OffsetIndex (IxFun num) num | Rebase (IxFun num) (IxFun num) deriving (Eq, Show) @@ -53,13 +50,16 @@ ppr (Direct dims) = text "Direct" <> parens (commasep $ map ppr dims) ppr (Permute fun perm) = ppr fun <> ppr perm - ppr (Rotate fun offsets) = ppr fun <> brackets (commasep $ map ((text "+" <>) . ppr) offsets) ppr (Index fun is) = ppr fun <> ppr is ppr (FlatIndex fun is) = ppr fun <> ppr is ppr (Reshape fun oldshape) = ppr fun <> text "->reshape" - <> parens (commasep (map ppr oldshape)) + <> parens (ppr oldshape) + ppr (Coerce fun oldshape) = + ppr fun + <> text "->coerce" + <> parens (ppr oldshape) ppr (OffsetIndex fun i) = ppr fun <> text "->offset_index" <> parens (ppr i) ppr (Rebase new_base fun) = @@ -74,9 +74,6 @@ permute :: IxFun num -> Permutation -> IxFun num permute = Permute -rotate :: IxFun num -> Indices num -> IxFun num -rotate = Rotate - slice :: IxFun num -> Slice num -> IxFun num slice = Index @@ -86,9 +83,12 @@ rebase :: IxFun num -> IxFun num -> IxFun num rebase = Rebase -reshape :: IxFun num -> ShapeChange num -> IxFun num +reshape :: IxFun num -> Shape num -> IxFun num reshape = Reshape +coerce :: IxFun num -> Shape num -> IxFun num +coerce = Reshape + shape :: IntegralExp num => IxFun num -> @@ -97,14 +97,14 @@ dims shape (Permute ixfun perm) = rearrangeShape perm $ shape ixfun -shape (Rotate ixfun _) = - shape ixfun shape (Index _ how) = sliceDims how shape (FlatIndex ixfun how) = flatSliceDims how <> tail (shape ixfun) shape (Reshape _ dims) = - map newDim dims + dims +shape (Coerce _ dims) = + dims shape (OffsetIndex ixfun _) = shape ixfun shape (Rebase _ ixfun) = @@ -123,10 +123,6 @@ index fun is_old where is_old = rearrangeShape (rearrangeInverse perm) is_new -index (Rotate fun offsets) is = - index fun $ zipWith mod (zipWith (+) is offsets) dims - where - dims = shape fun index (Index fun (Slice js)) is = index fun (adjust js is) where @@ -138,8 +134,10 @@ where f i (FlatDimIndex _ s) = i * s index (Reshape fun newshape) is = - let new_indices = reshapeIndex (shape fun) (newDims newshape) is + let new_indices = reshapeIndex (shape fun) newshape is in index fun new_indices +index (Coerce fun _) is = + index fun is index (OffsetIndex fun i) is = case shape fun of d : ds -> @@ -150,17 +148,17 @@ Direct old_shape -> if old_shape == shape new_base then new_base - else reshape new_base $ map DimCoercion old_shape + else reshape new_base old_shape Permute ixfun perm -> permute (rebase new_base ixfun) perm - Rotate ixfun offsets -> - rotate (rebase new_base ixfun) offsets Index ixfun iis -> slice (rebase new_base ixfun) iis FlatIndex ixfun iis -> flatSlice (rebase new_base ixfun) iis Reshape ixfun new_shape -> reshape (rebase new_base ixfun) new_shape + Coerce ixfun new_shape -> + coerce (rebase new_base ixfun) new_shape OffsetIndex ixfun s -> offsetIndex (rebase new_base ixfun) s r@Rebase {} -> diff --git a/unittests/Futhark/IR/Mem/IxFunTests.hs b/unittests/Futhark/IR/Mem/IxFunTests.hs --- a/unittests/Futhark/IR/Mem/IxFunTests.hs +++ b/unittests/Futhark/IR/Mem/IxFunTests.hs @@ -100,12 +100,6 @@ test_slice_iota, test_reshape_slice_iota1, test_permute_slice_iota, - test_rotate_rotate_permute_slice_iota, - test_slice_rotate_permute_slice_iota1, - test_slice_rotate_permute_slice_iota2, - test_slice_rotate_permute_slice_iota3, - test_permute_rotate_slice_permute_slice_iota, - test_reshape_rotate_iota, test_reshape_permute_iota, test_reshape_slice_iota2, test_reshape_slice_iota3, @@ -119,10 +113,7 @@ test_slice_flatSlice_iota, test_flatSlice_flatSlice_iota, test_flatSlice_slice_iota, - test_flatSlice_rotate_iota, - test_flatSlice_rotate_slice_iota, - test_flatSlice_transpose_slice_iota, - test_rotate_flatSlice_transpose_slice_iota + test_flatSlice_transpose_slice_iota ] singleton :: TestTree -> [TestTree] @@ -149,7 +140,7 @@ compareOps $ reshape (slice (iota [n, n, n]) slice3) - [DimNew (n `P.div` 2), DimNew (n `P.div` 3)] + [n `P.div` 2, n `P.div` 3] test_permute_slice_iota :: [TestTree] test_permute_slice_iota = @@ -158,114 +149,13 @@ compareOps $ permute (slice (iota [n, n, n]) slice3) [1, 0] -test_rotate_rotate_permute_slice_iota :: [TestTree] -test_rotate_rotate_permute_slice_iota = - singleton $ - testCase "rotate . rotate . permute . slice . iota" $ - compareOps $ - let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] - in rotate (rotate ixfun [2, 1]) [1, 2] - -test_slice_rotate_permute_slice_iota1 :: [TestTree] -test_slice_rotate_permute_slice_iota1 = - singleton $ - testCase "slice . rotate . permute . slice . iota 1" $ - compareOps $ - let slice2 = - Slice - [ DimSlice 0 n 1, - DimSlice 1 (n `P.div` 2) 2, - DimSlice 0 n 1 - ] - slice13 = - Slice - [ DimSlice 2 (n `P.div` 3) 3, - DimSlice 0 (n `P.div` 2) 1, - DimSlice 1 (n `P.div` 2) 2 - ] - ixfun = permute (slice (iota [n, n, n]) slice2) [2, 1, 0] - ixfun' = slice (rotate ixfun [3, 1, 2]) slice13 - in ixfun' - -test_slice_rotate_permute_slice_iota2 :: [TestTree] -test_slice_rotate_permute_slice_iota2 = - singleton $ - testCase "slice . rotate . permute . slice . iota 2" $ - compareOps $ - let slice2 = - Slice - [ DimSlice 0 (n `P.div` 2) 1, - DimFix (n `P.div` 2), - DimSlice 0 (n `P.div` 3) 1 - ] - slice13 = - Slice - [ DimSlice 2 (n `P.div` 3) 3, - DimSlice 0 n 1, - DimSlice 1 (n `P.div` 2) 2 - ] - ixfun = permute (slice (iota [n, n, n]) slice13) [2, 1, 0] - ixfun' = slice (rotate ixfun [3, 1, 2]) slice2 - in ixfun' - -test_slice_rotate_permute_slice_iota3 :: [TestTree] -test_slice_rotate_permute_slice_iota3 = - singleton $ - testCase "slice . rotate . permute . slice . iota 3" $ - compareOps $ - -- full-slice of (-1) stride - let ixfun = permute (slice (iota [n, n, n]) slice3) [1, 0] - ixfun' = rotate ixfun [2, 1] - - (n1, m1) = case IxFunLMAD.shape (fst ixfun') of - [a, b] -> (a, b) - _ -> error "expecting 2 dimensions at this point!" - negslice = - Slice - [ DimSlice 0 n1 1, - DimSlice (m1 - 1) m1 (-1) - ] - ixfun'' = rotate (slice ixfun' negslice) [1, 2] - in ixfun'' - -test_permute_rotate_slice_permute_slice_iota :: [TestTree] -test_permute_rotate_slice_permute_slice_iota = - singleton $ - testCase "permute . rotate . slice . permute . slice . iota" $ - compareOps $ - -- contiguousness - let slice33 = - Slice - [ DimFix (n `P.div` 2), - DimSlice (n - 1) (n `P.div` 3) (-1), - DimSlice 0 n 1 - ] - ixfun = permute (slice (iota [n, n, n]) slice33) [1, 0] - m = n `P.div` 3 - slice1 = - Slice - [ DimSlice (n - 1) n (-1), - DimSlice 2 (m - 2) 1 - ] - ixfun' = permute (rotate (slice ixfun slice1) [1, 2]) [1, 0] - in ixfun' - -test_reshape_rotate_iota :: [TestTree] -test_reshape_rotate_iota = - -- negative reshape test - singleton $ - testCase "reshape . rotate . iota" $ - compareOps $ - let newdims = [DimNew (n * n), DimCoercion n] - in reshape (rotate (iota [n, n, n]) [1, 0, 0]) newdims - test_reshape_permute_iota :: [TestTree] test_reshape_permute_iota = -- negative reshape test singleton $ testCase "reshape . permute . iota" $ compareOps $ - let newdims = [DimNew (n * n), DimCoercion n] + let newdims = [n * n, n] in reshape (permute (iota [n, n, n]) [1, 2, 0]) newdims test_reshape_slice_iota2 :: [TestTree] @@ -274,7 +164,7 @@ singleton $ testCase "reshape . slice . iota 2" $ compareOps $ - let newdims = [DimNew (n * n), DimCoercion n] + let newdims = [n * n, n] slc = Slice [ DimFix (n `P.div` 2), @@ -290,7 +180,7 @@ singleton $ testCase "reshape . slice . iota 3" $ compareOps $ - let newdims = [DimNew (n * n), DimCoercion n] + let newdims = [n * n, n] slc = Slice [ DimFix (n `P.div` 2), @@ -303,13 +193,13 @@ test_complex1 :: [TestTree] test_complex1 = singleton $ - testCase "reshape . permute . rotate . slice . permute . slice . iota 1" $ + testCase "reshape . permute . slice . permute . slice . iota 1" $ compareOps $ let newdims = - [ DimCoercion n, - DimCoercion n, - DimNew n, - DimCoercion ((n `P.div` 3) - 2) + [ n, + n, + n, + (n `P.div` 3) - 2 ] slice33 = Slice @@ -327,18 +217,18 @@ DimSlice (n - 1) n (-1), DimSlice 1 (m - 2) (-1) ] - ixfun' = reshape (rotate (slice ixfun slice1) [1, 2, 3, 4]) newdims + ixfun' = reshape (slice ixfun slice1) newdims in ixfun' test_complex2 :: [TestTree] test_complex2 = singleton $ - testCase "reshape . permute . rotate . slice . permute . slice . iota 2" $ + testCase "reshape . permute . slice . permute . slice . iota 2" $ compareOps $ let newdims = - [ DimCoercion n, - DimNew (n * n), - DimCoercion ((n `P.div` 3) - 2) + [ n, + n * n, + (n `P.div` 3) - 2 ] slc2 = Slice @@ -357,7 +247,7 @@ DimSlice (n - 1) n (-1), DimSlice 1 (m - 2) (-1) ] - ixfun' = reshape (rotate (slice ixfun slice1) [1, 0, 0, 2]) newdims + ixfun' = reshape (slice ixfun slice1) newdims in ixfun' test_rebase1 :: [TestTree] @@ -371,8 +261,8 @@ DimSlice 2 (n - 2) 1, DimSlice 3 (n - 3) 1 ] - ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] - ixfn_orig = rotate (permute (iota [n - 3, n - 2]) [1, 0]) [1, 2] + ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0] + ixfn_orig = permute (iota [n - 3, n - 2]) [1, 0] ixfn_rebase = rebase ixfn_base ixfn_orig in ixfn_rebase @@ -392,8 +282,8 @@ [ DimSlice (n - 4) (n - 3) (-1), DimSlice (n - 3) (n - 2) (-1) ] - ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] - ixfn_orig = rotate (permute (slice (iota [n - 3, n - 2]) slice_orig) [1, 0]) [1, 2] + ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0] + ixfn_orig = permute (slice (iota [n - 3, n - 2]) slice_orig) [1, 0] ixfn_rebase = rebase ixfn_base ixfn_orig in ixfn_rebase @@ -415,8 +305,8 @@ [ DimSlice (n3 - 1) n3 (-1), DimSlice (n2 - 1) n2 (-1) ] - ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] - ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] + ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0] + ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0] ixfn_rebase = rebase ixfn_base ixfn_orig in ixfn_rebase @@ -435,8 +325,8 @@ [ DimSlice (n3 - 1) n3 (-1), DimSlice 0 n2 1 ] - ixfn_base = rotate (permute (slice (iota [n, n, n]) slice_base) [1, 0]) [2, 1] - ixfn_orig = rotate (permute (slice (iota [n3, n2]) slice_orig) [1, 0]) [1, 2] + ixfn_base = permute (slice (iota [n, n, n]) slice_base) [1, 0] + ixfn_orig = permute (slice (iota [n3, n2]) slice_orig) [1, 0] in [ testCase "rebase mixed monotonicities" $ compareOps $ rebase ixfn_base ixfn_orig @@ -479,38 +369,11 @@ where flat_slice_1 = FlatSlice 17 [FlatDimIndex 3 27, FlatDimIndex 3 10, FlatDimIndex 3 1] -test_flatSlice_rotate_iota :: [TestTree] -test_flatSlice_rotate_iota = - singleton $ - testCase "flatSlice . rotate . iota " $ - compareOps $ - flatSlice (rotate (iota [10, 10]) [2, 5]) flat_slice_1 - where - flat_slice_1 = FlatSlice 3 [FlatDimIndex 2 2, FlatDimIndex 2 1] - -test_flatSlice_rotate_slice_iota :: [TestTree] -test_flatSlice_rotate_slice_iota = - singleton $ - testCase "flatSlice . rotate . slice . iota " $ - compareOps $ - flatSlice (rotate (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 0 5 2]) [2, 3]) flat_slice_1 - where - flat_slice_1 = FlatSlice 1 [FlatDimIndex 2 2] - test_flatSlice_transpose_slice_iota :: [TestTree] test_flatSlice_transpose_slice_iota = singleton $ testCase "flatSlice . transpose . slice . iota " $ compareOps $ flatSlice (permute (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 0 5 2]) [1, 0]) flat_slice_1 - where - flat_slice_1 = FlatSlice 1 [FlatDimIndex 2 2] - -test_rotate_flatSlice_transpose_slice_iota :: [TestTree] -test_rotate_flatSlice_transpose_slice_iota = - singleton $ - testCase "flatSlice . transpose . slice . iota " $ - compareOps $ - rotate (flatSlice (permute (slice (iota [20, 20]) $ Slice [DimSlice 1 5 2, DimSlice 1 5 2]) [1, 0]) flat_slice_1) [2, 1] where flat_slice_1 = FlatSlice 1 [FlatDimIndex 2 2] diff --git a/unittests/Futhark/IR/Mem/IxFunWrapper.hs b/unittests/Futhark/IR/Mem/IxFunWrapper.hs --- a/unittests/Futhark/IR/Mem/IxFunWrapper.hs +++ b/unittests/Futhark/IR/Mem/IxFunWrapper.hs @@ -4,8 +4,8 @@ ( IxFun, iota, permute, - rotate, reshape, + coerce, slice, flatSlice, rebase, @@ -14,13 +14,11 @@ import qualified Futhark.IR.Mem.IxFun as I import qualified Futhark.IR.Mem.IxFun.Alg as IA -import Futhark.IR.Syntax (FlatSlice, ShapeChange, Slice) +import Futhark.IR.Syntax (FlatSlice, Slice) import Futhark.Util.IntegralExp type Shape num = [num] -type Indices num = [num] - type Permutation = [Int] type IxFun num = (I.IxFun num, IA.IxFun num) @@ -38,19 +36,19 @@ IxFun num permute (l, a) x = (I.permute l x, IA.permute a x) -rotate :: +reshape :: (Eq num, IntegralExp num) => IxFun num -> - Indices num -> + Shape num -> IxFun num -rotate (l, a) x = (I.rotate l x, IA.rotate a x) +reshape (l, a) x = (I.reshape l x, IA.reshape a x) -reshape :: +coerce :: (Eq num, IntegralExp num) => IxFun num -> - ShapeChange num -> + Shape num -> IxFun num -reshape (l, a) x = (I.reshape l x, IA.reshape a x) +coerce (l, a) x = (I.coerce l x, IA.coerce a x) slice :: (Eq num, IntegralExp num) => diff --git a/unittests/Futhark/IR/Prop/ReshapeTests.hs b/unittests/Futhark/IR/Prop/ReshapeTests.hs --- a/unittests/Futhark/IR/Prop/ReshapeTests.hs +++ b/unittests/Futhark/IR/Prop/ReshapeTests.hs @@ -5,93 +5,40 @@ ) where -import Control.Applicative import Futhark.IR.Prop.Constants import Futhark.IR.Prop.Reshape import Futhark.IR.Syntax import Test.Tasty import Test.Tasty.HUnit -import Test.Tasty.QuickCheck -import Prelude -tests :: TestTree -tests = - testGroup "ReshapeTests" $ - fuseReshapeTests - ++ informReshapeTests - ++ reshapeOuterTests - ++ reshapeInnerTests - ++ [ fuseReshapeProp, - informReshapeProp - ] - -fuseReshapeTests :: [TestTree] -fuseReshapeTests = - [ testCase (unwords ["fuseReshape ", show d1, show d2]) $ - fuseReshape (d1 :: ShapeChange Int) d2 @?= dres -- type signature to avoid warning - | (d1, d2, dres) <- - [ ([DimCoercion 1], [DimNew 1], [DimCoercion 1]), - ([DimNew 1], [DimCoercion 1], [DimNew 1]), - ([DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2]), - ([DimNew 1, DimNew 2], [DimCoercion 1, DimNew 2], [DimNew 1, DimNew 2]) - ] - ] - -informReshapeTests :: [TestTree] -informReshapeTests = - [ testCase (unwords ["informReshape ", show shape, show sc, show sc_res]) $ - informReshape (shape :: [Int]) sc @?= sc_res -- type signature to avoid warning - | (shape, sc, sc_res) <- - [ ([1, 2], [DimNew 1, DimNew 3], [DimCoercion 1, DimNew 3]), - ([2, 2], [DimNew 1, DimNew 3], [DimNew 1, DimNew 3]) - ] - ] - reshapeOuterTests :: [TestTree] reshapeOuterTests = [ testCase (unwords ["reshapeOuter", show sc, show n, show shape, "==", show sc_res]) $ - reshapeOuter (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res + reshapeOuter (intShape sc) n (intShape shape) @?= intShape sc_res | (sc, n, shape, sc_res) <- - [ ([DimNew 1], 1, [4, 3], [DimNew 1, DimCoercion 3]), - ([DimNew 1], 2, [4, 3], [DimNew 1]), - ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 2, DimNew 2, DimNew 3]), - ([DimNew 2, DimNew 2], 2, [4, 3], [DimNew 2, DimNew 2]) + [ ([1], 1, [4, 3], [1, 3]), + ([1], 2, [4, 3], [1]), + ([2, 2], 1, [4, 3], [2, 2, 3]), + ([2, 2], 2, [4, 3], [2, 2]) ] ] reshapeInnerTests :: [TestTree] reshapeInnerTests = [ testCase (unwords ["reshapeInner", show sc, show n, show shape, "==", show sc_res]) $ - reshapeInner (intShapeChange sc) n (intShape shape) @?= intShapeChange sc_res + reshapeInner (intShape sc) n (intShape shape) @?= intShape sc_res | (sc, n, shape, sc_res) <- - [ ([DimNew 1], 1, [4, 3], [DimCoercion 4, DimNew 1]), - ([DimNew 1], 0, [4, 3], [DimNew 1]), - ([DimNew 2, DimNew 2], 1, [4, 3], [DimNew 4, DimNew 2, DimNew 2]), - ([DimNew 2, DimNew 2], 0, [4, 3], [DimNew 2, DimNew 2]) + [ ([1], 1, [4, 3], [4, 1]), + ([1], 0, [4, 3], [1]), + ([2, 2], 1, [4, 3], [4, 2, 2]), + ([2, 2], 0, [4, 3], [2, 2]) ] ] intShape :: [Int] -> Shape intShape = Shape . map (intConst Int32 . toInteger) -intShapeChange :: ShapeChange Int -> ShapeChange SubExp -intShapeChange = map (fmap $ intConst Int32 . toInteger) - -fuseReshapeProp :: TestTree -fuseReshapeProp = testProperty "fuseReshape result matches second argument" prop - where - prop :: ShapeChange Int -> ShapeChange Int -> Bool - prop sc1 sc2 = map newDim (fuseReshape sc1 sc2) == map newDim sc2 - -informReshapeProp :: TestTree -informReshapeProp = testProperty "informReshape result matches second argument" prop - where - prop :: [Int] -> ShapeChange Int -> Bool - prop sc1 sc2 = map newDim (informReshape sc1 sc2) == map newDim sc2 - -instance Arbitrary d => Arbitrary (DimChange d) where - arbitrary = - oneof - [ DimNew <$> arbitrary, - DimCoercion <$> arbitrary - ] +tests :: TestTree +tests = + testGroup "ReshapeTests" $ + reshapeOuterTests ++ reshapeInnerTests