mie/mie/pass/builtin/convert-scf-to-cf.c

#include <mie/ctx.h>
#include <mie/dialect/arith.h>
#include <mie/dialect/builtin.h>
#include <mie/dialect/cf.h>
#include <mie/dialect/dialect.h>
#include <mie/ir/block.h>
#include <mie/ir/op-definition.h>
#include <mie/ir/op.h>
#include <mie/ir/region.h>
#include <mie/ir/walk.h>
#include <mie/macros.h>
#include <mie/pass/pass-definition.h>
#include <mie/pass/pass.h>
#include <mie/rewrite/convert.h>
#include <mie/rewrite/pattern.h>
#include <mie/rewrite/rewriter.h>

static bool op_is_flat(const struct mie_op *op)
{
	return false;
}

static enum mie_match_result if_match(const struct mie_op *op)
{
	if (!mie_op_is(op, "scf", "if")) {
		return MIE_NO_MATCH_FOUND;
	}

	return MIE_MATCH_FOUND;
}

static struct mie_rewrite_result if_rewrite(
	struct mie_op *op, struct mie_rewriter *rewriter)
{
	printf("if: rewriting %p %s.%s\n", op, op->op_info->op_parent->d_name,
	       op->op_info->op_name);

	struct mie_register *cond = mie_op_get_arg(op, 0);
	struct mie_region *parent_region = op->op_container->b_parent;

	struct mie_region *if_region = mie_op_get_first_region(op);
	struct mie_block *if_start = mie_region_get_first_block(if_region);
	struct mie_block *if_end = mie_region_get_last_block(if_region);
	struct mie_region *else_region = mie_op_get_next_region(op, if_region);
	struct mie_block *else_start = mie_region_get_first_block(else_region);

	struct mie_block *pre_block = op->op_container;
	struct mie_block *end_block
		= mie_rewriter_split_block(rewriter, pre_block, op, "if.end");

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_result); i++) {
		struct mie_register *old_reg = &op->op_result.items[i];
		struct mie_register *new_reg = mie_block_add_param(end_block);
		new_reg->reg_type = old_reg->reg_type;
		char *name = b_strdup(old_reg->reg_name.n_str);

		mie_name_destroy(&old_reg->reg_name);
		mie_rewriter_rename_register(rewriter, new_reg, name);
		mie_rewriter_replace_register(rewriter, old_reg, new_reg);
		free(name);
	}

	struct mie_walker walker;
	mie_walker_begin(&walker, op, MIE_WALKER_F_INCLUDE_OPS);

	do {
		const struct mie_walk_item *item = mie_walker_get(&walker);
		if (!mie_op_is(item->i_op, "scf", "yield")) {
			continue;
		}

		printf("if: found scf.yield %p\n", item->i_op);

		struct mie_op *br = mie_rewriter_replace_op(
			rewriter, item->i_op, "cf", "br");
		struct mie_op_successor *s = mie_rewriter_add_op_successor(
			rewriter, br, end_block, NULL, 0);
		mie_rewriter_move_op_args_to_successor(rewriter, br, s);

	} while (mie_walker_step(&walker) == MIE_SUCCESS);

	mie_walker_end(&walker);

	mie_rewriter_move_blocks_after(
		rewriter, if_region, parent_region, pre_block);
	mie_rewriter_rename_block(rewriter, if_start, "if.then");

	if (else_region) {
		mie_rewriter_move_blocks_after(
			rewriter, else_region, parent_region, if_end);
		mie_rewriter_rename_block(rewriter, else_start, "if.else");
	}

	mie_rewriter_erase_op(rewriter, op);

	mie_rewriter_set_insertion_block(rewriter, pre_block);
	mie_rewriter_set_insertion_point(rewriter, NULL);
	mie_cf_br_cond_put(
		MIE_EMITTER(rewriter), cond, if_start, NULL, 0, else_start,
		NULL, 0);

	return MIE_REWRITE_RESULT(MIE_REWRITE_SUCCESS, MIE_SUCCESS);
}

static enum mie_match_result for_match(const struct mie_op *op)
{
	if (!mie_op_is(op, "scf", "for")) {
		return MIE_NO_MATCH_FOUND;
	}

	return MIE_MATCH_FOUND;
}

static struct mie_rewrite_result for_rewrite(
	struct mie_op *op, struct mie_rewriter *rewriter)
{
	printf("for: rewriting %p %s.%s\n", op, op->op_info->op_parent->d_name,
	       op->op_info->op_name);

	struct mie_region *parent_region = op->op_container->b_parent;
	struct mie_region *for_body = mie_op_get_first_region(op);
	struct mie_block *for_entry = mie_region_get_first_block(for_body);
	struct mie_block *pre_block = op->op_container;
	struct mie_block *end_block
		= mie_rewriter_split_block(rewriter, pre_block, op, "for.end");

	struct mie_register *entry_iv = &for_entry->b_params.items[0];
	struct mie_register *lb = op->op_args.items[0].arg_value.u_reg;
	struct mie_register *ub = op->op_args.items[1].arg_value.u_reg;
	struct mie_register *step = op->op_args.items[2].arg_value.u_reg;

	MIE_VECTOR_DEFINE(struct mie_register *, initial_args);
	mie_vector_push_back(initial_args, &lb, NULL);
	for (size_t i = 3; i < MIE_VECTOR_COUNT(op->op_args); i++) {
		struct mie_register *arg = op->op_args.items[i].arg_value.u_reg;
		mie_vector_push_back(initial_args, &arg, NULL);
	}

	mie_rewriter_set_insertion_block(rewriter, pre_block);
	mie_cf_br_put(
		MIE_EMITTER(rewriter), for_entry, initial_args.items,
		initial_args.count);

	char iv_next_name[64];
	snprintf(
		iv_next_name, sizeof iv_next_name, "%s.next",
		entry_iv->reg_name.n_str);

	mie_rewriter_set_insertion_block(rewriter, for_entry);
	struct mie_block *for_cond
		= mie_rewriter_create_block(rewriter, end_block, "for.cond");

	for (size_t i = 0; i < MIE_VECTOR_COUNT(for_entry->b_params); i++) {
		const char *var_name = for_entry->b_params.items[i].reg_name.n_str;
		const struct mie_type *var_type
			= for_entry->b_params.items[i].reg_type;

		mie_rewriter_add_block_parameter(
			rewriter, for_cond, var_name, entry_iv->reg_type);
	}

	for (size_t i = 0; i < MIE_VECTOR_COUNT(op->op_result); i++) {
		struct mie_register *old_reg = &op->op_result.items[i];
		struct mie_register *new_reg = mie_block_add_param(end_block);
		new_reg->reg_type = old_reg->reg_type;
		char *name = b_strdup(old_reg->reg_name.n_str);

		mie_name_destroy(&old_reg->reg_name);
		mie_rewriter_rename_register(rewriter, new_reg, name);
		mie_rewriter_replace_register(rewriter, old_reg, new_reg);
		free(name);
	}

	struct mie_walker walker;
	mie_walker_begin(&walker, op, MIE_WALKER_F_INCLUDE_OPS);

	do {
		const struct mie_walk_item *item = mie_walker_get(&walker);
		if (!mie_op_is(item->i_op, "scf", "yield")) {
			continue;
		}

		printf("for: found scf.yield %p\n", item->i_op);

		struct mie_op *br = mie_rewriter_replace_op(
			rewriter, item->i_op, "cf", "br");
		struct mie_op_successor *s = mie_rewriter_add_op_successor(
			rewriter, br, for_cond, NULL, 0);
		struct mie_op_arg *iv_arg = mie_rewriter_add_op_successor_arg(
			rewriter, br, s, entry_iv);
		mie_rewriter_move_op_args_to_successor(rewriter, br, s);

	} while (mie_walker_step(&walker) == MIE_SUCCESS);

	mie_walker_end(&walker);

	mie_rewriter_move_blocks_after(rewriter, for_body, parent_region, pre_block);

	mie_rewriter_set_insertion_block(rewriter, for_cond);

	struct mie_register *cond_iv = &for_cond->b_params.items[0];

	struct mie_register *iv_next = mie_arith_addi_put(
		MIE_EMITTER(rewriter), cond_iv, step, iv_next_name);
	struct mie_register *iv_cmp = mie_arith_cmpi_put(
		MIE_EMITTER(rewriter), MIE_ARITH_CMPI_UGE, iv_next, ub, "stop");

	MIE_VECTOR_DEFINE(struct mie_register *, true_args);
	MIE_VECTOR_DEFINE(struct mie_register *, false_args);

	mie_vector_push_back(false_args, &iv_next, NULL);

	for (size_t i = 1; i < MIE_VECTOR_COUNT(for_cond->b_params); i++) {
		struct mie_register *param = &for_cond->b_params.items[i];
		mie_vector_push_back(true_args, &param, NULL);
		mie_vector_push_back(false_args, &param, NULL);
	}

	mie_cf_br_cond_put(
		MIE_EMITTER(rewriter), iv_cmp, end_block, true_args.items,
		true_args.count, for_entry, false_args.items, false_args.count);

	mie_vector_destroy(true_args, NULL);
	mie_vector_destroy(false_args, NULL);

	mie_rewriter_erase_op(rewriter, op);

	return MIE_REWRITE_RESULT(MIE_REWRITE_SUCCESS, MIE_SUCCESS);
}

MIE_REWRITE_PATTERN_BEGIN(if_pattern)
	MIE_REWRITE_PATTERN_ROOT("scf", "if");
	MIE_REWRITE_PATTERN_MATCH(if_match);
	MIE_REWRITE_PATTERN_REWRITE(if_rewrite);
MIE_REWRITE_PATTERN_END()

MIE_REWRITE_PATTERN_BEGIN(for_pattern)
	MIE_REWRITE_PATTERN_ROOT("scf", "for");
	MIE_REWRITE_PATTERN_MATCH(for_match);
	MIE_REWRITE_PATTERN_REWRITE(for_rewrite);
MIE_REWRITE_PATTERN_END()

static struct mie_pass_result transform(
	struct mie_pass *pass, struct mie_op *op, struct mie_pass_args *args)
{
	printf("%s: taking a look at %p %s.%s\n", mie_pass_get_name(pass), op,
	       op->op_info->op_parent->d_name, op->op_info->op_name);

	struct mie_convert_config *cfg = mie_convert_config_create(args->p_ctx);
	mie_convert_config_add_illegal_op(cfg, "scf", "if");
	mie_convert_config_add_illegal_op(cfg, "scf", "for");

	struct mie_pattern_set patterns = {};
	if_pattern_create(&patterns);
	for_pattern_create(&patterns);
	mie_convert_apply(op, cfg, &patterns);

	mie_pattern_set_cleanup(&patterns);
	mie_convert_config_destroy(cfg);

	return MIE_PASS_CONTINUE;
}

MIE_PASS_DEFINITION_BEGIN(convert_scf_to_cf)
	MIE_PASS_NAME("convert-scf-to-cf");
	MIE_PASS_DESCRIPTION(
		"Convert high-level SCF constructs to low-level CF operations "
		"and blocks.");
	MIE_PASS_TRANSFORM(transform);
	MIE_PASS_FILTER_OP("func", "func");
MIE_PASS_DEFINITION_END()