| # 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 |
| |
| import argparse |
| import jaxlib.mlir.ir as mlir_ir |
| import jax._src.interpreters.mlir as mlir |
| import multiprocessing |
| import os |
| import re |
| |
| parser = argparse.ArgumentParser( |
| prog="triage_jaxtest.py", description="Triage the jax tests" |
| ) |
| parser.add_argument("-l", "--logdir", default="/tmp/jaxtest") |
| parser.add_argument("-d", "--delete", default=False) |
| parser.add_argument("-j", "--jobs", default=None) |
| args = parser.parse_args() |
| |
| tests = set(os.listdir(args.logdir)) |
| |
| |
| def filter_to_failures(tests): |
| failures = list() |
| for test in tests: |
| files = os.listdir(f"{args.logdir}/{test}") |
| if "error.txt" in files or "CRASH_MARKER" in files: |
| failures.append(test) |
| failures = sorted(failures) |
| return failures |
| |
| |
| def check_custom_call(errortxt, mlirbc, __): |
| return "stablehlo.custom_call" in errortxt or "stablehlo.custom_call" in mlirbc |
| |
| |
| def check_load_ui(errortxt, _, __): |
| return ( |
| "'flow.tensor.load' op result #0 must be index or signless integer or floating-point or complex-type or vector of any type values, but got 'ui32'" |
| in errortxt |
| ) |
| |
| |
| def check_splat_ui(errortxt, _, __): |
| return ( |
| "'flow.tensor.splat' op failed to verify that value type matches element type of result" |
| in errortxt |
| and "xui" in errortxt |
| ) |
| |
| |
| def check_degenerate_tensor(_, mlirbc, __): |
| return "tensor<0x" in mlirbc or "x0x" in mlirbc |
| |
| |
| def check_topk_bf16(_, mlirbc, __): |
| return "bf16" in mlirbc and "hlo.top_k" in mlirbc |
| |
| |
| def check_cross_replica(errortxt, _, __): |
| return "cross-replica" in errortxt |
| |
| |
| def check_collective(errortxt, _, __): |
| return ( |
| "stablehlo.collective" in errortxt |
| or "UNIMPLEMENTED; collectives not implemented" in errortxt |
| ) |
| |
| |
| def check_sort_shape(errortxt, _, __): |
| return ( |
| "'iree_linalg_ext.sort' op expected operand 1 to have same shape as other operands" |
| in errortxt |
| ) |
| |
| |
| def check_reverse_i1(_, mlirbc, __): |
| for line in mlirbc.split("\n"): |
| if "stablehlo.reverse" in line and "xui" in line: |
| return True |
| return False |
| |
| |
| def check_complex(errortxt, mlirbc, __): |
| return "complex<" in errortxt or "complex<" in mlirbc |
| |
| |
| def check_timeout(errortxt, _, __): |
| return "jaxlib.xla_extension.XlaRuntimeError: ABORTED: ABORTED" in errortxt |
| |
| |
| def check_rng_bit_i8(_, mlirbc, __): |
| lines = mlirbc.split("\n") |
| for line in lines: |
| if "stablehlo.rng_bit_generator" in line and "i8" in line: |
| return True |
| return False |
| |
| |
| def check_min_max_f16(errortxt, _, __): |
| if "undefined symbol: fminf" in errortxt: |
| return True |
| lines = errortxt.split("\n") |
| for line in lines: |
| has_fmax = "llvm.intr.vector.reduce.fmax" in line |
| has_fmin = "llvm.intr.vector.reduce.fmin" in line |
| has_f16 = "f16" in line |
| if (has_fmax or has_fmin) and has_f16: |
| return True |
| return False |
| |
| |
| def check_scatter_ui(errortxt, _, __): |
| lines = errortxt.split("\n") |
| for line in lines: |
| has_scatter = "iree_linalg_ext.scatter" in line |
| has_operand = "expected type of `outs` operand #0" in line |
| has_type = "xui" in line |
| if has_scatter and has_operand and has_type: |
| return True |
| return False |
| |
| |
| def check_bitcast_bf16(errortxt, _, __): |
| return "bf16" in errortxt and "`arith.bitcast` op operand type" in errortxt |
| |
| |
| def check_constant_bf16(errortxt, _, __): |
| return "FloatAttr does not match expected type of the constant" in errortxt |
| |
| |
| def check_triangular_solve(errortxt, _, __): |
| return "stablehlo.triangular_solve" in errortxt |
| |
| |
| def check_cholesky(errortxt, _, __): |
| return "stablehlo.cholesky" in errortxt |
| |
| |
| def check_fft(_, mlirbc, __): |
| return "stablehlo.fft" in mlirbc |
| |
| |
| def check_schedule_allocation(errortxt, _, __): |
| return "Pipeline failed while executing [`ScheduleAllocation`" in errortxt |
| |
| |
| def check_dot_i1(_, mlirbc, __): |
| for line in mlirbc.split("\n"): |
| has_i1 = re.search("tensor<([0-9]*x)*i1>", line) |
| has_dot = re.search("stablehlo.dot", line) |
| if has_i1 and has_dot: |
| return True |
| return False |
| |
| |
| def check_vectorize(errortxt, _, __): |
| return ( |
| "arith.truncf' op operand #0 must be floating-point-like, but got 'vector<f32>" |
| in errortxt |
| ) |
| |
| |
| def check_roundeven(errortxt, _, __): |
| return "roundeven" in errortxt |
| |
| |
| def check_numerical(errortxt, _, __): |
| return "Mismatched elements" in errortxt |
| |
| |
| def check_compilation(errortxt, _, __): |
| return "iree/integrations/pjrt/common/api_impl.cc:1085" in errortxt |
| |
| |
| def check_scatter(errortxt, _, __): |
| return ( |
| "'iree_linalg_ext.scatter' op mismatch in shape of indices and update value at dim#0" |
| in errortxt |
| ) |
| |
| |
| def check_shape_cast(errortxt, _, __): |
| return ( |
| "'vector.shape_cast' op source/result number of elements must match" in errortxt |
| ) |
| |
| |
| def check_scatter_crash(_, mlirbc, runtime_crash): |
| return "stablehlo.scatter" in mlirbc and runtime_crash |
| |
| |
| def check_eigen_decomposition(errortxt, _, __): |
| return ( |
| "Nonsymmetric eigendecomposition is only implemented on the CPU backend" |
| in errortxt |
| ) |
| |
| |
| def check_jax_unimplemented(errortxt, _, __): |
| return "NotImplementedError: MLIR translation rule for primitive" in errortxt |
| |
| |
| def check_serialize_exe(errortxt, _, __): |
| return "UNIMPLEMENTED; PJRT_Executable_Serialize" in errortxt |
| |
| |
| def check_optimized_prgrm(errortxt, _, __): |
| return "UNIMPLEMENTED; PJRT_Executable_OptimizedProgram" in errortxt |
| |
| |
| def check_optimized_program(errortxt, _, __): |
| return "UNIMPLEMENTED; PJRT_Executable_OptimizedProgram" in errortxt |
| |
| |
| def check_donation(errortxt, _, __): |
| return "Donation is not implemented for iree_cpu" in errortxt |
| |
| |
| def check_semaphore_overload(errortxt, _, __): |
| return ( |
| "OUT_OF_RANGE; semaphore values must be monotonically increasing;" in errortxt |
| ) |
| |
| |
| def check_python_callback(errortxt, _, __): |
| return "ValueError: `EmitPythonCallback` not supported" in errortxt |
| |
| |
| def check_complex_convolution(errortxt, mlirbc, __): |
| if "failed to legalize operation 'complex.constant'" in errortxt: |
| return True |
| |
| for line in mlirbc.split("\n"): |
| has_i1 = re.search("tensor<([0-9]*x)*complex<f[0-9]*>>", line) |
| has_conv = re.search("stablehlo.convolution", line) |
| has_dot = re.search("stablehlo.dot", line) |
| if has_i1 and (has_conv or has_dot): |
| return True |
| return False |
| |
| |
| def check_subspan(errortxt, _, __): |
| return ( |
| "failed to legalize operation 'hal.interface.binding.subspan' that was explicitly marked illegal" |
| in errortxt |
| ) |
| |
| |
| def check_from_tensor(errortxt, _, __): |
| return "error: 'tensor.from_elements' op unhandled tensor operation" in errortxt |
| |
| |
| def check_unknown_backend(errortxt, _, __): |
| return "RuntimeError: Unknown backend" in errortxt |
| |
| |
| def check_unsigned_topk(_, mlirbc, __): |
| for line in mlirbc.split("\n"): |
| if "xui" in line and "chlo.top_k" in line: |
| return True |
| return False |
| |
| |
| def check_runtime_crash(__, _, runtime_crash): |
| return runtime_crash |
| |
| |
| def check_aborted(errortxt, _, __): |
| return "ABORTED" in errortxt |
| |
| |
| def check_bounds_indexing(errortxt, _, __): |
| return "out-of-bounds indexing for array of shape" in errortxt |
| |
| |
| def check_nan_correctness(errortxt, _, __): |
| return "nan location mismatch" in errortxt |
| |
| |
| def check_pointer_mismatch(errortxt, _, __): |
| return "unsafe_buffer_pointer()" in errortxt |
| |
| |
| def check_select_and_scatter(errortxt, _, __): |
| return "failed to legalize operation 'stablehlo.select_and_scatter'" in errortxt |
| |
| |
| def check_degenerate_scatter(errortxt, _, __): |
| return ( |
| "'iree_linalg_ext.scatter' op operand #2 must be ranked tensor or memref of any type values" |
| in errortxt |
| ) |
| |
| |
| def check_cost_analysis(errortxt, _, __): |
| return "cost_analysis()" in errortxt |
| |
| |
| def check_invalid_option(errortxt, _, __): |
| return "No such compile option: 'invalid_key'" in errortxt |
| |
| |
| def check_inf_mismatch(errortxt, _, __): |
| return "inf location mismatch" in errortxt |
| |
| |
| def check_shape_assertion(errortxt, _, __): |
| for line in errortxt.split("\n"): |
| if "assertEqual" in line and ".shape" in line: |
| return True |
| return False |
| |
| |
| def check_vector_contract(errortxt, _, __): |
| return ( |
| "'vector.contract' op failed to verify that lhs and rhs have same element type" |
| in errortxt |
| ) |
| |
| |
| def check_subbyte_read(errortxt, _, __): |
| return "opaque and sub-byte aligned element types cannot be indexed" in errortxt |
| |
| |
| def check_buffer_usage(errortxt, _, __): |
| return ( |
| "requested buffer usage is not supported" in errortxt |
| or "tensor requested usage was not specified when the buffer" in errortxt |
| or "PERMISSION_DENIED; requested usage was not specified when the buffer was allocated; buffer allows DISPATCH_INDIRECT_PARAMS" |
| in errortxt |
| ) |
| |
| |
| def check_subbyte_singleton(errortxt, _, __): |
| return "does not have integral number of total bytes" in errortxt |
| |
| |
| def check_max_arg(errortxt, _, __): |
| return "max() arg is an empty sequence" in errortxt |
| |
| |
| def check_double_support(errortxt, _, __): |
| return "expected f32 (21000020) but have f64 (21000040)" in errortxt |
| |
| |
| def check_stablehlo_degenerate(_, mlirbc, __): |
| for line in mlirbc.split("\n"): |
| if "stablehlo" in line and ("x0x" in line or "<0x" in line): |
| return True |
| return False |
| |
| |
| def check_stablehlo_allreduce(errortxt, _, __): |
| return "failed to legalize operation 'stablehlo.all_reduce'" in errortxt |
| |
| |
| def check_dot_shape(errortxt, _, __): |
| for line in errortxt.split("\n"): |
| if ( |
| "error: inferred shape" in line |
| and "is incompatible with return type of operation " in line |
| ): |
| return True |
| return False |
| |
| |
| KnownChecks = { |
| "https://github.com/iree-org/iree/issues/14255 (detensoring)": check_from_tensor, |
| "https://github.com/iree-org/iree/issues/????? (unknown)": check_jax_unimplemented, |
| "https://github.com/iree-org/iree/issues/13726 (collective)": check_collective, |
| "https://github.com/iree-org/iree/issues/12410 (custom call)": check_custom_call, |
| "https://github.com/iree-org/iree/issues/11018 (triangle)": check_triangular_solve, |
| "https://github.com/iree-org/iree/issues/12263 (fft)": check_fft, |
| "https://github.com/iree-org/iree/issues/14072 (complex convolution)": check_complex_convolution, |
| "https://github.com/iree-org/iree/issues/10816 (cholesky)": check_cholesky, |
| "https://github.com/iree-org/iree/issues/11761 (rng bit gen i8)": check_rng_bit_i8, |
| "https://github.com/iree-org/iree/issues/????? (eigen decomp)": check_eigen_decomposition, |
| "https://github.com/iree-org/iree/issues/13579 (scatter ui)": check_scatter_ui, |
| "https://github.com/iree-org/iree/issues/13725 (cross repl)": check_cross_replica, |
| "https://github.com/iree-org/iree/issues/13493 (dot i1)": check_dot_i1, |
| "https://github.com/iree-org/iree/issues/13522 (roundeven)": check_roundeven, |
| "https://github.com/iree-org/iree/issues/13577 (max/min f16)": check_min_max_f16, |
| "https://github.com/iree-org/iree/issues/13523 (scatter)": check_scatter, |
| "https://github.com/iree-org/iree/issues/13580 (scatter crash)": check_scatter_crash, |
| "https://github.com/iree-org/iree/issues/14079 (shape_cast)": check_shape_cast, |
| "https://github.com/iree-org/iree/issues/????? (optimized prgrm)": check_optimized_program, |
| "https://github.com/iree-org/iree/issues/????? (donation)": check_donation, |
| "https://github.com/iree-org/iree/issues/????? (python callback)": check_python_callback, |
| "https://github.com/iree-org/iree/issues/????? (subspan)": check_subspan, |
| "https://github.com/iree-org/iree/issues/14098 (unsigned topk)": check_unsigned_topk, |
| "https://github.com/iree-org/iree/issues/????? (bounds indexing)": check_bounds_indexing, |
| "https://github.com/iree-org/iree/issues/????? (nan correctness)": check_nan_correctness, |
| "https://github.com/iree-org/iree/issues/????? (pointer mismatch)": check_pointer_mismatch, |
| "https://github.com/iree-org/iree/issues/10841 (select and scatter)": check_select_and_scatter, |
| "https://github.com/iree-org/iree/issues/????? (degenerate scatter)": check_degenerate_scatter, |
| "https://github.com/iree-org/iree/issues/????? (cost analysis)": check_cost_analysis, |
| "https://github.com/iree-org/iree/issues/????? (invalid option)": check_invalid_option, |
| "https://github.com/iree-org/iree/issues/????? (inf mismatch)": check_inf_mismatch, |
| "https://github.com/iree-org/iree/issues/????? (shape assertion)": check_shape_assertion, |
| "https://github.com/iree-org/iree/issues/????? (vector contract)": check_vector_contract, |
| "https://github.com/iree-org/iree/issues/????? (subbyte indexed)": check_subbyte_read, |
| "https://github.com/iree-org/iree/issues/????? (buffer usage)": check_buffer_usage, |
| "https://github.com/iree-org/iree/issues/????? (subbyte singleton)": check_subbyte_singleton, |
| "https://github.com/iree-org/iree/issues/????? (max arg)": check_max_arg, |
| "https://github.com/iree-org/iree/issues/????? (double support)": check_double_support, |
| "https://github.com/iree-org/iree/issues/????? (zero extent)": check_stablehlo_degenerate, |
| "https://github.com/iree-org/iree/issues/????? (all reduce)": check_stablehlo_allreduce, |
| "https://github.com/iree-org/iree/issues/????? (stablehlo dot_general)": check_dot_shape, |
| "(unknown backend)": check_unknown_backend, |
| "(semaphore)": check_semaphore_overload, |
| "Aborted (possible timeout)": check_aborted, |
| "Runtime Crash": check_runtime_crash, |
| "Compilation Failure": check_compilation, |
| "Numerical Failures": check_numerical, |
| "Untriaged": lambda _, __, ___: True, |
| } |
| |
| |
| def triage_test(test): |
| files = sorted(os.listdir(f"{args.logdir}/{test}")) |
| # Load the error.txt if it is available. |
| error = "" |
| if "error.txt" in files: |
| with open(f"{args.logdir}/{test}/error.txt") as f: |
| error = "".join(f.readlines()) |
| |
| # Load the last bytecode file that was attempted to be compiled: |
| mlirbc_count = len([f for f in files if "mlirbc" in f]) |
| mlirbc_name = f"{mlirbc_count - 1}-program.mlirbc" |
| vmfb_name = f"{mlirbc_count - 1}-program.vmfb" |
| |
| runtime_crash = "CRASH_MARKER" in files |
| |
| mlirbc = "" |
| if mlirbc_count > 0: |
| with mlir.make_ir_context() as ctx: |
| with open(f"{args.logdir}/{test}/{mlirbc_name}", "rb") as f: |
| mlirbc = f.read() |
| mlirbc = str(mlir_ir.Module.parse(mlirbc)) |
| |
| for checkname in KnownChecks: |
| if KnownChecks[checkname](error, mlirbc, runtime_crash): |
| return checkname |
| |
| return "unknown error" |
| |
| |
| def filter_error_mapping(tests): |
| error_mapping = {} |
| with multiprocessing.Pool(int(args.jobs) if args.jobs else args.jobs) as p: |
| results = p.map(triage_test, tests) |
| |
| for test, result in zip(tests, results): |
| error_mapping[test] = result |
| return error_mapping |
| |
| |
| def generate_summary(mapping): |
| summary = {} |
| for err in KnownChecks.keys(): |
| summary[err] = [] |
| for test in mapping: |
| summary[mapping[test]].append(test) |
| return summary |
| |
| |
| def print_summary(summary): |
| maxlen = 0 |
| for error in summary: |
| maxlen = max(len(error), maxlen) |
| for error in summary: |
| print(f"{error:<{maxlen}} : {len(summary[error])}") |
| |
| passstr = "Passing" |
| failstr = "Failing" |
| print(f"{passstr:<{maxlen}} : {len(tests) - len(failing)}") |
| print(f"{failstr:<{maxlen}} : {len(failing)}") |
| |
| |
| failing = filter_to_failures(tests) |
| mapping = filter_error_mapping(failing) |
| summary = generate_summary(mapping) |
| print_summary(summary) |
