diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,203 @@ +Copyright 2016 The TensorFlow Authors. All rights reserved. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2016, The 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. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,3 @@ +import Distribution.Simple + +main = defaultMain diff --git a/src/TensorFlow/OpGen.hs b/src/TensorFlow/OpGen.hs new file mode 100644 --- /dev/null +++ b/src/TensorFlow/OpGen.hs @@ -0,0 +1,452 @@ +-- 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 (def, 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 + , 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)" + , "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" + + 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_HALF = "Data.Word.Word16" -- TODO(gnezdo): make unique +dtTypeToHaskell DT_UINT8 = "Data.Word.Word8" +dtTypeToHaskell DT_RESOURCE = "ResourceHandle" +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 diff --git a/src/TensorFlow/OpGen/ParsedOp.hs b/src/TensorFlow/OpGen/ParsedOp.hs new file mode 100644 --- /dev/null +++ b/src/TensorFlow/OpGen/ParsedOp.hs @@ -0,0 +1,343 @@ +-- | This module helps parse the proto OpDef into a Haskell type which is more +-- descriptive of how the attributes and arguments will be used in the +-- generated code. +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +module TensorFlow.OpGen.ParsedOp + ( ParsedOp(..) + , Name(..) + , HaskellName(..) + , TFName(..) + , Attr(..) + , AttrType(..) + , AttrBaseType(..) + , TypeParam(..) + , ParsedArg(..) + , ParsedArgCase(..) + , ArgType(..) + , ArgKind(..) + , parseOp + , camelCase + ) where + +import Data.Char (toUpper, toLower) +import Data.List (sortBy) +import Data.List.NonEmpty (NonEmpty, nonEmpty) +import Data.Maybe (mapMaybe) +import Data.Monoid ((<>)) +import Data.Ord (comparing) +import qualified Data.Set as Set +import Data.Text (Text) +import qualified Data.Text as Text +import Lens.Family2 ((^.)) +import Proto.Tensorflow.Core.Framework.AttrValue (list) +import Proto.Tensorflow.Core.Framework.OpDef + ( OpDef + , OpDef'ArgDef + , OpDef'AttrDef + , allowedValues + , attr + , maybe'defaultValue + , description + , name + , inputArg + , isRef + , isStateful + , outputArg + , summary + , typeListAttr + , numberAttr + , typeAttr + , type' + ) +import Proto.Tensorflow.Core.Framework.Types (DataType(DT_RESOURCE)) + +data ParsedOp = ParsedOp + { parsedOpName :: Name + , parsedOpSummary :: Text + , parsedOpDescription :: Text + , parsedInputs :: [ParsedArg] + , parsedOutputs :: [ParsedArg] + , explicitInputAttrs :: [Attr AttrType] + -- ^ Attributes that must be set explicitly when creating the op. + -- Associated with the type of the attribute. + , inferredTypeAttrs :: [Attr TypeParam] + -- ^ Attributes that are type parameters. + , inferredListSizeAttrs :: [Attr (NonEmpty Name)] + -- Attributes which are list sizes (ints) that are inferred automatically + -- from one or more of the input tensors. + -- Associated with the list of tensors whose size it describes. + , parsedOpIsMonadic :: Bool + -- ^ Whether this op is stateful or takes a stateful input. Such ops + -- should not be CSE'd and must be monadic in our API (i.e., return a + -- Build action). + } + +data Name = Name + { haskellName :: HaskellName + , tfName :: TFName + } + +-- | A raw name as specified in the OpDef proto. +newtype TFName = TFName { unTFName :: Text } + deriving (Eq, Ord) + +-- | A name that's appropriate for a variable in a Haskell source file. +newtype HaskellName = HaskellName { unHaskellName :: Text } + +-- | A named attribute, associated with some information about it. +data Attr a = Attr + { attrName :: Name + , attrDescription :: Text + , attrInfo :: a + } + +-- | The type of an attribute. +data AttrType = AttrSingle AttrBaseType + | AttrList AttrBaseType + deriving Eq + +data AttrBaseType = AttrBytes | AttrInt64 | AttrFloat | AttrBool + | AttrType | AttrShape | AttrTensor + deriving Eq + +data TypeParam = TypeParam + { typeParamIsList :: Bool + , typeParamRestrictions :: Maybe (NonEmpty DataType) + -- ^ The list of allowed types (see: TensorFlow.Types.OneOf). + -- If 'Nothing', then any type is acceptable. + } + +-- | An input or output argument (Tensor) for an op. +data ParsedArg = ParsedArg + { parsedArgName :: Name + , parsedArgDescription :: Text + , parsedArgCase :: ParsedArgCase + } + +data ParsedArgCase + = SimpleArg { argType :: ArgType, argKind :: ArgKind } + | ListArg + { argLength :: Name -- ^ The attribute that specifies this list's length. + , argType :: ArgType + , argKind :: ArgKind + } + | MixedListArg { argTypeAttr :: Name, argKind :: ArgKind } + -- ^ A heterogeneous list. + +maybeArgType :: ParsedArgCase -> Maybe ArgType +maybeArgType MixedListArg{} = Nothing +maybeArgType a = Just $ argType a + +-- | The type of an argument. +data ArgType + = ArgTypeFixed DataType -- ^ A fixed type. + | ArgTypeAttr Name -- ^ A type that depends on an attribute. + +-- The kind of an op input or output (not including the argument type `a`). +data ArgKind + = ArgTensorRef -- Tensor Ref a + | ArgTensorValue -- Tensor Value a + | ArgTensorBuild -- Tensor Build a + | ArgSomeTensor Text -- Tensor v a; the Text is the variable 'v'. + deriving (Eq) + +isRefCase :: ParsedArgCase -> Bool +isRefCase a + | ArgTensorRef <- argKind a = True + | Just (ArgTypeFixed DT_RESOURCE) <- maybeArgType a = True + | otherwise = False + +makeName :: Text -> Name +makeName n = Name + { haskellName = HaskellName $ fixReservedName $ lowCase n + , tfName = TFName n + } + +-- | Change a name so it doesn't conflict with any Haskell keywords. +fixReservedName :: Text -> Text +fixReservedName n + | n `Set.member` reservedKeywords = n <> "'" + | otherwise = n + +reservedKeywords :: Set.Set Text +reservedKeywords = Set.fromList $ + -- Haskell2010 keywords: + -- https://www.haskell.org/onlinereport/haskell2010/haskellch2.html#x7-180002.4 + -- We don't include keywords that are allowed to be variable names, + -- in particular: "as", "forall", and "hiding". + [ "case" + , "class" + , "data" + , "default" + , "deriving" + , "do" + , "else" + , "foreign" + , "if" + , "import" + , "in" + , "infix" + , "infixl" + , "infixr" + , "instance" + , "let" + , "module" + , "newtype" + , "of" + , "then" + , "type" + , "where" + ] + ++ -- Nonstandard extensions + [ "mdo" -- RecursiveDo + , "rec" -- Arrows, RecursiveDo + , "proc" -- Arrows + ] + +-- | Lower-case the given text. +lowCase :: Text -> Text +lowCase = forceCase toLower + +forceCase :: (Char -> Char) -> Text -> Text +forceCase convert s = maybe "" (\(c, cs) -> Text.cons (convert c) cs) + (Text.uncons s) + +camelCase :: Text -> Text +camelCase s = Text.concat $ map upCase + $ Text.splitOn "_" s + +-- | Upper-case the given text. +upCase :: Text -> Text +upCase = forceCase toUpper + + +parseOp :: OpDef -> ParsedOp +parseOp o = ParsedOp + { parsedOpName = makeName $ o ^. name + , parsedOpSummary = o ^. summary + , parsedOpDescription = o ^. description + , .. + } + where + parsedOpIsMonadic = o ^. isStateful + || any (isRefCase . parsedArgCase) parsedInputs + || null (o ^. outputArg) + parsedInputs = zipWith (\t a -> parseArg a (inputTensorKind t a)) + tensorKindParams (o ^. inputArg) + tensorKindParams = ["v'" <> Text.pack (show x) | x <- [1::Integer ..]] + parsedOutputs = map (\a -> parseArg a (outputTensorKind parsedOpIsMonadic a)) + (o ^. outputArg) + -- Integer attributes that can be inferred from the size of at least one + -- input list. + inferredListSizeAttrs = mapMaybeAttrs (getInferredListSizeAttr parsedInputs) + $ o ^. attr + implicitAttrs = Set.fromList $ map tfName $ + map attrName inferredTypeAttrs + ++ map attrName inferredListSizeAttrs + inferredTypeAttrs = mapMaybeAttrs (getInferredTypeAttr argTypeParams) $ o ^. attr + argTypeParams = Set.fromList $ map tfName $ + mapMaybe (getArgTypeParam . parsedArgCase) $ + parsedInputs ++ parsedOutputs + -- Attributes that can't be inferred and don't have defaults, so must be + -- passed as separate arguments to the op. + explicitInputAttrs = sortBy (comparing (tfName . attrName)) + $ mapMaybeAttrs (getExplicitInputAttr o implicitAttrs) + $ o ^. attr + +-- TODO(judahjacobson): Some arguments should be refs. +inputTensorKind :: Text -> OpDef'ArgDef -> ArgKind +inputTensorKind v a + | a ^. isRef = ArgTensorRef + | otherwise = ArgSomeTensor v + +outputTensorKind :: Bool -> OpDef'ArgDef -> ArgKind +outputTensorKind isMonadic a + | a ^. isRef = ArgTensorRef + | isMonadic = ArgTensorValue + | otherwise = ArgTensorBuild + +getExplicitInputAttr :: OpDef -> Set.Set TFName -> OpDef'AttrDef -> Maybe AttrType +getExplicitInputAttr o implicitAttrs a + | TFName (a ^. name) `Set.notMember` implicitAttrs + , a ^. maybe'defaultValue == Nothing + , t <- parseAttrType o (a ^. type') + , t `elem` map AttrSingle + [AttrBool, AttrInt64, AttrFloat, AttrType, AttrShape] + ++ [AttrList AttrType] = Just t + | otherwise = Nothing + +getInferredTypeAttr :: Set.Set TFName -> OpDef'AttrDef -> Maybe TypeParam +getInferredTypeAttr argTypeParams a + | TFName (a ^. name) `notElem` argTypeParams = Nothing + | a ^. type' == "type" = Just $ TypeParam False allowed + | a ^. type' == "list(type)" = Just $ TypeParam True allowed + | otherwise = Nothing + where + allowed = nonEmpty (a ^. allowedValues . list . type') + +getArgTypeParam :: ParsedArgCase -> Maybe Name +getArgTypeParam SimpleArg { argType = ArgTypeAttr n} = Just n +getArgTypeParam ListArg { argType = ArgTypeAttr n} = Just n +getArgTypeParam MixedListArg { argTypeAttr = n } = Just n +getArgTypeParam _ = Nothing + +getInferredListSizeAttr :: [ParsedArg] -> OpDef'AttrDef -> Maybe (NonEmpty Name) +getInferredListSizeAttr inputs a + | a ^. type' == "int" + = nonEmpty [t | ParsedArg { parsedArgName = t + , parsedArgCase + = ListArg { argLength = n } + } <- inputs + , TFName (a ^. name) == tfName n] + | otherwise = Nothing + +-- | Like mapMaybe, but associates the attribute name/description with the given info. +mapMaybeAttrs :: (OpDef'AttrDef -> Maybe a) -> [OpDef'AttrDef] -> [Attr a] +mapMaybeAttrs f = mapMaybe $ \a -> do + x <- f a + Just Attr + { attrName = makeName (a ^. name) + , attrDescription = a ^. description + , attrInfo = x + } + +parseArg :: OpDef'ArgDef -> ArgKind -> ParsedArg +parseArg a tKind = ParsedArg + { parsedArgName = makeName (a ^. name) + , parsedArgDescription = a ^. description + , parsedArgCase = parseArgCase a tKind + } + +parseArgCase :: OpDef'ArgDef -> ArgKind -> ParsedArgCase +parseArgCase a tKind + | Just n <- maybeAttr (a ^. typeListAttr) = MixedListArg n tKind + | Just n <- maybeAttr (a ^. numberAttr) = ListArg n thisArgType tKind + | otherwise = SimpleArg thisArgType tKind + where + thisArgType + | Just n <- maybeAttr (a ^. typeAttr) = ArgTypeAttr n + | otherwise = ArgTypeFixed (a ^. type') + maybeAttr :: Text -> Maybe Name + maybeAttr "" = Nothing + maybeAttr t = Just $ makeName t + +parseAttrType :: OpDef -> Text -> AttrType +parseAttrType o = \case + "string" -> AttrSingle AttrBytes + "int" -> AttrSingle AttrInt64 + "float" -> AttrSingle AttrFloat + "bool" -> AttrSingle AttrBool + "type" -> AttrSingle AttrType + "shape" -> AttrSingle AttrShape + "tensor" -> AttrSingle AttrTensor + "list(string)" -> AttrList AttrBytes + "list(int)" -> AttrList AttrInt64 + "list(float)" -> AttrList AttrFloat + "list(bool)" -> AttrList AttrBool + "list(type)" -> AttrList AttrType + "list(shape)" -> AttrList AttrShape + "list(tensor)" -> AttrList AttrTensor + t -> error $ "parseAttrType: unrecognized type " ++ show t + ++ " for op " ++ show (o ^. name) diff --git a/tensorflow-opgen.cabal b/tensorflow-opgen.cabal new file mode 100644 --- /dev/null +++ b/tensorflow-opgen.cabal @@ -0,0 +1,34 @@ +name: tensorflow-opgen +version: 0.1.0.0 +synopsis: Code generation for TensorFlow operations. +description: Please see README.md +homepage: https://github.com/tensorflow/haskell#readme +license: Apache +license-file: LICENSE +author: TensorFlow authors +maintainer: tensorflow-haskell@googlegroups.com +copyright: Google Inc. +category: Machine Learning +build-type: Simple +cabal-version: >=1.22 + +library + hs-source-dirs: src + exposed-modules: TensorFlow.OpGen.ParsedOp + , TensorFlow.OpGen + build-depends: proto-lens == 0.2.* + , tensorflow-proto == 0.1.* + , base >= 4.7 && < 5 + , bytestring + , containers + , filepath + , lens-family + , mainland-pretty + , optparse-applicative + , semigroups + , text + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/tensorflow/haskell