def test_set_dispatch(ir_docsifier):
IRDocsifier.set_dispatch(Var, _get_id_doc_printer("x2"), f"{__file__}-2")
with ir_docsifier.dispatch_token(f"{__file__}-2"):
doc = ir_docsifier.as_doc(Var("x", dtype="int8"), ObjectPath.root())
assert doc.name == "x2"
doc = ir_docsifier.as_doc(Var("x", dtype="int8"), ObjectPath.root())
assert doc.name == "x"
def test_is_var_defined():
var_table = VarTable()
a = Var("a", dtype="int32")
object_path = ObjectPath.root().attr("a")
frame = VarDefFrame()
var_table.define(a, "a", object_path, frame)
assert var_table.is_var_defined(a)
assert a in var_table
def var(dtype, span):
assert isinstance(
self.node, ast.Assign
), f"VarDef expected ast.Assign but got {type(self.node)}"
names = [x.id.name for x in self.node.lhs]
if len(names) != 1:
self.context.report_error(
f"VarDef expected assign to only one var, but got {names}", span
)
v = Var(names[0], dtype, span=span)
self.context.update_symbol(v.name, v, self.node)
def env_thread(env_name, span):
assert isinstance(
self.node, ast.Assign
), f"EnvThread expected ast.Assign but got {type(self.node)}"
names = [x.id.name for x in self.node.lhs]
if len(names) != 1:
self.context.report_error(
f"VarDef expected assign to only one var, but got {names}", span
)
v = Var(names[0], dtype="int32", span=span)
self.context.func_var_env_dict[v] = env_name
self.context.update_symbol(v.name, v, self.node)
def buffer_var(dtype, storage_scope, span):
assert isinstance(
self.node, ast.Assign
), f"BufferVarDef expected ast.Assign but got {type(self.node)}"
names = [x.id.name for x in self.node.lhs]
if len(names) != 1:
self.context.report_error(
f"VarDef expected assign to only one var, but got {names}", span
)
ptr_type = tvm.ir.PointerType(tvm.ir.PrimType(dtype), storage_scope)
v = Var(names[0], ptr_type, span=span)
self.context.update_symbol(v.name, v, self.node)
def test_var_out_of_scope():
var_table = VarTable()
var_name = "a"
var_obj = Var(var_name, dtype="int32")
object_path = ObjectPath.root().attr("a")
frame = VarDefFrame()
var_table.define(var_obj, var_name, object_path, frame)
with frame:
assert var_obj in var_table
assert var_obj not in var_table
assert var_table.get_var_doc(var_obj, object_path) is None
def test_define_by_doc():
var_table = VarTable()
var_name = "a"
var_obj = Var(var_name, dtype="int32")
object_path = ObjectPath.root().attr("a")
frame = VarDefFrame()
var_table.define_by_doc(var_obj, lambda: LiteralDoc(var_name), frame)
var_doc = var_table.get_var_doc(var_obj, object_path)
assert isinstance(var_doc, LiteralDoc)
assert var_doc.value == var_name
assert list(var_doc.source_paths) == [object_path]
def test_define():
var_table = VarTable()
var_name = "a"
var_obj = Var(var_name, dtype="int32")
object_path = ObjectPath.root().attr("a")
frame = VarDefFrame()
id_doc = var_table.define(var_obj, var_name, object_path, frame)
assert id_doc.name == "a"
assert list(id_doc.source_paths) == [object_path]
id_doc = var_table.get_var_doc(var_obj, object_path)
assert id_doc.name == "a"
assert list(id_doc.source_paths) == [object_path]
def _make_vars(*args: str) -> List[Var]:
return [Var(arg, dtype="int32") for arg in args]
def test_as_doc(ir_docsifier):
object_path = ObjectPath.root()
doc = ir_docsifier.as_doc(Var("x", "int8"), ObjectPath.root())
assert doc.name == "x"
assert list(doc.source_paths) == [object_path]