    module Futhark.Optimise.ArrayLayout.Layout
      ( layoutTableFromIndexTable,
        Layout,
        Permutation,
        LayoutTable,

        -- * Exposed for testing
        commonPermutationEliminators,
      )
    where

    import Control.Monad (join)
    import Data.List qualified as L
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Futhark.Analysis.AccessPattern
    import Futhark.Analysis.PrimExp.Table (PrimExpTable)
    import Futhark.IR.Aliases
    import Futhark.IR.GPU
    import Futhark.IR.MC
    import Futhark.IR.MCMem
    import Futhark.Util (mininum)

    type Permutation = [Int]

    type LayoutTable =
      M.Map
        SegOpName
        ( M.Map
            ArrayName
            (M.Map IndexExprName Permutation)
        )

    class Layout rep where
      -- | Produce a coalescing permutation that will be used to create a
      -- manifest of the array. Returns Nothing if the array is already in
      -- the optimal layout or if the array access is too complex to
      -- confidently determine the optimal layout. Map each list of
      -- 'DimAccess' in the IndexTable to a permutation in a generic way
      -- that can be handled uniquely by each backend.
      permutationFromDimAccess ::
        PrimExpTable ->
        SegOpName ->
        ArrayName ->
        IndexExprName ->
        [DimAccess rep] ->
        Maybe Permutation

    isInscrutableExp :: PrimExp VName -> Bool
    isInscrutableExp :: PrimExp VName -> Bool
    isInscrutableExp (LeafExp VName
    _ PrimType
    _) = Bool
    False
    isInscrutableExp (ValueExp PrimValue
    _) = Bool
    False
    isInscrutableExp (BinOpExp BinOp
    _ PrimExp VName
    a PrimExp VName
    b) =
      PrimExp VName -> Bool
    isInscrutableExp PrimExp VName
    a Bool -> Bool -> Bool
    || PrimExp VName -> Bool
    isInscrutableExp PrimExp VName
    b
    isInscrutableExp (UnOpExp UnOp
    _ PrimExp VName
    a) =
      PrimExp VName -> Bool
    isInscrutableExp PrimExp VName
    a
    isInscrutableExp PrimExp VName
    _ = Bool
    True

    isInscrutable :: PrimExp VName -> Bool -> Bool
    isInscrutable :: PrimExp VName -> Bool -> Bool
    isInscrutable op :: PrimExp VName
    op@(BinOpExp {}) Bool
    counter =
      if Bool
    counter
        then -- Calculate stride and offset for loop-counters and thread-IDs
        case PrimExp VName -> Maybe (Int, Int)
    forall l. PrimExp l -> Maybe (Int, Int)
    reduceStrideAndOffset PrimExp VName
    op of
          -- Maximum allowable stride, might need tuning.
          Just (Int
    s, Int
    _) -> Int
    s Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    8
          Maybe (Int, Int)
    Nothing -> PrimExp VName -> Bool
    isInscrutableExp PrimExp VName
    op
        else PrimExp VName -> Bool
    isInscrutableExp PrimExp VName
    op
    isInscrutable PrimExp VName
    op Bool
    _ = PrimExp VName -> Bool
    isInscrutableExp PrimExp VName
    op

    reduceStrideAndOffset :: PrimExp l -> Maybe (Int, Int)
    reduceStrideAndOffset :: forall l. PrimExp l -> Maybe (Int, Int)
    reduceStrideAndOffset (LeafExp l
    _ PrimType
    _) = (Int, Int) -> Maybe (Int, Int)
    forall a. a -> Maybe a
    Just (Int
    1, Int
    0)
    reduceStrideAndOffset (BinOpExp BinOp
    oper PrimExp l
    a PrimExp l
    b) = case (PrimExp l
    a, PrimExp l
    b) of
      (ValueExp (IntValue IntValue
    v), PrimExp l
    _) -> IntValue -> PrimExp l -> Maybe (Int, Int)
    forall {l}. IntValue -> PrimExp l -> Maybe (Int, Int)
    reduce IntValue
    v PrimExp l
    b
      (PrimExp l
    _, ValueExp (IntValue IntValue
    v)) -> IntValue -> PrimExp l -> Maybe (Int, Int)
    forall {l}. IntValue -> PrimExp l -> Maybe (Int, Int)
    reduce IntValue
    v PrimExp l
    a
      (PrimExp l, PrimExp l)
    _ -> Maybe (Int, Int)
    forall a. Maybe a
    Nothing
      where
        reduce :: IntValue -> PrimExp l -> Maybe (Int, Int)
    reduce IntValue
    v (LeafExp l
    _ PrimType
    _) =
          case BinOp
    oper of
            Add IntType
    _ Overflow
    _ -> (Int, Int) -> Maybe (Int, Int)
    forall a. a -> Maybe a
    Just (Int
    1, IntValue -> Int
    forall int. Integral int => IntValue -> int
    valueIntegral IntValue
    v)
            Sub IntType
    _ Overflow
    _ -> (Int, Int) -> Maybe (Int, Int)
    forall a. a -> Maybe a
    Just (Int
    1, -IntValue -> Int
    forall int. Integral int => IntValue -> int
    valueIntegral IntValue
    v)
            Mul IntType
    _ Overflow
    _ -> (Int, Int) -> Maybe (Int, Int)
    forall a. a -> Maybe a
    Just (IntValue -> Int
    forall int. Integral int => IntValue -> int
    valueIntegral IntValue
    v, Int
    0)
            BinOp
    _ -> Maybe (Int, Int)
    forall a. Maybe a
    Nothing
        reduce IntValue
    v op :: PrimExp l
    op@(BinOpExp {}) =
          case PrimExp l -> Maybe (Int, Int)
    forall l. PrimExp l -> Maybe (Int, Int)
    reduceStrideAndOffset PrimExp l
    op of
            Maybe (Int, Int)
    Nothing -> Maybe (Int, Int)
    forall a. Maybe a
    Nothing
            Just (Int
    s, Int
    o) -> case BinOp
    oper of
              Add IntType
    _ Overflow
    _ -> (Int, Int) -> Maybe (Int, Int)
    forall a. a -> Maybe a
    Just (Int
    s, Int
    o Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + IntValue -> Int
    forall int. Integral int => IntValue -> int
    valueIntegral IntValue
    v)
              Sub IntType
    _ Overflow
    _ -> (Int, Int) -> Maybe (Int, Int)
    forall a. a -> Maybe a
    Just (Int
    s, Int
    o Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - IntValue -> Int
    forall int. Integral int => IntValue -> int
    valueIntegral IntValue
    v)
              Mul IntType
    _ Overflow
    _ -> (Int, Int) -> Maybe (Int, Int)
    forall a. a -> Maybe a
    Just (Int
    s Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * IntValue -> Int
    forall int. Integral int => IntValue -> int
    valueIntegral IntValue
    v, Int
    o Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * IntValue -> Int
    forall int. Integral int => IntValue -> int
    valueIntegral IntValue
    v)
              BinOp
    _ -> Maybe (Int, Int)
    forall a. Maybe a
    Nothing
        reduce IntValue
    _ (UnOpExp (Neg PrimType
    Bool) PrimExp l
    _) = Maybe (Int, Int)
    forall a. Maybe a
    Nothing
        reduce IntValue
    _ (UnOpExp (Complement IntType
    _) PrimExp l
    _) = Maybe (Int, Int)
    forall a. Maybe a
    Nothing
        reduce IntValue
    _ (UnOpExp (Abs IntType
    _) PrimExp l
    _) = Maybe (Int, Int)
    forall a. Maybe a
    Nothing
        reduce IntValue
    _ (UnOpExp UnOp
    _ PrimExp l
    sub_op) = PrimExp l -> Maybe (Int, Int)
    forall l. PrimExp l -> Maybe (Int, Int)
    reduceStrideAndOffset PrimExp l
    sub_op
        reduce IntValue
    _ (ConvOpExp ConvOp
    _ PrimExp l
    sub_op) = PrimExp l -> Maybe (Int, Int)
    forall l. PrimExp l -> Maybe (Int, Int)
    reduceStrideAndOffset PrimExp l
    sub_op
        reduce IntValue
    _ PrimExp l
    _ = Maybe (Int, Int)
    forall a. Maybe a
    Nothing
    reduceStrideAndOffset PrimExp l
    _ = Maybe (Int, Int)
    forall a. Maybe a
    Nothing

    -- | Reasons common to all backends to not manifest an array.
    commonPermutationEliminators :: [Int] -> [BodyType] -> Bool
    commonPermutationEliminators :: [Int] -> [BodyType] -> Bool
    commonPermutationEliminators [Int]
    perm [BodyType]
    nest = do
      -- Don't manifest if the permutation is the permutation is invalid
      let is_invalid_perm :: Bool
    is_invalid_perm = Bool -> Bool
    not ([Int] -> [Int]
    forall a. Ord a => [a] -> [a]
    L.sort [Int]
    perm [Int] -> [Int] -> Bool
    forall a. Eq a => [a] -> [a] -> Bool
    `L.isPrefixOf` [Int
    0 ..])
          -- Don't manifest if the permutation is the identity permutation
          is_identity :: Bool
    is_identity = [Int]
    perm [Int] -> [Int] -> Bool
    forall a. Eq a => [a] -> [a] -> Bool
    `L.isPrefixOf` [Int
    0 ..]
          -- or is not a transpose.
          inefficient_transpose :: Bool
    inefficient_transpose = Maybe (Int, Int, Int) -> Bool
    forall a. Maybe a -> Bool
    isNothing (Maybe (Int, Int, Int) -> Bool) -> Maybe (Int, Int, Int) -> Bool
    forall a b. (a -> b) -> a -> b
    $ [Int] -> Maybe (Int, Int, Int)
    isMapTranspose [Int]
    perm
          -- or if the last idx remains last
          static_last_idx :: Bool
    static_last_idx = [Int] -> Int
    forall a. HasCallStack => [a] -> a
    last [Int]
    perm Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Int] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Int]
    perm Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1
          -- Don't manifest if the array is defined inside a segOp
          inside_undesired :: Bool
    inside_undesired = (BodyType -> Bool) -> [BodyType] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any BodyType -> Bool
    undesired [BodyType]
    nest

      Bool
    is_invalid_perm
        Bool -> Bool -> Bool
    || Bool
    is_identity
        Bool -> Bool -> Bool
    || Bool
    inefficient_transpose
        Bool -> Bool -> Bool
    || Bool
    static_last_idx
        Bool -> Bool -> Bool
    || Bool
    inside_undesired
      where
        undesired :: BodyType -> Bool
        undesired :: BodyType -> Bool
    undesired BodyType
    bodyType = case BodyType
    bodyType of
          SegOpName SegOpName
    _ -> Bool
    True
          BodyType
    _ -> Bool
    False

    sortMC :: [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    sortMC :: forall {k} (rep :: k).
    [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    sortMC =
      ((Int, DimAccess rep) -> (Int, DimAccess rep) -> Ordering)
    -> [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    forall a. (a -> a -> Ordering) -> [a] -> [a]
    L.sortBy (Int, DimAccess rep) -> (Int, DimAccess rep) -> Ordering
    forall {k} {k} {b} {rep :: k} {rep :: k}.
    Ord b =>
    (b, DimAccess rep) -> (b, DimAccess rep) -> Ordering
    dimdexMCcmp
      where
        dimdexMCcmp :: (b, DimAccess rep) -> (b, DimAccess rep) -> Ordering
    dimdexMCcmp (b
    ia, DimAccess rep
    a) (b
    ib, DimAccess rep
    b) = do
          let aggr1 :: Maybe (VarType, Int, b)
    aggr1 =
                (Maybe (VarType, Int, b)
     -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b))
    -> Maybe (VarType, Int, b)
    -> [Maybe (VarType, Int, b)]
    -> Maybe (VarType, Int, b)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Maybe (VarType, Int, b)
    -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b)
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b)
    -> Maybe (VarType, a, b) -> Maybe (VarType, a, b)
    max' Maybe (VarType, Int, b)
    forall a. Maybe a
    Nothing ([Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b))
    -> [Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b)
    forall a b. (a -> b) -> a -> b
    $ ((VName, Dependency) -> Maybe (VarType, Int, b))
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> [a] -> [b]
    map (b -> Dependency -> Maybe (VarType, Int, b)
    forall {c}. c -> Dependency -> Maybe (VarType, Int, c)
    f b
    ia (Dependency -> Maybe (VarType, Int, b))
    -> ((VName, Dependency) -> Dependency)
    -> (VName, Dependency)
    -> Maybe (VarType, Int, b)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Dependency) -> Dependency
    forall a b. (a, b) -> b
    snd) ([(VName, Dependency)] -> [Maybe (VarType, Int, b)])
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> a -> b
    $ Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName Dependency -> [(VName, Dependency)])
    -> Map VName Dependency -> [(VName, Dependency)]
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    a
              aggr2 :: Maybe (VarType, Int, b)
    aggr2 =
                (Maybe (VarType, Int, b)
     -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b))
    -> Maybe (VarType, Int, b)
    -> [Maybe (VarType, Int, b)]
    -> Maybe (VarType, Int, b)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Maybe (VarType, Int, b)
    -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b)
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b)
    -> Maybe (VarType, a, b) -> Maybe (VarType, a, b)
    max' Maybe (VarType, Int, b)
    forall a. Maybe a
    Nothing ([Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b))
    -> [Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b)
    forall a b. (a -> b) -> a -> b
    $ ((VName, Dependency) -> Maybe (VarType, Int, b))
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> [a] -> [b]
    map (b -> Dependency -> Maybe (VarType, Int, b)
    forall {c}. c -> Dependency -> Maybe (VarType, Int, c)
    f b
    ib (Dependency -> Maybe (VarType, Int, b))
    -> ((VName, Dependency) -> Dependency)
    -> (VName, Dependency)
    -> Maybe (VarType, Int, b)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Dependency) -> Dependency
    forall a b. (a, b) -> b
    snd) ([(VName, Dependency)] -> [Maybe (VarType, Int, b)])
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> a -> b
    $ Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName Dependency -> [(VName, Dependency)])
    -> Map VName Dependency -> [(VName, Dependency)]
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    b
          Maybe (VarType, Int, b) -> Maybe (VarType, Int, b) -> Ordering
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    cmpIdxPat Maybe (VarType, Int, b)
    aggr1 Maybe (VarType, Int, b)
    aggr2
          where
            cmpIdxPat :: Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    cmpIdxPat Maybe (VarType, a, b)
    Nothing Maybe (VarType, a, b)
    Nothing = Ordering
    EQ
            cmpIdxPat (Just (VarType, a, b)
    _) Maybe (VarType, a, b)
    Nothing = Ordering
    GT
            cmpIdxPat Maybe (VarType, a, b)
    Nothing (Just (VarType, a, b)
    _) = Ordering
    LT
            cmpIdxPat
              (Just (VarType
    iterL, a
    lvlL, b
    original_lvl_L))
              (Just (VarType
    iterR, a
    lvlR, b
    original_lvl_R)) =
                case (VarType
    iterL, VarType
    iterR) of
                  (VarType
    ThreadID, VarType
    ThreadID) -> (a
    lvlL, b
    original_lvl_L) (a, b) -> (a, b) -> Ordering
    forall a. Ord a => a -> a -> Ordering
    `compare` (a
    lvlR, b
    original_lvl_R)
                  (VarType
    ThreadID, VarType
    _) -> Ordering
    LT
                  (VarType
    _, VarType
    ThreadID) -> Ordering
    GT
                  (VarType, VarType)
    _ -> (a
    lvlL, b
    original_lvl_L) (a, b) -> (a, b) -> Ordering
    forall a. Ord a => a -> a -> Ordering
    `compare` (a
    lvlR, b
    original_lvl_R)

            max' :: Maybe (VarType, a, b)
    -> Maybe (VarType, a, b) -> Maybe (VarType, a, b)
    max' Maybe (VarType, a, b)
    lhs Maybe (VarType, a, b)
    rhs =
              case Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    cmpIdxPat Maybe (VarType, a, b)
    lhs Maybe (VarType, a, b)
    rhs of
                Ordering
    LT -> Maybe (VarType, a, b)
    rhs
                Ordering
    _ -> Maybe (VarType, a, b)
    lhs

            f :: c -> Dependency -> Maybe (VarType, Int, c)
    f c
    og (Dependency Int
    lvl VarType
    varType) = (VarType, Int, c) -> Maybe (VarType, Int, c)
    forall a. a -> Maybe a
    Just (VarType
    varType, Int
    lvl, c
    og)

    multicorePermutation :: PrimExpTable -> SegOpName -> ArrayName -> IndexExprName -> [DimAccess rep] -> Maybe Permutation
    multicorePermutation :: forall {k} (rep :: k).
    PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess rep]
    -> Maybe [Int]
    multicorePermutation PrimExpTable
    primExpTable SegOpName
    _segOpName (VName
    _arr_name, [BodyType]
    nest, [Int]
    arr_layout) VName
    _idx_name [DimAccess rep]
    dimAccesses = do
      -- Dont accept indices where the last index is invariant
      let lastIdxIsInvariant :: Bool
    lastIdxIsInvariant = DimAccess rep -> Bool
    forall {k} (rep :: k). DimAccess rep -> Bool
    isInvariant (DimAccess rep -> Bool) -> DimAccess rep -> Bool
    forall a b. (a -> b) -> a -> b
    $ [DimAccess rep] -> DimAccess rep
    forall a. HasCallStack => [a] -> a
    last [DimAccess rep]
    dimAccesses

      -- Check if any of the dependencies are too complex to reason about
      let dimAccesses' :: [DimAccess rep]
    dimAccesses' = (DimAccess rep -> Bool) -> [DimAccess rep] -> [DimAccess rep]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Maybe VName -> Bool
    forall a. Maybe a -> Bool
    isJust (Maybe VName -> Bool)
    -> (DimAccess rep -> Maybe VName) -> DimAccess rep -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar) [DimAccess rep]
    dimAccesses
          deps :: [VName]
    deps = (DimAccess rep -> Maybe VName) -> [DimAccess rep] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar [DimAccess rep]
    dimAccesses'
          counters :: [Bool]
    counters = (DimAccess rep -> [Bool]) -> [DimAccess rep] -> [Bool]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (((VName, Dependency) -> Bool) -> [(VName, Dependency)] -> [Bool]
    forall a b. (a -> b) -> [a] -> [b]
    map (VarType -> Bool
    isCounter (VarType -> Bool)
    -> ((VName, Dependency) -> VarType) -> (VName, Dependency) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Dependency -> VarType
    varType (Dependency -> VarType)
    -> ((VName, Dependency) -> Dependency)
    -> (VName, Dependency)
    -> VarType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Dependency) -> Dependency
    forall a b. (a, b) -> b
    snd) ([(VName, Dependency)] -> [Bool])
    -> (DimAccess rep -> [(VName, Dependency)])
    -> DimAccess rep
    -> [Bool]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName Dependency -> [(VName, Dependency)])
    -> (DimAccess rep -> Map VName Dependency)
    -> DimAccess rep
    -> [(VName, Dependency)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies) [DimAccess rep]
    dimAccesses'
          primExps :: Maybe [PrimExp VName]
    primExps = (VName -> Maybe (PrimExp VName))
    -> [VName] -> Maybe [PrimExp VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (Maybe (Maybe (PrimExp VName)) -> Maybe (PrimExp VName)
    forall (m :: * -> *) a. Monad m => m (m a) -> m a
    join (Maybe (Maybe (PrimExp VName)) -> Maybe (PrimExp VName))
    -> (VName -> Maybe (Maybe (PrimExp VName)))
    -> VName
    -> Maybe (PrimExp VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName -> PrimExpTable -> Maybe (Maybe (PrimExp VName))
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` PrimExpTable
    primExpTable)) [VName]
    deps
          inscrutable :: Bool
    inscrutable = Bool -> ([PrimExp VName] -> Bool) -> Maybe [PrimExp VName] -> Bool
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Bool
    True (((PrimExp VName, Bool) -> Bool) -> [(PrimExp VName, Bool)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any ((PrimExp VName -> Bool -> Bool) -> (PrimExp VName, Bool) -> Bool
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry PrimExp VName -> Bool -> Bool
    isInscrutable) ([(PrimExp VName, Bool)] -> Bool)
    -> ([PrimExp VName] -> [(PrimExp VName, Bool)])
    -> [PrimExp VName]
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([PrimExp VName] -> [Bool] -> [(PrimExp VName, Bool)])
    -> [Bool] -> [PrimExp VName] -> [(PrimExp VName, Bool)]
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip [PrimExp VName] -> [Bool] -> [(PrimExp VName, Bool)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Bool]
    counters) Maybe [PrimExp VName]
    primExps

      -- Create a candidate permutation
      let perm :: [Int]
    perm = ((Int, DimAccess rep) -> Int) -> [(Int, DimAccess rep)] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int, DimAccess rep) -> Int
    forall a b. (a, b) -> a
    fst ([(Int, DimAccess rep)] -> [Int])
    -> [(Int, DimAccess rep)] -> [Int]
    forall a b. (a -> b) -> a -> b
    $ [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    forall {k} (rep :: k).
    [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    sortMC ([Int] -> [DimAccess rep] -> [(Int, DimAccess rep)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Int]
    arr_layout [DimAccess rep]
    dimAccesses)

      -- Check if we want to manifest this array with the permutation
      if Bool
    lastIdxIsInvariant Bool -> Bool -> Bool
    || Bool
    inscrutable Bool -> Bool -> Bool
    || [Int] -> [BodyType] -> Bool
    commonPermutationEliminators [Int]
    perm [BodyType]
    nest
        then Maybe [Int]
    forall a. Maybe a
    Nothing
        else [Int] -> Maybe [Int]
    forall a. a -> Maybe a
    Just [Int]
    perm

    instance Layout MC where
      permutationFromDimAccess :: PrimExpTable
    -> SegOpName -> ArrayName -> VName -> [DimAccess MC] -> Maybe [Int]
    permutationFromDimAccess = PrimExpTable
    -> SegOpName -> ArrayName -> VName -> [DimAccess MC] -> Maybe [Int]
    forall {k} (rep :: k).
    PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess rep]
    -> Maybe [Int]
    multicorePermutation

    sortGPU :: [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    sortGPU :: forall {k} (rep :: k).
    [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    sortGPU =
      ((Int, DimAccess rep) -> (Int, DimAccess rep) -> Ordering)
    -> [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    forall a. (a -> a -> Ordering) -> [a] -> [a]
    L.sortBy (Int, DimAccess rep) -> (Int, DimAccess rep) -> Ordering
    forall {k} {k} {b} {rep :: k} {rep :: k}.
    Ord b =>
    (b, DimAccess rep) -> (b, DimAccess rep) -> Ordering
    dimdexGPUcmp
      where
        dimdexGPUcmp :: (b, DimAccess rep) -> (b, DimAccess rep) -> Ordering
    dimdexGPUcmp (b
    ia, DimAccess rep
    a) (b
    ib, DimAccess rep
    b) = do
          let aggr1 :: Maybe (VarType, Int, b)
    aggr1 =
                (Maybe (VarType, Int, b)
     -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b))
    -> Maybe (VarType, Int, b)
    -> [Maybe (VarType, Int, b)]
    -> Maybe (VarType, Int, b)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Maybe (VarType, Int, b)
    -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b)
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b)
    -> Maybe (VarType, a, b) -> Maybe (VarType, a, b)
    max' Maybe (VarType, Int, b)
    forall a. Maybe a
    Nothing ([Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b))
    -> [Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b)
    forall a b. (a -> b) -> a -> b
    $ ((VName, Dependency) -> Maybe (VarType, Int, b))
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> [a] -> [b]
    map (b -> Dependency -> Maybe (VarType, Int, b)
    forall {c}. c -> Dependency -> Maybe (VarType, Int, c)
    f b
    ia (Dependency -> Maybe (VarType, Int, b))
    -> ((VName, Dependency) -> Dependency)
    -> (VName, Dependency)
    -> Maybe (VarType, Int, b)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Dependency) -> Dependency
    forall a b. (a, b) -> b
    snd) ([(VName, Dependency)] -> [Maybe (VarType, Int, b)])
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> a -> b
    $ Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName Dependency -> [(VName, Dependency)])
    -> Map VName Dependency -> [(VName, Dependency)]
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    a
              aggr2 :: Maybe (VarType, Int, b)
    aggr2 =
                (Maybe (VarType, Int, b)
     -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b))
    -> Maybe (VarType, Int, b)
    -> [Maybe (VarType, Int, b)]
    -> Maybe (VarType, Int, b)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Maybe (VarType, Int, b)
    -> Maybe (VarType, Int, b) -> Maybe (VarType, Int, b)
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b)
    -> Maybe (VarType, a, b) -> Maybe (VarType, a, b)
    max' Maybe (VarType, Int, b)
    forall a. Maybe a
    Nothing ([Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b))
    -> [Maybe (VarType, Int, b)] -> Maybe (VarType, Int, b)
    forall a b. (a -> b) -> a -> b
    $ ((VName, Dependency) -> Maybe (VarType, Int, b))
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> [a] -> [b]
    map (b -> Dependency -> Maybe (VarType, Int, b)
    forall {c}. c -> Dependency -> Maybe (VarType, Int, c)
    f b
    ib (Dependency -> Maybe (VarType, Int, b))
    -> ((VName, Dependency) -> Dependency)
    -> (VName, Dependency)
    -> Maybe (VarType, Int, b)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Dependency) -> Dependency
    forall a b. (a, b) -> b
    snd) ([(VName, Dependency)] -> [Maybe (VarType, Int, b)])
    -> [(VName, Dependency)] -> [Maybe (VarType, Int, b)]
    forall a b. (a -> b) -> a -> b
    $ Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName Dependency -> [(VName, Dependency)])
    -> Map VName Dependency -> [(VName, Dependency)]
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies DimAccess rep
    b
          Maybe (VarType, Int, b) -> Maybe (VarType, Int, b) -> Ordering
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    cmpIdxPat Maybe (VarType, Int, b)
    aggr1 Maybe (VarType, Int, b)
    aggr2
          where
            cmpIdxPat :: Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    cmpIdxPat Maybe (VarType, a, b)
    Nothing Maybe (VarType, a, b)
    Nothing = Ordering
    EQ
            cmpIdxPat (Just (VarType, a, b)
    _) Maybe (VarType, a, b)
    Nothing = Ordering
    GT
            cmpIdxPat Maybe (VarType, a, b)
    Nothing (Just (VarType, a, b)
    _) = Ordering
    LT
            cmpIdxPat
              (Just (VarType
    iterL, a
    lvlL, b
    original_lvl_L))
              (Just (VarType
    iterR, a
    lvlR, b
    original_lvl_R)) = case (VarType
    iterL, VarType
    iterR) of
                (VarType
    ThreadID, VarType
    ThreadID) -> (a
    lvlL, b
    original_lvl_L) (a, b) -> (a, b) -> Ordering
    forall a. Ord a => a -> a -> Ordering
    `compare` (a
    lvlR, b
    original_lvl_R)
                (VarType
    ThreadID, VarType
    _) -> Ordering
    GT
                (VarType
    _, VarType
    ThreadID) -> Ordering
    LT
                (VarType, VarType)
    _ -> (a
    lvlL, b
    original_lvl_L) (a, b) -> (a, b) -> Ordering
    forall a. Ord a => a -> a -> Ordering
    `compare` (a
    lvlR, b
    original_lvl_R)

            max' :: Maybe (VarType, a, b)
    -> Maybe (VarType, a, b) -> Maybe (VarType, a, b)
    max' Maybe (VarType, a, b)
    lhs Maybe (VarType, a, b)
    rhs =
              case Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    forall {a} {b}.
    (Ord a, Ord b) =>
    Maybe (VarType, a, b) -> Maybe (VarType, a, b) -> Ordering
    cmpIdxPat Maybe (VarType, a, b)
    lhs Maybe (VarType, a, b)
    rhs of
                Ordering
    LT -> Maybe (VarType, a, b)
    rhs
                Ordering
    _ -> Maybe (VarType, a, b)
    lhs

            f :: c -> Dependency -> Maybe (VarType, Int, c)
    f c
    og (Dependency Int
    lvl VarType
    varType) = (VarType, Int, c) -> Maybe (VarType, Int, c)
    forall a. a -> Maybe a
    Just (VarType
    varType, Int
    lvl, c
    og)

    gpuPermutation :: PrimExpTable -> SegOpName -> ArrayName -> IndexExprName -> [DimAccess rep] -> Maybe Permutation
    gpuPermutation :: forall {k} (rep :: k).
    PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess rep]
    -> Maybe [Int]
    gpuPermutation PrimExpTable
    primExpTable SegOpName
    _segOpName (VName
    _arr_name, [BodyType]
    nest, [Int]
    arr_layout) VName
    _idx_name [DimAccess rep]
    dimAccesses = do
      -- Find the outermost parallel level. XXX: this is a bit hacky. Why
      -- don't we simply know at this point the nest in which this index
      -- occurs?
      let outermost_par :: Int
    outermost_par = [Int] -> Int
    forall a (f :: * -> *). (Num a, Ord a, Foldable f) => f a -> a
    mininum ([Int] -> Int) -> [Int] -> Int
    forall a b. (a -> b) -> a -> b
    $ (DimAccess rep -> [Int]) -> [DimAccess rep] -> [Int]
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap ((Dependency -> Int) -> [Dependency] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map Dependency -> Int
    lvl ([Dependency] -> [Int])
    -> (DimAccess rep -> [Dependency]) -> DimAccess rep -> [Int]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> [Dependency]
    forall {k} {rep :: k}. DimAccess rep -> [Dependency]
    parDeps) [DimAccess rep]
    dimAccesses
          invariantToPar :: Dependency -> Bool
    invariantToPar = (Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    outermost_par) (Int -> Bool) -> (Dependency -> Int) -> Dependency -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Dependency -> Int
    lvl

      -- Do nothing if last index is invariant to segop.
      let lastIdxIsInvariant :: Bool
    lastIdxIsInvariant = (Dependency -> Bool) -> Map VName Dependency -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all Dependency -> Bool
    invariantToPar (Map VName Dependency -> Bool) -> Map VName Dependency -> Bool
    forall a b. (a -> b) -> a -> b
    $ DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies (DimAccess rep -> Map VName Dependency)
    -> DimAccess rep -> Map VName Dependency
    forall a b. (a -> b) -> a -> b
    $ [DimAccess rep] -> DimAccess rep
    forall a. HasCallStack => [a] -> a
    last [DimAccess rep]
    dimAccesses

      -- Do nothing if any index is constant, because otherwise we can end
      -- up transposing a too-large array.
      let anyIsConstant :: Bool
    anyIsConstant = (DimAccess rep -> Bool) -> [DimAccess rep] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (Map VName Dependency -> Bool
    forall a. Map VName a -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null (Map VName Dependency -> Bool)
    -> (DimAccess rep -> Map VName Dependency) -> DimAccess rep -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies) [DimAccess rep]
    dimAccesses

      -- Check if any of the dependencies are too complex to reason about
      let dimAccesses' :: [DimAccess rep]
    dimAccesses' = (DimAccess rep -> Bool) -> [DimAccess rep] -> [DimAccess rep]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Maybe VName -> Bool
    forall a. Maybe a -> Bool
    isJust (Maybe VName -> Bool)
    -> (DimAccess rep -> Maybe VName) -> DimAccess rep -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar) [DimAccess rep]
    dimAccesses
          deps :: [VName]
    deps = (DimAccess rep -> Maybe VName) -> [DimAccess rep] -> [VName]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe DimAccess rep -> Maybe VName
    forall {k} (rep :: k). DimAccess rep -> Maybe VName
    originalVar [DimAccess rep]
    dimAccesses'
          counters :: [Bool]
    counters = (DimAccess rep -> [Bool]) -> [DimAccess rep] -> [Bool]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (((VName, Dependency) -> Bool) -> [(VName, Dependency)] -> [Bool]
    forall a b. (a -> b) -> [a] -> [b]
    map (VarType -> Bool
    isCounter (VarType -> Bool)
    -> ((VName, Dependency) -> VarType) -> (VName, Dependency) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Dependency -> VarType
    varType (Dependency -> VarType)
    -> ((VName, Dependency) -> Dependency)
    -> (VName, Dependency)
    -> VarType
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Dependency) -> Dependency
    forall a b. (a, b) -> b
    snd) ([(VName, Dependency)] -> [Bool])
    -> (DimAccess rep -> [(VName, Dependency)])
    -> DimAccess rep
    -> [Bool]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map VName Dependency -> [(VName, Dependency)]
    forall k a. Map k a -> [(k, a)]
    M.toList (Map VName Dependency -> [(VName, Dependency)])
    -> (DimAccess rep -> Map VName Dependency)
    -> DimAccess rep
    -> [(VName, Dependency)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies) [DimAccess rep]
    dimAccesses'
          primExps :: Maybe [PrimExp VName]
    primExps = (VName -> Maybe (PrimExp VName))
    -> [VName] -> Maybe [PrimExp VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (Maybe (Maybe (PrimExp VName)) -> Maybe (PrimExp VName)
    forall (m :: * -> *) a. Monad m => m (m a) -> m a
    join (Maybe (Maybe (PrimExp VName)) -> Maybe (PrimExp VName))
    -> (VName -> Maybe (Maybe (PrimExp VName)))
    -> VName
    -> Maybe (PrimExp VName)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName -> PrimExpTable -> Maybe (Maybe (PrimExp VName))
    forall k a. Ord k => k -> Map k a -> Maybe a
    `M.lookup` PrimExpTable
    primExpTable)) [VName]
    deps
          inscrutable :: Bool
    inscrutable = Bool -> ([PrimExp VName] -> Bool) -> Maybe [PrimExp VName] -> Bool
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Bool
    True (((PrimExp VName, Bool) -> Bool) -> [(PrimExp VName, Bool)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any ((PrimExp VName -> Bool -> Bool) -> (PrimExp VName, Bool) -> Bool
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry PrimExp VName -> Bool -> Bool
    isInscrutable) ([(PrimExp VName, Bool)] -> Bool)
    -> ([PrimExp VName] -> [(PrimExp VName, Bool)])
    -> [PrimExp VName]
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([PrimExp VName] -> [Bool] -> [(PrimExp VName, Bool)])
    -> [Bool] -> [PrimExp VName] -> [(PrimExp VName, Bool)]
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip [PrimExp VName] -> [Bool] -> [(PrimExp VName, Bool)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Bool]
    counters) Maybe [PrimExp VName]
    primExps

      -- Create a candidate permutation
      let perm :: [Int]
    perm = ((Int, DimAccess rep) -> Int) -> [(Int, DimAccess rep)] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int, DimAccess rep) -> Int
    forall a b. (a, b) -> a
    fst ([(Int, DimAccess rep)] -> [Int])
    -> [(Int, DimAccess rep)] -> [Int]
    forall a b. (a -> b) -> a -> b
    $ [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    forall {k} (rep :: k).
    [(Int, DimAccess rep)] -> [(Int, DimAccess rep)]
    sortGPU ([Int] -> [DimAccess rep] -> [(Int, DimAccess rep)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Int]
    arr_layout [DimAccess rep]
    dimAccesses)

      -- Check if we want to manifest this array with the permutation
      if Bool
    lastIdxIsInvariant
        Bool -> Bool -> Bool
    || Bool
    anyIsConstant
        Bool -> Bool -> Bool
    || Bool
    inscrutable
        Bool -> Bool -> Bool
    || [Int] -> [BodyType] -> Bool
    commonPermutationEliminators [Int]
    perm [BodyType]
    nest
        then Maybe [Int]
    forall a. Maybe a
    Nothing
        else [Int] -> Maybe [Int]
    forall a. a -> Maybe a
    Just [Int]
    perm
      where
        parDeps :: DimAccess rep -> [Dependency]
    parDeps = (Dependency -> Bool) -> [Dependency] -> [Dependency]
    forall a. (a -> Bool) -> [a] -> [a]
    filter ((VarType -> VarType -> Bool
    forall a. Eq a => a -> a -> Bool
    == VarType
    ThreadID) (VarType -> Bool) -> (Dependency -> VarType) -> Dependency -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Dependency -> VarType
    varType) ([Dependency] -> [Dependency])
    -> (DimAccess rep -> [Dependency]) -> DimAccess rep -> [Dependency]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Map VName Dependency -> [Dependency]
    forall k a. Map k a -> [a]
    M.elems (Map VName Dependency -> [Dependency])
    -> (DimAccess rep -> Map VName Dependency)
    -> DimAccess rep
    -> [Dependency]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . DimAccess rep -> Map VName Dependency
    forall {k} (rep :: k). DimAccess rep -> Map VName Dependency
    dependencies

    instance Layout GPU where
      permutationFromDimAccess :: PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess GPU]
    -> Maybe [Int]
    permutationFromDimAccess = PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess GPU]
    -> Maybe [Int]
    forall {k} (rep :: k).
    PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess rep]
    -> Maybe [Int]
    gpuPermutation

    -- | like mapMaybe, but works on nested maps. Eliminates "dangling"
    -- maps / rows with missing (Nothing) values.
    tableMapMaybe ::
      (k0 -> k1 -> k2 -> a -> Maybe b) ->
      M.Map k0 (M.Map k1 (M.Map k2 a)) ->
      M.Map k0 (M.Map k1 (M.Map k2 b))
    tableMapMaybe :: forall k0 k1 k2 a b.
    (k0 -> k1 -> k2 -> a -> Maybe b)
    -> Map k0 (Map k1 (Map k2 a)) -> Map k0 (Map k1 (Map k2 b))
    tableMapMaybe k0 -> k1 -> k2 -> a -> Maybe b
    f =
      (k0 -> Map k1 (Map k2 a) -> Maybe (Map k1 (Map k2 b)))
    -> Map k0 (Map k1 (Map k2 a)) -> Map k0 (Map k1 (Map k2 b))
    forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
    M.mapMaybeWithKey ((k0 -> Map k1 (Map k2 a) -> Maybe (Map k1 (Map k2 b)))
     -> Map k0 (Map k1 (Map k2 a)) -> Map k0 (Map k1 (Map k2 b)))
    -> (k0 -> Map k1 (Map k2 a) -> Maybe (Map k1 (Map k2 b)))
    -> Map k0 (Map k1 (Map k2 a))
    -> Map k0 (Map k1 (Map k2 b))
    forall a b. (a -> b) -> a -> b
    $ \k0
    key0 -> (k1 -> Map k2 a -> Maybe (Map k2 b))
    -> Map k1 (Map k2 a) -> Maybe (Map k1 (Map k2 b))
    forall {k} {a} {a}.
    (k -> a -> Maybe a) -> Map k a -> Maybe (Map k a)
    mapToMaybe ((k1 -> Map k2 a -> Maybe (Map k2 b))
     -> Map k1 (Map k2 a) -> Maybe (Map k1 (Map k2 b)))
    -> (k1 -> Map k2 a -> Maybe (Map k2 b))
    -> Map k1 (Map k2 a)
    -> Maybe (Map k1 (Map k2 b))
    forall a b. (a -> b) -> a -> b
    $ (k2 -> a -> Maybe b) -> Map k2 a -> Maybe (Map k2 b)
    forall {k} {a} {a}.
    (k -> a -> Maybe a) -> Map k a -> Maybe (Map k a)
    mapToMaybe ((k2 -> a -> Maybe b) -> Map k2 a -> Maybe (Map k2 b))
    -> (k1 -> k2 -> a -> Maybe b) -> k1 -> Map k2 a -> Maybe (Map k2 b)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . k0 -> k1 -> k2 -> a -> Maybe b
    f k0
    key0
      where
        maybeMap :: M.Map k a -> Maybe (M.Map k a)
        maybeMap :: forall k a. Map k a -> Maybe (Map k a)
    maybeMap Map k a
    val = if Map k a -> Bool
    forall a. Map k a -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null Map k a
    val then Maybe (Map k a)
    forall a. Maybe a
    Nothing else Map k a -> Maybe (Map k a)
    forall a. a -> Maybe a
    Just Map k a
    val

        mapToMaybe :: (k -> a -> Maybe a) -> Map k a -> Maybe (Map k a)
    mapToMaybe k -> a -> Maybe a
    g = Map k a -> Maybe (Map k a)
    forall k a. Map k a -> Maybe (Map k a)
    maybeMap (Map k a -> Maybe (Map k a))
    -> (Map k a -> Map k a) -> Map k a -> Maybe (Map k a)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (k -> a -> Maybe a) -> Map k a -> Map k a
    forall k a b. (k -> a -> Maybe b) -> Map k a -> Map k b
    M.mapMaybeWithKey k -> a -> Maybe a
    g

    -- | Given an ordering function for `DimAccess`, and an IndexTable,
    -- return a LayoutTable. We remove entries with no results after
    -- `permutationFromDimAccess`
    layoutTableFromIndexTable ::
      (Layout rep) =>
      PrimExpTable ->
      IndexTable rep ->
      LayoutTable
    layoutTableFromIndexTable :: forall {k} (rep :: k).
    Layout rep =>
    PrimExpTable -> IndexTable rep -> LayoutTable
    layoutTableFromIndexTable = (SegOpName -> ArrayName -> VName -> [DimAccess rep] -> Maybe [Int])
    -> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
    -> LayoutTable
    forall k0 k1 k2 a b.
    (k0 -> k1 -> k2 -> a -> Maybe b)
    -> Map k0 (Map k1 (Map k2 a)) -> Map k0 (Map k1 (Map k2 b))
    tableMapMaybe ((SegOpName
      -> ArrayName -> VName -> [DimAccess rep] -> Maybe [Int])
     -> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
     -> LayoutTable)
    -> (PrimExpTable
        -> SegOpName
        -> ArrayName
        -> VName
        -> [DimAccess rep]
        -> Maybe [Int])
    -> PrimExpTable
    -> Map SegOpName (Map ArrayName (Map VName [DimAccess rep]))
    -> LayoutTable
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess rep]
    -> Maybe [Int]
    forall {k} (rep :: k).
    Layout rep =>
    PrimExpTable
    -> SegOpName
    -> ArrayName
    -> VName
    -> [DimAccess rep]
    -> Maybe [Int]
    permutationFromDimAccess