[TIR] Ramp and Broadcast lanes fixed to int32 dtype

[quic-sanirudh]

When Ramp and Broadcast nodes are created with fixed length lanes,
they're fixed to int32 dtype since DLDataType always supports only
uint16 lanes.

[tqchen]

In this particular case (deep equality), i think type do matter, so it would be great instead to fix the cases that would depend on the relaxed behavior.

I know we had some i64/i32 issues, and general rule of thumb now is to try to be explicit as much as possible and that helps to reduce errors

[quic-sanirudh]

In this particular case (deep equality), i think type do matter, so it would be great instead to fix the cases that would depend on the relaxed behavior.

I know we had some i64/i32 issues, and general rule of thumb now is to try to be explicit as much as possible and that helps to reduce errors

Oh okay, thanks for the feedback @tqchen. The cases we started seeing was that some expressions were not getting simplified properly after RampNode lanes were changed to PrimExpr. I narrowed down the exact simplification that was failing was actually a rewrite simplify rule here.

I realized that the simplification was not happening because lanes between broadcast and RampNode in my case had different types, so a couple other solutions I thought would apply here is to either fix the RampNode constructor to stick to some fixed dtype for lanes (something like int32/int16, since lanes in DLDataType anyways only supports uint16 dtype), or to update the simplify rules here to try the same rules with an PVar<IntImm> lanes type in case of fixed length vectors.

But if dtype does matter, then should we update the PEqualChecked<IntImm> to also check for dtypes?

[tqchen]

ah oK, i think in this case we should try to come up with a rule for lanes. In this particular case i think it was related to SVE changes, so also cc @ekalda

I think having a fixed dtype probably makes sense then we handle cast for related cases (keep i32 and i64 iterators in mind

[ekalda]

Yeah I think fixing the dtype is a good idea, it would hopefully avoid this kind of problems in the future as well. Out of interest, what were the mismatching dtypes of the two compared IntImmNodes that you observed @quic-sanirudh?

[quic-sanirudh]

Yeah I think fixing the dtype is a good idea, it would hopefully avoid this kind of problems in the future as well. Out of interest, what were the mismatching dtypes of the two compared IntImmNodes that you observed @quic-sanirudh?

Thanks @ekalda. I'll update the PR to fix the dtypes in RampNode (and perhaps the broadcast node as well).

The dtypes in my case were int32 and int64. The expression I saw was something like this (slightly simpler version)
T.Broadcast(c, 128) + T.Ramp(T.int64(0), T.int64(1), T.int64(128))

The RampNode seems to get the int64 lanes because the all the iterators in our case is by default int64, but the broadcast seems to be inserted during the evaluation of AddNode in op.cc here

[quic-sanirudh]

I've updated the PR to fix the dtype of lanes as we discussed. Let me know if this looks good.

[ekalda]

Looks good to me, thank you for the fix @quic-sanirudh! Hopefully we won't see these kind of silent simplification failures in the future now 🤞

[quic-sanirudh]

Thanks for the review @ekalda

[tqchen]

@quic-sanirudh can you add a test case about i64 iterator type with Ramp? I know fixing to i32 likely will resolve issues when most programs are i32 iterator, we should make sure it is i64 compatible

[tqchen]

I actually meant that the value being broadcasted and ramp base being i64. but lanes remains i32, this would be a more common case in our setting

[quic-sanirudh]

ah okay, sorry for the confusion. I'll update the test cases to check that as well.

[quic-sanirudh]

@tqchen I've updated the test case, let me know if this looks good.

[quic-sanirudh]

@tvm-bot rerun