<commit_before>#!/usr/bin/env python3
"""Helper functions to instrument assembly code for and analyze with IACA."""
# Version check
import sys
import re
import subprocess
import os
from copy import copy
import argparse
from pprint import pformat, pprint
import textwrap
from collections import OrderedDict, defaultdict
import io
from hashlib import md5
from os.path import expanduser
from itertools import chain

from distutils.spawn import find_executable
from osaca import osaca
from osaca.parser import get_parser
from osaca.semantics import MachineModel, ISASemantics
from osaca.semantics.marker_utils import find_basic_loop_bodies, get_marker

from kerncraft import iaca_get, __version__


def itemsEqual(lst):
   return lst[1:] == lst[:-1]


class IncoreModel:
    def __init__(self, isa='x86'):
        isa


class IACA(IncoreModel):
    pass


class OSACA(IncoreModel):
    pass


class LlvmMCA(IncoreModel):
    pass


class ISA:
    @staticmethod
    def get_isa(isa='x86'):
        if isa.lower() == 'x86':
            return x86
        elif isa.lower() == 'aarch64':
            return AArch64

    @staticmethod
    def compute_block_metric(block):
        """Compute sortable metric to rank blocks."""
        return NotImplementedError

    @classmethod
    def select_best_block(cls, blocks):
        """
        Return best block label selected based on simple heuristic.

        :param blocks: OrderedDict map of label to list of instructions
        """
        # TODO make this cleverer with more stats
        if not blocks:
            raise ValueError("No suitable blocks were found in assembly.")

        best_block_label = next(iter(blocks))
        best_metric = cls.compute_block_metric(blocks[best_block_label])

        for label, block in list(blocks.items())[1:]:
            metric = cls.compute_block_metric(block)
            if best_metric < metric:
                best_block_label = label
                best_metric = metric

        return best_block_label

    @staticmethod
    def get_pointer_increment(block):
        """Return pointer increment."""
        raise NotImplementedError


class x86(ISA):
    @staticmethod
    def compute_block_metric(block):
        """Return comparable metric on block information."""
        register_class_usage = {'zmm': [], 'ymm': [], 'xmm': []}
        packed_instruction_ctr, avx_instruction_ctr, instruction_ctr = 0, 0, 0
        # Analyze code to determine metric
        for line in block:
            # Skip non-instruction lines (e.g., comments)
            if line.instruction is None:
                continue
            # Count all instructions
            instruction_ctr += 1

            # Count registers used
            for prefix in register_class_usage:
                for op in line.operands:
                    if 'register' in op:
                        if op.register.name.startswith(prefix):
                            register_class_usage[prefix].append(op.register.name)

            # Identify and count packed and avx instructions
            if re.match(r"^[v]?(movu|mul|add|sub|div|fmadd(132|213|231)?)[h]?p[ds]",
                    line.instruction):
                if line.instruction.startswith('v'):
                    avx_instruction_ctr += 1
                packed_instruction_ctr += 1

        # Build metric
        return (packed_instruction_ctr, avx_instruction_ctr, instruction_ctr,
                len(set(register_class_usage['zmm'])),
                len(set(register_class_usage['ymm'])),
                len(set(register_class_usage['xmm'])))

    @staticmethod
    def get_pointer_increment(block):
        """Return pointer increment."""
        increments = {}
        mem_references = []
        stores_only = False

        modified_registers = []
        for line in block:
            # Skip non-instruction lines (e.g., comments)
            if line.instruction is None:
                continue

            # Extract destination references, ignoring var(%rip)
            dst_mem_references = [op.memory for op in line.semantic_operands.destination
                                  if 'memory' in op and
                                  (op.memory.base is None or op.memory.base.name != 'rip')]
            if dst_mem_references:
                if not stores_only:
                    stores_only = True
                    mem_references = []
                mem_references += dst_mem_references

            # If no destination references were found sofar, include source references
            if not stores_only:
                mem_references += [op.memory for op in line.semantic_operands.source
                                   if 'memory' in op]
            if re.match(r'^inc[bwlq]?$', line.instruction):
                reg = line.operands[0].register.name
                modified_registers.append(reg)
                increments[reg] = 1
            elif re.match(r'^add[bwlq]?$', line.instruction) and 'immediate' in line.operands[0] \
                    and 'register' in line.operands[1]:
                increments[line.operands[1].register.name] = int(line.operands[0].immediate.value)
            elif re.match(r'^dec[bwlq]?$', line.instruction):
                reg = line.operands[0].register.name
                modified_registers.append(reg)
                increments[reg] = -1
            elif re.match(r'^sub[bwlq]?$', line.instruction) and 'immediate' in line.operands[0] \
                    and 'register' in line.operands[1]:
                reg = line.operands[1].register.name
                modified_registers.append(reg)
                increments[reg] = -int(line.operands[0].immediate.value)
            elif re.match(r'^lea[bwlq]?$', line.instruction):
                # `lea 1(%r11), %r11` is the same as `add $1, %r11`
                if line.operands[0].memory.base.name  == line.operands[1].register.name and \
                        line.operands[0].memory.index is None:
                    reg = line.operands[1].register.name
                    modified_registers.append(reg)
                    increments[reg] = int(
                        line.operands[0].memory.offset.value)

        # deduce loop increment from memory index register
        pointer_increment = None  # default -> can not decide, let user choose
        possible_idx_regs = None
        if mem_references:
            # we found memory references to work with
            possible_idx_regs = list(set(increments.keys()).intersection(
                set([mref.base.name for mref in mem_references if mref.base is not None] +
                    [mref.index.name for mref in mem_references if mref.index is not None])))
            for mref in mem_references:
                for reg in list(possible_idx_regs):
                    # Only consider references with two registers, where one could be an
                    # index
                    if None not in [mref.base, mref.index]:<commit_msg>fixed problem in pointer increment detection<commit_after>#!/usr/bin/env python3
"""Helper functions to instrument assembly code for and analyze with IACA."""
# Version check
import sys
import re
import subprocess
import os
from copy import copy
import argparse
from pprint import pformat, pprint
import textwrap
from collections import OrderedDict, defaultdict
import io
from hashlib import md5
from os.path import expanduser
from itertools import chain

from distutils.spawn import find_executable
from osaca import osaca
from osaca.parser import get_parser
from osaca.semantics import MachineModel, ISASemantics
from osaca.semantics.marker_utils import find_basic_loop_bodies, get_marker

from kerncraft import iaca_get, __version__


def itemsEqual(lst):
   return lst[1:] == lst[:-1]


class IncoreModel:
    def __init__(self, isa='x86'):
        isa


class IACA(IncoreModel):
    pass


class OSACA(IncoreModel):
    pass


class LlvmMCA(IncoreModel):
    pass


class ISA:
    @staticmethod
    def get_isa(isa='x86'):
        if isa.lower() == 'x86':
            return x86
        elif isa.lower() == 'aarch64':
            return AArch64

    @staticmethod
    def compute_block_metric(block):
        """Compute sortable metric to rank blocks."""
        return NotImplementedError

    @classmethod
    def select_best_block(cls, blocks):
        """
        Return best block label selected based on simple heuristic.

        :param blocks: OrderedDict map of label to list of instructions
        """
        # TODO make this cleverer with more stats
        if not blocks:
            raise ValueError("No suitable blocks were found in assembly.")

        best_block_label = next(iter(blocks))
        best_metric = cls.compute_block_metric(blocks[best_block_label])

        for label, block in list(blocks.items())[1:]:
            metric = cls.compute_block_metric(block)
            if best_metric < metric:
                best_block_label = label
                best_metric = metric

        return best_block_label

    @staticmethod
    def get_pointer_increment(block):
        """Return pointer increment."""
        raise NotImplementedError


class x86(ISA):
    @staticmethod
    def compute_block_metric(block):
        """Return comparable metric on block information."""
        register_class_usage = {'zmm': [], 'ymm': [], 'xmm': []}
        packed_instruction_ctr, avx_instruction_ctr, instruction_ctr = 0, 0, 0
        # Analyze code to determine metric
        for line in block:
            # Skip non-instruction lines (e.g., comments)
            if line.instruction is None:
                continue
            # Count all instructions
            instruction_ctr += 1

            # Count registers used
            for prefix in register_class_usage:
                for op in line.operands:
                    if 'register' in op:
                        if op.register.name.startswith(prefix):
                            register_class_usage[prefix].append(op.register.name)

            # Identify and count packed and avx instructions
            if re.match(r"^[v]?(movu|mul|add|sub|div|fmadd(132|213|231)?)[h]?p[ds]",
                    line.instruction):
                if line.instruction.startswith('v'):
                    avx_instruction_ctr += 1
                packed_instruction_ctr += 1

        # Build metric
        return (packed_instruction_ctr, avx_instruction_ctr, instruction_ctr,
                len(set(register_class_usage['zmm'])),
                len(set(register_class_usage['ymm'])),
                len(set(register_class_usage['xmm'])))

    @staticmethod
    def get_pointer_increment(block):
        """Return pointer increment."""
        increments = {}
        mem_references = []
        stores_only = False

        modified_registers = []
        for line in block:
            # Skip non-instruction lines (e.g., comments)
            if line.instruction is None:
                continue

            # Extract destination references, ignoring var(%rip)
            dst_mem_references = [op.memory for op in line.semantic_operands.destination
                                  if 'memory' in op and
                                  (op.memory.base is None or op.memory.base.name != 'rip')]
            if dst_mem_references:
                if not stores_only:
                    stores_only = True
                    mem_references = []
                mem_references += dst_mem_references

            # If no destination references were found sofar, include source references
            if not stores_only:
                mem_references += [op.memory for op in line.semantic_operands.source
                                   if 'memory' in op]
            if re.match(r'^inc[bwlq]?$', line.instruction):
                reg = line.operands[0].register.name
                modified_registers.append(reg)
                increments[reg] = 1
            elif re.match(r'^add[bwlq]?$', line.instruction) and 'immediate' in line.operands[0] \
                    and 'register' in line.operands[1]:
                reg = line.operands[1].register.name
                increments[reg] = int(line.operands[0].immediate.value)
                modified_registers.append(reg)
            elif re.match(r'^dec[bwlq]?$', line.instruction):
                reg = line.operands[0].register.name
                modified_registers.append(reg)
                increments[reg] = -1
            elif re.match(r'^sub[bwlq]?$', line.instruction) and 'immediate' in line.operands[0] \
                    and 'register' in line.operands[1]:
                reg = line.operands[1].register.name
                modified_registers.append(reg)
                increments[reg] = -int(line.operands[0].immediate.value)
            elif re.match(r'^lea[bwlq]?$', line.instruction):
                # `lea 1(%r11), %r11` is the same as `add $1, %r11`
                if line.operands[0].memory.base.name  == line.operands[1].register.name and \
                        line.operands[0].memory.index is None:
                    reg = line.operands[1].register.name
                    modified_registers.append(reg)
                    increments[reg] = int(
                        line.operands[0].memory.offset.value)

        # deduce loop increment from memory index register
        pointer_increment = None  # default -> can not decide, let user choose
        possible_idx_regs = None
        if mem_references:
            # we found memory references to work with
            possible_idx_regs = list(set(increments.keys()).intersection(
                set([mref.base.name for mref in mem_references if mref.base is not None] +
                    [mref.index.name for mref in mem_references if mref.index is not None])))
            for mref in mem_references:
                for reg in list(possible_idx_regs):
                    # Only consider references with two registers, where one could be an
                    # index
                    if None not in [mref.base, mref.index]:
{'fraction_non_alphanumeric': 0.08838599487617421, 'fraction_numerical': 0.008397381155707373, 'mean_word_length': 1.9982931512694688, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7556490', 'n_tokens_mistral': 3986, 'n_tokens_neox': 3745, 'n_words': 1029}