main.py

from __future__ import annotations
import argparse
import claripy
import json
import os
import sys
import itertools
from typing import Dict, List, Tuple, NamedTuple
from idem_checker.state import ProgramState, StateTree, MemoryOperation
from idem_checker.executor import SimpleExecutor
from idem_checker.exception import *
from idem_checker.parser.parser import SASSParser
from idem_checker.program import SASSFunction
from idem_checker.preprocess import *
from idem_checker.utils import *
from idem_checker.dependency_analysis import *
from idem_checker.verifier import InstanceValidator
from idem_checker.monotonicity import MonoChecker, get_memop_iterator_bounds
from idem_checker.codegen import MemOpCodeGen, MemOpIterator, gen_global_cond_code, combine_all_range_funcs
from idem_checker.expr import expr_contains_sr, find_matching_memop, merge_ranges, expr_similar, expr_equal, SymbolicRange
from idem_checker.context import Context


class TraceInstance(NamedTuple):
    grid_dims: List[int]
    block_dims: List[int]
    launch_args: bytearray


class MemOp(NamedTuple):
    addr: int
    memop: MemoryOperation
    state_tree: StateTree


def load_trace_instances(trace_path: str, kernel_name: str = None) -> List[TraceInstance]:
    f = open(trace_path, "r")
    trace_json = json.loads(f.read())
    f.close()
    instances = []
    for item in trace_json:
        if kernel_name is not None and item.get("name") != kernel_name:
            continue
        grid_dims = item["grid_dims"]
        block_dims = item["block_dims"]
        launch_args = bytearray(item["params"]) if item["params"] is not None else bytearray()
        instances.append(TraceInstance(grid_dims, block_dims, launch_args))
    return instances


def transform_new_addr_to_old_addr(old_addrs, new_addr):
    new_addr_offset = addr_to_index(new_addr)
    return old_addrs[new_addr_offset]


def step_exec(func: SASSFunction, args) -> Dict[int, List[MemOp]]:
    ldsts = func.get_ldsts()
    memops = []
    for idx, ldst in enumerate(ldsts):
        print(f"[EXEC] ldst: {idx} / {len(ldsts)}, {ldst}")
        dda_result = func.perform_data_dependency_analysis_on_instr(ldst_instr=ldst, print_trace=False)
        new_func, old_addrs = func.get_simplified_function(dda_result)
        if not args.no_match_loop:
            new_func.match_loops(force=True)
        executor = SimpleExecutor(new_func.addr_instrs)
        state_tree = executor.execute()
        state_memops = state_tree.get_deduplicated_memops(args.split_memop)
        for memop in state_memops:
            memops.append(MemOp(old_addrs[addr_to_index(memop.instr_offset)], memop, state_tree))
    memops = sorted(memops, key=lambda x: x.addr)
    memops_groups: Dict[int, List[MemOp]] = {addr: list(memop) for addr, memop in itertools.groupby(memops, key=lambda x: x.addr)}
    for key in memops_groups.keys():
        memop_list = memops_groups[key]
        deduplicated_list = []
        memop_hashes = set()
        for memop in memop_list:
            min_expr = memop.state_tree.symbol_store.replace_loop_symbol_min_expr(memop.memop.base_mem_addr)
            max_expr = memop.state_tree.symbol_store.replace_loop_symbol_max_expr(memop.memop.base_mem_addr)
            cond_hashes = tuple(sorted([hash(cond) for cond in memop.memop.extract_validatable_conditions()]))
            memop_hash = (hash(min_expr), hash(max_expr), cond_hashes)
            if memop_hash not in memop_hashes:
                memop_hashes.add(memop_hash)
                deduplicated_list.append(memop)
        memops_groups[key] = deduplicated_list
    return memops_groups


def step_enumerate(memop_groups: Dict[int, List[MemOp]]) -> Dict[int, List[MemOpCodeGen]]:
    memop_funcs = {}
    for op_idx, key in enumerate(sorted(memop_groups.keys())):
        memops = memop_groups[key]
        candidate_funcs = []
        for candidate_idx, memop in enumerate(memops):
            memaddr = memop.memop.base_mem_addr
            start_expr = memaddr
            end_expr = memaddr
            memop_iters = []
            sr_vars = [i for i in get_expr_var_set(memaddr) if i.args[0].startswith("SR_")]
            for var in sr_vars:
                memop_iter = MemOpIterator(var, 0, memop.state_tree.symbol_store.replace_symbol_by_max_expr(var))
                memop_iters.append(memop_iter)
            validatable_conditions = memop.memop.extract_validatable_conditions()
            codegen = MemOpCodeGen(op_idx, candidate_idx, memop.memop.op, start_expr, end_expr,
                        validatable_conditions, memop_iters)
            candidate_funcs.append(codegen)
        memop_funcs[key] = candidate_funcs
    return memop_funcs 


def step_mono(memop_groups: Dict[int, List[MemOp]], 
              param_constraints: List[claripy.Base],
              early_exit=False, mono_mode="check") -> Tuple[Dict[int, List[MemOpCodeGen]], bool]:
    memop_funcs = {}
    all_mono = True
    for op_idx, key in enumerate(sorted(memop_groups.keys())):
        memops = memop_groups[key]
        candidate_funcs = []
        for candidate_idx, memop in enumerate(memops):
            print(f"[Mono] checking {op_idx} - {candidate_idx}")
            is_validatable = memop.memop.is_validatable()
            if not is_validatable:
                start_expr = claripy.BVV(0, memop.memop.base_mem_addr.length)
                end_expr = claripy.BVV(-1, memop.memop.base_mem_addr.length)
                codegen = MemOpCodeGen(op_idx, candidate_idx, memop.memop.op, start_expr, end_expr, 
                                       memop.memop.extract_validatable_conditions(), [])
                candidate_funcs.append(codegen)
                continue

            memaddr, var_bounds = get_memop_iterator_bounds(memop.state_tree, memop.memop)
            range_conditions = []
            mono_check_constraints = param_constraints + memop.memop.conditions
            vars = []
            for var, lower, upper in var_bounds:
                vars.append(var)
                mono_check_constraints.append(lower <= var)
                mono_check_constraints.append(var <= upper)
                mono_check_constraints.append(var < 0xffffffff)
                range_condition = lower <= upper
                if not awalys_true_for_dims(range_condition):
                    range_conditions.append(range_condition)
            
            checker = MonoChecker(mono_check_constraints)
            start_expr = memaddr
            end_expr = memaddr
            
            non_mono_iters = []
            if mono_mode == "transform":
                transform_result = checker.transform_to_mono_expr(memaddr, vars)
                start_expr, end_expr = transform_result.replace_min_max()
                for var, lower, upper in var_bounds:
                    start_expr = start_expr.replace(var, lower)
                    end_expr = end_expr.replace(var, upper)
                if not transform_result.all_monotonic():
                    all_mono = False
            else:
                for var, lower, upper in var_bounds:
                    is_monotonic = True
                    negated = False
                    if mono_mode == "check":
                        is_monotonic = checker.check(memaddr, var)
                    elif mono_mode == "negate":
                        is_monotonic = checker.check(memaddr, var) 
                        if not is_monotonic:
                            is_monotonic = checker.check(memaddr, var, negated=True)
                            negated = True
                    elif mono_mode == "all":
                        is_monotonic = True
                    elif mono_mode == "none":
                        is_monotonic = False
                    
                    if is_monotonic:
                        if negated:
                            start_expr = start_expr.replace(var, upper)
                            end_expr = end_expr.replace(var, lower)
                        else:
                            start_expr = start_expr.replace(var, lower)
                            end_expr = end_expr.replace(var, upper)
                    else:
                        all_mono = False
                        if early_exit:
                            return {}, all_mono
                        non_mono_iters.append(MemOpIterator(var, lower, upper))
            codegen = MemOpCodeGen(op_idx, candidate_idx, memop.memop.op, start_expr, end_expr,
                                   (memop.memop.extract_validatable_conditions()+range_conditions), non_mono_iters) 
            candidate_funcs.append(codegen)
        memop_funcs[key] = candidate_funcs
    return memop_funcs, all_mono


def step_compact(memop_groups: Dict[int, List[MemOpCodeGen]], compact=False):
    all_memopcgs: List[MemOpCodeGen] = []
    ordered_keys = sorted(memop_groups.keys())
    for key in ordered_keys:
        for memopcg in memop_groups[key]:
            all_memopcgs.append(memopcg)
    
    if compact:
        print("[Compact] compacting memops...")
        read_memopcgs = []
        write_memopcgs = []
        read_symbolic_ranges = []
        write_symbolic_ranges = []

        for idx, memopcg in enumerate(all_memopcgs):
                if memopcg.op == "READ":
                    memopcgs = read_memopcgs
                    symbolic_ranges = read_symbolic_ranges
                else:
                    memopcgs = write_memopcgs
                    symbolic_ranges = write_symbolic_ranges
                
                current_symbolic_range = SymbolicRange(memopcg.start_expr, memopcg.end_expr, memopcg.conds)
                match = find_matching_memop(symbolic_ranges, current_symbolic_range)
                if match is not None:
                    i, result = match
                    symbolic_ranges[i] = result
                    memopcgs[i] = MemOpCodeGen(memopcgs[i].op_idx, memopcgs[i].candidate_idx, memopcgs[i].op, result.start, result.end, result.conds, [])
                    print(f"[Compact] {memopcg.op_idx}-{memopcg.candidate_idx} to {memopcgs[i].op_idx}-{memopcgs[i].candidate_idx}")
                else:
                    memopcgs.append(memopcg)
                    symbolic_ranges.append(current_symbolic_range)
        print(f"[Compact] compacting memops done {len(all_memopcgs)} to {len(read_memopcgs) + len(write_memopcgs)}")
        all_memopcgs = read_memopcgs + write_memopcgs
    return all_memopcgs


def step_codegen(all_memopcgs: List[MemOpCodeGen], global_constraints=[], func_name="") -> str:
    memop_funcs = []
    memop_func_names = []
    for memopcg in all_memopcgs:
        memop_funcs.append(memopcg.gen_funcs())
        memop_func_names.append(memopcg.fun_name_prefix())
    def constraint_validable(expr):
        return all([var.startswith("CMEM[0]") or \
                                var.startswith("GDIM") or \
                                var.startswith("BDIM")
                                for var in expr.variables])
        
    validatable_constraints = [c for c in global_constraints if constraint_validable(c)]
    global_cond = claripy.And(*validatable_constraints)
    global_cond_code = gen_global_cond_code(global_cond)
    return combine_all_range_funcs(memop_funcs, memop_func_names, func_name=func_name, global_cond_code=global_cond_code)


def step_kernel_level(all_memops: List[MemOp]):
    all_read = all(memop.memop.op == "READ" for memop in all_memops)
    all_write = all(memop.memop.op == "WRITE" for memop in all_memops)
    if all_read or all_write:
        return "Idempotent", "None"
    def get_signature(memop: MemOp):
        cond_signature = tuple([cond.cache_key for cond in memop.memop.extract_validatable_conditions()])
        return (memop.memop.base_mem_addr.cache_key, cond_signature)
    read_sets = set([get_signature(memop) for memop in all_memops if memop.memop.op == "READ"])
    write_sets = set([get_signature(memop) for memop in all_memops if memop.memop.op == "WRITE"])
    
    if not read_sets.isdisjoint(write_sets):
        return "Non-Idempotent", "OVERLAP"

    return "Conditionally-idempotent", "None"


def get_all_global_constraints(memop_groups: Dict[int, List[MemOp]]):
    state_trees = []
    state_tree_hashes = set()
    for memop_list in memop_groups.values():
        for memop in memop_list:
            state_tree = memop.state_tree
            state_tree_hash = hash(state_tree)
            if state_tree_hash not in state_tree_hashes:
                state_tree_hashes.add(state_tree_hash)
                state_trees.append(state_tree)
    all_global_constrains = []
    for state_tree in state_trees:
        all_global_constrains += state_tree.get_global_conditions()
    return all_global_constrains


def setup_context(args):
    Context.instance().print_exec_state = True
    Context.instance().enable_loop_unrolling = args.enable_unroll > 0
    Context.instance().loop_unrolling_factor = args.unroll_factor
    Context.instance().loop_split = args.loop_split > 0
    return


def analyze(func: SASSFunction, instances: List[TraceInstance], args) -> Tuple[str, str]:
    constraints = []
    if instances:
        constraints = func.get_constraints_from_instances(
            [i.launch_args for i in instances],
            [i.grid_dims for i in instances],
            [i.block_dims for i in instances],
        )
    else:
        constraints = func.get_constraints_from_instances([], [], [])
    print(f"Infered constraints from trace: ")
    [print(i) for i in constraints]

    memop_groups = step_exec(func, args)
    
    kernel_level, k_reason = step_kernel_level(list(itertools.chain.from_iterable(list(memop_groups.values()))))
    print(f"Kernel level: {kernel_level}, reason: {k_reason}")
    
    if args.skip_mono:
        memop_cgs = step_enumerate(memop_groups)
        all_mono = True
    else:
        memop_cgs, all_mono = step_mono(memop_groups, constraints, mono_mode=args.mono_mode)
    
    print(f"All monotonic: {all_mono}")
    
    memop_cgs = step_compact(memop_cgs, args.compact_memop > 0)
    print(f"Total memops after compact: {len(memop_cgs)}")
    
    code = step_codegen(memop_cgs, global_constraints=get_all_global_constraints(memop_groups), func_name=func.name)
    
    return code, kernel_level, k_reason


def main():
    sys.setrecursionlimit(10000)
    
    parser = argparse.ArgumentParser(description="GPU kernel idempotence checker code generator")
    parser.add_argument("--sass", "-s", type=str, required=True, help="SASS assembly file path")
    parser.add_argument("--output", "-o", type=str, required=True, help="Output CPP file path")
    parser.add_argument("--trace", "-t", type=str, default=None, help="Trace JSON file path (optional, for constraint inference)")
    parser.add_argument("--kernel", "-k", type=str, default=None, help="Kernel name to analyze (used with --trace)")
    parser.add_argument("--skip-mono", action="store_true", help="Skip monotonicity analysis")
    parser.add_argument("--no-match-loop", action="store_true", help="Disable loop matching")
    parser.add_argument("--mono-mode", type=str, default="all", 
                        choices=["check", "transform", "all", "none", "negate"],
                        help="Monotonicity analysis mode")
    parser.add_argument("--split-memop", type=int, default=1, help="Memop split parameter")
    parser.add_argument("--compact-memop", type=int, default=1, help="Memop compact parameter")
    parser.add_argument("--enable-unroll", type=int, default=0, help="Enable loop unrolling")
    parser.add_argument("--loop-split", type=int, default=0, help="Enable loop split")
    parser.add_argument("--unroll-factor", type=int, default=32, help="Loop unroll factor")

    args = parser.parse_args()
    
    setup_context(args)
    
    with open(args.sass, "r") as f:
        sasslines = f.readlines()
    
    parser_obj = SASSParser("sm70")
    funcs = parser_obj.parse_programs(sasslines)
    name, instrs = funcs.popitem()
    func = SASSFunction(name, instrs)
    print(f"Analyzing function: {name}")
    
    instances = []
    if args.trace:
        instances = load_trace_instances(args.trace, args.kernel)
        print(f"Loaded {len(instances)} instances from trace")
    
    try:
        code, kernel_level, reason = analyze(func, instances, args)
        
        output_dir = os.path.dirname(args.output)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        with open(args.output, "w") as f:
            f.write(code)
        print(f"Generated code saved to: {args.output}")
        print(f"Result: kernel_level={kernel_level}, reason={reason}")
        
    except GlobalMemoryDependencyError as e:
        print(f"Error: GlobalMemoryDependencyError - {e}")
        print("Result: kernel_level=N, reason=GMEM")
        sys.exit(1)
    except LoopNotMatchError as e:
        print(f"Error: LoopNotMatchError - {e}")
        print("Result: kernel_level=N, reason=LOOP")
        sys.exit(1)
    except (RecursiveSubroutineError, IndirectSubroutineError) as e:
        print(f"Error: SubroutineError - {e}")
        print("Result: kernel_level=N, reason=FUNC")
        sys.exit(1)
    except AtomicOperationError as e:
        print(f"Error: AtomicOperationError - {e}")
        print("Result: kernel_level=N, reason=ATOMIC")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {type(e).__name__} - {e}")
        raise e


if __name__ == '__main__':
    main()