load("//devtools/python/blaze:pytype.bzl", "pytype_strict_library")
load("//learning/brain/google/xla:libtpu_build_defs.bzl", "tpu_cc_library")
load("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_test")

package(
    default_visibility = [
        "//platforms/xla:__subpackages__",
        "//platforms/xla/mosaic:__subpackages__",
    ],
)

[
    py_test(
        name = "multidevice_test" + suffix,
        srcs = ["multidevice_test.py"],
        args = args,
        main = "multidevice_test.py",
        tags = tags,
        deps = [
            ":test_util",
            "//learning/brain/research/jax:tpu_simulator_support",
            "//learning/brain/research/jax:tpu_support",
            "//platforms/xla/mosaic/python/dialects:llo_dialect",
            "//testing/pybase",
            "//testing/pybase:parameterized",
            "//third_party/llvm/llvm-project/mlir:IR",
            "//third_party/py/jax",
            "//third_party/py/mlir:arithmetic_dialect",
            "//third_party/py/mlir:builtin_dialect",
            "//third_party/py/mlir:ir",
            "//third_party/py/mlir:memref_dialect",
            "//third_party/py/mlir:vector_dialect",
            "//third_party/py/mlir/_mlir_libs:_mlirRegisterEverything",
            "//third_party/py/numpy",
        ],
    )
    for suffix, args, tags in [
        (
            "_vl_grm",
            [
                "--deepsea_chips_per_host_bounds=2,2,1",
                "--deepsea_version=viperfish",
                "--deepsea_variant=lite",
            ],
            [
                "noasan",
                "notsan",
                "nodebug",
            ],  # Takes way too long with sanitizers or debug mode.
        ),
        (
            "_vl",
            [],
            ["requires-viperlite:8"],
        ),
    ]
]

[
    py_test(
        name = "tpu_custom_call_test" + suffix + py_suffix,
        timeout = "long",
        srcs = ["tpu_custom_call_test.py"],
        args = args,
        env = env,
        main = "tpu_custom_call_test.py",
        shard_count = shard_count,
        tags = tags,
        deps = [
            ":test_util",
            "//learning/brain/research/jax:tpu_iss_support",
            "//learning/brain/research/jax:tpu_support",
            "//platforms/xla/mosaic/python/dialects:llo_dialect",
            "//platforms/xla/mosaic/tests:mlir_interpreter",
            "//testing/pybase",
            "//testing/pybase:parameterized",
            "//third_party/llvm/llvm-project/mlir:IR",
            "//third_party/py/absl/flags",
            "//third_party/py/absl/testing:flagsaver",
            "//third_party/py/hypothesis",
            "//third_party/py/jax",
            "//third_party/py/mlir:arithmetic_dialect",
            "//third_party/py/mlir:builtin_dialect",
            "//third_party/py/mlir:control_flow_dialect",
            "//third_party/py/mlir:ir",
            "//third_party/py/mlir:memref_dialect",
            "//third_party/py/mlir:scf_dialect",
            "//third_party/py/mlir:vector_dialect",
            "//third_party/py/mlir/_mlir_libs:_mlirRegisterEverything",
            "//third_party/py/numpy",
            "//third_party/tensorflow/compiler/xla/python:xla_extension",
        ],
    )
    for py_suffix, env in [
        (
            "_py",
            {"MOSAIC_USE_PYTHON_PIPELINE": "True"},
        ),
        (
            "",
            {"MOSAIC_USE_PYTHON_PIPELINE": "False"},
        ),
    ]
    for suffix, args, tags, shard_count in [
        (
            "_jf_iss",
            [
                "--xla_tpu_mosaic_fusion",
                "--xla_tpu_autofdo=false",
            ],
            [
                "noasan",
                "notsan",
                "nodebug",
            ],
            8,
        ),
        (
            "_pf_iss",
            [
                "--deepsea_version=pufferfish",
                "--xla_tpu_autofdo=false",
            ],
            [
                "noasan",
                "notsan",
                "nodebug",
            ],
            10,
        ),
        (
            "_vl_iss",
            [
                "--deepsea_version=viperfish",
                "--deepsea_variant=lite",
                "--xla_tpu_autofdo=false",
            ],
            [
                "notap",  # vl_iss on TAP drops precision http://shortn/_6XHhU0fZml.
            ],
            10,
        ),
        (
            "_gl_iss",
            [
                "--deepsea_version=ghostlite",
                "--xla_tpu_autofdo=false",
            ],
            [
                "noasan",
                "notsan",
                "nodebug",
            ],
            10,
        ),
        (
            "_gf_iss",
            [
                "--deepsea_version=ghostfish",
                "--xla_tpu_autofdo=false",
            ],
            ["notap"],
            10,
        ),
        (
            "_pf_megacore_iss",
            [
                "--deepsea_version=pufferfish",
                "--deepsea_chip_config_name=megacore",
                "--iss_topology_mode=ISS_TOPOLOGY_MODE_MULTI_TC",
                "--xla_tpu_autofdo=false",
            ],
            [
                "noasan",
                "notsan",
                "nodebug",
            ],
            10,
        ),
        (
            "_jf",
            ["--xla_tpu_mosaic_fusion"],
            ["requires-jellyfish"],
            1,
        ),
        (
            "_pf",
            [],
            ["requires-pufferfish"],
            4,
        ),
        (
            "_vl",
            [],
            ["requires-viperlite"],
            10,
        ),
        (
            "_pf_fusion",
            ["--xla_tpu_mosaic_fusion"],
            ["requires-pufferfish"],
            4,
        ),
        (
            "_vl_fusion",
            ["--xla_tpu_mosaic_fusion"],
            ["requires-viperlite"],
            10,
        ),
    ]
]

py_test(
    name = "apply_vector_layout_test",
    timeout = "long",
    srcs = ["apply_vector_layout_test.py"],
    main = "apply_vector_layout_test.py",
    shard_count = 4,
    deps = [
        ":test_util",
        "//platforms/xla/mosaic/tests:mlir_interpreter",
        "//testing/pybase",
        "//testing/pybase:parameterized",
        "//third_party/py/hypothesis",
        "//third_party/py/jax",
        "//third_party/py/mlir:arithmetic_dialect",
        "//third_party/py/mlir:func_dialect",
        "//third_party/py/mlir:ir",
        "//third_party/py/mlir:math_dialect",
        "//third_party/py/mlir:vector_dialect",
        "//third_party/py/numpy",
    ],
)

tpu_cc_library(
    name = "megacore_adjuster_base",
    hdrs = ["megacore_adjuster_base.h"],
    compatible_with = ["//buildenv/target:libtpu"],
    deps_for_viperfish = [
        "//platforms/xla/sparse_core:debug_info",
        "//platforms/xla/sparse_core:lowering_util",
    ],
    deps = [
        "//learning/brain/tpu/runtime:libtpu_support_macros",
        "//platforms/xla/service/jellyfish:llo_ir",
        "//platforms/xla/service/jellyfish:target_base",
        "//platforms/xla/service/jellyfish/lowering:math_util",
        "//third_party/absl/algorithm:container",
        "//third_party/absl/log",
        "//third_party/absl/log:check",
        "//third_party/absl/types:optional",
        "//third_party/absl/types:span",
        "//third_party/llvm/llvm-project/llvm:Support",
        "//third_party/llvm/llvm-project/mlir:ArithDialect",
        "//third_party/llvm/llvm-project/mlir:IR",
        "//third_party/py/jax/jaxlib/mosaic:tpu_dialect",
        "//third_party/tensorflow/compiler/xla:status",
        "//third_party/tensorflow/compiler/xla:status_macros",
        "//third_party/tensorflow/compiler/xla:util",
        "//third_party/tensorflow/compiler/xla:xla_data_proto_cc",
        "//third_party/tensorflow/tsl/platform:statusor",
    ],
)

cc_library(
    name = "custom_call_emitter",
    srcs = ["custom_call_emitter.cc"],
    compatible_with = ["//buildenv/target:libtpu"],
    deps = [
        ":megacore_adjuster_base",
        "//base:googleinit",
        "//base:vlog",
        "//learning/brain/tpu/runtime:libtpu_support_macros",
        "//learning/brain/tpu/runtime:tpu_version",
        "//platforms/xla/mosaic:llo_dialect",
        "//platforms/xla/mosaic:tpu_passes",
        "//platforms/xla/service:fusion_util",
        "//platforms/xla/service/jellyfish:async_collective_fusion_util",
        "//platforms/xla/service/jellyfish:custom_call_registration",
        "//platforms/xla/service/jellyfish:fusion_options",
        "//platforms/xla/service/jellyfish:hardware_layout",
        "//platforms/xla/service/jellyfish:hlo_deduplication",
        "//platforms/xla/service/jellyfish:llo_ir",
        "//platforms/xla/service/jellyfish:llo_program_shared_registry",
        "//platforms/xla/service/jellyfish:llo_sdc_reporter",
        "//platforms/xla/service/jellyfish:memory_space_enum",
        "//platforms/xla/service/jellyfish:target_base",
        "//platforms/xla/service/jellyfish:tpu_compilation_environment",
        "//platforms/xla/service/jellyfish:tpu_compilation_environment_cc_proto",
        "//platforms/xla/service/jellyfish:tpu_instruction_fusion",
        "//platforms/xla/service/jellyfish:transfer_size_util",
        "//platforms/xla/service/jellyfish/cost_model",
        "//platforms/xla/service/jellyfish/cost_model:cost_model_util",
        "//platforms/xla/service/jellyfish/lowering:backend_config_util",
        "//platforms/xla/service/jellyfish/lowering:backend_configs_cc_proto",
        "//platforms/xla/service/jellyfish/lowering:deep_copy_util",
        "//platforms/xla/service/jellyfish/lowering:fusion_emitter",
        "//platforms/xla/service/jellyfish/lowering:fusion_util",
        "//platforms/xla/service/jellyfish/lowering:group_utils",
        "//platforms/xla/service/jellyfish/lowering:lowering_util",
        "//platforms/xla/service/jellyfish/lowering:net_util",
        "//platforms/xla/service/jellyfish/lowering:op_emitter",
        "//platforms/xla/service/jellyfish/lowering:param_input",
        "//platforms/xla/service/jellyfish/lowering:pipeline_emitter",
        "//platforms/xla/service/jellyfish/lowering:windowing_util",
        "//third_party/absl/algorithm:container",
        "//third_party/absl/container:flat_hash_map",
        "//third_party/absl/flags:flag",
        "//third_party/absl/log",
        "//third_party/absl/log:check",
        "//third_party/absl/status",
        "//third_party/absl/strings",
        "//third_party/absl/strings:str_format",
        "//third_party/absl/strings:string_view",
        "//third_party/absl/types:optional",
        "//third_party/absl/types:span",
        "//third_party/llvm/llvm-project/llvm:Support",
        "//third_party/llvm/llvm-project/mlir:FuncDialect",
        "//third_party/llvm/llvm-project/mlir:FuncExtensions",
        "//third_party/llvm/llvm-project/mlir:IR",
        "//third_party/llvm/llvm-project/mlir:LinalgDialect",
        "//third_party/llvm/llvm-project/mlir:Parser",
        "//third_party/llvm/llvm-project/mlir:Pass",
        "//third_party/llvm/llvm-project/mlir:Support",
        "//third_party/llvm/llvm-project/mlir:TensorDialect",
        "//third_party/llvm/llvm-project/mlir:Transforms",
        "//third_party/py/jax/jaxlib/mosaic:tpu_dialect",
        "//third_party/tensorflow/compiler/xla:shape_util",
        "//third_party/tensorflow/compiler/xla:status",
        "//third_party/tensorflow/compiler/xla:statusor",
        "//third_party/tensorflow/compiler/xla:util",
        "//third_party/tensorflow/compiler/xla:xla_data_proto_cc",
        "//third_party/tensorflow/compiler/xla/hlo/ir:hlo",
        "//third_party/tensorflow/compiler/xla/mlir/utils:type_util",
        "//third_party/tensorflow/compiler/xla/mlir_hlo",
        "//third_party/tensorflow/compiler/xla/mlir_hlo:mhlo_passes",
        "//third_party/tensorflow/compiler/xla/service:hlo_cost_analysis",
        "//third_party/tensorflow/core/platform:path",
        "//third_party/tensorflow/tsl/platform:env",
        "//third_party/tensorflow/tsl/platform:errors",
        "//third_party/tensorflow/tsl/platform:status",
        "//third_party/tensorflow/tsl/platform:statusor",
        "//util/task:status",
    ],
    alwayslink = True,
)

pytype_strict_library(
    name = "test_util",
    testonly = True,
    srcs = ["test_util.py"],
    deps = [
        "//platforms/xla/mosaic/python/dialects:llo_dialect",
        "//testing/pybase:parameterized",
        "//third_party/py/hypothesis",
        "//third_party/py/jax",
        "//third_party/py/jax:mosaic",
        "//third_party/py/mlir:func_dialect",
        "//third_party/py/mlir:ir",
        "//third_party/py/mlir:mhlo_dialect",
        "//third_party/py/mlir/_mlir_libs:_mlirRegisterEverything",
        "//third_party/py/numpy",
    ],
)