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 + using namespace mlir; using namespace mlir::torch; using namespace mlir::torch::Torch; @@ -1702,6 +1704,20 @@ void Aten__Getitem__TOp::getCanonicalizationPatterns( rewriter.replaceOpWithNewOp(op, sizeOp.self(), op.idx()); return success(); }); + + patterns.add(+[](Aten__Getitem__TOp op, PatternRewriter &rewriter) { + auto potentialUnbindOp = op.getOperand(0); + auto unbindOp = potentialUnbindOp.getDefiningOp(); + if (!unbindOp) + return failure(); + std::cout << "XXX: Found unbind op!" << std::endl; + + auto tensor = unbindOp.getOperand(0); + //rewriter.replaceOpWithNewOp(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)")