add argmax lowering and integer return types

[dan-garvey]

Add argmax lowering from torch to linalg.

Borrows heavily from the tosa lowering.

Has an issue at the moment, when I run something like:

npcomp-opt sample_argmax.mlir -torchscript-to-npcomp-backend-pipeline

I get:
sample_argmax.mlir:6:10: error: 'std.constant' op requires integer result types to be signless %1 = torch.aten.argmax %arg1, %int0, %true_0 : !torch.tensor, !torch.int, !torch.bool -> !torch.tensor ^ sample_argmax.mlir:6:10: note: see current operation: %6 = "std.constant"() {value = 0 : si64} : () -> si64

I struggled for awhile to realize that I had to make a case in refinetypes that changed the output dtype to int, but I'm not sure if I did that correctly. I'm probably missing something really obvious, but I'm not sure how to address this error.

[silvasean]

You are the first to step on the landmine of integer dtypes in TorchToLinalg :) You will need to add a case for it in setupValueTensorToBuiltinTensorConversion which transitively calls into ValueTensorType::toBuiltinTensor -- that function needs to be updated to translate the dtypes to builtin integer types. Specifically, float types get passed through, and integer types get converted to a signless integer version of the same bitwidth.

[cathyzhyi]

Looks good! Just need to delete a few lines of redundant code.

[cathyzhyi]

I just realize similar to the OpConversionPattern we can't assume else block is always None. It could be variable. For integer variable, we can't known statically the sizes of each dim so use:

else if (op.dim().getType().isa<IntegerType>())
     knowledge.sizes.resize(keepdimBool ? inputRank : inputRank - 1,
                             kUnknownSize);

For NoneType case, nothing is needed as knowledge.hasSizes = true; is set above already. It's nice to keep comment about None case though.

[dan-garvey]

Sorry, I should have caught this after your other comment. Thank you for the very thorough review

[cathyzhyi]

@dan-garvey Hey Daniel, sry that the op.dim().getType().isa<IntegerType>() above was wrong it should be Torch::IntType. We need to check against torch type. I will fix this as part of my PR.