scx_lavd 1.0.20

#include "scxtest/scx_test.h"
#include <scx/common.bpf.h>
#include <lib/sdt_task.h>

#include <lib/cpumask.h>
#include <lib/topology.h>

volatile topo_ptr topo_all;

/*
 * XXXETSAL: This is a (hopefully) temporary measure that
 * makes it easier to integrate with existing schedulers that
 * use arbitraty IDs to index CPUs/LLCs/nodes. In the future we 
 * will just keep a CPU id to CPU topology node array, but for
 * now we will have an array for each level.
 */
u64 topo_nodes[TOPO_MAX_LEVEL][NR_CPUS];

int nr_topo_nodes[TOPO_MAX_LEVEL];

__hidden
int topo_contains(topo_ptr topo, u32 cpu)
{
	return scx_bitmap_test_cpu(cpu, topo->mask);
}

static
int topo_subset(topo_ptr topo, scx_bitmap_t mask)
{
	return scx_bitmap_subset(topo->mask, mask);
}

static
topo_ptr topo_node(topo_ptr parent, scx_bitmap_t mask, u64 id)
{
	topo_ptr topo;

	topo = scx_static_alloc(sizeof(struct topology), 1);
	if (!topo) {
		bpf_printk("static allocation failed");
		return NULL;
	}

	topo->parent = parent;
	topo->nr_children = 0;
	topo->level = parent ? topo->parent->level + 1 : 0;
	topo->id = id;
	/*
	* The passed-in mask is deliberately consumed; topo_node takes ownership.
	* Do not reuse the same mask elsewhere after this call.
	*/
	topo->mask = mask;

	if (topo->level >= TOPO_MAX_LEVEL) {
		bpf_printk("topology is too deep");
		return NULL;
	}

	if (id >= NR_CPUS) {
		bpf_printk("invalid node id");
		return NULL;
	}

	topo_nodes[topo->level][topo->id] = (u64)topo;

	return topo;
}


static
int topo_add(topo_ptr parent, scx_bitmap_t mask, u64 id)
{
	topo_ptr child;

	if (unlikely(!mask)) {
		bpf_printk("NULL mask pointer");
		return -EINVAL;
	}

	child = topo_node(parent, mask, id);
	if (!child)
		return -ENOMEM;

	if (parent->nr_children >= TOPO_MAX_CHILDREN) {
		bpf_printk("topology fanout is too large");
		return -EINVAL;
	}

	parent->children[parent->nr_children++] = child;

	return 0;
}

__weak
int topo_init(scx_bitmap_t __arg_arena mask, u64 data_size, u64 id)
{
	/* Initializing the child to appease the verifier. */
	topo_ptr topo, child = NULL;
	int i, j;

	topo = topo_all;
	if (!topo_all) {
		topo_all = topo_node(NULL, mask, id);
		if (!topo_all) {
			bpf_printk("couldn't initialize topology");
			return -EINVAL;
		}

		return 0;
	}

	for (i = 0; i < TOPO_MAX_LEVEL && can_loop; i++) {
		if (!topo_subset(topo, mask)) {
			bpf_printk("mask not a subset of a topology node");
			topo_print();
			return -EINVAL;
		}

		for (j = 0; j < topo->nr_children && can_loop; j++) {
			child = topo->children[j];
			if (topo_subset(child, mask))
				break;

			if (scx_bitmap_intersects(child->mask, mask)) {
				bpf_printk("partially intersecting topology nodes");
				return -EINVAL;
			}
		}

		/*
		 * If we don't fit in any child, we belong right below the
		 * parent topology node.
		 */
		if (j == topo->nr_children) {
			topo_add(topo, mask, id);
			nr_topo_nodes[i]++;
			return 0;
		}

		if (!child) {
			bpf_printk("child is not valid");
			return 0;
		}

		topo = child;
	}

	topo->data = NULL;
	if (data_size) {
		topo->data = scx_static_alloc(data_size, 1);
		if (!topo->data)
			return -ENOMEM;
	}

	bpf_printk("topology is too deep");
	return -EINVAL;
}

__weak
topo_ptr topo_find_descendant(topo_ptr topo, u32 cpu)
{
	topo_ptr child;
	int lvl, i;

	if (!topo_contains(topo, cpu)) {
		bpf_printk("missing cpu from topology");
		return NULL;
	}

	for (lvl = 0; lvl < TOPO_MAX_LEVEL && can_loop; lvl++) {
		if (topo->nr_children == 0)
			return topo;

		for (i = 0; i < topo->nr_children && can_loop; i++) {
			child = topo->children[i];
			if (topo_contains(child, cpu))
				break;
		}

		if (i == topo->nr_children) {
			bpf_printk("missing cpu from inner topology nodes");
			return NULL;
		}

		topo = child;
	}

	return topo;
}

__weak
topo_ptr topo_find_ancestor(topo_ptr topo, u32 cpu)
{
	while (topo->parent && !topo_contains(topo, cpu))
		topo = topo->parent;

	if (!topo_contains(topo, cpu))
		bpf_printk("could not find cpu");

	return topo;

}

__weak
topo_ptr topo_find_sibling(topo_ptr topo, u32 cpu)
{
	topo_ptr parent = topo->parent;
	topo_ptr child;
	int i;

	if (!parent) {
		bpf_printk("parent has no sibling");
		return NULL;
	}

	for (i = 0; i < topo->nr_children && can_loop; i++) {
		child = topo->children[i];
		if (topo_contains(child, cpu))
			return child;
	}

	return NULL;

}

__weak
u64 topo_mask_level_internal(topo_ptr topo, enum topo_level level)
{
	if (unlikely(level < 0 || level >= TOPO_MAX_LEVEL)) {
		bpf_printk("invalid topology level %d", level);
		return (u64)NULL;
	}

	if (unlikely(topo->level < level)) {
		bpf_printk("requesting cpumask from lower level %d, starting from %d", level, topo->level);
		return (u64)NULL;
	}

	while (topo->level > level && can_loop)
		topo = topo->parent;

	return (u64)topo->mask;
}

static int __maybe_unused
topo_iter_start_from(struct topo_iter *iter, topo_ptr topo)
{
	topo_ptr parent;
	enum topo_level lvl;
	int ind;

	if (!topo_all)
		return -EINVAL;

	iter->topo = topo;
	bpf_for(ind, 0, TOPO_MAX_LEVEL)
		iter->indices[ind] = -1;

	parent = topo->parent;
	for (lvl = topo->level; lvl > 0 && can_loop; lvl--) {
		for (ind = 0; ind < parent->nr_children && can_loop; ind++) {
			if (parent->children[ind] == topo)
				break;
		}

		if (ind == parent->nr_children) {
			bpf_printk("could not find topology node in parent");
			return -EINVAL;
		}


		if (unlikely(lvl >= TOPO_MAX_LEVEL)) {
			bpf_printk("invalid level %d", lvl);
			return -EINVAL;
		}

		iter->indices[lvl] = ind;

		/* Go one level up. */
		parent = parent->parent;
		topo = topo->parent;
	}

	/* We know we only have one root topology node. */
	iter->indices[0] = 0;

	return 0;
}

/* We choose in-order traversal. */
__weak bool
topo_iter_next(struct topo_iter *iter)
{
	topo_ptr parent;
	enum topo_level lvl;

	if (unlikely(!iter)) {
		bpf_printk("passing NULL iterator");
		return false;
	}

	/* Case 1: We have children. Go one step down and get the left most one. */
	if (iter->topo->nr_children > 0) {
		iter->topo = iter->topo->children[0];

		lvl = iter->topo->level;
		if (unlikely(lvl < 0 || lvl >= TOPO_MAX_LEVEL)) {
			/*
			 * XXXETSAL: We have both bpf_printk and bpf_printk
			 * out of an abundance of caution: In some cases bpf_printk
			 * does not fire at all, making debugging more difficult.
			 */
			bpf_printk("invalid child level %d", lvl);
			bpf_printk("invalid child level %d", lvl);
			return false;
		}

		iter->indices[lvl] = 0;

		return true;
	}


	/*
	 * Case 2: We have no children. Go up until we find a rightmost sibling,
	 * then choose that sibling.
	 */
	while (iter->topo->level > 0 && can_loop) {
		lvl = iter->topo->level;
		if (unlikely(lvl < 0 || lvl >= TOPO_MAX_LEVEL)) {
			bpf_printk("invalid level %d", lvl);
			return false;
		}

		iter->indices[lvl] += 1;

		parent = iter->topo->parent;
		if (iter->indices[lvl] == parent->nr_children) {
			/* Done with parent, go up a level. */
			iter->indices[lvl] = -1;
			iter->topo = parent;
			continue;
		}

		iter->topo = parent->children[iter->indices[lvl]];
		return true;
	}

	/* Could not find a right sibling for any ancestor. */

	return false;
}

__weak u64
topo_iter_level_internal(struct topo_iter *iter, enum topo_level lvl)
{
	if (!iter)
		return (u64)NULL;

	do  {
		if (!topo_iter_next(iter))
			return (u64)NULL;
	} while (iter->topo->level != lvl && can_loop);

	return (u64)iter->topo;
}

__weak int
topo_cpu_to_llc_id(u32 cpu)
{
	topo_ptr topo;
	u32 id;

	if (cpu >= nr_cpu_ids) {
		bpf_printk("invalid cpu id: %u", cpu);
		return -EINVAL;
	}

	topo = (topo_ptr)topo_nodes[TOPO_CPU][cpu];
	if (!topo) {
		bpf_printk("cpu %u has no topology node set", cpu);
		return -EINVAL;
	}

	/* TOPO_CPU -> TOPO_CORE -> TOPO_LLC */
	id = topo->parent->parent->id;
	return id;
}

volatile u64 a;

__weak __maybe_unused
int topo_print(void)
{
	struct topo_iter iter;
	char indent[TOPO_MAX_LEVEL];
	int ret;
	int i;

	if (!topo_all) {
		bpf_printk("[NO TOPOLOGY]");
		return 0;
	}

	ret = topo_iter_start(&iter);
	if (ret)
		return ret;

	do {
		bpf_for(i, 0, TOPO_MAX_LEVEL) {
			if (i == iter.topo->level) {
				indent[i] = '\0';
				break;
			}

			indent[i] = '\t';
		}
		bpf_printk("%s (LEVEL %d) [%d, %d, %d ,%d, %d]",
			indent,
			iter.topo->level,
			iter.indices[TOPO_TOP],
			iter.indices[TOPO_NODE],
			iter.indices[TOPO_LLC],
			iter.indices[TOPO_CORE],
			iter.indices[TOPO_CPU]);

		scx_bitmap_print(iter.topo->mask);
	} while (topo_iter_next(&iter) && can_loop);

	return 0;
}


__weak __maybe_unused
int topo_print_by_level(void)
{
	struct topo_iter iter;
	topo_ptr topo;

	bpf_printk("TOP-LEVEL MASK");
	scx_bitmap_print(topo_all->mask);
	bpf_printk("\n");

	bpf_printk("NODE MASKS");
	TOPO_FOR_EACH_NODE(&iter, topo)
		scx_bitmap_print(topo->mask);
	bpf_printk("\n");

	bpf_printk("LLC MASKS");
	TOPO_FOR_EACH_LLC(&iter, topo)
		scx_bitmap_print(topo->mask);
	bpf_printk("\n");

	bpf_printk("CORE MASKS");
	TOPO_FOR_EACH_CORE(&iter, topo)
		scx_bitmap_print(topo->mask);
	bpf_printk("\n");

	bpf_printk("CPU MASKS");
	TOPO_FOR_EACH_CPU(&iter, topo)
		scx_bitmap_print(topo->mask);
	bpf_printk("\n");

	return 0;
}