Add fast path for large arrays. (#2162)

Add fast path for large arrays.

Currently only handles one dimensional arrays, but could be extended.
Since this is only useful for generated code, I don't think wider
applicability is important.

CHANGELOG.md

@@ -11,7 +11,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 * New prelude function: `rep`, an implicit form of `replicate`.
 
-* Improved handling of large array literals in the C backends (#2160).
+* Improved handling of large monomorphic single-dimensional array
+  literals (#2160).
 
 ### Removed
 

docs/language-reference.rst

@@ -743,6 +743,14 @@ This holds only if ``n`` is a variable or constant.
 Create an array containing the indicated elements.  Each element must
 have the same type and shape.
 
+**Large array optimisation**: as a special case, large one-dimensional
+array literal consisting *entirely* of monomorphic constants (i.e.,
+numbers must have a type suffix) are handled with specialised
+fast-path code by the compiler. To keep compile times manageable, make
+sure that all very large array literals (more than about ten thousand
+elements) are of this form. This is likely relevant only for generated
+code.
+
 .. _range:
 
 ``x..y...z``

src/Futhark/AD/Rev.hs

@@ -107,6 +107,9 @@ diffBasicOp pat aux e m =
     Assert {} ->
       void $ commonBasicOp pat aux e m
     --
+    ArrayVal {} ->
+      void $ commonBasicOp pat aux e m
+    --
     ArrayLit elems _ -> do
       (_pat_v, pat_adj) <- commonBasicOp pat aux e m
       t <- lookupType pat_adj

src/Futhark/Analysis/AccessPattern.hs

@@ -465,6 +465,7 @@ analyseBasicOp ctx expression pats =
   let ctx_val = case expression of
         SubExp se -> varInfoFromSubExp se
         Opaque _ se -> varInfoFromSubExp se
+        ArrayVal _ _ -> (varInfoFromNames ctx mempty) {variableType = ConstType}
         ArrayLit ses _t -> concatVariableInfos mempty ses
         UnOp _ se -> varInfoFromSubExp se
         BinOp _ lsubexp rsubexp -> concatVariableInfos mempty [lsubexp, rsubexp]

src/Futhark/Analysis/Metrics.hs

@@ -123,6 +123,7 @@ expMetrics (Op op) =
 basicOpMetrics :: BasicOp -> MetricsM ()
 basicOpMetrics (SubExp _) = seen "SubExp"
 basicOpMetrics (Opaque _ _) = seen "Opaque"
+basicOpMetrics ArrayVal {} = seen "ArrayVal"
 basicOpMetrics ArrayLit {} = seen "ArrayLit"
 basicOpMetrics BinOp {} = seen "BinOp"
 basicOpMetrics UnOp {} = seen "UnOp"

src/Futhark/CLI/Misc.hs

@@ -128,7 +128,7 @@ mainTokens = mainWithOptions () [] "program" $ \args () ->
         Just (Right (Left e)) -> do
           hPrint stderr e
           exitWith $ ExitFailure 2
-        Just (Right (Right (tokens, _))) ->
+        Just (Right (Right tokens)) ->
           mapM_ printToken tokens
     _ -> Nothing
   where

src/Futhark/CodeGen/ImpGen.hs

@@ -954,17 +954,19 @@ defCompileBasicOp (Pat [pe]) (Concat i (x :| ys) _) = do
         destslice = skip_slices ++ [DimSlice (tvExp offs_glb) rows 1]
     copyDWIM (patElemName pe) destslice (Var y) []
     offs_glb <-- tvExp offs_glb + rows
+defCompileBasicOp (Pat [pe]) (ArrayVal vs t) = do
+  dest_mem <- entryArrayLoc <$> lookupArray (patElemName pe)
+  static_array <- newVNameForFun "static_array"
+  emit $ Imp.DeclareArray static_array t $ Imp.ArrayValues vs
+  let static_src =
+        MemLoc static_array [intConst Int64 $ fromIntegral $ length vs] $
+          LMAD.iota 0 [fromIntegral $ length vs]
+  addVar static_array $ MemVar Nothing $ MemEntry DefaultSpace
+  copy t dest_mem static_src
 defCompileBasicOp (Pat [pe]) (ArrayLit es _)
   | Just vs@(v : _) <- mapM isLiteral es = do
-      dest_mem <- entryArrayLoc <$> lookupArray (patElemName pe)
       let t = primValueType v
-      static_array <- newVNameForFun "static_array"
-      emit $ Imp.DeclareArray static_array t $ Imp.ArrayValues vs
-      let static_src =
-            MemLoc static_array [intConst Int64 $ fromIntegral $ length es] $
-              LMAD.iota 0 [fromIntegral $ length es]
-      addVar static_array $ MemVar Nothing $ MemEntry DefaultSpace
-      copy t dest_mem static_src
+      defCompileBasicOp (Pat [pe]) (ArrayVal vs t)
   | otherwise =
       forM_ (zip [0 ..] es) $ \(i, e) ->
         copyDWIMFix (patElemName pe) [fromInteger i] e []

src/Futhark/IR/Parse.hs

@@ -314,6 +314,10 @@ pBasicOp =
           <*> pFlatSlice
           <* lexeme "="
           <*> pVName,
+      try $
+        ArrayVal
+          <$> brackets (pPrimValue `sepBy` pComma)
+          <*> (lexeme ":" *> "[]" *> pPrimType),
       ArrayLit
         <$> brackets (pSubExp `sepBy` pComma)
         <*> (lexeme ":" *> "[]" *> pType),

src/Futhark/IR/Pretty.hs

@@ -198,6 +198,11 @@ instance Pretty BasicOp where
       <+> colon
       <+> "[]"
       <> pretty rt
+  pretty (ArrayVal vs t) =
+    brackets (commasep $ map pretty vs)
+      <+> colon
+      <+> "[]"
+      <> pretty t
   pretty (BinOp bop x y) = pretty bop <> parens (pretty x <> comma <+> pretty y)
   pretty (CmpOp op x y) = pretty op <> parens (pretty x <> comma <+> pretty y)
   pretty (ConvOp conv x) =

src/Futhark/IR/Prop/Aliases.hs

@@ -57,6 +57,7 @@ subExpAliases (Var v) = vnameAliases v
 basicOpAliases :: BasicOp -> [Names]
 basicOpAliases (SubExp se) = [subExpAliases se]
 basicOpAliases (Opaque _ se) = [subExpAliases se]
+basicOpAliases (ArrayVal _ _) = [mempty]
 basicOpAliases (ArrayLit _ _) = [mempty]
 basicOpAliases BinOp {} = [mempty]
 basicOpAliases ConvOp {} = [mempty]

src/Futhark/IR/Prop/TypeOf.hs

@@ -63,6 +63,10 @@ basicOpType (SubExp se) =
   pure <$> subExpType se
 basicOpType (Opaque _ se) =
   pure <$> subExpType se
+basicOpType (ArrayVal vs t) =
+  pure [arrayOf (Prim t) (Shape [n]) NoUniqueness]
+  where
+    n = intConst Int64 $ toInteger $ length vs
 basicOpType (ArrayLit es rt) =
   pure [arrayOf rt (Shape [n]) NoUniqueness]
   where

src/Futhark/IR/Syntax.hs

@@ -328,6 +328,13 @@ data BasicOp
   | -- | Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@.
     -- Second arg is the element type of the rows of the array.
     ArrayLit [SubExp] Type
+  | -- | A one-dimensional array literal that contains only constants.
+    -- This is a fast-path for representing very large array literals
+    -- that show up in some programs. The key rule for processing this
+    -- in compiler passes is that you should never need to look at the
+    -- individual elements. Has exactly the same semantics as an
+    -- 'ArrayLit'.
+    ArrayVal [PrimValue] PrimType
   | -- | Unary operation.
     UnOp UnOp SubExp
   | -- | Binary operation.

src/Futhark/IR/Traversals.hs

@@ -89,6 +89,8 @@ mapExpM ::
   m (Exp trep)
 mapExpM tv (BasicOp (SubExp se)) =
   BasicOp <$> (SubExp <$> mapOnSubExp tv se)
+mapExpM _ (BasicOp (ArrayVal vs t)) =
+  pure $ BasicOp $ ArrayVal vs t
 mapExpM tv (BasicOp (ArrayLit els rowt)) =
   BasicOp
     <$> ( ArrayLit
@@ -279,6 +281,8 @@ walkOnLambda tv (Lambda params ret body) = do
 walkExpM :: (Monad m) => Walker rep m -> Exp rep -> m ()
 walkExpM tv (BasicOp (SubExp se)) =
   walkOnSubExp tv se
+walkExpM _ (BasicOp ArrayVal {}) =
+  pure ()
 walkExpM tv (BasicOp (ArrayLit els rowt)) =
   mapM_ (walkOnSubExp tv) els >> walkOnType tv rowt
 walkExpM tv (BasicOp (BinOp _ x y)) =

src/Futhark/IR/TypeCheck.hs

@@ -837,6 +837,9 @@ checkBasicOp (SubExp es) =
   void $ checkSubExp es
 checkBasicOp (Opaque _ es) =
   void $ checkSubExp es
+checkBasicOp ArrayVal {} =
+  -- We assume this is never changed, so no need to check it.
+  pure ()
 checkBasicOp (ArrayLit [] _) =
   pure ()
 checkBasicOp (ArrayLit (e : es') t) = do

src/Futhark/Internalise/Defunctionalise.hs

@@ -507,6 +507,8 @@ defuncExp (RecordLit fs loc) = do
         _ ->
           let tp = Info $ structTypeFromSV sv
            in pure (RecordFieldImplicit vn tp loc', (baseName vn, sv))
+defuncExp e@(ArrayVal vs t loc) =
+  pure (ArrayVal vs t loc, Dynamic $ toParam Observe $ typeOf e)
 defuncExp (ArrayLit es t@(Info t') loc) = do
   es' <- mapM defuncExp' es
   pure (ArrayLit es' t loc, Dynamic $ toParam Observe t')

src/Futhark/Internalise/Exps.hs

@@ -628,6 +628,9 @@ internaliseExp desc (E.RecordLit orig_fields _) =
           (baseName name)
           (E.Var (E.qualName name) t loc)
           loc
+internaliseExp desc (E.ArrayVal vs t _) =
+  fmap pure . letSubExp desc . I.BasicOp $
+    I.ArrayVal (map internalisePrimValue vs) (internalisePrimType t)
 internaliseExp desc (E.ArrayLit es (Info arr_t) loc)
   -- If this is a multidimensional array literal of primitives, we
   -- treat it specially by flattening it out followed by a reshape.

src/Futhark/Internalise/FullNormalise.hs

@@ -172,9 +172,18 @@ getOrdering _ (RecordLit fs loc) = do
       pure $ RecordFieldExplicit n e' floc
     f (RecordFieldImplicit v t _) =
       f $ RecordFieldExplicit (baseName v) (Var (qualName v) t loc) loc
-getOrdering _ (ArrayLit es ty loc) = do
-  es' <- mapM (getOrdering False) es
-  pure $ ArrayLit es' ty loc
+getOrdering _ (ArrayVal vs t loc) =
+  pure $ ArrayVal vs t loc
+getOrdering _ (ArrayLit es ty loc)
+  | Just vs <- mapM isLiteral es,
+    Info (Array _ (Shape [_]) (Prim t)) <- ty =
+      pure $ ArrayVal vs t loc
+  | otherwise = do
+      es' <- mapM (getOrdering False) es
+      pure $ ArrayLit es' ty loc
+  where
+    isLiteral (Literal v _) = Just v
+    isLiteral _ = Nothing
 getOrdering _ (Project n e ty loc) = do
   e' <- getOrdering False e
   pure $ Project n e' ty loc

src/Futhark/Internalise/Monomorphise.hs

@@ -638,6 +638,8 @@ transformExp (RecordLit fs loc) =
           (baseName v)
           (Var (qualName v) t' loc)
           loc
+transformExp (ArrayVal vs t loc) =
+  pure $ ArrayVal vs t loc
 transformExp (ArrayLit es t loc) =
   ArrayLit <$> mapM transformExp es <*> traverse transformType t <*> pure loc
 transformExp (AppExp e res) =

src/Futhark/Optimise/ArrayShortCircuiting/DataStructs.hs

@@ -336,6 +336,7 @@ createsNewArrOK (BasicOp Iota {}) = True
 createsNewArrOK (BasicOp Manifest {}) = True
 createsNewArrOK (BasicOp Concat {}) = True
 createsNewArrOK (BasicOp ArrayLit {}) = True
+createsNewArrOK (BasicOp ArrayVal {}) = True
 createsNewArrOK (BasicOp Scratch {}) = True
 createsNewArrOK _ = False
 

src/Futhark/Optimise/ArrayShortCircuiting/MemRefAggreg.hs

@@ -93,6 +93,9 @@ getUseSumFromStm td_env coal_tab (Let _ _ (BasicOp (Index arr (Slice slc))))
     isFix _ = False
 getUseSumFromStm _ _ (Let Pat {} _ (BasicOp Index {})) = Just ([], []) -- incomplete slices
 getUseSumFromStm _ _ (Let Pat {} _ (BasicOp FlatIndex {})) = Just ([], []) -- incomplete slices
+getUseSumFromStm td_env coal_tab (Let (Pat pes) _ (BasicOp ArrayVal {})) =
+  let wrts = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) pes
+   in Just (wrts, wrts)
 getUseSumFromStm td_env coal_tab (Let (Pat pes) _ (BasicOp (ArrayLit ses _))) =
   let rds = mapMaybe (getDirAliasedIxfn td_env coal_tab) $ mapMaybe seName ses
       wrts = mapMaybe (getDirAliasedIxfn td_env coal_tab . patElemName) pes

src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs

@@ -482,6 +482,9 @@ graphStm stm = do
       -- Whether the rows are primitive constants or arrays, without any scalar
       -- variable operands such ArrayLit cannot directly prevent a scalar read.
       graphHostOnly e
+    BasicOp ArrayVal {} ->
+      -- As above.
+      graphHostOnly e
     BasicOp Update {} ->
       graphHostOnly e
     BasicOp Concat {} ->

src/Language/Futhark/FreeVars.hs

@@ -54,6 +54,7 @@ freeInExp expr = case expr of
     where
       freeInExpField (RecordFieldExplicit _ e _) = freeInExp e
       freeInExpField (RecordFieldImplicit vn t _) = ident $ Ident vn t mempty
+  ArrayVal {} -> mempty
   ArrayLit es t _ ->
     foldMap freeInExp es <> freeInType (unInfo t)
   AppExp (Range e me incl _) _ ->

src/Language/Futhark/Interpreter.hs

@@ -966,6 +966,9 @@ eval env (ArrayLit (v : vs) _ _) = do
   v' <- eval env v
   vs' <- mapM (eval env) vs
   pure $ toArray' (valueShape v') (v' : vs')
+eval _ (ArrayVal vs _ _) =
+  -- Probably will not ever be used.
+  pure $ toArray' ShapeLeaf $ map ValuePrim vs
 eval env (AppExp e (Info (AppRes t retext))) = do
   let t' = expandType env $ toStruct t
   v <- evalAppExp env e

src/Language/Futhark/Parser/Lexer.x

@@ -196,20 +196,19 @@ getToken state@(pos,c,s,n) =
       let x = action (BS.take (n'-n) s)
       x `seq` Right (state', (pos, pos', x))
 
-scanTokens :: Pos -> BS.ByteString -> Either LexerError ([L Token], Pos)
-scanTokens pos str = loop $ initialLexerState pos str
+scanTokens :: Pos -> BS.ByteString -> Either LexerError [L Token]
+scanTokens pos str = fmap reverse $ loop [] $ initialLexerState pos str
   where
-   loop s = do
+   loop toks s = do
      (s', tok) <- getToken s
      case tok of
        (start, end, EOF) ->
-         pure ([], end)
-       (start, end, t) -> do
-         (rest, endpos) <- loop s'
-         pure (L (Loc start end) t : rest, endpos)
+         pure toks
+       (start, end, t) ->
+         loop (L (Loc start end) t:toks) s'
 
 -- | Given a starting position, produce tokens from the given text (or
 -- a lexer error).  Returns the final position.
-scanTokensText :: Pos -> T.Text -> Either LexerError ([L Token], Pos)
+scanTokensText :: Pos -> T.Text -> Either LexerError [L Token]
 scanTokensText pos = scanTokens pos . BS.fromStrict . T.encodeUtf8
 }

src/Language/Futhark/Parser/Monad.hs

@@ -17,6 +17,7 @@ module Language.Futhark.Parser.Monad
     mustBe,
     primNegate,
     applyExp,
+    arrayLitExp,
     patternExp,
     addDocSpec,
     addAttrSpec,
@@ -111,6 +112,17 @@ type ParserMonad = ExceptT SyntaxError (State ParserState)
 arrayFromList :: [a] -> Array Int a
 arrayFromList l = listArray (0, length l - 1) l
 
+arrayLitExp :: [UncheckedExp] -> SrcLoc -> UncheckedExp
+arrayLitExp es loc
+  | Just (v : vs) <- mapM isLiteral es,
+    all ((primValueType v ==) . primValueType) vs =
+      ArrayVal (v : vs) (primValueType v) loc
+  | otherwise =
+      ArrayLit es NoInfo loc
+  where
+    isLiteral (Literal v _) = Just v
+    isLiteral _ = Nothing
+
 applyExp :: NE.NonEmpty UncheckedExp -> ParserMonad UncheckedExp
 applyExp all_es@((Constr n [] _ loc1) NE.:| es) =
   pure $ Constr n es NoInfo (srcspan loc1 (NE.last all_es))

src/Language/Futhark/Parser/Parser.y

@@ -605,8 +605,8 @@ Atom : PrimLit        { Literal (fst $1) (srclocOf (snd $1)) }
      | '(' Exp ')'            { Parens $2 (srcspan $1 $>) }
      | '(' Exp ',' Exps1 ')'  { TupLit ($2 : $4) (srcspan $1 $>) }
      | '('      ')'           { TupLit [] (srcspan $1 $>) }
-     | '[' Exps1 ']'          { ArrayLit $2 NoInfo (srcspan $1 $>) }
-     | '['       ']'          { ArrayLit [] NoInfo (srcspan $1 $>) }
+     | '[' Exps1 ']'          { arrayLitExp $2 (srcspan $1 $>) }
+     | '['       ']'          { arrayLitExp [] (srcspan $1 $>) }
 
      | id { let L loc (ID v)  = $1 in Var (QualName [] v) NoInfo (srclocOf loc) }
 

src/Language/Futhark/Pretty.hs

@@ -353,6 +353,8 @@ prettyExp _ (RecordLit fs _)
   where
     fieldArray (RecordFieldExplicit _ e _) = hasArrayLit e
     fieldArray RecordFieldImplicit {} = False
+prettyExp _ (ArrayVal vs _ _) =
+  brackets (commasep $ map pretty vs)
 prettyExp _ (ArrayLit es t _) =
   brackets (commasep $ map pretty es) <> prettyInst t
 prettyExp _ (StringLit s _) =

src/Language/Futhark/Prop.hs

@@ -463,6 +463,8 @@ typeOf (RecordLit fs _) =
     record (RecordFieldExplicit name e _) = (name, typeOf e)
     record (RecordFieldImplicit name (Info t) _) = (baseName name, t)
 typeOf (ArrayLit _ (Info t) _) = t
+typeOf (ArrayVal vs t loc) =
+  Array mempty (Shape [sizeFromInteger (genericLength vs) loc]) (Prim t)
 typeOf (StringLit vs loc) =
   Array
     mempty

src/Language/Futhark/Syntax.hs

@@ -806,6 +806,12 @@ data ExpBase f vn
   | -- | Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@.
     -- Second arg is the row type of the rows of the array.
     ArrayLit [ExpBase f vn] (f StructType) SrcLoc
+  | -- | Array value constants, where the elements are known to be
+    -- constant primitives. This is a fast-path variant of 'ArrayLit'
+    -- that will never be constructed by the parser, but may result
+    -- from normalisation later on. Has exactly the same semantics as
+    -- an 'ArrayLit'.
+    ArrayVal [PrimValue] PrimType SrcLoc
   | -- | An attribute applied to the following expression.
     Attr (AttrInfo vn) (ExpBase f vn) SrcLoc
   | Project Name (ExpBase f vn) (f StructType) SrcLoc
@@ -877,6 +883,7 @@ instance Located (ExpBase f vn) where
   locOf (RecordLit _ pos) = locOf pos
   locOf (Project _ _ _ pos) = locOf pos
   locOf (ArrayLit _ _ pos) = locOf pos
+  locOf (ArrayVal _ _ loc) = locOf loc
   locOf (StringLit _ loc) = locOf loc
   locOf (Var _ _ loc) = locOf loc
   locOf (Ascript _ _ loc) = locOf loc

src/Language/Futhark/Traversals.hs

@@ -146,6 +146,8 @@ instance ASTMappable (ExpBase Info VName) where
     TupLit <$> mapM (mapOnExp tv) els <*> pure loc
   astMap tv (RecordLit fields loc) =
     RecordLit <$> astMap tv fields <*> pure loc
+  astMap _ (ArrayVal vs t loc) =
+    pure $ ArrayVal vs t loc
   astMap tv (ArrayLit els t loc) =
     ArrayLit <$> mapM (mapOnExp tv) els <*> traverse (mapOnStructType tv) t <*> pure loc
   astMap tv (Ascript e tdecl loc) =
@@ -456,6 +458,7 @@ bareExp (QualParens name e loc) = QualParens name (bareExp e) loc
 bareExp (TupLit els loc) = TupLit (map bareExp els) loc
 bareExp (StringLit vs loc) = StringLit vs loc
 bareExp (RecordLit fields loc) = RecordLit (map bareField fields) loc
+bareExp (ArrayVal vs t loc) = ArrayVal vs t loc
 bareExp (ArrayLit els _ loc) = ArrayLit (map bareExp els) NoInfo loc
 bareExp (Ascript e te loc) = Ascript (bareExp e) (bareTypeExp te) loc
 bareExp (Coerce e te _ loc) = Coerce (bareExp e) (bareTypeExp te) NoInfo loc

src/Language/Futhark/TypeChecker/Consumption.hs

@@ -954,6 +954,7 @@ checkExp e@IntLit {} = noAliases e
 checkExp e@FloatLit {} = noAliases e
 checkExp e@Literal {} = noAliases e
 checkExp e@StringLit {} = noAliases e
+checkExp e@ArrayVal {} = noAliases e
 checkExp e@ArrayLit {} = noAliases e
 checkExp e@Negate {} = noAliases e
 checkExp e@Not {} = noAliases e