ConvBatchModelGenerator.py

#!/usr/bin/python3
#
# ConvBatchModelGenerator.py
#
# Generates a batch>1 Conv ONNX model and its reference output for the SOFIE
# alpaka GPU test oracle (ConvBatch4). The reference is computed with
# onnxruntime so it reflects the exact ONNX Conv semantics, independent of
# anything SOFIE does. This model is the correctness oracle for the
# strided-batched Conv GEMM work (Conv batch>1 path).
#
# Usage:   python3 ConvBatchModelGenerator.py
# Needs:   pip install onnx numpy
# Writes:  input_models/ConvBatch4.onnx
#          input_models/references/ConvBatch4.ref.hxx          (expected output)
#          input_models/references/ConvBatch4_input.ref.hxx    (input data)

import os
import numpy as np
import onnx
from onnx import helper, TensorProto, numpy_helper


def conv2d_ref(X, W, Bs, pad, stride):
    # Plain cross-correlation, matching ONNX Conv semantics for the
    # symmetric-pad, unit-stride, group=1, no-dilation case used here.
    Bn, Cin, H, Wd = X.shape
    Cout, _, K, _ = W.shape
    Xp = np.pad(X, ((0, 0), (0, 0), (pad, pad), (pad, pad)), mode="constant")
    oH = (H + 2 * pad - K) // stride + 1
    oW = (Wd + 2 * pad - K) // stride + 1
    Y = np.empty((Bn, Cout, oH, oW), np.float32)
    for b in range(Bn):
        for co in range(Cout):
            for i in range(oH):
                for j in range(oW):
                    patch = Xp[b, :, i * stride:i * stride + K, j * stride:j * stride + K]
                    Y[b, co, i, j] = np.float32(np.sum(patch * W[co]) + Bs[co])
    return Y

NAME = "ConvBatch4"
B, Cin, Cout, H, W = 4, 2, 3, 5, 5      # batch=4 is the point of this oracle
K, PAD, STRIDE = 3, 1, 1

OH = (H + 2 * PAD - K) // STRIDE + 1
OW = (W + 2 * PAD - K) // STRIDE + 1

np.random.seed(42)
X  = np.random.randn(B, Cin, H, W).astype(np.float32)
Wt = (np.random.randn(Cout, Cin, K, K) * 0.2).astype(np.float32)
Bs = (np.random.randn(Cout) * 0.1).astype(np.float32)

node = helper.make_node(
    "Conv", ["input", "W", "B"], ["output"],
    kernel_shape=[K, K], pads=[PAD, PAD, PAD, PAD],
    strides=[STRIDE, STRIDE], dilations=[1, 1], group=1)

graph = helper.make_graph(
    [node], NAME,
    [helper.make_tensor_value_info("input", TensorProto.FLOAT, [B, Cin, H, W])],
    [helper.make_tensor_value_info("output", TensorProto.FLOAT, [B, Cout, OH, OW])],
    [numpy_helper.from_array(Wt, "W"), numpy_helper.from_array(Bs, "B")])

model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
model.ir_version = 8
onnx.checker.check_model(model)

# Reference output computed independently of SOFIE.
Y = conv2d_ref(X, Wt, Bs, PAD, STRIDE)

here   = os.path.dirname(os.path.abspath(__file__))
imdir  = os.path.join(here, "input_models")
refdir = os.path.join(imdir, "references")
os.makedirs(refdir, exist_ok=True)
onnx.save(model, os.path.join(imdir, NAME + ".onnx"))


def emit(path, ns, decl, arr):
    body = ", ".join("{:.8f}f".format(v) for v in arr.reshape(-1))
    with open(path, "w") as f:
        f.write("// Auto-generated by ConvBatchModelGenerator.py - DO NOT EDIT\n")
        f.write("#pragma once\n")
        f.write("namespace {} {{\n".format(ns))
        f.write("   {} = {{{}}};\n".format(decl, body))
        f.write("}} // namespace {}\n".format(ns))


emit(os.path.join(refdir, NAME + "_input.ref.hxx"),
     NAME + "_Input", "static float data[{}]".format(X.size), X)
emit(os.path.join(refdir, NAME + ".ref.hxx"),
     NAME + "_ExpectedOutput", "float output[]", Y)

print("wrote {}.onnx and references; input {} output {}".format(NAME, X.shape, Y.shape))