def simple_if_else(cond): if_op = scf.IfOp(cond, [i32, i32], hasElse=True) with InsertionPoint(if_op.then_block): x_true = arith.ConstantOp(i32, 0) y_true = arith.ConstantOp(i32, 1) scf.YieldOp([x_true, y_true]) with InsertionPoint(if_op.else_block): x_false = arith.ConstantOp(i32, 2) y_false = arith.ConstantOp(i32, 3) scf.YieldOp([x_false, y_false]) add = arith.AddIOp(if_op.results[0], if_op.results[1]) return
def simple_if(cond): if_op = scf.IfOp(cond) with InsertionPoint(if_op.then_block): one = arith.ConstantOp(i32, 1) add = arith.AddIOp(one, one) scf.YieldOp([]) return
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