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

Add in-place variant of torch.scatter_add (#1836)

pull/1851/head snapshot-20230204.739
Zachary Cetinic 2023-02-03 12:54:28 -05:00 committed by GitHub
parent f58ba19448
commit 2c2009a13d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 26 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -7613,6 +7613,31 @@ def Torch_AtenScatterAddOp : Torch_Op<"aten.scatter_add", [
}];
}
def Torch_AtenScatterAdd_Op : Torch_Op<"aten.scatter_add_", [
IsTrailingUnderscoreInplaceVariant,
AllowsTypeRefinement
]> {
let summary = "Generated op for `aten::scatter_add_ : (Tensor, int, Tensor, Tensor) -> (Tensor)`";
let arguments = (ins
AnyTorchTensorType:$self,
Torch_IntType:$dim,
AnyTorchTensorType:$index,
AnyTorchTensorType:$src
);
let results = (outs
AnyTorchTensorType:$result
);
let hasCustomAssemblyFormat = 1;
let extraClassDefinition = [{
ParseResult AtenScatterAdd_Op::parse(OpAsmParser &parser, OperationState &result) {
return parseDefaultTorchOp(parser, result, 4, 1);
}
void AtenScatterAdd_Op::print(OpAsmPrinter &printer) {
printDefaultTorchOp(printer, *this, 4, 1);
}
}];
}
def Torch_AtenIntImplicitOp : Torch_Op<"aten.IntImplicit", [
AllowsTypeRefinement,
HasValueSemantics,

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

@ -509,7 +509,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
emit("aten::len.Tensor : (Tensor) -> (int)")
emit("aten::cpu : (Tensor) -> (Tensor)")
emit("aten::gather : (Tensor, int, Tensor, bool) -> (Tensor)")
emit("aten::scatter_add : (Tensor, int, Tensor, Tensor) -> (Tensor)")
emit_with_mutating_variants("aten::scatter_add : (Tensor, int, Tensor, Tensor) -> (Tensor)")
emit("aten::IntImplicit : (Tensor) -> (int)")
emit("aten::FloatImplicit : (Tensor) -> (float)")
emit("aten::tensor.float : (float, int?, Device?, bool) -> (Tensor)")