content: module Futhark.Optimise.ArrayLayout.AnalyseTests (tests) where
  
  import Data.Map.Strict qualified as M
  import Futhark.Analysis.AccessPattern
  import Futhark.IR.GPU
  import Futhark.IR.GPUTests ()
  import Futhark.IR.SyntaxTests ()
  import Test.Tasty
  import Test.Tasty.HUnit
  
  tests :: TestTree
  tests = testGroup "Analyse" [analyseStmTests]
  
  analyseStmTests :: TestTree
  analyseStmTests =
    testGroup
      "analyseStm"
      [analyseIndexTests, analyseDimAccessesTests]
  
  analyseIndexTests :: TestTree
  analyseIndexTests =
    testGroup
      "analyseIndex"
      $ do
        let arr_name = "xss_5144"
        -- ============================= TestCase0 =============================
        -- Most simple case where we want to manifest an array, hence, we record
        -- the Index in the IndexTable.
        let testCase0 = testCase "2D manifest" $ do
              let ctx =
                    mempty
                      { parents =
                          [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),
                            LoopBodyName "defunc_0_f_res_5208"
                          ],
                        assignments =
                          M.fromList
                            [ ("gtid_5205", VariableInfo mempty 0 mempty ThreadID),
                              ("i_5209", VariableInfo mempty 1 mempty LoopVar)
                            ]
                      }
              let patternNames = ["b_5211"]
              let dimFixes =
                    [ DimFix (Var "gtid_5205"),
                      DimFix (Var "i_5209")
                    ]
              let indexTable =
                    M.fromList
                      [ ( SegmentedMap "defunc_0_map_res_5204",
                          M.fromList
                            [ ( (arr_name, [], [0 .. 1]),
                                M.fromList
                                  [ ( "b_5211",
                                      [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "gtid_5205"),
                                        DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "i_5209")
                                      ]
                                    )
                                  ]
                              )
                            ]
                        )
                      ]
              let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes
              indexTable' @?= indexTable
  
        -- ============================= TestCase2 =============================
        -- We don't want to manifest an array with only one dimension, so we don't
        -- record anything in the IndexTable.
        let testCase1 = testCase "1D manifest" $ do
              let ctx =
                    mempty
                      { parents =
                          [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),
                            LoopBodyName "defunc_0_f_res_5208"
                          ]
                      }
              let patternNames = ["b_5211"]
              let dimFixes = [DimFix "i_5209"]
  
              let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes
              indexTable' @?= mempty
  
        -- ============================= TestCase1 =============================
        -- We don't want to record anything to the IndexTable when the array is
        -- not accessed inside a SegMap
        -- TODO: Create a similar one for MC with loops
        let testCase2 = testCase "Not inside SegMap" $ do
              let ctx = mempty
              let patternNames = ["b_5211"]
              let dimFixes =
                    [ DimFix "gtid_5205",
                      DimFix "i_5209"
                    ]
              let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes
              indexTable' @?= mempty
  
        -- ============================= TestCase3 =============================
        -- If an array is allocated inside a loop or SegMap, we want to record that
        -- information in the ArrayName of the IndexTable.
        let testCase3 = testCase "Allocated inside SegMap" $ do
              let parents' =
                    [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),
                      LoopBodyName "defunc_0_f_res_5208"
                    ]
              let ctx =
                    mempty
                      { parents = parents',
                        assignments =
                          M.fromList
                            [ ("gtid_5205", VariableInfo mempty 0 mempty ThreadID),
                              ("i_5209", VariableInfo mempty 1 mempty LoopVar),
                              (arr_name, VariableInfo mempty 0 parents' Variable)
                            ]
                      }
              let patternNames = ["b_5211"]
              let dimFixes =
                    [ DimFix "gtid_5205",
                      DimFix "i_5209"
                    ]
              let indexTable =
                    M.fromList
                      [ ( SegmentedMap "defunc_0_map_res_5204",
                          M.fromList
                            [ ( (arr_name, parents', [0 .. 1]),
                                M.fromList
                                  [ ( "b_5211",
                                      [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "gtid_5205"),
                                        DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "i_5209")
                                      ]
                                    )
                                  ]
                              )
                            ]
                        )
                      ]
              let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes
              indexTable' @?= indexTable
  
        -- ============================= TestCase4 =============================
        -- If the vars in the index are temporaries, we want to reduce them to
        -- to the thread IDs and or loop counters they are functions of.
        let testCase4 = testCase "Reduce dependencies" $ do
              let ctx =
                    mempty
                      { parents =
                          [ SegOpName (SegmentedMap "defunc_0_map_res_5204"),
                            LoopBodyName "defunc_0_f_res_5208"
                          ],
                        assignments =
                          M.fromList
                            [ ("gtid_5205", VariableInfo mempty 0 mempty ThreadID),
                              ("i_5209", VariableInfo mempty 1 mempty LoopVar),
                              ("tmp0_5210", VariableInfo (namesFromList ["gtid_5205"]) 2 mempty Variable),
                              ("tmp1_5211", VariableInfo (namesFromList ["i_5209"]) 3 mempty Variable),
                              ("k_5212", VariableInfo mempty 1 mempty ConstType)
                            ]
                      }
              let patternNames = ["b_5211"]
              let dimFixes =
                    [ DimFix "tmp0_5210",
                      DimFix "tmp1_5211",
                      DimFix "k_5212"
                    ]
              let indexTable =
                    M.fromList
                      [ ( SegmentedMap "defunc_0_map_res_5204",
                          M.fromList
                            [ ( (arr_name, [], [0 .. 2]),
                                M.fromList
                                  [ ( "b_5211",
                                      [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "tmp0_5210"),
                                        DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "tmp1_5211"),
                                        DimAccess mempty (Just "k_5212")
                                      ]
                                    )
                                  ]
                              )
                            ]
                        )
                      ]
              let (_, indexTable') = analyseIndex ctx patternNames arr_name dimFixes
              indexTable' @?= indexTable
  
        [testCase0, testCase1, testCase2, testCase3, testCase4]
  
  analyseDimAccessesTests :: TestTree
  analyseDimAccessesTests = testGroup
    "analyseDimAccesses"
    $ do
      let testCase0 = testCase "Fold" $ do
            let indexTable =
                  M.fromList
                    [ ( SegmentedMap "defunc_0_map_res_5204",
                        M.fromList
                          [ ( ("xss_5144", [], [0, 1]),
                              M.fromList
                                [ ( "b_5211",
                                    [ DimAccess (M.fromList [("gtid_5205", Dependency 0 ThreadID)]) (Just "gtid_5205"),
                                      DimAccess (M.fromList [("i_5209", Dependency 1 LoopVar)]) (Just "i_5209")
                                    ]
                                  )
                                ]
                            )
                          ]
                      )
                    ]
            let indexTable' = (analyseDimAccesses @GPU) prog0
            indexTable' @?= indexTable
  
      [testCase0]
    where
      prog0 :: Prog GPU
      prog0 =
        "\
        \entry(\"main\",\
        \      {xss: [][]i64},\
        \      {[]i64})\
        \  entry_main (n_5142 : i64,\
        \              m_5143 : i64,\
        \              xss_5144 : [n_5142][m_5143]i64)\
        \  : {[n_5142]i64#([2], [0])} = {\
        \  let {segmap_group_size_5202 : i64} =\
        \    get_size(segmap_group_size_5190, thread_block_size)\
        \  let {segmap_usable_groups_5203 : i64} =\
        \    sdiv_up64(n_5142, segmap_group_size_5202)\
        \  let {defunc_0_map_res_5204 : [n_5142]i64} =\
        \    segmap(thread; ; grid=segmap_usable_groups_5203; blocksize=segmap_group_size_5202)\
        \    (gtid_5205 < n_5142) (~phys_tid_5206) : {i64} {\
        \      let {defunc_0_f_res_5208 : i64} =\
        \        loop {acc_5210 : i64} = {0i64}\
        \        for i_5209:i64 < m_5143 do {\
        \          let {b_5211 : i64} =\
        \            xss_5144[gtid_5205, i_5209]\
        \          let {defunc_0_f_res_5212 : i64} =\
        \            add64(acc_5210, b_5211)\
        \          in {defunc_0_f_res_5212}\
        \        }\
        \      return {returns defunc_0_f_res_5208}\
        \    }\
        \  in {defunc_0_map_res_5204}\
        \}"
  