# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from typing import Dict, List, Sequence, Union
from contextlib import contextmanager
import functools
import inspect
import threading
from ..... import ir
from ...._ods_common import (
get_op_result_or_value as _get_op_result_or_value,
get_op_results_or_values as _get_op_results_or_values,
)
from .comprehension import *
from .config import *
from .emitter import *
_CONTEXT = threading.local()
StructuredOpOuts = Union[
ir.Operation,
ir.OpView,
ir.OpResultList,
Sequence[Union[ir.Value, ir.Operation, ir.OpView]],
]
@contextmanager
def bind_op_def(op_def: LinalgOpDef):
if hasattr(_CONTEXT, "current_op_def"):
raise ValueError("Cannot recursively define an operation")
_CONTEXT.current_op_def = op_def
try:
yield op_def
finally:
del _CONTEXT.current_op_def
def current_op_def() -> LinalgOpDef:
try:
return _CONTEXT.current_op_def
except AttributeError:
raise ValueError(
"Attempt to access the current op definition being defined "
"but none is set. Did you mean to call this in an op definition?"
)
def _prepare_structured_op_outs(outs: StructuredOpOuts) -> ValueList:
if isinstance(outs, (ir.Operation, ir.OpView)):
return _get_op_results_or_values(outs)
elif isinstance(outs, ir.OpResultList):
return outs
return [_get_op_result_or_value(o) for o in outs]
class DefinedOpCallable:
"""Callable that wraps any defined op function."""
def __init__(self, op_name: str, op_def: LinalgOpDef):
self.op_name = op_name
self.op_def = op_def
def __call__(
self,
*ins: Union[ir.Operation, ir.OpView, ir.Value],
outs: StructuredOpOuts,
**kwargs,
):
"""Emits the corresponding op definition as IR.
Most arguments are passed through to the underlying emitter. The following
keyword argument is interpreted here:
emit_generic: Emits a generic form as appropriate (default True). If
False, a named form is emitted (which must have been built in to the
compiler).
"""
emit_generic = kwargs.pop("emit_generic", False)
if not isinstance(emit_generic, bool):
raise ValueError(
f"The named argument 'emit_generic' needs to be "
f" of type bool but got {type(emit_generic)}"
)
op_configs = LinalgOpConfig.from_linalg_op_def(
self.op_def, context=ir.Context.current
)
if len(op_configs) != 1:
# TODO: Support composite ops.
raise NotImplementedError(
f"Emission of composite linalg ops not supported: {op_configs}"
)
ctx = ir.Context.current
linalgDialect = ctx.get_dialect_descriptor("linalg")
fully_qualified_name = "linalg." + self.op_name
emit_generic = emit_generic or not ctx.is_registered_operation(
fully_qualified_name
)
op_config = op_configs[0]
out_values = _prepare_structured_op_outs(outs)
in_values = [_get_op_result_or_value(i) for i in ins]
if op_config.structured_op:
if emit_generic:
return emit_generic_structured_op(
op_config.structured_op, *in_values, outs=out_values, **kwargs
)
else:
return emit_named_structured_op(
op_config.structured_op,
self.op_name,
self.op_def.metadata.cpp_class_name,
*in_values,
outs=out_values,
**kwargs,
)
raise NotImplementedError(
f"Emission of linalg op type not supported: {op_config}"
)
def linalg_structured_op(
dsl_func=None, *, op_name=None, op_class_name=None
) -> DefinedOpCallable:
if dsl_func is None:
# Curry the keyword args in for delayed application.
return functools.partial(
linalg_structured_op, op_name=op_name, op_class_name=op_class_name
)
# Determine default names by introspecting the function.
if op_name is None:
op_name = dsl_func.__name__
if op_class_name is None:
# Camel case it.
op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
op_def = LinalgOpDef(
name=op_name, cpp_class_name=op_class_name, doc=inspect.getdoc(dsl_func)
)
# Extract arguments and TensorDefs from the signature.
dsl_func_args = list()
sig = inspect.signature(dsl_func)
for param_name, param in sig.parameters.items():
param_default = param.default
if isinstance(
param_default,
(
TensorDef,
ScalarDef,
IndexAttrDef,
UnaryFnAttrDef,
BinaryFnAttrDef,
TypeFnAttrDef,
),
):
op_def.add_operand(param_name, param_default.operand_def)
else:
raise ValueError(
f"@linalg_structured_op function parameters must be defaulted as "
f"TensorDef(...), ScalarDef(...), or IndexAttrDef(...): "
f"Found {param_name}: {param_default}"
)
dsl_func_args.append(param_default)
# Invoke the DSL func to finish populating the op definition.
with bind_op_def(op_def):
dsl_func(*dsl_func_args)
# TODO: The returned callable should be an IR emitter but that is not
# upstreamed yet.
return DefinedOpCallable(op_name, op_def)
def domain(*dimensions: DimDef):
if any(not isinstance(d, DimDef) for d in dimensions):
raise ValueError(f"Expected dimensions of type DimDef but got {dimensions}")
current_op_def().domain.extend(dimensions)
def implements(*interfaces: OpInterfaceDef):
if any(not isinstance(intr, OpInterfaceDef) for intr in interfaces):
raise ValueError(
f"Expected interfaces of type OpInterfaceDef but got {interfaces}"
)
current_op_def().metadata.implements.extend(interfaces)
def defines(*definitions: OpDefinitionDef):
if any(not isinstance(defi, OpDefinitionDef) for defi in definitions):
raise ValueError(
f"Expected definitions of type OpDefinitionDef but got {definitions}"
)
current_op_def().metadata.defines.extend(definitions)
Codebase Browser
llvm