459 lines
17 KiB
Haskell
459 lines
17 KiB
Haskell
-- Copyright 2016 TensorFlow authors.
|
|
--
|
|
-- Licensed under the Apache License, Version 2.0 (the "License");
|
|
-- you may not use this file except in compliance with the License.
|
|
-- You may obtain a copy of the License at
|
|
--
|
|
-- http://www.apache.org/licenses/LICENSE-2.0
|
|
--
|
|
-- Unless required by applicable law or agreed to in writing, software
|
|
-- distributed under the License is distributed on an "AS IS" BASIS,
|
|
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
-- See the License for the specific language governing permissions and
|
|
-- limitations under the License.
|
|
|
|
{-# LANGUAGE CPP #-}
|
|
{-# LANGUAGE FlexibleContexts #-}
|
|
{-# LANGUAGE LambdaCase #-}
|
|
{-# LANGUAGE OverloadedStrings #-}
|
|
{-# LANGUAGE TypeFamilies #-}
|
|
{- | Rendering of TensorFlow operations as Haskell functions.
|
|
|
|
The basic type signature generated for each op is:
|
|
|
|
> {constraints} => {mandatory attrs} -> {input tensors} -> {output tensors}
|
|
|
|
where:
|
|
|
|
* @{mandatory attrs}@ is of the form @A_1 -> ... -> A_N@, where each @A@ is an
|
|
op attribute that doesn't have a default and can't be inferred from other
|
|
inputs.
|
|
|
|
* @{constraints}@ restrict the type parameters of the input and output tensors
|
|
(for example: 'TensorType' or 'OneOf').
|
|
|
|
* @{input tensors}@ is of the form @T_1 -> ... -> T_N@, where each @T@ is of
|
|
the form @Tensor Ref a@ or @Tensor v a@ (or a list of one of those types),
|
|
and @a@ is either a concrete type or a (constrained) type variable.
|
|
|
|
* @{output tensors}@ is of the form @(T_1,...,T_N)@ for "pure" ops, and
|
|
@Build (T_1,...,T_N)@ for "stateful" ops. An op is considered "stateful" if
|
|
it takes a @Tensor Ref@ or @Tensor v ResourceHandle@ as input, or if it's
|
|
explicitly marked \"Stateful\" in its @REGISTER_OP@ definition. (If there
|
|
are no outputs, it is either @ControlNode@ or @Build ControlNode@.)
|
|
-}
|
|
|
|
module TensorFlow.OpGen
|
|
( OpGenFlags(..)
|
|
, docOpList
|
|
, flagParser)
|
|
where
|
|
|
|
import Data.Foldable (toList)
|
|
import Data.Maybe (fromMaybe)
|
|
import Data.ProtoLens.Default(def)
|
|
import Data.ProtoLens (showMessage)
|
|
import Data.List (sortOn)
|
|
import Data.List.NonEmpty (NonEmpty)
|
|
import qualified Data.List.NonEmpty as NE
|
|
import Lens.Family2 ((^.), (.~), (&), view)
|
|
import Options.Applicative (Parser, help, long, strOption, value)
|
|
import Proto.Tensorflow.Core.Framework.OpDef
|
|
( OpList
|
|
, OpDef
|
|
)
|
|
import Proto.Tensorflow.Core.Framework.OpDef_Fields
|
|
( attr
|
|
, inputArg
|
|
, name
|
|
, op
|
|
, outputArg
|
|
)
|
|
import Proto.Tensorflow.Core.Framework.Types (DataType(..))
|
|
import System.FilePath (takeBaseName)
|
|
import TensorFlow.OpGen.ParsedOp
|
|
import Text.PrettyPrint.Mainland
|
|
( Doc
|
|
, (<+>)
|
|
, (</>)
|
|
, (<+/>)
|
|
, brackets
|
|
, comma
|
|
, commasep
|
|
, dquotes
|
|
, empty
|
|
, enclose
|
|
, flatten
|
|
, folddoc
|
|
, hang
|
|
, indent
|
|
, parens
|
|
, sep
|
|
, stack
|
|
, strictText
|
|
, tuple
|
|
)
|
|
import qualified Data.Set as Set
|
|
import qualified Data.Text as Text
|
|
|
|
data OpGenFlags = OpGenFlags
|
|
{ outputFile :: String
|
|
, prefix :: String
|
|
, excludeList :: String
|
|
}
|
|
|
|
flagParser :: Parser OpGenFlags
|
|
flagParser = OpGenFlags
|
|
<$> strOption (mconcat [ long "output"
|
|
, help "File to write."
|
|
])
|
|
<*> strOption (mconcat [ long "prefix"
|
|
, help "Haskell package prefix to use"
|
|
])
|
|
<*> strOption (mconcat [ long "exclude_list"
|
|
, value ""
|
|
, help "Comma separated Ops names to ignore"
|
|
])
|
|
|
|
|
|
docOpList :: OpGenFlags -> OpList -> Doc
|
|
docOpList flags opList =
|
|
stack [ "{-# LANGUAGE ConstraintKinds #-}"
|
|
, "{-# LANGUAGE DataKinds #-}"
|
|
, "{-# LANGUAGE FlexibleContexts #-}"
|
|
, "{-# LANGUAGE FlexibleInstances #-}"
|
|
, "{-# LANGUAGE OverloadedStrings #-}"
|
|
, "{-# LANGUAGE ScopedTypeVariables #-}"
|
|
-- Avoids reports about shadowing standard library names.
|
|
, "{-# OPTIONS_GHC -fno-warn-name-shadowing #-}"
|
|
-- eqLengthGuard never returns false and dies instead.
|
|
, "{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}"
|
|
, "module" <+> strictText moduleName <+> "where"
|
|
, empty
|
|
, imports
|
|
, empty
|
|
, folddoc (\x y -> x </> empty </> y)
|
|
(map renderOpAndExtras $
|
|
sortOn (view name) $
|
|
filter (not . flip elem exclusions . view name) $
|
|
toList $ opList ^. op)
|
|
]
|
|
where moduleName =
|
|
Text.pack (prefix flags) <> "." <> camelCase
|
|
-- Discards the optional trailing _ops_op_lib
|
|
(fromMaybe shortName (Text.stripSuffix "_ops_op_lib" shortName))
|
|
shortName = Text.pack (takeBaseName $ outputFile flags)
|
|
exclusions = Text.splitOn "," $ Text.pack $ excludeList flags
|
|
renderOpAndExtras o = renderOp (parseOp o) </> extras o
|
|
|
|
imports :: Doc
|
|
imports = stack [
|
|
"import Data.ByteString (ByteString)"
|
|
, "import Data.Complex (Complex)"
|
|
, "import Data.Int (Int8, Int16, Int32, Int64)"
|
|
, "import Data.Proxy (Proxy(Proxy))"
|
|
, "import Data.Word (Word8, Word16, Word32, Word64)"
|
|
, "import Lens.Family2 ((.~), (&))"
|
|
, "import TensorFlow.Build"
|
|
, "import TensorFlow.BuildOp"
|
|
, "import TensorFlow.Tensor"
|
|
, "import TensorFlow.Types"
|
|
]
|
|
|
|
renderHaskellName, renderTFName, renderQuotedTFName :: Name -> Doc
|
|
renderHaskellName = strictText . unHaskellName . haskellName
|
|
renderTFName = strictText . unTFName . tfName
|
|
renderQuotedTFName = dquotes . renderTFName
|
|
|
|
|
|
-- | Generate the source code for a single op.
|
|
-- For example:
|
|
--
|
|
-- -- | {haddock comment}
|
|
-- foo :: {type sig}
|
|
-- foo attr1 attr2 input1 input2 | eqLengthGuard [...] = {function body}
|
|
renderOp :: ParsedOp -> Doc
|
|
renderOp pOp = stack $
|
|
[ haddocks
|
|
-- Prevent unreasonably long compilation times on ghc-7.10, due
|
|
-- to stack calling "-dump-hi" which (unnecessarily) includes the
|
|
-- inlining information, and is large for ops with many arguments.
|
|
#if __GLASGOW_HASKELL__ < 800
|
|
, "{-# NOINLINE" <+> n <+> "#-}"
|
|
#endif
|
|
, n <+> "::" <+> hang 0 (typeSig empty pOp)
|
|
, n <+> "=" <+> n <> "' id"
|
|
, n' <+> "::" <+> hang 0 (typeSig "OpParams ->" pOp)
|
|
, n' <+> hang 0 args <+> "|" <+> funcGuard listSizeAttrs
|
|
<+> "=" </> -- args are indented
|
|
-- the body needs to be indented wrt the name
|
|
indent indentation (functionBody pOp)
|
|
] ++ whereClause listSizeAttrs
|
|
where
|
|
n = renderHaskellName $ parsedOpName pOp
|
|
n' = n <> "'"
|
|
listSizeAttrs = inferredListSizeAttrs pOp
|
|
args = sep $ "op'options"
|
|
: (map renderHaskellName
|
|
$ map attrName (explicitInputAttrs pOp)
|
|
++ map parsedArgName (parsedInputs pOp))
|
|
haddocks = "-- |" <+> multilineComment (parsedOpSummary pOp) (parsedOpDescription pOp)
|
|
|
|
-- | A check that all lists of the given size have the given length.
|
|
-- For example:
|
|
-- eqLengthGuard [("N", [("input1", length input1), ("input2", length input2)])]
|
|
funcGuard :: [Attr (NonEmpty Name)] -> Doc
|
|
funcGuard attrs = "eqLengthGuard" <+> brackets (commasep entries)
|
|
where
|
|
entries =
|
|
[ parens $ nAttr <> comma <+>
|
|
brackets (commasep $ toList $
|
|
map renderTensorName (toList $ attrInfo a))
|
|
| a <- attrs
|
|
, let nAttr = renderQuotedTFName (attrName a)
|
|
]
|
|
renderTensorName x = parens $ renderQuotedTFName x <> comma <+>
|
|
"length" <+> renderHaskellName x
|
|
|
|
-- | Define the implicit list length attributes.
|
|
-- For example:
|
|
-- where
|
|
-- n1 = fromIntegral (length input1) :: Int64
|
|
-- n2 = fromIntegral (length input2) :: Int64
|
|
whereClause :: [Attr (NonEmpty Name)] -> [Doc]
|
|
whereClause [] = []
|
|
whereClause as = [indent 2 $ "where" </> indent 2 (stack $ map defineLengthAttr as)]
|
|
where
|
|
defineLengthAttr a = renderHaskellAttrName a <+> "="
|
|
<+> "fromIntegral (length"
|
|
<+> renderHaskellName (NE.head $ attrInfo a)
|
|
<> ") :: Int64"
|
|
|
|
renderHaskellAttrName :: Attr a -> Doc
|
|
renderHaskellAttrName = renderHaskellName . attrName
|
|
|
|
functionBody :: ParsedOp -> Doc
|
|
functionBody pOp
|
|
| parsedOpIsMonadic pOp
|
|
= "build $ do"
|
|
</> indent indentation (bindOpInputsVar
|
|
</> "buildOp" <+> outputListsSizes <+> opDef)
|
|
| otherwise
|
|
= "pureOp" <+> outputListsSizes <+> "$ do"
|
|
</> indent indentation (bindOpInputsVar </> "return" <+> opDef)
|
|
where
|
|
outputListsSizes = brackets $ commasep
|
|
[ renderHaskellName a
|
|
| ParsedArg { parsedArgCase = ListArg { argLength = a } }
|
|
<- parsedOutputs pOp
|
|
]
|
|
opInputsVar = "op'inputs"
|
|
bindOpInputsVar = opInputsVar <+> "<- fmap Prelude.concat $ Prelude.sequence"
|
|
<+> brackets (commasep $ map (\a -> "buildInputs" <+> a) tensorArgs)
|
|
opDef = parens $ hang 0 $ stack $
|
|
"opDef" <+> renderQuotedTFName (parsedOpName pOp) :
|
|
-- Renders type parameter arguments.
|
|
[ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> inferredTypeExpr a
|
|
| a <- inferredTypeAttrs pOp, let n = attrName a
|
|
] ++
|
|
-- Renders mandatory attributes as function parameters.
|
|
[ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
|
|
| a <- explicitInputAttrs pOp, let n = attrName a
|
|
] ++
|
|
-- Renders sizes of tensor list types having number_attr.
|
|
[ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
|
|
| a <- inferredListSizeAttrs pOp, let n = attrName a
|
|
] ++
|
|
["& op'options & opInputs .~" <+> opInputsVar]
|
|
tensorArgs = renderTensorArg <$> parsedInputs pOp
|
|
renderTensorArg = renderHaskellName . parsedArgName
|
|
inferredTypeExpr a
|
|
| typeParamIsList $ attrInfo a
|
|
= "fromTensorTypes (Proxy :: Proxy" <+> renderHaskellAttrName a
|
|
<> ")"
|
|
| otherwise = "tensorType (undefined ::" <+> renderHaskellAttrName a
|
|
<> ")"
|
|
|
|
-- | Write a comment with the inputs/outputs/attributes in proto format, for
|
|
-- debugging.
|
|
extras :: OpDef -> Doc
|
|
extras d = enclose "{-\n" "\n-}" $
|
|
strictText $ Text.pack $
|
|
showMessage ((def :: OpDef)
|
|
& inputArg .~ (d ^. inputArg)
|
|
& outputArg .~ (d ^. outputArg)
|
|
& attr .~ (d ^. attr))
|
|
|
|
-- | The type signature for an op.
|
|
-- Of the form:
|
|
-- forall t1 t2 v1 v2 . (TensorType t1, TensorType t2)
|
|
-- => {pre} Float -> Tensor t1 v1 -> Tensor t2 v2
|
|
-- where "Float" is an explicit input attribute, "Tensor t1 v1" is an input, and
|
|
-- "Tensor t2 v2" is an output.
|
|
typeSig :: Doc -> ParsedOp -> Doc
|
|
typeSig pre pOp = constraints
|
|
<+/> pre </> signatureFold (map attrInput (explicitInputAttrs pOp)
|
|
++ map tensorArgAndComment (parsedInputs pOp)
|
|
++ [outputs])
|
|
where
|
|
constraints
|
|
| null classConstraints = empty
|
|
| otherwise = "forall" <+> sep typeParams <+> "." <+> tuple classConstraints <+> "=>"
|
|
typeParams = [strictText v | k <- parsedInputs pOp ++ parsedOutputs pOp,
|
|
ArgSomeTensor v <- [argKind $ parsedArgCase k]]
|
|
++ [renderHaskellAttrName n | n <- inferredTypeAttrs pOp]
|
|
++ if parsedOpIsMonadic pOp then ["m'"] else []
|
|
-- Use m' as the type parameter to avoid clashing with an attribute name.
|
|
monadConstraint
|
|
| parsedOpIsMonadic pOp = ["MonadBuild m'"]
|
|
| otherwise = []
|
|
classConstraints = monadConstraint ++ map tensorArgConstraint
|
|
(inferredTypeAttrs pOp)
|
|
signatureFold = folddoc (\x y -> x </> "->" <+> y)
|
|
attrInput a = renderAttrType (attrInfo a) <+> hang 0 ("-- ^" <+> attrComment a)
|
|
renderAttrType (AttrSingle a) = renderAttrBaseType a
|
|
renderAttrType (AttrList a) = brackets $ renderAttrBaseType a
|
|
renderAttrBaseType = \case
|
|
AttrBytes -> "ByteString"
|
|
AttrInt64 -> "Data.Int.Int64"
|
|
AttrFloat -> "Float"
|
|
AttrBool -> "Bool"
|
|
AttrType -> "DataType"
|
|
AttrShape -> "Shape"
|
|
AttrTensor -> "TensorProto"
|
|
AttrFunc -> error "AttrFunc not supported"
|
|
|
|
tensorArgAndComment t = tensorArg t <+> hang 0 ("-- ^" <+> argComment t)
|
|
outputs = case parsedOutputs pOp of
|
|
[] -> wrapOutput "ControlNode"
|
|
-- TODO(judahjacobson): To improve indentation: `tensorArgAndComment a`
|
|
[a] -> wrapOutput (tensorArg a) <+> "-- ^" <+> argComment a
|
|
as -> wrapOutput (tuple (map tensorArg as)) <+/> resultComment as
|
|
wrapOutput o
|
|
| parsedOpIsMonadic pOp = "m'" <+> parens o
|
|
| otherwise = o
|
|
|
|
-- | Render an op input or output.
|
|
-- For example: "Tensor Ref Int64", "Tensor v t"
|
|
tensorArg :: ParsedArg -> Doc
|
|
tensorArg p = case parsedArgCase p of
|
|
SimpleArg { argType = t, argKind = k } -> tensorType t k
|
|
ListArg { argType = t, argKind = k } -> brackets $ tensorType t k
|
|
MixedListArg {argTypeAttr = t, argKind = k}
|
|
-> "TensorList" <+> parens (kind k) <+> renderHaskellName t
|
|
where
|
|
kind k = case k of
|
|
ArgTensorRef -> "Ref"
|
|
ArgTensorValue -> "Value"
|
|
ArgTensorBuild -> "Build"
|
|
ArgSomeTensor v -> strictText v
|
|
tensorType t k = let
|
|
a = case t of
|
|
ArgTypeFixed dt -> strictText $ dtTypeToHaskell dt
|
|
ArgTypeAttr n -> renderHaskellName n
|
|
in "Tensor" <+> kind k <+> a
|
|
|
|
attrComment :: Attr a -> Doc
|
|
attrComment a = argComment' (attrName a) (attrDescription a)
|
|
|
|
argComment :: ParsedArg -> Doc
|
|
argComment a = argComment' (parsedArgName a) (parsedArgDescription a)
|
|
|
|
argComment' :: Name -> Text.Text -> Doc
|
|
argComment' argName argDesc =
|
|
bold (renderTFName argName) <> splitMultilineText (":" <+>) argDesc
|
|
|
|
bold :: Doc -> Doc
|
|
bold n = "__" <> n <> "__"
|
|
|
|
-- | Comment for the outputs of an op.
|
|
-- For example:
|
|
-- -- ^ (__output1__, __output2__)
|
|
-- --
|
|
-- -- * __output1__: description1
|
|
-- --
|
|
-- -- * __output2__: description2
|
|
resultComment :: [ParsedArg] -> Doc
|
|
resultComment os = stack $ flatten commentSummary : map commentDetails os
|
|
where
|
|
commentSummary = "-- ^" <+> tuple [bold (renderTFName $ parsedArgName o) | o <- os]
|
|
commentDetails o =
|
|
stack [ "--"
|
|
, "-- *" <+> argComment o
|
|
]
|
|
|
|
-- | Constraints for a given type parameter.
|
|
-- E.g.: "TensorType t" or "OneOf [Int64, Float] t"
|
|
-- or "TensorTypes ts" or "OneOfs [..] ts".
|
|
tensorArgConstraint :: Attr TypeParam -> Doc
|
|
tensorArgConstraint a = case attrInfo a of
|
|
TypeParam False Nothing -> "TensorType" <+> n
|
|
TypeParam False (Just as) -> "OneOf" <+> typeList as <+> n
|
|
TypeParam True Nothing -> "TensorTypes" <+> n
|
|
TypeParam True (Just as) -> "OneOfs" <+> typeList as <+> n
|
|
where
|
|
n = renderHaskellAttrName a
|
|
-- Produces a type-level list, e.g.: '[Int32,Int64,Float]
|
|
typeList = ("'" <>) . brackets . commasep . map strictText .
|
|
Set.toList . Set.fromList .
|
|
map dtTypeToHaskell . toList
|
|
|
|
-- NOTE: The cases of this function should be kept in sync with
|
|
-- TensorFlow.Types.AllTensorTypes.
|
|
dtTypeToHaskell :: DataType -> Text.Text
|
|
dtTypeToHaskell DT_BOOL = "Bool"
|
|
dtTypeToHaskell DT_BFLOAT16 = "Data.Word.Word16"
|
|
dtTypeToHaskell DT_COMPLEX128 = "(Data.Complex.Complex Double)"
|
|
dtTypeToHaskell DT_COMPLEX64 = "(Data.Complex.Complex Float)"
|
|
dtTypeToHaskell DT_DOUBLE = "Double"
|
|
dtTypeToHaskell DT_FLOAT = "Float"
|
|
dtTypeToHaskell DT_INT16 = "Data.Int.Int16"
|
|
dtTypeToHaskell DT_INT32 = "Data.Int.Int32"
|
|
dtTypeToHaskell DT_INT64 = "Data.Int.Int64"
|
|
dtTypeToHaskell DT_INT8 = "Data.Int.Int8"
|
|
dtTypeToHaskell DT_QINT32 = "Data.Int.Int32" -- TODO(gnezdo): make unique
|
|
dtTypeToHaskell DT_QINT8 = "Data.Word.Word8" -- TODO(gnezdo): make unique
|
|
dtTypeToHaskell DT_QINT16 = "Data.Int.Int16" -- TODO(gnezdo): make unique
|
|
dtTypeToHaskell DT_QUINT16 = "Data.Word.Word16" -- TODO(gnezdo): make unique
|
|
dtTypeToHaskell DT_QUINT8 = "Data.Word.Word8" -- TODO(gnezdo): make unique
|
|
dtTypeToHaskell DT_STRING = "Data.ByteString.ByteString"
|
|
dtTypeToHaskell DT_UINT16 = "Data.Word.Word16"
|
|
dtTypeToHaskell DT_UINT32 = "Data.Word.Word32"
|
|
dtTypeToHaskell DT_UINT64 = "Data.Word.Word64"
|
|
dtTypeToHaskell DT_HALF = "Data.Word.Word16" -- TODO(gnezdo): make unique
|
|
dtTypeToHaskell DT_UINT8 = "Data.Word.Word8"
|
|
dtTypeToHaskell DT_RESOURCE = "ResourceHandle"
|
|
dtTypeToHaskell DT_VARIANT = "Variant"
|
|
dtTypeToHaskell x =
|
|
Text.pack $ "Unsupported type in dtTypeToHaskell: " ++ show x
|
|
|
|
-- | haddockComment escapes TensorFlow doc strings into haddock.
|
|
-- TODO(gnezdo): deal with the markup.
|
|
haddockComment :: Text.Text -> Doc
|
|
haddockComment = strictText
|
|
|
|
-- | Generate a multiline comment. For example:
|
|
-- summary'
|
|
-- --
|
|
-- -- detail_line1
|
|
-- -- detail_line2
|
|
-- -- ...
|
|
multilineComment :: Text.Text -> Text.Text -> Doc
|
|
multilineComment summary' detail =
|
|
haddockComment summary' </>
|
|
splitMultilineText insertParagraphAndComment detail
|
|
where insertParagraphAndComment x = "--" </> "--" <+> x
|
|
|
|
-- | Converts the given multi-line detail string into
|
|
-- a multi-line haddock. Applies the given lead to the
|
|
-- first line. Returns an empty document for empty detail.
|
|
splitMultilineText :: (Doc -> Doc) -> Text.Text -> Doc
|
|
splitMultilineText lead detail =
|
|
case Text.lines detail of
|
|
[] -> empty
|
|
(l : ls) -> stack $ lead (haddockComment l)
|
|
: map (("--" <+>) . haddockComment) ls
|
|
|
|
indentation :: Int
|
|
indentation = 4
|