fei-base-1.0.0: cmdline/Main.hs
module Main where
import Prelude
import RIO (first, forM, on, zipWithM_, (<>))
import RIO.Char (isAlphaNum, isSpace, isUpper,
toLower)
import RIO.Directory
import RIO.FilePath (FilePath, (</>))
import RIO.List (sortBy)
import qualified RIO.Text as T
import RIO.Writer (Writer, execWriter, tell)
import Data.Tuple.Ops (uncons)
import Language.Haskell.Exts
import Options.Applicative hiding (optional)
import System.Log.Logger
import Text.ParserCombinators.ReadP
import Text.Printf (printf)
import MXNet.Base.Raw.Common
import MXNet.Base.Raw.DataIter
import MXNet.Base.Raw.NDArray
import MXNet.Base.Raw.Symbol
_module_ = "Main"
data Arguments = Arguments
{ output_dir :: FilePath
}
args_spec = Arguments
<$> strOption (long "output" <> short 'o' <> value "tmp" <> metavar "OUTPUT-DIR")
main = do
updateGlobalLogger _module_ (setLevel INFO)
args <- execParser opts
let base = output_dir args </> "MXNet" </> "Base" </> "Operators"
createDirectoryIfMissing True base
ops <- mxSymbolListAtomicSymbolCreators
op_names <- mapM getOpName ops
let ops_sorted = map snd $ sortBy (compare `on` fst) $ zip op_names ops
infoM _module_ "Generating NDArray/Symbol operators..."
ops <- concat <$> mapM genTensorOp ops_sorted
writeFile (base </> "Tensor.hs") $ prettyPrint (modTensor ops)
dataitercreators <- mxListDataIters
infoM _module_ "Generating DataIter..."
dataiters <- concat <$> mapM genDataIter (zip dataitercreators [0..])
let base = output_dir args </> "MXNet" </> "Base"
writeFile (base </> "DataIter.hs") $ prettyPrint (modDataIter dataiters)
where
opts = info (args_spec <**> helper) (fullDesc <> progDesc "Generate MXNet operators")
modTensor = Module () (Just $ ModuleHead () (ModuleName () "MXNet.Base.Operators.Tensor") Nothing Nothing) []
[ simpleImport "RIO"
, simpleImport "RIO.List"
, simpleImport "MXNet.Base.Raw"
, simpleImport "MXNet.Base.Spec.Operator"
, simpleImport "MXNet.Base.Spec.HMap"
, simpleImport "MXNet.Base.Tensor"
, simpleImportVars "Data.Maybe" ["catMaybes", "fromMaybe"]]
modDataIter = Module () (Just $ ModuleHead () (ModuleName () "MXNet.Base.DataIter") Nothing Nothing) []
[ simpleImport "RIO"
, simpleImport "RIO.List"
, simpleImportVars "RIO.List.Partial" ["(!!)"]
, simpleImport "MXNet.Base.Raw"
, simpleImport "MXNet.Base.Spec.Operator"
, simpleImport "MXNet.Base.Spec.HMap"
, simpleImportVars "Data.Maybe" ["catMaybes", "fromMaybe"]]
getOpName :: AtomicSymbolCreator -> IO String
getOpName = fmap T.unpack . mxSymbolGetAtomicSymbolName
tensorVar = tyVarIdent "t"
tensorVar2 = tyVarIdent "o"
makeParamInst :: String -> [ResolvedType] -> Bool -> Decl ()
makeParamInst symname typs symbolapi =
TypeInsDecl () (tyApp (tyApp (tyCon $ unQual $ name "ParameterList") (tyPromotedStr symname)) tensorVar)
(tyPromotedList paramList)
where
paramList = map (\(name, typ1, typ2) -> tyPromotedTuple [tyPromotedStr name, tyApp typ1 typ2]) typs
data GenFlag = GenSymbolOp
| GenNDArrayReturn
| GenNDArrayUpdate
makeSignature :: String -> Type () -> Type () -> [Asst ()] -> Decl ()
makeSignature symname ti to extra_constraints =
let cxfullfill = appA (name "Fullfilled") [tyPromotedStr symname, ti, tyVarIdent "a"]
cxtensor = appA (name "TensorOp") [ti, to]
callstack = appA (name "HasCallStack") []
cx_all = cxTuple (cxtensor : cxfullfill : callstack : extra_constraints)
fun = tyFun (let hmap = tyCon $ unQual $ name "ArgsHMap"
in hmap `tyApp` tyPromotedStr symname
`tyApp` ti
`tyApp` tyVarIdent "a")
((tyCon $ unQual $ name "TensorApply") `tyApp ` to)
vars = map (unkindedVar . name) $ "a" : "t" : if ti == to then [] else ["o"]
in tySig [name symname] $ tyForall vars cx_all fun
genTensorOp :: AtomicSymbolCreator -> IO [Decl ()]
genTensorOp sc = do
(symname_t, _, argname_t, argtype_t, _, key_var_num_args_t, rettyp) <- mxSymbolGetAtomicSymbolInfo sc
let symname = T.unpack symname_t
argname = map T.unpack argname_t
argtype = map T.unpack argtype_t
key_var_num_args = T.unpack key_var_num_args_t
if symname `elem` ["_Native", "_NDArray"] then
return []
else do
let symname_ = normalizeName symname
out_new = symname `elem` ["Cast"]
out_var = if out_new then tensorVar2 else tensorVar
(errs, scalarTypes, tensorKeyTypes, tensorVarTypes) = execWriter $ zipWithM_ (resolveHaskellType ResolveTensor) argname argtype
if not (null errs) then do
forM errs $ \(name, msg) ->
errorM _module_ (printf "Function: %s %s" symname msg)
return []
else if (length tensorVarTypes >= 2) then do
errorM _module_ (printf "Function: %s has more than one NDArray/Symbol[] argument." symname)
return []
else if (not (null tensorKeyTypes) && not (null tensorVarTypes)) then do
errorM _module_ (printf "Function: %s has both NDArray/Symbol and NDArray/Symbol[] arguments." symname)
return []
else do
let paramListInst = makeParamInst symname_ (scalarTypes ++ tensorKeyTypes ++ tensorVarTypes) True
sig = makeSignature symname_ tensorVar out_var $
-- the "Custom" is a little special, because it allow extra arguments
if symname == "Custom"
-- PopKey (ArgOf "_Custom(symbol)") args "data",
-- Dump (PopResult (ArgOf "_Custom(symbol)") args "data"))
then let argOfCustom = tyParen $ (tyCon $ unQual $ name "ArgOf") `tyApp`
(tyPromotedStr "_Custom") `tyApp`
tensorVar
in [appA (name "PopKey") [argOfCustom, tyVarIdent "a", tyPromotedStr "data"],
appA (name "Dump") [ tyParen $
(tyCon $ unQual $ name "PopResult") `tyApp`
argOfCustom `tyApp`
tyVarIdent "a" `tyApp`
tyPromotedStr "data"]]
else []
let fun = sfun (name symname_) [name "args"] (UnGuardedRhs () body) Nothing
make_scalar_values =
if symname == "Custom"
-- dump (pop args #data)
then function "dump" `app` (
function "pop" `app`
(var $ name "args") `app`
OverloadedLabel () "data")
-- catMaybes
-- [("KEY",) . showValue <$> (args !? #KEY :: Maybe VALUE_TYPE),
-- ... ]
else function "catMaybes" `app` listE
([ infixApp (infixApp (tupleSection [Just $ strE argkey, Nothing])
(op $ sym ".")
(function "showValue"))
(op $ sym "<$>") $
ExpTypeSig () (infixApp (var $ name "args")
(op $ sym "!?")
(OverloadedLabel () argkey))
(tyApp (tyCon $ unQual $ name "Maybe") typ)
| (argkey, _, typ) <- scalarTypes] ++
if null key_var_num_args || null tensorVarTypes || key_var_num_args `elem` [k | (k,_,_) <- scalarTypes]
then []
else [(con $ unQual $ name "Just") `app` tuple [
strE key_var_num_args,
function "showValue" `app` (function "length" `app` (var $ name "tensorVarArgs"))]])
body = letE ([
patBind (pvar $ name "scalarArgs") make_scalar_values
, patBind (pvar $ name "tensorKeyArgs") (function "catMaybes"
`app` listE [
infixApp (tupleSection [Just $ strE argkey, Nothing]) (op $ sym "<$>") $
ExpTypeSig ()
(infixApp (var $ name "args") (op $ sym "!?") (OverloadedLabel () argkey))
(tyApp (tyCon $ unQual $ name "Maybe") typ) | (argkey, _, typ) <- tensorKeyTypes])
] ++
case tensorVarTypes of
[(argkey,_,_)] ->
[patBind (pvar $ name "tensorVarArgs") $ function "fromMaybe"
`app` eList
`app` ExpTypeSig ()
(infixApp (var $ name "args") (op $ sym "!?") (OverloadedLabel () argkey))
(tyApp (tyCon $ unQual $ name "Maybe") (tyList tensorVar))]
_ -> [])
(function "apply"
`app` (strE symname)
`app` (var $ name "scalarArgs")
`app` (if null tensorVarTypes
then function "Left" `app` (expTypeSig
(var $ name "tensorKeyArgs")
(tyList $ tyTuple [tyCon $ unQual $ name "Text", tensorVar]))
else function "Right" `app` (var $ name "tensorVarArgs")))
return [paramListInst, sig, fun]
genDataIter :: (DataIterCreator, Integer) -> IO [Decl ()]
genDataIter (dataitercreator, index) = do
(diname_t, _, argnames_t, argtypes_t, _) <- mxDataIterGetIterInfo dataitercreator
let diname = normalizeName $ T.unpack diname_t
argnames = map T.unpack argnames_t
argtypes = map T.unpack argtypes_t
(errs, scalarTypes, _, _) = execWriter $ zipWithM_ (resolveHaskellType ResolveDataIter) argnames argtypes
-- parameter list
paramList = map (\(name, typ1, typ2) -> tyPromotedTuple [tyPromotedStr name, tyApp typ1 typ2]) scalarTypes
paramInst = TypeInsDecl ()
(tyApp (tyApp (tyCon $ unQual $ name "ParameterList") (tyPromotedStr diname)) (tyVarIdent "dummy"))
(tyPromotedList paramList)
-- signature
void = tyTuple []
cxfullfill = appA (name "Fullfilled") [tyPromotedStr diname, void, tyVarIdent "a"]
tyfun = tyFun
(let hmap = tyCon $ unQual $ name "ArgsHMap"
in hmap `tyApp` tyPromotedStr diname `tyApp` void `tyApp` tyVarIdent "a")
(tyApp (tyCon $ unQual $ name "IO") (tyCon $ unQual $ name "DataIterHandle"))
tysig = tySig [name diname] $ tyForall [unkindedVar (name "a")] (cxSingle cxfullfill) tyfun
-- function
fun = sfun (name diname) [name "args"] (UnGuardedRhs () body) Nothing
body = letE ([
patBind (pvar $ name "allargs") (function "catMaybes"
`app` listE [
infixApp (infixApp (tupleSection [Just $ strE argkey, Nothing]) (op $ sym ".") (function "showValue")) (op $ sym "<$>") $
ExpTypeSig () (infixApp (var $ name "args") (op $ sym "!?") (OverloadedLabel () argkey)) (tyApp (tyCon $ unQual $ name "Maybe") typ) | (argkey, _, typ) <- scalarTypes])
, patBind (pTuple [pvar $ name "keys", pvar $ name "vals"]) (app (function "unzip") $ var $ name "allargs")
]) (doE $ [
genStmt (pvar $ name "dis") $ function "mxListDataIters",
genStmt (pvar $ name "di") $ function "return" `app` (infixApp (var $ name "dis") (op $ sym "!!") (intE index)),
qualStmt $ function "mxDataIterCreateIter" `app` (var $ name "di") `app` (var $ name "keys") `app` (var $ name "vals")
])
if not (null errs) then do
forM errs $ \(name, msg) ->
errorM _module_ (printf "Function: %s %s" diname msg)
return []
else
return [paramInst, tysig, fun]
normalizeName :: String -> String
normalizeName name = '_': name
uncapitalize :: String -> String
uncapitalize [] = []
uncapitalize (x:xs) = toLower x: xs
data ParamDesc = ParamDescItem String
| ParamDescList Bool [String]
deriving (Eq, Show)
data ResolveMode = ResolveTensor
| ResolveDataIter
deriving Eq
type ResolvedType = (String, Type (), Type ())
resolveHaskellType :: ResolveMode -> String -> String -> Writer ([(String, String)], [ResolvedType], [ResolvedType], [ResolvedType]) ()
resolveHaskellType mode symname desc = do
let fail msg = tell ([(symname, msg)], [], [], [])
case readP_to_S typedesc desc of
[(fields, "")] -> do
let required = ParamDescItem "required" `elem` fields
attr = tyCon $ unQual $ name $ if required then "AttrReq" else "AttrOpt"
scalar hstyp = tell ([], [(symname_, attr, hstyp)], [], [])
tensor_karg hstyp = tell ([], [], [(symname_, attr, hstyp)], [])
tensor_varg hstyp = tell ([], [], [], [(symname_, attr, hstyp)])
symname_ = uncapitalize symname
--
-- tensor operators
--
handleTensor (ParamDescItem "NDArray") = tensor_karg $ tensorVar
handleTensor (ParamDescItem "NDArray-or-Symbol") = tensor_karg $ tensorVar
handleTensor (ParamDescItem "NDArray-or-Symbol[]") = tensor_varg $ tyList $ tensorVar
handleTensor (ParamDescItem "Symbol") = tensor_karg $ tensorVar
handleTensor (ParamDescItem "Symbol[]") = tensor_varg $ tyList $ tensorVar
handleTensor (ParamDescItem "Symbol or Symbol[]") = tensor_varg $ tyList $ tensorVar
handleTensor t = fallThrough t
fallThrough t = fail $ printf "Unknown type: arg %s(%s)." symname desc
case head fields of
ParamDescItem "Shape(tuple)" -> scalar $ tyList $ tyCon $ unQual $ name "Int"
ParamDescItem "int" -> scalar $ tyCon $ unQual $ name "Int"
ParamDescItem "int (non-negative)" -> scalar $ tyCon $ unQual $ name "Int"
ParamDescItem "long (non-negative)" -> scalar $ tyCon $ unQual $ name "Int"
ParamDescItem "boolean" -> scalar $ tyCon $ unQual $ name "Bool"
ParamDescItem "float" -> scalar $ tyCon $ unQual $ name "Float"
ParamDescItem "double" -> scalar $ tyCon $ unQual $ name "Double"
ParamDescItem "float32" -> scalar $ tyCon $ unQual $ name "Float"
-- real_t (from mshadow) is by default float.
ParamDescItem "real_t" -> scalar $ tyCon $ unQual $ name "Float"
ParamDescItem "string" -> scalar $ tyCon $ unQual $ name "Text"
ParamDescItem "int or None" -> scalar $ tyApp (tyCon $ unQual $ name "Maybe") (tyCon $ unQual $ name "Int")
ParamDescItem "float or None" -> scalar $ tyApp (tyCon $ unQual $ name "Maybe") (tyCon $ unQual $ name "Float")
ParamDescItem "double or None" -> scalar $ tyApp (tyCon $ unQual $ name "Maybe") (tyCon $ unQual $ name "Double")
ParamDescItem "boolean or None" -> scalar $ tyApp (tyCon $ unQual $ name "Maybe") (tyCon $ unQual $ name "Bool")
ParamDescItem "Shape or None" -> scalar $ tyApp (tyCon $ unQual $ name "Maybe") (tyList $ tyCon $ unQual $ name "Int")
ParamDescItem "tuple of <float>" -> scalar $ tyList $ tyCon $ unQual $ name "Float"
ParamDescItem "tuple of <double>" -> scalar $ tyList $ tyCon $ unQual $ name "Double"
ParamDescList hasnone vs -> do
let vsprom = map tyPromotedStr vs
typ1 = tyApp (tyCon $ unQual $ name "EnumType") (tyPromotedList vsprom)
typ2 = tyApp (tyCon $ unQual $ name "Maybe") typ1
scalar $ if hasnone then typ2 else typ1
t | mode == ResolveTensor -> handleTensor t
| mode == ResolveDataIter -> fallThrough t
other -> fail (printf "cannot parse type description: %s. Result: %s" desc (show other))
typedesc = do
-- since 1.3, there are types starting with ',', and it implies 'int' type.
def <- (char ',' >> return [ParamDescItem "int"]) <++ return []
ds <- sepBy (skipSpaces >> (list1 +++ list2 +++ item)) (char ',')
eof
return $ def ++ ds
where
list1 = ParamDescList True <$> between (string "{None,") (char '}') (sepBy (skipSpaces >> strItem) (char ','))
list2 = ParamDescList False <$> between (string "{") (char '}') (sepBy (skipSpaces >> strItem) (char ','))
strItem = between (char '\'') (char '\'') (munch1 (\c -> isAlphaNum c || oneOf c "_-+-./"))
item = ParamDescItem <$> munch1 (\c -> isAlphaNum c || oneOf c " _-+()=[]<>./'")
oneOf :: Eq a => a -> [a] -> Bool
oneOf c wl = c `elem` wl
unQual = UnQual ()
unkindedVar = UnkindedVar ()
tyCon = TyCon ()
tyVarSymbol = TyVar () . Symbol ()
tyVarIdent = TyVar () . Ident ()
tyApp = TyApp ()
tyFun = TyFun ()
tySig names types = TypeSig () names types
tyList = TyList ()
tyTuple = TyTuple () Boxed
tyVar = TyVar ()
tyParen = TyParen ()
tyPromotedInteger s = TyPromoted () (PromotedInteger () s (show s))
tyPromotedStr s = TyPromoted () (PromotedString () s s)
tyPromotedList s = TyPromoted () (PromotedList () True s)
tyPromotedTuple s = TyPromoted () (PromotedTuple () s)
tyForall vars cxt typ = TyForall () vars_ cxt_ typ
where
vars_ = if null vars then Nothing else Just vars
cxt_ = if cxt == CxEmpty () then Nothing else Just cxt
cxSingle = CxSingle ()
cxTuple = CxTuple ()
appA = AppA ()
tupleSection = TupleSection () Boxed
con = Con ()
expTypeSig = ExpTypeSig ()
simpleImport mod = ImportDecl {
importAnn = (),
importModule = ModuleName () mod,
importQualified = False,
importSrc = False,
importSafe = False,
importPkg = Nothing,
importAs = Nothing,
importSpecs = Nothing
}
simpleImportVars mod vars = ImportDecl {
importAnn = (),
importModule = ModuleName () mod,
importQualified = False,
importSrc = False,
importSafe = False,
importPkg = Nothing,
importAs = Nothing,
importSpecs = Just $ ImportSpecList () False [IVar () $ Ident () var | var <- vars]
}