# 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)