def yield_outputs(self, *output_names: str):
output_values = []
for n in output_names:
try:
output_values.append(self.yield_mapping[n])
except KeyError:
raise ValueError(
f"Body assignments do not assign all outputs: "
f"missing '{n}'")
linalg.YieldOp(output_values)
def build_matmul_tensors_func(func_name, m, k, n, dtype):
lhs_type = RankedTensorType.get([m, k], dtype)
rhs_type = RankedTensorType.get([k, n], dtype)
result_type = RankedTensorType.get([m, n], dtype)
# TODO: There should be a one-liner for this.
func_type = FunctionType.get([lhs_type, rhs_type], [result_type])
_, entry = FuncOp(func_name, func_type)
lhs, rhs = entry.arguments
with InsertionPoint(entry):
op = linalg.MatmulOp([lhs, rhs], results=[result_type])
# TODO: Implement support for SingleBlockImplicitTerminator
block = op.regions[0].blocks.append()
with InsertionPoint(block):
linalg.YieldOp(values=[])
std.ReturnOp([op.result])