Implement folder for unused empty.memory_formats (#139)

Xilinx · Sep 8, 2023 · de998c0 · de998c0
1 parent 75ec45a
commit de998c0
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 1 deletion.
diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
@@ -7348,6 +7348,7 @@ def Torch_AtenEmptyMemoryFormatOp : Torch_Op<"aten.empty.memory_format", [
       printDefaultTorchOp(printer, *this, 6, 1);
     }
   }];
+  let hasCanonicalizer = 1;
 }
 
 def Torch_AtenExpandOp : Torch_Op<"aten.expand", [

diff --git a/lib/Dialect/Torch/IR/TorchOps.cpp b/lib/Dialect/Torch/IR/TorchOps.cpp
@@ -1337,6 +1337,21 @@ static OpFoldResult intComparatorFoldHelper(OpTy op,
 
 OpFoldResult AtenDetachOp::fold(FoldAdaptor adaptor) { return getSelf(); }
 
+//===----------------------------------------------------------------------===//
+// AtenEmptyMemoryFormatOp
+//===----------------------------------------------------------------------===//
+
+void AtenEmptyMemoryFormatOp::getCanonicalizationPatterns(
+    RewritePatternSet &patterns, MLIRContext *context) {
+  patterns.add(+[](AtenEmptyMemoryFormatOp op, PatternRewriter &rewriter) {
+    if (!op->use_empty()) {
+      return failure();
+    }
+    rewriter.eraseOp(op);
+    return success();
+  });
+}
+
 //===----------------------------------------------------------------------===//
 // AtenNeIntOp
 //===----------------------------------------------------------------------===//

diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py b/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py
@@ -494,7 +494,7 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::new_empty_strided : (Tensor, int[], int[], int?, int?, Device?, bool?) -> (Tensor)")
     emit("aten::zeros_like : (Tensor, int?, int?, Device?, bool?, int?) -> (Tensor)")
     emit("aten::ones_like : (Tensor, int?, int?, Device?, bool?, int?) -> (Tensor)")
-    emit("aten::empty.memory_format : (int[], int?, int?, Device?, bool?, int?) -> (Tensor)")
+    emit("aten::empty.memory_format : (int[], int?, int?, Device?, bool?, int?) -> (Tensor)", has_canonicalizer=True)
     emit("aten::expand : (Tensor, int[], bool) -> (Tensor)")
     emit("aten::expand_as : (Tensor, Tensor) -> (Tensor)")
     emit("aten::broadcast_to : (Tensor, int[]) -> (Tensor)", has_canonicalizer=True)
-Original file line number
+Diff line change
@@ Expand Up @@
           printDefaultTorchOp(printer, *this, 6, 1);
         }
       }];
+      let hasCanonicalizer = 1;
     }
     def Torch_AtenExpandOp : Torch_Op<"aten.expand", [
@@ Expand Down @@