def only_functions_and_definitions_file(self, defns_and_fns):
assert isinstance(defns_and_fns, list)
assert all(isinstance(el, tuple) for el in defns_and_fns)
defns = sum([defns for defns, fns in defns_and_fns], [])
fns = sum([fns for defns, fns in defns_and_fns], [])
return astnodes.MLIRFile(
defns, astnodes.Module(None, None, astnodes.Region(fns)))
def make_mlir_file(module: Optional[mast.Module] = None) -> mast.MLIRFile:
"""
Returns an instance of :class:`mlir.astnodes.MLIRFile` for *module*. If
*module* is *None*, defaults it with an empty :class:`mlir.astnodes.Module`.
"""
if module is None:
module = mast.Module(None, None, mast.Region([]))
return mast.MLIRFile([], [module])
def mlir_file_as_definition_and_module_list(self, defns_and_mods):
assert isinstance(defns_and_mods, list)
assert all(isinstance(el, tuple) for el in defns_and_mods)
defns = sum([defns for defns, mods in defns_and_mods], [])
mods = sum([mods for defns, mods in defns_and_mods], [])
if len(mods) == 1:
module = mods[0]
else:
module = astnodes.Module(None, None, astnodes.Region(mods))
return astnodes.MLIRFile(defns, module)
def linalg_generic(self, iterator_types: List[str]):
assert all(iterator in ["parallel", "reduction"]
for iterator in iterator_types)
array_attr = ast.ArrayAttr([
ast.StringAttr(ast.StringLiteral(it), None)
for it in iterator_types
])
op = linalg.LinalgGeneric(
0, None, None, ast.Region([]), None, None, None, None, None,
ast.AttributeDict(
[ast.AttributeEntry("iterator_types", array_attr)]))
self._insert_op_in_block([], op)
return op
def only_functions_and_definitions_file(self, defns_and_fns):
assert isinstance(defns_and_fns, list)
assert all(isinstance(el, tuple) for el in defns_and_fns)
defns = sum([defns for defns, fns in defns_and_fns], [])
fns = sum([fns for defns, fns in defns_and_fns], [])
if len(fns) == 0:
return astnodes.MLIRFile(defns, [])
else:
fns = [astnodes.Operation([], fn) for fn in fns]
return astnodes.MLIRFile(defns, [
astnodes.Module(None, None,
astnodes.Region([astnodes.Block(None, fns)]))
])
def function(self, name: Optional[str] = None) -> mast.Function:
"""
Inserts a :class:`mlir.astnodes.Function` with name *name* into *block*.
Returns the inserted function.
"""
if name is None:
name = self.name_gen("fn")
op = mast.Function(mast.SymbolRefId(value=name), [], [], None,
mast.Region([]))
self._insert_op_in_block([], op)
return op
def module(self, name: Optional[str] = None) -> mast.Module:
"""
Inserts a :class:`mlir.astnodes.Module` with name *name* into *block*.
Returns the inserted module.
"""
if name is None:
name = None
else:
name = mast.SymbolRefId(name)
op = mast.Module(name, None, mast.Region([]))
self._insert_op_in_block([], op)
return op
def for_(self,
lower_bound: Union[int, mast.SsaId],
upper_bound: Union[int, mast.SsaId],
step: Optional[int] = None,
indexname: Optional[str] = None):
if indexname is None:
indexname = self.core_builder.name_gen("i")
index = mast.AffineSsa(indexname)
if step is None:
match = 0
else:
match = 1
op = affine.AffineForOp(match=match,
index=index,
begin=lower_bound,
end=upper_bound,
step=step,
region=mast.Region(body=[]))
self.core_builder._insert_op_in_block([], op)
return op