content: -- Copyright 2021 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  module MLIR.Lower (coreToMLIR) where
  
  import qualified MLIR.AST as AST
  import qualified MLIR.AST.Builder as AST
  import qualified MLIR.AST.Dialect.LLVM   as LLVM
  import qualified MLIR.AST.Dialect.Std    as Std
  import qualified MLIR.AST.Dialect.Tensor as Tensor
  
  import Data.Functor
  import Control.Monad.Reader
  import Control.Monad.State.Strict
  import GHC.Stack
  import qualified Data.Map.Strict    as M
  import qualified Data.Text.Encoding as T
  
  import Cat (extendR)
  import Env
  import Syntax
  import PPrint
  import Type
  import Util (bindM2)
  import Builder
  
  
  data MLIRAtomF a = Value a
                   | Pair (MLIRAtomF a) (MLIRAtomF a)
                   | Unit
                     deriving (Functor, Foldable, Traversable)
  type MLIRAtom = MLIRAtomF AST.Value
  
  type MonadLower m = (MonadReader (Env MLIRAtom) m, AST.MonadBlockBuilder m)
  
  
  coreToMLIR :: Block -> (AST.Operation, Abs (Nest Binder) Atom)
  coreToMLIR block@(Block decls resultExpr) = (moduleOp, resultRecon)
    where
      attrs = AST.namedAttribute "llvm.emit_c_interface" AST.UnitAttr
      (moduleOp, Just resultRecon) = flip runReader mempty $ flip runStateT Nothing $
        AST.buildModule $ do
          AST.buildSimpleFunction "entry" [] attrs $ do
            -- NB: All args are inlined in simplification for now
            unless (null $ freeVars block) $ error "MLIR backend can't handle free variables"
            _rawArgsPtr   <- AST.blockArgument $ LLVM.Ptr i64
            rawResultsPtr <- AST.blockArgument $ LLVM.Ptr i64
            blockEnv <- lowerDecls decls
            (resultVals, recon) <- case resultExpr of
              -- The special case on Atom lets us handle results that are potentially
              -- more complicated than what MLIRAtom can represent (for example
              -- type-class dictionaries).
              Atom a -> do
                let resultVars = bindingsAsVars $ freeVars a
                let resultMAtom = foldr Pair Unit $ (blockEnv !) <$> resultVars
                (resultMValues, Abs bs consList) <-
                      toCoreAtom (mkConsListTy $ varType <$> resultVars) resultMAtom
                let Right resultAtoms = fromConsList consList
                return ( resultMValues
                       , Abs bs $ subst (newEnv resultVars resultAtoms, mempty) a)
              _ -> toCoreAtom (getType resultExpr) =<< extendR blockEnv (lowerExpr resultExpr)
            put $ Just recon  -- Smuggle the reconstruction out of this block
            forM_ (zip [0..] resultVals) \(i, result) -> do
              iv <- Std.constant i32 $ AST.IntegerAttr i32 i
              rawResultPtr <- LLVM.getelementptr (LLVM.Ptr i64) rawResultsPtr [iv]
              resultPtr    <- LLVM.bitcast (LLVM.Ptr (AST.typeOf result)) rawResultPtr
              LLVM.store result resultPtr
            Std.return []
  
  
  lowerDecls :: MonadLower m => Nest Decl -> m (Env MLIRAtom)
  lowerDecls decls = case decls of
      Empty    -> ask
      Nest d t -> withLoweredDecl d $ lowerDecls t
    where
      withLoweredDecl :: MonadLower m => Decl -> m a -> m a
      withLoweredDecl (Let _ b bound) m = do
        vals <- lowerExpr bound
        extendR (b @> vals) m
  
  
  lowerExpr :: MonadLower m => Expr -> m MLIRAtom
  lowerExpr expr = case expr of
    App  _ _   -> error "Applications (array indexing) not supported in MLIR backend"
    Case _ _ _ -> error "Case expressions not supported in MLIR backend"
    Atom a     -> lowerAtom a
    Op   op    -> lowerOp op
    Hof  _     -> error "Higher-order-functions not supported in MLIR backend"
  
  
  lowerOp :: MonadLower m => Op -> m MLIRAtom
  lowerOp op = case op of
    TabCon (TabTy _ elTy@(BaseTy _)) els ->
      liftM Value . Tensor.from_elements tensorType =<< traverse lowerValue els
      where tensorType = AST.RankedTensorType [Just $ length els] (toMLIRType elTy) Nothing
    TabCon _ _ -> error "Multi-dimensional array literals not supported in MLIR backend"
    ScalarBinOp bop x y -> liftM Value $ bindM2 (getBinOpLowering bop) (lowerValue x) (lowerValue y)
    ScalarUnOp  uop x   -> liftM Value $ getUnOpLowering uop =<< lowerValue x
    CastOp      ty  x   -> do
      xv <- lowerValue x
      Value <$> case (ty, getType x) of
        (BaseTy (PtrType (Heap CPU, _)), BaseTy (Scalar i)) | Just _ <- intWidth i -> LLVM.inttoptr targetTy xv
        (BaseTy (Scalar st), BaseTy (Scalar ss)) -> case (st, ss) of
          -- Integral <-> Integral casts
          (t, s) | Just (_, tw) <- intWidth t, Just (_, sw) <- intWidth s, tw < sw  -> Std.trunci targetTy xv
          (t, s) | Just (_, tw) <- intWidth t, Just (_, sw) <- intWidth s, tw == sw -> return xv
          (t, s) | Just (True , _) <- intWidth t, Just (True , _) <- intWidth s -> Std.sexti targetTy xv
          (t, s) | Just (False, _) <- intWidth t, Just (False, _) <- intWidth s -> Std.zexti targetTy xv
          (t, s) | Just (False, _) <- intWidth t, Just (True , _) <- intWidth s -> Std.zexti targetTy xv
          -- Integral <-> Floating point casts
          (fp, i) | Just _ <- fpWidth fp, Just (True, _) <- intWidth i -> Std.sitofp targetTy xv
          (i, fp) | Just _ <- fpWidth fp, Just (True, _) <- intWidth i -> Std.fptosi targetTy xv
          -- Floating point <-> Floating point casts
          (fpt, fps) | Just tw <- fpWidth fpt, Just sw <- fpWidth fps -> case compare tw sw of
            GT -> Std.fpext   targetTy xv
            EQ -> return xv
            LT -> Std.fptrunc targetTy xv
          _ -> unsupported
        _ -> unsupported
      where
        targetTy = toMLIRType ty
        unsupported = error $ "Unsupported cast in MLIR lowering: " ++ pprint (getType x) ++ " to " ++ pprint ty
    _ -> error $ "Unsupported op in MLIR lowering: " ++ pprint op
    where
      getBinOpLowering bop = case bop of
        IAdd -> Std.addi
        ISub -> Std.subi
        IMul -> Std.muli
        IDiv -> Std.divi_signed
        FAdd -> Std.addf
        FSub -> Std.subf
        FMul -> Std.mulf
        FDiv -> Std.divf
        _    -> error $ "Unsupported binary operation in MLIR lowering: " ++ show bop
  
      getUnOpLowering uop = case uop of
        _    -> error $ "Unsupported unary operation in MLIR lowering: " ++ show uop
  
      fpWidth :: ScalarBaseType -> Maybe Int
      fpWidth ty = case ty of
        Float32Type  -> Just 32
        Float64Type  -> Just 64
        _            -> Nothing
  
      intWidth :: ScalarBaseType -> Maybe (Bool, Int)
      intWidth ty = case ty of
        Word8Type  -> Just (False, 8 )
        Word32Type -> Just (False, 32)
        Word64Type -> Just (False, 64)
        Int32Type  -> Just (True , 32)
        Int64Type  -> Just (True , 64)
        _          -> Nothing
  
  
  lowerAtom :: MonadLower m => Atom -> m MLIRAtom
  lowerAtom atom = case atom of
    Var n -> asks (!n)
    Con con -> case con of
      Lit (Float32Lit f) -> constant $ AST.FloatAttr   mlirType $ realToFrac f
      Lit (Float64Lit f) -> constant $ AST.FloatAttr   mlirType $ realToFrac f
      Lit (Int64Lit   i) -> constant $ AST.IntegerAttr mlirType $ fromIntegral i
      Lit (Int32Lit   i) -> constant $ AST.IntegerAttr mlirType $ fromIntegral i
      PairCon l r        -> Pair <$> lowerAtom l <*> lowerAtom r
      UnitCon            -> return Unit
      _ -> unsupported
    _ -> unsupported
    where
      constant = liftM Value . Std.constant mlirType
      unsupported = error $ "Unsupported atom in MLIR lowering: " ++ pprint atom
      mlirType = toMLIRType $ getType atom
  
  
  lowerValue :: (HasCallStack, MonadLower m)=> Atom -> m AST.Value
  lowerValue atom = lowerAtom atom <&> \case Value v -> v
                                             _ -> error "Expected a value"
  
  
  toMLIRType :: HasCallStack => Type -> AST.Type
  toMLIRType ty = case ty of
    TabTy  _ _ | Just (ns, elTy) <- rectScalarTable ty -> do
      let shape = reverse ns <&> \case
            FixedIntRange low high -> Just $ max 0 $ (fromIntegral high) - (fromIntegral low)
            _                      -> Nothing
      AST.RankedTensorType shape (toMLIRType elTy) Nothing
    BaseTy (Scalar Word8Type  )       -> AST.IntegerType AST.Signless 8
    BaseTy (Scalar Word32Type )       -> AST.IntegerType AST.Signless 32
    BaseTy (Scalar Word64Type )       -> AST.IntegerType AST.Signless 64
    BaseTy (Scalar Int32Type  )       -> AST.IntegerType AST.Signless 32
    BaseTy (Scalar Int64Type  )       -> AST.IntegerType AST.Signless 64
    BaseTy (Scalar Float32Type)       -> AST.Float32Type
    BaseTy (Scalar Float64Type)       -> AST.Float64Type
    BaseTy (PtrType (Heap CPU, elTy)) -> LLVM.Ptr $ toMLIRType $ BaseTy elTy
    _ -> unsupported
    where
      unsupported = error $ "Unsupported type in MLIR lowering: " ++ pprint ty
  
  
  rectScalarTable :: Type -> Maybe ([Type], Type)
  rectScalarTable ty = go [] ty
    where
      go ns elTy = case elTy of
        TabTy  (Ignore n) a -> go (n : ns) a
        BaseTy (Scalar _)   -> Just (reverse ns, elTy)
        _                   -> Nothing
  
  
  toCoreAtom :: AST.MonadBlockBuilder m => Type -> MLIRAtom -> m ([AST.Value], Abs (Nest Binder) Atom)
  toCoreAtom fullTy matom = do
    (neededVals, atom) <- go fullTy matom
    let valsWithTys = M.toList neededVals <&> \(name, (mty, ty)) -> (name AST.:> mty, ty)
        reconAtom = Abs (toNest [Bind (valueName v :> ty) | (v, ty) <- valsWithTys]) atom
    return (fst <$> valsWithTys, reconAtom)
    where
      --go :: MonadBlockBuilder m => Type -> MLIRAtom -> m
      go ty ma = case (ty, ma) of
        (BaseTy _    , Value v ) -> return (singleton v ty, Var $ valueName v :> ty)
        (TabTy  _ _  , Value v ) | Just (ns, BaseTy elTy) <- rectScalarTable ty -> do
          ~[vMemref] <- AST.emitOp $ AST.Operation
            { AST.opName = "memref.buffer_cast"
            , AST.opLocation = AST.UnknownLocation
            , AST.opResultTypes = AST.Explicit [bufferizeType $ AST.typeOf v]
            , AST.opOperands = [AST.operand v]
            , AST.opRegions = []
            , AST.opSuccessors = []
            , AST.opAttributes = AST.NoAttrs
            }
          vs <- LLVM.mlir_cast (llvmType $ AST.typeOf vMemref) vMemref
          -- TODO: Assert that allocPtr and basePtr are equal, or else we will be leaking data!
          ~[_allocPtr, basePtr, dataOff, _sizeArr, stridesArr] <- unpackStruct vs
          -- TODO: Substitute into ns
          unless (null $ freeVars ns) $ error "Unsupported dependent array result"
          unless (toMLIRType IdxRepTy == i32) $ error "MLIR lowering assumes that indices are i32"
          strides <- traverse (Std.trunci i32) =<< unpackArray stridesArr
          dataOffI32 <- Std.trunci i32 dataOff
          let arrayAtom = fst $ flip runBuilder mempty do
                buildNaryLam TabArrow (toNest $ Ignore <$> ns) \idxs -> do
                  let strideAtoms = strides <&> \s -> Var $ valueName s :> IdxRepTy
                  let dataOffAtom = Var $ valueName dataOffI32 :> IdxRepTy
                  let basePtrAtom = Var $ valueName basePtr :> PtrTy (Heap CPU, elTy)
                  ordinals <- traverse indexToIntE idxs
                  offset   <- foldM iadd dataOffAtom =<< zipWithM imul ordinals strideAtoms
                  unsafePtrLoad =<< ptrOffset basePtrAtom offset
          let mlirResults =
                [ (basePtr   , PtrTy (Heap CPU, elTy))
                , (dataOffI32, IdxRepTy              )
                ] <> (strides <&> (,IdxRepTy))
          return (results mlirResults, arrayAtom)
        (PairTy lt rt, Pair l r) -> do
          (lv, la) <- go lt l
          (rv, ra) <- go rt r
          return (lv <> rv, PairVal la ra)
        (UnitTy, Unit) -> return (mempty, UnitVal)
        _    -> error $ "Unsupported output MLIR atom for type: " ++ pprint fullTy
      singleton v ty = M.singleton (AST.operand v) ((AST.typeOf v), ty)
      results vtys = foldMap (uncurry singleton) vtys
      valueName = (\n -> Name GenName n 0) . T.decodeUtf8 . AST.operand
  
  
  bufferizeType :: AST.Type -> AST.Type
  bufferizeType ty = case ty of
    AST.RankedTensorType shape elTy Nothing ->
      AST.MemRefType shape elTy [] Nothing
    _ -> error "Don't know how to bufferize a type"
  
  llvmType :: AST.Type -> AST.Type
  llvmType ty = case ty of
    AST.MemRefType shape elTy [] Nothing ->
      LLVM.LiteralStruct
        [ LLVM.Ptr elTy                  -- allocated pointer
        , LLVM.Ptr elTy                  -- aligned pointer
        , i64                            -- offset between aligned pointer and data
        , LLVM.Array (length shape) i64  -- array of sizes
        , LLVM.Array (length shape) i64  -- array of strides
        ]
    _ -> error "Don't know how to get the llvm representation type"
  
  
  unpackElements :: AST.MonadBlockBuilder m => AST.Value -> [AST.Type] -> m [AST.Value]
  unpackElements val tys =
    forM (zip [0..] tys) \(i, ty) -> LLVM.extractvalue ty val [AST.IntegerAttr i32 i]
  
  unpackStruct :: AST.MonadBlockBuilder m => AST.Value -> m [AST.Value]
  unpackStruct val = case AST.typeOf val of
    LLVM.LiteralStruct tys -> unpackElements val tys
    _ -> error "Not a struct"
  
  unpackArray :: AST.MonadBlockBuilder m => AST.Value -> m [AST.Value]
  unpackArray val = case AST.typeOf val of
    LLVM.Array size ty -> unpackElements val $ replicate size ty
    _ -> error "Not an array"
  
  i32 :: AST.Type
  i32 = AST.IntegerType AST.Signless 32
  
  i64 :: AST.Type
  i64 = AST.IntegerType AST.Signless 64
  