def testBuildFuncOp():
ctx = Context()
with Location.unknown(ctx) as loc:
m = builtin.ModuleOp()
f32 = F32Type.get()
tensor_type = RankedTensorType.get((2, 3, 4), f32)
with InsertionPoint.at_block_begin(m.body):
f = func.FuncOp(name="some_func",
type=FunctionType.get(
inputs=[tensor_type, tensor_type],
results=[tensor_type]),
visibility="nested")
# CHECK: Name is: "some_func"
print("Name is: ", f.name)
# CHECK: Type is: (tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
print("Type is: ", f.type)
# CHECK: Visibility is: "nested"
print("Visibility is: ", f.visibility)
try:
entry_block = f.entry_block
except IndexError as e:
# CHECK: External function does not have a body
print(e)
with InsertionPoint(f.add_entry_block()):
func.ReturnOp([f.entry_block.arguments[0]])
pass
try:
f.add_entry_block()
except IndexError as e:
# CHECK: The function already has an entry block!
print(e)
# Try the callback builder and passing type as tuple.
f = func.FuncOp(name="some_other_func",
type=([tensor_type, tensor_type], [tensor_type]),
visibility="nested",
body_builder=lambda f: func.ReturnOp(
[f.entry_block.arguments[0]]))
# CHECK: module {
# CHECK: func nested @some_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
# CHECK: return %arg0 : tensor<2x3x4xf32>
# CHECK: }
# CHECK: func nested @some_other_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
# CHECK: return %arg0 : tensor<2x3x4xf32>
# CHECK: }
print(m)
def testFunctionCalls():
foo = func.FuncOp("foo", ([], []))
foo.sym_visibility = StringAttr.get("private")
bar = func.FuncOp("bar", ([], [IndexType.get()]))
bar.sym_visibility = StringAttr.get("private")
qux = func.FuncOp("qux", ([], [F32Type.get()]))
qux.sym_visibility = StringAttr.get("private")
with InsertionPoint(func.FuncOp("caller", ([], [])).add_entry_block()):
func.CallOp(foo, [])
func.CallOp([IndexType.get()], "bar", [])
func.CallOp([F32Type.get()], FlatSymbolRefAttr.get("qux"), [])
func.ReturnOp([])
def testOpsAsArguments():
index_type = IndexType.get()
callee = func.FuncOp("callee", ([], [index_type, index_type]),
visibility="private")
f = func.FuncOp("ops_as_arguments", ([], []))
with InsertionPoint(f.add_entry_block()):
lb = arith.ConstantOp.create_index(0)
ub = arith.ConstantOp.create_index(42)
step = arith.ConstantOp.create_index(2)
iter_args = func.CallOp(callee, [])
loop = scf.ForOp(lb, ub, step, iter_args)
with InsertionPoint(loop.body):
scf.YieldOp(loop.inner_iter_args)
func.ReturnOp([])
def build(self, types: List[ir.Type]):
"""Builds the ir.Module. The module has only the @main function,
which will convert the input through the list of types and then back
to the initial type. The roundtrip type must be a dense tensor."""
assert self._module is None, 'StressTest: must not call build() repeatedly'
self._module = ir.Module.create()
with ir.InsertionPoint(self._module.body):
tp0 = types.pop(0)
self._roundtripTp = tp0
# TODO: assert dense? assert element type is recognised by the TypeConverter?
types.append(tp0)
funcTp = ir.FunctionType.get(inputs=[tp0], results=[tp0])
funcOp = func.FuncOp(name='main', type=funcTp)
funcOp.attributes['llvm.emit_c_interface'] = ir.UnitAttr.get()
with ir.InsertionPoint(funcOp.add_entry_block()):
arg0 = funcOp.entry_block.arguments[0]
self._assertEqualsRoundtripTp(arg0.type)
v = st.ConvertOp(types.pop(0), arg0)
for tp in types:
w = st.ConvertOp(tp, v)
# Release intermediate tensors before they fall out of scope.
st.ReleaseOp(v.result)
v = w
self._assertEqualsRoundtripTp(v.result.type)
func.ReturnOp(v)
return self
def testBlockCreation():
with Context() as ctx, Location.unknown():
module = Module.create()
with InsertionPoint(module.body):
f_type = FunctionType.get(
[IntegerType.get_signless(32),
IntegerType.get_signless(16)], [])
f_op = func.FuncOp("test", f_type)
entry_block = f_op.add_entry_block()
i32_arg, i16_arg = entry_block.arguments
successor_block = entry_block.create_after(i32_arg.type)
with InsertionPoint(successor_block) as successor_ip:
assert successor_ip.block == successor_block
func.ReturnOp([])
middle_block = successor_block.create_before(i16_arg.type)
with InsertionPoint(entry_block) as entry_ip:
assert entry_ip.block == entry_block
cf.BranchOp([i16_arg], dest=middle_block)
with InsertionPoint(middle_block) as middle_ip:
assert middle_ip.block == middle_block
cf.BranchOp([i32_arg], dest=successor_block)
print(module.operation)
# Ensure region back references are coherent.
assert entry_block.region == middle_block.region == successor_block.region
def testTransferReadOp():
module = Module.create()
with InsertionPoint(module.body):
vector_type = VectorType.get([2, 3], F32Type.get())
memref_type = MemRefType.get([-1, -1], F32Type.get())
index_type = IndexType.get()
mask_type = VectorType.get(vector_type.shape,
IntegerType.get_signless(1))
identity_map = AffineMap.get_identity(vector_type.rank)
identity_map_attr = AffineMapAttr.get(identity_map)
f = func.FuncOp("transfer_read",
([memref_type, index_type,
F32Type.get(), mask_type], []))
with InsertionPoint(f.add_entry_block()):
A, zero, padding, mask = f.arguments
vector.TransferReadOp(vector_type,
A, [zero, zero],
identity_map_attr,
padding,
mask=mask)
vector.TransferReadOp(vector_type, A, [zero, zero],
identity_map_attr, padding)
func.ReturnOp([])
# CHECK: @transfer_read(%[[MEM:.*]]: memref<?x?xf32>, %[[IDX:.*]]: index,
# CHECK: %[[PAD:.*]]: f32, %[[MASK:.*]]: vector<2x3xi1>)
# CHECK: vector.transfer_read %[[MEM]][%[[IDX]], %[[IDX]]], %[[PAD]], %[[MASK]]
# CHECK: vector.transfer_read %[[MEM]][%[[IDX]], %[[IDX]]], %[[PAD]]
# CHECK-NOT: %[[MASK]]
print(module)
def testFuncArgumentAccess():
with Context() as ctx, Location.unknown():
ctx.allow_unregistered_dialects = True
module = Module.create()
f32 = F32Type.get()
f64 = F64Type.get()
with InsertionPoint(module.body):
f = func.FuncOp("some_func", ([f32, f32], [f32, f32]))
with InsertionPoint(f.add_entry_block()):
func.ReturnOp(f.arguments)
f.arg_attrs = ArrayAttr.get([
DictAttr.get({
"custom_dialect.foo": StringAttr.get("bar"),
"custom_dialect.baz": UnitAttr.get()
}),
DictAttr.get({"custom_dialect.qux": ArrayAttr.get([])})
])
f.result_attrs = ArrayAttr.get([
DictAttr.get({"custom_dialect.res1": FloatAttr.get(f32, 42.0)}),
DictAttr.get({"custom_dialect.res2": FloatAttr.get(f64, 256.0)})
])
other = func.FuncOp("other_func", ([f32, f32], []))
with InsertionPoint(other.add_entry_block()):
func.ReturnOp([])
other.arg_attrs = [
DictAttr.get({"custom_dialect.foo": StringAttr.get("qux")}),
DictAttr.get()
]
# CHECK: [{custom_dialect.baz, custom_dialect.foo = "bar"}, {custom_dialect.qux = []}]
print(f.arg_attrs)
# CHECK: [{custom_dialect.res1 = 4.200000e+01 : f32}, {custom_dialect.res2 = 2.560000e+02 : f64}]
print(f.result_attrs)
# CHECK: func @some_func(
# CHECK: %[[ARG0:.*]]: f32 {custom_dialect.baz, custom_dialect.foo = "bar"},
# CHECK: %[[ARG1:.*]]: f32 {custom_dialect.qux = []}) ->
# CHECK: f32 {custom_dialect.res1 = 4.200000e+01 : f32},
# CHECK: f32 {custom_dialect.res2 = 2.560000e+02 : f64})
# CHECK: return %[[ARG0]], %[[ARG1]] : f32, f32
#
# CHECK: func @other_func(
# CHECK: %{{.*}}: f32 {custom_dialect.foo = "qux"},
# CHECK: %{{.*}}: f32)
print(module)
def emit_timer_func() -> func.FuncOp:
"""Returns the declaration of nanoTime function. If nanoTime function is
used, the `MLIR_RUNNER_UTILS` and `MLIR_C_RUNNER_UTILS` must be included.
"""
i64_type = ir.IntegerType.get_signless(64)
nanoTime = func.FuncOp("nanoTime", ([], [i64_type]), visibility="private")
nanoTime.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
return nanoTime
def testFirstBlockCreation():
with Context() as ctx, Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
f = func.FuncOp("test", ([f32], []))
entry_block = Block.create_at_start(f.operation.regions[0], [f32])
with InsertionPoint(entry_block):
func.ReturnOp([])
print(module)
assert module.operation.verify()
assert f.body.blocks[0] == entry_block
def emit_benchmark_wrapped_main_func(func, timer_func):
"""Takes a function and a timer function, both represented as FuncOp
objects, and returns a new function. This new function wraps the call to
the original function between calls to the timer_func and this wrapping
in turn is executed inside a loop. The loop is executed
len(func.type.results) times. This function can be used to create a
"time measuring" variant of a function.
"""
i64_type = ir.IntegerType.get_signless(64)
memref_of_i64_type = ir.MemRefType.get([-1], i64_type)
wrapped_func = func.FuncOp(
# Same signature and an extra buffer of indices to save timings.
"main",
(func.arguments.types + [memref_of_i64_type], func.type.results),
visibility="public"
)
wrapped_func.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
num_results = len(func.type.results)
with ir.InsertionPoint(wrapped_func.add_entry_block()):
timer_buffer = wrapped_func.arguments[-1]
zero = arith.ConstantOp.create_index(0)
n_iterations = memref.DimOp(ir.IndexType.get(), timer_buffer, zero)
one = arith.ConstantOp.create_index(1)
iter_args = list(wrapped_func.arguments[-num_results - 1:-1])
loop = scf.ForOp(zero, n_iterations, one, iter_args)
with ir.InsertionPoint(loop.body):
start = func.CallOp(timer_func, [])
call = func.CallOp(
func,
wrapped_func.arguments[:-num_results - 1] + loop.inner_iter_args
)
end = func.CallOp(timer_func, [])
time_taken = arith.SubIOp(end, start)
memref.StoreOp(time_taken, timer_buffer, [loop.induction_variable])
scf.YieldOp(list(call.results))
func.ReturnOp(loop)
return wrapped_func
