TheRock/rocm-systems/emulation/rocjitsu/lib/python/amdisa/legalization_codegen.py

# Copyright (c) 2026 Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""Emit C++20 legalization tables from ``LegalizationGenerator`` output.

Generates:
- ``legalization_types.h`` — shared types (``Action``, ``InstructionLegalization``,
  ``lookup()``).
- One ``legalization_<src>_<dst>.h`` header per ISA pair containing an
  ``inline constexpr InstructionLegalization[]`` array.
"""

from __future__ import annotations

from pathlib import Path
import shutil
import subprocess
import sys

from amdisa.legalization import LegalizationEntry

_COPYRIGHT = """\
// Copyright (c) 2025-2026 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: MIT
//
// AUTO-GENERATED by the amdisa codegen pipeline. DO NOT EDIT.
"""


def _action_enum(entry: LegalizationEntry) -> str:
    return {
        'identity': 'Action::Identity',
        'substitute': 'Action::Substitute',
        'lower': 'Action::Lower',
        'expand': 'Action::Expand',
        'illegal': 'Action::Illegal',
    }[entry.action.kind]


def _pair_name(src: str, dst: str) -> str:
    return f'{src}_to_{dst}'


def _clang_format(path: Path) -> None:
    """Format emitted C++ headers when clang-format is available.

    The generator emits simple 4-space templates; formatting here keeps the
    checked-in generated headers aligned with rocjitsu's 2-space C++ style.
    """
    exe = shutil.which('clang-format')
    if exe is None:
        print(
            'warning: clang-format not found; emitted files may need manual '
            'formatting',
            file=sys.stderr,
        )
        return
    try:
        subprocess.run([exe, '-i', str(path)], check=True)
    except subprocess.CalledProcessError as exc:
        print(f'warning: clang-format failed for {path}: {exc}', file=sys.stderr)


def emit_all(
    output_dir: str | Path,
    pairs: list[tuple[str, str, list[LegalizationEntry]]],
) -> list[str]:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    generated: list[str] = []

    types_path = output_dir / 'legalization_types.h'
    _emit_types(types_path)
    _clang_format(types_path)
    generated.append(str(types_path))

    for src, dst, entries in pairs:
        # Duplicate entries for encoding formats with narrow (<9-bit) encoding
        # fields. The 9-bit encoding_id (raw[0]>>23) has don't-care low bits
        # for these formats. We emit one entry per alias so that binary search
        # works without encoding_id normalization.
        expanded = []
        for e in entries:
            enc_id = e.src_encoding_order if e.src_encoding_order >= 0 else 0xFFFF
            dont_care = 9 - e.src_encoding_bits
            if dont_care > 0 and enc_id < 0xFFFF:
                for i in range(1 << dont_care):
                    alias_id = enc_id | i
                    expanded.append(
                        LegalizationEntry(
                            src_mnemonic=e.src_mnemonic,
                            src_encoding=e.src_encoding,
                            src_encoding_order=alias_id,
                            src_encoding_bits=9,
                            src_opcode=e.src_opcode,
                            action=e.action,
                        )
                    )
            else:
                expanded.append(e)

        sorted_entries = sorted(
            expanded,
            key=lambda e: (
                e.src_encoding_order if e.src_encoding_order >= 0 else 0xFFFF,
                e.src_opcode,
            ),
        )
        pair_path = output_dir / f'legalization_{_pair_name(src, dst)}.h'
        _emit_pair(pair_path, src, dst, sorted_entries)
        _clang_format(pair_path)
        generated.append(str(pair_path))

    return generated


def _emit_types(path: Path) -> None:
    lines = [
        _COPYRIGHT,
        '',
        '#pragma once',
        '',
        '#include <algorithm>',
        '#include <compare>',
        '#include <cstdint>',
        '#include <span>',
        '',
        'namespace rocjitsu {',
        '',
        'enum class Action : uint8_t {',
        '    Identity   = 0,',
        '    Substitute = 1,',
        '    Lower      = 2,',
        '    Expand     = 3,',
        '    Illegal    = 4,',
        '};',
        '',
        'struct InstructionLegalization {',
        '    uint16_t src_opcode;',
        '    uint16_t src_encoding_id;',
        '    Action   action;',
        '    uint16_t target_opcode;',
        '',
        '    constexpr auto operator<=>(const InstructionLegalization &rhs) const {',
        '        if (auto cmp = src_encoding_id <=> rhs.src_encoding_id; cmp != 0)',
        '            return cmp;',
        '        return src_opcode <=> rhs.src_opcode;',
        '    }',
        '    constexpr bool operator==(const InstructionLegalization &rhs) const {',
        '        return src_encoding_id == rhs.src_encoding_id && src_opcode == rhs.src_opcode;',
        '    }',
        '};',
        '',
        'inline const InstructionLegalization *lookup(',
        '    std::span<const InstructionLegalization> table,',
        '    uint16_t encoding_id, uint16_t opcode) {',
        '    InstructionLegalization key{.src_opcode = opcode, .src_encoding_id = encoding_id,',
        '                               .action = Action::Illegal, .target_opcode = 0};',
        '    auto it = std::lower_bound(table.begin(), table.end(), key);',
        '    if (it != table.end() && *it == key)',
        '        return &*it;',
        '    return nullptr;',
        '}',
        '',
        '}  // namespace rocjitsu',
        '',
    ]

    path.write_text('\n'.join(lines))


def _emit_pair(
    path: Path,
    src: str,
    dst: str,
    entries: list[LegalizationEntry],
) -> None:
    name = _pair_name(src, dst)
    lines = [
        _COPYRIGHT,
        '',
        '#pragma once',
        '',
        '#include "legalization_types.h"',
        '',
        'namespace rocjitsu {',
        '',
        f'inline constexpr InstructionLegalization kLegalization_{name}[] = {{',
    ]

    for e in entries:
        tgt = e.action.target_opcode
        enc_id = e.src_encoding_order if e.src_encoding_order >= 0 else 0xFFFF
        lines.append(
            f'    {{{e.src_opcode:>5}, {enc_id:>5}, '
            f'{_action_enum(e):<20s}, {tgt:>5}}},'
        )

    lines += ['};', '', '}  // namespace rocjitsu', '']

    path.write_text('\n'.join(lines))