Add canonicalizer in torch_ods_gen.py + update error message

Author: sahas3

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -5708,8 +5708,8 @@ struct CanonicalizeAvgPoolWithSingleIntTuple
     // Attempt to expand params if necessary.
     if (failed(expandPoolParams(op, numSpatialDims, rewriter, kernel, stride,
                                 pad, dilations)))
-      return rewriter.notifyMatchFailure(op,
-                                         "Failed to expand params for pooling");
+      return rewriter.notifyMatchFailure(
+          op, "Failed to expand params for AvgPooling");
 
     rewriter.replaceOpWithNewOp<AvgPoolOpT>(
         op, op.getResult().getType(), op.getSelf(), kernel, stride, pad,
@@ -5736,8 +5736,8 @@ struct CanonicalizeMaxPoolWithSingleIntTuple
     // Attempt to expand params if necessary.
     if (failed(expandPoolParams(op, numSpatialDims, rewriter, kernel, stride,
                                 pad, dilations)))
-      return rewriter.notifyMatchFailure(op,
-                                         "Failed to expand params for pooling");
+      return rewriter.notifyMatchFailure(
+          op, "Failed to expand params for MaxPooling");
 
     rewriter.replaceOpWithNewOp<MaxPoolOpT>(op, op.getResult().getType(),
                                             op.getSelf(), kernel, stride, pad,

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -657,7 +657,10 @@ def emit_with_mutating_variants(key, **kwargs):
     emit(
         "aten::max_pool1d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)"
     )
-    emit("aten::max_pool2d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)")
+    emit(
+        "aten::max_pool2d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)",
+        has_canonicalizer=True,
+    )
     emit("aten::max_unpool2d : (Tensor, Tensor, int[]) -> (Tensor)")
     emit(
         "aten::max_pool2d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)",
@@ -666,7 +669,10 @@ def emit_with_mutating_variants(key, **kwargs):
     emit(
         "aten::max_pool2d_with_indices_backward : (Tensor, Tensor, int[], int[], int[], int[], bool, Tensor) -> (Tensor)"
     )
-    emit("aten::max_pool3d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)")
+    emit(
+        "aten::max_pool3d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)",
+        has_canonicalizer=True,
+    )
     emit("aten::max_unpool3d : (Tensor, Tensor, int[], int[], int[]) -> (Tensor)")
     emit(
         "aten::max_pool3d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)",
@@ -677,13 +683,15 @@ def emit_with_mutating_variants(key, **kwargs):
     )
     emit("aten::avg_pool1d : (Tensor, int[], int[], int[], bool, bool) -> (Tensor)")
     emit(
-        "aten::avg_pool2d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"
+        "aten::avg_pool2d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)",
+        has_canonicalizer=True,
     )
     emit(
         "aten::avg_pool2d_backward : (Tensor, Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"
     )
     emit(
-        "aten::avg_pool3d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"
+        "aten::avg_pool3d : (Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)",
+        has_canonicalizer=True,
     )
     emit(
         "aten::avg_pool3d_backward : (Tensor, Tensor, int[], int[], int[], bool, bool, int?) -> (Tensor)"