# Copyright 2023 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

# Tests of end-to-end IREE support for individual ops in the StableHLO dialect.
# Each test file should have a name matching the corresponding StableHLO op and test only the
# functionality of that op (though may make use of other ops where necessary). Tests should be
# written using the IREE Check framework and should always pass on the reference VMVX backend.
# See https://github.com/openxla/iree/blob/main/docs/developers/developing_iree/testing_guide.md#iree-core-end-to-end-tests.

load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite")

package(
    features = ["layering_check"],
    licenses = ["notice"],  # Apache 2.0
)

iree_check_single_backend_test_suite(
    name = "check_llvm-cpu_local-task",
    srcs = enforce_glob(
        # keep sorted
        [
            "abs.mlir",
            "add.mlir",
            "batch_norm_inference.mlir",
            "bitcast_convert.mlir",
            "broadcast.mlir",
            "broadcast_add.mlir",
            "broadcast_in_dim.mlir",
            "clamp.mlir",
            "compare.mlir",
            "complex.mlir",
            "concatenate.mlir",
            "constant.mlir",
            "convert.mlir",
            "convolution.mlir",
            "cosine.mlir",
            "divide.mlir",
            "dot.mlir",
            "dot_bf16.mlir",
            "dot_general.mlir",
            "dynamic_slice.mlir",
            "dynamic_update_slice.mlir",
            "exponential.mlir",
            "exponential_fp16.mlir",
            "exponential_minus_one.mlir",
            "fft.mlir",
            "finite.mlir",
            "floor.mlir",
            "gather.mlir",
            "iota.mlir",
            "log.mlir",
            "log_plus_one.mlir",
            "maximum.mlir",
            "minimum.mlir",
            "multiply.mlir",
            "negate.mlir",
            "pad.mlir",
            "pow.mlir",
            "reduce.mlir",
            "reduce_window.mlir",
            "remainder.mlir",
            "reshape.mlir",
            "reverse.mlir",
            "rng_normal.mlir",
            "rng_uniform.mlir",
            "round.mlir",
            "rsqrt.mlir",
            "scatter.mlir",
            "scatter_dynamic.mlir",
            "select.mlir",
            "sine.mlir",
            "slice.mlir",
            "sort.mlir",
            "sqrt.mlir",
            "subtract.mlir",
            "tanh.mlir",
            "torch_index_select.mlir",
            "transpose.mlir",
            "while.mlir",
        ],
        include = ["*.mlir"],
        exclude = [],
    ),
    compiler_flags = ["--iree-input-type=stablehlo"],
    driver = "local-task",
    target_backend = "llvm-cpu",
)

iree_check_single_backend_test_suite(
    name = "check_vmvx_local-task",
    srcs = enforce_glob(
        # keep sorted
        [
            "abs.mlir",
            "add.mlir",
            "batch_norm_inference.mlir",
            "bitcast_convert.mlir",
            "broadcast.mlir",
            "broadcast_add.mlir",
            "broadcast_in_dim.mlir",
            "clamp.mlir",
            "compare.mlir",
            "complex.mlir",
            "concatenate.mlir",
            "constant.mlir",
            "convert.mlir",
            "convolution.mlir",
            "cosine.mlir",
            "divide.mlir",
            "dot.mlir",
            "dot_bf16.mlir",
            "dot_general.mlir",
            "dynamic_slice.mlir",
            "dynamic_update_slice.mlir",
            "exponential.mlir",
            "exponential_minus_one.mlir",
            "fft.mlir",
            "finite.mlir",
            "floor.mlir",
            "gather.mlir",
            "iota.mlir",
            "log.mlir",
            "log_plus_one.mlir",
            "maximum.mlir",
            "minimum.mlir",
            "multiply.mlir",
            "negate.mlir",
            "pad.mlir",
            "pow.mlir",
            "reduce.mlir",
            "reduce_window.mlir",
            "remainder.mlir",
            "reshape.mlir",
            "reverse.mlir",
            "rng_normal.mlir",
            "rng_uniform.mlir",
            "round.mlir",
            "rsqrt.mlir",
            "scatter.mlir",
            "scatter_dynamic.mlir",
            "select.mlir",
            "sine.mlir",
            "slice.mlir",
            "sort.mlir",
            "sqrt.mlir",
            "subtract.mlir",
            "tanh.mlir",
            "torch_index_select.mlir",
            "transpose.mlir",
            "while.mlir",
        ],
        include = ["*.mlir"],
        exclude = [
            "exponential_fp16.mlir",
        ],
    ),
    compiler_flags = ["--iree-input-type=stablehlo"],
    driver = "local-task",
    target_backend = "vmvx",
)

iree_check_single_backend_test_suite(
    name = "check_vulkan-spirv_vulkan",
    srcs = enforce_glob(
        # keep sorted
        [
            "abs.mlir",
            "add.mlir",
            "batch_norm_inference.mlir",
            "bitcast_convert.mlir",
            "broadcast.mlir",
            "broadcast_add.mlir",
            "broadcast_in_dim.mlir",
            "clamp.mlir",
            "compare.mlir",
            "complex.mlir",
            "concatenate.mlir",
            "constant.mlir",
            "convert.mlir",
            "convolution.mlir",
            "cosine.mlir",
            "divide.mlir",
            "dot.mlir",
            "dot_bf16.mlir",
            "dot_general.mlir",
            "dynamic_slice.mlir",
            "dynamic_update_slice.mlir",
            "exponential.mlir",
            "exponential_minus_one.mlir",
            "finite.mlir",
            "floor.mlir",
            "gather.mlir",
            "iota.mlir",
            "log.mlir",
            "log_plus_one.mlir",
            "maximum.mlir",
            "minimum.mlir",
            "multiply.mlir",
            "negate.mlir",
            "pad.mlir",
            "pow.mlir",
            "reduce.mlir",
            "reduce_window.mlir",
            "remainder.mlir",
            "reshape.mlir",
            "rng_normal.mlir",
            "rng_uniform.mlir",
            "round.mlir",
            "rsqrt.mlir",
            "scatter.mlir",
            "scatter_dynamic.mlir",
            "select.mlir",
            "sine.mlir",
            "slice.mlir",
            "sort.mlir",
            "sqrt.mlir",
            "subtract.mlir",
            "tanh.mlir",
            "torch_index_select.mlir",
            "transpose.mlir",
            "while.mlir",
        ],
        include = ["*.mlir"],
        exclude = [
            "exponential_fp16.mlir",
            "fft.mlir",  # TODO(#9583)
            "reverse.mlir",  # TODO(#12415): disabled due to miscompilation on Pixel 6.
        ],
    ),
    compiler_flags = ["--iree-input-type=stablehlo"],
    driver = "vulkan",
    target_backend = "vulkan-spirv",
)

# Check host features compilation (LLVM backend with host cpu features).
iree_check_single_backend_test_suite(
    name = "check_llvm-cpu-host_local-task",
    srcs = enforce_glob(
        # keep sorted
        [
            "abs.mlir",
            "add.mlir",
            "batch_norm_inference.mlir",
            "bitcast_convert.mlir",
            "broadcast.mlir",
            "broadcast_add.mlir",
            "broadcast_in_dim.mlir",
            "clamp.mlir",
            "compare.mlir",
            "complex.mlir",
            "concatenate.mlir",
            "constant.mlir",
            "convert.mlir",
            "convolution.mlir",
            "cosine.mlir",
            "divide.mlir",
            "dot.mlir",
            "dot_bf16.mlir",
            "dot_general.mlir",
            "dynamic_slice.mlir",
            "dynamic_update_slice.mlir",
            "exponential.mlir",
            "exponential_fp16.mlir",
            "exponential_minus_one.mlir",
            "fft.mlir",
            "finite.mlir",
            "floor.mlir",
            "gather.mlir",
            "iota.mlir",
            "log.mlir",
            "log_plus_one.mlir",
            "maximum.mlir",
            "minimum.mlir",
            "multiply.mlir",
            "negate.mlir",
            "pad.mlir",
            "pow.mlir",
            "reduce.mlir",
            "reduce_window.mlir",
            "remainder.mlir",
            "reshape.mlir",
            "reverse.mlir",
            "rng_normal.mlir",
            "rng_uniform.mlir",
            "round.mlir",
            "rsqrt.mlir",
            "scatter.mlir",
            "scatter_dynamic.mlir",
            "select.mlir",
            "sine.mlir",
            "slice.mlir",
            "sort.mlir",
            "sqrt.mlir",
            "subtract.mlir",
            "tanh.mlir",
            "torch_index_select.mlir",
            "transpose.mlir",
            "while.mlir",
        ],
        include = ["*.mlir"],
        exclude = [],
    ),
    compiler_flags = [
        "--iree-input-type=stablehlo",
        "--iree-llvmcpu-target-cpu-features=host",
    ],
    driver = "local-task",
    # Building and testing must be on the same architecture, which doesn't work
    # with remote execution in general.
    tags = [
        "hostonly",
        "local",
    ],
    target_backend = "llvm-cpu",
)

test_suite(
    name = "check",
    tests = [
        ":check_llvm-cpu-host_local-task",
        ":check_llvm-cpu_local-task",
        ":check_vmvx_local-task",
        ":check_vulkan-spirv_vulkan",
    ],
)

iree_check_single_backend_test_suite(
    name = "check_cuda_graph",
    srcs = enforce_glob(
        # keep sorted
        [
            "abs.mlir",
            "add.mlir",
            "batch_norm_inference.mlir",
            "bitcast_convert.mlir",
            "broadcast.mlir",
            "broadcast_add.mlir",
            "broadcast_in_dim.mlir",
            "clamp.mlir",
            "compare.mlir",
            "complex.mlir",
            "concatenate.mlir",
            "constant.mlir",
            "convert.mlir",
            "convolution.mlir",
            "cosine.mlir",
            "divide.mlir",
            "dot.mlir",
            "dot_bf16.mlir",
            "dot_general.mlir",
            "dynamic_slice.mlir",
            "dynamic_update_slice.mlir",
            "exponential.mlir",
            "exponential_fp16.mlir",
            "exponential_minus_one.mlir",
            "fft.mlir",
            "finite.mlir",
            "floor.mlir",
            "gather.mlir",
            "iota.mlir",
            "log.mlir",
            "log_plus_one.mlir",
            "maximum.mlir",
            "minimum.mlir",
            "multiply.mlir",
            "negate.mlir",
            "pad.mlir",
            "pow.mlir",
            "reduce.mlir",
            "reduce_window.mlir",
            "remainder.mlir",
            "reshape.mlir",
            "reverse.mlir",
            "rng_normal.mlir",
            "rng_uniform.mlir",
            "round.mlir",
            "rsqrt.mlir",
            "scatter.mlir",
            "scatter_dynamic.mlir",
            "select.mlir",
            "sine.mlir",
            "slice.mlir",
            "sort.mlir",
            "sqrt.mlir",
            "subtract.mlir",
            "tanh.mlir",
            "torch_index_select.mlir",
            "transpose.mlir",
            "while.mlir",
        ],
        include = ["*.mlir"],
        exclude = [],
    ),
    compiler_flags = ["--iree-input-type=stablehlo"],
    driver = "cuda",
    runner_args = ["--cuda_use_streams=false"],
    tags = [
        # CUDA cuInit fails with sanitizer on.
        "noasan",
        "nomsan",
        "notsan",
        "noubsan",
        "requires-gpu-nvidia",
    ],
    target_backend = "cuda",
)

# Run cuda tests using stream command buffer
iree_check_single_backend_test_suite(
    name = "check_cuda_streams",
    srcs = enforce_glob(
        # keep sorted
        [
            "abs.mlir",
            "add.mlir",
            "batch_norm_inference.mlir",
            "bitcast_convert.mlir",
            "broadcast.mlir",
            "broadcast_add.mlir",
            "broadcast_in_dim.mlir",
            "clamp.mlir",
            "compare.mlir",
            "complex.mlir",
            "concatenate.mlir",
            "constant.mlir",
            "convert.mlir",
            "convolution.mlir",
            "cosine.mlir",
            "divide.mlir",
            "dot.mlir",
            "dot_bf16.mlir",
            "dot_general.mlir",
            "dynamic_slice.mlir",
            "dynamic_update_slice.mlir",
            "exponential.mlir",
            "exponential_fp16.mlir",
            "exponential_minus_one.mlir",
            "fft.mlir",
            "finite.mlir",
            "floor.mlir",
            "gather.mlir",
            "iota.mlir",
            "log.mlir",
            "log_plus_one.mlir",
            "maximum.mlir",
            "minimum.mlir",
            "multiply.mlir",
            "negate.mlir",
            "pad.mlir",
            "pow.mlir",
            "reduce.mlir",
            "reduce_window.mlir",
            "remainder.mlir",
            "reshape.mlir",
            "reverse.mlir",
            "rng_normal.mlir",
            "rng_uniform.mlir",
            "round.mlir",
            "rsqrt.mlir",
            "scatter.mlir",
            "scatter_dynamic.mlir",
            "select.mlir",
            "sine.mlir",
            "slice.mlir",
            "sort.mlir",
            "sqrt.mlir",
            "subtract.mlir",
            "tanh.mlir",
            "torch_index_select.mlir",
            "transpose.mlir",
            "while.mlir",
        ],
        include = ["*.mlir"],
        exclude = [],
    ),
    compiler_flags = ["--iree-input-type=stablehlo"],
    driver = "cuda",
    runner_args = ["--cuda_use_streams=true"],
    tags = [
        # CUDA cuInit fails with sanitizer on.
        "noasan",
        "nomsan",
        "notsan",
        "noubsan",
        "requires-gpu-nvidia",
    ],
    target_backend = "cuda",
)