1#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2#  See https://llvm.org/LICENSE.txt for license information.
3#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5from typing import Dict, List, Sequence, Union
6
7from contextlib import contextmanager
8import functools
9import inspect
10import threading
11
12from ..... import ir
13from ...._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
14from .comprehension import *
15from .config import *
16from .emitter import *
17
18_CONTEXT = threading.local()
19
20StructuredOpOuts = Union[ir.Operation, ir.OpView, ir.OpResultList,
21                         Sequence[Union[ir.Value, ir.Operation, ir.OpView]]]
22
23
24@contextmanager
25def bind_op_def(op_def: LinalgOpDef):
26  if hasattr(_CONTEXT, "current_op_def"):
27    raise ValueError("Cannot recursively define an operation")
28  _CONTEXT.current_op_def = op_def
29  try:
30    yield op_def
31  finally:
32    del _CONTEXT.current_op_def
33
34
35def current_op_def() -> LinalgOpDef:
36  try:
37    return _CONTEXT.current_op_def
38  except AttributeError:
39    raise ValueError(
40        "Attempt to access the current op definition being defined "
41        "but none is set. Did you mean to call this in an op definition?")
42
43
44def _prepare_structured_op_outs(outs: StructuredOpOuts) -> ValueList:
45  if isinstance(outs, (ir.Operation, ir.OpView)):
46    return _get_op_results_or_values(outs)
47  elif isinstance(outs, ir.OpResultList):
48    return outs
49
50  return [_get_op_result_or_value(o) for o in outs]
51
52
53class DefinedOpCallable:
54  """Callable that wraps any defined op function."""
55
56  def __init__(self, op_name: str, op_def: LinalgOpDef):
57    self.op_name = op_name
58    self.op_def = op_def
59
60  def __call__(self, *ins: Union[ir.Operation, ir.OpView, ir.Value],
61               outs: StructuredOpOuts, **kwargs):
62    """Emits the corresponding op definition as IR.
63
64    Most arguments are passed through to the underlying emitter. The following
65    keyword argument is interpreted here:
66      emit_generic: Emits a generic form as appropriate (default True). If
67        False, a named form is emitted (which must have been built in to the
68        compiler).
69    """
70    emit_generic = kwargs.pop("emit_generic", False)
71    if not isinstance(emit_generic, bool):
72      raise ValueError(f"The named argument 'emit_generic' needs to be "
73                       f" of type bool but got {type(emit_generic)}")
74
75    op_configs = LinalgOpConfig.from_linalg_op_def(
76        self.op_def, context=ir.Context.current)
77
78    if len(op_configs) != 1:
79      # TODO: Support composite ops.
80      raise NotImplementedError(
81          f"Emission of composite linalg ops not supported: {op_configs}")
82
83    ctx = ir.Context.current
84    linalgDialect = ctx.get_dialect_descriptor("linalg")
85    fully_qualified_name = "linalg." + self.op_name
86    emit_generic = (
87        emit_generic or not ctx.is_registered_operation(fully_qualified_name))
88
89    op_config = op_configs[0]
90    out_values = _prepare_structured_op_outs(outs)
91    in_values = [_get_op_result_or_value(i) for i in ins]
92    if op_config.structured_op:
93      if emit_generic:
94        return emit_generic_structured_op(
95            op_config.structured_op, *in_values, outs=out_values, **kwargs)
96      else:
97        return emit_named_structured_op(
98            op_config.structured_op,
99            self.op_name,
100            self.op_def.metadata.cpp_class_name,
101            *in_values,
102            outs=out_values,
103            **kwargs)
104
105    raise NotImplementedError(
106        f"Emission of linalg op type not supported: {op_config}")
107
108
109def linalg_structured_op(dsl_func=None,
110                         *,
111                         op_name=None,
112                         op_class_name=None) -> DefinedOpCallable:
113  if dsl_func is None:
114    # Curry the keyword args in for delayed application.
115    return functools.partial(
116        linalg_structured_op, op_name=op_name, op_class_name=op_class_name)
117  # Determine default names by introspecting the function.
118  if op_name is None:
119    op_name = dsl_func.__name__
120  if op_class_name is None:
121    # Camel case it.
122    op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
123
124  op_def = LinalgOpDef(
125      name=op_name, cpp_class_name=op_class_name, doc=inspect.getdoc(dsl_func))
126
127  # Extract arguments and TensorDefs from the signature.
128  dsl_func_args = list()
129  sig = inspect.signature(dsl_func)
130  for param_name, param in sig.parameters.items():
131    param_default = param.default
132    if isinstance(param_default,
133                  (TensorDef, ScalarDef, IndexAttrDef, UnaryFnAttrDef,
134                   BinaryFnAttrDef, TypeFnAttrDef)):
135      op_def.add_operand(param_name, param_default.operand_def)
136    else:
137      raise ValueError(
138          f"@linalg_structured_op function parameters must be defaulted as "
139          f"TensorDef(...), ScalarDef(...), or IndexAttrDef(...): "
140          f"Found {param_name}: {param_default}")
141    dsl_func_args.append(param_default)
142
143  # Invoke the DSL func to finish populating the op definition.
144  with bind_op_def(op_def):
145    dsl_func(*dsl_func_args)
146
147  # TODO: The returned callable should be an IR emitter but that is not
148  # upstreamed yet.
149  return DefinedOpCallable(op_name, op_def)
150
151
152def domain(*dimensions: DimDef):
153  if any(not isinstance(d, DimDef) for d in dimensions):
154    raise ValueError(f"Expected dimensions of type DimDef but got {dimensions}")
155  current_op_def().domain.extend(dimensions)
156
157
158def implements(*interfaces: OpInterfaceDef):
159  if any(not isinstance(intr, OpInterfaceDef) for intr in interfaces):
160    raise ValueError(
161        f"Expected interfaces of type OpInterfaceDef but got {interfaces}")
162  current_op_def().metadata.implements.extend(interfaces)
163
164
165def defines(*definitions: OpDefinitionDef):
166  if any(not isinstance(defi, OpDefinitionDef) for defi in definitions):
167    raise ValueError(
168        f"Expected definitions of type OpDefinitionDef but got {definitions}")
169  current_op_def().metadata.defines.extend(definitions)
170