Attempt to rewrite unbind / getitem combo.

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

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();
+  });
+
 }
 
 //===----------------------------------------------------------------------===//

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)")