リポジトリ / emssr / independent_common_compiler

# -*- coding: utf-8 -*-

"""lane64 の三本めの道。W 枚の lane を numpy の配列に束ね、block を一度に走らせる。

doc: casm.lane64.sem casm.lane64.merge casm.lane64.cost casm.simd

    解釈器   vcpu.Interp     式 (2.1) の mask 付き意味。これが定義になる

    翻訳器   vcpu.run_lane   block を Python の code object へ落とす。一枚

    束ね     ここ            同じ block を走る lane を配列の一次元に並べる

命題 6.1 が束ねの根拠になる。lane の作業領域は式 (1.2) で交わらないので、

(mask, pc) の組を pc 昇順に進める観測は、lane を一本ずつ走らせた観測に一致する。

だから組を「lane の添字」で持ってよく、mask の bit 演算は要らない。

20 章 §0.4 は「大整数の欄へ W 個を詰める」形を退けている。乗算が欄をまたいで

畳み込み、LD と ST は lane ごとに番地が違うためになる。配列はその制限を受けない。

gather と scatter が LD と ST をそのまま運ぶので、頭打ちの式 (0.3) は効かない。

実測（CPython 3.12 + numpy 2.5、一 lane あたり ns）

    種      list   W=256   W=1024      倍

    ADD     37.3    3.54     1.43     26

    MUL     35.5    3.48     1.44     25

    LD 定   26.9    1.46     0.54     50      番地が定数。M の行がそのまま連続になる

    ST 定   28.9    1.56     0.53     55

    LD 変   26.9   10.21     5.87      4.6    番地がワイヤの値。gather になる

    ST 変   28.9    8.56     4.94      5.9

    比較    30.8    3.17     1.18     26

作業領域は (語数, W) に取る。(W, 語数) では定数番地が stride になり、可変番地の

gather も遅くなる（W=1024 で 5.95 対 7.80 ns）。

    重さ = 語数 x W x 8 byte。W=256 で 語数 70000 なら 143 MB

語は Z（多倍長）が仕様（§0.3）だが、配列は int64 になる。だから `overflow` が

溢れを見る。溢れたら翻訳器の道へ落とす。証とソルバーの語は在庫・金・手数・

区画の番号なので 2^62 に遠く及ばない。commonc の compile はここへ載せない。

"""

import numpy as np

import vcpu

from vcpu import (ADD, EQ, HALT, JMP, KO, LD, LT, MAT, MUL, NEG, OK, PUT, ST,

                  LANDMARKS, NAME, R_END, R_KO, R_OK, blocks)

GUARD = 1 << 62          # これを越えた語が出たら int64 の巻き戻しを疑う

class LaneVec(object):

    """W 枚の lane（式 1.1 の Σ_lane64 を W 本）。

    R[a] が landmark a の lane 方向のベクトル、M[c] が語 c の lane 方向のベクトルになる。

    どちらも第一軸が番号なので、定数番地は連続した W 語になる。

    """

    def __init__(self, width, words):

        self.W = width

        self.R = np.zeros((LANDMARKS, width), dtype=np.int64)

        self.M = np.zeros((words, width), dtype=np.int64)

        self.ix = np.arange(width, dtype=np.int64)

        self.OUT = [[] for _ in range(width)]

        self.report = [None] * width

        self.blocks = 0

    def seed(self, addr, values):

        """語 addr へ lane ごとの値を置く。put_input と違い、こちらは一語だけ動かす"""

        self.M[addr] = values

    def put_input(self, data, base):

        """18 章 §7.1 と同じ置き方（式 8.2）。全 lane へ同じ入力を配る"""

        n = min(len(data), self.M.shape[0] - base)

        self.M[0] = n

        self.M[1] = base

        self.M[base:base + n] = np.asarray(data[:n], dtype=np.int64)[:, None]

        return n

    def overflow(self):

        """int64 の巻き戻しを疑う語があるか。無ければ False"""

        return bool(np.abs(self.R).max() >= GUARD or np.abs(self.M).max() >= GUARD)

# ===== 分岐 =====

#

# 閉じの判定は lane ごとに割れる。三通りを返す。

#

#   全部が真     行き先ひとつ

#   全部が偽     行き先ひとつ

#   割れた       (真の lane の印, 真の行き先, 偽の行き先)

#

# 割れた先は run_vec が組へ足す。命題 6.1 により pc 昇順で進めれば観測は変わらない。

def _br(cond, yes, no):

    if cond.all():

        return yes

    if not cond.any():

        return no

    return (cond, yes, no)

# ===== block を numpy の文へ落とす =====

#

# 局所の値番号付けは翻訳器（vcpu._block）と同じ。違いは三つだけになる。

#

#   landmark   R[a][s]        s は組の印。全 lane なら slice(None) で view になる

#   作業領域   M[c][s]        番地が定数のとき。連続した W 語

#              M[addr, ixs]   番地がワイヤの値のとき。gather / scatter

#   即値       そのまま置く。numpy が broadcast する

def _block_vec(prog, k, body, term, stop, idx, n, out_live=None, dead=None, scan=None):

    steps = list(body) + ([] if term is None else [term])

    reg, mem, tmp = {}, {}, [0]

    out = ['def _v%d(R, M, OUT, ixs, s):' % k]

    # 定数番地の Load は M[c][s] になる。s が slice(None) のとき、これは写しではなく

    # view になる。同じ block が後で M[c][s] へ書くと、先に読んだ名前の中身が一緒に

    # 動く。翻訳器の道では M[c] が Python の int なので起きない。

    #

    #   t1 = M[A][s]        view

    #   M[A][s] = t2

    #   M[B][s] = t1 + t2   t1 は既に t2。Fibonacci が 2 の冪へ化ける

    #

    # 後ろに Store がある Load だけ写しを取る。番地がワイヤの値のときは

    # fancy indexing が写しを返すので、そのままでよい。

    last_st = -1

    for _p, _i in enumerate(body):

        if prog[_i][0] == ST:

            last_st = _p

    def get(r):

        if r not in reg:

            reg[r] = 'R[%d][s]' % r

        return reg[r]

    def lit(e):

        try:

            return int(e)

        except ValueError:

            return None

    def shifted(r, sh):

        e = get(r)

        v = lit(e)

        if sh == 0:

            return e

        if v is not None:

            return str(v << sh)

        return '(%s<<%d)' % (e, sh)

    def fresh(expr):

        tmp[0] += 1

        name = 't%d' % tmp[0]

        out.append('    %s=%s' % (name, expr))

        return name

    def fold2(x, y, f, sym):

        vx, vy = lit(x), lit(y)

        if vx is not None and vy is not None:

            return str(f(vx, vy))

        return fresh('%s%s%s' % (x, sym, y))

    for pos, i in enumerate(body):

        op, a, b, c, sb, sc, im = prog[i]

        if op == MAT:

            reg[a] = str(im)

        elif op == ADD:

            reg[a] = fold2(shifted(b, sb), shifted(c, sc), lambda x, y: x + y, '+')

        elif op == MUL:

            reg[a] = fold2(get(b), get(c), lambda x, y: x * y, '*')

        elif op == NEG:

            e = get(b)

            v = lit(e)

            reg[a] = str(-v) if v is not None else fresh('-%s' % e)

        elif op == LD:

            addr = get(b)

            v = lit(addr)

            if v is not None and v in mem:

                reg[a] = mem[v]

            else:

                if scan is not None:

                    scan.add(v)

                reg[a] = fresh(('M[%s][s].copy()' if pos < last_st else 'M[%s][s]') % addr

                               if v is not None else 'M[%s,ixs]' % addr)

                if v is not None:

                    mem[v] = reg[a]

        elif op == ST:

            addr, val = get(b), get(c)

            v = lit(addr)

            if not (dead is not None and v in dead):

                out.append('    M[%s][s]=%s' % (addr, val) if v is not None

                           else '    M[%s,ixs]=%s' % (addr, val))

            if v is not None:

                mem[v] = val

            else:

                mem.clear()

        elif op == PUT:

            # 出力は lane ごとに別の列になるので、ここだけは Python の loop になる。

            # ixs が lane の実番号を持つので、組が割れていても行き先を取り違えない

            out.append('    _o=%s' % get(b))

            out.append('    for _p in range(len(ixs)): OUT[int(ixs[_p])].append(int(_o[_p])&255)')

        else:

            raise ValueError('直線部分に置けない箱: %s' % NAME[op])

    keep = set(reg) if out_live is None else (set(reg) & out_live)

    for r in sorted(keep):

        if reg[r] != 'R[%d][s]' % r:

            out.append('    R[%d][s]=%s' % (r, reg[r]))

    def to(t):

        return idx[t] if 0 <= t < n and t in idx else R_END

    if term is None:

        out.append('    return %d' % (idx[stop] if stop < n else R_END))

    else:

        op, _a, b, c, sb, sc, im = prog[term]

        if op == OK:

            out.append('    return %d' % R_OK)

        elif op == KO:

            out.append('    return %d' % R_KO)

        elif op in (HALT,):

            out.append('    return %d' % R_END)

        elif op == JMP:

            out.append('    return %d' % to(term + im))

        else:

            x, y = shifted(b, sb), shifted(c, sc)

            vx, vy = lit(x), lit(y)

            if vx is not None and vy is not None:

                taken = (vx == vy) if op == EQ else (vx < vy)

                out.append('    return %d' % (to(term + im) if taken else to(term + 1)))

            else:

                rel = '==' if op == EQ else '<'

                out.append('    return _br(%s%s%s,%d,%d)' % (x, rel, y, to(term + im), to(term + 1)))

    return out + ['']

def translate(prog, wires=None):

    """image を numpy の block 関数の表へ落とす。引数は vcpu.translate と同じ意味になる。

    wires = (lo, hi) はワイヤの区間。他のどの block も読まない Store を落とす。

    """

    ent, idx, bl = blocks(prog)

    n = len(prog)

    live = vcpu._liveness(prog, bl, idx, n)

    dead = None

    if wires is not None:

        lo, hi = wires

        read = set()

        for k, (pc, body, term, stop) in enumerate(bl):

            _block_vec(prog, k, body, term, stop, idx, n, live[k], scan=read)

        dead = set(c for c in range(lo, hi) if c not in read)

    lines = []

    for k, (pc, body, term, stop) in enumerate(bl):

        lines += _block_vec(prog, k, body, term, stop, idx, n, live[k], dead=dead)

    lines.append('TAB = (' + ''.join('_v%d,' % k for k in range(len(bl))) + ')')

    src = chr(10).join(lines)

    ns = {'_br': _br, 'np': np}

    exec(compile(src, '<lane64-vec>', 'exec'), ns)

    return src, ns['TAB']

# ===== 組を進める（命題 6.1）=====

#

# 組は (lane の印, block 番号) の対で持つ。印は slice(None) か添字の配列になる。

# slice(None) のときは R も M も view を返すので、fancy indexing の写しが出ない。

# 組が割れたら添字の配列になり、同じ block 番号の組は一つへ合わせる。

#

# 「pc 昇順」は block 番号の昇順で足りる。block の入口は pc の昇順に並ぶので

# （vcpu.entries が sorted を返す）、番号の順序が pc の順序に一致する。

_ALL = slice(None)

def _join(groups, fresh):

    """同じ block 番号の組を合わせる（式 3.3 の reconvergence）"""

    for sel, b in fresh:

        for k, (s2, b2) in enumerate(groups):

            if b2 == b:

                if s2 is _ALL or sel is _ALL:

                    groups[k] = (_ALL, b)

                else:

                    groups[k] = (np.union1d(s2, sel), b)

                break

        else:

            groups.append((sel, b))

    return groups

def run(tab, lv, budget=1 << 32):

    """W 枚を束ねて走らせる。返りは lane ごとの報告の列。

    報告は翻訳器と同じ値を使う。0 が tu、42 が te、None が停止になる。

    """

    R, M, OUT, ix = lv.R, lv.M, lv.OUT, lv.ix

    groups = [(_ALL, 0)]

    steps = 0

    while groups:

        steps += 1

        if steps > budget:

            raise RuntimeError('組の歩数上限 %d を越えた' % budget)

        # 最も小さい block 番号の組を取る（命題 6.1）

        k = 0 if len(groups) == 1 else min(range(len(groups)), key=lambda i: groups[i][1])

        sel, b = groups.pop(k)

        ixs = ix if sel is _ALL else ix[sel]

        r = tab[b](R, M, OUT, ixs, sel)

        if isinstance(r, tuple):

            cond, yes, no = r

            base = ix if sel is _ALL else ix[sel]

            groups = _join(groups, [(base[cond], yes), (base[~cond], no)])

            continue

        if r >= 0:

            groups = _join(groups, [(sel, r)])

            continue

        rep = 0 if r == R_OK else (42 if r == R_KO else None)

        if sel is _ALL:

            lv.report = [rep] * lv.W

        else:

            for l in sel:

                lv.report[int(l)] = rep

    lv.blocks = steps

    return lv.report

# ===== 突き合わせ（命題 6.1 の検査）=====

def agrees(image, seeds, words, need=None, watch=()):

    """束ねの道と翻訳器の道が同じ観測を返すか。

    命題 6.1 の最後が「二つの観測が一致することが、この package の検査の一つになる」

    と言う。seeds は lane ごとの {番地: 値}。返りは (合うか, 束ねの報告, 翻訳器の報告)。

    """

    prog = vcpu.decode(image)

    W = len(seeds)

    lv = LaneVec(W, words)

    for l, s in enumerate(seeds):

        for a, v in s.items():

            lv.M[a, l] = v

    _src, tabv = translate(prog, wires=(64, need) if need else None)

    rep_v = list(run(tabv, lv))

    val_v = [[int(lv.M[a, l]) for a in watch] for l in range(W)]

    _src2, tab = vcpu.translate(prog, wires=(64, need) if need else None)

    rep_s, val_s = [], []

    for s in seeds:

        lane = vcpu.Lane(words)

        for a, v in s.items():

            lane.M[a] = v

        rep_s.append(vcpu.run_lane(tab, lane))

        val_s.append([lane.M[a] for a in watch])

    ok = (rep_v == rep_s and val_v == val_s and not lv.overflow())

    return ok, rep_v, rep_s