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

Emit underscore version of aten.sqrt (#1072)

pull/1083/head snapshot-20220719.538
Kevin Kiningham 2022-07-18 23:57:47 -07:00 committed by GitHub
parent c73a39e40a
commit 21f905afbe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 46 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -2024,6 +2024,51 @@ def Torch_AtenLog2_Op : Torch_Op<"aten.log2_", [
}];
}
def Torch_AtenSqrtOp : Torch_Op<"aten.sqrt", [
AllowsTypeRefinement,
HasValueSemantics,
ReadOnly
]> {
let summary = "Generated op for `aten::sqrt : (Tensor) -> (Tensor)`";
let arguments = (ins
AnyTorchTensorType:$self
);
let results = (outs
AnyTorchTensorType:$result
);
let hasCustomAssemblyFormat = 1;
let extraClassDefinition = [{
ParseResult AtenSqrtOp::parse(OpAsmParser &parser, OperationState &result) {
return parseDefaultTorchOp(parser, result, 1, 1);
}
void AtenSqrtOp::print(OpAsmPrinter &printer) {
printDefaultTorchOp(printer, *this, 1, 1);
}
}];
}
def Torch_AtenSqrt_Op : Torch_Op<"aten.sqrt_", [
IsTrailingUnderscoreInplaceVariant,
AllowsTypeRefinement
]> {
let summary = "Generated op for `aten::sqrt_ : (Tensor) -> (Tensor)`";
let arguments = (ins
AnyTorchTensorType:$self
);
let results = (outs
AnyTorchTensorType:$result
);
let hasCustomAssemblyFormat = 1;
let extraClassDefinition = [{
ParseResult AtenSqrt_Op::parse(OpAsmParser &parser, OperationState &result) {
return parseDefaultTorchOp(parser, result, 1, 1);
}
void AtenSqrt_Op::print(OpAsmPrinter &printer) {
printDefaultTorchOp(printer, *this, 1, 1);
}
}];
}
def Torch_AtenRsqrtOp : Torch_Op<"aten.rsqrt", [
AllowsTypeRefinement,
HasValueSemantics,
@ -3764,29 +3809,6 @@ def Torch_Aten__And__TensorOp : Torch_Op<"aten.__and__.Tensor", [
}];
}
def Torch_AtenSqrtOp : Torch_Op<"aten.sqrt", [
AllowsTypeRefinement,
HasValueSemantics,
ReadOnly
]> {
let summary = "Generated op for `aten::sqrt : (Tensor) -> (Tensor)`";
let arguments = (ins
AnyTorchTensorType:$self
);
let results = (outs
AnyTorchTensorType:$result
);
let hasCustomAssemblyFormat = 1;
let extraClassDefinition = [{
ParseResult AtenSqrtOp::parse(OpAsmParser &parser, OperationState &result) {
return parseDefaultTorchOp(parser, result, 1, 1);
}
void AtenSqrtOp::print(OpAsmPrinter &printer) {
printDefaultTorchOp(printer, *this, 1, 1);
}
}];
}
def Torch_Aten_SoftmaxOp : Torch_Op<"aten._softmax", [
AllowsTypeRefinement,
HasValueSemantics,

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

@ -281,6 +281,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
"aten::clamp_min : (Tensor, Scalar) -> (Tensor)",
"aten::clamp_max : (Tensor, Scalar) -> (Tensor)",
"aten::log2 : (Tensor) -> (Tensor)",
"aten::sqrt : (Tensor) -> (Tensor)",
"aten::rsqrt : (Tensor) -> (Tensor)",
"aten::abs : (Tensor) -> (Tensor)",
"aten::reciprocal : (Tensor) -> (Tensor)",
@ -374,7 +375,6 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
emit("aten::logsumexp : (Tensor, int[], bool) -> (Tensor)")
emit("aten::mean.dim : (Tensor, int[], bool, int?) -> (Tensor)")
emit("aten::__and__.Tensor : (Tensor, Tensor) -> (Tensor)")
emit("aten::sqrt : (Tensor) -> (Tensor)")
emit("aten::_softmax : (Tensor, int, bool) -> (Tensor)")
emit("aten::mean : (Tensor, int?) -> (Tensor)")
emit("aten::std : (Tensor, bool) -> (Tensor)")