#!/usr/bin/env python3
"""Generate kern.s: 604e latency/throughput kernels in AIX assembler syntax.

Every kernel is  name(long n, ...): mtctr n; loop: UNROLL x body; bdnz loop
so one call executes n * UNROLL body instructions.  Arguments arrive in
r3.. (integers) and fp1.. (doubles), per the AIX calling convention.
"""
UNROLL = 128

KERNELS = [
    # name, body lines (one unrolled step; may be several instructions), result
    ("k_add",    ["add 4,4,5"],                       "mr 3,4"),
    ("k_add4",   ["add 6,6,5", "add 7,7,5", "add 8,8,5", "add 9,9,5"], "mr 3,6"),
    ("k_mullw",  ["mullw 4,4,5"],                     "mr 3,4"),
    ("k_divw",   ["divw 4,4,5"],                      "mr 3,4"),
    ("k_fadd",   ["fadd 1,1,2"],                      None),
    ("k_fmadd",  ["fmadd 1,1,2,3"],                   None),
    ("k_fdiv",   ["fdiv 1,1,2"],                      None),
    ("k_fdivs",  ["fdivs 1,1,2"],                     None),
    ("k_chase",  ["lwz 4,0(4)"],                      "mr 3,4"),
]

def kernel(name, body, result):
    out = []
    out.append("\t.csect\t.%s{PR}" % name)
    out.append("\t.globl\t.%s" % name)
    out.append(".%s:" % name)
    if name in ("k_add", "k_mullw", "k_divw", "k_chase"):
        pass                                   # r4 = start value / pointer, r5 = operand
    if name == "k_add4":
        out += ["\tmr 6,4", "\tmr 7,4", "\tmr 8,4", "\tmr 9,4"]
    out.append("\tmtctr 3")
    out.append("L.%s:" % name)
    steps = UNROLL // len(body) if name == "k_add4" else UNROLL
    for _ in range(steps):
        for b in body:
            out.append("\t" + b)
    out.append("\tbdnz L.%s" % name)
    if result:
        out.append("\t" + result)
    out.append("\tblr")
    # function descriptor
    out.append("\t.csect\t%s{DS}" % name)
    out.append("\t.globl\t%s{DS}" % name)
    out.append("\t.long\t.%s, TOC{TC0}, 0" % name)
    return out

lines = ['\t.machine\t"604"', '\t.file\t"kern.s"']
for k in KERNELS:
    lines += kernel(*k)
# k_tb: read the 64-bit time base (upper, lower, upper again until stable)
lines += [
    "\t.csect\t.k_tb{PR}", "\t.globl\t.k_tb", ".k_tb:",
    "L.k_tb:", "\tmftbu 3", "\tmftb 4", "\tmftbu 5", "\tcmpw 3,5", "\tbne L.k_tb", "\tblr",
    "\t.csect\tk_tb{DS}", "\t.globl\tk_tb{DS}", "\t.long\t.k_tb, TOC{TC0}, 0",
]
lines.append("\t.toc")
open("kern.s", "w").write("\n".join(lines) + "\n")
print("kern.s:", len(lines), "lines,", len(KERNELS) + 1, "functions, unroll", UNROLL)