Add Rsqrt

e2e_testing/torchscript/elementwise.py

@@ -445,3 +445,20 @@ class ElementwiseLog2Module(torch.nn.Module):
 @register_test_case(module_factory=lambda: ElementwiseLog2Module())
 def ElementwiseLog2Module_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
+
+class ElementwiseRsqrtModule(torch.nn.Module):
+    def __init__(self): 
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+    ])
+
+    def forward(self, a):
+        return torch.rsqrt(a)
+
+@register_test_case(module_factory=lambda: ElementwiseRsqrtModule())
+def ElementwiseRsqrtModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4))

include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td

@@ -878,6 +878,34 @@ def Torch_AtenLog2_Op : Torch_Op<"aten.log2_", [
   let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
 }
 
+def Torch_AtenRsqrtOp : Torch_Op<"aten.rsqrt", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::rsqrt : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
+}
+
+def Torch_AtenRsqrt_Op : Torch_Op<"aten.rsqrt_", [
+    IsTrailingUnderscoreInplaceVariant,
+    AllowsTypeRefinement
+  ]> {
+  let summary = "Generated op for `aten::rsqrt_ : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
+}
+
 def Torch_AtenMaximumOp : Torch_Op<"aten.maximum", [
     AllowsTypeRefinement,
     HasValueSemantics

lib/Conversion/TorchToLinalg/TorchToLinalg.cpp

@@ -1320,6 +1320,8 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
     return b.create<math::LogOp>(loc, payloadArgs[0]);
   if (isa<AtenSqrtOp>(op))
     return b.create<math::SqrtOp>(loc, payloadArgs[0]);
+  if (isa<AtenRsqrtOp>(op))
+    return b.create<math::RsqrtOp>(loc, payloadArgs[0]);
   if (isa<AtenLog2Op>(op))
     return b.create<math::Log2Op>(loc, payloadArgs[0]);
   if (isa<AtenSigmoidOp>(op)) {
@@ -1724,7 +1726,8 @@ struct ConvertElementwiseOp : ConversionPattern {
              AtenMulTensorOp, AtenDivTensorOp, AtenSubTensorOp,
              AtenLerpTensorOp, AtenSigmoidOp, AtenExpOp, AtenMinimumOp,
              AtenMaximumOp, AtenToDtypeOp, AtenClampOp, AtenRsubScalarOp,
-             AtenLogOp, AtenSqrtOp, AtenFloorOp, AtenPowTensorScalarOp, AtenLog2Op>(op))
+             AtenLogOp, AtenSqrtOp, AtenFloorOp, AtenPowTensorScalarOp,
+             AtenLog2Op, AtenRsqrtOp>(op))
       return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
 
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
@@ -2885,7 +2888,7 @@ public:
                         AtenLerpTensorOp, AtenSigmoidOp, AtenMinimumOp,
                         AtenMaximumOp, AtenToDtypeOp, AtenClampOp,
                         AtenRsubScalarOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
-                        AtenPowTensorScalarOp, AtenLog2Op>();
+                        AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp>();
     patterns.add<ConvertElementwiseOp>(typeConverter, context);
     target.addIllegalOp<AtenUnsqueezeOp>();
     patterns.add<ConvertAtenUnsqueezeOp>(typeConverter, context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -230,7 +230,7 @@ public:
             AtenToPrimDeviceOp, AtenCpuOp, AtenContiguousOp, AtenFill_ScalarOp,
             AtenDetachOp, AtenMaskedFill_ScalarOp, AtenCopy_Op, AtenIndexPut_Op,
             AtenCumsumOp, AtenLayerNormOp, AtenClampOp, AtenLogOp, AtenSqrtOp,
-            AtenFloorOp, AtenLog2Op, Aten_SoftmaxBackwardDataOp>(op)) {
+            AtenFloorOp, AtenLog2Op, Aten_SoftmaxBackwardDataOp, AtenRsqrtOp>(op)) {
       return getLatticeElement(op->getResult(0)).join(*operands[0]);
     }
 

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -466,6 +466,7 @@ def emit_aten_ops(torch_ir_dir: str, registry: Registry):
                 "aten::masked_fill.Scalar : (Tensor, Tensor, Scalar) -> (Tensor)",
                 "aten::clamp : (Tensor, Scalar?, Scalar?) -> (Tensor)",
                 "aten::log2 : (Tensor) -> (Tensor)",
+                "aten::rsqrt : (Tensor) -> (Tensor)",
         ]:
             emit_with_mutating_variants(key)
         # Elementwise tensor compute ops that don't have the standard mutating