More readable output

2020-10-08 11:39:08 +02:00 · 2020-10-08 11:39:08 +02:00 · 9a2ae6e533
--- a/cli/cli.ts
+++ b/cli/cli.ts
@ -4,6 +4,7 @@ import * as U from '../src/util'
 import * as path from 'path'
 import * as child_process from 'child_process'
 import { program as commander } from "commander"
+import { compileAndTest } from '../src/main'

 interface CmdOptions {
    debug?: boolean;
@ -121,10 +122,12 @@ async function loadModel(modelPath: string): Promise<tf.io.ModelArtifacts> {
 }

 async function processModelFile(modelFile: string) {
+    tf.setBackend("cpu")
+
    mkdirP(options.output)
    const model = loadModel(modelFile)
    const m = await tf.loadLayersModel({ load: () => model })
-    m.summary()
+    compileAndTest(m)
 }

 export async function mainCli() {
--- a/src/compiler.ts
+++ b/src/compiler.ts
@ -24,6 +24,7 @@ interface LayerInfo {
 }

 enum OpCode {
+    comment,
    repeat,
    loadWeightAddr,
    loadDataAddr,
@ -73,6 +74,11 @@ function unsupported(msg: string) {
    throw new Error("Unsupported operator or config: " + msg)
 }

+function stringifyComment(msg: string) {
+    if (!msg) return ""
+    return "// " + msg.replace(/\n/g, "\n// ")
+}
+
 function assert(cond: boolean, msg = "assertion failed") {
    if (!cond)
        unsupported(msg)
@ -100,6 +106,8 @@ function numCycles(ops: Op[]): number {
    const addConst = (k: number) => k < (1 << 12) ? 1 : 2
    for (const op of ops) {
        switch (op.opcode) {
+            case OpCode.comment:
+                break
            case OpCode.repeat:
                cycles += (numCycles(op.body) + 4 + (op.isDef ? 1 : 0)) * op.num + 1
                break
@ -366,6 +374,9 @@ _header:
        const incr = op.increment ? "!" : ""

        switch (op.opcode) {
+            case OpCode.comment:
+                write(stringifyComment(op.fname))
+                break
            case OpCode.repeat:
                assert(op.num >= 1)
                alloc(op.dst, () => {
@ -477,6 +488,8 @@ function toJS(op: Op): string {
    const srcAlt = op.srcAlt == null ? null : reg(op.srcAlt)

    switch (op.opcode) {
+        case OpCode.comment:
+            return stringifyComment(op.fname) + "\n"
        case OpCode.repeat:
            return `for (let ${dst} = 0; ${dst} < ${op.num}; ${dst}++) {\n${indent(toJSs(op.body))}}\n`
        case OpCode.loadWeightAddr:
@ -580,6 +593,13 @@ function repeat(n: number, f: () => Op[]): Op {
    return r
 }

+function comment(str: string): Op {
+    return {
+        opcode: OpCode.comment,
+        fname: str
+    }
+}
+
 function loadWeightAddr(dst: Reg, idx: number): Op {
    assert(idx >= 0)
    return {
@ -942,8 +962,8 @@ function compileDense(layer: tf.layers.Layer) {
    const memReg0 = Reg.S1
    const flashReg0 = memReg0 + maxChunk

-    if (info.model.opts.verbose)
-        console.log(info.inputShape, info.outputShape, config)
+    //if (info.model.opts.verbose)
+    //    console.log(info.inputShape, info.outputShape, config)

    if (info.inputShape.length != 2)
        unsupported("inputShape: " + info.inputShape.length)
@ -1099,6 +1119,7 @@ const compilers: SMap<(l: tf.layers.Layer) => Op[]> = {
    Dense: compileDense,
    Dropout: noop,
    Flatten: noop,
+    InputLayer: noop,
 }

 function isInPlace(layer: tf.layers.Layer) {
@ -1119,7 +1140,13 @@ export interface Options {

 export interface CompileResult {
    execute: (inp: ArrayLike<number>) => Float32Array
+    js: string
    thumb: string
+    machineCode: Uint8Array
+}
+
+function shapeToString(shape: tf.Shape) {
+    return `[${shape.filter(x => x != null).join(",")}]`
 }

 export function compileModel(m: tf.LayersModel, opts: Options = {}) {
@ -1128,7 +1155,6 @@ export function compileModel(m: tf.LayersModel, opts: Options = {}) {
    if (opts.verbose)
        m.summary()

-
    const inputShape = m.layers[0].batchInputShape

    const modelInfo: ModelInfo = {
@ -1186,8 +1212,14 @@ export function compileModel(m: tf.LayersModel, opts: Options = {}) {
            const c0 = numCycles(tmp)
            tmp = optimize(tmp)
            const c1 = numCycles(tmp)
+            const optRate = 100 * (c0 - c1) / c0
+            const optinfo = c0 ? `${c1} cycles (${optRate.toFixed(1)}% opt)` : "(no computation)"
+            const shapeinfo = `data: ${shapeToString(info.inputShape)} => ${shapeToString(info.outputShape)}`
+            const meminfo = `mem: @${info.inputOff} -> @${info.outputOff}`
+            const infostr = `Layer: ${l.getClassName()}; ${optinfo} ${shapeinfo} ${meminfo}`
+            tmp.unshift(comment(infostr))
            if (opts.verbose)
-                console.log(l.getClassName(), c0, c1, info.inputShape, info.inputOff, info.outputOff)
+                console.log(infostr)
            ops.push(tmp)
        } else
            console.log("unsupported layer: ", l.getClassName())
@ -1198,7 +1230,14 @@ export function compileModel(m: tf.LayersModel, opts: Options = {}) {
    const lastInfo = getLayerInfo(m.layers[m.layers.length - 1])
    modelInfo.outputOffset = lastInfo.outputOff

-    let fn = `
+    const cycles = numCycles(flat)
+    const cycleinfo = `total cycles: ${cycles} (${(cycles / 84000).toFixed(3)}ms at 84MHz)`
+    flat.unshift(comment(cycleinfo))
+
+    if (opts.verbose)
+        console.log(cycleinfo)
+
+    const js = `
 (weights => {
    "use strict";
    const weightOff = ${arenaSize}
@ -1236,14 +1275,12 @@ ${toJSs(flat)}
 })
 `

-    if (opts.verbose) {
-        console.log(fn)
-        console.log("cycles:", numCycles(flat))
-    }
-
+    const thumb = toThumb(modelInfo, flat)
    const res: CompileResult = {
-        execute: (eval(fn))(modelInfo.weights),
-        thumb: toThumb(modelInfo, flat)
+        execute: (eval(js))(modelInfo.weights),
+        js,
+        thumb,
+        machineCode: null
    }

    return res
--- a/src/driver.ts
+++ b/src/driver.ts
@ -0,0 +1,104 @@
+import * as tf from '@tensorflow/tfjs'
+import { ThumbProcessor } from './thumb';
+import * as assembler from './assembler'
+import * as U from './util'
+import { compileModel, CompileResult, shapeElts } from './compiler';
+
+function mkProcessorFile() {
+    const b = new assembler.File(new ThumbProcessor())
+
+    b.ei.testAssembler(); // just in case
+
+    //   b.disablePeepHole = true
+
+    b.lookupExternalLabel = _name => null;
+    b.normalizeExternalLabel = s => s;
+    b.throwOnError = true;
+
+    return b
+}
+
+function throwAssemblerErrors(b: assembler.File) {
+    if (b.errors.length > 0) {
+        throw new Error(b.errors[0].message)
+    }
+}
+
+export function assemble(src: string) {
+    let b = mkProcessorFile()
+    b.emit(src);
+
+    throwAssemblerErrors(b)
+
+    const buf = new Uint8Array(b.buf.length << 1)
+    for (let i = 0; i < b.buf.length; ++i) {
+        buf[i << 1] = b.buf[i] & 0xff
+        buf[(i << 1) + 1] = (b.buf[i] >> 8) & 0xff
+    }
+
+    return buf
+}
+
+function randomTensor(shape: tf.Shape, mult = 1) {
+    shape = shape.map(s => s == null ? 1 : s)
+    const num = shapeElts(shape)
+    return tf.tidy(() => tf.tensor(U.range(num).map(_ => mult * U.randomSFloat())).reshape(shape))
+}
+
+export function setRandomWeights(l: tf.layers.Layer) {
+    let idx = 0
+    for (const w of l.weights) {
+        const mult = 1
+        w.write(randomTensor(w.shape, mult))
+        idx++
+    }
+}
+
+const eps = 0.00002
+function isNear(a: number, b: number) {
+    const diff = Math.abs(a - b)
+    if (diff < eps)
+        return true
+    if (diff / (Math.abs(a) + Math.abs(b)) < eps)
+        return true
+    return false
+}
+
+export function compileAndTest(m: tf.LayersModel, desc: string = "") {
+    const verbose = true
+    try {
+        const randomInput = randomTensor(m.inputs[0].shape)
+        const resTensor = m.predict(randomInput) as tf.Tensor
+        const res = resTensor.flatten().arraySync()
+        const cres = compileModel(m, {
+            verbose,
+            testInput: randomInput.flatten().arraySync(),
+            testOutput: res
+        })
+        //console.log(res)
+        const res2 = cres.execute(randomInput.flatten().arraySync())
+        //console.log(res2)
+
+        let numerr = 0
+        for (let i = 0; i < res2.length; ++i) {
+            if (!isNear(res[i], res2[i])) {
+                console.log(`at ${i} ${res[i]} - ${res2[i]} = ${res[i] - res2[i]}`)
+                numerr++
+                if (numerr > 5) break
+            }
+        }
+
+        if (numerr)
+            throw new Error("mismatch")
+
+        cres.machineCode = assemble(cres.thumb)
+
+        return cres
+    } catch (e) {
+        if (desc)
+            console.log(desc)
+        if (!verbose)
+            compileModel(m, { verbose: true })
+        throw e
+    }
+}
--- a/src/main.ts
+++ b/src/main.ts
@ -1,5 +1,5 @@
 export * from "./assembler"
 export * from "./compiler"
 export * from "./library"
-export * from "./toplevel"
+export * from "./driver"
 export * from "./util"
--- a/src/testing.ts
+++ b/src/testing.ts
@ -0,0 +1,121 @@
+import * as tf from '@tensorflow/tfjs'
+import * as U from './util'
+import { CompileResult } from './compiler';
+import { compileAndTest, setRandomWeights } from './driver';
+
+function sampleModel() {
+    const model = tf.sequential();
+
+    model.add(tf.layers.conv2d({
+        inputShape: [50, 3, 1],
+        kernelSize: [4, 3],
+        filters: 16,
+        strides: [1, 1],
+        padding: 'valid',
+        activation: 'relu',
+        kernelInitializer: 'varianceScaling'
+    }));
+
+    model.add(tf.layers.maxPooling2d({ poolSize: [2, 1], strides: [2, 1] }));
+    model.add(tf.layers.dropout({ rate: 0.1 }));
+
+    model.add(tf.layers.conv2d({
+        kernelSize: [2, 1],
+        filters: 16,
+        strides: 1,
+        activation: 'relu',
+        kernelInitializer: 'varianceScaling'
+    }));
+    //model.add(tf.layers.maxPooling2d({ poolSize: [2, 1], strides: [2, 1] }));
+    model.add(tf.layers.dropout({ rate: 0.1 }));
+
+    model.add(tf.layers.conv2d({
+        kernelSize: [2, 1],
+        filters: 16,
+        strides: 1,
+        activation: 'relu',
+        kernelInitializer: 'varianceScaling'
+    }));
+    model.add(tf.layers.dropout({ rate: 0.1 }));
+
+    model.add(tf.layers.flatten());
+
+    model.add(tf.layers.dense({
+        units: 4,
+        kernelInitializer: 'varianceScaling',
+        activation: 'softmax'
+    }));
+
+    const optimizer = tf.train.adam();
+    model.compile({
+        optimizer: optimizer,
+        loss: 'categoricalCrossentropy',
+        metrics: ['accuracy'],
+    });
+
+    // make sure weights are deterministic
+    for (const l of model.layers)
+        setRandomWeights(l)
+
+    return model;
+}
+
+
+function randomModel() {
+    const model = tf.sequential();
+
+    const inputShape = [U.randomInclusive(1, 100), U.randomInclusive(1, 50), U.randomInclusive(1, 32)]
+    const kernelSize = [U.randomInclusive(1, 5), U.randomInclusive(1, 5)]
+    const strides = [U.randomInclusive(1, 3), U.randomInclusive(1, 3)]
+    const filters = U.randomInclusive(1, 5)
+
+    kernelSize[0] = Math.min(kernelSize[0], inputShape[0])
+    kernelSize[1] = Math.min(kernelSize[1], inputShape[1])
+
+    const desc = `Conv2D ${inputShape} X ${kernelSize} @${strides} -> ${filters}`
+
+    model.add(tf.layers.conv2d({
+        inputShape,
+        kernelSize,
+        filters,
+        strides,
+        padding: 'valid',
+        activation: 'relu',
+        kernelInitializer: 'varianceScaling'
+    }));
+
+    // make sure weights are deterministic
+    for (const l of model.layers)
+        setRandomWeights(l)
+
+    return { model, desc }
+}
+
+function logThumb(cres: CompileResult) {
+    let str = cres.thumb
+    function hex2(n: number) {
+        return ("0" + n.toString(16)).slice(-2)
+    }
+    str += "// BUF: "
+    for (const v of cres.machineCode) str += hex2(v)
+    console.log(str)
+}
+
+
+export async function runBrowser() {
+    tf.setBackend('cpu');
+    const t0 = Date.now()
+    U.seedRandom(220)
+    // const m = await tf.loadLayersModel("./models/gestures.tfjsmodel.json")
+    const sample = sampleModel()
+    // compileModel(sample, { verbose: true }) 
+    logThumb(compileAndTest(sample, "sample"))
+
+    for (let i = 0; i < 0; ++i) {
+        const { model, desc } = randomModel()
+        console.log(desc)
+        compileAndTest(model, desc)
+    }
+
+    console.log(Date.now() - t0 + "ms")
+}
--- a/src/toplevel.ts
+++ b/src/toplevel.ts
@ -1,222 +0,0 @@
-import * as tf from '@tensorflow/tfjs'
-import { ThumbProcessor } from './thumb';
-import * as assembler from './assembler'
-import * as U from './util'
-import { compileModel, shapeElts } from './compiler';
-
-function mkProcessorFile() {
-    const b = new assembler.File(new ThumbProcessor())
-
-    b.ei.testAssembler(); // just in case
-
-    //   b.disablePeepHole = true
-
-    b.lookupExternalLabel = _name => null;
-    b.normalizeExternalLabel = s => s;
-    b.throwOnError = true;
-
-    return b
-}
-
-function throwAssemblerErrors(b: assembler.File) {
-    if (b.errors.length > 0) {
-        throw new Error(b.errors[0].message)
-    }
-}
-
-export function assemble(src: string) {
-    let b = mkProcessorFile()
-    b.emit(src);
-
-    throwAssemblerErrors(b)
-
-    return {
-        src: src,
-        buf: b.buf,
-        thumbFile: b
-    }
-}
-
-function randomTensor(shape: tf.Shape, mult = 1) {
-    shape = shape.map(s => s == null ? 1 : s)
-    const num = shapeElts(shape)
-    return tf.tidy(() => tf.tensor(U.range(num).map(_ => mult * U.randomSFloat())).reshape(shape))
-}
-
-function setRandomWeights(l: tf.layers.Layer) {
-    let idx = 0
-    for (const w of l.weights) {
-        const mult = 1
-        w.write(randomTensor(w.shape, mult))
-        idx++
-    }
-}
-
-function sampleModel() {
-    const model = tf.sequential();
-
-    model.add(tf.layers.conv2d({
-        inputShape: [50, 3, 1],
-        kernelSize: [4, 3],
-        filters: 16,
-        strides: [1, 1],
-        padding: 'valid',
-        activation: 'relu',
-        kernelInitializer: 'varianceScaling'
-    }));
-
-    model.add(tf.layers.maxPooling2d({ poolSize: [2, 1], strides: [2, 1] }));
-    model.add(tf.layers.dropout({ rate: 0.1 }));
-
-    model.add(tf.layers.conv2d({
-        kernelSize: [2, 1],
-        filters: 16,
-        strides: 1,
-        activation: 'relu',
-        kernelInitializer: 'varianceScaling'
-    }));
-    //model.add(tf.layers.maxPooling2d({ poolSize: [2, 1], strides: [2, 1] }));
-    model.add(tf.layers.dropout({ rate: 0.1 }));
-
-    model.add(tf.layers.conv2d({
-        kernelSize: [2, 1],
-        filters: 16,
-        strides: 1,
-        activation: 'relu',
-        kernelInitializer: 'varianceScaling'
-    }));
-    model.add(tf.layers.dropout({ rate: 0.1 }));
-
-    model.add(tf.layers.flatten());
-
-    model.add(tf.layers.dense({
-        units: 4,
-        kernelInitializer: 'varianceScaling',
-        activation: 'softmax'
-    }));
-
-    const optimizer = tf.train.adam();
-    model.compile({
-        optimizer: optimizer,
-        loss: 'categoricalCrossentropy',
-        metrics: ['accuracy'],
-    });
-
-    // make sure weights are deterministic
-    for (const l of model.layers)
-        setRandomWeights(l)
-
-    return model;
-}
-
-
-function randomModel() {
-    const model = tf.sequential();
-
-    const inputShape = [U.randomInclusive(1, 100), U.randomInclusive(1, 50), U.randomInclusive(1, 32)]
-    const kernelSize = [U.randomInclusive(1, 5), U.randomInclusive(1, 5)]
-    const strides = [U.randomInclusive(1, 3), U.randomInclusive(1, 3)]
-    const filters = U.randomInclusive(1, 5)
-
-    kernelSize[0] = Math.min(kernelSize[0], inputShape[0])
-    kernelSize[1] = Math.min(kernelSize[1], inputShape[1])
-
-    const desc = `Conv2D ${inputShape} X ${kernelSize} @${strides} -> ${filters}`
-
-    model.add(tf.layers.conv2d({
-        inputShape,
-        kernelSize,
-        filters,
-        strides,
-        padding: 'valid',
-        activation: 'relu',
-        kernelInitializer: 'varianceScaling'
-    }));
-
-    /*
-    const optimizer = tf.train.adam();
-    model.compile({
-        optimizer: optimizer,
-        loss: 'categoricalCrossentropy',
-        metrics: ['accuracy'],
-    });
-    */
-
-    // make sure weights are deterministic
-    for (const l of model.layers)
-        setRandomWeights(l)
-
-    return { model, desc }
-}
-
-const eps = 0.00002
-function isNear(a: number, b: number) {
-    const diff = Math.abs(a - b)
-    if (diff < eps)
-        return true
-    if (diff / (Math.abs(a) + Math.abs(b)) < eps)
-        return true
-    return false
-}
-
-export async function runBrowser() {
-    tf.setBackend('cpu');
-    const t0 = Date.now()
-    U.seedRandom(220)
-    // const m = await tf.loadLayersModel("./models/gestures.tfjsmodel.json")
-    const sample = sampleModel()
-    // compileModel(sample, { verbose: true })
-    compareModel(sample, "sample")
-
-    for (let i = 0; i < 0; ++i) {
-        const { model, desc } = randomModel()
-        console.log(desc)
-        compareModel(model, desc)
-    }
-
-    console.log(Date.now() - t0 + "ms")
-}
-
-function compareModel(m: tf.LayersModel, desc: string) {
-    const verbose = true
-    try {
-        const randomInput = randomTensor(m.inputs[0].shape)
-        const resTensor = m.predict(randomInput) as tf.Tensor
-        const res = resTensor.flatten().arraySync()
-        const cres = compileModel(m, {
-            verbose,
-            testInput: randomInput.flatten().arraySync(),
-            testOutput: res
-        })
-        //console.log(res)
-        const res2 = cres.execute(randomInput.flatten().arraySync())
-        //console.log(res2)
-
-        let numerr = 0
-        for (let i = 0; i < res2.length; ++i) {
-            if (!isNear(res[i], res2[i])) {
-                console.log(`at ${i} ${res[i]} - ${res2[i]} = ${res[i] - res2[i]}`)
-                numerr++
-                if (numerr > 5) break
-            }
-        }
-
-        if (numerr)
-            throw new Error("mismatch")
-
-        const asmr = assemble(cres.thumb)
-        function hex2(n: number) {
-            return ("0" + n.toString(16)).slice(-2)
-        }
-        cres.thumb += "// BUF: " + asmr.buf.map(k => hex2(k & 0xff) + hex2(k >> 8)).join("") + "\n"
-
-        if (verbose)
-            console.log(cres.thumb)
-
-    } catch (e) {
-        console.log(desc)
-        if (!verbose)
-            compileModel(m, { verbose: true })
-        throw e
-    }
-}