Explicit Branch for Halide version for an std::unordered_set/std::set

[skhiat]

Various algorithm need some unordered storing.

My attempt:

Expr	IsKeyExist(Func& keyArray, Expr key)
{
	RDom rDom{0, keyArray.output_buffer().width() };

	Tuple result = argmin(cast<int>(keyArray(rDom)) - cast<int>(key) == 0, "IsKeyExist");
	return select(result[1], true, false);
}
Expr	GetIndexFromKey(Func& keyArray, Expr key)
{
	RDom rDom{ 0, keyArray.output_buffer().width() };

	Tuple result = argmin(cast<int>(keyArray(rDom)) - cast<int>(key) == 0, "GetIndexFromKey");
	return select(result[1], result[0], -1);
}

The main issue is the "Insert" Function:

Expr	InsertKey(Func& keyArray, Expr key, Expr lastElementIndex)
{
	Expr index = GetIndexFromKey(keyArray, key);

	Expr newIndex = select(index != -1, lastElementIndex + 1, lastElementIndex);

	keyArray(newIndex) = key;

	return lastElementIndex;
}

If key the key exist we write multiple time with at the same index which is not allowed with Halide (keyArray(newIndex)).
Another issue is the lastElementIndex + 1 must be "atomic on the algorithmic side" not on the scheduling.
A proposition for a "fix" would be to have an explicit branching:

Expr	InsertKey(Func& keyArray, Expr key, Expr lastElementIndex)
{
	Expr index = GetIndexFromKey(keyArray, key);

	Halide::if_then_else(
		index != -1,
		[](){
			lastElementIndex = algorithmic_atomic_increment(lastElementIndex);
			keyArray(lastElementIndex) = key;
		},
		[](){
			unreachable();
		}
	};

	return lastElementIndex;
}

That's a proposition to face this situation, What would be the alternative?