예제 #1
0
def test_inline_into_multiple_consumers():
    mod = SoftmaxBeforeInline
    target = Target("cuda", host="llvm")
    ctx = _create_context(
        mod=mod,
        target=target,
        rule=auto_inline(target=target),
    )
    (space, ) = ctx.space_generator.generate_design_space(mod=mod)
    tvm.ir.assert_structural_equal(lhs=space.mod, rhs=SoftmaxAfterInline)
예제 #2
0
def test_inline_consumer_chain():
    mod = Conv2DBiasBnReLU
    target = Target("llvm")
    ctx = _create_context(
        mod=mod,
        target=target,
        rule=auto_inline(target=target),
    )
    (space, ) = ctx.space_generator.generate_design_space(mod=mod)
    tvm.ir.assert_structural_equal(lhs=space.mod, rhs=Conv2DBiasBnReLUInlined)
예제 #3
0
def test_inline_pure_spatial():
    mod = BeforePureSpatial
    target = Target("llvm")
    ctx = _create_context(
        mod=mod,
        target=target,
        rule=auto_inline(target=target),
    )
    (space, ) = ctx.space_generator.generate_design_space(mod=mod)
    tvm.ir.assert_structural_equal(lhs=space.mod, rhs=AfterPureSpatial)
예제 #4
0
def test_inline_into_cache():
    mod = MultiLevelTiledConv2D
    target = Target("cuda", host="llvm")
    ctx = _create_context(
        mod=mod,
        target=target,
        rule=auto_inline(target=target),
    )
    (space, ) = ctx.space_generator.generate_design_space(mod=mod)
    tvm.ir.assert_structural_equal(lhs=space.mod,
                                   rhs=MultiLevelTiledConv2DAfterInline)