Folder and Canonicalizer for PrimsConvertElementTypeOp and AtenMaxPool2dWithIndicesOp

[zezhang]

While playing with TorchDynamo on ResNet18. I notice following issues:

• prims.convert_element_type can’t be canonicalized even if the input and the output share the same type

• aten.max_pool2d_with_indices is always used instead of aten.max_pool2d, even if the second returned output (indices) has no user

This PR fixes above issues by adding a folder to the PrimsConvertElementTypeOp and a canonicalizer to the AtenMaxPool2dWithIndicesOp

Lit test:

cmake --build build --target check-torch-mlir-all

[zezhang]

Looks like adding a folder for PrimsConvertElementTypeOp makes the linalg PrimsConvertElementTypeModule_basic test fail.

repro: python -m e2e_testing.main --config=linalg -v -f 'PrimsConvertElementTypeModule_basic'

Because the input and output type could be a simple torch.tensor (no element dtype is available). We will have to use the dtype argument of the op to specify the output type. Any idea on this? @sjain-stanford

[sjain-stanford]

Because the input and output type could be a simple torch.tensor (no element dtype is available). We will have to use the dtype argument of the op to specify the output type.

I think you want to prevent it from being folded when element dtype is unknown. I haven't validated it but something like this should work:

OpFoldResult PrimsConvertElementTypeOp::fold(FoldAdaptor adaptor) {
  auto inputType = cast<BaseTensorType>(getA().getType());
  auto outputType = cast<BaseTensorType>(getResult().getType());
  if (inputType != outputType)
    return nullptr;
  if (!inputType.hasDtype() || !outputType.hasDtype())
    return nullptr;
  if (inputType.getDtype() != outputType.getDtype())
    return nullptr;
  return getA();
}

[sjain-stanford]

Good news is your PR fixes 4 tests which previously failed. You may update the xfail set.

  ****** Unexpectedly Passed tests - 4 tests
      XPASS - "MaxPool2dEmptyStrideStaticModule_basic"
      XPASS - "MaxPool2dStaticCeilModeTrueModule_basic"
      XPASS - "MaxPool2dStaticModule_basic"
      XPASS - "ResNet18StaticModule_basic"

[sjain-stanford]

Feel free to land once CI is clean.