[TorchToTosa] Refine argmax integer casting and update tests

[catcor01]

• Cast only i32/i64 inputs to f32 before tosa.argmax
• Update argmax e2e test to exercise integer casting

Change-Id: I5f8847cc7400152c20001905dfacc92af0c1583c

[catcor01]

@sahas3 @Lallapallooza

[sahas3]

Thanks for the change @catcor01. The code changes look fine to me but can you clarify why the cast to fp32 is required?

[catcor01]

Thanks for the change @catcor01. The code changes look fine to me but can you clarify why the cast to fp32 is required?

The cast is there to satisfy the TOSA type constraints for tosa.argmax. In the base TOSA floating‑point profile, argmax only accepts fp16/fp32 inputs, so int32/int64 isn’t legal so we normalize to fp32 to keep the op spec‑compliant across targets.

EDIT: I have removed reference to bf16 casting now in the patch as it should be targeted through the enabling of the BF15 extension set in TOSA which should be in a separate MR i believe.

[sahas3]

Thanks for the clarification.

I agree with enhancing the bf16 support as part of a separate PR. For hardware targets that can support bf16, there should be no need to cast it to f32.

[Lallapallooza]

LGTM: small nit

[Lallapallooza]

Could we rename ensureF32Input?

The current name implies "always convert to f32", but the function actually leaves floats and i8 unchanged and only casts some integer types. Something like legalizeArgMaxInputType or normalizeArgMaxInputForTosa may be clearer.