TRIBITS_CONFIGURE_FILE(${PACKAGE_NAME}_config.h)
TRIBITS_INCLUDE_DIRECTORIES(${CMAKE_CURRENT_BINARY_DIR})

TRIBITS_ADD_EXECUTABLE(ex1 SOURCES ex1.cpp)

IF(Tpetra_INST_DOUBLE)
  TRIBITS_ADD_EXECUTABLE(RelaxationWithEquilibration SOURCES RelaxationWithEquilibration.cpp)
ENDIF()

TRIBITS_ADD_EXECUTABLE(
  BlockTriDiagonalSolver
  SOURCES BlockTriDi.cpp
  CATEGORIES BASIC PERFORMANCE
)

TRIBITS_ADD_TEST(
  BlockTriDiagonalSolver
  NAME BlockTriDiExample
  ARGS ""
  COMM serial mpi
  NUM_MPI_PROCS 1-4
  STANDARD_PASS_OUTPUT
)

ASSERT_DEFINED(${PACKAGE_NAME}_ENABLE_Galeri)

IF(${PACKAGE_NAME}_ENABLE_Galeri)
  # Correctness test with maximum block size (32)
  # Use a small grid so that GPU memory requirement isn't too large
  # Block TriDi
  TRIBITS_ADD_TEST(
    BlockTriDiagonalSolver
    NAME BlockTriDiLargeBlock
    ARGS "--matrixType=Laplace3D --blockSize=32 --nx=20 --ny=20 --nz=20"
    COMM serial mpi
    NUM_MPI_PROCS 1-4
    STANDARD_PASS_OUTPUT
  )
  # Block TriDi with Schur line splitting
  TRIBITS_ADD_TEST(
    BlockTriDiagonalSolver
    NAME BlockTriDiLargeBlockSchur
    ARGS "--matrixType=Laplace3D --blockSize=32 --nx=20 --ny=20 --nz=20 --sublinesPerLine=1 --sublinesPerLineSchur=2"
    COMM serial mpi
    NUM_MPI_PROCS 1-4
    STANDARD_PASS_OUTPUT
  )
  # Block TriDi with Schur line splitting + multiple RHS vectors
  TRIBITS_ADD_TEST(
    BlockTriDiagonalSolver
    NAME BlockTriDiSchurMultivector
    ARGS "--matrixType=Laplace3D --blockSize=7 --nx=20 --ny=20 --nz=20 --sublinesPerLine=1 --sublinesPerLineSchur=2 --numVecs=3"
    COMM serial mpi
    NUM_MPI_PROCS 1-4
    STANDARD_PASS_OUTPUT
  )
  # Block Jacobi
  TRIBITS_ADD_TEST(
    BlockTriDiagonalSolver
    NAME BlockTriDiLargeBlockJacobi
    ARGS "--matrixType=Laplace3D --blockSize=32 --nx=20 --ny=20 --nz=20 --sublinesPerLine=-1"
    COMM serial mpi
    NUM_MPI_PROCS 1-4
    STANDARD_PASS_OUTPUT
  )
ENDIF()

IF(${PACKAGE_NAME}_ENABLE_Galeri)

  SET(blockSize 11)
  SET(nx 256)
  SET(ny 128)
  SET(nz 128)

  SET(COMMON_ARGS "--withoutOverlapCommAndComp --matrixType=Laplace3D --blockSize=${blockSize} --nx=${nx} --ny=${ny} --nz=${nz} --tol=0 --numIters=5 --withStackedTimer --numRepeats=5")
  SET(BTDS_ARGS "${COMMON_ARGS} --sublinesPerLine=1 --sublinesPerLineSchur=1")
  SET(BTDS_SCHUR_ARGS "${COMMON_ARGS} --sublinesPerLine=1 --sublinesPerLineSchur=2")
  SET(BTDS_DEFAULT_SCHUR_ARGS "${COMMON_ARGS} --sublinesPerLine=1 --sublinesPerLineSchur=-1")
  SET(BLOCK_JACOBI_ARGS "${COMMON_ARGS} --sublinesPerLine=-1")

  SET(TEST_ARGS
      ${BTDS_ARGS}
      ${BTDS_SCHUR_ARGS}
      ${BTDS_DEFAULT_SCHUR_ARGS}
      ${BLOCK_JACOBI_ARGS}
  )
  SET(TEST_NAMES
      "BTDS"
      "BTDS_Schur"
      "BTDS_Schur_Default"
      "BTDS_Jacobi"
  )
  SET(ARCHS "CUDA" "HIP")

  LIST(LENGTH TEST_ARGS len1)
  MATH(EXPR N_TESTS "${len1} - 1")

  ASSERT_DEFINED(Tpetra_INST_CUDA Tpetra_INST_HIP)

  FOREACH(ARCH ${ARCHS})
    IF(Tpetra_INST_${ARCH})
      MESSAGE(STATUS "Ifpack2: BlockTriDiagonalSolver ${ARCH} test ENABLED")
      SET(COUNTER 2)
      SET(MAX_COUNT 8)

      WHILE(COUNTER LESS MAX_COUNT)
        MATH(EXPR NP "1 << ${COUNTER}")
        FOREACH(I_TEST RANGE ${N_TESTS})
          LIST(
            GET
            TEST_ARGS
            ${I_TEST}
            TEST_ARG
          )
          LIST(
            GET
            TEST_NAMES
            ${I_TEST}
            TEST_NAME
          )
          TRIBITS_ADD_TEST(
            BlockTriDiagonalSolver
            NAME "${TEST_NAME}_${ARCH}"
            COMM mpi
            ARGS "${TEST_ARG} --problemName=\"${TEST_NAME} ${ARCH} ${blockSize}x${nx}x${ny}x${nz} ${NP} ranks\""
            NUM_MPI_PROCS ${NP}
            RUN_SERIAL
            CATEGORIES PERFORMANCE
          )

        ENDFOREACH()

        MATH(EXPR COUNTER "${COUNTER} + 1")
      ENDWHILE()
    ENDIF()
  ENDFOREACH()
ENDIF()
