def test_trace_construct_append_1():
trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV())
trace.append(inst=_make_get_block("block2", BlockRV()))
assert str(trace) == "\n".join((
'b0 = sch.get_block(name="block", func_name="main")',
"l1, l2 = sch.get_loops(block=b0)",
'b3 = sch.get_block(name="block2", func_name="main")',
))
def test_trace_simplified_1():
trace = _make_trace_3(BlockRV(), BlockRV(), add_postproc=True)
assert str(trace) == "\n".join((
'b0 = sch.get_block(name="B", func_name="main")',
"sch.compute_inline(block=b0)",
'b1 = sch.get_block(name="C", func_name="main")',
"sch.enter_postproc()",
"sch.compute_inline(block=b1)",
))
trace = trace.simplified(remove_postproc=True)
assert str(trace) == "\n".join((
'b0 = sch.get_block(name="B", func_name="main")',
"sch.compute_inline(block=b0)",
))
def test_trace_construct_1():
trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV())
assert str(trace) == "\n".join((
'b0 = sch.get_block(name="block", func_name="main")',
"l1, l2 = sch.get_loops(block=b0)",
))
assert len(trace.insts) == 2
assert len(trace.decisions) == 0
def test_trace_simplified_3():
trace = _make_trace_4(BlockRV(), LoopRV(), LoopRV(),
LoopRV()).simplified(remove_postproc=False)
assert str(trace) == "\n".join((
'b0 = sch.get_block(name="B", func_name="main")',
"l1, = sch.get_loops(block=b0)",
"l2, l3 = sch.split(loop=l1, factors=[None, 32])",
))
def test_trace_as_json_1():
trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV())
obj = trace.as_json()
assert obj == [
[
["GetBlock", [], ["block", "main"], ["b0"]],
["GetLoops", ["b0"], [], ["l1", "l2"]],
],
[],
]
def test_inst_construct_2():
block = BlockRV()
inst = Instruction(
kind=InstructionKind.get("ComputeInline"),
inputs=[block],
attrs=[],
outputs=[],
)
assert str(inst) == "sch.compute_inline(block=_)"
assert len(inst.inputs) == 1
assert len(inst.attrs) == 0
assert len(inst.outputs) == 0
assert inst.kind.same_as(InstructionKind.get("ComputeInline"))
assert inst.inputs[0].same_as(block)
def test_inst_construct_1():
block = BlockRV()
loop0 = LoopRV()
loop1 = LoopRV()
inst = Instruction(
kind=InstructionKind.get("GetLoops"),
inputs=[block],
attrs=[],
outputs=[loop0, loop1],
)
assert str(inst) == "_, _ = sch.get_loops(block=_)"
assert len(inst.inputs) == 1
assert len(inst.attrs) == 0
assert len(inst.outputs) == 2
assert inst.kind.same_as(InstructionKind.get("GetLoops"))
assert inst.inputs[0].same_as(block)
assert inst.outputs[0].same_as(loop0)
assert inst.outputs[1].same_as(loop1)
def test_apply_json_to_schedule_1():
trace = _make_trace_2(BlockRV())
json_obj = trace.as_json()
sch = tir.Schedule(elementwise, debug_mask="all")
Trace.apply_json_to_schedule(json_obj, sch)
tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"])
def test_trace_apply_to_schedule():
trace = _make_trace_2(BlockRV())
sch = tir.Schedule(elementwise, debug_mask="all")
trace.apply_to_schedule(sch, remove_postproc=False, decision_provider=None)
tvm.ir.assert_structural_equal(elementwise_inlined, sch.mod["main"])
def test_trace_construct_pop_1():
trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV())
last_inst = trace.insts[-1]
assert trace.pop().same_as(last_inst)
assert str(trace) == 'b0 = sch.get_block(name="block", func_name="main")'
def test_trace_construct_get_decision_1():
trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV())
assert trace.get_decision(trace.insts[0]) is None
assert trace.get_decision(trace.insts[1]) is None
