| # 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 |
| |
| try: |
| from typing import Optional, Sequence, Union |
| |
| import inspect |
| |
| from ..ir import * |
| except ImportError as e: |
| raise RuntimeError("Error loading imports from extension module") from e |
| |
| ARGUMENT_ATTRIBUTE_NAME = "arg_attrs" |
| RESULT_ATTRIBUTE_NAME = "res_attrs" |
| |
| class ModuleOp: |
| """Specialization for the module op class.""" |
| |
| def __init__(self, *, loc=None, ip=None): |
| super().__init__(self.build_generic(results=[], operands=[], loc=loc, |
| ip=ip)) |
| body = self.regions[0].blocks.append() |
| |
| @property |
| def body(self): |
| return self.regions[0].blocks[0] |
| |
| |
| class FuncOp: |
| """Specialization for the func op class.""" |
| |
| def __init__(self, |
| name, |
| type, |
| *, |
| visibility=None, |
| body_builder=None, |
| loc=None, |
| ip=None): |
| """ |
| Create a FuncOp with the provided `name`, `type`, and `visibility`. |
| - `name` is a string representing the function name. |
| - `type` is either a FunctionType or a pair of list describing inputs and |
| results. |
| - `visibility` is a string matching `public`, `private`, or `nested`. None |
| implies private visibility. |
| - `body_builder` is an optional callback, when provided a new entry block |
| is created and the callback is invoked with the new op as argument within |
| an InsertionPoint context already set for the block. The callback is |
| expected to insert a terminator in the block. |
| """ |
| sym_name = StringAttr.get(str(name)) |
| |
| # If the type is passed as a tuple, build a FunctionType on the fly. |
| if isinstance(type, tuple): |
| type = FunctionType.get(inputs=type[0], results=type[1]) |
| |
| type = TypeAttr.get(type) |
| sym_visibility = StringAttr.get( |
| str(visibility)) if visibility is not None else None |
| super().__init__(sym_name, type, sym_visibility, loc=loc, ip=ip) |
| if body_builder: |
| entry_block = self.add_entry_block() |
| with InsertionPoint(entry_block): |
| body_builder(self) |
| |
| @property |
| def is_external(self): |
| return len(self.regions[0].blocks) == 0 |
| |
| @property |
| def body(self): |
| return self.regions[0] |
| |
| @property |
| def type(self): |
| return FunctionType(TypeAttr(self.attributes["type"]).value) |
| |
| @property |
| def visibility(self): |
| return self.attributes["sym_visibility"] |
| |
| @property |
| def name(self) -> StringAttr: |
| return StringAttr(self.attributes["sym_name"]) |
| |
| @property |
| def entry_block(self): |
| if self.is_external: |
| raise IndexError('External function does not have a body') |
| return self.regions[0].blocks[0] |
| |
| def add_entry_block(self): |
| """ |
| Add an entry block to the function body using the function signature to |
| infer block arguments. |
| Returns the newly created block |
| """ |
| if not self.is_external: |
| raise IndexError('The function already has an entry block!') |
| self.body.blocks.append(*self.type.inputs) |
| return self.body.blocks[0] |
| |
| @property |
| def arg_attrs(self): |
| return ArrayAttr(self.attributes[ARGUMENT_ATTRIBUTE_NAME]) |
| |
| @arg_attrs.setter |
| def arg_attrs(self, attribute: Union[ArrayAttr, list]): |
| if isinstance(attribute, ArrayAttr): |
| self.attributes[ARGUMENT_ATTRIBUTE_NAME] = attribute |
| else: |
| self.attributes[ARGUMENT_ATTRIBUTE_NAME] = ArrayAttr.get( |
| attribute, context=self.context) |
| |
| @property |
| def arguments(self): |
| return self.entry_block.arguments |
| |
| @property |
| def result_attrs(self): |
| return self.attributes[RESULT_ATTRIBUTE_NAME] |
| |
| @result_attrs.setter |
| def result_attrs(self, attribute: ArrayAttr): |
| self.attributes[RESULT_ATTRIBUTE_NAME] = attribute |
| |
| @classmethod |
| def from_py_func(FuncOp, |
| *inputs: Type, |
| results: Optional[Sequence[Type]] = None, |
| name: Optional[str] = None): |
| """Decorator to define an MLIR FuncOp specified as a python function. |
| |
| Requires that an `mlir.ir.InsertionPoint` and `mlir.ir.Location` are |
| active for the current thread (i.e. established in a `with` block). |
| |
| When applied as a decorator to a Python function, an entry block will |
| be constructed for the FuncOp with types as specified in `*inputs`. The |
| block arguments will be passed positionally to the Python function. In |
| addition, if the Python function accepts keyword arguments generally or |
| has a corresponding keyword argument, the following will be passed: |
| * `func_op`: The `func` op being defined. |
| |
| By default, the function name will be the Python function `__name__`. This |
| can be overriden by passing the `name` argument to the decorator. |
| |
| If `results` is not specified, then the decorator will implicitly |
| insert a `ReturnOp` with the `Value`'s returned from the decorated |
| function. It will also set the `FuncOp` type with the actual return |
| value types. If `results` is specified, then the decorated function |
| must return `None` and no implicit `ReturnOp` is added (nor are the result |
| types updated). The implicit behavior is intended for simple, single-block |
| cases, and users should specify result types explicitly for any complicated |
| cases. |
| |
| The decorated function can further be called from Python and will insert |
| a `CallOp` at the then-current insertion point, returning either None ( |
| if no return values), a unary Value (for one result), or a list of Values). |
| This mechanism cannot be used to emit recursive calls (by construction). |
| """ |
| |
| def decorator(f): |
| from . import std |
| # Introspect the callable for optional features. |
| sig = inspect.signature(f) |
| has_arg_func_op = False |
| for param in sig.parameters.values(): |
| if param.kind == param.VAR_KEYWORD: |
| has_arg_func_op = True |
| if param.name == "func_op" and (param.kind |
| == param.POSITIONAL_OR_KEYWORD or |
| param.kind == param.KEYWORD_ONLY): |
| has_arg_func_op = True |
| |
| # Emit the FuncOp. |
| implicit_return = results is None |
| symbol_name = name or f.__name__ |
| function_type = FunctionType.get( |
| inputs=inputs, results=[] if implicit_return else results) |
| func_op = FuncOp(name=symbol_name, type=function_type) |
| with InsertionPoint(func_op.add_entry_block()): |
| func_args = func_op.entry_block.arguments |
| func_kwargs = {} |
| if has_arg_func_op: |
| func_kwargs["func_op"] = func_op |
| return_values = f(*func_args, **func_kwargs) |
| if not implicit_return: |
| return_types = list(results) |
| assert return_values is None, ( |
| "Capturing a python function with explicit `results=` " |
| "requires that the wrapped function returns None.") |
| else: |
| # Coerce return values, add ReturnOp and rewrite func type. |
| if return_values is None: |
| return_values = [] |
| elif isinstance(return_values, tuple): |
| return_values = list(return_values) |
| elif isinstance(return_values, Value): |
| # Returning a single value is fine, coerce it into a list. |
| return_values = [return_values] |
| elif isinstance(return_values, OpView): |
| # Returning a single operation is fine, coerce its results a list. |
| return_values = return_values.operation.results |
| elif isinstance(return_values, Operation): |
| # Returning a single operation is fine, coerce its results a list. |
| return_values = return_values.results |
| else: |
| return_values = list(return_values) |
| std.ReturnOp(return_values) |
| # Recompute the function type. |
| return_types = [v.type for v in return_values] |
| function_type = FunctionType.get(inputs=inputs, results=return_types) |
| func_op.attributes["type"] = TypeAttr.get(function_type) |
| |
| def emit_call_op(*call_args): |
| call_op = std.CallOp(return_types, FlatSymbolRefAttr.get(symbol_name), |
| call_args) |
| if return_types is None: |
| return None |
| elif len(return_types) == 1: |
| return call_op.result |
| else: |
| return call_op.results |
| |
| wrapped = emit_call_op |
| wrapped.__name__ = f.__name__ |
| wrapped.func_op = func_op |
| return wrapped |
| |
| return decorator |
