# Copyright 2022 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

# End-to-end matrix multiplication tests.

load("@rules_python//python:defs.bzl", "py_binary")
load("//build_tools/bazel:iree_e2e_generated_runner_test.bzl", "iree_generated_e2e_runner_test")

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

py_binary(
    name = "generate_e2e_matmul_tests",
    srcs = [
        "common.py",
        "compilation_info.py",
        "generate_code.py",
        "generate_code_mx.py",
        "generate_e2e_matmul_tests.py",
    ],
)

###########################################################################
##
## LLVMCPU backend
##
###########################################################################

# LLVMCPU, non-data-tiling, no microkernels
[iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_nondt_%s_%s" % (lhs_rhs_type, acc_type),
    compiler_flags = [
        "--iree-opt-data-tiling=false",
        "--iree-llvmcpu-enable-ukernels=none",
        "--iree-llvmcpu-enable-scalable-vectorization",
        "--iree-llvmcpu-target-triple=aarch64-unknown-unknown",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
    ],
    tags = [
        # f16/bf16 trigger internal LLVM assertion errors on riscv and wasm.
        "noriscv",
        "nowasm",
    ] if (lhs_rhs_type == "f16" or lhs_rhs_type == "bf16") else [],
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = ["generic"] +
                                   # Widening matmuls fail to lower for SVE.
                                   (["arm_64:sve:+sve"] if lhs_rhs_type == acc_type else []),
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for (lhs_rhs_type, acc_type) in [
    # ("i8", "i32"),  # TODO(#15800): enable once compile time is reasonable
    # ("f32", "f32"),  # TODO(#15800): enable once compile time is reasonable
    # ("f16", "f16"),  # TODO(#15800): enable once compile time is reasonable
    # ("f16", "f32"),  # TODO(#15800): enable once compile time is reasonable
    # TODO(#15258): enable bf16 tests when that bug is fixed.
    # ("bf16", "bf16"),
    # ("bf16", "f32"),
]]

PREPROCESSING_TRANSPOSE_LHS = "--iree-preprocessing-pass-pipeline=builtin.module\\(util.func\\(iree-preprocessing-transpose-matmul-pass{input=lhs}\\)\\)"

PREPROCESSING_PEEL = "--iree-llvmcpu-vector-pproc-strategy=peel"

# LLVMCPU, non-data-tiling, no microkernels, ArmSME
[iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_arm_sme_nondt_%s_%s%s" % (
        dtype,
        "_transpose_lhs" if transpose_lhs else "",
        "_peel" if peel else "",
    ),
    compiler_flags = [
                         "--iree-opt-data-tiling=false",
                         "--iree-llvmcpu-enable-scalable-vectorization",
                         "--iree-llvmcpu-target-triple=aarch64-unknown-unknown",
                     ] + ([PREPROCESSING_TRANSPOSE_LHS] if transpose_lhs else []) +
                     ([PREPROCESSING_PEEL] if peel else []),
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % dtype,
        "--acc_type=%s" % dtype,
    ],
    tags = [
        "requires-arm-sme",
    ],
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = ["arm_64:sme:+sme"],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for dtype in [
    "f32",
    # f64 disabled because it wasn't supported by the test generator at the time
    # this was added. When adding it in the future, consider passing
    # --iree-input-demote-f64-to-f32=false to the compiler.
    # "f64"
] for transpose_lhs in [
    True,
    False,
] for peel in [
    True,
    False,
]]

X86_64_AVX2 = [
    "+avx",
    "+avx2",
    "+fma",
    "+f16c",
]

X86_64_AVX512 = X86_64_AVX2 + [
    "+avx512f",
    "+avx512vl",
    "+avx512cd",
    "+avx512bw",
    "+avx512dq",
]

X86_64_AVX512_VNNI = X86_64_AVX512 + [
    "+avx512vnni",
]

X86_64_AVX512_BF16 = X86_64_AVX512 + [
    "+avx512bf16",
]

# LLVMCPU, data-tiling, data-tiling + ukernels.
[iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_dt%s_%s_%s" % (
        ("_uk" if use_uk else ""),
        lhs_rhs_type,
        acc_type,
    ),
    compiler_flags = [
        "--iree-opt-data-tiling",
    ] + [
        "--iree-llvmcpu-enable-ukernels=%s" % ("all" if use_uk else "none"),
    ] + (["--iree-input-demote-f64-to-f32=false"] if acc_type == "f64" else []),
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
    ] + (["--shapes=small"] if acc_type == "f64" else []),
    tags = ([
        # f16/bf16 trigger internal LLVM assertion errors on riscv and wasm.
        "noriscv",
        "nowasm",
    ] if (lhs_rhs_type == "f16" or lhs_rhs_type == "bf16") else []) + ([
        # The non-ukernel (generic-vectorized) data-tiling path is not yet
        # optimized on RISC-V: with the large default shapes it runs for tens
        # of seconds under the QEMU emulator and trips the e2e timeout.
        # Skip it on RISC-V and re-add small-shape RISC-V coverage via the
        # dedicated riscv-only rules below.
        "noriscv",
    ] if (not use_uk and acc_type != "f64" and lhs_rhs_type not in [
        "f16",
        "bf16",
    ]) else []),
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = ["generic"] +
                                   ([
                                       "arm_64:dotprod:+dotprod",
                                       "arm_64:i8mm:+i8mm",
                                       "x86_64:avx512vnni:" + ",".join(X86_64_AVX512_VNNI),
                                       "riscv_64:ime:+v,+zvl256b,+xsmtvdot",
                                   ] if lhs_rhs_type == "i8" and acc_type == "i32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                   ] if lhs_rhs_type == "f32" and acc_type == "f32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "arm_64:fullfp16:+fullfp16",
                                   ] if lhs_rhs_type == "f16" and acc_type == "f16" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "arm_64:fp16fml:+fp16fml",
                                   ] if lhs_rhs_type == "f16" and acc_type == "f32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "x86_64:avx512bf16:" + ",".join(X86_64_AVX512_BF16),
                                       "arm_64:bf16:+bf16",
                                   ] if lhs_rhs_type == "bf16" and acc_type == "bf16" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "x86_64:avx512bf16:" + ",".join(X86_64_AVX512_BF16),
                                       "arm_64:bf16:+bf16",
                                   ] if lhs_rhs_type == "bf16" and acc_type == "f32" else []),
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for use_uk in [
    False,
    True,
] for (lhs_rhs_type, acc_type) in (
    [
        ("i8", "i32"),
        ("f32", "f32"),
        ("f64", "f64"),
        ("f16", "f16"),
        ("f16", "f32"),
        ("bf16", "bf16"),
        ("bf16", "f32"),
    ]
)]

# LLVMCPU, data-tiling lowering matmul to iree_codegen.inner_tiled
# (the new MMA-intrinsic path) instead of linalg.mmt4d.
[iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_dt_inner_tiled_%s_%s" % (lhs_rhs_type, acc_type),
    timeout = "moderate" if lhs_rhs_type == "f32" and acc_type == "f32" else None,
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-llvmcpu-enable-inner-tiled",
    ] + (["--iree-input-demote-f64-to-f32=false"] if acc_type == "f64" else []),
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
    ] + (["--shapes=small"] if acc_type == "f64" else []),
    tags = ([
        # f16/bf16 trigger internal LLVM assertion errors on riscv and wasm.
        "noriscv",
        "nowasm",
    ] if (lhs_rhs_type == "f16" or lhs_rhs_type == "bf16") else []),
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = ["generic"] +
                                   ([
                                       "x86_64:avx512vnni:" + ",".join(X86_64_AVX512_VNNI),
                                   ] if lhs_rhs_type == "i8" and acc_type == "i32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                   ] if lhs_rhs_type == "f32" and acc_type == "f32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                   ] if lhs_rhs_type == "f16" and acc_type == "f16" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                   ] if lhs_rhs_type == "f16" and acc_type == "f32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "x86_64:avx512bf16:" + ",".join(X86_64_AVX512_BF16),
                                   ] if lhs_rhs_type == "bf16" and acc_type == "bf16" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "x86_64:avx512bf16:" + ",".join(X86_64_AVX512_BF16),
                                   ] if lhs_rhs_type == "bf16" and acc_type == "f32" else []),
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for (lhs_rhs_type, acc_type) in [
    ("i8", "i32"),
    ("f32", "f32"),
    ("f64", "f64"),
    ("f16", "f16"),
    ("f16", "f32"),
    ("bf16", "bf16"),
    ("bf16", "f32"),
]]

# LLVMCPU, data-tiling lowering to iree_codegen.inner_tiled AND routing the
# inner_tiled op through a built-in C-bitcode microkernel via
# `--iree-llvmcpu-enable-llvm-ukernels=inner_tiled`. Restricted to the
# avx512bf16 variant, which is the one for which a bf16 ukernel
# (`iree_uk_mma_x86_avx512bf16_1x16x2_f32_bf16`) exists; this is the
# numerical end-to-end check that the ukernel (and the generic
# `intrinsics_{m,n,k}` unrolling + operand threading) computes correct
# results, complementing the IR-level lit tests under
# compiler/plugins/target/LLVMCPU/builtins/ukernel/test/.
iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_dt_inner_tiled_llvm_ukernel_bf16_f32",
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-llvmcpu-enable-inner-tiled",
        "--iree-llvmcpu-enable-llvm-ukernels=inner_tiled",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=bf16",
        "--acc_type=f32",
    ],
    tags = [
        # bf16 triggers internal LLVM assertion errors on riscv and wasm.
        "noriscv",
        "nowasm",
    ],
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = [
        "x86_64:avx512bf16:" + ",".join(X86_64_AVX512_BF16),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
)

# Companion of the bf16 test above for the symmetric signed-i8 -> i32 path:
# routes the data-tiled inner_tiled op through the C-bitcode microkernel
# `iree_uk_mma_x86_avx512vnni_16x16x2_i32_i8_casti16` (x86 AVX-512 VNNI
# `vpdpwssd`, with the s8 inputs widened to i16). Restricted to the
# avx512vnni variant, which is the one for which that ukernel exists; this
# is the numerical end-to-end check that the ukernel computes correct
# results, complementing the IR-level lit tests under
# compiler/plugins/target/LLVMCPU/builtins/ukernel/test/.
iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_dt_inner_tiled_llvm_ukernel_i8_i32",
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-llvmcpu-enable-inner-tiled",
        "--iree-llvmcpu-enable-llvm-ukernels=inner_tiled",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=i8",
        "--acc_type=i32",
    ],
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = [
        "x86_64:avx512vnni:" + ",".join(X86_64_AVX512_VNNI),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
)

# Mixed unsigned-LHS / signed-RHS i8 matmul: exercises x86 AVX-512 VNNI
# `vpdpbusd`, our first asymmetric MMA. The `generic` variant exercises
# the standard codegen fallback (no matching MMA intrinsic).
iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_dt_inner_tiled_ui8_i8_i32",
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-llvmcpu-enable-inner-tiled",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs=ui8",
        "--rhs=i8",
        "--acc_type=i32",
        # easy_large_static (M=N=512): mixed LHS/RHS types lower to a
        # `linalg.generic`, and unit-dim folding rewrites narrow-M/N
        # forms of that into rank-reduced contractions that the
        # inner_tiled lowering doesn't handle.
        "--shapes=easy_large_static",
    ],
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = [
        "generic",
        "x86_64:avx512vnni:" + ",".join(X86_64_AVX512_VNNI),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
)

# LLVMCPU, data-tiling, data-tiling + ukernels + late materialization.
[iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_experimental_dt%s_%s_%s" % (
        ("_uk" if use_uk else ""),
        lhs_rhs_type,
        acc_type,
    ),
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-global-opt-enable-early-materialization=false",
    ] + ["--iree-llvmcpu-enable-ukernels=%s" % ("all" if use_uk else "none")],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
    ],
    tags = ([
        # f16/bf16 trigger internal LLVM assertion errors on riscv and wasm.
        "noriscv",
        "nowasm",
    ] if (lhs_rhs_type == "f16" or lhs_rhs_type == "bf16") else []) + ([
        # The non-ukernel (generic-vectorized)
        # data-tiling path is not yet optimized on RISC-V and times out under
        # the QEMU emulator with the large default shapes. Skip it on RISC-V
        # here and re-add small-shape RISC-V
        # coverage via the dedicated riscv-only rules below.
        "noriscv",
    ] if (not use_uk and acc_type != "f64" and lhs_rhs_type not in [
        "f16",
        "bf16",
    ]) else []),
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = ["generic"] +
                                   ([
                                       "arm_64:dotprod:+dotprod",
                                       "arm_64:i8mm:+i8mm",
                                       "x86_64:avx512vnni:" + ",".join(X86_64_AVX512_VNNI),
                                   ] if lhs_rhs_type == "i8" and acc_type == "i32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                   ] if lhs_rhs_type == "f32" and acc_type == "f32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "arm_64:fullfp16:+fullfp16",
                                   ] if lhs_rhs_type == "f16" and acc_type == "f16" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "arm_64:fp16fml:+fp16fml",
                                   ] if lhs_rhs_type == "f16" and acc_type == "f32" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "x86_64:avx512bf16:" + ",".join(X86_64_AVX512_BF16),
                                       "arm_64:bf16:+bf16",
                                   ] if lhs_rhs_type == "bf16" and acc_type == "bf16" else [
                                       "x86_64:avx2:" + ",".join(X86_64_AVX2),
                                       "x86_64:avx512:" + ",".join(X86_64_AVX512),
                                       "x86_64:avx512bf16:" + ",".join(X86_64_AVX512_BF16),
                                       "arm_64:bf16:+bf16",
                                   ] if lhs_rhs_type == "bf16" and acc_type == "f32" else []),
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for use_uk in [
    False,
    True,
] for (lhs_rhs_type, acc_type) in (
    [
        ("i8", "i32"),
        ("f32", "f32"),
        ("f16", "f16"),
        ("f16", "f32"),
        ("bf16", "bf16"),
        ("bf16", "f32"),
    ]
)]

# RISC-V only: small-shape coverage of the non-ukernel (generic-vectorized)
# data-tiling path. That path is not yet optimized on RISC-V, so the dt /
# experimental_dt rules above carry "noriscv" and skip RISC-V with the large
# default shapes, would time out under the QEMU emulator.
# These rules re-add RISC-V coverage with small shapes only; the "riscv_64:"
# variant restricts them to RISC-V builds, so no tests are created on other
# architectures (which keep their full large-shape coverage above). Keep the
# type list in sync with the "noriscv" condition above; drop these once proper
# RVV codegen for the non-ukernel path is tested.
[iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_%sdt_riscv_small_%s_%s" % (
        variant_prefix,
        lhs_rhs_type,
        acc_type,
    ),
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-llvmcpu-enable-ukernels=none",
    ] + extra_compiler_flags,
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
        "--shapes=small",
    ],
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    # The riscv_64: prefix enables this variant only on RISC-V builds; the
    # actual RVV cpu features are supplied by the toolchain's default test
    # flags, so the features listed here are only a label.
    target_cpu_features_variants = ["riscv_64:rvv:+v,+zvl512b"],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for (variant_prefix, extra_compiler_flags) in [
    ("", []),
    (
        "experimental_",
        ["--iree-global-opt-enable-early-materialization=false"],
    ),
] for (lhs_rhs_type, acc_type) in [
    ("i8", "i32"),
    ("f32", "f32"),
]]

[iree_generated_e2e_runner_test(
    name = "e2e_matmul_vmvx_experimental_dt%s_%s_%s" % (
        ("_uk" if use_uk else ""),
        lhs_rhs_type,
        acc_type,
    ),
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-global-opt-enable-early-materialization=false",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
        "--shapes=small",
    ],
    tags = [],
    target_backends_and_drivers = [
        ("vmvx", "local-task"),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for use_uk in [
    False,
    True,
] for (lhs_rhs_type, acc_type) in (
    [
        ("f32", "f32"),
    ]
)]

[iree_generated_e2e_runner_test(
    name = "e2e_matmul_spirv_experimental_dt_%s_%s" % (
        lhs_rhs_type,
        acc_type,
    ),
    compiler_flags = [
        "--iree-opt-data-tiling",
        "--iree-global-opt-enable-early-materialization=false",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
        "--shapes=small",
    ],
    tags = [],
    target_backends_and_drivers = [
        ("vulkan-spirv", "vulkan"),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for (lhs_rhs_type, acc_type) in (
    [
        ("f32", "f32"),
    ]
)]

# LLVMCPU, MX (MXFP4) matmul with data-tiling. Optimized MX support is not
# implemented for CPU, so encodings are dropped and it effectively runs through
# the non-data-tiling path. This ensures the compiler path works and provides a
# reference for other backends.
iree_generated_e2e_runner_test(
    name = "e2e_matmul_cpu_dt_mxfp4_f32",
    compiler_flags = [
        "--iree-opt-data-tiling",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=f4E2M1FN",
        "--acc_type=f32",
        "--mx_scale_type=f8E8M0FNU",
        "--mx_block_size=32",
        "--shapes=easy_large_static",
        "--transpose_rhs",
    ],
    target_backends_and_drivers = [
        ("llvm-cpu", "local-task"),
    ],
    target_cpu_features_variants = ["generic"],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
)

###########################################################################
##
## VMVX backend
##
###########################################################################

# VMVX, data-tiling + microkernels.
[iree_generated_e2e_runner_test(
    name = "e2e_matmul_vmvx_dt_uk_%s_small" % lhs_rhs_type,
    compiler_flags = [
        "--iree-vmvx-enable-microkernels",
        "--iree-opt-data-tiling",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
        "--shapes=small",
    ],
    target_backends_and_drivers = [
        ("vmvx", "local-task"),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for (lhs_rhs_type, acc_type) in [
    ("i8", "i32"),
    ("f32", "f32"),
]]

###########################################################################
##
## Vulkan backend legacy tests. We don't add GPU tests in Bazel anymore.
##
###########################################################################

[iree_generated_e2e_runner_test(
    name = "e2e_matmul_vulkan_{0}_large_valhall".format(lhs_rhs_type),
    compiler_flags = [
        "--iree-vulkan-target=valhall",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
        "--shapes=easy_large_static",
        "--compilation_info=SPIRVVectorizeMali",
    ],
    tags = [
        # Nvidia GPUs support a superset of Valhall features
        "requires-gpu-nvidia",
        "vulkan_uses_vk_khr_shader_float16_int8",
    ],
    target_backends_and_drivers = [
        ("vulkan-spirv", "vulkan"),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for (lhs_rhs_type, acc_type) in [
    ("i8", "i32"),
    ("f16", "f32"),
    ("f32", "f32"),
]]

[iree_generated_e2e_runner_test(
    name = "e2e_matmul_vulkan_{0}_large_ampere".format(lhs_rhs_type),
    compiler_flags = [
        "--iree-vulkan-target=ampere",
    ],
    generator = ":generate_e2e_matmul_tests",
    generator_args = [
        "--lhs_rhs_type=%s" % lhs_rhs_type,
        "--acc_type=%s" % acc_type,
        "--shapes=easy_large_static",
        "--compilation_info=SPIRVVectorizeNVIDIA",
    ],
    tags = [
        "requires-gpu-sm80",
        "vulkan_uses_vk_khr_shader_float16_int8",
    ],
    target_backends_and_drivers = [
        ("vulkan-spirv", "vulkan"),
    ],
    test_runner = "//tools/testing/e2e:iree-e2e-matmul-test",
    test_type = "matmul",
) for (lhs_rhs_type, acc_type) in [
    ("i8", "i32"),
    ("f16", "f32"),
    ("f32", "f32"),
]]

# GPU tests go to CMakeLists.txt directly because they need to be conditioned on
# target-GPU variables that exist only in CMake.