llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Rename 'macro' nomenclature to 'partial eval'.

pull/1/head
Stella Laurenzo 2020-06-26 13:50:51 -07:00
parent dd6a4e638b
commit e45287d83e
4 changed files with 137 additions and 115 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

0
pytest/Compiler/macro_getattr.py → pytest/Compiler/partial_eval_getattr.py
View File

123
python/npcomp/compiler/environment.py
View File

@ -17,13 +17,13 @@ from .target import *
__all__ = [
"BuiltinsValueCoder",
"Environment",
"MacroEvalResult",
"MacroEvalType",
"MacroResolver",
"MacroValueRef",
"LiveValueRef",
"NameReference",
"NameResolver",
"ResolveAttrMacroValueRef",
"PartialEvalResult",
"PartialEvalType",
"PartialEvalHook",
"ResolveAttrLiveValueRef",
"ValueCoder",
"ValueCoderChain",
]
@ -59,7 +59,7 @@ class NameReference:
self.name = name
def load(self, env: "Environment",
ir_h: ir.DialectHelper) -> "MacroEvalResult":
ir_h: ir.DialectHelper) -> "PartialEvalResult":
"""Loads the IR Value associated with the name.
The load may either be direct, returning an existing value or
@ -68,9 +68,9 @@ class NameReference:
Args:
ir_h: The dialect helper used to emit code.
Returns:
A macro evaluation result.
A partial evaluation result.
"""
return MacroEvalResult.not_evaluated()
return PartialEvalResult.not_evaluated()
def store(self, env: "Environment", value: ir.Value, ir_h: ir.DialectHelper):
"""Stores a new value into the name.
@ -103,57 +103,64 @@ class NameResolver:
################################################################################
# Macro evaluation
# Partial evaluation
# When the compiler is extracting from a running program, it is likely that
# evaluations produce live values which can be further partially evaluated
# at import time, in the context of the running instance (versus emitting
# program IR to do so). This facility is called macro evaluation and is
# a pluggable component on the environment.
# program IR to do so). This behavior is controlled through a PartialEvalHook
# on the environment.
################################################################################
class MacroEvalType(Enum):
# The macro could not be evaluated immediately and the operation should
# be code-generated. yields NotImplemented.
class PartialEvalType(Enum):
# Could not be evaluated immediately and the operation should be
# code-generated. yields NotImplemented.
NOT_EVALUATED = 0
# The macro yields a LiveValueRef
# Yields a LiveValueRef
YIELDS_LIVE_VALUE = 1
# The macro yields an IR value
# Yields an IR value
YIELDS_IR_VALUE = 2
# Evaluation yielded an error (yields contains exc_info from sys.exc_info()).
ERROR = 3
class MacroEvalResult(namedtuple("MacroEvalResult", "type,yields")):
"""Encapsulates the result of a macro evaluation."""
class PartialEvalResult(namedtuple("PartialEvalResult", "type,yields")):
"""Encapsulates the result of a partial evaluation."""
@classmethod
def not_evaluated(cls):
return cls(MacroEvalType.NOT_EVALUATED, NotImplemented)
return cls(PartialEvalType.NOT_EVALUATED, NotImplemented)
@classmethod
def yields_live_value(cls, live_value):
assert isinstance(live_value, MacroValueRef)
return cls(MacroEvalType.YIELDS_LIVE_VALUE, live_value)
assert isinstance(live_value, LiveValueRef)
return cls(PartialEvalType.YIELDS_LIVE_VALUE, live_value)
@classmethod
def yields_ir_value(cls, ir_value):
assert isinstance(ir_value, ir.Value)
return cls(MacroEvalType.YIELDS_IR_VALUE, ir_value)
return cls(PartialEvalType.YIELDS_IR_VALUE, ir_value)
@classmethod
def error(cls):
return cls(MacroEvalType.ERROR, sys.exc_info())
return cls(PartialEvalType.ERROR, sys.exc_info())
@classmethod
def error_message(cls, message):
try:
raise RuntimeError(message)
except RuntimeError:
return cls.error()
class MacroValueRef:
class LiveValueRef:
"""Wraps a live value from the containing environment.
Typically, when expressions encounter a live value, a limited number of
"macro" expansions can be done against it in place (versus emitting the code
partial evaluations can be done against it in place (versus emitting the code
to import it and perform the operation). This default base class will not
perform any static evaluations.
"""
@ -165,31 +172,31 @@ class MacroValueRef:
super().__init__()
self.live_value = live_value
def resolve_getattr(self, env: "Environment", attr_name) -> MacroEvalResult:
def resolve_getattr(self, env: "Environment", attr_name) -> PartialEvalResult:
"""Gets a named attribute from the live value."""
return MacroEvalResult.not_evaluated()
return PartialEvalResult.not_evaluated()
def __repr__(self):
return "MacroValueRef({}, {})".format(self.__class__.__name__,
self.live_value)
class ResolveAttrMacroValueRef(MacroValueRef):
class ResolveAttrLiveValueRef(LiveValueRef):
"""Custom MacroValueRef that will resolve attributes via getattr."""
__slots__ = []
def resolve_getattr(self, env: "Environment", attr_name) -> MacroEvalResult:
def resolve_getattr(self, env: "Environment", attr_name) -> PartialEvalResult:
logging.debug("RESOLVE_GETATTR '{}' on {}".format(attr_name,
self.live_value))
try:
attr_py_value = getattr(self.live_value, attr_name)
except:
return MacroEvalResult.error()
return env.macro_resolver.resolve(attr_py_value)
return PartialEvalResult.error()
return env.partial_eval_hook.resolve(attr_py_value)
class MacroResolver:
"""Owned by an environment and performs system-wide macro resolution."""
class PartialEvalHook:
"""Owned by an environment to customize partial evaluation."""
__slots__ = [
"_value_map",
]
@ -198,26 +205,26 @@ class MacroResolver:
super().__init__()
self._value_map = PyValueMap()
def resolve(self, py_value) -> MacroEvalResult:
"""Performs macro resolution on a python value."""
def resolve(self, py_value) -> PartialEvalResult:
"""Performs partial evaluation on a python value."""
binding = self._value_map.lookup(py_value)
if binding is None:
logging.debug("MACRO RESOLVE {}: Passthrough", py_value)
return MacroEvalResult.yields_live_value(MacroValueRef(py_value))
if isinstance(binding, MacroValueRef):
logging.debug("MACRO RESOLVE {}: {}", py_value, binding)
return MacroEvalResult.yields_live_value(binding)
if isinstance(binding, MacroEvalResult):
logging.debug("PARTIAL EVAL RESOLVE {}: Passthrough", py_value)
return PartialEvalResult.yields_live_value(LiveValueRef(py_value))
if isinstance(binding, LiveValueRef):
logging.debug("PARTIAL EVAL RESOLVE {}: {}", py_value, binding)
return PartialEvalResult.yields_live_value(binding)
if isinstance(binding, PartialEvalResult):
return binding
# Attempt to call.
try:
binding = binding(py_value)
assert isinstance(binding, MacroEvalResult), (
"Expected MacroEvalResult but got {}".format(binding))
logging.debug("MACRO RESOLVE {}: {}", py_value, binding)
assert isinstance(binding, PartialEvalResult), (
"Expected PartialEvalResult but got {}".format(binding))
logging.debug("PARTIAL EVAL RESOLVE {}: {}", py_value, binding)
return binding
except:
return MacroEvalResult.error()
return PartialEvalResult.error()
def _bind(self,
binding,
@ -236,10 +243,10 @@ class MacroResolver:
"Must specify one of 'for_ref', 'for_type' or 'for_predicate")
def enable_getattr(self, **kwargs):
"""Enables macro attribute resolution."""
"""Enables partial evaluation of getattr."""
self._bind(
lambda pv: MacroEvalResult.yields_live_value(
ResolveAttrMacroValueRef(pv)), **kwargs)
lambda pv: PartialEvalResult.yields_live_value(
ResolveAttrLiveValueRef(pv)), **kwargs)
################################################################################
@ -258,7 +265,7 @@ class Environment(NameResolver):
"_name_resolvers",
"target",
"value_coder",
"macro_resolver",
"partial_eval_hook",
]
def __init__(self,
@ -267,13 +274,14 @@ class Environment(NameResolver):
target: Target,
name_resolvers=(),
value_coder,
macro_resolver=None):
partial_eval_hook=None):
super().__init__()
self.ir_h = ir_h
self.target = target
self._name_resolvers = name_resolvers
self.value_coder = value_coder
self.macro_resolver = macro_resolver if macro_resolver else MacroResolver()
self.partial_eval_hook = partial_eval_hook if partial_eval_hook else PartialEvalHook(
)
@classmethod
def for_const_global_function(cls, ir_h: ir.DialectHelper, f, *,
@ -332,12 +340,11 @@ class LocalNameReference(NameReference):
super().__init__(name)
self._current_value = initial_value
def load(self, env: "Environment") -> MacroEvalResult:
def load(self, env: "Environment") -> PartialEvalResult:
if self._current_value is None:
return MacroEvalResult.error(
RuntimeError("Attempt to access local '{}' before assignment".format(
self.name)))
return MacroEvalResult.yields_ir_value(self._current_value)
return PartialEvalResult.error_message(
"Attempt to access local '{}' before assignment".format(self.name))
return PartialEvalResult.yields_ir_value(self._current_value)
def store(self, env: "Environment", value: ir.Value):
self._current_value = value
@ -374,8 +381,8 @@ class ConstNameReference(NameReference):
super().__init__(name)
self._py_value = py_value
def load(self, env: "Environment") -> MacroEvalResult:
return env.macro_resolver.resolve(self._py_value)
def load(self, env: "Environment") -> PartialEvalResult:
return env.partial_eval_hook.resolve(self._py_value)
def __repr__(self):
return "<ConstNameReference({}={})>".format(self.name, self._py_value)

18
python/npcomp/compiler/frontend.py
View File

@ -31,7 +31,7 @@ class ImportFrontend:
"_helper",
"_target_factory",
"_value_coder",
"_macro_resolver",
"_partial_eval_hook",
]
def __init__(self,
@ -39,15 +39,15 @@ class ImportFrontend:
*,
target_factory: TargetFactory = GenericTarget64,
value_coder: Optional[ValueCoder] = None,
macro_resolver: Optional[MacroResolver] = None):
partial_eval_hook: Optional[PartialEvalHook] = None):
self._ir_context = ir.MLIRContext() if not ir_context else ir_context
self._ir_module = self._ir_context.new_module()
self._helper = AllDialectHelper(self._ir_context,
ir.OpBuilder(self._ir_context))
self._target_factory = target_factory
self._value_coder = value_coder if value_coder else BuiltinsValueCoder()
self._macro_resolver = (macro_resolver if macro_resolver else
build_default_macro_resolver())
self._partial_eval_hook = (partial_eval_hook if partial_eval_hook else
build_default_partial_eval_hook())
@property
def ir_context(self):
@ -62,8 +62,8 @@ class ImportFrontend:
return self._helper
@property
def macro_resolver(self):
return self._macro_resolver
def partial_eval_hook(self):
return self._partial_eval_hook
def import_global_function(self, f):
"""Imports a global function.
@ -121,7 +121,7 @@ class ImportFrontend:
parameter_bindings=zip(f_params.keys(), ir_f.first_block.args),
value_coder=self._value_coder,
target=target,
macro_resolver=self._macro_resolver)
partial_eval_hook=self._partial_eval_hook)
fctx = FunctionContext(ir_c=ir_c,
ir_f=ir_f,
ir_h=h,
@ -166,8 +166,8 @@ class AllDialectHelper(Numpy.DialectHelper, ScfDialectHelper):
ScfDialectHelper.__init__(self, *args, **kwargs)
def build_default_macro_resolver() -> MacroResolver:
mr = MacroResolver()
def build_default_partial_eval_hook() -> PartialEvalHook:
mr = PartialEvalHook()
### Modules
mr.enable_getattr(for_type=ast.__class__) # The module we use is arbitrary.

111
python/npcomp/compiler/importer.py
View File

@ -46,16 +46,16 @@ class FunctionContext:
ir.emit_error(loc, message)
raise EmittedError(loc, message)
def check_macro_evaluated(self, result: MacroEvalResult):
"""Checks that a macro has evaluated without error."""
if result.type == MacroEvalType.ERROR:
def check_partial_evaluated(self, result: PartialEvalResult):
"""Checks that a PartialEvalResult has evaluated without error."""
if result.type == PartialEvalType.ERROR:
exc_info = result.yields
loc = self.current_loc
message = ("Error while evaluating value from environment:\n" +
"".join(traceback.format_exception(*exc_info)))
ir.emit_error(loc, message)
raise EmittedError(loc, message)
if result.type == MacroEvalType.NOT_EVALUATED:
if result.type == PartialEvalType.NOT_EVALUATED:
self.abort("Unable to evaluate expression")
@property
@ -82,22 +82,26 @@ class FunctionContext:
self.abort("Cannot code python value as constant: {}".format(py_value))
return result
def emit_macro_result(self, macro_result: MacroEvalResult) -> ir.Value:
"""Emits a macro result either as a direct IR value or a constant."""
self.check_macro_evaluated(macro_result)
if macro_result.type == MacroEvalType.YIELDS_IR_VALUE:
def emit_partial_eval_result(self,
partial_result: PartialEvalResult) -> ir.Value:
"""Emits a partial eval result either as a direct IR value or a constant."""
self.check_partial_evaluated(partial_result)
if partial_result.type == PartialEvalType.YIELDS_IR_VALUE:
# Return directly.
return macro_result.yields
elif macro_result.type == MacroEvalType.YIELDS_LIVE_VALUE:
return partial_result.yields
elif partial_result.type == PartialEvalType.YIELDS_LIVE_VALUE:
# Import constant.
return self.emit_const_value(macro_result.yields.live_value)
return self.emit_const_value(partial_result.yields.live_value)
else:
self.abort("Unhandled macro result type {}".format(macro_result))
self.abort("Unhandled partial eval result type {}".format(partial_result))
class BaseNodeVisitor(ast.NodeVisitor):
"""Base class of a node visitor that aborts on unhandled nodes."""
IMPORTER_TYPE = "<unknown>"
__slots__ = [
"fctx",
]
def __init__(self, fctx):
super().__init__()
@ -119,6 +123,10 @@ class FunctionDefImporter(BaseNodeVisitor):
Handles nodes that are direct children of a FunctionDef.
"""
IMPORTER_TYPE = "statement"
__slots__ = [
"ast_fd",
"_last_was_return",
]
def __init__(self, fctx, ast_fd):
super().__init__(fctx)
@ -186,6 +194,9 @@ class ExpressionImporter(BaseNodeVisitor):
IR value that the expression lowers to.
"""
IMPORTER_TYPE = "expression"
__slots__ = [
"value",
]
def __init__(self, fctx):
super().__init__(fctx)
@ -209,12 +220,13 @@ class ExpressionImporter(BaseNodeVisitor):
self.value = ir_const_value
def visit_Attribute(self, ast_node):
# Import the attribute's value recursively as a macro if possible.
macro_importer = MacroImporter(self.fctx)
macro_importer.visit(ast_node)
if macro_importer.macro_result:
self.fctx.check_macro_evaluated(macro_importer.macro_result)
self.value = self.fctx.emit_macro_result(macro_importer.macro_result)
# Import the attribute's value recursively as a partial eval if possible.
pe_importer = PartialEvalImporter(self.fctx)
pe_importer.visit(ast_node)
if pe_importer.partial_eval_result:
self.fctx.check_partial_evaluated(pe_importer.partial_eval_result)
self.value = self.fctx.emit_partial_eval_result(
pe_importer.partial_eval_result)
return
self.fctx.abort("unhandled attribute access mode: {}".format(
@ -331,9 +343,9 @@ class ExpressionImporter(BaseNodeVisitor):
self.fctx.abort("Unsupported expression name context type %s" %
ast_node.ctx.__class__.__name__)
name_ref = self.fctx.lookup_name(ast_node.id)
macro_result = name_ref.load(self.fctx.environment)
logging.debug("LOAD {} -> {}", name_ref, macro_result)
self.value = self.fctx.emit_macro_result(macro_result)
pe_result = name_ref.load(self.fctx.environment)
logging.debug("LOAD {} -> {}", name_ref, pe_result)
self.value = self.fctx.emit_partial_eval_result(pe_result)
def visit_UnaryOp(self, ast_node):
ir_h = self.fctx.ir_h
@ -371,52 +383,55 @@ class ExpressionImporter(BaseNodeVisitor):
self.emit_constant(ast_node.value)
class MacroImporter(BaseNodeVisitor):
"""Importer for expressions that can resolve through the environment's macro
system.
class PartialEvalImporter(BaseNodeVisitor):
"""Importer for performing greedy partial evaluation.
Concretely this is used for Attribute.value and Call resolution.
Attribute resolution is not just treated as a normal expression because it
is first subject to "macro expansion", allowing the environment's macro
resolution facility to operate on live python values from the containing
environment versus naively emitting code for attribute resolution from
is first subject to "partial evaluation", allowing the environment's partial
eval hook to operate on live python values from the containing
environment versus naively emitting code for attribute resolution for
entities that can/should be considered constants from the hosting context.
This is used, for example, to resolve attributes from modules without
by immediately dereferencing/transforming the intervening chain of attributes.
immediately dereferencing/transforming the intervening chain of attributes.
"""
IMPORTER_TYPE = "macro"
IMPORTER_TYPE = "partial_eval"
__slots__ = [
"partial_eval_result",
]
def __init__(self, fctx):
super().__init__(fctx)
self.macro_result = None
self.partial_eval_result = None
def visit_Attribute(self, ast_node):
# Sub-evaluate the 'value'.
sub_macro = MacroImporter(self.fctx)
sub_macro.visit(ast_node.value)
sub_eval = PartialEvalImporter(self.fctx)
sub_eval.visit(ast_node.value)
if sub_macro.macro_result:
# Macro sub-evaluation successful.
sub_result = sub_macro.macro_result
if sub_eval.partial_eval_result:
# Partial sub-evaluation successful.
sub_result = sub_eval.partial_eval_result
else:
# Need to evaluate it as an expression.
sub_expr = ExpressionImporter(self.fctx)
sub_expr.visit(ast_node.value)
assert sub_expr.value, (
"Macro sub expression did not return a value: %r" % (ast_node.value))
sub_result = MacroEvalResult.yields_ir_value(sub_expr.value)
"Evaluated sub expression did not return a value: %r" %
(ast_node.value))
sub_result = PartialEvalResult.yields_ir_value(sub_expr.value)
# Attempt to perform a static getattr as a macro if still operating on a
# live value.
self.fctx.check_macro_evaluated(sub_result)
if sub_result.type == MacroEvalType.YIELDS_LIVE_VALUE:
# Attempt to perform a static getattr as a partial eval if still operating
# on a live value.
self.fctx.check_partial_evaluated(sub_result)
if sub_result.type == PartialEvalType.YIELDS_LIVE_VALUE:
logging.debug("STATIC getattr '{}' on {}", ast_node.attr, sub_result)
getattr_result = sub_result.yields.resolve_getattr(
self.fctx.environment, ast_node.attr)
if getattr_result.type != MacroEvalType.NOT_EVALUATED:
self.fctx.check_macro_evaluated(getattr_result)
self.macro_result = getattr_result
if getattr_result.type != PartialEvalType.NOT_EVALUATED:
self.fctx.check_partial_evaluated(getattr_result)
self.partial_eval_result = getattr_result
return
# If a non-statically evaluable live value, then convert to a constant
# and dynamic dispatch.
@ -424,7 +439,7 @@ class MacroImporter(BaseNodeVisitor):
else:
ir_value = sub_result.yields
# Yielding an IR value from a recursive macro evaluation means that the
# Yielding an IR value from a recursive partial evaluation means that the
# entire chain needs to be hoisted to IR.
# TODO: Implement.
self.fctx.abort("dynamic-emitted getattr not yet supported: %r" %
@ -432,9 +447,9 @@ class MacroImporter(BaseNodeVisitor):
def visit_Name(self, ast_node):
name_ref = self.fctx.lookup_name(ast_node.id)
macro_result = name_ref.load(self.fctx.environment)
logging.debug("LOAD MACRO {} -> {}", name_ref, macro_result)
self.macro_result = macro_result
partial_eval_result = name_ref.load(self.fctx.environment)
logging.debug("PARTIAL EVAL {} -> {}", name_ref, partial_eval_result)
self.partial_eval_result = partial_eval_result
class EmittedError(Exception):