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