def validate(shape, value, dtype): def before_left(x, elem_op, full): return elem_op(full, x) def after_left(x, elem_op, value): return elem_op(relay.const(value, dtype), x) def before_right(x, elem_op, full): return elem_op(x, full) def after_right(x, elem_op, value): return elem_op(x, relay.const(value, dtype)) x = relay.var("x", shape=shape, dtype=dtype) elem_ops = [relay.add, relay.multiply, relay.subtract, relay.divide] full_ops = [] if value == 0: full_ops.append(relay.zeros(shape, dtype)) full_ops.append(relay.zeros_like(x)) if value == 1: full_ops.append(relay.ones(shape, dtype)) full_ops.append(relay.ones_like(x)) else: full_ops.append(relay.full(relay.const(value, dtype), shape)) full_ops.append(relay.full_like(x, relay.const(value, dtype))) for op in elem_ops: for full in full_ops: z = before_left(x, op, full) zz = run_opt_pass(z, transform.SimplifyExpr()) after = run_opt_pass(after_left(x, op, value), transform.InferType()) assert tvm.ir.structural_equal(zz, after) z = before_right(x, op, full) zz = run_opt_pass(z, transform.SimplifyExpr()) after = run_opt_pass(after_right(x, op, value), transform.InferType()) assert tvm.ir.structural_equal(zz, after) # Test the case in which x is broadcast to full's shape full_ops = [] if value == 0: full_ops.append(relay.zeros(shape * 2, dtype)) if value == 1: full_ops.append(relay.ones(shape * 2, dtype)) else: full_ops.append(relay.full(relay.const(value, dtype), shape * 2)) for op in elem_ops: for full in full_ops: z = before_left(x, op, full) zz = run_opt_pass(z, transform.SimplifyExpr()) after = run_opt_pass(before_left(x, op, full), transform.InferType()) assert tvm.ir.structural_equal(zz, after) z = before_right(x, op, full) zz = run_opt_pass(z, transform.SimplifyExpr()) after = run_opt_pass(before_right(x, op, full), transform.InferType()) assert tvm.ir.structural_equal(zz, after)
def check(x, y=None, do_nothing=False): expected = run_infer_type(x) if do_nothing: actual = run_opt_pass(x, transform.SimplifyExpr()) assert tvm.ir.structural_equal(actual, expected) else: assert y is not None actual = run_opt_pass(y, transform.SimplifyExpr()) assert tvm.ir.structural_equal(actual, expected)
def partition_for_dnnl(mod, params=None): """Partition the graph greedily offloading supported operators to DNNL. Parameters ---------- mod : Module The module to run passes on. params : Optional[Dict[str, NDArray]] Constant input parameters. Returns ------- mod : Module Annotated and partitioned module. """ if params: mod["main"] = bind_params_by_name(mod["main"], params) seq = tvm.transform.Sequential([ transform.CanonicalizeOps(), transform.InferType(), transform.SimplifyInference(), transform.FoldConstant(), transform.FoldScaleAxis(), # fold consecutive add ops to simplify pattern `conv2d-bias_add-bn-relu` transform.SimplifyExpr(), transform.FoldConstant(), transform.MergeComposite(pattern_table()), transform.AnnotateTarget("dnnl"), transform.MergeCompilerRegions(), transform.PartitionGraph(), ]) with tvm.transform.PassContext(opt_level=3): mod = seq(mod) return mod
def test_simplify_transpose(): # Test a series of transpose and layout_transform ops def before1(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.transpose(x, axes=[0, 2, 3, 1]) # To NHWC y = relay.layout_transform(y, "NHWC", "HWCN") # To HWCN y = relay.transpose(y, axes=[3, 0, 1, 2]) # To NHWC return relay.Function([x], y) def expected1(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.transpose(x, axes=[0, 2, 3, 1]) # To NHWC return relay.Function([x], y) # Test that all transpose ops can be cancelled def before2(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) y = relay.transpose(y, axes=[0, 2, 3, 1]) # To NHWC y = relay.transpose(y, axes=[1, 2, 3, 0]) # To HWCN y = relay.transpose(y, axes=[3, 2, 0, 1]) # To NCHW return relay.Function([x], y) def expected2(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) return relay.Function([x], y) # Test default axis (reverse) and negative axis def before3(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) y = relay.transpose(y) # Reverse y = relay.transpose(y) # Reverse y = relay.transpose(y, axes=[0, 2, -1, 1]) y = relay.transpose(y) # Reverse y = relay.transpose(y) # Reverse return relay.Function([x], y) def expected3(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) y = relay.transpose(y, axes=[0, 2, 3, 1]) return relay.Function([x], y) for before, expected in [ [before1(), expected1()], [before2(), expected2()], [before3(), expected3()], ]: after = run_opt_pass(before, transform.SimplifyExpr()) expected = run_opt_pass(expected, transform.InferType()) assert tvm.ir.structural_equal( after, expected), "\nafter: {} \nexpected: {}".format(after, expected)
def test_simplify_consecutive_add(): shape = (32, 1, 1) c_data = np.empty(shape).astype("float32") c1 = relay.const(c_data) c2 = relay.const(c_data) def before_const_right(): x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32") w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32") y = relay.nn.conv2d(x, w, padding=(1, 1)) y = relay.add(y, c1) y = relay.add(y, c2) y = relay.nn.relu(y) return relay.Function([x, w], y) def before_const_left(): x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32") w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32") y = relay.nn.conv2d(x, w, padding=(1, 1)) y = relay.add(c1, y) y = relay.add(c2, y) y = relay.nn.relu(y) return relay.Function([x, w], y) def expected(): x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32") w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32") y = relay.nn.conv2d(x, w, padding=(1, 1)) c3 = relay.add(c1, c2) y = relay.add(y, c3) y = relay.nn.relu(y) return relay.Function([x, w], y) zr = before_const_right() zl = before_const_left() zzr = run_opt_pass(zr, transform.SimplifyExpr()) zzl = run_opt_pass(zl, transform.SimplifyExpr()) after = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal(zzr, after) assert tvm.ir.structural_equal(zzl, after)
def test_simplify_reshape(): def before(): x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32") w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32") y = relay.nn.conv2d(x, w, padding=(1, 1)) y = relay.reshape(y, newshape=(1, 16, -1)) y = relay.reshape(y, newshape=(4, 8, -1, 16)) y = relay.reverse_reshape(y, newshape=(32, 0, -1)) return relay.Function([x, w], y) def expected(): x = relay.var("x", shape=(1, 16, 16, 16), dtype="float32") w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32") y = relay.nn.conv2d(x, w, padding=(1, 1)) y = relay.reshape(y, newshape=(32, 16, 16)) return relay.Function([x, w], y) def symbolic(): b = tvm.te.size_var("b") x = relay.var("x", shape=(b, 16, 16, 16), dtype="float32") w = relay.var("w", shape=(32, 16, 3, 3), dtype="float32") y = relay.nn.conv2d(x, w, padding=(1, 1)) y = relay.reshape(y, newshape=(1, 16, -1)) y = relay.reshape(y, newshape=(4, 8, -1, 16)) y = relay.reverse_reshape(y, newshape=(32, 0, -1)) return relay.Function([x, w], y) z = before() zz = run_opt_pass(z, transform.SimplifyExpr()) after = run_opt_pass(expected(), transform.InferType()) assert tvm.ir.structural_equal(zz, after) z = symbolic() zz = run_opt_pass(z, transform.SimplifyExpr()) after = run_opt_pass(symbolic(), transform.InferType()) assert tvm.ir.structural_equal(zz, after)
def validate(ndim, pad_width, pad_value, pad_mode, orig_padding, layout): if layout[1] == "C": shape = [1, 3] + [10] * ndim wshape = [8, 3] + [3] * ndim elif layout[-1] == "C": shape = [1] + [10] * ndim + [3] wshape = [8] + [3] * ndim + [3] else: raise ValueError("This test only supports NC* and N*C") x = relay.var("x", shape=shape, dtype="float32") w = relay.var("w", shape=wshape, dtype="float32") pad = relay.nn.pad(x, pad_width, pad_value, pad_mode) if layout[1] == "C": conv = convs[ndim - 1](pad, w, padding=orig_padding) else: conv = convs[ndim - 1](pad, w, padding=orig_padding, data_layout=layout, kernel_layout="DHWIO"[3 - ndim:]) if pad_mode == "constant" and pad_value == 0: new_padding = [] for j in range(2): for i in range(len(pad_width)): if layout[i] in ["D", "H", "W"]: new_padding.append(pad_width[i][j]) for i in range(len(new_padding)): new_padding[i] += orig_padding[i] if layout[1] == "C": after = convs[ndim - 1](x, w, padding=new_padding) else: after = convs[ndim - 1](x, w, padding=new_padding, data_layout=layout, kernel_layout="DHWIO"[3 - ndim:]) else: after = conv zz = run_opt_pass(conv, transform.SimplifyExpr()) expected = run_opt_pass(after, transform.InferType()) assert tvm.ir.structural_equal(zz, expected) mod1 = tvm.IRModule.from_expr(conv) mod2 = tvm.IRModule.from_expr(zz) with tvm.transform.PassContext(disabled_pass="******"): ex1 = relay.create_executor("vm", mod=mod1, ctx=tvm.cpu(), target="llvm") ex2 = relay.create_executor("vm", mod=mod2, ctx=tvm.cpu(), target="llvm") x_np = np.random.rand(*shape).astype("float32") w_np = np.random.rand(*wshape).astype("float32") result1 = ex1.evaluate()(x_np, w_np) result2 = ex2.evaluate()(x_np, w_np) tvm.testing.assert_allclose(result1.asnumpy(), result2.asnumpy())
def test_simplify_transpose(): # Test a series of transpose and layout_transform ops def before1(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.transpose(x, axes=[0, 2, 3, 1]) # To NHWC y = relay.layout_transform(y, "NHWC", "HWCN") # To HWCN y = relay.transpose(y, axes=[3, 0, 1, 2]) # To NHWC return relay.Function([x], y) def expected1(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.transpose(x, axes=[0, 2, 3, 1]) # To NHWC return relay.Function([x], y) # Test that all transpose ops can be cancelled def before2(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) y = relay.transpose(y, axes=[0, 2, 3, 1]) # To NHWC y = relay.transpose(y, axes=[1, 2, 3, 0]) # To HWCN y = relay.transpose(y, axes=[3, 2, 0, 1]) # To NCHW return relay.Function([x], y) def expected2(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) return relay.Function([x], y) # Test default axis (reverse) and negative axis def before3(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) y = relay.transpose(y) # Reverse y = relay.transpose(y) # Reverse y = relay.transpose(y, axes=[0, 2, -1, 1]) y = relay.transpose(y) # Reverse y = relay.transpose(y) # Reverse return relay.Function([x], y) def expected3(): x = relay.var("x", shape=(1, 3, 224, 224), dtype="float32") # NCHW y = relay.nn.relu(x) y = relay.transpose(y, axes=[0, 2, 3, 1]) return relay.Function([x], y) # Test a series of transpose and rank changing layout_transform def before4(): """ Simplify transpose->layout_transform and its inverse. Input: NHWC -> NCHW -> NCHW4c -> op -> NCHW4c -> NCHW -> NHWC Simplified: NHWC -> NCHW4c -> op -> NCHW4c -> NHWC """ x = relay.var("x", shape=(1, 56, 56, 128), dtype="float32") y = relay.transpose(x, axes=[0, 3, 1, 2]) y = relay.layout_transform(y, "NCHW", "NCHW4c") y = relay.nn.relu(y) y = relay.layout_transform(y, "NCHW4c", "NCHW") y = relay.transpose(y, axes=[0, 2, 3, 1]) return relay.Function([x], y) def expected4(): x = relay.var("x", shape=(1, 56, 56, 128), dtype="float32") # NHWC y = relay.layout_transform(x, "NHWC", "NCHW4c") # To NCHW4c y = relay.nn.relu(y) y = relay.layout_transform(y, "NCHW4c", "NHWC") # To NHWC return relay.Function([x], y) def before5(): """ Simplify layout_transform->layout_transform and its inverse. Input: NHWC -> NCHW -> NCHW4c -> op -> NCHW4c -> NCHW -> NHWC Simplified: NHWC -> NCHW4c -> op -> NCHW4c -> NHWC """ x = relay.var("x", shape=(1, 56, 56, 128), dtype="float32") # NHWC y = relay.layout_transform(x, "NHWC", "NCHW") # To NCHW y = relay.layout_transform(y, "NCHW", "NCHW4c") # To NCHW4c y = relay.nn.relu(y) y = relay.layout_transform(y, "NCHW4c", "NCHW") # To NCHW y = relay.layout_transform(y, "NCHW", "NHWC") # To NHWC return relay.Function([x], y) def expected5(): x = relay.var("x", shape=(1, 56, 56, 128), dtype="float32") # NHWC y = relay.layout_transform(x, "NHWC", "NCHW4c") # To NCHW4c y = relay.nn.relu(y) y = relay.layout_transform(y, "NCHW4c", "NHWC") # To NHWC return relay.Function([x], y) def before6(): """ Remove trivial layout_transform->layout_transform. Input: NCHW -> NHWC -> NCHW -> op Simplified: NHWC -> op """ x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32") y = relay.layout_transform(x, "NCHW", "NHWC") y = relay.layout_transform(y, "NHWC", "NCHW") y = relay.nn.relu(y) return relay.Function([x], y) def expected6(): x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32") y = relay.nn.relu(x) return relay.Function([x], y) def before7(): """ Remove trivial layout_transform->layout_transform. Input: NCHW4c -> NCHW8c -> NCHW4c -> op Simplified: NCHW4c -> op """ x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32") y = relay.layout_transform(x, "NCHW4c", "NCHW8c") y = relay.layout_transform(y, "NCHW8c", "NCHW4c") y = relay.nn.relu(y) return relay.Function([x], y) def expected7(): x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32") y = relay.nn.relu(x) return relay.Function([x], y) def before8(): """ Simplify layout_transform->layout_transform with rank contraction and expansion Input: NCHW4c -> NCHW -> NCHW8c -> op Simplified: NCHW4c -> NCHW8c -> op """ x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32") y = relay.layout_transform(x, "NCHW4c", "NCHW") y = relay.layout_transform(y, "NCHW", "NCHW8c") y = relay.nn.relu(y) return relay.Function([x], y) def expected8(): x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32") y = relay.layout_transform(x, "NCHW4c", "NCHW8c") y = relay.nn.relu(y) return relay.Function([x], y) def before9(): """ Remove trivial layout_transform->layout_transform. Input: NCHW -> NCHW4c -> NCHW -> op Simplified: NCHW -> op """ x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32") y = relay.layout_transform(x, "NCHW", "NCHW4c") y = relay.layout_transform(y, "NCHW4c", "NCHW") y = relay.nn.relu(y) return relay.Function([x], y) def expected9(): x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32") y = relay.nn.relu(x) return relay.Function([x], y) def before10(): """ Simplify layout_transform->layout_transform without rank change to transpose. Input: NCHW -> NHWC -> CHWN -> op Simplified: NCHW -> CHWN -> op """ x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32") y = relay.layout_transform(x, "NCHW", "NHWC") y = relay.layout_transform(y, "NHWC", "CHWN") y = relay.nn.relu(y) return relay.Function([x], y) def expected10(): x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32") y = relay.transpose(x, axes=[1, 2, 3, 0]) y = relay.nn.relu(y) return relay.Function([x], y) for before, expected in [ [before1(), expected1()], [before2(), expected2()], [before3(), expected3()], [before4(), expected4()], [before5(), expected5()], [before6(), expected6()], [before7(), expected7()], [before8(), expected8()], [before9(), expected9()], [before10(), expected10()], ]: after = run_opt_pass(before, transform.SimplifyExpr()) expected = run_opt_pass(expected, transform.InferType()) assert tvm.ir.structural_equal( after, expected), "\nafter: {} \nexpected: {}".format(after, expected)
def partition_for_dnnl(mod, params=None, alter_layout=True): """Partition the graph greedily offloading supported operators to DNNL. Parameters ---------- mod : Module The module to run passes on. params : Optional[Dict[str, NDArray]] Constant input parameters. Returns ------- mod : Module Annotated and partitioned module. """ if params: mod["main"] = bind_params_by_name(mod["main"], params) with TempOpAttr("nn.conv2d", "FTVMLegalize", dnnl.legalize_group_conv): with TempOpAttr("nn.conv2d_transpose", "FTVMLegalize", dnnl.legalize_group_conv): seq = tvm.transform.Sequential([ transform.CanonicalizeOps(), transform.InferType(), transform.SimplifyInference(), transform.FoldConstant(), transform.FoldScaleAxis(), # fold consecutive add ops to simplify pattern `conv2d-bias_add-bn-relu` transform.SimplifyExpr(), transform.FoldConstant(), # alter group conv /conv_transpose layout to `GOIHW` / `GIOHW` transform.Legalize(), transform.FoldConstant(), ]) with tvm.transform.PassContext(opt_level=3): mod = seq(mod) if alter_layout: with TempOpAttr("nn.conv1d", "FTVMAlterOpLayout", dnnl.alter_conv): with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", dnnl.alter_conv): with TempOpAttr("nn.conv3d", "FTVMAlterOpLayout", dnnl.alter_conv): with TempOpAttr("nn.conv2d_transpose", "FTVMAlterOpLayout", dnnl.alter_conv_transpose): with TempOpAttr("nn.conv3d_transpose", "FTVMAlterOpLayout", dnnl.alter_conv_transpose): alter_layout_seq = tvm.transform.Sequential([ transform.AlterOpLayout(), transform.FoldConstant(), ]) with tvm.transform.PassContext(opt_level=3): mod = alter_layout_seq(mod) byoc_seq = tvm.transform.Sequential([ transform.MergeComposite(dnnl.pattern_table()), transform.AnnotateTarget("dnnl"), transform.MergeCompilerRegions(), transform.PartitionGraph(), ]) with tvm.transform.PassContext(opt_level=3): mod = byoc_seq(mod) return mod