def test_lower_build_tir_module():
func = matmul.with_attr("global_symbol", "main")
func = func.with_attr("tir.noalias", True)
ir_mod = IRModule({"main": func})
# check lowering
lowered_mod = tvm.lower(ir_mod)
tvm.ir.assert_structural_equal(lowered_mod, LoweredModule())
# check building
mod = tvm.build(ir_mod, target="llvm")
_check_module_with_numpy(mod)
def test_lower_build_tir_module():
func = matmul.with_attr("global_symbol", "main")
func = func.with_attr("tir.noalias", True)
ir_mod = IRModule({"main": func})
# check lowering with the CSE pass disabled as otherwise it would do some commoning
with tvm.transform.PassContext(opt_level=3,
disabled_pass=["tir.CommonSubexprElimTIR"]):
lowered_mod = tvm.lower(ir_mod)
tvm.ir.assert_structural_equal(lowered_mod, LoweredTIRModule)
# check building
mod = tvm.build(ir_mod, target="llvm")
_check_module_with_numpy(mod)
def test_meta_schedule_integration_apply_history_best():
@derived_object
class DummyDatabase(PyDatabase):
def __init__(self):
super().__init__()
self.records = []
self.workload_reg = []
def has_workload(self, mod: IRModule) -> Workload:
for workload in self.workload_reg:
if tvm.ir.structural_equal(workload.mod, mod):
return True
return False
def commit_tuning_record(self, record: TuningRecord) -> None:
self.records.append(record)
def commit_workload(self, mod: IRModule) -> Workload:
for workload in self.workload_reg:
if tvm.ir.structural_equal(workload.mod, mod):
return workload
workload = Workload(mod)
self.workload_reg.append(workload)
return workload
def get_top_k(self, workload: Workload,
top_k: int) -> List[TuningRecord]:
return list(
filter(
lambda x: x.workload == workload,
sorted(self.records,
key=lambda x: sum(x.run_secs) / len(x.run_secs)),
))[:int(top_k)]
def __len__(self) -> int:
return len(self.records)
def print_results(self) -> None:
print("\n".join([str(r) for r in self.records]))
mod, _, _ = get_network(name="resnet_18", input_shape=[1, 3, 224, 224])
database = DummyDatabase()
env = ApplyHistoryBest(database)
target = Target("llvm")
workload = database.commit_workload(MockModule)
database.commit_tuning_record(
TuningRecord(Schedule(MockModule).trace, [1.0], workload, target, []))
mod = env.query(task_name="mock-task",
mod=mod,
target=target,
dispatched=[MockModule])
mod = IRModule({"main": mod})
assert tvm.ir.structural_equal(mod, workload.mod)
ir_module = MyModule
print(type(ir_module))
print(ir_module.script())
################################################################################################
# Besides, we can also use tensor expression DSL to write simple operators, and convert them
# to an IRModule.
#
from tvm import te
A = te.placeholder((8,), dtype="float32", name="A")
B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B")
func = te.create_prim_func([A, B])
ir_module_from_te = IRModule({"main": func})
print(ir_module_from_te.script())
################################################################################################
# Build and Run an IRModule
# -------------------------
# We can build the IRModule into a runnable module with specific target backends.
#
mod = tvm.build(ir_module, target="llvm") # The module for CPU backends.
print(type(mod))
################################################################################################
# Prepare the input array and output array, then run the module.
#
