| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| import timeit |
| import itertools |
| import argparse |
| import os |
| |
| class OpArgMngr(object): |
| """Operator argument manager for storing operator workloads.""" |
| args = {} |
| |
| @staticmethod |
| def add_workload(funcname, *args, **kwargs): |
| if "_specifier" not in kwargs: |
| _specifier = funcname |
| else: |
| _specifier = kwargs["_specififer"] |
| del kwargs["_specififer"] |
| if _specifier in OpArgMngr.args: |
| raise ValueError(f"duplicate {_specifier}") |
| OpArgMngr.args[_specifier] = {'args': args, 'kwargs': kwargs, 'funcname': funcname} |
| |
| |
| def generate_workloads(): |
| array_pool = {} |
| shapes = [] |
| for ndim in range(4): |
| shapes.extend(list(itertools.product(range(4), repeat=ndim))) |
| for shape in shapes: |
| name = 'x'.join(str(i) for i in shape) |
| if name in array_pool: |
| raise ValueError(f"duplicate array {name}") |
| array_pool[name] = dnp.ones(shape) |
| return array_pool |
| |
| |
| def prepare_workloads(): |
| pool = generate_workloads() |
| OpArgMngr.add_workload("zeros", (2, 2)) |
| OpArgMngr.add_workload("full", (2, 2), 10) |
| OpArgMngr.add_workload("identity", 3) |
| OpArgMngr.add_workload("ones", (2, 2)) |
| OpArgMngr.add_workload("einsum", "ii", pool['2x2'], optimize=False) |
| OpArgMngr.add_workload("unique", pool['1'], return_index=True, return_inverse=True, return_counts=True, axis=-1) |
| OpArgMngr.add_workload("dstack", (pool['2x1'], pool['2x1'], pool['2x1'], pool['2x1'])) |
| OpArgMngr.add_workload("polyval", dnp.arange(10), pool['2x2']) |
| OpArgMngr.add_workload("ediff1d", pool['2x2'], pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("nan_to_num", pool['2x2']) |
| OpArgMngr.add_workload("tri", 2, 3, 4) |
| OpArgMngr.add_workload("tensordot", pool['2x2'], pool['2x2'], ((1, 0), (0, 1))) |
| OpArgMngr.add_workload("cumsum", pool['3x2'], axis=0, out=pool['3x2']) |
| OpArgMngr.add_workload("random.shuffle", pool['3']) |
| OpArgMngr.add_workload("equal", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("not_equal", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("less", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("greater_equal", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("less_equal", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("maximum", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("minimum", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("sum", pool['2x2'], axis=0, keepdims=True, out=pool['1x2']) |
| OpArgMngr.add_workload("std", pool['2x2'], axis=0, ddof=0, keepdims=True, out=pool['1x2']) |
| OpArgMngr.add_workload("var", pool['2x2'], axis=0, ddof=1, keepdims=True, out=pool['1x2']) |
| OpArgMngr.add_workload("average", pool['2x2'], weights=pool['2'], axis=1, returned=True) |
| OpArgMngr.add_workload("histogram", pool['2x2'], bins=10, range=(0.0, 10.0)) |
| OpArgMngr.add_workload("add", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("cross", pool['2'], pool['2']) |
| OpArgMngr.add_workload("linalg.eig", pool['3x3']) |
| OpArgMngr.add_workload("linalg.eigh", pool['3x3']) |
| OpArgMngr.add_workload("linalg.det", pool['3x3']) |
| OpArgMngr.add_workload("linalg.slogdet", pool['3x3']) |
| OpArgMngr.add_workload("linalg.matrix_rank", pool['3x3'], pool['1'], hermitian=False) |
| OpArgMngr.add_workload("linalg.svd", pool['3x3']) |
| OpArgMngr.add_workload("linalg.cholesky", pool['1x1']) |
| OpArgMngr.add_workload("linalg.qr", pool['3x3']) |
| OpArgMngr.add_workload("linalg.lstsq", pool['2x1'], pool['2'], rcond=None) |
| OpArgMngr.add_workload("linalg.eigvals", pool['1x1']) |
| OpArgMngr.add_workload("linalg.eigvalsh", pool['1x1'], UPLO='L') |
| OpArgMngr.add_workload("linalg.inv", pool['1x1']) |
| OpArgMngr.add_workload("linalg.pinv", pool['2x3x3'], pool['1'], hermitian=False) |
| OpArgMngr.add_workload("linalg.solve", pool['1x1'], pool['1']) |
| OpArgMngr.add_workload("linalg.tensorinv", pool['1x1'], ind=2) |
| OpArgMngr.add_workload("linalg.norm", pool['3x3']) |
| OpArgMngr.add_workload("linalg.tensorsolve", pool['1x1x1'], pool['1x1x1'], (2, 0, 1)) |
| OpArgMngr.add_workload("tile", pool['2x2'], 1) |
| OpArgMngr.add_workload("trace", pool['2x2']) |
| OpArgMngr.add_workload("transpose", pool['2x2']) |
| OpArgMngr.add_workload("split", pool['3x3'], (0, 1, 2), axis=1) |
| OpArgMngr.add_workload("vstack", (pool['3x3'], pool['3x3'], pool['3x3'])) |
| OpArgMngr.add_workload("argmax", pool['3x2'], axis=-1) |
| OpArgMngr.add_workload("argmin", pool['3x2'], axis=-1) |
| OpArgMngr.add_workload("atleast_1d", pool['2'], pool['2x2']) |
| OpArgMngr.add_workload("atleast_2d", pool['2'], pool['2x2']) |
| OpArgMngr.add_workload("atleast_3d", pool['2'], pool['2x2']) |
| OpArgMngr.add_workload("argsort", pool['3x2'], axis=-1) |
| OpArgMngr.add_workload("sort", pool['3x2'], axis=-1) |
| OpArgMngr.add_workload("indices", dimensions=(1, 2, 3)) |
| OpArgMngr.add_workload("subtract", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("multiply", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("mod", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("remainder", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("divide", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("true_divide", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("power", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("lcm", pool['2x2'].astype('int32'), pool['2x2'].astype('int32')) |
| OpArgMngr.add_workload("diff", pool['2x2'], n=1, axis=-1) |
| OpArgMngr.add_workload("inner", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("random.multinomial", n=2, pvals=[1/6.]*6, size=(2,2)) |
| OpArgMngr.add_workload("random.rand", 3, 2) |
| OpArgMngr.add_workload("random.randn", 2, 2) |
| OpArgMngr.add_workload("nonzero", pool['2x2']) |
| OpArgMngr.add_workload("tril", pool['2x2'], k=0) |
| OpArgMngr.add_workload("random.choice", pool['2'], size=(2, 2)) |
| OpArgMngr.add_workload("take", pool['2'], dnp.array([1,0], dtype='int64')) |
| OpArgMngr.add_workload("clip", pool['2x2'], 0, 1) |
| OpArgMngr.add_workload("expand_dims", pool['2x2'], axis=0) |
| OpArgMngr.add_workload("broadcast_to", pool['2x2'], (2, 2, 2)) |
| OpArgMngr.add_workload("full_like", pool['2x2'], 2) |
| OpArgMngr.add_workload("zeros_like", pool['2x2']) |
| OpArgMngr.add_workload("ones_like", pool['2x2']) |
| OpArgMngr.add_workload("bitwise_and", pool['2x2'].astype(int), pool['2x2'].astype(int)) |
| OpArgMngr.add_workload("bitwise_xor", pool['2x2'].astype(int), pool['2x2'].astype(int)) |
| OpArgMngr.add_workload("bitwise_or", pool['2x2'].astype(int), pool['2x2'].astype(int)) |
| OpArgMngr.add_workload("copysign", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("arctan2", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("hypot", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("ldexp", pool['2x2'].astype(int), pool['2x2'].astype(int)) |
| OpArgMngr.add_workload("logical_and", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("logical_or", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("logical_xor", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("random.uniform", low=0, high=1, size=1) |
| OpArgMngr.add_workload("random.exponential", scale=2, size=(2,2)) |
| OpArgMngr.add_workload("random.rayleigh", scale=2, size=(2,2)) |
| OpArgMngr.add_workload("random.weibull", a=2, size=(2,2)) |
| OpArgMngr.add_workload("random.pareto", a=2, size=(2,2)) |
| OpArgMngr.add_workload("random.power", a=2, size=(2,2)) |
| OpArgMngr.add_workload("random.logistic", loc=2, scale=2, size=(2,2)) |
| OpArgMngr.add_workload("random.gumbel", loc=2, scale=2, size=(2,2)) |
| OpArgMngr.add_workload("where", pool['2x3'], pool['2x3'], pool['2x1']) |
| OpArgMngr.add_workload("may_share_memory", pool['2x3'][:0], pool['2x3'][:1]) |
| OpArgMngr.add_workload('squeeze', pool['2x2'], axis=None) |
| OpArgMngr.add_workload("pad", pool['2x2'], pad_width=((1,2),(1,2)), mode="constant") |
| OpArgMngr.add_workload("prod", pool['2x2'], axis=1, dtype="float64", keepdims=False) |
| OpArgMngr.add_workload("around", pool['2x2'], decimals=0) |
| OpArgMngr.add_workload("round", pool['2x2'], decimals=1) |
| OpArgMngr.add_workload("repeat", pool['2x2'], repeats=1, axis=None) |
| OpArgMngr.add_workload("diagflat", pool['2x2'], k=1) |
| OpArgMngr.add_workload("diag", pool['2x2'], k=1) |
| OpArgMngr.add_workload("diagonal", pool['2x2x2'], offset=-1, axis1=0, axis2=1) |
| OpArgMngr.add_workload("diag_indices_from", pool['2x2']) |
| OpArgMngr.add_workload("bincount", dnp.arange(3, dtype=int), pool['3'], minlength=4) |
| OpArgMngr.add_workload("percentile", pool['2x2x2'], 80, axis=0, out=pool['2x2'],\ |
| interpolation='midpoint') |
| OpArgMngr.add_workload("quantile", pool['2x2x2'], 0.8, axis=0, out=pool['2x2'],\ |
| interpolation='midpoint') |
| OpArgMngr.add_workload("all", pool['2x2x2'], axis=(0, 1),\ |
| out=dnp.array([False, False], dtype=bool), keepdims=False) |
| OpArgMngr.add_workload("any", pool['2x2x2'], axis=(0, 1),\ |
| out=dnp.array([False, False], dtype=bool), keepdims=False) |
| OpArgMngr.add_workload("roll", pool["2x2"], 1, axis=0) |
| OpArgMngr.add_workload("rot90", pool["2x2"], 2) |
| OpArgMngr.add_workload("column_stack", (pool['3x3'], pool['3x3'], pool['3x3'])) |
| OpArgMngr.add_workload("hstack", (pool['3x3'], pool['3x3'], pool['3x3'])) |
| OpArgMngr.add_workload("triu", pool['3x3']) |
| OpArgMngr.add_workload("array_split", pool['2x2'], 2, axis=1) |
| OpArgMngr.add_workload("vsplit", pool['2x2'], 2) |
| OpArgMngr.add_workload("hsplit", pool['2x2'], 2) |
| OpArgMngr.add_workload("dsplit", pool['2x2x2'], 2) |
| OpArgMngr.add_workload("arange", 10) |
| OpArgMngr.add_workload("concatenate", (pool['1x2'], pool['1x2'], pool['1x2']), axis=0) |
| OpArgMngr.add_workload("append", pool['2x2'], pool['1x2'], axis=0) |
| OpArgMngr.add_workload("insert", pool['3x2'], 1, pool['1x1'], axis=0) |
| OpArgMngr.add_workload("delete", pool['3x2'], 1, axis=0) |
| OpArgMngr.add_workload("blackman", 12) |
| OpArgMngr.add_workload("eye", 5) |
| OpArgMngr.add_workload("hamming", 12) |
| OpArgMngr.add_workload("hanning", 12) |
| OpArgMngr.add_workload("linspace", 0, 10, 8, endpoint=False) |
| OpArgMngr.add_workload("logspace", 2.0, 3.0, num=4, base=2.0, dtype=onp.float32) |
| OpArgMngr.add_workload("matmul", pool['2x2'], pool['2x2']) |
| OpArgMngr.add_workload("mean", pool['2x2'], axis=0, keepdims=True) |
| OpArgMngr.add_workload("random.gamma", 1, size=(2, 3)) |
| OpArgMngr.add_workload("random.normal", 1, size=(2, 3)) |
| OpArgMngr.add_workload("max", pool["2x2"], axis=0, out=pool['2'], keepdims=False) |
| OpArgMngr.add_workload("min", pool["2x2"], axis=0, out=pool['2'], keepdims=False) |
| OpArgMngr.add_workload("amax", pool["2x2"], axis=1, out=pool['2'], keepdims=False) |
| OpArgMngr.add_workload("amin", pool["2x2"], axis=1, out=pool['2'], keepdims=False) |
| |
| unary_ops = ['negative', 'reciprocal', 'abs', 'sign', 'rint', 'ceil', 'floor', |
| 'bitwise_not', 'trunc', 'fix', 'square', 'sqrt', 'cbrt', 'exp', |
| 'log', 'log10', 'log2', 'log1p', 'expm1', 'logical_not', 'isnan', |
| 'isinf', 'isposinf', 'isneginf', 'isfinite', 'sin', 'cos', 'tan', |
| 'arcsin', 'arccos', 'arctan', 'degrees', 'radians', 'sinh', 'cosh', |
| 'tanh', 'arcsinh', 'arccosh', 'arctanh'] # 'rad2deg', 'deg2rad' cannot run without tvm |
| for unary_op in unary_ops: |
| if unary_op == "bitwise_not": |
| OpArgMngr.add_workload(unary_op, dnp.ones((2, 2), dtype=int)) |
| else: |
| OpArgMngr.add_workload(unary_op, pool['2x2']) |
| |
| |
| def benchmark_helper(f, *args, **kwargs): |
| number = 10000 |
| return timeit.timeit(lambda: f(*args, **kwargs), number=number) / number |
| |
| |
| def get_op(module, funcname): |
| funcname = funcname.split(".") |
| for fname in funcname: |
| module = getattr(module, fname) |
| return module |
| |
| |
| def run_benchmark(packages): |
| results = {} |
| for (k, v) in OpArgMngr.args.items(): |
| result = {} |
| for (name, package) in packages.items(): |
| print(f'{name}.{k} running...') |
| op = get_op(package["module"], v["funcname"]) |
| args = [package["data"](arg) for arg in v["args"]] |
| kwargs = {k: package["data"](v) for (k, v) in v["kwargs"].items()} |
| benchmark = benchmark_helper(op, *args, **kwargs) |
| result[name] = benchmark |
| results[k] = result |
| return results |
| |
| |
| def show_results(results): |
| print(f'{"name":>24}{"package":>24}{"time(us)":>24}') |
| for (specifier, d) in results.items(): |
| for (k, v) in d.items(): |
| print(f"{specifier:>24}{k:>24}{v * 10 ** 6:>24}") |
| |
| |
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument('ffi_type') |
| parsed = parser.parse_args() |
| if parsed.ffi_type == "cython": |
| os.environ['MXNET_ENABLE_CYTHON'] = '1' |
| os.environ['MXNET_ENFORCE_CYTHON'] = '1' |
| elif parsed.ffi_type == "ctypes": |
| os.environ['MXNET_ENABLE_CYTHON'] = '0' |
| else: |
| raise ValueError("unknown ffi_type {}",format(parsed.ffi_type)) |
| os.environ["MXNET_ENGINE_TYPE"] = "NaiveEngine" |
| import mxnet as mx |
| import numpy as onp |
| from mxnet import np as dnp |
| |
| mx.npx.set_np(dtype=False) |
| packages = { |
| "onp": { |
| "module": onp, |
| "data": lambda arr: arr.asnumpy() if isinstance(arr, dnp.ndarray) else arr |
| }, |
| "dnp": { |
| "module": dnp, |
| "data": lambda arr: arr |
| } |
| } |
| prepare_workloads() |
| results = run_benchmark(packages) |
| show_results(results) |
