Fold __is__ and unchecked_cast of derefine

The added e2e maxpool testcase from #545 was not getting a static shape
due to an unfolded prim.If when RefineTypes was called. This was because
of unfolded torch.iaten.__is__ and torch.prim.unchecked_cast operators
with torch.derefine operands.

Author: ljfitz

e2e_testing/torchscript/basic.py

@@ -268,6 +268,28 @@ def forward(self, x):
 def MaxPool2dModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(1, 1, 20, 20) - 0.5)
 
+class MaxPool2dStaticModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mp2d = torch.nn.MaxPool2d(kernel_size=[3, 3],
+                                       stride=[2, 2],
+                                       padding=[1, 1],
+                                       dilation=[1, 1])
+
+    @export
+    @annotate_args([
+        None,
+        ([1, 64, 112, 112], torch.float32, True),
+    ])
+    def forward(self, x):
+        return self.mp2d(x)
+
+
+@register_test_case(module_factory=lambda: MaxPool2dStaticModule())
+def MaxPool2dStaticModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(1, 64, 112, 112))
+
+# ==============================================================================
 
 class ConstantPad2dStaticModule(torch.nn.Module):
     def __init__(self):

include/torch-mlir/Dialect/Torch/IR/GeneratedPrimOps.td

@@ -201,6 +201,7 @@ def Torch_PrimUncheckedCastOp : Torch_Op<"prim.unchecked_cast", [
     AnyTorchType:$result
   );
   let assemblyFormat = "$x attr-dict `:` qualified(type($x)) `->` qualified(type($result))";
+  let hasFolder = 1;
 }
 
 def Torch_PrimPrintOp : Torch_Op<"prim.Print", [

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -390,12 +390,19 @@ bool DerefineOp::areCastCompatible(mlir::TypeRange inputs,
 
 template <typename OpTy>
 static OpFoldResult atenIsOrIsNotFoldHelper(OpTy op, bool equalIsTrue) {
-  Type lhsType = op.self().getType();
-  Type rhsType = op.obj().getType();
+  Value lhs = op.self();
+  Value rhs = op.obj();
 
-  // If either type is a NoneType, make it be the lhsType.
-  if (rhsType.template isa<Torch::NoneType>())
-    std::swap(lhsType, rhsType);
+  // If either value is typed NoneType, make it be the lhs.
+  if (rhs.getType().template isa<Torch::NoneType>())
+    std::swap(lhs, rhs);
+
+  if (rhs.getType().template isa<Torch::OptionalType>())
+    if (auto derefine = rhs.getDefiningOp<Torch::DerefineOp>())
+      rhs = derefine.operand();
+
+  Type lhsType = lhs.getType();
+  Type rhsType = rhs.getType();
   // TODO: Implement and use subtype infra for this.
   // If neither type is a subtype of the other, then the result is false.
   if (lhsType.template isa<Torch::NoneType>() &&
@@ -986,6 +993,14 @@ bool PrimUncheckedCastOp::areCastCompatible(mlir::TypeRange inputs,
   return isValidSubtype(outputs[0], inputs[0]);
 }
 
+OpFoldResult PrimUncheckedCastOp::fold(ArrayRef<Attribute> operands) {
+  if (auto derefineOp = x().getDefiningOp<Torch::DerefineOp>()) {
+    if (derefineOp.operand().getType() == getType())
+      return derefineOp.operand();
+  }
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // Aten__Getitem__TOp
 //===----------------------------------------------------------------------===//

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -416,7 +416,7 @@ def emit(key, **kwargs):
         emit("prim::max.int : (int, int) -> (int)")
         emit("prim::RaiseException : (str, str?) -> ()")
         emit("prim::Uninitialized : () -> (Any)", traits=["NoSideEffect"])
-        emit("prim::unchecked_cast : (t) -> (t)",
+        emit("prim::unchecked_cast : (t) -> (t)", has_folder=True,
              traits=["DeclareOpInterfaceMethods<CastOpInterface>"])
         emit("prim::Print : (...) -> ()")
         emit("prim::tolist : (...) -> (...)")

test/Dialect/Torch/canonicalize.mlir

@@ -8,6 +8,15 @@ func @torch.aten.__is__(%arg0: !torch.list<!torch.int>, %arg1: !torch.none) -> !
   return %0 : !torch.bool
 }
 
+// CHECK-LABEL:   func @torch.aten.__is__$derefine_is_none
+// CHECK:           %[[FALSE:.*]] = torch.constant.bool false
+// CHECK:           return %[[FALSE]] : !torch.bool
+func @torch.aten.__is__$derefine_is_none(%arg0: !torch.list<!torch.int>, %arg1: !torch.none) -> !torch.bool {
+  %0 = torch.derefine %arg0 : !torch.list<!torch.int> to !torch.optional<!torch.list<!torch.int>>
+  %1 = torch.aten.__is__ %0, %arg1 : !torch.optional<!torch.list<!torch.int>>, !torch.none -> !torch.bool
+  return %1 : !torch.bool
+}
+
 // CHECK-LABEL:   func @torch.aten.__is__$none_is_none
 // CHECK:           %[[TRUE:.*]] = torch.constant.bool true
 // CHECK:           return %[[TRUE]] : !torch.bool
@@ -644,6 +653,13 @@ func @torch.prim.TupleIndex$out_of_bound(%t0: !torch.tensor, %t1: !torch.tensor,
     return %1 : !torch.tensor
 }
 
+// CHECK-LABEL:   func @torch.prim.unchecked_cast$derefine
+// CHECK-next:      return %arg0 : !torch.list<!torch.int>
+func @torch.prim.unchecked_cast$derefine(%arg0: !torch.list<!torch.int>) -> !torch.list<!torch.int> {
+  %0 = torch.derefine %arg0 : !torch.list<!torch.int> to !torch.optional<!torch.list<!torch.int>>
+  %1 = torch.prim.unchecked_cast %0 : !torch.optional<!torch.list<!torch.int>> -> !torch.list<!torch.int>
+  return %1 : !torch.list<!torch.int>
+}
 
 // CHECK-LABEL:   func @torch.aten.Int.Tensor(
 // CHECK-SAME:            %[[NUM:.*]]: !torch.int) -> !torch.int {