Flatten hierarchy of primops

Author: dougalm

src/lib/Builder.hs

@@ -59,7 +59,7 @@ emit e = case e of
 {-# INLINE emit #-}
 
 idExpr :: Atom n -> Expr n
-idExpr x = PrimOp (getType x) (UnOp Identity x)
+idExpr x = PrimOp (getType x) Identity [x]
 
 declsToExpr :: RNest Decl n l -> Atom l -> Expr n
 declsToExpr (RNest ds (Let b e)) (Var v _) | v == binderName b = maybeBlock ds e

src/lib/ConcreteSyntax.hs

@@ -460,11 +460,11 @@ leafGroup = leafGroup' >>= appendPostfixGroups
       '['  -> cBrackets
       '\"' -> toCLeaf CString <$> strLit
       '\'' -> toCLeaf CChar   <$> charLit
-      '%'  -> do
-        WithSrc sid name <- primName
-        case strToPrimName name of
-          Just prim -> WithSrcs sid [] <$> CPrim prim <$> argList
-          Nothing   -> fail $ "Unrecognized primitive: " ++ bs2str name
+      '%'  -> undefined
+        -- WithSrc sid name <- primName
+        -- case strToPrimName name of
+        --   Just prim -> WithSrcs sid [] <$> CPrim prim <$> argList
+        --   Nothing   -> fail $ "Unrecognized primitive: " ++ bs2str name
       _ | isDigit next -> (    toCLeaf CNat   <$> natLit
                            <|> toCLeaf CFloat <$> doubleLit)
       '\\' -> withSrcs (cNullaryLam <|> cLam)

src/lib/Inference.hs

@@ -541,9 +541,7 @@ bottomUpExplicit (WithSrcE sid expr) = case expr of
 --   UPrim UProjNewtype [x] -> do
 --     x' <- bottomUp x >>= unwrapNewtype
 --     return $ SigmaAtom Nothing x'
-  UPrim prim xs -> do
-    xs' <- mapM bottomUp xs
-    liftM (SigmaAtom Nothing) $ matchPrimApp prim xs'
+  UPrim prim xs -> throw sid $ MiscTypeErr "primitive ops must have result type annotations"
   -- UNatLit l -> liftM (SigmaAtom Nothing) $ fromNatLit sid l NatTy
   -- UIntLit l -> liftM (SigmaAtom Nothing) $ fromIntLit sid l (BaseTy $ Scalar Int32Type)
   UFloatLit x -> return $ SigmaAtom Nothing $ CLit  $ Float32Lit $ realToFrac x
@@ -988,66 +986,8 @@ checkOrInferApp appSrcId funSrcId f' posArgs namedArgs reqTy = undefined
 --   when (not $ null unrecognizedNames) do
 --     throw sid $ UnrecognizedOptionalArgs (map pprint unrecognizedNames) (map pprint acceptedNames)
 
-matchPrimApp :: PrimName -> [CAtom o] -> InfererM i o (CExpr o)
-matchPrimApp = \case
---  UNat                -> \case ~[]  -> return $ toAtom $ NewtypeTyCon Nat
---  UFin                -> \case ~[n] -> return $ toAtom $ NewtypeTyCon (Fin n)
---  UBaseType b         -> \case ~[]  -> return $ toAtomR $ BaseType b
---  UNatCon             -> \case ~[x] -> return $ toAtom $ NewtypeCon NatCon x
---  UPrimTC tc -> case tc of
---    S.ProdType -> \ts -> return $ toAtom $ ProdType $ map (fromJust . toMaybeType) ts
---    S.SumType  -> \ts -> return $ toAtom $ SumType  $ map (fromJust . toMaybeType) ts
---    S.RefType  -> \case ~[h, a] -> undefined -- return $ toAtom $ RefType h (fromJust $ toMaybeType a)
---    S.TypeKind -> \case ~[] -> return $ toAtom $ Kind $ TypeKind
---  UCon con -> case con of
---    S.ProdCon -> \xs -> return $ toAtom $ ProdCon xs
---    S.SumCon _ -> error "not supported"
- MiscOp op -> \xs -> matchMiscOp op xs
- -- MemOp  op -> \xs -> CPrimOp <$> MemOp  <$> matchGenericOp op xs
- UnOp  op ()    -> \case ~[x]    -> return $ CPrimOp (getCType x) $ UnOp op x
- BinOp op () () -> \case ~[x, y] -> return $ CPrimOp (getCType x) $ BinOp op x y
---  UMGet      -> \case ~[r]    -> emitRefOp r MGet
---  UMPut      -> \case ~[r, x] -> emitRefOp r $ MPut x
---  UIndexRef  -> \case ~[r, i] -> indexRef r i
---  UApplyMethod i -> \case ~(d:args) -> emit =<< mkApplyMethod (fromJust $ toMaybeDict d) i args
---  ULinearize -> \case ~[f, x]  -> do f' <- lam1 f; emitHof $ Linearize f' x
---  UTranspose -> \case ~[f, x]  -> do f' <- lam1 f; emitHof $ Transpose f' x
---  p -> \case xs -> throwInternal $ "Bad primitive application: " ++ show (p, xs)
-
-cUnitTy :: CType n
-cUnitTy = CTyCon $ CProdType []
-
-matchMiscOp :: MiscOp () -> [CAtom o] -> InfererM i o (CExpr o)
-matchMiscOp = \case
-  DebugPrintInt () -> \case ~[x] -> return $ CPrimOp cUnitTy $ MiscOp $ DebugPrintInt x
---                            where
--- --    lam2 :: Fallible m => CAtom n -> m (LamExpr n)
--- --    lam2 x = do
--- --      ExplicitCoreLam (BinaryNest b1 b2) body <- return x
--- --      return $ BinaryLamExpr b1 b2 body
-
--- --    lam1 :: Fallible m => CAtom n -> m (LamExpr n)
--- --    lam1 x = do
--- --      ExplicitCoreLam (UnaryNest b) body <- return x
--- --      return $ UnaryLamExpr b body
-
---    matchGenericOp :: Functor op => op () -> [CAtom n] -> InfererM i n (op (CAtom n))
---    matchGenericOp op xs = undefined
--- do
---      (tyArgs, dataArgs) <- partitionEithers <$> forM xs \x -> do
---        case getType x of
---          TyCon (Kind TypeKind) -> do
---            Just x' <- return $ toMaybeType x
---            return $ Left x'
---          _ -> return $ Right x
---      let tyArgs' = case tyArgs of
---            [] -> Nothing
---            [t] -> Just t
---            _ -> error "Expected at most one type arg"
---      return $ fromJust $ toOp $ GenericOpRep op tyArgs' dataArgs
-
--- pattern ExplicitCoreLam :: Nest CBinder n l -> CExpr l -> CAtom n
--- pattern ExplicitCoreLam bs body <- Con (Lam (CoreLamExpr _ (LamExpr bs body)))
+matchPrimApp :: PrimOp -> [CAtom o] -> InfererM i o (CExpr o)
+matchPrimApp = undefined
 
 -- -- === n-ary applications ===
 

src/lib/QueryTypePure.hs

@@ -28,23 +28,6 @@ litType v = case v of
   Float32Lit _ -> Scalar Float32Type
   PtrLit ty _  -> PtrType ty
 
-typeBinOp :: BinOp -> BaseType -> BaseType
-typeBinOp binop xTy = case binop of
-  IAdd   -> xTy;  ISub   -> xTy
-  IMul   -> xTy;  IDiv   -> xTy
-  IRem   -> xTy;
-  ICmp _ -> Scalar Word8Type
-  FAdd   -> xTy;  FSub   -> xTy
-  FMul   -> xTy;  FDiv   -> xTy;
-  FPow   -> xTy
-  FCmp _ -> Scalar Word8Type
-  BAnd   -> xTy;  BOr    -> xTy
-  BXor   -> xTy
-  BShL   -> xTy;  BShR   -> xTy
-
-typeUnOp :: UnOp -> BaseType -> BaseType
-typeUnOp = const id  -- All unary ops preserve the type of the input
-
 getKind :: Type n -> Kind
 getKind = undefined
 -- getKind = \case
@@ -105,7 +88,7 @@ instance HasType Expr where
   getType = \case
     Block ty _ -> ty
     TopApp ty _ _ -> ty
-    PrimOp ty _ -> ty
+    PrimOp ty _ _ -> ty
     Case ty _ _ -> ty
     For _ _ -> undefined
     While _ -> undefined
@@ -130,7 +113,7 @@ instance HasCType CExpr where
     CBlock ty _ -> ty
     CVar   _ ty -> ty
     CLit   l    -> CTyCon $ CBaseType $ litType l
-    CPrimOp ty _ -> ty
+    CPrimOp ty _ _ -> ty
     -- CTyCon  (CTyCon n)
     -- Lam         (CoreLamExpr n)
     -- NewtypeCon  (NewtypeCon n) (CExpr n)

src/lib/Simplify.hs

@@ -373,10 +373,10 @@ simplifyLam (CoreLamExpr _ ab) = go ab
 simplifyExpr :: Emits o => CExpr i -> SimplifyM i o (Atom o)
 simplifyExpr = \case
   CLit val -> return $ Lit val
-  CPrimOp ty op -> do
-    op' <- mapM simplifyExpr op
+  CPrimOp ty op xs -> do
+    xs' <- mapM simplifyExpr xs
     ty <- simplifyType ty
-    emit $ PrimOp ty op'
+    emit $ PrimOp ty op xs'
   e -> error $ show e
 
 simplifyType :: CType i -> SimplifyM i o (Type o)

src/lib/ToLLVM.hs

@@ -18,6 +18,7 @@ import Types.Simple
 import Types.Primitives
 import PPrint
 
+import Err
 import Debug.Trace
 import QueryTypePure
 import Util
@@ -52,13 +53,13 @@ data TranslateState i = TranslateState
 type TranslateSubst i = Subst (LiftE L.Operand) i VoidS
 
 newtype TranslateM (i::S) (a:: *) =
-  TranslateM { inner :: State (TranslateState i) a }
-  deriving (Functor, Applicative, Monad)
+  TranslateM { inner :: StateT (TranslateState i) Except a }
+  deriving (Functor, Applicative, Monad, MonadFail)
 
 runTranslateM :: Monad m => TranslateM VoidS a -> m (TranslateState VoidS)
 runTranslateM cont = do
   let initState = TranslateState [] [] (L.Name "__entry__") 0 voidSubst
-  return $ execState cont.inner initState
+  return $ ignoreExcept $ execStateT cont.inner initState
 
 emitInstr :: L.Type -> L.Instruction -> TranslateM i L.Operand
 emitInstr resultTy instr = do
@@ -76,7 +77,7 @@ extendEnv :: NameBinder i i' -> L.Operand -> TranslateM i' a -> TranslateM i a
 extendEnv b x cont = TranslateM do
   prevState <- get
   let subst' = prevState.subst <>> (b @> LiftE x)
-  let (ans, newState) = runState (cont.inner) $ updateSubst prevState subst'
+  let (ans, newState) = ignoreExcept $ runStateT (cont.inner) $ updateSubst prevState subst'
   put $ updateSubst newState prevState.subst
   return ans
 
@@ -117,10 +118,10 @@ toLLVMEntryFun' (TopLamExpr (Abs Empty body)) = do
 trExpr :: Expr i -> TranslateM i L.Operand
 trExpr = \case
   Block resultTy block -> trBlock block
-  PrimOp resultTy op -> do
+  PrimOp resultTy op xs -> do
     resultTy' <- trType resultTy
-    op' <- forM op trAtom
-    trPrimOp resultTy' op'
+    xs' <- forM xs trAtom
+    trPrimOp resultTy' op xs'
 
 trType :: Type i -> TranslateM i L.Type
 trType = \case
@@ -142,11 +143,12 @@ trBlock (Abs decls result) = case decls of
     val <- trExpr expr
     extendEnv b val $ trBlock $ Abs rest result
 
-trPrimOp :: L.Type -> PrimOp L.Operand -> TranslateM i L.Operand
-trPrimOp resultTy op = case op of
-  BinOp b x y -> case b of
-    FAdd -> emitInstr resultTy $ L.FAdd x y
-  MiscOp op' -> case op' of
-    DebugPrintInt x -> do
-      emitStatement $ L.Call floatTy  "printfloat" [x]
-      return unitOperand
+trPrimOp :: L.Type -> PrimOp -> [L.Operand] -> TranslateM i L.Operand
+trPrimOp resultTy op xs = case op of
+  FAdd -> do
+    [x, y] <- return xs
+    emitInstr resultTy $ L.FAdd x y
+  DebugPrintInt -> do
+    [x] <- return xs
+    emitStatement $ L.Call floatTy  "printfloat" [x]
+    return unitOperand

src/lib/Types/Complicated.hs

@@ -28,7 +28,7 @@ data CExpr (n::S) =
    CBlock  (CType n) (CBlock n)
  | CVar    (Name n) (CType n)
  | CLit    LitVal
- | CPrimOp (CType n) (PrimOp (CExpr n))
+ | CPrimOp (CType n) PrimOp [CExpr n]
  | CTyCon  (CTyCon n)
  | Lam         (CoreLamExpr n)
  | NewtypeCon  (NewtypeCon n) (CExpr n)
@@ -179,7 +179,7 @@ instance Pretty (CExpr n) where
     CBlock  _ b -> pr b
     CVar v _ -> pr v
     CLit l -> pr l
-    CPrimOp _ op -> pr op
+    CPrimOp _ op args -> app (pr op) (map pr args)
     CTyCon _ -> undefined
     Lam _ -> undefined
     NewtypeCon _ _ -> undefined