[TIR] enhance tir signed-unsigned cast

[ganler]

This is related to an issue here. In a word, when == and != binary operators are applied to tir.const(1) and tir.const(True), the return value will be an "EQ" expression instead of a folded IntImm variable (e.g., tir.const(True) after constant folding). This issue will bring some limitations when integrating TIR to some python native operations (e.g., ==, !=, python sets, etc.).

To reproduce:

from tvm import tir
assert tir.const(1) == tir.const(True)
#  Or: assert tir.const(1, "int32").__eq__(tir.const(1, "uint1")).__bool__()
# `uint1` means `bool`

The root cause is that during BinaryOpMatchTypes (https://github.com/apache/tvm/blob/main/src/tir/op/op.cc#L145), singed and unsigned variable are not compatible and IntImmNode would be converted to CastNode. However, TVM's constant folding did not support CastNode so that tir.const(1, "int32").__eq__(tir.const(1, "uint1")) would return an expression instead of another folded IntImm variable.

This PR fixed this problem by enhancing the unsigned-signed casting:

If the bits of signed type is greater than that of unsigned type, it means that signed type has a bigger value range than that of unsigned type. In this case, we can promote the unsigned type to a signed type so that there's no need to employ CastNode.

A more thorough enhancement would be supporting constant folding for CastNode.

@YuchenJin @junrushao1994 @tqchen

[junrushao]

@YuchenJin would you like to shepherd this PR? Thanks a lot!

[YuchenJin]

@YuchenJin would you like to shepherd this PR? Thanks a lot!

Of course! Will take a deeper look later today. :)

[YuchenJin]

Hi @ganler, Overall LGTM! Thanks for elaborating on the issue and implementing the fix!

One question is that whether we want to also make assert tir.const(1, "uint2") == tir.const(1, "int2") work? I think since lhs and rhs have the same bits, SimpleCast will be applied to both sides, but the same error throws.

[ganler]

@YuchenJin Good question to discuss.

As I said a more thorough fix would be supporting CastNode in constant folding, but it can be tricky. In this PR, I only made a safe/harmless change by casting lower-bit uint constant value to higher-bit int.

However, when we have, say int32 and uint32, the question is that how should we perform casting? i.e., what is the logic of op(int32(lhs), uint32(rhs))? In TVM, we will cast all types to a signed integer with maximum bits (regarding the operands). i.e., op(int32(lhs), int32(uint32(rhs))). But here, int32(uint32(rhs)) is not IntImm but CastNode... Then your question would be: shall we directly evaluate the CastNode during casting (which requires changing the value of IntImm).

[YuchenJin]

@ganler Thanks for going through the detail! I am not quite sure whether we want to support evaluating CastNode now, maybe we can add this support when there are use cases. I think this PR can be merged. Thanks for the effort!

[Johnson9009]

@ganler @YuchenJin @tqchen I happens to file a PR relevant to this change, after reading all of your preceding discussing, I think there maybe have something not clear.

From reading the PR, I found, for the test case here tir.const(1, "int32").__eq__(tir.const(1, "uint1")), the old behavior of BinaryOpMatchTypes is int32 == uint1 -> int32 == int32(uint1), it is exactly same with the job that this PR does.

Then I feel it is strange why this PR and my own PR can make assert tir.const(1) == tir.const(True) work, after debug it step-by-step I found the reason is the function we use "cast" instead the "tir::Cast" used by "SimpleCast".

tvm/src/tir/op/op.cc

Lines 273 to 284 in a06863a

	PrimExpr cast(const DataType& t, PrimExpr value, Span span) {
	using tir::FloatImmNode;
	if (value.dtype() == t) return value;
	// const fold IntImm as they are used in index computations
	if (t.lanes() == 1) {
	if (const IntImmNode* op = value.as<IntImmNode>()) {
	return make_const(t, op->value, op->span);
	} else if (const FloatImmNode* op = value.as<FloatImmNode>()) {
	return make_const(t, op->value, op->span);
	}
	return tir::Cast(t, value, span);
	} else {

We can see it is the logic around line 279 that eliminate the CastNode, so the function cast is do some constant fold optimization.

Hi @ganler, Overall LGTM! Thanks for elaborating on the issue and implementing the fix!

One question is that whether we want to also make assert tir.const(1, "uint2") == tir.const(1, "int2") work? I think since lhs and rhs have the same bits, SimpleCast will be applied to both sides, but the same error throws.

@YuchenJin For the scenario that unsigned and signed integer have the same width, I think we should consistent with C standard (e.g., "int32 + uint32" -> "uint32(int32) + uint32") instead of current logic "int32 + uint32" -> "int32 + int32(uint32)", please have a look of my PR #8733, we implement a new AI chip backend of TVM, now we meet this issue, we need make a pass to change the implicit conversion consistent with C standard.

[YuchenJin]

@Johnson9009 Thanks for the effort! I will read your PR tomorrow.