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 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
def induction_var(lb, ub, step):
loop = scf.ForOp(lb, ub, step, [lb])
with InsertionPoint(loop.body):
scf.YieldOp([loop.induction_variable])
return
def simple_loop(lb, ub, step):
loop = scf.ForOp(lb, ub, step, [lb, lb])
with InsertionPoint(loop.body):
scf.YieldOp(loop.inner_iter_args)
return