E2E tests for math ops. (#20169)

This is thought to be needed before going ahead with a batch of math ops
codegen changes
(#19970 (review))
and this also discovered a few bugs:
#20163
#20164
#20165.

Signed-off-by: Benoit Jacob <[email protected]>

Author: bjacob

CMakeLists.txt

@@ -568,6 +568,7 @@ include(iree_cc_binary_benchmark)
 include(iree_hal_cts_test_suite)
 include(iree_static_linker_test)
 include(iree_plugin_register)
+include(iree_genrule)
 
 # Default any sub-tree which doesn't provide its own package namespacing
 # to derive it relative to this directory and prefixed with iree/.

build_tools/bazel/iree_check_test.bzl

@@ -25,6 +25,7 @@ def iree_check_test(
         runner_args = [],
         tags = [],
         timeout = None,
+        deps = [],
         **kwargs):
     """Creates an iree-check-module test for the specified source file.
 
@@ -53,11 +54,13 @@ def iree_check_test(
         "--iree-hal-target-backends=%s" % target_backend,
     ] + compiler_flags + input_type_flags
     bytecode_module_name = name + "_bytecode_module"
+
     iree_bytecode_module(
         name = bytecode_module_name,
         src = src,
         flags = flags,
         tags = ["target=%s" % target_backend],
+        deps = deps,
         visibility = ["//visibility:private"],
     )
 
@@ -86,6 +89,7 @@ def iree_check_single_backend_test_suite(
         input_type = None,
         runner_args = [],
         tags = [],
+        deps = [],
         timeout = None,
         **kwargs):
     """Creates a test suite of iree-check-module tests for a single backend/driver pair.
@@ -124,7 +128,7 @@ def iree_check_single_backend_test_suite(
 
     tests = []
     for src in srcs:
-        test_name = "_".join([name, src])
+        test_name = "_".join([name, src]).replace("/", "_").replace(":", "_")
         iree_check_test(
             name = test_name,
             src = src,
@@ -135,6 +139,7 @@ def iree_check_single_backend_test_suite(
             runner_args = runner_args,
             tags = tags,
             timeout = timeout,
+            deps = deps,
             **kwargs
         )
         tests.append(test_name)
@@ -161,6 +166,7 @@ def iree_check_test_suite(
         runner_args = [],
         tags = [],
         target_cpu_features_variants = [],
+        deps = [],
         **kwargs):
     """Creates a test suite of iree-check-module tests.
 
@@ -205,6 +211,7 @@ def iree_check_test_suite(
             input_type = input_type,
             runner_args = runner_args,
             tags = tags,
+            deps = deps,
             **kwargs
         )
         tests.append(suite_name)

build_tools/bazel/iree_genrule.bzl

@@ -0,0 +1,40 @@
+# Copyright 2025 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
+
+"""A minimal subset of Bazel genrule for common use cases"""
+
+def iree_genrule(
+        name,
+        srcs,
+        outs,
+        cmd,
+        **kwargs):
+    """A minimal subset of Bazel genrule for common use cases.
+
+    Args:
+        name: Name of the target.
+        srcs: Source files, including any script run in the command.
+              Unlike Bazel's genrule, we do not try to distinguish between the
+              two. The distinction is needed when tools need to be compiled for
+              host, but that doesn't concern us if we only need to run python
+              scripts.
+        outs: Files generated by the command.
+        cmd: The command to be executed. The only supported special Bazel
+             genrule syntax is:
+               * "$(rootpath x)", which expands to the path to a file in the
+                 source tree.
+               * "$(execpath x)", which expands to the path to a file in the
+                 directory where Bazel runs the build action.
+        **kwargs: any additional attributes to pass to the underlying rules.
+    """
+
+    native.genrule(
+        name = name,
+        srcs = srcs,
+        outs = outs,
+        cmd = cmd,
+        **kwargs
+    )

build_tools/bazel_to_cmake/bazel_to_cmake_converter.py

@@ -789,6 +789,7 @@ def iree_check_single_backend_test_suite(
         runner_args=None,
         tags=None,
         timeout=None,
+        deps=None,
         **kwargs,
     ):
         if self._should_skip_target(tags=tags, **kwargs):
@@ -806,6 +807,7 @@ def iree_check_single_backend_test_suite(
         runner_args_block = self._convert_string_list_block("RUNNER_ARGS", runner_args)
         labels_block = self._convert_string_list_block("LABELS", tags)
         timeout_block = self._convert_timeout_arg_block("TIMEOUT", timeout)
+        deps_block = self._convert_string_list_block("DEPS", deps)
 
         self._converter.body += (
             f"iree_check_single_backend_test_suite(\n"
@@ -818,6 +820,7 @@ def iree_check_single_backend_test_suite(
             f"{runner_args_block}"
             f"{labels_block}"
             f"{timeout_block}"
+            f"{deps_block}"
             f")\n\n"
         )
 
@@ -832,6 +835,7 @@ def iree_check_test_suite(
         tags=None,
         target_cpu_features_variants=None,
         timeout=None,
+        deps=None,
         **kwargs,
     ):
         if self._should_skip_target(tags=tags, **kwargs):
@@ -858,6 +862,7 @@ def iree_check_test_suite(
             "TARGET_CPU_FEATURES_VARIANTS", target_cpu_features_variants
         )
         timeout_block = self._convert_timeout_arg_block("TIMEOUT", timeout)
+        deps_block = self._convert_string_list_block("DEPS", deps)
 
         self._converter.body += (
             f"iree_check_test_suite(\n"
@@ -871,6 +876,7 @@ def iree_check_test_suite(
             f"{labels_block}"
             f"{target_cpu_features_variants_block}"
             f"{timeout_block}"
+            f"{deps_block}"
             f")\n\n"
         )
 
@@ -1007,6 +1013,21 @@ def iree_cmake_extra_content(self, content, inline=False):
         else:
             self._converter.header += f"\n{content}\n"
 
+    def iree_genrule(self, name, srcs, outs, cmd):
+        name_block = self._convert_string_arg_block("NAME", name, quote=False)
+        srcs_block = self._convert_srcs_block(srcs)
+        outs_block = self._convert_target_list_block("OUTS", outs)
+        cmd_block = self._convert_string_arg_block("CMD", cmd, quote=True)
+
+        self._converter.body += (
+            f"iree_genrule(\n"
+            f"{name_block}"
+            f"{srcs_block}"
+            f"{outs_block}"
+            f"{cmd_block}"
+            f")\n\n"
+        )
+
 
 class Converter(object):
     """Conversion state tracking and full file template substitution."""

build_tools/cmake/iree_genrule.cmake

@@ -0,0 +1,74 @@
+# Copyright 2025 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
+
+include(CMakeParseArguments)
+
+# A wrapper around add_custom_command and a minimal subset of Bazel genrule.
+#
+# Parameters:
+# NAME: Name of the target.
+# SRCS: Source files, including any script run in the command.
+#       Unlike Bazel's genrule, we do not try to distinguish between the
+#       two. The distinction is needed when tools need to be compiled for
+#       host, but that doesn't concern us if we only need to run python
+#       scripts.
+# OUTS: Files generated by the command.
+# CMD:  The command to be executed. The only supported special Bazel genrule
+#       syntax is:
+#            * "$(rootpath x)", only supported for source files. In conversion
+#              to CMake, this expands to the path of x relatively to the current
+#              source dir.
+#            * "$(execpath x)", only supported for generated files. In
+#              conversion to CMake, this expands to just x, as the binary dir is
+#              the default working dir for custom commands anyway.
+function(iree_genrule)
+  cmake_parse_arguments(
+    _RULE
+    ""
+    "NAME"
+    "SRCS;OUTS;CMD"
+    ${ARGN}
+  )
+
+  set(_CMD "${_RULE_CMD}")
+
+  # Replace Bazel syntax $(rootpath x) with the path into the source dir.
+  string(REGEX REPLACE "\\$\\(rootpath ([^)]+)\\)" "${CMAKE_CURRENT_SOURCE_DIR}/\\1" _CMD "${_CMD}")
+
+  # Simply drop Bazel syntax $(execpath x) as Bazel custom commands are executed
+  # by default in the build directory.
+  string(REGEX REPLACE "\\$\\(execpath ([^)]+)\\)" "\\1" _CMD "${_CMD}")
+
+  # Convert CMake/Unix-style paths with forward slashes to Windows-style with
+  # backslashes. It is a bit incorrect to do it as a single cmake_path command
+  # on the whole command string, which isn't technically a path, but this should
+  # not matter if all what this does is this character substitution.
+  # It is not worth implementing a cumbersome fix here, when CMake 4.0 brings
+  # the $<PATH:NATIVE_PATH,...> generator expression which is a simpler, better
+  # fix here. TODO(bjacob): use that generator expression in the above string
+  # replace command directly, whenever we can rely on CMake 4.0.
+  cmake_path(NATIVE_PATH _CMD _CMD)
+
+  # CMake add_custom_command expects a list as the command, so we replace spaces
+  # by semicolon here. Careful to avoid replacing backslash-escaped spaces.
+  string(REGEX REPLACE "([^\\]) " "\\1;" _CMD "${_CMD}")
+
+  add_custom_command(
+    OUTPUT
+      "${_RULE_OUTS}"
+    COMMAND
+      ${_CMD}
+    DEPENDS
+      "${_RULE_SRCS}"
+    COMMENT
+      "Generating ${_RULE_OUTS}"
+    VERBATIM
+  )
+
+  add_custom_target("${_RULE_NAME}"
+    DEPENDS "${_RULE_OUTS}"
+  )
+endfunction()

compiler/src/iree/compiler/Modules/Check/IR/CheckOps.td

@@ -138,11 +138,14 @@ def CHECK_ExpectAlmostEqOp :
   let summary = [{Checks that the operands are almost equal}];
   let description = [{
     Verifies that the buffer view or tensor operands with float elements satisfy
-    the Numpy-style fuzzy-comparision condition with pararameters `atol`,
-    `rtol`, which is the following element-wise on array elements `lhs`, `rhs`:
-    ```
-    abs(lhs - rhs) <= atol + rtol * abs(rhs).
-    ```
+    the Numpy-style fuzzy-comparision condition with parameters `atol`,
+    `rtol`, defined exactly as in NumPy isclose():
+    https://github.com/numpy/numpy/blob/7297f3117d84745bfade1e2f9aec3531e5917500/numpy/_core/numeric.py#L2447-L2449
+
+    The condition being verified on each lhs and rhs value is:
+      lhs == rhs || (isfinite(rhs) && abs(lhs - rhs) <= atol + rtol * abs(rhs)).
+    Note that the `lhs == rhs` part is needed for the case (lhs=+inf, rhs+inf)
+    to return true. Indeed, in that case, lhs-rhs is NaN.
 
     Issues a non-fatal failure if the verification fails.
 

runtime/src/iree/modules/check/module.cc

@@ -69,13 +69,17 @@ bool EqByteSpan(iree_byte_span_t lhs_bytes, iree_byte_span_t rhs_bytes) {
 // Numpy-compatible fuzzy comparison of floating-point values lhs, rhs with
 // respect to tolerance parameters atol, rtol.
 //
-// The meaning of the tolerance parameters atol and rtol is exactly as in:
-//   https://numpy.org/doc/stable/reference/generated/numpy.testing.assert_allclose.html
+// The meaning of the tolerance parameters atol and rtol is exactly as in NumPy
+// isclose():
+// https://github.com/numpy/numpy/blob/7297f3117d84745bfade1e2f9aec3531e5917500/numpy/_core/numeric.py#L2447-L2449
 // The condition being verified on each lhs and rhs value is:
-//   abs(lhs - rhs) <= atol + rtol * abs(rhs).
+//   lhs == rhs || (isfinite(rhs) && abs(lhs - rhs) <= atol + rtol * abs(rhs)).
+// Note that the `lhs == rhs` part is needed for the case (lhs=+inf, rhs+inf)
+// to return true. Indeed, in that case, lhs-rhs is NaN.
 template <typename T>
 bool NumpyFuzzyCompare(T lhs, T rhs, float atol, float rtol) {
-  return std::abs(lhs - rhs) <= atol + rtol * std::abs(rhs);
+  return lhs == rhs || (std::isfinite(rhs) &&
+                        std::abs(lhs - rhs) <= atol + rtol * std::abs(rhs));
 }
 
 // Records information about some LHS/RHS scalars that failed a fuzzy comparison

tests/e2e/math/BUILD.bazel

@@ -0,0 +1,55 @@
+# Copyright 2025 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
+
+load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite")
+load("//build_tools/bazel:iree_genrule.bzl", "iree_genrule")
+
+package(
+    features = ["layering_check"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+testcases = [
+    (
+        # Input JSON file describing testcases
+        "math_ops_%s.json" % backend,
+        # Output generated MLIR test file.
+        "math_ops_%s.mlir" % backend,
+    )
+    for backend in [
+        "llvm-cpu",
+        "rocm",
+    ]
+]
+
+[iree_genrule(
+    name = "gen_%s" % generated_src,
+    srcs = [
+        "generate.py",
+        testcases_json,
+    ],
+    outs = [generated_src],
+    cmd = " ".join([
+        "python3",
+        "$(rootpath generate.py)",
+        "--testcases=$(rootpath %s)" % testcases_json,
+        "> $(execpath %s)" % generated_src,
+    ]),
+) for testcases_json, generated_src in testcases]
+
+[iree_check_single_backend_test_suite(
+    name = "math_ops_%s" % backend,
+    srcs = ["//tests/e2e/math:math_ops_%s.mlir" % backend],
+    compiler_flags = ["--iree-llvmcpu-target-cpu=generic"] if backend == "llvm-cpu" else [],
+    driver = driver,
+    target_backend = backend,
+    deps = [
+        "gen_math_ops_%s.mlir" % backend,
+    ],
+) for backend, driver in [
+    ("llvm-cpu", "local-task"),
+    ("rocm", "hip"),
+]]