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)
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)
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)
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)