From bb5876e9f6f9029d78b6e9374dc6f30a7d5bd544 Mon Sep 17 00:00:00 2001
From: Daniel Ellis <1346302+dellis23@users.noreply.github.com>
Date: Tue, 18 Oct 2022 20:47:46 +0000
Subject: [PATCH] Attempt to rewrite unbind / getitem combo.

---
 .../Dialect/Torch/IR/GeneratedTorchOps.td     | 23 +++++++++++++++++++
 lib/Dialect/Torch/IR/TorchOps.cpp             | 16 +++++++++++++
 .../jit_ir/build_tools/torch_ods_gen.py       |  1 +
 3 files changed, 40 insertions(+)

diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
index 55e29441f..e2ec90c85 100644
--- a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
+++ b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
@@ -6442,6 +6442,29 @@ def Torch_AtenTypeAsOp : Torch_Op<"aten.type_as", [
   let hasFolder = 1;
 }
 
+def Torch_AtenUnbindIntOp : Torch_Op<"aten.unbind.int", [
+    AllowsTypeRefinement,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::unbind.int : (Tensor, int) -> (Tensor[])`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    Torch_IntType:$dim
+  );
+  let results = (outs
+    AnyTorchListOfTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenUnbindIntOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void AtenUnbindIntOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+}
+
 def Torch_AtenViewOp : Torch_Op<"aten.view", [
     AllowsTypeRefinement,
     ReadOnly
diff --git a/lib/Dialect/Torch/IR/TorchOps.cpp b/lib/Dialect/Torch/IR/TorchOps.cpp
index e8e52f410..fe11b5271 100644
--- a/lib/Dialect/Torch/IR/TorchOps.cpp
+++ b/lib/Dialect/Torch/IR/TorchOps.cpp
@@ -20,6 +20,8 @@
 #include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Casting.h"
 
+#include <iostream>
+
 using namespace mlir;
 using namespace mlir::torch;
 using namespace mlir::torch::Torch;
@@ -1702,6 +1704,20 @@ void Aten__Getitem__TOp::getCanonicalizationPatterns(
     rewriter.replaceOpWithNewOp<AtenSizeIntOp>(op, sizeOp.self(), op.idx());
     return success();
   });
+
+  patterns.add(+[](Aten__Getitem__TOp op, PatternRewriter &rewriter) {
+    auto potentialUnbindOp = op.getOperand(0);
+    auto unbindOp = potentialUnbindOp.getDefiningOp<AtenUnbindIntOp>();
+    if (!unbindOp)
+      return failure();
+    std::cout << "XXX: Found unbind op!" << std::endl;
+
+    auto tensor = unbindOp.getOperand(0);
+    //rewriter.replaceOpWithNewOp<Aten__Getitem__TOp>(op, tensor, op.getOperand(1));
+    rewriter.replaceOp(op, {tensor, op.getOperand(1)});
+    return success();
+  });
+
 }
 
 //===----------------------------------------------------------------------===//
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
index 9f3a601b8..996678ac4 100644
--- 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
@@ -477,6 +477,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
     emit("aten::to.prim_Device : (Tensor, Device?, int?, bool, bool) -> (Tensor)")
     emit("aten::to.device : (Tensor, Device, int, bool, bool, int?) -> (Tensor)")
     emit("aten::type_as : (Tensor, Tensor) -> (Tensor)", has_folder=True)
+    emit("aten::unbind.int : (Tensor, int) -> (Tensor[])")
     emit("aten::view : (Tensor, int[]) -> (Tensor)", has_folder=True)
     emit("aten::_unsafe_view : (Tensor, int[]) -> (Tensor)")
     emit("aten::where.self : (Tensor, Tensor, Tensor) -> (Tensor)")