slyubomirsky commented on code in PR #14394:
URL: https://github.com/apache/tvm/pull/14394#discussion_r1148662616
##########
python/tvm/script/parser/relax/parser.py:
##########
@@ -220,7 +220,30 @@ def visit_tvm_declare_function(self: Parser, node:
doc.FunctionDef) -> GlobalVar
param_sinfo = eval_struct_info(self, arg.annotation, eval_str=True)
params.append(relax.Var(arg.arg, param_sinfo))
- func_signature = relax.Function.create_empty(params, ret_sinfo)
+ # find a call to R.func_attr to see if purity should be indicated
+ # namely, find a call to R.func_attr({..., "IsPure": val, ...})
+ # (we don't need any other attributes at the function declaration stage)
+ attrs = None
+ for item in node.body:
+ if (
+ isinstance(item.value, doc.Call)
+ and isinstance(item.value.func, doc.Attribute)
+ and item.value.func.attr == "func_attr"
+ and len(item.value.args) == 1
+ and isinstance(item.value.args[0], doc.Dict)
+ ):
+ index = None
+ for i, key in enumerate(item.value.args[0].keys):
+ if isinstance(key, doc.Constant) and key.value == "IsPure":
+ index = i
+ break
+ if index is not None:
+ val = item.value.args[0].values[index]
+ if isinstance(val, doc.Constant):
+ purity = bool(val.value)
+ attrs = make_node("DictAttrs", IsPure=purity)
+
+ func_signature = relax.Function.create_empty(params, ret_sinfo,
attrs=attrs)
Review Comment:
This seems very messy, but it was needed to prevent a bug where a function
would be parsed and the global var would have the wrong `StructInfo` (the GV
would be pure even if the function was impure). It turned out to be due to the
parsing of function declarations.
We could consider some alternatives, like using a decorator for
purity/impurity instead of an attribute, or checking the types of global vars
during normalization.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
