注解区域

注解区域#

import tvm
from tvm import relay
from tvm.relay.op.annotation import compiler_begin, compiler_end

def check_region(region_set, target, args, nodes, rets):
    region = region_set.get_region(args[0])
    assert region
    assert target == region.target
    assert set(args) == set(region.args)
    assert set(nodes) == set(region.nodes)
    assert set(rets) == set(region.rets)

测试区域集创建者菱形#

data = relay.var("data", shape=(10, 10))
cb_1 = compiler_begin(data, "test_target")
O_1 = relay.abs(cb_1)
ce_1 = compiler_end(O_1, "test_target")
ce_2 = compiler_end(O_1, "test_target")
cb_2 = compiler_begin(ce_1, "test_target")
O_2 = relay.nn.relu(cb_2)
ce_3 = compiler_end(O_2, "test_target")
cb_d = compiler_begin(ce_2, "default")
X = relay.tanh(cb_d)
ce_d = compiler_end(X, "default")
cb_3 = compiler_begin(ce_3, "test_target")
cb_4 = compiler_begin(ce_d, "test_target")
O_3 = relay.add(cb_3, cb_4)
ce_4 = compiler_end(O_3, "test_target")
diamond = relay.Function([data], ce_4)
print(diamond)

fn (%data: Tensor[(10, 10), float32]) {
  %0 = annotation.compiler_begin(%data, compiler="test_target");
  %1 = abs(%0);
  %2 = annotation.compiler_end(%1, compiler="test_target");
  %3 = annotation.compiler_begin(%2, compiler="test_target");
  %4 = nn.relu(%3);
  %5 = annotation.compiler_end(%4, compiler="test_target");
  %6 = annotation.compiler_end(%1, compiler="test_target");
  %7 = annotation.compiler_begin(%6, compiler="default");
  %8 = tanh(%7);
  %9 = annotation.compiler_end(%8, compiler="default");
  %10 = annotation.compiler_begin(%5, compiler="test_target");
  %11 = annotation.compiler_begin(%9, compiler="test_target");
  %12 = add(%10, %11);
  annotation.compiler_end(%12, compiler="test_target")
}

region_set = relay.analysis.AnnotatedRegionSet(
    diamond, relay.op.get("annotation.compiler_begin"), relay.op.get("annotation.compiler_end")
)
assert len(region_set) == 4

relay.analysis.AnnotatedRegionSet?

Init signature: relay.analysis.AnnotatedRegionSet(expr, region_begin_op, region_end_op)
Docstring:      Class to represent a relay expression split into regions.
Init docstring:
Construct regions from an expression.

Parameters
----------
expr : tvm.relay.Expr
    The expression from which to construct the regions.
region_begin_op : tvm.ir.Op
    The region begin annotation.
region_end_op : tvm.ir.Op
    The region end annotation.
File:           /media/pc/data/lxw/ai/tvm/python/tvm/relay/analysis/annotated_regions.py
Type:           type
Subclasses:     

check_region(
    region_set,
    "test_target",
    [cb_1],
    [cb_1, O_1, ce_1, ce_2],
    [ce_1, ce_2],
)
check_region(
    region_set,
    "test_target",
    [cb_2],
    [cb_2, O_2, ce_3],
    [ce_3],
)
check_region(
    region_set,
    "default",
    [cb_d],
    [cb_d, X, ce_d],
    [ce_d],
)
check_region(
    region_set,
    "test_target",
    [cb_3, cb_4],
    [cb_3, cb_4, O_3, ce_4],
    [ce_4],
)

测试区域集创建器合并#

data = relay.var("data", shape=(10, 10))
cb_1 = compiler_begin(data, "test_target")
O_1 = relay.abs(cb_1)
ce_2 = compiler_end(O_1, "test_target")
O_2 = relay.nn.relu(O_1)
ce_3 = compiler_end(O_2, "test_target")
cb_d = compiler_begin(ce_2, "default")
X = relay.tanh(cb_d)
ce_d = compiler_end(X, "default")
cb_3 = compiler_begin(ce_3, "test_target")
cb_4 = compiler_begin(ce_d, "test_target")
O_3 = relay.add(cb_3, cb_4)
O_4 = relay.add(cb_3, cb_4)
O_5 = relay.Tuple([O_3, O_4])
ce_4 = compiler_end(O_5, "test_target")
merged = relay.Function([data], ce_4)
print(merged)

fn (%data: Tensor[(10, 10), float32]) {
  %0 = annotation.compiler_begin(%data, compiler="test_target");
  %1 = abs(%0);
  %2 = nn.relu(%1);
  %3 = annotation.compiler_end(%2, compiler="test_target");
  %4 = annotation.compiler_end(%1, compiler="test_target");
  %5 = annotation.compiler_begin(%4, compiler="default");
  %6 = tanh(%5);
  %7 = annotation.compiler_end(%6, compiler="default");
  %8 = annotation.compiler_begin(%3, compiler="test_target");
  %9 = annotation.compiler_begin(%7, compiler="test_target");
  %10 = add(%8, %9);
  %11 = add(%8, %9);
  %12 = (%10, %11);
  annotation.compiler_end(%12, compiler="test_target")
}

region_set = relay.analysis.AnnotatedRegionSet(
    merged, relay.op.get("annotation.compiler_begin"), relay.op.get("annotation.compiler_end")
)
assert len(region_set) == 3
check_region(
    region_set,
    "test_target",
    [cb_1],
    [cb_1, O_1, O_2, ce_2, ce_3],
    [ce_2, ce_3],
)
check_region(
    region_set,
    "default",
    [cb_d],
    [cb_d, X, ce_d],
    [ce_d],
)
check_region(
    region_set,
    "test_target",
    [cb_3, cb_4],
    [cb_3, cb_4, O_3, O_4, O_5, ce_4],
    [ce_4],
)
目录