TVM Quantization

def _percentile_scale(mod, dataset):
    cfg = quantize.current_qconfig()
    chunk_by = cfg.calibrate_chunk_by
    scales = []
    for samples in collect_stats(mod, dataset, chunk_by):
        logging.info("finding threshold with percentile for calibration...")
        with mp.Pool() as pool:
            scales += list(pool.map(_find_scale_by_percentile, samples))

    def func(_):
        # func 返回每一层的 scale (实际上相当于每层的`最大值`, 后续会通过
        # scale/valid_range 获得最终的 scale)
        scale = scales[func.scale_idx]
        func.scale_idx += 1
        return scale

    func.scale_idx = 0

    return func


# 使用 calibrate_dataset 获得每一层的输出
def collect_stats(mod, dataset, chunk_by=-1):
    runtime = _get_profile_runtime(mod)
    num_outputs = runtime.get_num_outputs()

    for i in range(0, num_outputs, chunk_by):
        outputs = [[] for i in range(min(chunk_by, num_outputs - i))]
        for batch in dataset:
            runtime.set_input(**batch)
            runtime.run()
            for j in range(i, min(i + chunk_by, num_outputs)):
                outputs[j - i].append(runtime.get_output(j).numpy())
        yield [np.concatenate(output).reshape(-1) for output in outputs]


# 这个函数相当于一个 get_k_largest 函数, 其中 k=percentile*(arr.size), 用来避免
# 个别极大的值造成量化误差
# file:~/Gitbox/code/leetcode/go/src/kth_largest_element_in_an_array/kth_largest_element_in_an_array.go
def _find_scale_by_percentile(arr, percentile=0.99999):
    assert isinstance(arr, np.ndarray)
    x = np.abs(arr)
    max_k = int(x.size * percentile)
    return np.partition(x, max_k)[max_k]

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-08-23 22:25
import numpy as np


def kl_divergence(p, q):
    return np.sum(p * np.log(p / q))


def check(P, target_bins):
    Q = np.array([np.average(x) for x in np.split(P, target_bins)])
    Q = np.repeat(Q, len(P) // target_bins)

    return kl_divergence(P / np.sum(P), Q / np.sum(Q))


if __name__ == "__main__":
    print(check(np.array([1, 2, 3, 4, 5, 6, 7, 8]), 2))
    print(check(np.array([1, 3, 5, 7, 2, 4, 6, 8]), 2))
    print(check(np.array([2, 2, 2, 2, 6, 6, 6, 6]), 2))

def _global_scale(sq_call):
    cfg = quantize.current_qconfig()
    return cfg.global_scale

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-09-22 16:20

import tvm
from tvm import te
import numpy as np
from tvm import relay

x_data = np.array((1, 153, 2, 178)).reshape((1, 4))
y_data = np.array((204, 178, 1, 8)).reshape((1, 4))

x_scale = y_scale = z_scale = 0.00784314
x_zero_point = y_zero_point = z_zero_point = 127


def dequant(data, scale, zp):
    return scale * (np.asarray(data) - zp)


def quant(data, scale, zp):
    z = np.around(data / scale + zp)
    q_min = np.iinfo(np.uint8).min
    q_max = np.iinfo(np.uint8).max
    return np.clip(z, q_min, q_max)


def mul_manually():
    return quant(
        dequant(x_data, x_scale, x_zero_point) * dequant(y_data, y_scale, y_zero_point),
        z_scale,
        z_zero_point,
    )


if __name__ == "__main__":
    x = relay.var("x", shape=(1, 4), dtype="uint8")
    y = relay.var("y", shape=(1, 4), dtype="uint8")
    z = relay.qnn.op.mul(
        lhs=x,
        rhs=y,
        lhs_scale=relay.const(x_scale, "float32"),
        lhs_zero_point=relay.const(x_zero_point, "int32"),
        rhs_scale=relay.const(y_scale, "float32"),
        rhs_zero_point=relay.const(y_zero_point, "int32"),
        output_scale=relay.const(z_scale, "float32"),
        output_zero_point=relay.const(z_zero_point, "int32"),
    )

    func = relay.Function([x, y], z)
    mod = tvm.IRModule.from_expr(func)
    print("----------before qnn transform----------")
    print(mod)
    mod = relay.transform.InferType()(mod)
    mod = relay.qnn.transform.CanonicalizeOps()(mod)
    print("----------after qnn transform----------")
    print(mod)
    func = mod["main"]

    intrp = relay.create_executor("graph", device=tvm.cpu(0), target="llvm")
    op_res = intrp.evaluate(func)(x_data, y_data)

    print("----------")
    print(op_res.numpy())
    print("----------")
    golden = mul_manually()
    print(golden.astype("uint8"))

QNN_REGISTER_BINARY_OP("mul")
    .describe("Elementwise mul with with broadcasting for quantized tensors.")
    .set_support_level(11)
    .set_attr<FTVMLegalize>("FTVMQnnCanonicalize", QnnMulCanonicalize);

Pass Legalize() {
  Array<Pass> pass_seqs;
  pass_seqs.push_back(relay::transform::Legalize("FTVMQnnLegalize"));
  pass_seqs.push_back(relay::transform::Legalize("FTVMQnnCanonicalize"));
  relay::transform::Pass seq = relay::transform::Sequential(pass_seqs);
  return seq;
}

Array<Pass> GetPassPrefix(const Map<tvm::Integer, tvm::Target>& targets, bool is_vm) {
  Array<Pass> pass_seqs;
  Array<runtime::String> entry_functions{"main"};
  // ..
  // Run all dialect legalization passes.
  pass_seqs.push_back(relay::qnn::transform::Legalize());
  // ...
}