fix when binary search find invalid value

[qelk123]

This is a simple fix up for the situation when binary searh find invalid indices,which contains:
1.An IfThenElse stmt after binary search,if the index is invalid,then return -1 as an invalid index value.
2.add a postprocess pass after lower sparse buffer,for those sparse data buffer,if the index value is a invalid value,then return a 0 from the sparse data buffer.

[yzh119]

Good job @qelk123 !
Thanks for your contribution, it's a behavior I should have fixed but never did, I'll leave my feedback later.

[qelk123]

Dear Zihao Ye, Thank you for your reply.This is the first time I propose a PR,so although this modification works fine in my case,I believe there are inconsiderate parts in my code.Please give me your feedback so that I could refine it later. Michael

[yzh119]

Overall a nice contribution, I have a few comments.

Please format the code per suggestions, that would not take too much of your time.

I suggest adding a binary_search_vaild_check key in block annotation, indicating whether the user wants to perform a such check or not. If set to false, then we will examine whether the search is invalid or not, to reduce some overhead.

[yzh119]

The data types are not only float32 and int32.
0 is also not only the valid choice as initial values, suppose we want to perform max reduction, -inf would be a better choice.

[yzh119]

We can let user specify their desired initial values as block annotations, if such value is not set by user, we turn to use 0 instead.

[qelk123]

What kind of data types do we need to support currently(besides float and int)? Type info is needed when the user doesn't offer a default value,so we need to set the default value to 0 with a specific type according to the dtype of the sparse buffer.

[yzh119]

We might not need to emit this IfThenElse if we can confirm the lookup will not fail.

[yzh119]

Do we really need to store BufferLoad? How about we only store Buffer?

[qelk123]

BufferLoad here is to construct the if cond expr，if we only store the buffer how should we get the indices of the load buffer?

[yzh119]

I had a proposal months before: #74, the really tricky part is INVALID flag propagation:

Suppose lookup_result stores the binary search results, and we then use this value as an index to access another buffer: A[lookup_result[...]], the whole expression should be marked as invalid. However, I accept your solution for now and we can improve it later on.

[qelk123]

From my perspective，even the binary search results is not a valid index in the indices buffer,the buffer load result like A[lookup_result[...]] probably should not directly return a invalid flag,but a default value about the sparse buffer.In TACO,this default value can be spcified when declare a tensor.
In an Union operation like add,when we do co-iteration on two sparse axis,as long as current index is valid in one axis,the result in related index of the output tensor should be set to the correct value，instead of invalid value.
Currently,it only support 0 as the default value,but maybe we can improve it by letting users to specify the default value later on?

[yzh119]

From my perspective，even the binary search results is not a valid index in the indices buffer,the buffer load result like A[lookup_result[...]] probably should not directly return a invalid flag,but a default value about the sparse buffer.In TACO,this default value can be spcified when declare a tensor.

Yes we should return a default value, what do I mean by "invalid flag" is a indicator that whole expression containing it should be invalidated and return its default value.

Back to my example A[lookup_result[...]], if we return the default value 0 for an unsuccessful lookup_result, then we will get A[0], but what we really need is default value for buffer A, and such information need to be propagated in some way.

We can leave it for another PR.

[yzh119]

It looks much better than its initial form, thank you so much, I have some suggestions for further improvements.

[qelk123]

From my perspective，even the binary search results is not a valid index in the indices buffer,the buffer load result like A[lookup_result[...]] probably should not directly return a invalid flag,but a default value about the sparse buffer.In TACO,this default value can be spcified when declare a tensor.

Yes we should return a default value, what do I mean by "invalid flag" is a indicator that whole expression containing it should be invalidated and return its default value.

Back to my example A[lookup_result[...]], if we return the default value 0 for an unsuccessful lookup_result, then we will get A[0], but what we really need is default value for buffer A, and such information need to be propagated in some way.

We can leave it for another PR.

Actually, current solution is let lookup_result[...] return -1 as an invalid index.Do you mean return "-1" is not a proper way to represent invalid index?

[yzh119]

In other original implementation we are using half open interval ([a, b)) semantics, please check our loop invariant(

SparseTIR/src/tir/transforms/lower_sparse_iter.cc

Lines 830 to 844 in ca59cbe

	/* Algorithm:
	* - when left = true
	* - pre-condition
	* lb < ub, and the last dimension of buf is sorted.
	* - loop-invariant
	* low <= mid < high, buf[..., lb:low] < val, buf[..., high:ub] >= val
	* - post-condition
	* low = mid = high, buf[..., lb:low] < val, buf[..., high:ub] >= val
	* - when left = false
	* - pre-condition
	* lb < ub, and the last dimension of buf is sorted.
	* - loop-invariant
	* low <= mid < high, buf[..., lb:low] <= val, buf[..., high:ub] > val
	* - post-condition
	* low = mid = high, buf[..., lb:low] <= val, buf[..., high:ub] > val

). So I suppose this change is not necessary?

[qelk123]

SparseTIR/src/tir/transforms/lower_sparse_iter.cc

Lines 828 to 829 in ca59cbe

	PrimExpr BinarySearch(SparseBuffer buf, Array<PrimExpr> prefix_indices, PrimExpr lb, PrimExpr ub,
	PrimExpr val, bool left, bool minus_one = false) {

In the function you set the upper bound as indptr[i+1]-indptr[i]
from these line

SparseTIR/src/tir/transforms/lower_sparse_iter.cc

Lines 923 to 925 in ca59cbe

	PrimExpr pivot_cmp_cond = left ? (pivot < val) : (pivot > val);
	Stmt if_true = left ? BufferStore(low, mid_val + 1, {Integer(0)})
	: BufferStore(high, mid_val, {Integer(0)});

we can see if the val is greater than indices[i,indptr[i+1]-indptr[i]-1],it's possable that the mid buffer we return contains the value of indices[i,indptr[i+1]-indptr[i]],which is wrong.

[qelk123]

What's more,if the val of index just equal to indices[i,indptr[i+1]-indptr[i]] then it will return a valid index,and cause the wrong result.

[yzh119]

Our semantic follows the lower_bound and upper_bound API in both C++ (https://en.cppreference.com/w/cpp/algorithm/lower_bound and https://en.cppreference.com/w/cpp/algorithm/upper_bound) and Numpy: (https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html).
Both of them would return the ub when the key being searched is out of bound, which aligns with our implementation. Note that none of lower_bound nor upper_bound are designed for exact match.

I know the purpose of your J_indptr_data[vi + 1] - J_indptr_data[vi] - 1 change is to avoid the case where the returned value in mid is J_indptr_data[vi + 1] - J_indptr_data[vi] and we want to use mid buffer to index original buffer for valid check. However, such change would break our loop invariant.

A better idea is not to check the validity in caller side, but creating another buffer valid to store the state whether the check is valid or not when during the binary search. I'll create another PR for that and you can rebase on it.

[qelk123]

OK，I can fix this later.For simplicity, I will probably add a judgment condition when deciding whether set a invalid index in mid

[yzh119]

I have created the PR: #96 let me know if it looks good to you.

[qelk123]

It looks like you choose the same way (by adding a judgment condition) to fix this problem like me in my newest commit

Stmt then_stmt = BufferStore(mid, -1, mid_indices);
PrimExpr if_stmt = (pivot != val || mid_val == ub);
body_stmts.push_back(IfThenElse(if_stmt, then_stmt));

However,I don't know why you return a pair in binary_search and it seems the second item only contain the valid flag instead of exact value.Do you mean we should leave the valid check outside the binary search function?

[yzh119]

Sorry I didn't notice your commit.
The valid flag is also returned to tell the caller of BinarySearch function whether the search is successful, so that we don't need to access buf[..., mid[...]] outside the binary search body.

We don't need exact value because the valid flag indicate such information, if the search is successful then the value equals buf[..., ,mid[...]].

[qelk123]

Hi @qelk123 do you need any assistance on rebase? It seems the codebase disappears and the PR is closed.

It's OK ,I will open a new pr later.

[yzh119]

I suppose this annotation would not be emitted because we will only emit "check_invalid_binary_search": True when invalid binary search check is enabled.

[qelk123]

actually,check_invalid_binary_search is a sign to tell the lower_sparse_buffer pass whether we should consider invalid index created by the lower_iter_pass.This is to reduce the overhead.Otherwise, a sign should add to the lower_sparse_buffer pass in other ways.

[yzh119]

I mean we can choose not to print this annotation when check_invalid_binary_search = False (not enabled).

[yzh119]

From my perspective，even the binary search results is not a valid index in the indices buffer,the buffer load result like A[lookup_result[...]] probably should not directly return a invalid flag,but a default value about the sparse buffer.In TACO,this default value can be spcified when declare a tensor.

Yes we should return a default value, what do I mean by "invalid flag" is a indicator that whole expression containing it should be invalidated and return its default value.
Back to my example A[lookup_result[...]], if we return the default value 0 for an unsuccessful lookup_result, then we will get A[0], but what we really need is default value for buffer A, and such information need to be propagated in some way.
We can leave it for another PR.

Actually, current solution is let lookup_result[...] return -1 as an invalid index.Do you mean return "-1" is not a proper way to represent invalid index?

We support non-affine indices, and the buffer access indices could be value loaded from another buffer access (e.g. A[B[C[...]]]), my example shows that we need to return default value for the entire expression A[B[C[...]]] if binary search is invalid, not only C[...] or B[C[...]].

[qelk123]

From my perspective，even the binary search results is not a valid index in the indices buffer,the buffer load result like A[lookup_result[...]] probably should not directly return a invalid flag,but a default value about the sparse buffer.In TACO,this default value can be spcified when declare a tensor.

Yes we should return a default value, what do I mean by "invalid flag" is a indicator that whole expression containing it should be invalidated and return its default value.
Back to my example A[lookup_result[...]], if we return the default value 0 for an unsuccessful lookup_result, then we will get A[0], but what we really need is default value for buffer A, and such information need to be propagated in some way.
We can leave it for another PR.

Actually, current solution is let lookup_result[...] return -1 as an invalid index.Do you mean return "-1" is not a proper way to represent invalid index?

We support non-affine indices, and the buffer access indices could be value loaded from another buffer access (e.g. A[B[C[...]]]), my example shows that we need to return default value for the entire expression A[B[C[...]]] if binary search is invalid, not only C[...] or B[C[...]].

I thought when facing the situation like A[B[C[...]]] if the binary search is failed in C,then C[...] will return a default value,and we can use this value to index B,the expression is still available?

[yzh119]

From my perspective，even the binary search results is not a valid index in the indices buffer,the buffer load result like A[lookup_result[...]] probably should not directly return a invalid flag,but a default value about the sparse buffer.In TACO,this default value can be spcified when declare a tensor.

Yes we should return a default value, what do I mean by "invalid flag" is a indicator that whole expression containing it should be invalidated and return its default value.
Back to my example A[lookup_result[...]], if we return the default value 0 for an unsuccessful lookup_result, then we will get A[0], but what we really need is default value for buffer A, and such information need to be propagated in some way.
We can leave it for another PR.

Actually, current solution is let lookup_result[...] return -1 as an invalid index.Do you mean return "-1" is not a proper way to represent invalid index?

We support non-affine indices, and the buffer access indices could be value loaded from another buffer access (e.g. A[B[C[...]]]), my example shows that we need to return default value for the entire expression A[B[C[...]]] if binary search is invalid, not only C[...] or B[C[...]].

I thought when facing the situation like A[B[C[...]]] if the binary search is failed in C,then C[...] will return a default value,and we can use this value to index B,the expression is still available?

I believe you're correct. Let's stick with this solution until rare cases.

[yzh119]

Hi @qelk123 do you need any assistance on rebase? It seems the codebase disappears and the PR is closed.

[yzh119]

Thank you @qelk123, this version looks good to me!