llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Mark unbind as pure.

rewrite-getitem
Daniel Ellis 2022-10-25 18:51:49 +00:00
parent 1391bb6794
commit 0c5f209fb5
3 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
View File

@ -6443,6 +6443,7 @@ def Torch_AtenTypeAsOp : Torch_Op<"aten.type_as", [
}
def Torch_AtenUnbindIntOp : Torch_Op<"aten.unbind.int", [
Pure,
AllowsTypeRefinement,
ReadOnly
]> {

2
lib/Dialect/Torch/IR/TorchOps.cpp
View File

@ -1737,7 +1737,7 @@ void Aten__Getitem__TOp::getCanonicalizationPatterns(
auto indexAsTensorType = rewriter.getType<Torch::NonValueTensorType>(
llvm::makeArrayRef({1L}), op.getOperand(1).getType());
auto indexAsTensorMLIRType = rewriter.getType<Torch::NonValueTensorType>(
llvm::makeArrayRef({1L}), IntegerType::get(op->getContext(), 64, IntegerType::Signed));
ArrayRef<int64_t>(), IntegerType::get(op->getContext(), 64, IntegerType::Signed));
auto indexAsTensor = rewriter.create<Torch::AtenTensorIntOp>(
op->getLoc(),

2
python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py
View File

@ -477,7 +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::unbind.int : (Tensor, int) -> (Tensor[])", traits=["Pure"])
emit("aten::view : (Tensor, int[]) -> (Tensor)", has_folder=True)
emit("aten::_unsafe_view : (Tensor, int[]) -> (Tensor)")
emit("aten::where.self : (Tensor, Tensor, Tensor) -> (Tensor)")