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
4from mlir.ir import *
5
6
7class ModuleOp:
8  """Specialization for the module op class."""
9
10  def __init__(self, *, loc=None, ip=None):
11    super().__init__(
12        self.build_generic(results=[], operands=[], loc=loc, ip=ip))
13    body = self.regions[0].blocks.append()
14    with InsertionPoint(body):
15      Operation.create("module_terminator")
16
17  @property
18  def body(self):
19    return self.regions[0].blocks[0]
20
21
22class FuncOp:
23  """Specialization for the func op class."""
24
25  def __init__(self,
26               name,
27               type,
28               *,
29               visibility=None,
30               body_builder=None,
31               loc=None,
32               ip=None):
33    """
34    Create a FuncOp with the provided `name`, `type`, and `visibility`.
35    - `name` is a string representing the function name.
36    - `type` is either a FunctionType or a pair of list describing inputs and
37      results.
38    - `visibility` is a string matching `public`, `private`, or `nested`. None
39      implies private visibility.
40    - `body_builder` is an optional callback, when provided a new entry block
41      is created and the callback is invoked with the new op as argument within
42      an InsertionPoint context already set for the block. The callback is
43      expected to insert a terminator in the block.
44    """
45    sym_name = StringAttr.get(str(name))
46
47    # If the type is passed as a tuple, build a FunctionType on the fly.
48    if isinstance(type, tuple):
49      type = FunctionType.get(inputs=type[0], results=type[1])
50
51    type = TypeAttr.get(type)
52    sym_visibility = StringAttr.get(
53        str(visibility)) if visibility is not None else None
54    super().__init__(sym_name, type, sym_visibility, loc=loc, ip=ip)
55    if body_builder:
56      entry_block = self.add_entry_block()
57      with InsertionPoint(entry_block):
58        body_builder(self)
59
60  @property
61  def is_external(self):
62    return len(self.regions[0].blocks) == 0
63
64  @property
65  def body(self):
66    return self.regions[0]
67
68  @property
69  def type(self):
70    return FunctionType(TypeAttr(self.attributes["type"]).value)
71
72  @property
73  def visibility(self):
74    return self.attributes["sym_visibility"]
75
76  @property
77  def name(self):
78    return self.attributes["sym_name"]
79
80  @property
81  def entry_block(self):
82    if self.is_external:
83      raise IndexError('External function does not have a body')
84    return self.regions[0].blocks[0]
85
86  def add_entry_block(self):
87    '''
88    Add an entry block to the function body using the function signature to infer block arguments
89    Returns the newly created block
90    '''
91    if not self.is_external:
92      raise IndexError('The function already has an entry block!')
93    self.body.blocks.append(*self.type.inputs)
94    return self.body.blocks[0]
95