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