aten.abs and aten.reciprocal to linalg

e2e_testing/torchscript/elementwise.py

@@ -499,6 +499,43 @@ def ElementwiseRsqrtModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
 
 # ==============================================================================
+
+class ElementwiseAbsModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float32, True),
+    ])
+
+    def forward(self, a):
+        return torch.abs(a)
+
+@register_test_case(module_factory=lambda: ElementwiseAbsModule())
+def ElementwiseAbsModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4, 5, low=-1.0, high=1.0))
+
+# ==============================================================================
+
+class ElementwiseReciprocalModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    @export
+    @annotate_args([
+        None,
+        ([-1], torch.float32, True),
+    ])
+
+    def forward(self, a):
+        return torch.reciprocal(a)
+
+@register_test_case(module_factory=lambda: ElementwiseReciprocalModule())
+def ElementwiseReciprocalModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(4))
+
+# ==============================================================================
+
 class ElementwiseDivScalarModule(torch.nn.Module):
     def __init__(self):
         super().__init__()

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

@@ -936,6 +936,62 @@ def Torch_AtenRsqrt_Op : Torch_Op<"aten.rsqrt_", [
   let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
 }
 
+def Torch_AtenAbsOp : Torch_Op<"aten.abs", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::abs : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
+}
+
+def Torch_AtenAbs_Op : Torch_Op<"aten.abs_", [
+    IsTrailingUnderscoreInplaceVariant,
+    AllowsTypeRefinement
+  ]> {
+  let summary = "Generated op for `aten::abs_ : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
+}
+
+def Torch_AtenReciprocalOp : Torch_Op<"aten.reciprocal", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::reciprocal : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` type($self) `->` type($result)";
+}
+
+def Torch_AtenReciprocal_Op : Torch_Op<"aten.reciprocal_", [
+    IsTrailingUnderscoreInplaceVariant,
+    AllowsTypeRefinement
+  ]> {
+  let summary = "Generated op for `aten::reciprocal_ : (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

@@ -1396,6 +1396,8 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
     return b.create<math::RsqrtOp>(loc, payloadArgs[0]);
   if (isa<AtenLog2Op>(op))
     return b.create<math::Log2Op>(loc, payloadArgs[0]);
+  if (isa<AtenAbsOp>(op))
+    return b.create<math::AbsOp>(loc, payloadArgs[0]);
   if (isa<AtenSigmoidOp>(op)) {
     Type elementType = payloadArgs[0].getType();
     auto one = b.create<arith::ConstantOp>(loc, FloatAttr::get(elementType, 1));
@@ -1661,6 +1663,28 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
     Value other = convertScalarToDtype(b, loc, operands[1], dtype);
     return b.create<arith::DivFOp>(loc, self, other);
   }
+  if (auto reciprocal = dyn_cast<AtenReciprocalOp>(op)) {
+    if (!reciprocal.getType()
+             .cast<ValueTensorType>()
+             .getDtype()
+             .isa<mlir::FloatType>()) {
+      reciprocal.emitError("unimplemented: non-floating point dtype");
+      return nullptr;
+    }
+
+    Type elementType = payloadArgs[0].getType();
+    // assert(element != 0)
+    auto zero =
+        b.create<arith::ConstantOp>(loc, FloatAttr::get(elementType, 0.0));
+    auto pred = b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::ONE,
+                                        payloadArgs[0], zero);
+    b.create<AssertOp>(
+        loc, pred, b.getStringAttr("unimplemented: tensor with zero element"));
+
+    auto one =
+        b.create<arith::ConstantOp>(loc, FloatAttr::get(elementType, 1.0));
+    return b.create<arith::DivFOp>(loc, one, payloadArgs[0]);
+  }
 
   op->emitError("unimplemented lowering in "
                 "createLinalgPayloadCalculationForElementwiseOp");
@@ -1872,8 +1896,8 @@ struct ConvertElementwiseOp : ConversionPattern {
              AtenLerpTensorOp, AtenSigmoidOp, AtenExpOp, AtenMinimumOp,
              AtenMaximumOp, AtenToDtypeOp, AtenClampOp, AtenRsubScalarOp,
              AtenMulScalarOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
-             AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp, AtenDivScalarOp>(
-            op))
+             AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp, AtenDivScalarOp,
+             AtenAbsOp, AtenReciprocalOp>(op))
       return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
 
     if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
@@ -3055,7 +3079,8 @@ public:
         AtenMulTensorOp, AtenDivTensorOp, AtenSubTensorOp, AtenLerpTensorOp,
         AtenSigmoidOp, AtenMinimumOp, AtenMaximumOp, AtenToDtypeOp, AtenClampOp,
         AtenRsubScalarOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
-        AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp>();
+        AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp, AtenAbsOp,
+        AtenReciprocalOp>();
     patterns.add<ConvertElementwiseOp>(typeConverter, context);
     target.addIllegalOp<AtenUnsqueezeOp>();
     patterns.add<ConvertAtenUnsqueezeOp>(typeConverter, context);

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -232,8 +232,8 @@ public:
             AtenMaskedFill_ScalarOp, AtenCopy_Op, AtenIndexPut_Op, AtenCumsumOp,
             AtenLayerNormOp, AtenClampOp, AtenLogOp, AtenSqrtOp, AtenFloorOp,
             AtenLog2Op, Aten_SoftmaxBackwardDataOp, AtenRsqrtOp, AtenDropoutOp,
-            AtenTanhBackwardOp, Aten_LogSoftmaxBackwardDataOp, AtenAddIntOp>(
-            op)) {
+            AtenTanhBackwardOp, Aten_LogSoftmaxBackwardDataOp, AtenAddIntOp,
+            AtenAbsOp, AtenReciprocalOp>(op)) {
       return getLatticeElement(op->getResult(0)).join(*operands[0]);
     }
 

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

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