import torch import torch._C._te as te import time import numpy as np import pandas as pd import matplotlib.pyplot as plt import seaborn as sns import argparse class kernel_arena_scope(object): def __enter__(self): self.scope = te.KernelScope() def __exit__(self, typ, val, traceback): self.scope = None unary_ops = [ ("sin", torch.sin), ("cos", torch.cos), ("tan", torch.tan), ("asin", torch.asin), ("acos", torch.acos), ("atan", torch.atan), ("sinh", torch.sinh), ("cosh", torch.cosh), ("tanh", torch.tanh), ("sigmoid", torch.sigmoid), ("exp", torch.exp), ("expm1", torch.expm1), ("expm1", torch.expm1), ("abs", torch.abs), ("log", torch.log), ("fast_log", torch.log), ("log2", torch.log2), ("log10", torch.log10), ("log1p", torch.log1p), ("erf", torch.erf), ("erfc", torch.erfc), ("sqrt", torch.sqrt), ("rsqrt", torch.rsqrt), ("ceil", torch.ceil), ("floor", torch.floor), ("round", torch.round), ("trunc", torch.trunc), ("lgamma", torch.lgamma), # ("frac", torch.frac), # seems unimplemented # ("isnan", torch.isnan), # no out variant ] def gen_unary_nnc_fun(nnc_name): def nnc_fun(A, B): def compute(i, j): return getattr(A.load([i, j]), nnc_name)() return compute return nnc_fun def gen_unary_torch_fun(torch_op): def torch_fun(a, b, out): def fun(): return torch_op(a, out=out) return fun return torch_fun def gen_binary_nnc_fun(fn): def nnc_fun(A, B): def compute(i, j): return fn(A.load([i, j]), B.load([i, j])) return compute return nnc_fun def gen_binary_torch_fun(fn): def pt_fun(a, b, out): def fun(): return fn(a, b, out=out) return fun return pt_fun def gen_int_comparison_tensors(N, M): return (torch.randint(0, 3, (N, M)), torch.randint(0, 3, (N, M)), torch.empty((N, M), dtype=torch.bool)) def gen_float_comparison_tensors(N, M): return (torch.rand(N, M), torch.rand(N, M), torch.empty((N, M), dtype=torch.bool)) te_bool = te.Dtype.Bool binary_ops = [ ('add', (lambda a, b: a + b), torch.add), ('mul', (lambda a, b: a * b), torch.mul), ('sub', (lambda a, b: a - b), torch.sub), ('div', (lambda a, b: a / b), torch.div), ('eq', (lambda a, b: te.Cast.make(te_bool, a == b)), torch.eq, gen_int_comparison_tensors), ('gt', (lambda a, b: te.Cast.make(te_bool, a > b)), torch.gt, gen_float_comparison_tensors), ('lt', (lambda a, b: te.Cast.make(te_bool, a < b)), torch.lt, gen_float_comparison_tensors), ('gte', (lambda a, b: te.Cast.make(te_bool, a >= b)), torch.greater_equal, gen_float_comparison_tensors), ('lte', (lambda a, b: te.Cast.make(te_bool, a <= b)), torch.less_equal, gen_float_comparison_tensors), # ('neq', (lambda a, b: a != b), None)), # no one-op equivalent # ('&', (lambda a, b: a & b), torch.bitwise_and), # requires more work to test ] def nnc_relu(A, B): def f(i, j): return torch._C._te.ifThenElse(A.load([i, j]) < torch._C._te.ExprHandle.float(0), torch._C._te.ExprHandle.float(0), A.load([i, j])) return f def pt_relu(a, b, c): return torch.relu(a) custom_ops = [ ('relu', nnc_relu, pt_relu), # ('nnc_mul_relu', nnc_mul_relu, pt_mul_relu) # ('manual_sigmoid', nnc_manual_sigmoid, lambda a, b, c: torch.sigmoid(a, out=c)) ] def gen_custom_torch_fun(fn): def pt_fun(a, b, out): def fun(): return fn(a, b, out) return fun return pt_fun def normalize_benchmarks(ops): return [i + (None,) if len(i) == 3 else i for i in ops] names = [] nnc_fns = [] pt_fns = [] shape_fns = [] for nnc_name, pt_op in unary_ops: names.append(nnc_name) nnc_fns.append(gen_unary_nnc_fun(nnc_name)) pt_fns.append(gen_unary_torch_fun(pt_op)) shape_fns.append(None) for name, lmbda, pt_fn, shape_fn in normalize_benchmarks(binary_ops): names.append(name) nnc_fns.append(gen_binary_nnc_fun(lmbda)) pt_fns.append(gen_binary_torch_fun(pt_fn)) shape_fns.append(shape_fn) for name, lmbda, pt_fn, shape_fn in normalize_benchmarks(custom_ops): names.append(name) nnc_fns.append(lmbda) pt_fns.append(gen_custom_torch_fun(pt_fn)) shape_fns.append(shape_fn) benchmarks = list(zip(names, nnc_fns, pt_fns, shape_fns)) def run_benchmarks(benchmarks, sizes): df = pd.DataFrame(columns=['name', 'N', 'M', 'nnc_time', 'torch_time', 'ratio']) with torch.no_grad(): for name, nnc_fun, torch_fun, shape_fn in benchmarks: for N, M in sizes: iters = int(1e6 / (N + M)) with kernel_arena_scope(): if shape_fn is None: tA = torch.rand(M, N).clamp(0.01, 0.99) tB = torch.rand(M, N).clamp(0.01, 0.99) tX = torch.empty(M, N) tR = torch.empty(M, N) else: tA, tB, tX = shape_fn(M, N) tR = tX.clone() def get_nnc_type(dtype): if dtype == torch.float: return torch._C._te.Dtype.Float elif dtype == torch.long: return torch._C._te.Dtype.Long dtype = get_nnc_type(tA.dtype) dM = torch._C._te.ExprHandle.int(M) dN = torch._C._te.ExprHandle.int(N) A = torch._C._te.Placeholder('A', dtype, [dM, dN]) B = torch._C._te.Placeholder('B', dtype, [dM, dN]) dim_args = [torch._C._te.DimArg(*args) for args in [(dM, 'm'), (dN, 'n')]] compute = nnc_fun(A, B) X = torch._C._te.Compute('X', dim_args, compute) loopnest = torch._C._te.LoopNest([X]) loopnest.prepare_for_codegen() stmt = torch._C._te.simplify(loopnest.root_stmt()) cg = torch._C._te.construct_codegen('llvm', stmt, [torch._C._te.BufferArg(x) for x in [A, B, X]]) # warmup for _ in range(10): cg.call([tA, tB, tX]) start = time.time() for it in range(iters): cg.call([tA, tB, tX]) time1 = time.time() - start fn = torch_fun(tA, tB, tR) # warmup for _ in range(10): tR = fn() start = time.time() for it in range(iters): tR = fn() time2 = time.time() - start df = df.append({'name': name, 'N': N, 'M': M, 'nnc_time': time1, 'torch_time': time2, 'ratio': time2 / time1}, ignore_index=True) print(name, N, M) print(time2 / time1, time1, time2) print() def check_correctness(a, b): if not np.allclose(a, b): print(name) assert(np.allclose(a, b)) check_correctness(tX, tR) return df def dump_plot(df, sizes): keys = [] vals = [] indexed = df[df['N'] == df['M']] for index, row in indexed.iterrows(): keys.append(row['name']) vals.append(row['ratio']) keys = keys[::len(sizes)] sns.set(rc={'figure.figsize' : (5.0, len(keys) * 0.5)}) cmap = sns.diverging_palette(10, 120, n=9, as_cmap=True) np_vals = np.array([vals]).reshape(-1, len(sizes)) g = sns.heatmap(np_vals, annot=True, cmap=cmap, center=1.0, yticklabels=True) plt.yticks(rotation=0) plt.title('PyTorch performance divided by NNC performance (single core)') plt.xlabel('Size of NxN matrix') plt.ylabel('Operation') g.set_yticklabels(keys) g.set_xticklabels(sizes) plt.savefig('nnc.png') if __name__ == "__main__": parser = argparse.ArgumentParser(description='Runs NNC microbenchmarks') parser.add_argument('--multi_threaded', action='store_true', help='Run with more than one thread') args = parser.parse_args() if not args.multi_threaded: torch.set_num_threads(1) sizes = [1, 4, 16, 64, 256, 1024] df = run_benchmarks(benchmarks, [(i, i) for i in sizes]) dump_plot(df, sizes)